Merge pull request #258 from avik-pal/ap/fix_gauss_newton

Gauss Newton & LM Robustness Fixes

.github/workflows/CI.yml

@@ -14,12 +14,13 @@ jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
         group:
-          - Core
-          - 23TestProblems
+          - All
         version:
           - '1'
+          - '~1.10.0-0'
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "2.4.0"
+version = "2.5.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -44,7 +44,7 @@ FiniteDiff = "2"
 ForwardDiff = "0.10.3"
 LeastSquaresOptim = "0.8"
 LineSearches = "7"
-LinearSolve = "2"
+LinearSolve = "2.12"
 NonlinearProblemLibrary = "0.1"
 PrecompileTools = "1"
 RecursiveArrayTools = "2"
@@ -65,6 +65,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -76,4 +77,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt"]
+test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath"]

docs/src/solvers/NonlinearLeastSquaresSolvers.md

@@ -19,10 +19,8 @@ Solves the nonlinear least squares problem defined by `prob` using the algorithm
     handling of sparse matrices via colored automatic differentiation and preconditioned
     linear solvers. Designed for large-scale and numerically-difficult nonlinear least squares
     problems.
-  - `SimpleNewtonRaphson()`: Newton Raphson implementation that uses Linear Least Squares
-    solution at every step to compute the descent direction. **WARNING**: This method is not
-    a robust solver for nonlinear least squares problems. The computed delta step might not
-    be the correct descent direction!
+  - `SimpleNewtonRaphson()`: Simple Gauss Newton Implementation with `QRFactorization` to
+    solve a linear least squares problem at each step!
 
 ## Example usage
 

docs/src/solvers/NonlinearSystemSolvers.md

@@ -71,6 +71,10 @@ features, but have a bit of overhead on very small problems.
   - `GeneralKlement()`: Generalization of Klement's Quasi-Newton Method with Line Search and
     Automatic Jacobian Resetting. This is a fast method but unstable when the condition number of
     the Jacobian matrix is sufficiently large.
+  - `LimitedMemoryBroyden()`: An advanced version of `LBroyden` which uses a limited memory
+    Broyden method. This is a fast method but unstable when the condition number of
+    the Jacobian matrix is sufficiently large. It is recommended to use `GeneralBroyden` or
+    `GeneralKlement` instead unless the memory usage is a concern.
 
 ### SimpleNonlinearSolve.jl
 

src/broyden.jl

@@ -66,8 +66,8 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralBroyde
                       alg.reset_tolerance
     reset_check = x -> abs(x) ≤ reset_tolerance
     return GeneralBroydenCache{iip}(f, alg, u, _mutable_zero(u), fu, zero(fu),
-        zero(fu), p, J⁻¹, zero(_vec(fu)'), _mutable_zero(u), false, 0, alg.max_resets,
-        maxiters, internalnorm, ReturnCode.Default, abstol, reset_tolerance,
+        zero(fu), p, J⁻¹, zero(_reshape(fu, 1, :)), _mutable_zero(u), false, 0,
+        alg.max_resets, maxiters, internalnorm, ReturnCode.Default, abstol, reset_tolerance,
         reset_check, prob, NLStats(1, 0, 0, 0, 0),
         init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)))
 end
@@ -78,7 +78,7 @@ function perform_step!(cache::GeneralBroydenCache{true})
 
     mul!(_vec(du), J⁻¹, -_vec(fu))
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(fu2, u, p)
 
     cache.internalnorm(fu2) < cache.abstol && (cache.force_stop = true)
