Merge pull request #414 from SciML/ap/linearsolve_retcode

Robust Singular Jacobian Handling

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.10.1"
+version = "3.11.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -60,7 +60,7 @@ NonlinearSolveZygoteExt = "Zygote"
 ADTypes = "0.2.6"
 Aqua = "0.8"
 ArrayInterface = "7.9"
-BandedMatrices = "1.4"
+BandedMatrices = "1.5"
 BenchmarkTools = "1.4"
 CUDA = "5.2"
 ConcreteStructs = "0.2.3"
@@ -76,7 +76,7 @@ LazyArrays = "1.8.2"
 LeastSquaresOptim = "0.8.5"
 LineSearches = "7.2"
 LinearAlgebra = "1.10"
-LinearSolve = "2.21"
+LinearSolve = "2.29"
 MINPACK = "1.2"
 MaybeInplace = "0.1.1"
 NLSolvers = "0.5"

docs/src/devdocs/internal_interfaces.md

@@ -15,6 +15,12 @@ NonlinearSolve.AbstractDescentAlgorithm
 NonlinearSolve.AbstractDescentCache
 ```
 
+## Descent Results
+
+```@docs
+NonlinearSolve.DescentResult
+```
+
 ## Approximate Jacobian
 
 ```@docs

src/NonlinearSolve.jl

@@ -12,8 +12,8 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
           LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf, SciMLBase,
           SimpleNonlinearSolve, SparseArrays, SparseDiffTools
 
-    import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing,
-                           ismutable
+    import ArrayInterface: ArrayInterface, undefmatrix, can_setindex, restructure,
+                           fast_scalar_indexing, ismutable
     import DiffEqBase: AbstractNonlinearTerminationMode,
                        AbstractSafeNonlinearTerminationMode,
                        AbstractSafeBestNonlinearTerminationMode,
@@ -50,6 +50,7 @@ include("adtypes.jl")
 include("timer_outputs.jl")
 include("internal/helpers.jl")
 
+include("descent/common.jl")
 include("descent/newton.jl")
 include("descent/steepest.jl")
 include("descent/dogleg.jl")
@@ -135,10 +136,10 @@ include("default.jl")
 
     @compile_workload begin
         for prob in probs_nls, alg in nls_algs
-            solve(prob, alg; abstol = 1e-2)
+            solve(prob, alg; abstol = 1e-2, verbose = false)
         end
         for prob in probs_nlls, alg in nlls_algs
-            solve(prob, alg; abstol = 1e-2)
+            solve(prob, alg; abstol = 1e-2, verbose = false)
         end
     end
 end

src/abstract_types.jl

@@ -66,7 +66,7 @@ Abstract Type for all Descent Caches.
 ### `__internal_solve!` specification
 
 ```julia
-δu, success, intermediates = __internal_solve!(
+descent_result = __internal_solve!(
     cache::AbstractDescentCache, J, fu, u, idx::Val; skip_solve::Bool = false, kwargs...)
 ```
 
@@ -81,12 +81,7 @@ Abstract Type for all Descent Caches.
 
 #### Returned values
 
-  - `δu`: the descent direction.
-  - `success`: Certain Descent Algorithms can reject a descent direction for example
-    `GeodesicAcceleration`.
-  - `intermediates`: A named tuple containing intermediates computed during the solve.
-    For example, `GeodesicAcceleration` returns `NamedTuple{(:v, :a)}` containing the
-    "velocity" and "acceleration" terms.
+  - `descent_result`: Result in a [`DescentResult`](@ref).
 
 ### Interface Functions
 

src/core/approximate_jacobian.jl

@@ -114,6 +114,7 @@ end
     retcode::ReturnCode.T
     force_stop::Bool
     force_reinit::Bool
+    kwargs
 end
 
 store_inverse_jacobian(::ApproximateJacobianSolveCache{INV}) where {INV} = INV
