Merge pull request #379 from SciML/ap/nlsolvers

Wrapper for NLSolvers.jl
SciML · Feb 18, 2024 · b731362 · b731362
2 parents d383660 + 9a6bdfb
commit b731362
Show file tree

Hide file tree

Showing 8 changed files with 170 additions and 10 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.5.6"
+version = "3.6.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -36,6 +36,7 @@ FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+NLSolvers = "337daf1e-9722-11e9-073e-8b9effe078ba"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
@@ -48,6 +49,7 @@ NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
 NonlinearSolveMINPACKExt = "MINPACK"
 NonlinearSolveNLsolveExt = "NLsolve"
+NonlinearSolveNLSolversExt = "NLSolvers"
 NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
 NonlinearSolveSpeedMappingExt = "SpeedMapping"
 NonlinearSolveSymbolicsExt = "Symbolics"
@@ -77,6 +79,7 @@ LinearSolve = "2.21"
 MINPACK = "1.2"
 MaybeInplace = "0.1.1"
 NLsolve = "4.5"
+NLSolvers = "0.5"
 NaNMath = "1"
 NonlinearProblemLibrary = "0.1.2"
 OrdinaryDiffEq = "6.63"
@@ -120,6 +123,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+NLSolvers = "337daf1e-9722-11e9-073e-8b9effe078ba"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
@@ -139,4 +143,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration", "SIAMFANLEquations", "Sundials", "ReTestItems", "Reexport", "CUDA"]
+test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration", "SIAMFANLEquations", "Sundials", "ReTestItems", "Reexport", "CUDA", "NLSolvers"]
diff --git a/docs/src/api/nlsolve.md b/docs/src/api/nlsolve.md
@@ -7,7 +7,7 @@ using these solvers:
 ```julia
 using Pkg
 Pkg.add("NLsolve")
-using NLSolve, NonlinearSolve
+using NLsolve, NonlinearSolve
 ```
 
 ## Solver API

diff --git a/docs/src/api/nlsolvers.md b/docs/src/api/nlsolvers.md
@@ -0,0 +1,17 @@
+# NLSolvers.jl
+
+This is a extension for importing solvers from NLSolvers.jl into the SciML interface. Note
+that these solvers do not come by default, and thus one needs to install the package before
+using these solvers:
+
+```julia
+using Pkg
+Pkg.add("NLSolvers")
+using NLSolvers, NonlinearSolve
+```
+
+## Solver API
+
+```@docs
+NLSolversJL
+```
diff --git a/docs/src/solvers/nonlinear_system_solvers.md b/docs/src/solvers/nonlinear_system_solvers.md
@@ -168,3 +168,12 @@ SIAMFANLEquations.jl is a wrapper for the methods in the SIAMFANLEquations.jl li
 Other solvers listed in [Fixed Point Solvers](@ref fixed_point_methods_full_list),
 [FastLevenbergMarquardt.jl](@ref fastlm_wrapper_summary) and
 [LeastSquaresOptim.jl](@ref lso_wrapper_summary) can also solve nonlinear systems.