@@ -101,7 +101,8 @@ function perform_step!(cache::GeneralBroydenCache{true})
     else
         mul!(_vec(J⁻¹df), J⁻¹, _vec(dfu))
         mul!(J⁻¹₂, _vec(du)', J⁻¹)
-        du .= (du .- J⁻¹df) ./ (dot(du, J⁻¹df) .+ T(1e-5))
+        denom = dot(du, J⁻¹df)
+        du .= (du .- J⁻¹df) ./ ifelse(iszero(denom), T(1e-5), denom)
         mul!(J⁻¹, _vec(du), J⁻¹₂, 1, 1)
     end
     fu .= fu2
@@ -136,7 +137,8 @@ function perform_step!(cache::GeneralBroydenCache{false})
     else
         cache.J⁻¹df = _restructure(cache.J⁻¹df, cache.J⁻¹ * _vec(cache.dfu))
         cache.J⁻¹₂ = _vec(cache.du)' * cache.J⁻¹
-        cache.du = (cache.du .- cache.J⁻¹df) ./ (dot(cache.du, cache.J⁻¹df) .+ T(1e-5))
+        denom = dot(cache.du, cache.J⁻¹df)
+        cache.du = (cache.du .- cache.J⁻¹df) ./ ifelse(iszero(denom), T(1e-5), denom)
         cache.J⁻¹ = cache.J⁻¹ .+ _vec(cache.du) * cache.J⁻¹₂
     end
     cache.fu = cache.fu2

src/default.jl

@@ -159,8 +159,8 @@ end
         ]
     else
         [
-            :(GeneralBroyden()),
             :(GeneralKlement()),
+            :(GeneralBroyden()),
             :(NewtonRaphson(; linsolve, precs, adkwargs...)),
             :(NewtonRaphson(; linsolve, precs, linesearch = BackTracking(), adkwargs...)),
             :(TrustRegion(; linsolve, precs, adkwargs...)),

src/gaussnewton.jl

@@ -46,7 +46,7 @@ function set_ad(alg::GaussNewton{CJ}, ad) where {CJ}
     return GaussNewton{CJ}(ad, alg.linsolve, alg.precs)
 end
 
-function GaussNewton(; concrete_jac = nothing, linsolve = CholeskyFactorization(),
+function GaussNewton(; concrete_jac = nothing, linsolve = nothing,
     precs = DEFAULT_PRECS, adkwargs...)
     ad = default_adargs_to_adtype(; adkwargs...)
     return GaussNewton{_unwrap_val(concrete_jac)}(ad, linsolve, precs)
@@ -81,15 +81,25 @@ function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_::
     kwargs...) where {uType, iip}
     alg = get_concrete_algorithm(alg_, prob)
     @unpack f, u0, p = prob
+
+    linsolve_with_JᵀJ = Val(_needs_square_A(alg, u0))
+
     u = alias_u0 ? u0 : deepcopy(u0)
     if iip
         fu1 = f.resid_prototype === nothing ? zero(u) : f.resid_prototype
         f(fu1, u, p)
     else
         fu1 = f(u, p)
     end
-    uf, linsolve, J, fu2, jac_cache, du, JᵀJ, Jᵀf = jacobian_caches(alg, f, u, p, Val(iip);
-        linsolve_with_JᵀJ = Val(true))
+
+    if SciMLBase._unwrap_val(linsolve_with_JᵀJ)
+        uf, linsolve, J, fu2, jac_cache, du, JᵀJ, Jᵀf = jacobian_caches(alg, f, u, p,
+            Val(iip); linsolve_with_JᵀJ)
+    else
+        uf, linsolve, J, fu2, jac_cache, du = jacobian_caches(alg, f, u, p,
+            Val(iip); linsolve_with_JᵀJ)
+        JᵀJ, Jᵀf = nothing, nothing
+    end
 
     return GaussNewtonCache{iip}(f, alg, u, fu1, fu2, zero(fu1), du, p, uf, linsolve, J,
         JᵀJ, Jᵀf, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol,
@@ -99,12 +109,20 @@ end
 function perform_step!(cache::GaussNewtonCache{true})
     @unpack u, fu1, f, p, alg, J, JᵀJ, Jᵀf, linsolve, du = cache
     jacobian!!(J, cache)
-    __matmul!(JᵀJ, J', J)
-    __matmul!(Jᵀf, J', fu1)
 
-    # u = u - J \ fu
-    linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(JᵀJ), b = _vec(Jᵀf),
-        linu = _vec(du), p, reltol = cache.abstol)
+    if JᵀJ !== nothing
+        __matmul!(JᵀJ, J', J)
+        __matmul!(Jᵀf, J', fu1)
+    end
+
+    # u = u - JᵀJ \ Jᵀfu
+    if cache.JᵀJ === nothing
+        linres = dolinsolve(alg.precs, linsolve; A = J, b = _vec(fu1), linu = _vec(du),
+            p, reltol = cache.abstol)
+    else
+        linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(JᵀJ), b = _vec(Jᵀf),
+            linu = _vec(du), p, reltol = cache.abstol)
+    end
     cache.linsolve = linres.cache
     @. u = u - du
     f(cache.fu_new, u, p)
@@ -125,14 +143,22 @@ function perform_step!(cache::GaussNewtonCache{false})
 
     cache.J = jacobian!!(cache.J, cache)
 
-    cache.JᵀJ = cache.J' * cache.J
-    cache.Jᵀf = cache.J' * fu1
+    if cache.JᵀJ !== nothing
+        cache.JᵀJ = cache.J' * cache.J
+        cache.Jᵀf = cache.J' * fu1
+    end
+
     # u = u - J \ fu
     if linsolve === nothing
         cache.du = fu1 / cache.J
     else
-        linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(cache.JᵀJ),
-            b = _vec(cache.Jᵀf), linu = _vec(cache.du), p, reltol = cache.abstol)
+        if cache.JᵀJ === nothing
+            linres = dolinsolve(alg.precs, linsolve; A = cache.J, b = _vec(fu1),
+                linu = _vec(cache.du), p, reltol = cache.abstol)
+        else
+            linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(cache.JᵀJ),
+                b = _vec(cache.Jᵀf), linu = _vec(cache.du), p, reltol = cache.abstol)
+        end
         cache.linsolve = linres.cache
     end
     cache.u = @. u - cache.du  # `u` might not support mutation

src/jacobian.jl

@@ -50,8 +50,8 @@ jacobian!!(::Number, cache) = last(value_derivative(cache.uf, cache.u))
 
 # Build Jacobian Caches
 function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{iip};
-    linsolve_kwargs = (;),
-    linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false)) where {iip, needsJᵀJ}
+    linsolve_kwargs = (;), lininit::Val{linsolve_init} = Val(true),
+    linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false)) where {iip, needsJᵀJ, linsolve_init}
     uf = JacobianWrapper{iip}(f, p)
 
     haslinsolve = hasfield(typeof(alg), :linsolve)