@@ -211,10 +212,10 @@ function SciMLBase.__init(
 
         return ApproximateJacobianSolveCache{INV, GB, iip, maxtime !== nothing}(
             fu, u, u_cache, p, du, J, alg, prob, initialization_cache,
-            descent_cache, linesearch_cache, trustregion_cache,
-            update_rule_cache, reinit_rule_cache, inv_workspace, 0, 0, 0,
-            alg.max_resets, maxiters, maxtime, alg.max_shrink_times, 0, timer,
-            0.0, termination_cache, trace, ReturnCode.Default, false, false)
+            descent_cache, linesearch_cache, trustregion_cache, update_rule_cache,
+            reinit_rule_cache, inv_workspace, 0, 0, 0, alg.max_resets,
+            maxiters, maxtime, alg.max_shrink_times, 0, timer, 0.0,
+            termination_cache, trace, ReturnCode.Default, false, false, kwargs)
     end
 end
 
@@ -282,16 +283,38 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip};
     @static_timeit cache.timer "descent" begin
         if cache.trustregion_cache !== nothing &&
            hasfield(typeof(cache.trustregion_cache), :trust_region)
-            δu, descent_success, descent_intermediates = __internal_solve!(
+            descent_result = __internal_solve!(
                 cache.descent_cache, J, cache.fu, cache.u; new_jacobian,
-                trust_region = cache.trustregion_cache.trust_region)
+                trust_region = cache.trustregion_cache.trust_region, cache.kwargs...)
         else
-            δu, descent_success, descent_intermediates = __internal_solve!(
-                cache.descent_cache, J, cache.fu, cache.u; new_jacobian)
+            descent_result = __internal_solve!(
+                cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...)
         end
     end
 
-    if descent_success
+    if !descent_result.linsolve_success
+        if new_jacobian && cache.steps_since_last_reset == 0
+            # Extremely pathological case. Jacobian was just reset and linear solve
+            # failed. Should ideally never happen in practice unless true jacobian init
+            # is used.
+            cache.retcode = LinearSolveFailureCode
+            cache.force_stop = true
+            return
+        else
+            # Force a reinit because the problem is currently un-solvable
+            if !haskey(cache.kwargs, :verbose) || cache.kwargs[:verbose]
+                @warn "Linear Solve Failed but Jacobian Information is not current. \
+                       Retrying with reinitialized Approximate Jacobian."
+            end
+            cache.force_reinit = true
+            __step!(cache; recompute_jacobian = true)
+            return
+        end
+    end
+
+    δu, descent_intermediates = descent_result.δu, descent_result.extras
+
+    if descent_result.success
         if GB === :LineSearch
             @static_timeit cache.timer "linesearch" begin
                 needs_reset, α = __internal_solve!(cache.linesearch_cache, cache.u, δu)

src/core/generalized_first_order.jl

@@ -111,6 +111,7 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ
     trace
     retcode::ReturnCode.T
     force_stop::Bool
+    kwargs
 end
 
 SymbolicIndexingInterface.state_values(cache::GeneralizedFirstOrderAlgorithmCache) = cache.u
@@ -201,8 +202,8 @@ function SciMLBase.__init(
 
         return GeneralizedFirstOrderAlgorithmCache{iip, GB, maxtime !== nothing}(
             fu, u, u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache,
-            trustregion_cache, 0, 0, maxiters, maxtime, alg.max_shrink_times,
-            timer, 0.0, true, termination_cache, trace, ReturnCode.Default, false)
+            trustregion_cache, 0, 0, maxiters, maxtime, alg.max_shrink_times, timer,
+            0.0, true, termination_cache, trace, ReturnCode.Default, false, kwargs)
     end
 end
 
@@ -221,16 +222,38 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB};
     @static_timeit cache.timer "descent" begin
         if cache.trustregion_cache !== nothing &&
            hasfield(typeof(cache.trustregion_cache), :trust_region)
-            δu, descent_success, descent_intermediates = __internal_solve!(
+            descent_result = __internal_solve!(
                 cache.descent_cache, J, cache.fu, cache.u; new_jacobian,
-                trust_region = cache.trustregion_cache.trust_region)
+                trust_region = cache.trustregion_cache.trust_region, cache.kwargs...)
         else
