Merge pull request #341 from SciML/ap/fastlm_jac

Standardize the Extension Algorithms

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -18,8 +18,8 @@ jobs:
       fail-fast: false
       matrix:
         group:
-          - Core
-          - NLLS
+          - RootFinding
+          - NLLSSolvers
           - 23TestProblems
           - Wrappers
           - Miscellaneous

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.3.0"
+version = "3.4.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -10,6 +10,7 @@ ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
+FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
@@ -64,6 +65,7 @@ DiffEqBase = "6.144"
 EnumX = "1"
 Enzyme = "0.11.11"
 FastBroadcast = "0.2.8"
+FastClosures = "0.3"
 FastLevenbergMarquardt = "0.1"
 FiniteDiff = "2.21"
 FixedPointAcceleration = "0.3"
@@ -85,16 +87,17 @@ Printf = "1.9"
 Random = "1.91"
 RecursiveArrayTools = "3.2"
 Reexport = "1.2"
+SIAMFANLEquations = "1.0.1"
 SafeTestsets = "0.1"
 SciMLBase = "2.11"
 SciMLOperators = "0.3.7"
-SIAMFANLEquations = "1.0.1"
 SimpleNonlinearSolve = "1.0.2"
 SparseArrays = "1.9"
 SparseDiffTools = "2.14"
 SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.7"
+Sundials = "4.23.1"
 Symbolics = "5.13"
 Test = "1"
 UnPack = "1.0"
@@ -120,15 +123,16 @@ NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 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"]
+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"]

docs/make.jl

@@ -14,7 +14,7 @@ makedocs(; sitename = "NonlinearSolve.jl",
         DiffEqBase, SciMLBase],
     clean = true, doctest = false, linkcheck = true,
     linkcheck_ignore = ["https://twitter.com/ChrisRackauckas/status/1544743542094020615"],
-    warnonly = [:cross_references], checkdocs = :export,
+    checkdocs = :export,
     format = Documenter.HTML(assets = ["assets/favicon.ico"],
         canonical = "https://docs.sciml.ai/NonlinearSolve/stable/"),
     pages)

docs/src/api/siamfanlequations.md

@@ -1,6 +1,7 @@
 # SIAMFANLEquations.jl
 
