Add implementation for NLLS Forward Mode

Author: avik-pal

Project.toml

@@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -32,26 +33,41 @@ SimpleNonlinearSolveZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "0.2.6"
+AllocCheck = "0.1.1"
 ArrayInterface = "7.7"
+Aqua = "0.8"
+CUDA = "5.2"
 ChainRulesCore = "1.21"
 ConcreteStructs = "0.2.3"
 DiffEqBase = "6.146"
+DiffResults = "1.1"
 FastClosures = "0.3"
 FiniteDiff = "2.22"
 ForwardDiff = "0.10.36"
 LinearAlgebra = "1.10"
+LinearSolve = "2.25"
 MaybeInplace = "0.1.1"
+NonlinearProblemLibrary = "0.1.2"
+Pkg = "1.10"
+PolyesterForwardDiff = "0.1.1"
 PrecompileTools = "1.2"
+Random = "1.10"
+ReTestItems = "1.23"
 Reexport = "1.2"
 SciMLBase = "2.23"
+SciMLSensitivity = "7.56"
 StaticArrays = "1.9"
 StaticArraysCore = "1.4.2"
+Test = "1.10"
+Zygote = "0.6.69"
 julia = "1.10"
 
 [extras]
 AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
@@ -67,4 +83,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["AllocCheck", "DiffEqBase", "ForwardDiff", "LinearAlgebra", "LinearSolve", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SciMLSensitivity", "StaticArrays", "Zygote", "CUDA", "PolyesterForwardDiff", "Reexport", "Test"]
+test = ["Aqua", "AllocCheck", "DiffEqBase", "ForwardDiff", "LinearAlgebra", "LinearSolve", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SciMLSensitivity", "StaticArrays", "Zygote", "CUDA", "PolyesterForwardDiff", "Reexport", "Test", "FiniteDiff"]

ext/SimpleNonlinearSolveZygoteExt.jl

@@ -1,7 +1,12 @@
 module SimpleNonlinearSolveZygoteExt
 
-import SimpleNonlinearSolve
+import SimpleNonlinearSolve, Zygote
 
 SimpleNonlinearSolve.__is_extension_loaded(::Val{:Zygote}) = true
 
+function SimpleNonlinearSolve.__zygote_compute_nlls_vjp(f::F, u, p) where {F}
+    y, pb = Zygote.pullback(Base.Fix2(f, p), u)
+    return 2 .* only(pb(y))
+end
+
 end

src/SimpleNonlinearSolve.jl

@@ -11,6 +11,7 @@ import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidat
                        AbstractSafeBestNonlinearTerminationMode,
                        NonlinearSafeTerminationReturnCode, get_termination_mode,
                        NONLINEARSOLVE_DEFAULT_NORM
+    import DiffResults
     import ForwardDiff: Dual
     import MaybeInplace: @bb, setindex_trait, CanSetindex, CannotSetindex
     import SciMLBase: AbstractNonlinearAlgorithm, build_solution, isinplace, _unwrap_val

src/ad.jl

@@ -5,8 +5,17 @@ function SciMLBase.solve(
     sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...)
     dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p)
     return SciMLBase.build_solution(
-        prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats,
-        sol.original)
+        prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original)
+end
+
+function SciMLBase.solve(
+        prob::NonlinearLeastSquaresProblem{<:AbstractArray,
+            iip, <:Union{<:AbstractArray{<:Dual{T, V, P}}}},
+        alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...) where {T, V, P, iip}
+    sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...)
+    dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p)
+    return SciMLBase.build_solution(
+        prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original)
 end
 
 for algType in (Bisection, Brent, Alefeld, Falsi, ITP, Ridder)
