All extension algorithms are working again

Author: avik-pal

ext/NonlinearSolveFastLevenbergMarquardtExt.jl

@@ -2,8 +2,9 @@ module NonlinearSolveFastLevenbergMarquardtExt
 
 using ArrayInterface, NonlinearSolve, SciMLBase
 import ConcreteStructs: @concrete
+import FastClosures: @closure
 import FastLevenbergMarquardt as FastLM
-import FiniteDiff, ForwardDiff
+import StaticArraysCore: StaticArray
 
 @inline function _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, x) where {linsolve}
     if linsolve === :cholesky
@@ -28,31 +29,25 @@ end
 
 function SciMLBase.__init(prob::NonlinearLeastSquaresProblem,
         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! = NonlinearSolve.__make_inplace{iip}(prob.f, nothing)
+        reltol = nothing, maxiters = 1000, termination_condition = nothing, kwargs...)
+    NonlinearSolve.__test_termination_condition(termination_condition,
+        :FastLevenbergMarquardt)
+    if prob.u0 isa StaticArray  # FIXME
+        error("FastLevenbergMarquardtJL does not support StaticArrays yet.")
+    end
 
+    _f!, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0)
+    f! = @closure (du, u, p) -> _f!(du, u)
     abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u))
     reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(u))
 
-    if prob.f.jac === nothing
-        alg = NonlinearSolve.get_concrete_algorithm(alg, prob)
-        J! = NonlinearSolve.__construct_jac(prob, alg, u;
-            can_handle_arbitrary_dims = Val(true))
-    else
-        J! = NonlinearSolve.__make_inplace{iip}(prob.f.jac, nothing)
-    end
-
-    J = similar(u, length(fu), length(u))
+    _J! = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; alg.autodiff)
+    J! = @closure (J, u, p) -> _J!(J, u)
+    J = prob.f.jac_prototype === nothing ? similar(u, length(resid), length(u)) :
+        zero(prob.f.jac_prototype)
 
     solver = _fast_lm_solver(alg, u)
-    LM = FastLM.LMWorkspace(u, fu, J)
+    LM = FastLM.LMWorkspace(u, resid, J)
 
     return FastLevenbergMarquardtJLCache(f!, J!, prob, alg, LM, solver,
         (; xtol = reltol, ftol = reltol, gtol = abstol, maxit = maxiters, alg.factor,
@@ -62,7 +57,7 @@ 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...)
+        cache.lmworkspace; cache.solver, cache.kwargs...)
     stats = SciMLBase.NLStats(nfev, njev, -1, -1, iter)
     retcode = info == -1 ? ReturnCode.MaxIters : ReturnCode.Success
     return SciMLBase.build_solution(cache.prob, cache.alg, res, fx;

ext/NonlinearSolveLeastSquaresOptimExt.jl

@@ -4,14 +4,13 @@ using NonlinearSolve, SciMLBase
 import ConcreteStructs: @concrete
 import LeastSquaresOptim as LSO
 
-@inline function _lso_solver(::LeastSquaresOptimJL{alg, linsolve}) where {alg, linsolve}
-    ls = linsolve === :qr ? LSO.QR() :
-         (linsolve === :cholesky ? LSO.Cholesky() :
-          (linsolve === :lsmr ? LSO.LSMR() : nothing))
+@inline function _lso_solver(::LeastSquaresOptimJL{alg, ls}) where {alg, ls}
+    linsolve = ls === :qr ? LSO.QR() :
+               (ls === :cholesky ? LSO.Cholesky() : (ls === :lsmr ? LSO.LSMR() : nothing))
     if alg === :lm
-        return LSO.LevenbergMarquardt(ls)
+        return LSO.LevenbergMarquardt(linsolve)
     elseif alg === :dogleg
-        return LSO.Dogleg(ls)
+        return LSO.Dogleg(linsolve)
     else
         throw(ArgumentError("Unknown LeastSquaresOptim Algorithm: $alg"))
     end
@@ -25,24 +24,26 @@ end
     kwargs
 end
 
-function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, alg::LeastSquaresOptimJL,
-        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 = NonlinearSolve.__maybe_unaliased(prob.u0, alias_u0)
+function SciMLBase.__init(prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem},
+        alg::LeastSquaresOptimJL, args...; alias_u0 = false, abstol = nothing,
+        show_trace::Val{ShT} = Val(false), trace_level = TraceMinimal(), reltol = nothing,
+        store_trace::Val{StT} = Val(false), maxiters = 1000,
+        termination_condition = nothing, kwargs...) where {ShT, StT}
+    NonlinearSolve.__test_termination_condition(termination_condition, :LeastSquaresOptim)
 