-            δu, descent_success, descent_intermediates = __internal_solve!(
-                cache.descent_cache, J, cache.fu, cache.u; new_jacobian)
+            descent_result = __internal_solve!(
+                cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...)
         end
     end
 
-    if descent_success
+    if !descent_result.linsolve_success
+        if new_jacobian
+            # Jacobian Information is current and linear solve failed terminate the solve
+            cache.retcode = LinearSolveFailureCode
+            cache.force_stop = true
+            return
+        else
+            # Jacobian Information is not current and linear solve failed, recompute
+            # Jacobian
+            if !haskey(cache.kwargs, :verbose) || cache.kwargs[:verbose]
+                @warn "Linear Solve Failed but Jacobian Information is not current. \
+                       Retrying with updated Jacobian."
+            end
+            # In the 2nd call the `new_jacobian` is guaranteed to be `true`.
+            cache.make_new_jacobian = true
+            __step!(cache; recompute_jacobian = true, kwargs...)
+            return
+        end
+    end
+
+    δu, descent_intermediates = descent_result.δu, descent_result.extras
+
+    if descent_result.success
         cache.make_new_jacobian = true
         if GB === :LineSearch
             @static_timeit cache.timer "linesearch" begin

src/core/spectral_methods.jl

@@ -74,6 +74,7 @@ concrete_jac(::GeneralizedDFSane) = nothing
     trace
     retcode::ReturnCode.T
     force_stop::Bool
+    kwargs
 end
 
 function __reinit_internal!(
@@ -150,7 +151,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane
         return GeneralizedDFSaneCache{isinplace(prob), maxtime !== nothing}(
             fu, fu_cache, u, u_cache, prob.p, du, alg, prob, σ_n, T(alg.σ_min),
             T(alg.σ_max), linesearch_cache, 0, 0, maxiters, maxtime,
-            timer, 0.0, tc_cache, trace, ReturnCode.Default, false)
+            timer, 0.0, tc_cache, trace, ReturnCode.Default, false, kwargs)
     end
 end
 

src/descent/common.jl

@@ -0,0 +1,30 @@
+"""
+    DescentResult(; δu = missing, u = missing, success::Bool = true,
+        linsolve_success::Bool = true, extras = (;))
+
+Construct a `DescentResult` object.
+
+### Keyword Arguments
+
+  - `δu`: The descent direction.
+  - `u`: The new iterate. This is provided only for multi-step methods currently.
+  - `success`: Certain Descent Algorithms can reject a descent direction for example
+    [`GeodesicAcceleration`](@ref).
+  - `linsolve_success`: Whether the line search was successful.
+  - `extras`: A named tuple containing intermediates computed during the solve.
+    For example, [`GeodesicAcceleration`](@ref) returns `NamedTuple{(:v, :a)}` containing
+    the "velocity" and "acceleration" terms.
+"""
+@concrete struct DescentResult
+    δu
+    u
+    success::Bool
+    linsolve_success::Bool
+    extras
+end
+
+function DescentResult(; δu = missing, u = missing, success::Bool = true,
+        linsolve_success::Bool = true, extras = (;))
+    @assert δu !== missing || u !== missing
+    return DescentResult(δu, u, success, linsolve_success, extras)
+end

src/descent/damped_newton.jl

@@ -135,7 +135,7 @@ function __internal_solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu,
         u, idx::Val{N} = Val(1); skip_solve::Bool = false,
         new_jacobian::Bool = true, kwargs...) where {INV, N, mode}
     δu = get_du(cache, idx)
-    skip_solve && return δu, true, (;)
+    skip_solve && return DescentResult(; δu)
 
     recompute_A = idx === Val(1)
 
@@ -200,15 +200,19 @@ function __internal_solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu,
     end
 
     @static_timeit cache.timer "linear solve" begin