@@ -95,25 +95,28 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii
         Jᵀfu = J' * _vec(fu)
     end
 
-    linprob = LinearProblem(needsJᵀJ ? __maybe_symmetric(JᵀJ) : J,
-        needsJᵀJ ? _vec(Jᵀfu) : _vec(fu); u0 = _vec(du))
-
-    if alg isa PseudoTransient
-        alpha = convert(eltype(u), alg.alpha_initial)
-        J_new = J - (1 / alpha) * I
-        linprob = LinearProblem(J_new, _vec(fu); u0 = _vec(du))
+    if linsolve_init
+        linprob_A = alg isa PseudoTransient ?
+                    (J - (1 / (convert(eltype(u), alg.alpha_initial))) * I) :
+                    (needsJᵀJ ? __maybe_symmetric(JᵀJ) : J)
+        linsolve = __setup_linsolve(linprob_A, needsJᵀJ ? Jᵀfu : fu, du, p, alg)
+    else
+        linsolve = nothing
     end
 
+    needsJᵀJ && return uf, linsolve, J, fu, jac_cache, du, JᵀJ, Jᵀfu
+    return uf, linsolve, J, fu, jac_cache, du
+end
+
+function __setup_linsolve(A, b, u, p, alg)
+    linprob = LinearProblem(A, _vec(b); u0 = _vec(u))
+
     weight = similar(u)
     recursivefill!(weight, true)
 
-    Pl, Pr = wrapprecs(alg.precs(needsJᵀJ ? __maybe_symmetric(JᵀJ) : J, nothing, u, p,
-            nothing, nothing, nothing, nothing, nothing)..., weight)
-    linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr,
-        linsolve_kwargs...)
-
-    needsJᵀJ && return uf, linsolve, J, fu, jac_cache, du, JᵀJ, Jᵀfu
-    return uf, linsolve, J, fu, jac_cache, du
+    Pl, Pr = wrapprecs(alg.precs(A, nothing, u, p, nothing, nothing, nothing, nothing,
+            nothing)..., weight)
+    return init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr)
 end
 
 __get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff()

src/klement.jl

@@ -27,6 +27,10 @@ solves.
     linesearch
 end
 
+function set_linsolve(alg::GeneralKlement, linsolve)
+    return GeneralKlement(alg.max_resets, linsolve, alg.precs, alg.linesearch)
+end
+
 function GeneralKlement(; max_resets::Int = 5, linsolve = nothing,
     linesearch = LineSearch(), precs = DEFAULT_PRECS)
     linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch)
@@ -60,30 +64,27 @@ end
 
 get_fu(cache::GeneralKlementCache) = cache.fu
 
-function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralKlement, args...;
+function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralKlement, args...;
     alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM,
     linsolve_kwargs = (;), kwargs...) where {uType, iip}
     @unpack f, u0, p = prob
     u = alias_u0 ? u0 : deepcopy(u0)
     fu = evaluate_f(prob, u)
     J = __init_identity_jacobian(u, fu)
