diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index e5e6960e2..4328af345 100644 --- a/src/core/approximate_jacobian.jl +++ b/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 diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index fd1eb5b3f..7d20b0bfc 100644 --- a/src/core/generalized_first_order.jl +++ b/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 diff --git a/src/descent/common.jl b/src/descent/common.jl index 2a614d84a..3d53573ea 100644 --- a/src/descent/common.jl +++ b/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 diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index aaabc6c4e..792f56773 100644 --- a/src/descent/damped_newton.jl +++ b/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 diff --git a/src/descent/newton.jl b/src/descent/newton.jl index bcb83686f..36edbf86c 100644 --- a/src/descent/newton.jl +++ b/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) diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 4d4877a1b..227fdb436 100644 --- a/src/descent/steepest.jl +++ b/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)`." diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 080de05ef..f2277a165 100644 --- a/src/internal/linear_solve.jl +++ b/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