-This is an extension for importing solvers from [SIAMFANLEquations.jl](https://github.com/ctkelley/SIAMFANLEquations.jl) into the SciML
+This is an extension for importing solvers from
+[SIAMFANLEquations.jl](https://github.com/ctkelley/SIAMFANLEquations.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:
 

docs/src/basics/solve.md

@@ -14,8 +14,8 @@ solve(prob::SciMLBase.NonlinearProblem, args...; kwargs...)
 ## Iteration Controls
 
   - `maxiters::Int`: The maximum number of iterations to perform. Defaults to `1000`.
-  - `abstol::Number`: The absolute tolerance.
-  - `reltol::Number`: The relative tolerance.
+  - `abstol::Number`: The absolute tolerance. Defaults to `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`.
+  - `reltol::Number`: The relative tolerance. Defaults to `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`.
   - `termination_condition`: Termination Condition from DiffEqBase. Defaults to
     `AbsSafeBestTerminationMode()` for `NonlinearSolve.jl` and `AbsTerminateMode()` for
     `SimpleNonlinearSolve.jl`.

docs/src/solvers/FixedPointSolvers.md

@@ -40,3 +40,10 @@ We are only listing the methods that natively solve fixed point problems.
 
   - `FixedPointAccelerationJL()`: accelerates the convergence of a mapping to a fixed point
     by the Anderson acceleration algorithm and a few other methods.
+
+### NLsolve.jl
+
+In our tests, we have found the anderson method implemented here to NOT be the most
+robust.
+
+  - `NLsolveJL(; method = :anderson)`: Anderson acceleration for fixed point problems.

docs/src/solvers/NonlinearSystemSolvers.md

@@ -143,3 +143,10 @@ Newton-Krylov form. However, KINSOL is known to be less stable than some other
 implementations, as it has no line search or globalizer (trust region).
 
   - `KINSOL()`: The KINSOL method of the SUNDIALS C library
+
+### SIAMFANLEquations.jl
+
+SIAMFANLEquations.jl is a wrapper for the methods in the SIAMFANLEquations.jl library.
+
+  - `SIAMFANLEquationsJL()`: A wrapper for using the methods in
+    [SIAMFANLEquations.jl](https://github.com/ctkelley/SIAMFANLEquations.jl)

ext/NonlinearSolveFastLevenbergMarquardtExt.jl

@@ -5,7 +5,7 @@ import ConcreteStructs: @concrete
 import FastLevenbergMarquardt as FastLM
 import FiniteDiff, ForwardDiff
 
-function _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, x) where {linsolve}
+@inline function _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, x) where {linsolve}
     if linsolve === :cholesky
         return FastLM.CholeskySolver(ArrayInterface.undefmatrix(x))
     elseif linsolve === :qr
@@ -15,6 +15,7 @@ function _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, x) where {linsolv
     end
 end
 
+# TODO: Implement reinit
 @concrete struct FastLevenbergMarquardtJLCache
     f!
     J!
@@ -25,68 +26,27 @@ end
     kwargs
 end
 
-@concrete struct InplaceFunction{iip} <: Function
-    f
-end
-
-(f::InplaceFunction{true})(fx, x, p) = f.f(fx, x, p)
-(f::InplaceFunction{false})(fx, x, p) = (fx .= f.f(x, p))
-
 function SciMLBase.__init(prob::NonlinearLeastSquaresProblem,
-        alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = 1e-8,
-        reltol = 1e-8, maxiters = 1000, kwargs...)
+        alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = nothing,
+        reltol = nothing, maxiters = 1000, kwargs...)
+    # FIXME: Support scalar u0
+    prob.u0 isa Number &&
+        throw(ArgumentError("FastLevenbergMarquardtJL does not support scalar `u0`"))
     iip = SciMLBase.isinplace(prob)
     u = NonlinearSolve.__maybe_unaliased(prob.u0, alias_u0)
     fu = NonlinearSolve.evaluate_f(prob, u)
 
-    f! = InplaceFunction{iip}(prob.f)
+    f! = NonlinearSolve.__make_inplace{iip}(prob.f, nothing)
+
+    abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u))
+    reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(u))
 
     if prob.f.jac === nothing
-        use_forward_diff = if alg.autodiff === nothing
-            ForwardDiff.can_dual(eltype(u))
-        else
-            alg.autodiff isa AutoForwardDiff
-        end
-        uf = SciMLBase.JacobianWrapper{iip}(prob.f, prob.p)
-        if use_forward_diff
-            cache = iip ? ForwardDiff.JacobianConfig(uf, fu, u) :
-                    ForwardDiff.JacobianConfig(uf, u)
-        else
-            cache = FiniteDiff.JacobianCache(u, fu)
-        end
-        J! = if iip
-            if use_forward_diff
-                fu_cache = similar(fu)
-                function (J, x, p)
-                    uf.p = p
-                    ForwardDiff.jacobian!(J, uf, fu_cache, x, cache)
-                    return J
-                end
-            else
-                function (J, x, p)
-                    uf.p = p
-                    FiniteDiff.finite_difference_jacobian!(J, uf, x, cache)
-                    return J
-                end
-            end
-        else
-            if use_forward_diff
-                function (J, x, p)
-                    uf.p = p
-                    ForwardDiff.jacobian!(J, uf, x, cache)
-                    return J
-                end
-            else
-                function (J, x, p)
-                    uf.p = p
-                    J_ = FiniteDiff.finite_difference_jacobian(uf, x, cache)
-                    copyto!(J, J_)
-                    return J
-                end
-            end
-        end
+        alg = NonlinearSolve.get_concrete_algorithm(alg, prob)
+        J! = NonlinearSolve.__construct_jac(prob, alg, u;
+            can_handle_arbitrary_dims = Val(true))
     else
-        J! = InplaceFunction{iip}(prob.f.jac)
+        J! = NonlinearSolve.__make_inplace{iip}(prob.f.jac, nothing)
     end
 
     J = similar(u, length(fu), length(u))