+
+### NLSolvers.jl
+
+This is a wrapper package for importing solvers from NLSolvers.jl into the SciML interface.
+
+  - [`NLSolversJL()`](@ref): A wrapper for
+    [NLSolvers.jl](https://github.com/JuliaNLSolvers/NLSolvers.jl)
+
+For a list of possible solvers see the [NLSolvers.jl documentation](https://julianlsolvers.github.io/NLSolvers.jl/)
diff --git a/ext/NonlinearSolveNLSolversExt.jl b/ext/NonlinearSolveNLSolversExt.jl
@@ -0,0 +1,84 @@
+module NonlinearSolveNLSolversExt
+
+using ADTypes, FastClosures, NonlinearSolve, NLSolvers, SciMLBase, LinearAlgebra
+using FiniteDiff, ForwardDiff
+
+function SciMLBase.__solve(prob::NonlinearProblem, alg::NLSolversJL, args...;
+        abstol = nothing, reltol = nothing, maxiters = 1000, alias_u0::Bool = false,
+        termination_condition = nothing, kwargs...)
+    NonlinearSolve.__test_termination_condition(termination_condition, :NLSolversJL)
+
+    abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(prob.u0))
+    reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(prob.u0))
+
+    options = NEqOptions(; maxiter = maxiters, f_abstol = abstol, f_reltol = reltol)
+
+    if prob.u0 isa Number
+        f_scalar = @closure x -> prob.f(x, prob.p)
+
+        if alg.autodiff === nothing
+            if ForwardDiff.can_dual(typeof(prob.u0))
+                autodiff_concrete = :forwarddiff
+            else
+                autodiff_concrete = :finitediff
+            end
+        else
+            if alg.autodiff isa AutoForwardDiff || alg.autodiff isa AutoPolyesterForwardDiff
+                autodiff_concrete = :forwarddiff
+            elseif alg.autodiff isa AutoFiniteDiff
+                autodiff_concrete = :finitediff
+            else
+                error("Only ForwardDiff or FiniteDiff autodiff is supported.")
+            end
+        end
+
+        if autodiff_concrete === :forwarddiff
+            fj_scalar = @closure (Jx, x) -> begin
+                T = typeof(NonlinearSolve.NonlinearSolveTag())
+                x_dual = ForwardDiff.Dual{T}(x, one(x))
+                y = prob.f(x_dual, prob.p)
+                return ForwardDiff.value(y), ForwardDiff.extract_derivative(T, y)
+            end
+        else
+            fj_scalar = @closure (Jx, x) -> begin
+                _f = Base.Fix2(prob.f, prob.p)
+                return _f(x), FiniteDiff.finite_difference_derivative(_f, x)
+            end
+        end
+
+        prob_obj = NLSolvers.ScalarObjective(; f = f_scalar, fg = fj_scalar)
+        prob_nlsolver = NEqProblem(prob_obj; inplace = false)
+        res = NLSolvers.solve(prob_nlsolver, prob.u0, alg.method, options)
+
+        retcode = ifelse(norm(res.info.best_residual, Inf) ≤ abstol, ReturnCode.Success,
+            ReturnCode.MaxIters)
+        stats = SciMLBase.NLStats(-1, -1, -1, -1, res.info.iter)
+
+        return SciMLBase.build_solution(prob, alg, res.info.solution,
+            res.info.best_residual; retcode, original = res, stats)
+    end
+
+    f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0)
+
+    jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; alg.autodiff)
+
+    FJ_vector! = @closure (Fx, Jx, x) -> begin
+        f!(Fx, x)
+        jac!(Jx, x)
+        return Fx, Jx
+    end
+
+    prob_obj = NLSolvers.VectorObjective(; F = f!, FJ = FJ_vector!)
+    prob_nlsolver = NEqProblem(prob_obj)
+
+    res = NLSolvers.solve(prob_nlsolver, u0, alg.method, options)
+
+    retcode = ifelse(norm(res.info.best_residual, Inf) ≤ abstol, ReturnCode.Success,
+        ReturnCode.MaxIters)
+    stats = SciMLBase.NLStats(-1, -1, -1, -1, res.info.iter)
+
+    return SciMLBase.build_solution(prob, alg, res.info.solution,
+        res.info.best_residual; retcode, original = res, stats)
+end
+
+end
diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -144,7 +144,7 @@ export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPoly
        FastShortcutNLLSPolyalg
 
 # Extension Algorithms
-export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL,
+export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, NLSolversJL,
        FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL
 
 # Advanced Algorithms -- Without Bells and Whistles

diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl
@@ -279,6 +279,38 @@ function NLsolveJL(; method = :trust_region, autodiff = :central, store_trace =
         linsolve, factor, autoscale, m, beta, show_trace)
 end
 