-        δu = cache.lincache(;
+        linres = cache.lincache(;
             A, b, reuse_A_if_factorization = !new_jacobian && !recompute_A,
             kwargs..., linu = _vec(δu))
-        δu = _restructure(get_du(cache, idx), δu)
+        δu = _restructure(get_du(cache, idx), linres.u)
+        if !linres.success
+            set_du!(cache, δu, idx)
+            return DescentResult(; δu, success = false, linsolve_success = false)
+        end
     end
 
     @bb @. δu *= -1
     set_du!(cache, δu, idx)
-    return δu, true, (;)
+    return DescentResult(; δu)
 end
 
 # Define special concatenation for certain Array combinations

src/descent/dogleg.jl

@@ -81,14 +81,14 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} =
                                     want to use a Trust Region."
     δu = get_du(cache, idx)
     T = promote_type(eltype(u), eltype(fu))
-    δu_newton, _, _ = __internal_solve!(
-        cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...)
+    δu_newton = __internal_solve!(
+        cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...).δu
 
     # Newton's Step within the trust region
     if cache.internalnorm(δu_newton) ≤ trust_region
         @bb copyto!(δu, δu_newton)
         set_du!(cache, δu, idx)
-        return δu, true, (; δuJᵀJδu = T(NaN))
+        return DescentResult(; δu, extras = (; δuJᵀJδu = T(NaN)))
     end
 
     # Take intersection of steepest descent direction and trust region if Cauchy point lies
@@ -102,8 +102,8 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} =
         @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy)
         δuJᵀJδu = __dot(δu_cauchy, cache.δu_cache_mul)
     else
-        δu_cauchy, _, _ = __internal_solve!(
-            cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...)
+        δu_cauchy = __internal_solve!(
+            cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...).δu
         J_ = INV ? inv(J) : J
         l_grad = cache.internalnorm(δu_cauchy)
         @bb cache.JᵀJ_cache = J × vec(δu_cauchy)  # TODO: Rename
@@ -115,7 +115,7 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} =
         λ = trust_region / l_grad
         @bb @. δu = λ * δu_cauchy
         set_du!(cache, δu, idx)
-        return δu, true, (; δuJᵀJδu = λ^2 * δuJᵀJδu)
+        return DescentResult(; δu, extras = (; δuJᵀJδu = λ^2 * δuJᵀJδu))
     end
 
     # FIXME: For anything other than 2-norm a quadratic root will give incorrect results
@@ -133,5 +133,5 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} =
 
     @bb @. δu = cache.δu_cache_1 + τ * cache.δu_cache_2
     set_du!(cache, δu, idx)
-    return δu, true, (; δuJᵀJδu = T(NaN))
+    return DescentResult(; δu, extras = (; δuJᵀJδu = T(NaN)))
 end

src/descent/geodesic_acceleration.jl

@@ -105,9 +105,9 @@ end
 function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N} = Val(1);
         skip_solve::Bool = false, kwargs...) where {N}
     a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx)
-    skip_solve && return δu, true, (; a, v)
-    v, _, _ = __internal_solve!(
-        cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...)
+    skip_solve && return DescentResult(; δu, extras = (; a, v))
+    v = __internal_solve!(
+        cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...).δu
 
     @bb @. cache.u_cache = u + cache.h * v
     cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p)
@@ -116,8 +116,8 @@ function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{
     Jv = _restructure(cache.fu_cache, cache.Jv)
     @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv)
 
-    a, _, _ = __internal_solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N);
-        skip_solve, kwargs..., reuse_A_if_factorization = true)
+    a = __internal_solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N);
+        skip_solve, kwargs..., reuse_A_if_factorization = true).δu
 
     norm_v = cache.internalnorm(v)
     norm_a = cache.internalnorm(a)
@@ -130,5 +130,5 @@ function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{
         cache.last_step_accepted = false
     end
 
-    return δu, cache.last_step_accepted, (; a, v)
+    return DescentResult(; δu, success = cache.last_step_accepted, extras = (; a, v))
 end