Integrate the linear solve better into the algorithm

Author: avik-pal

src/core/approximate_jacobian.jl

@@ -291,6 +291,27 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip};
                 cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...)
         end
     end
+
+    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

src/core/generalized_first_order.jl

@@ -230,6 +230,27 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB};
                 cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...)
         end
     end
+
+    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

src/descent/common.jl

@@ -1,5 +1,6 @@
 """
-    DescentResult(; δu = missing, u = missing, success::Bool = true, extras = (;))
+    DescentResult(; δu = missing, u = missing, success::Bool = true,
+        linsolve_success::Bool = true, extras = (;))
 
 Construct a `DescentResult` object.
 
@@ -9,6 +10,7 @@ Construct a `DescentResult` object.
   - `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.
@@ -17,10 +19,12 @@ Construct a `DescentResult` object.
     δu
     u
     success::Bool
+    linsolve_success::Bool
     extras
 end
 
-function DescentResult(; δu = missing, u = missing, success::Bool = true, extras = (;))
+function DescentResult(; δu = missing, u = missing, success::Bool = true,
+        linsolve_success::Bool = true, extras = (;))
     @assert δu !== missing || u !== missing
-    return DescentResult(δu, u, success, extras)
+    return DescentResult(δu, u, success, linsolve_success, extras)
 end

src/descent/damped_newton.jl

@@ -200,10 +200,14 @@ 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

src/descent/newton.jl

@@ -81,10 +81,14 @@ function __internal_solve!(
         @bb δu = J × vec(fu)
     else
         @static_timeit cache.timer "linear solve" begin
-            δu = cache.lincache(;
+            linres = cache.lincache(;
                 A = J, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu),
                 reuse_A_if_factorization = !new_jacobian || (idx !== Val(1)))
-            δ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
     end
     @bb @. δu *= -1
@@ -102,10 +106,15 @@ function __internal_solve!(
     end
     @bb cache.Jᵀfu_cache = transpose(J) × vec(fu)
     @static_timeit cache.timer "linear solve" begin
-        δu = cache.lincache(; A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache,
+        linres = cache.lincache(;
+            A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache,
             kwargs..., linu = _vec(δu), du = _vec(δu),
             reuse_A_if_factorization = !new_jacobian || (idx !== Val(1)))
-        δ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)

src/descent/steepest.jl

@@ -54,10 +54,14 @@ function __internal_solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val
     if INV
         A = J === nothing ? nothing : transpose(J)
         @static_timeit cache.timer "linear solve" begin
-            δu = cache.lincache(;
+            linres = cache.lincache(;
                 A, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu),
                 reuse_A_if_factorization = !new_jacobian || idx !== Val(1))
-            δ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
     else
         @assert J!==nothing "`J` must be provided when `pre_inverted = Val(false)`."

src/internal/linear_solve.jl

@@ -105,6 +105,11 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...)
     return LinearSolverCache(lincache, linsolve, nothing, nothing, nothing, precs, 0, 0)
 end
 
+@kwdef @concrete struct LinearSolveResult
+    u
+    success::Bool = true
+end
+
 # Direct Linear Solve Case without Caching
 function (cache::LinearSolverCache{Nothing})(;
         A = nothing, b = nothing, linu = nothing, kwargs...)
@@ -119,7 +124,7 @@ function (cache::LinearSolverCache{Nothing})(;
     else
         res = cache.A \ cache.b
     end
-    return res
+    return LinearSolveResult(; u = res)
 end
 
 # Use LinearSolve.jl
@@ -154,11 +159,7 @@ function (cache::LinearSolverCache)(;
         cache.lincache.Pr = Pr
     end
 
-    # display(A)
-
     linres = solve!(cache.lincache)
-    # @show cache.lincache.cacheval
-    # @show LinearAlgebra.issuccess(cache.lincache.cacheval)
     cache.lincache = linres.cache
     # Unfortunately LinearSolve.jl doesn't have the most uniform ReturnCode handling
     if linres.retcode === ReturnCode.Failure
@@ -185,11 +186,14 @@ function (cache::LinearSolverCache)(;
             end
             linres = solve!(cache.additional_lincache)
             cache.additional_lincache = linres.cache
-            return linres.u
+            linres.retcode === ReturnCode.Failure &&
+                return LinearSolveResult(; u = linres.u, success = false)
+            return LinearSolveResult(; u = linres.u)
         end
+        return LinearSolveResult(; u = linres.u, success = false)
     end
 
-    return linres.u
+    return LinearSolveResult(; u = linres.u)
 end
 
 @inline __update_A!(cache::LinearSolverCache, ::Nothing, reuse) = cache