+    f!, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0)
     abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u))
     reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(u))
 
-    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
+    if prob.f.jac === nothing && alg.autodiff isa Symbol
+        lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid, alg.autodiff,
+            J = prob.f.jac_prototype, output_length = length(resid))
+    else
+        g! = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; alg.autodiff)
+        lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid, g!,
+            J = prob.f.jac_prototype, output_length = length(resid))
+    end
 
-    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,

ext/NonlinearSolveNLsolveExt.jl

@@ -11,7 +11,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::NLsolveJL, args...;
 
     f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0)
 
-    if prob.f.jac === nothing
+    if prob.f.jac === nothing && alg.autodiff isa Symbol
         df = OnceDifferentiable(f!, u0, resid; alg.autodiff)
     else
         jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; alg.autodiff)

src/algorithms/extension_algs.jl

@@ -1,97 +1,98 @@
 # This file only include the algorithm struct to be exported by NonlinearSolve.jl. The main
 # functionality is implemented as package extensions
-# """
-#     LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central)
-
-# Wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl)
-# for solving `NonlinearLeastSquaresProblem`.
-
-# ## Arguments:
-
-#   - `alg`: Algorithm to use. Can be `:lm` or `:dogleg`.
-#   - `linsolve`: Linear solver to use. Can be `:qr`, `:cholesky` or `:lsmr`. If `nothing`,
-#     then `LeastSquaresOptim.jl` will choose the best linear solver based on the Jacobian
-#     structure.
-#   - `autodiff`: Automatic differentiation / Finite Differences. Can be `:central` or
-#     `:forward`.
-
-# !!! note
-
-#     This algorithm is only available if `LeastSquaresOptim.jl` is installed.
-# """
-# struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm
-#     autodiff::Symbol
-# end
-
-# function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central)
-#     @assert alg in (:lm, :dogleg)
-#     @assert linsolve === nothing || linsolve in (:qr, :cholesky, :lsmr)
-#     @assert autodiff in (:central, :forward)
-
-#     if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing
-#         error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded")
-#     end
-
-#     return LeastSquaresOptimJL{alg, linsolve}(autodiff)
-# end
-
-# """
-#     FastLevenbergMarquardtJL(linsolve = :cholesky; autodiff = nothing)
-
-# Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl)
-# for solving `NonlinearLeastSquaresProblem`.
-
-# !!! warning
-
-#     This is not really the fastest solver. It is called that since the original package
-#     is called "Fast". `LevenbergMarquardt()` is almost always a better choice.
-
-# ## Arguments:
-
-#   - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`.
-#   - `autodiff`: determines the backend used for the Jacobian. Note that this argument is
-#     ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to
-#     `nothing` which means that a default is selected according to the problem specification!
-#     Valid choices are `nothing`, `AutoForwardDiff` or `AutoFiniteDiff`.
-
-# !!! note
-
-#     This algorithm is only available if `FastLevenbergMarquardt.jl` is installed.
-# """
-# @concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveExtensionAlgorithm
-#     ad
-#     factor
-#     factoraccept
-#     factorreject
-#     factorupdate::Symbol
-#     minscale
-#     maxscale
-#     minfactor
-#     maxfactor
-# end
-
-# function set_ad(alg::FastLevenbergMarquardtJL{linsolve}, ad) where {linsolve}
-#     return FastLevenbergMarquardtJL{linsolve}(ad, alg.factor, alg.factoraccept,
-#         alg.factorreject, alg.factorupdate, alg.minscale, alg.maxscale, alg.minfactor,
-#         alg.maxfactor)
-# end
-
-# function FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6,
-#         factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt,
-#         minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32,
-#         autodiff = nothing)
-#     @assert linsolve in (:qr, :cholesky)
-#     @assert factorupdate in (:marquardt, :nielson)
-#     @assert autodiff === nothing || autodiff isa AutoFiniteDiff ||
-#             autodiff isa AutoForwardDiff
-
-#     if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing
-#         error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded")
-#     end
-
-#     return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject,
-#         factorupdate, minscale, maxscale, minfactor, maxfactor)
-# end
+"""
+    LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central)
+
+Wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl)
+for solving `NonlinearLeastSquaresProblem`.
+
+### Arguments
+
+  - `alg`: Algorithm to use. Can be `:lm` or `:dogleg`.
+
+### Keyword Arguments
+
+  - `linsolve`: Linear solver to use. Can be `:qr`, `:cholesky` or `:lsmr`. If `nothing`,
+    then `LeastSquaresOptim.jl` will choose the best linear solver based on the Jacobian
+    structure.
+  - `autodiff`: Automatic differentiation / Finite Differences. Can be `:central` or
+    `:forward`.
+
+!!! note
+
+    This algorithm is only available if `LeastSquaresOptim.jl` is installed.
+"""
+struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm
+    autodiff
+end
+
+function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff = :central)
+    @assert alg in (:lm, :dogleg)
+    @assert linsolve === nothing || linsolve in (:qr, :cholesky, :lsmr)
+    autodiff isa Symbol && @assert autodiff in (:central, :forward)
+
+    if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing
+        error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded")
+    end
+
+    return LeastSquaresOptimJL{alg, linsolve}(autodiff)
+end
+
+"""
+    FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6,
+        factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt,
+        minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32,
+        autodiff = nothing)
+
+Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl)
+for solving `NonlinearLeastSquaresProblem`. For details about the other keyword arguments
+see the documentation for `FastLevenbergMarquardt.jl`.
+
+!!! warning
+
+    This is not really the fastest solver. It is called that since the original package
+    is called "Fast". `LevenbergMarquardt()` is almost always a better choice.
+
+### Arguments
+
+  - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`.
+
+### Keyword Arguments
+
+  - `autodiff`: determines the backend used for the Jacobian. Note that this argument is
+    ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to
+    `nothing` which means that a default is selected according to the problem specification!
+
+!!! note
+
+    This algorithm is only available if `FastLevenbergMarquardt.jl` is installed.
+"""
+@concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveExtensionAlgorithm
+    autodiff
+    factor
+    factoraccept
+    factorreject
+    factorupdate::Symbol
+    minscale
+    maxscale
+    minfactor
+    maxfactor
+end
+
+function FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6,
+        factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt,
+        minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32,
+        autodiff = nothing)
+    @assert linsolve in (:qr, :cholesky)
+    @assert factorupdate in (:marquardt, :nielson)
+
+    if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing
+        error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded")
+    end
+
+    return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject,
+        factorupdate, minscale, maxscale, minfactor, maxfactor)
+end
 
 """
     CMINPACK(; method::Symbol = :auto, autodiff = missing)
@@ -134,6 +135,10 @@ then the following methods are allowed:
 
 The default choice of `:auto` selects `:hybr` for NonlinearProblem and `:lm` for
 NonlinearLeastSquaresProblem.
+
+!!! note
+
+    This algorithm is only available if `MINPACK.jl` is installed.
 """
 @concrete struct CMINPACK <: AbstractNonlinearSolveExtensionAlgorithm
     show_trace::Bool