+"""
+    NLSolversJL(method; autodiff = nothing)
+    NLSolversJL(; method, autodiff = nothing)
+
+Wrapper over NLSolvers.jl Nonlinear Equation Solvers. We automatically construct the
+jacobian function and supply it to the solver.
+
+### Arguments
+
+  - `method`: the choice of method for solving the nonlinear system. See the documentation
+    for NLSolvers.jl for more information.
+  - `autodiff`: the choice of method for generating the Jacobian. Defaults to `nothing`
+    which means that a default is selected according to the problem specification. Can be
+    any valid ADTypes.jl autodiff type (conditional on that backend being supported in
+    NonlinearSolve.jl).
+"""
+struct NLSolversJL{M, AD} <: AbstractNonlinearSolveExtensionAlgorithm
+    method::M
+    autodiff::AD
+
+    function NLSolversJL(method, autodiff)
+        if Base.get_extension(@__MODULE__, :NonlinearSolveNLSolversExt) === nothing
+            error("NLSolversJL requires NLSolvers.jl to be loaded")
+        end
+
+        return new{typeof(method), typeof(autodiff)}(method, autodiff)
+    end
+end
+
+NLSolversJL(method; autodiff = nothing) = NLSolversJL(method, autodiff)
+NLSolversJL(; method, autodiff = nothing) = NLSolversJL(method, autodiff)
+
 """
     SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false,
         orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000)

diff --git a/test/wrappers/rootfind_tests.jl b/test/wrappers/rootfind_tests.jl
@@ -1,6 +1,9 @@
 @testsetup module WrapperRootfindImports
 using Reexport
-@reexport using LinearAlgebra, NLsolve, SIAMFANLEquations, MINPACK
+@reexport using LinearAlgebra
+import NLSolvers, NLsolve, SIAMFANLEquations, MINPACK
+
+export NLSolvers
 end
 
 @testitem "Steady State Problems" setup=[WrapperRootfindImports] begin
@@ -12,7 +15,8 @@ end
     u0 = zeros(2)
     prob_iip = SteadyStateProblem(f_iip, u0)
 
-    for alg in [NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL()]
+    for alg in [
+        NLSolversJL(NLSolvers.LineSearch(NLSolvers.Newton(), NLSolvers.Backtracking())), NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL()]
         sol = solve(prob_iip, alg)
         @test SciMLBase.successful_retcode(sol.retcode)
         @test maximum(abs, sol.resid) < 1e-6
@@ -23,7 +27,8 @@ end
     u0 = zeros(2)
     prob_oop = SteadyStateProblem(f_oop, u0)
 
-    for alg in [NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL()]
+    for alg in [
+        NLSolversJL(NLSolvers.LineSearch(NLSolvers.Newton(), NLSolvers.Backtracking())), NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL()]
         sol = solve(prob_oop, alg)
         @test SciMLBase.successful_retcode(sol.retcode)
         @test maximum(abs, sol.resid) < 1e-6
@@ -39,7 +44,8 @@ end
     u0 = zeros(2)
     prob_iip = NonlinearProblem{true}(f_iip, u0)
 
-    for alg in [NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL()]
+    for alg in [
+        NLSolversJL(NLSolvers.LineSearch(NLSolvers.Newton(), NLSolvers.Backtracking())), NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL()]
         local sol
         sol = solve(prob_iip, alg)
         @test SciMLBase.successful_retcode(sol.retcode)
@@ -50,7 +56,8 @@ end
     f_oop(u, p) = [2 - 2u[1], u[1] - 4u[2]]
     u0 = zeros(2)
     prob_oop = NonlinearProblem{false}(f_oop, u0)
-    for alg in [NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL()]
+    for alg in [
+        NLSolversJL(NLSolvers.LineSearch(NLSolvers.Newton(), NLSolvers.Backtracking())), NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL()]
         local sol
         sol = solve(prob_oop, alg)
         @test SciMLBase.successful_retcode(sol.retcode)
@@ -61,7 +68,10 @@ end
     f_tol(u, p) = u^2 - 2
     prob_tol = NonlinearProblem(f_tol, 1.0)
     for tol in [1e-1, 1e-3, 1e-6, 1e-10, 1e-15],
-        alg in [NLsolveJL(), CMINPACK(), SIAMFANLEquationsJL(; method = :newton),
+        alg in [
+            NLSolversJL(NLSolvers.LineSearch(NLSolvers.Newton(), NLSolvers.Backtracking())),
+            NLsolveJL(),
+            CMINPACK(), SIAMFANLEquationsJL(; method = :newton),
             SIAMFANLEquationsJL(; method = :pseudotransient),
             SIAMFANLEquationsJL(; method = :secant)]
 
@@ -107,6 +117,10 @@ end
 
     sol = solve(ProbN, NLsolveJL(); abstol = 1e-8)
     @test maximum(abs, sol.resid) < 1e-6
+    sol = solve(ProbN,
+        NLSolversJL(NLSolvers.LineSearch(NLSolvers.Newton(), NLSolvers.Backtracking()));
+        abstol = 1e-8)
+    @test maximum(abs, sol.resid) < 1e-6
     sol = solve(ProbN, SIAMFANLEquationsJL(; method = :newton); abstol = 1e-8)
     @test maximum(abs, sol.resid) < 1e-6
     sol = solve(ProbN, SIAMFANLEquationsJL(; method = :pseudotransient); abstol = 1e-8)