+    du = _mutable_zero(u)
 
     if u isa Number
         linsolve = nothing
+        alg = alg_
     else
         # For General Julia Arrays default to LU Factorization
-        linsolve_alg = alg.linsolve === nothing && u isa Array ? LUFactorization() :
+        linsolve_alg = alg_.linsolve === nothing && u isa Array ? LUFactorization() :
                        nothing
-        weight = similar(u)
-        recursivefill!(weight, true)
-        Pl, Pr = wrapprecs(alg.precs(J, nothing, u, p, nothing, nothing, nothing, nothing,
-                nothing)..., weight)
-        linprob = LinearProblem(J, _vec(fu); u0 = _vec(fu))
-        linsolve = init(linprob, linsolve_alg; alias_A = true, alias_b = true, Pl, Pr,
-            linsolve_kwargs...)
+        alg = set_linsolve(alg_, linsolve_alg)
+        linsolve = __setup_linsolve(J, _vec(fu), _vec(du), p, alg)
     end
 
-    return GeneralKlementCache{iip}(f, alg, u, fu, zero(fu), _mutable_zero(u), p, linsolve,
+    return GeneralKlementCache{iip}(f, alg, u, fu, zero(fu), du, p, linsolve,
         J, zero(J), zero(J), _vec(zero(fu)), _vec(zero(fu)), 0, false,
         maxiters, internalnorm, ReturnCode.Default, abstol, prob, NLStats(1, 0, 0, 0, 0),
         init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)))
@@ -114,7 +115,7 @@ function perform_step!(cache::GeneralKlementCache{true})
 
     # Line Search
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(cache.fu2, u, p)
 
     cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true)

src/lbroyden.jl

@@ -93,7 +93,7 @@ function perform_step!(cache::LimitedMemoryBroydenCache{true})
     T = eltype(u)
 
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(cache.fu2, u, p)
 
     cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true)
@@ -123,8 +123,8 @@ function perform_step!(cache::LimitedMemoryBroydenCache{true})
         __lbroyden_matvec!(_vec(cache.vᵀ_cache), cache.Ux, U_part, Vᵀ_part, _vec(cache.du))
         __lbroyden_rmatvec!(_vec(cache.u_cache), cache.xᵀVᵀ, U_part, Vᵀ_part,
             _vec(cache.dfu))
-        cache.u_cache .= (du .- cache.u_cache) ./
-                         (dot(cache.vᵀ_cache, cache.dfu) .+ T(1e-5))
+        denom = dot(cache.vᵀ_cache, cache.dfu)
+        cache.u_cache .= (du .- cache.u_cache) ./ ifelse(iszero(denom), T(1e-5), denom)
 
         idx = mod1(cache.iterations_since_reset + 1, size(cache.U, 1))
         selectdim(cache.U, 1, idx) .= _vec(cache.u_cache)
@@ -179,8 +179,8 @@ function perform_step!(cache::LimitedMemoryBroydenCache{false})
             __lbroyden_matvec(U_part, Vᵀ_part, _vec(cache.du)))
         cache.u_cache = _restructure(cache.u_cache,
             __lbroyden_rmatvec(U_part, Vᵀ_part, _vec(cache.dfu)))
-        cache.u_cache = (cache.du .- cache.u_cache) ./
-                        (dot(cache.vᵀ_cache, cache.dfu) .+ T(1e-5))
+        denom = dot(cache.vᵀ_cache, cache.dfu)
+        cache.u_cache = (cache.du .- cache.u_cache) ./ ifelse(iszero(denom), T(1e-5), denom)
 
         idx = mod1(cache.iterations_since_reset + 1, size(cache.U, 1))
         selectdim(cache.U, 1, idx) .= _vec(cache.u_cache)
@@ -249,12 +249,12 @@ end
         return nothing
     end
     mul!(xᵀVᵀ[:, 1:η], x', Vᵀ)
-    mul!(y', xᵀVᵀ[:, 1:η], U)
+    mul!(reshape(y, 1, :), xᵀVᵀ[:, 1:η], U)
     return nothing
 end
 
 @views function __lbroyden_rmatvec(U::AbstractMatrix, Vᵀ::AbstractMatrix, x::AbstractVector)
     # Computes xᵀ × Vᵀ × U
     size(U, 1) == 0 && return x
-    return (x' * Vᵀ) * U
+    return (reshape(x, 1, :) * Vᵀ) * U
 end