@@ -56,6 +65,79 @@ function __nlsolve_ad(
     return sol, partials
 end
 
+function __nlsolve_ad(prob::NonlinearLeastSquaresProblem, alg, args...; kwargs...)
+    p = value(prob.p)
+    u0 = value(prob.u0)
+    newprob = NonlinearLeastSquaresProblem(prob.f, u0, p; prob.kwargs...)
+
+    sol = solve(newprob, alg, args...; kwargs...)
+
+    uu = sol.u
+
+    if !SciMLBase.has_jac(prob.f)
+        if isinplace(prob)
+            _F = @closure (du, u, p) -> begin
+                resid = similar(du, length(sol.resid))
+                res = DiffResults.DiffResult(
+                    resid, similar(du, length(sol.resid), length(u)))
+                _f = @closure (du, u) -> prob.f(du, u, p)
+                ForwardDiff.jacobian!(res, _f, resid, u)
+                mul!(reshape(du, 1, :), vec(DiffResults.value(res))',
+                    DiffResults.jacobian(res), 2, false)
+                return nothing
+            end
+        else
+            # For small problems, nesting ForwardDiff is actually quite fast
+            if __is_extension_loaded(Val(:Zygote)) && (length(uu) + length(sol.resid) ≥ 50)
+                _F = @closure (u, p) -> __zygote_compute_nlls_vjp(prob.f, u, p)
+            else
+                _F = @closure (u, p) -> begin
+                    T = promote_type(eltype(u), eltype(p))
+                    res = DiffResults.DiffResult(
+                        similar(u, T, size(sol.resid)), similar(
+                            u, T, length(sol.resid), length(u)))
+                    ForwardDiff.jacobian!(res, Base.Fix2(prob.f, p), u)
+                    return reshape(
+                        2 .* vec(DiffResults.value(res))' * DiffResults.jacobian(res),
+                        size(u))
+                end
+            end
+        end
+    else
+        if isinplace(prob)
+            _F = @closure (du, u, p) -> begin
+                J = similar(du, length(sol.resid), length(u))
+                prob.jac(J, u, p)
+                resid = similar(du, length(sol.resid))
+                prob.f(resid, u, p)
+                mul!(reshape(du, 1, :), vec(resid)', J, 2, false)
+                return nothing
+            end
+        else
+            _F = @closure (u, p) -> begin
+                return reshape(2 .* vec(prob.f(u, p))' * prob.jac(u, p), size(u))
+            end
+        end
+    end
+
+    f_p = __nlsolve_∂f_∂p(prob, _F, uu, p)
+    f_x = __nlsolve_∂f_∂u(prob, _F, uu, p)
+
+    z_arr = -f_x \ f_p
+
+    pp = prob.p
+    sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z)
+    if uu isa Number
+        partials = sum(sumfun, zip(z_arr, pp))
+    elseif p isa Number
+        partials = sumfun((z_arr, pp))
+    else
+        partials = sum(sumfun, zip(eachcol(z_arr), pp))
+    end
+
+    return sol, partials
+end
+
 @inline function __nlsolve_∂f_∂p(prob, f::F, u, p) where {F}
     if isinplace(prob)
         __f = p -> begin

src/utils.jl

@@ -388,3 +388,6 @@ function __get_tolerance(x::Union{SArray, Number}, ::Nothing, ::Type{T}) where {
     η = real(oneunit(T)) * (eps(real(one(T))))^(real(T)(0.8))
     return T(η)
 end
+
+# Extension
+function __zygote_compute_nlls_vjp end

test/core/aqua_tests.jl

@@ -0,0 +1,9 @@
+@testitem "Aqua" begin
+    using Aqua
+
+    Aqua.test_all(SimpleNonlinearSolve; piracies = false, ambiguities = false)
+    Aqua.test_piracies(SimpleNonlinearSolve;
+        treat_as_own = [
+            NonlinearProblem, NonlinearLeastSquaresProblem, IntervalNonlinearProblem])
+    Aqua.test_ambiguities(SimpleNonlinearSolve; recursive = false)
+end

test/core/forward_ad_tests.jl

@@ -1,4 +1,4 @@
-@testsetup module ForwardADTesting
+@testsetup module ForwardADRootfindingTesting
 using Reexport
 @reexport using ForwardDiff, SimpleNonlinearSolve, StaticArrays, LinearAlgebra
 import SimpleNonlinearSolve: AbstractSimpleNonlinearSolveAlgorithm
@@ -40,7 +40,7 @@ __compatible(::SimpleHalley, ::Val{:iip}) = false
 export test_f, test_f!, jacobian_f, solve_with, __compatible
 end
 
-@testitem "ForwardDiff.jl Integration" setup=[ForwardADTesting] begin
+@testitem "ForwardDiff.jl Integration: Rootfinding" setup=[ForwardADRootfindingTesting] begin
     @testset "$(nameof(typeof(alg)))" for alg in (SimpleNewtonRaphson(),
         SimpleTrustRegion(), SimpleTrustRegion(; nlsolve_update_rule = Val(true)),
         SimpleHalley(), SimpleBroyden(), SimpleKlement(), SimpleDFSane())

test/core/rootfind_tests.jl

@@ -134,31 +134,33 @@ end
 end
 
 @testitem "Allocation Checks" setup=[RootfindingTesting] begin
-    @testset "$(nameof(typeof(alg)))" for alg in (SimpleNewtonRaphson(),
-        SimpleHalley(), SimpleBroyden(), SimpleKlement(), SimpleLimitedMemoryBroyden(),
-        SimpleTrustRegion(), SimpleTrustRegion(; nlsolve_update_rule = Val(true)),
-        SimpleDFSane(), SimpleBroyden(; linesearch = Val(true)),
-        SimpleLimitedMemoryBroyden(; linesearch = Val(true)))
-        @check_allocs nlsolve(prob, alg) = SciMLBase.solve(prob, alg; abstol = 1e-9)
-
-        nlprob_scalar = NonlinearProblem{false}(quadratic_f, 1.0, 2.0)
-        nlprob_sa = NonlinearProblem{false}(quadratic_f, @SVector[1.0, 1.0], 2.0)
-
-        try
-            nlsolve(nlprob_scalar, alg)
-            @test true
-        catch e
-            @error e
-            @test false
-        end
+    if Sys.islinux()  # Very slow on other OS
+        @testset "$(nameof(typeof(alg)))" for alg in (SimpleNewtonRaphson(),
+            SimpleHalley(), SimpleBroyden(), SimpleKlement(), SimpleLimitedMemoryBroyden(),
+            SimpleTrustRegion(), SimpleTrustRegion(; nlsolve_update_rule = Val(true)),
+            SimpleDFSane(), SimpleBroyden(; linesearch = Val(true)),
+            SimpleLimitedMemoryBroyden(; linesearch = Val(true)))
+            @check_allocs nlsolve(prob, alg) = SciMLBase.solve(prob, alg; abstol = 1e-9)
+
+            nlprob_scalar = NonlinearProblem{false}(quadratic_f, 1.0, 2.0)
+            nlprob_sa = NonlinearProblem{false}(quadratic_f, @SVector[1.0, 1.0], 2.0)
+
+            try
+                nlsolve(nlprob_scalar, alg)
+                @test true
+            catch e
+                @error e
+                @test false
+            end
 
-        # ForwardDiff allocates for hessian since we don't propagate the chunksize
-        try
-            nlsolve(nlprob_sa, alg)
-            @test true
-        catch e
-            @error e
-            @test false broken=(alg isa SimpleHalley)
+            # ForwardDiff allocates for hessian since we don't propagate the chunksize
+            try
+                nlsolve(nlprob_sa, alg)
+                @test true
+            catch e
+                @error e
+                @test false broken=(alg isa SimpleHalley)
+            end
         end
     end
 end