NLLS Poly Algorithm

Author: avik-pal

ext/NonlinearSolveFastLevenbergMarquardtExt.jl

@@ -32,23 +32,24 @@ end
 (f::InplaceFunction{false})(fx, x, p) = (fx .= f.f(x, p))
 
 function SciMLBase.__init(prob::NonlinearLeastSquaresProblem,
-        alg::FastLevenbergMarquardtJL, args...; abstol = 1e-8, reltol = 1e-8,
-        verbose = false, maxiters = 1000, kwargs...)
+        alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = 1e-8,
+        reltol = 1e-8, verbose = false, maxiters = 1000, kwargs...)
     iip = SciMLBase.isinplace(prob)
+    u0 = alias_u0 ? prob.u0 : deepcopy(prob.u0)
 
     @assert prob.f.jac!==nothing "FastLevenbergMarquardt requires a Jacobian!"
 
     f! = InplaceFunction{iip}(prob.f)
     J! = InplaceFunction{iip}(prob.f.jac)
 
     resid_prototype = prob.f.resid_prototype === nothing ?
-                      (!iip ? prob.f(prob.u0, prob.p) : zeros(prob.u0)) :
+                      (!iip ? prob.f(u0, prob.p) : zeros(u0)) :
                       prob.f.resid_prototype
 
-    J = similar(prob.u0, length(resid_prototype), length(prob.u0))
+    J = similar(u0, length(resid_prototype), length(u0))
 
-    solver = _fast_lm_solver(alg, prob.u0)
-    LM = FastLM.LMWorkspace(prob.u0, resid_prototype, J)
+    solver = _fast_lm_solver(alg, u0)
+    LM = FastLM.LMWorkspace(u0, resid_prototype, J)
 
     return FastLevenbergMarquardtJLCache(f!, J!, prob, alg, LM, solver,
         (; xtol = abstol, ftol = reltol, maxit = maxiters, alg.factor, alg.factoraccept,

ext/NonlinearSolveLeastSquaresOptimExt.jl

@@ -33,17 +33,19 @@ end
 (f::FunctionWrapper{false})(du, u) = (du .= f.f(u, f.p))
 
 function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, alg::LeastSquaresOptimJL,
-        args...; abstol = 1e-8, reltol = 1e-8, verbose = false, maxiters = 1000, kwargs...)
+        args...; alias_u0 = false, abstol = 1e-8, reltol = 1e-8, verbose = false,
+        maxiters = 1000, kwargs...)
     iip = SciMLBase.isinplace(prob)
+    u = alias_u0 ? prob.u0 : deepcopy(prob.u0)
 
     f! = FunctionWrapper{iip}(prob.f, prob.p)
     g! = prob.f.jac === nothing ? nothing : FunctionWrapper{iip}(prob.f.jac, prob.p)
 
     resid_prototype = prob.f.resid_prototype === nothing ?
-                      (!iip ? prob.f(prob.u0, prob.p) : zeros(prob.u0)) :
+                      (!iip ? prob.f(u, prob.p) : zeros(u)) :
                       prob.f.resid_prototype
 
-    lsoprob = LSO.LeastSquaresProblem(; x = prob.u0, f!, y = resid_prototype, g!,
+    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))
 

src/NonlinearSolve.jl

@@ -147,7 +147,8 @@ export RadiusUpdateSchemes
 export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient,
     GeneralBroyden, GeneralKlement, LimitedMemoryBroyden
 export LeastSquaresOptimJL, FastLevenbergMarquardtJL
-export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPolyalg
+export NonlinearSolvePolyAlgorithm,
+    RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg
 
 export LineSearch, LiFukushimaLineSearch
 

src/default.jl

@@ -250,6 +250,42 @@ function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothi
     return NonlinearSolvePolyAlgorithm(algs, Val(:NLS))
 end
 
+"""
+    FastShortcutNLLSPolyalg(; concrete_jac = nothing, linsolve = nothing,
+                              precs = DEFAULT_PRECS, adkwargs...)
+
+A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods
+for more performance and then tries more robust techniques if the faster ones fail.
+
+### 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
+    `AutoForwardDiff()`. Valid choices are types from ADTypes.jl.
+  - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used,
+    then the Jacobian will not be constructed and instead direct Jacobian-vector products
+    `J*v` are computed using forward-mode automatic differentiation or finite differencing
+    tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed,
+    for example for a preconditioner, `concrete_jac = true` can be passed in order to force
+    the construction of the Jacobian.
+  - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the
+    linear solves within the Newton method. Defaults to `nothing`, which means it uses the
+    LinearSolve.jl default algorithm choice. For more information on available algorithm
+    choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
+  - `precs`: the choice of preconditioners for the linear solver. Defaults to using no
+    preconditioners. For more information on specifying preconditioners for LinearSolve
+    algorithms, consult the
+    [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
+"""
+function FastShortcutNLLSPolyalg(; concrete_jac = nothing, linsolve = nothing,
+        precs = DEFAULT_PRECS, adkwargs...)
+    algs = (GaussNewton(; concrete_jac, linsolve, precs, adkwargs...),
+        GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(),
+            adkwargs...),
+        LevenbergMarquardt(; concrete_jac, linsolve, precs, adkwargs...))
+    return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS))
+end
+
 ## Defaults
 
 function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
@@ -263,10 +299,10 @@ end
 # FIXME: We default to using LM currently. But once we have line searches for GN implemented
 #        we should default to a polyalgorithm.
 function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...)
-    return SciMLBase.__init(prob, LevenbergMarquardt(), args...; kwargs...)
+    return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(), args...; kwargs...)
 end
 
 function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, ::Nothing, args...;
         kwargs...)
-    return SciMLBase.__solve(prob, LevenbergMarquardt(), args...; kwargs...)
+    return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(), args...; kwargs...)
 end

test/nonlinear_least_squares.jl

@@ -36,7 +36,7 @@ append!(solvers,
         LevenbergMarquardt(; linsolve = LUFactorization()),
         LeastSquaresOptimJL(:lm),
         LeastSquaresOptimJL(:dogleg),
-	nothing,
+        nothing,
     ])
 
 for prob in nlls_problems, solver in solvers
@@ -46,7 +46,8 @@ for prob in nlls_problems, solver in solvers
 end
 
 function jac!(J, θ, p)
-    ForwardDiff.jacobian!(J, resid -> loss_function(resid, θ, p), θ)
+    resid = zeros(length(p))
+    ForwardDiff.jacobian!(J, (resid, θ) -> loss_function(resid, θ, p), resid, θ)
     return J
 end
 
@@ -57,6 +58,5 @@ solvers = [FastLevenbergMarquardtJL(:cholesky), FastLevenbergMarquardtJL(:qr)]
 
 for solver in solvers
     @time sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8)
-    @test SciMLBase.successful_retcode(sol)
     @test norm(sol.resid) < 1e-6
 end