@@ -206,6 +211,10 @@ Choices for methods in `NLsolveJL`:
 
 For more information on these arguments, consult the
 [NLsolve.jl documentation](https://github.com/JuliaNLSolvers/NLsolve.jl).
+
+!!! note
+
+    This algorithm is only available if `NLsolve.jl` is installed.
 """
 @concrete struct NLsolveJL <: AbstractNonlinearSolveExtensionAlgorithm
     method::Symbol
@@ -289,6 +298,10 @@ Fixed Point Problems. We allow using this algorithm to solve root finding proble
 
 [1] N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms,
     arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974.
+
+!!! note
+
+    This algorithm is only available if `SpeedMapping.jl` is installed.
 """
 @concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm
     σ_min
@@ -337,6 +350,10 @@ problems as well.
     `:SEA` and `:VEA`. For `:SEA` and `:VEA`, this must be a multiple of `2`.
   - `replace_invalids`: The method to use for replacing invalid iterates. Can be
     `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`.
+
+!!! note
+
+    This algorithm is only available if `FixedPointAcceleration.jl` is installed.
 """
 @concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm
     algorithm::Symbol
@@ -411,6 +428,10 @@ end
   - `:pseudotransient`: Pseudo transient method.
   - `:secant`: Secant method for scalar equations.
   - `:anderson`: Anderson acceleration for fixed point iterations.
+
+!!! note
+
+    This algorithm is only available if `SIAMFANLEquations.jl` is installed.
 """
 @concrete struct SIAMFANLEquationsJL{L <: Union{Symbol, Nothing}} <:
                  AbstractNonlinearSolveExtensionAlgorithm