@@ -95,17 +55,16 @@ function SciMLBase.__init(prob::NonlinearLeastSquaresProblem,
     LM = FastLM.LMWorkspace(u, fu, J)
 
     return FastLevenbergMarquardtJLCache(f!, J!, prob, alg, LM, solver,
-        (; xtol = abstol, ftol = reltol, maxit = maxiters, alg.factor, alg.factoraccept,
-            alg.factorreject, alg.minscale, alg.maxscale, alg.factorupdate, alg.minfactor,
-            alg.maxfactor, kwargs...))
+        (; xtol = reltol, ftol = reltol, gtol = abstol, maxit = maxiters, alg.factor,
+            alg.factoraccept, alg.factorreject, alg.minscale, alg.maxscale,
+            alg.factorupdate, alg.minfactor, alg.maxfactor))
 end
 
 function SciMLBase.solve!(cache::FastLevenbergMarquardtJLCache)
     res, fx, info, iter, nfev, njev, LM, solver = FastLM.lmsolve!(cache.f!, cache.J!,
         cache.lmworkspace, cache.prob.p; cache.solver, cache.kwargs...)
     stats = SciMLBase.NLStats(nfev, njev, -1, -1, iter)
-    retcode = info == 1 ? ReturnCode.Success :
-              (info == -1 ? ReturnCode.MaxIters : ReturnCode.Default)
+    retcode = info == -1 ? ReturnCode.MaxIters : ReturnCode.Success
     return SciMLBase.build_solution(cache.prob, cache.alg, res, fx;
         retcode, original = (res, fx, info, iter, nfev, njev, LM, solver), stats)
 end

ext/NonlinearSolveFixedPointAccelerationExt.jl

@@ -7,43 +7,27 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::FixedPointAccelerationJL
         show_trace::Val{PrintReports} = Val(false), termination_condition = nothing,
         kwargs...) where {PrintReports}
     @assert (termination_condition ===
-             nothing)||(termination_condition isa AbsNormTerminationMode) "SpeedMappingJL does not support termination conditions!"
+             nothing)||(termination_condition isa AbsNormTerminationMode) "FixedPointAccelerationJL does not support termination conditions!"
 
-    u0 = NonlinearSolve.__maybe_unaliased(prob.u0, alias_u0)
-    u_size = size(u0)
-    T = eltype(u0)
-    iip = isinplace(prob)
-    p = prob.p
+    f, u0 = NonlinearSolve.__construct_f(prob; alias_u0, make_fixed_point = Val(true),
+        force_oop = Val(true))
 
-    if !iip && prob.u0 isa Number
-        # FixedPointAcceleration makes the scalar problem into a vector problem
-        f = (u) -> [prob.f(u[1], p) .+ u[1]]
-    elseif !iip && prob.u0 isa AbstractVector
-        f = (u) -> (prob.f(u, p) .+ u)
-    elseif !iip && prob.u0 isa AbstractArray
-        f = (u) -> vec(prob.f(reshape(u, u_size), p) .+ u)
-    elseif iip && prob.u0 isa AbstractVector
-        du = similar(u0)
-        f = (u) -> (prob.f(du, u, p); du .+ u)
-    else
-        du = similar(u0)
-        f = (u) -> (prob.f(du, reshape(u, u_size), p); vec(du) .+ u)
-    end
-
-    tol = abstol === nothing ? real(oneunit(T)) * (eps(real(one(T))))^(4 // 5) : abstol
+    tol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u0))
 
