diff --git a/lib/NonlinearSolveHomotopyContinuation/Project.toml b/lib/NonlinearSolveHomotopyContinuation/Project.toml index 4b4d3a512..f7882aca4 100644 --- a/lib/NonlinearSolveHomotopyContinuation/Project.toml +++ b/lib/NonlinearSolveHomotopyContinuation/Project.toml @@ -23,6 +23,7 @@ CommonSolve = "0.2.4" ConcreteStructs = "0.2.3" DifferentiationInterface = "0.6.27" DocStringExtensions = "0.9.3" +Enzyme = "0.13" HomotopyContinuation = "2.12.0" LinearAlgebra = "1.10" NonlinearSolve = "4" @@ -35,8 +36,9 @@ julia = "1.10" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Aqua", "Test", "NonlinearSolve"] +test = ["Aqua", "Test", "NonlinearSolve", "Enzyme"] diff --git a/lib/NonlinearSolveHomotopyContinuation/src/NonlinearSolveHomotopyContinuation.jl b/lib/NonlinearSolveHomotopyContinuation/src/NonlinearSolveHomotopyContinuation.jl index 5f68e55a8..26639e4f0 100644 --- a/lib/NonlinearSolveHomotopyContinuation/src/NonlinearSolveHomotopyContinuation.jl +++ b/lib/NonlinearSolveHomotopyContinuation/src/NonlinearSolveHomotopyContinuation.jl @@ -60,7 +60,12 @@ end HomotopyContinuationJL(; kwargs...) = HomotopyContinuationJL{false}(; kwargs...) +function HomotopyContinuationJL(alg::HomotopyContinuationJL{R}; kwargs...) where {R} + HomotopyContinuationJL{R}(; autodiff = alg.autodiff, alg.kwargs..., kwargs...) +end + include("interface_types.jl") +include("jacobian_handling.jl") include("solve.jl") end diff --git a/lib/NonlinearSolveHomotopyContinuation/src/interface_types.jl b/lib/NonlinearSolveHomotopyContinuation/src/interface_types.jl index 63a64ec36..70fcace78 100644 --- a/lib/NonlinearSolveHomotopyContinuation/src/interface_types.jl +++ b/lib/NonlinearSolveHomotopyContinuation/src/interface_types.jl @@ -4,47 +4,6 @@ struct Inplace <: HomotopySystemVariant end struct OutOfPlace <: HomotopySystemVariant end struct Scalar <: HomotopySystemVariant end -""" - $(TYPEDEF) - -A simple struct that wraps a polynomial function which takes complex input and returns -complex output in a form that supports automatic differentiation. If the wrapped -function if ``f: \\mathbb{C}^n \\rightarrow \\mathbb{C}^n`` then it is assumed -the input arrays are real-valued and have length ``2n``. They are `reinterpret`ed -into complex arrays and passed into the function. This struct has an in-place signature -regardless of the signature of ``f``. -""" -@concrete struct ComplexJacobianWrapper{variant <: HomotopySystemVariant} - f -end - -function (cjw::ComplexJacobianWrapper{Inplace})( - u::AbstractVector{T}, x::AbstractVector{T}, p) where {T} - x = reinterpret(Complex{T}, x) - u = reinterpret(Complex{T}, u) - cjw.f(u, x, p) - u = parent(u) - return u -end - -function (cjw::ComplexJacobianWrapper{OutOfPlace})( - u::AbstractVector{T}, x::AbstractVector{T}, p) where {T} - x = reinterpret(Complex{T}, x) - u_tmp = cjw.f(x, p) - u_tmp = reinterpret(T, u_tmp) - copyto!(u, u_tmp) - return u -end - -function (cjw::ComplexJacobianWrapper{Scalar})( - u::AbstractVector{T}, x::AbstractVector{T}, p) where {T} - x = reinterpret(Complex{T}, x) - u_tmp = cjw.f(x[1], p) - u[1] = real(u_tmp) - u[2] = imag(u_tmp) - return u -end - """ $(TYPEDEF) @@ -62,8 +21,10 @@ $(FIELDS) """ f """ - The jacobian function, if provided to the `NonlinearProblem` being solved. Otherwise, - a `ComplexJacobianWrapper` wrapping `f` used for automatic differentiation. + A function which calculates both the polynomial and the jacobian. Must be a function + of the form `f(u, U, x, p)` where `x` is the current unknowns and `p` is the parameter + object, writing the value of the polynomial to `u` and the jacobian to `U`. Must be able + to handle complex `x`. """ jac """ @@ -71,14 +32,6 @@ $(FIELDS) """ p """ - The ADType for automatic differentiation. - """ - autodiff - """ - The result from `DifferentiationInterface.prepare_jacobian`. - """ - prep - """ HomotopyContinuation.jl's symbolic variables for the system. """ vars @@ -86,10 +39,6 @@ $(FIELDS) The `TaylorDiff.TaylorScalar` objects used to compute the taylor series of `f`. """ taylorvars - """ - Preallocated intermediate buffers used for calculating the jacobian. - """ - jacobian_buffers end Base.size(sys::HomotopySystemWrapper) = (length(sys.vars), length(sys.vars)) @@ -112,54 +61,11 @@ function HC.ModelKit.evaluate!(u, sys::HomotopySystemWrapper{Scalar}, x, p = not end function HC.ModelKit.evaluate_and_jacobian!( - u, U, sys::HomotopySystemWrapper{Inplace}, x, p = nothing) - p = sys.p - sys.f(u, x, p) - sys.jac(U, x, p) - return u, U -end - -function HC.ModelKit.evaluate_and_jacobian!( - u, U, sys::HomotopySystemWrapper{OutOfPlace}, x, p = nothing) - p = sys.p - u_tmp = sys.f(x, p) - copyto!(u, u_tmp) - j_tmp = sys.jac(x, p) - copyto!(U, j_tmp) + u, U, sys::HomotopySystemWrapper, x, p = nothing) + sys.jac(u, U, x, sys.p) return u, U end -function HC.ModelKit.evaluate_and_jacobian!( - u, U, sys::HomotopySystemWrapper{Scalar}, x, p = nothing) - p = sys.p - u[1] = sys.f(x[1], p) - U[1] = sys.jac(x[1], p) - return u, U -end - -for V in (Inplace, OutOfPlace, Scalar) - @eval function HC.ModelKit.evaluate_and_jacobian!( - u, U, sys::HomotopySystemWrapper{$V, F, J}, x, - p = nothing) where {F, J <: ComplexJacobianWrapper} - p = sys.p - U_tmp = sys.jacobian_buffers - x = reinterpret(Float64, x) - u = reinterpret(Float64, u) - DI.value_and_jacobian!(sys.jac, u, U_tmp, sys.prep, sys.autodiff, x, DI.Constant(p)) - U = reinterpret(Float64, U) - @inbounds for j in axes(U, 2) - jj = 2j - 1 - for i in axes(U, 1) - U[i, j] = U_tmp[i, jj] - end - end - u = parent(u) - U = parent(U) - - return u, U - end -end - function update_taylorvars_from_taylorvector!( vars, x::HC.ModelKit.TaylorVector) for i in eachindex(x) diff --git a/lib/NonlinearSolveHomotopyContinuation/src/jacobian_handling.jl b/lib/NonlinearSolveHomotopyContinuation/src/jacobian_handling.jl new file mode 100644 index 000000000..730a4f3ee --- /dev/null +++ b/lib/NonlinearSolveHomotopyContinuation/src/jacobian_handling.jl @@ -0,0 +1,235 @@ +""" + $(TYPEDEF) + +A simple struct that wraps a polynomial function which takes complex input and returns +complex output in a form that supports automatic differentiation. If the wrapped +function if ``f: \\mathbb{C}^n \\rightarrow \\mathbb{C}^n`` then it is assumed +the input arrays are real-valued and have length ``2n``. They are `reinterpret`ed +into complex arrays and passed into the function. This struct matches the signature +of ``f``, except if ``f`` is scalar (in which case it acts like an in-place function). +""" +@concrete struct ComplexJacobianWrapper{variant <: HomotopySystemVariant} + f +end + +function (cjw::ComplexJacobianWrapper{Inplace})( + u::AbstractVector{T}, x::AbstractVector{T}, p) where {T} + x = reinterpret(Complex{T}, x) + u = reinterpret(Complex{T}, u) + cjw.f(u, x, p) + u = parent(u) + return u +end + +function (cjw::ComplexJacobianWrapper{OutOfPlace})(x::AbstractVector{T}, p) where {T} + x = reinterpret(Complex{T}, x) + u_tmp = cjw.f(x, p) + u_tmp = reinterpret(T, u_tmp) + return u_tmp +end + +function (cjw::ComplexJacobianWrapper{Scalar})( + u::AbstractVector{T}, x::AbstractVector{T}, p) where {T} + x = reinterpret(Complex{T}, x) + u_tmp = cjw.f(x[1], p) + u[1] = real(u_tmp) + u[2] = imag(u_tmp) + return u +end + +""" + $(TYPEDEF) + +A callable struct which calculates complex jacobians using `ComplexJacobianWrapper` and +DifferentiationInterface.jl. Follows the signature required by `HomotopySystemWrapper`. + +# Fields + +$(TYPEDFIELDS) +""" +@concrete struct ComplexDIJacobian{variant} + """ + The `ComplexJacobianWrapper`. + """ + f + """ + Preparation object from DifferentiationInterface.jl. + """ + prep + """ + Temporary buffer(s) required for improved AD performance. + """ + buffers + """ + Autodiff algorithm as an ADType. + """ + autodiff +end + +function (f::ComplexDIJacobian{Inplace})(u, U, x, p) + U_tmp = f.buffers + DI.value_and_jacobian!(f.f, reinterpret(Float64, u), reinterpret(Float64, U_tmp), f.prep, f.autodiff, reinterpret(Float64, x), DI.Constant(p)) + U = reinterpret(Float64, U) + @inbounds for j in axes(U, 2) + jj = 2j - 1 + for i in axes(U, 1) + U[i, j] = U_tmp[i, jj] + end + end +end + +function (f::ComplexDIJacobian{OutOfPlace})(u, U, x, p) + U_tmp = f.buffers + u_tmp, _ = DI.value_and_jacobian!(f.f, U_tmp, f.prep, f.autodiff, reinterpret(Float64, x), DI.Constant(p)) + copyto!(u, reinterpret(ComplexF64, u_tmp)) + U = reinterpret(Float64, U) + @inbounds for j in axes(U, 2) + jj = 2j - 1 + for i in axes(U, 1) + U[i, j] = U_tmp[i, jj] + end + end +end + +function (f::ComplexDIJacobian{Scalar})(u, U, x, p) + U_tmp = f.buffers + DI.value_and_jacobian!(f.f, reinterpret(Float64, u), U_tmp, f.prep, f.autodiff, reinterpret(Float64, x), DI.Constant(p)) + U[1] = U_tmp[1, 1] + im * U_tmp[2, 1] + return nothing +end + +""" + $(TYPEDSIGNATURES) + +Construct a jacobian function for the given function `f`, using `autodiff` as the AD +algorithm, `variant` being the `HomotopySystemVariant` of `f`, `u0` the state vector +and `p` the parameter object. + +The returned function must have the signature required by `HomotopySystemWrapper`. +""" +function construct_jacobian(f, autodiff, variant, u0, p) + if variant == Scalar + tmp = reinterpret(Float64, Vector{ComplexF64}(undef, 1)) + else + tmp = reinterpret(Float64, Vector{ComplexF64}(undef, length(u0))) + end + f = ComplexJacobianWrapper{variant}(f) + if variant == OutOfPlace + prep = DI.prepare_jacobian(f, autodiff, tmp, DI.Constant(p)) + else + prep = DI.prepare_jacobian(f, tmp, autodiff, copy(tmp), DI.Constant(p)) + end + + if variant == Scalar + U_tmp = Matrix{Float64}(undef, 2, 2) + else + U_tmp = Matrix{Float64}(undef, 2length(u0), 2length(u0)) + end + + return ComplexDIJacobian{variant}(f, prep, U_tmp, autodiff) +end + +""" + $(TYPEDEF) + +Efficient version of `ComplexDIJacobian` which directly computes complex jacobians since +Enzyme supports complex differentiation. + +# Fields + +$(TYPEDFIELDS) +""" +@concrete struct EnzymeJacobian{variant} + """ + The function to calculate the jacobian of. + """ + f + """ + Preparation object from DifferentiationInterface. + """ + prep + """ + AD algorithm specified as an ADType. + """ + autodiff +end + +function (f::EnzymeJacobian{Inplace})(u, U, x, p) + DI.value_and_jacobian!(f.f, u, U, f.prep, f.autodiff, x, DI.Constant(p)) + return nothing +end + +function (f::EnzymeJacobian{OutOfPlace})(u, U, x, p) + u_tmp, _ = DI.value_and_jacobian!(f.f, U, f.prep, f.autodiff, x, DI.Constant(p)) + copyto!(u, u_tmp) + return nothing +end + +function (f::EnzymeJacobian{Scalar})(u, U, x, p) + u_tmp, der_tmp = DI.value_and_derivative(f.f, f.prep, f.autodiff, x[1], DI.Constant(p)) + u[1] = u_tmp + U[1] = der_tmp + return nothing +end + +""" + $(TYPEDSIGNATURES) + +Construct an `EnzymeJacobian` function. +""" +function construct_jacobian(f, autodiff::AutoEnzyme, variant, u0, p) + if variant == Scalar + prep = DI.prepare_derivative(f, autodiff, u0, DI.Constant(p)) + else + tmp = Vector{ComplexF64}(undef, length(u0)) + if variant == Inplace + prep = DI.prepare_jacobian(f, tmp, autodiff, copy(tmp), DI.Constant(p)) + else + prep = DI.prepare_jacobian(f, autodiff, tmp, DI.Constant(p)) + end + end + return EnzymeJacobian{variant}(f, prep, autodiff) +end + + +""" + $(TYPEDEF) + +Jacobian function following the signature required by `HomotopySystemWrapper` using +a user-provided jacobian. + +# Fields + +$(TYPEDFIELDS) +""" +@concrete struct ExplicitJacobian{variant} + """ + The RHS function. + """ + f + """ + The jacobian function. + """ + jac +end + +function (f::ExplicitJacobian{Inplace})(u, U, x, p) + f.f(u, x, p) + f.jac(U, x, p) + return nothing +end + +function (f::ExplicitJacobian{OutOfPlace})(u, U, x, p) + u_tmp = f.f(x, p) + copyto!(u, u_tmp) + j_tmp = f.jac(x, p) + copyto!(U, j_tmp) + return nothing +end + +function (f::ExplicitJacobian{Scalar})(u, U, x, p) + u[1] = f.f(x[1], p) + U[1] = f.jac(x[1], p) + return nothing +end + diff --git a/lib/NonlinearSolveHomotopyContinuation/src/solve.jl b/lib/NonlinearSolveHomotopyContinuation/src/solve.jl index 2db789df9..721ff2704 100644 --- a/lib/NonlinearSolveHomotopyContinuation/src/solve.jl +++ b/lib/NonlinearSolveHomotopyContinuation/src/solve.jl @@ -23,17 +23,10 @@ function homotopy_continuation_preprocessing( # jacobian handling if SciMLBase.has_jac(f.f) # use if present - prep = nothing - jac = f.jac + jac = ExplicitJacobian{variant}(f.f.f, f.f.jac) else # prepare a DI jacobian if not - jac = ComplexJacobianWrapper{variant}(f.f.f) - tmp = if isscalar - Vector{Float64}(undef, 2) - else - similar(u0, Float64, 2length(u0)) - end - prep = DI.prepare_jacobian(jac, tmp, alg.autodiff, copy(tmp), DI.Constant(p)) + jac = construct_jacobian(f.f.f, alg.autodiff, variant, u0, p) end # variables for HC to use @@ -54,15 +47,9 @@ function homotopy_continuation_preprocessing( [TaylorScalar(ntuple(Returns(0.0), 4)) for _ in 1:2length(u0)] end - jacobian_buffers = if isscalar - Matrix{Float64}(undef, 2, 2) - else - similar(u0, Float64, 2length(u0), 2length(u0)) - end - # HC-compatible system hcsys = HomotopySystemWrapper{variant}( - f.f.f, jac, p, alg.autodiff, prep, vars, taylorvars, jacobian_buffers) + f.f.f, jac, p, vars, taylorvars) return f, hcsys end diff --git a/lib/NonlinearSolveHomotopyContinuation/test/allroots.jl b/lib/NonlinearSolveHomotopyContinuation/test/allroots.jl index c0d025e49..30955029a 100644 --- a/lib/NonlinearSolveHomotopyContinuation/test/allroots.jl +++ b/lib/NonlinearSolveHomotopyContinuation/test/allroots.jl @@ -1,6 +1,8 @@ using NonlinearSolve using NonlinearSolveHomotopyContinuation using SciMLBase: NonlinearSolution +using ADTypes +using Enzyme alg = HomotopyContinuationJL{true}(; threading = false) @@ -11,11 +13,19 @@ alg = HomotopyContinuationJL{true}(; threading = false) jac = function (u, p) return 2u - p[1] end - @testset "`NonlinearProblem` - $name" for (jac, name) in [ - (nothing, "no jac"), (jac, "jac")] + @testset "`NonlinearProblem` - $name" for (jac_or_autodiff, name) in [ + (AutoForwardDiff(), "no jac - forwarddiff"), (AutoEnzyme(), "no jac - enzyme"), (jac, "jac")] + if jac_or_autodiff isa Function + jac = jac_or_autodiff + autodiff = nothing + else + jac = nothing + autodiff = jac_or_autodiff + end fn = NonlinearFunction(rhs; jac) prob = NonlinearProblem(fn, 1.0, [5.0, 6.0]) - sol = solve(prob, alg) + _alg = HomotopyContinuationJL(alg; autodiff) + sol = solve(prob, _alg) @test sol isa EnsembleSolution @test sol.converged @@ -84,25 +94,33 @@ f = function (u, p) [u[1] * u[1] - p[1] * u[2] + u[2]^3 + 1, u[2]^3 + 2 * p[2] * u[1] * u[2] + u[2]] end -jac! = function (j, u, p) +fjac! = function (j, u, p) j[1, 1] = 2u[1] j[1, 2] = -p[1] + 3 * u[2]^2 j[2, 1] = 2 * p[2] * u[2] j[2, 2] = 3 * u[2]^2 + 2 * p[2] * u[1] + 1 end -jac = function (u, p) +fjac = function (u, p) [2u[1] -p[1]+3 * u[2]^2; 2*p[2]*u[2] 3*u[2]^2+2*p[2]*u[1]+1] end -@testset "vector u - $name" for (rhs, jac, name) in [ - (f, nothing, "oop"), (f, jac, "oop + jac"), - (f!, nothing, "iip"), (f!, jac!, "iip + jac")] +@testset "vector u - $name" for (rhs, jac_or_autodiff, name) in [ + (f, AutoForwardDiff(), "oop + forwarddiff"), (f, AutoEnzyme(), "oop + enzyme"), (f, fjac, "oop + jac"), + (f!, AutoForwardDiff(), "iip + forwarddiff"), (f!, AutoEnzyme(), "iip + enzyme"), (f!, fjac!, "iip + jac")] sol = nothing + if jac_or_autodiff isa Function + jac = jac_or_autodiff + autodiff = nothing + else + jac = nothing + autodiff = jac_or_autodiff + end + _alg = HomotopyContinuationJL(alg; autodiff) @testset "`NonlinearProblem`" begin fn = NonlinearFunction(rhs; jac) prob = NonlinearProblem(fn, [1.0, 2.0], [2.0, 3.0]) - sol = solve(prob, alg) + sol = solve(prob, _alg) @test sol isa EnsembleSolution @test sol.converged for nlsol in sol.u @@ -112,7 +130,7 @@ end @testset "no real solutions" begin _prob = remake(prob; p = zeros(2)) - _sol = solve(_prob, alg) + _sol = solve(_prob, _alg) @test !_sol.converged @test length(_sol) == 1 @test !SciMLBase.successful_retcode(_sol.u[1]) @@ -132,7 +150,7 @@ end nlfn = NonlinearFunction(rhs; jac) fn = HomotopyNonlinearFunction(nlfn; denominator, polynomialize, unpolynomialize) prob = NonlinearProblem(fn, [1.0, 2.0], [2.0, 3.0, 4.0, 5.0]) - sol2 = solve(prob, alg) + sol2 = solve(prob, _alg) @test sol2 isa EnsembleSolution @test sol2.converged @test length(sol.u) == length(sol2.u) @@ -144,7 +162,7 @@ end @testset "some invalid solutions" begin prob3 = remake(prob; p = [2.0, 3.0, polynomialize(sol2.u[1].u, prob.p)...]) - sol3 = solve(prob3, alg) + sol3 = solve(prob3, _alg) @test length(sol3.u) == length(sol2.u) - 1 end end diff --git a/lib/NonlinearSolveHomotopyContinuation/test/single_root.jl b/lib/NonlinearSolveHomotopyContinuation/test/single_root.jl index 049578b37..2da68e1ef 100644 --- a/lib/NonlinearSolveHomotopyContinuation/test/single_root.jl +++ b/lib/NonlinearSolveHomotopyContinuation/test/single_root.jl @@ -11,11 +11,19 @@ alg = HomotopyContinuationJL{false}(; threading = false) jac = function (u, p) return 2u - (p + 3) end - @testset "`NonlinearProblem` - $name" for (jac, name) in [ - (nothing, "no jac"), (jac, "jac")] + @testset "`NonlinearProblem` - $name" for (jac_or_autodiff, name) in [ + (AutoForwardDiff(), "no jac - forwarddiff"), (AutoEnzyme(), "no jac - enzyme"), (jac, "jac")] + if jac_or_autodiff isa Function + jac = jac_or_autodiff + autodiff = nothing + else + jac = nothing + autodiff = jac_or_autodiff + end fn = NonlinearFunction(rhs; jac) prob = NonlinearProblem(fn, 1.0, 2.0) - sol = solve(prob, alg) + _alg = HomotopyContinuationJL(alg; autodiff) + sol = solve(prob, _alg) @test sol isa NonlinearSolution @test sol.u≈2.0 atol=1e-10 @@ -86,19 +94,27 @@ jac = function (u, p) 2*p[2]*u[2] 3*u[2]^2+2*p[2]*u[1]+1] end -@testset "vector u - $name" for (rhs, jac, name) in [ - (f, nothing, "oop"), (f, jac, "oop + jac"), - (f!, nothing, "iip"), (f!, jac!, "iip + jac")] +@testset "vector u - $name" for (rhs, jac_or_autodiff, name) in [ + (f, AutoForwardDiff(), "oop + forwarddiff"), (f, AutoEnzyme(), "oop + enzyme"), (f, jac, "oop + jac"), + (f!, AutoForwardDiff(), "iip + forwarddiff"), (f!, AutoEnzyme(), "iip + enzyme"), (f!, jac!, "iip + jac")] + if jac_or_autodiff isa Function + jac = jac_or_autodiff + autodiff = nothing + else + jac = nothing + autodiff = jac_or_autodiff + end + _alg = HomotopyContinuationJL(alg; autodiff) @testset "`NonlinearProblem`" begin fn = NonlinearFunction(rhs; jac) prob = NonlinearProblem(fn, [1.0, 2.0], [2.0, 3.0]) - sol = solve(prob, alg) + sol = solve(prob, _alg) @test SciMLBase.successful_retcode(sol) @test f(sol.u, prob.p)≈[0.0, 0.0] atol=1e-10 @testset "no real solutions" begin prob2 = remake(prob; p = zeros(2)) - sol2 = solve(prob2, alg) + sol2 = solve(prob2, _alg) @test !SciMLBase.successful_retcode(sol2) end end @@ -116,13 +132,13 @@ end nlfn = NonlinearFunction(rhs; jac) fn = HomotopyNonlinearFunction(nlfn; denominator, polynomialize, unpolynomialize) prob = NonlinearProblem(fn, [1.0, 1.0], [2.0, 3.0, 4.0, 5.0]) - sol = solve(prob, alg) + sol = solve(prob, _alg) @test SciMLBase.successful_retcode(sol) @test f(polynomialize(sol.u, prob.p), prob.p)≈zeros(2) atol=1e-10 @testset "some invalid solutions" begin prob2 = remake(prob; p = [2.0, 3.0, polynomialize(sol.u, prob.p)...]) - sol2 = solve(prob2, alg) + sol2 = solve(prob2, _alg) @test !SciMLBase.successful_retcode(sol2) end end