-    sol = fixed_point(f, NonlinearSolve._vec(u0); Algorithm = alg.algorithm,
+    sol = fixed_point(f, u0; Algorithm = alg.algorithm,
         ConvergenceMetricThreshold = tol, MaxIter = maxiters, MaxM = alg.m,
         ExtrapolationPeriod = alg.extrapolation_period, Dampening = alg.dampening,
         PrintReports, ReplaceInvalids = alg.replace_invalids,
         ConditionNumberThreshold = alg.condition_number_threshold, quiet_errors = true)
 
-    res = prob.u0 isa Number ? first(sol.FixedPoint_) : sol.FixedPoint_
-    if res === missing
+    if sol.FixedPoint_ === missing
+        u0 = prob.u0 isa Number ? u0[1] : u0
         resid = NonlinearSolve.evaluate_f(prob, u0)
         res = u0
         converged = false
     else
+        res = prob.u0 isa Number ? first(sol.FixedPoint_) :
+              reshape(sol.FixedPoint_, size(prob.u0))
         resid = NonlinearSolve.evaluate_f(prob, res)
         converged = maximum(abs, resid) ≤ tol
     end

ext/NonlinearSolveLeastSquaresOptimExt.jl

@@ -4,7 +4,7 @@ using NonlinearSolve, SciMLBase
 import ConcreteStructs: @concrete
 import LeastSquaresOptim as LSO
 
-function _lso_solver(::LeastSquaresOptimJL{alg, linsolve}) where {alg, linsolve}
+@inline function _lso_solver(::LeastSquaresOptimJL{alg, linsolve}) where {alg, linsolve}
     ls = linsolve === :qr ? LSO.QR() :
          (linsolve === :cholesky ? LSO.Cholesky() :
           (linsolve === :lsmr ? LSO.LSMR() : nothing))
@@ -17,41 +17,37 @@ function _lso_solver(::LeastSquaresOptimJL{alg, linsolve}) where {alg, linsolve}
     end
 end
 
+# TODO: Implement reinit
 @concrete struct LeastSquaresOptimJLCache
     prob
     alg
     allocated_prob
     kwargs
 end
 
-@concrete struct FunctionWrapper{iip}
-    f
-    p
-end
-
-(f::FunctionWrapper{true})(du, u) = f.f(du, u, f.p)
-(f::FunctionWrapper{false})(du, u) = (du .= f.f(u, f.p))
-
 function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, alg::LeastSquaresOptimJL,
-        args...; alias_u0 = false, abstol = 1e-8, reltol = 1e-8, verbose = false,
-        maxiters = 1000, kwargs...)
+        args...; alias_u0 = false, abstol = nothing, show_trace::Val{ShT} = Val(false),
+        trace_level = TraceMinimal(), store_trace::Val{StT} = Val(false), maxiters = 1000,
+        reltol = nothing, kwargs...) where {ShT, StT}
     iip = SciMLBase.isinplace(prob)
-    u = alias_u0 ? prob.u0 : deepcopy(prob.u0)
+    u = NonlinearSolve.__maybe_unaliased(prob.u0, alias_u0)
+
+    abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u))
+    reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(u))
 
-    f! = FunctionWrapper{iip}(prob.f, prob.p)
-    g! = prob.f.jac === nothing ? nothing : FunctionWrapper{iip}(prob.f.jac, prob.p)
+    f! = NonlinearSolve.__make_inplace{iip}(prob.f, prob.p)
+    g! = NonlinearSolve.__make_inplace{iip}(prob.f.jac, prob.p)
 
     resid_prototype = prob.f.resid_prototype === nothing ?
-                      (!iip ? prob.f(u, prob.p) : zeros(u)) :
-                      prob.f.resid_prototype
+                      (!iip ? prob.f(u, prob.p) : zeros(u)) : prob.f.resid_prototype
 
     lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid_prototype, g!,
         J = prob.f.jac_prototype, alg.autodiff, output_length = length(resid_prototype))
     allocated_prob = LSO.LeastSquaresProblemAllocated(lsoprob, _lso_solver(alg))
 
     return LeastSquaresOptimJLCache(prob, alg, allocated_prob,
-        (; x_tol = abstol, f_tol = reltol, iterations = maxiters, show_trace = verbose,
-            kwargs...))
+        (; x_tol = reltol, f_tol = abstol, g_tol = abstol, iterations = maxiters,
+            show_trace = ShT, store_trace = StT, show_every = trace_level.print_frequency))
 end
 
 function SciMLBase.solve!(cache::LeastSquaresOptimJLCache)