Improve termination conditions

Author: avik-pal

Manifest.toml

@@ -164,7 +164,7 @@ uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
 
 [[deps.DiffEqBase]]
 deps = ["ArrayInterface", "ChainRulesCore", "DataStructures", "DocStringExtensions", "EnumX", "EnzymeCore", "FastBroadcast", "ForwardDiff", "FunctionWrappers", "FunctionWrappersWrappers", "LinearAlgebra", "Logging", "Markdown", "MuladdMacro", "Parameters", "PreallocationTools", "PrecompileTools", "Printf", "RecursiveArrayTools", "Reexport", "Requires", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Static", "StaticArraysCore", "Statistics", "Tricks", "TruncatedStacktraces", "ZygoteRules"]
-git-tree-sha1 = "e5049e32074cd22f86d74036caf6663637623003"
+git-tree-sha1 = "4e661d0beddac31da05e71b79afd769232622de8"
 repo-rev = "ap/tstable_termination"
 repo-url = "https://github.com/SciML/DiffEqBase.jl"
 uuid = "2b5f629d-d688-5b77-993f-72d75c75574e"
@@ -689,9 +689,9 @@ version = "0.1.0"
 
 [[deps.SLEEFPirates]]
 deps = ["IfElse", "Static", "VectorizationBase"]
-git-tree-sha1 = "897b39ec056c0619ea87adc7eeadba0bec0cf931"
+git-tree-sha1 = "f5c896d781486f1d67c8492f0e0ead2c3517208c"
 uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa"
-version = "0.6.40"
+version = "0.6.41"
 
 [[deps.SciMLBase]]
 deps = ["ADTypes", "ArrayInterface", "ChainRulesCore", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "FillArrays", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "PrecompileTools", "Preferences", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables", "TruncatedStacktraces", "ZygoteRules"]

src/NonlinearSolve.jl

@@ -30,7 +30,9 @@ PrecompileTools.@recompile_invalidations begin
 end
 
 @reexport using ADTypes, LineSearches, SciMLBase, SimpleNonlinearSolve
-import DiffEqBase: AbstractNonlinearTerminationMode
+import DiffEqBase: AbstractNonlinearTerminationMode,
+    AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode,
+    NonlinearSafeTerminationReturnCode, get_termination_mode
 
 const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences,
     ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode}
@@ -54,6 +56,8 @@ function not_terminated(cache::AbstractNonlinearSolveCache)
     return !cache.force_stop && cache.stats.nsteps < cache.maxiters
 end
 get_fu(cache::AbstractNonlinearSolveCache) = cache.fu1
+set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu1 = fu)
+get_u(cache::AbstractNonlinearSolveCache) = cache.u
 
 function SciMLBase.solve!(cache::AbstractNonlinearSolveCache)
     while not_terminated(cache)
@@ -70,7 +74,7 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache)
         end
     end
 
-    return SciMLBase.build_solution(cache.prob, cache.alg, cache.u, get_fu(cache);
+    return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache);
         cache.retcode, cache.stats)
 end
 
@@ -114,4 +118,10 @@ export RobustMultiNewton, FastShortcutNonlinearPolyalg
 
 export LineSearch, LiFukushimaLineSearch
 
+# Export the termination conditions from DiffEqBase
+export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode,
+    NormTerminationMode, RelTerminationMode, RelNormTerminationMode, AbsTerminationMode,
+    AbsNormTerminationMode, RelSafeTerminationMode, AbsSafeTerminationMode,
+    RelSafeBestTerminationMode, AbsSafeBestTerminationMode
+
 end # module

src/broyden.jl

@@ -53,11 +53,11 @@ end
     prob
     stats::NLStats
     ls_cache
-    termination_condition
-    tc_storage
+    tc_cache
 end
 
 get_fu(cache::GeneralBroydenCache) = cache.fu
+set_fu!(cache::GeneralBroydenCache, fu) = (cache.fu = fu)
 
 function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralBroyden, args...;
         alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing,
@@ -71,34 +71,26 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralBroyde
                       alg.reset_tolerance
     reset_check = x -> abs(x) ≤ reset_tolerance
 
-    abstol, reltol, termination_condition = _init_termination_elements(abstol, reltol,
-        termination_condition, eltype(u))
-
-    mode = DiffEqBase.get_termination_mode(termination_condition)
+    abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u,
+        termination_condition)
 
-    storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() :
-              nothing
     return GeneralBroydenCache{iip}(f, alg, u, zero(u), _mutable_zero(u), fu, zero(fu),
         zero(fu), p, J⁻¹, zero(_reshape(fu, 1, :)), _mutable_zero(u), false, 0,
         alg.max_resets, maxiters, internalnorm, ReturnCode.Default, abstol, reltol,
         reset_tolerance, reset_check, prob, NLStats(1, 0, 0, 0, 0),
-        init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), termination_condition,
-        storage)
+        init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache)
 end
 
 function perform_step!(cache::GeneralBroydenCache{true})
-    @unpack f, p, du, fu, fu2, dfu, u, u_prev, J⁻¹, J⁻¹df, J⁻¹₂, tc_storage = cache
-
-    termination_condition = cache.termination_condition(tc_storage)
+    @unpack f, p, du, fu, fu2, dfu, u, u_prev, J⁻¹, J⁻¹df, J⁻¹₂ = cache
     T = eltype(u)
 
     mul!(_vec(du), J⁻¹, _vec(fu))
     α = perform_linesearch!(cache.ls_cache, u, du)
     _axpy!(-α, du, u)
     f(fu2, u, p)
 
-    termination_condition(fu2, u, u_prev, cache.abstol, cache.reltol) &&
-        (cache.force_stop = true)
+    check_and_update!(cache, fu2, u, u_prev)
     cache.stats.nf += 1
 
     cache.force_stop && return nothing
@@ -130,9 +122,7 @@ function perform_step!(cache::GeneralBroydenCache{true})
 end
 
 function perform_step!(cache::GeneralBroydenCache{false})
-    @unpack f, p, tc_storage = cache
-
-    termination_condition = cache.termination_condition(tc_storage)
+    @unpack f, p = cache
 
     T = eltype(cache.u)
 
@@ -141,8 +131,7 @@ function perform_step!(cache::GeneralBroydenCache{false})
     cache.u = cache.u .- α * cache.du
     cache.fu2 = f(cache.u, p)
 
-    termination_condition(cache.fu2, cache.u, cache.u_prev, cache.abstol, cache.reltol) &&
-        (cache.force_stop = true)
+    check_and_update!(cache, cache.fu2, cache.u, cache.u_prev)
     cache.stats.nf += 1
 
     cache.force_stop && return nothing
@@ -172,9 +161,8 @@ function perform_step!(cache::GeneralBroydenCache{false})
 end
 
 function SciMLBase.reinit!(cache::GeneralBroydenCache{iip}, u0 = cache.u; p = cache.p,
-        abstol = cache.abstol, reltol = cache.reltol,
-        termination_condition = cache.termination_condition,
-        maxiters = cache.maxiters) where {iip}
+        abstol = cache.abstol, reltol = cache.reltol, maxiters = cache.maxiters,
+        termination_condition = get_termination_mode(cache.tc_cache)) where {iip}
     cache.p = p
     if iip
         recursivecopy!(cache.u, u0)
@@ -185,12 +173,12 @@ function SciMLBase.reinit!(cache::GeneralBroydenCache{iip}, u0 = cache.u; p = ca
         cache.fu = cache.f(cache.u, p)
     end
 
-    termination_condition = _get_reinit_termination_condition(cache, abstol, reltol,
+    abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fu, cache.u,
         termination_condition)
 
     cache.abstol = abstol
     cache.reltol = reltol
-    cache.termination_condition = termination_condition
+    cache.tc_cache = tc_cache
     cache.maxiters = maxiters
     cache.stats.nf = 1
     cache.stats.nsteps = 1

src/dfsane.jl

@@ -63,7 +63,8 @@ function DFSane(; σ_min = 1e-10, σ_max = 1e+10, σ_1 = 1.0, M = 10, γ = 1e-4,
         n_exp, η_strategy, max_inner_iterations)
 end
 
-@concrete mutable struct DFSaneCache{iip}
+# FIXME: Someone please make this code conform to the style of the remaining solvers
+@concrete mutable struct DFSaneCache{iip} <: AbstractNonlinearSolveCache{iip}
     alg
     uₙ
     uₙ₋₁
@@ -91,10 +92,13 @@ end
     reltol
     prob
     stats::NLStats
-    termination_condition
-    tc_storage
+    tc_cache
 end
 
+get_fu(cache::DFSaneCache) = cache.fuₙ
+set_fu!(cache::DFSaneCache, fu) = (cache.fuₙ = fu)
+get_u(cache::DFSaneCache) = cache.uₙ
+
 function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::DFSane, args...;
         alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing,
         termination_condition = nothing, internalnorm::F = DEFAULT_NORM,
@@ -124,24 +128,18 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::DFSane, args.
 
     ℋ = fill(f₍ₙₒᵣₘ₎ₙ₋₁, M)
 
-    abstol, reltol, termination_condition = _init_termination_elements(abstol, reltol,
-        termination_condition, T)
-
-    mode = DiffEqBase.get_termination_mode(termination_condition)
-
-    storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() :
-              nothing
+    abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fuₙ₋₁, uₙ₋₁,
+        termination_condition)
 
     return DFSaneCache{iip}(alg, uₙ, uₙ₋₁, fuₙ, fuₙ₋₁, 𝒹, ℋ, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀,
         M, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, p, false, maxiters,
         internalnorm, ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0),
-        termination_condition, storage)
+        tc_cache)
 end
 
 function perform_step!(cache::DFSaneCache{true})
-    @unpack alg, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, M, tc_storage = cache
+    @unpack alg, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, M = cache
 
-    termination_condition = cache.termination_condition(tc_storage)
     f = (dx, x) -> cache.prob.f(dx, x, cache.p)
 
     T = eltype(cache.uₙ)
@@ -184,9 +182,7 @@ function perform_step!(cache::DFSaneCache{true})
         f₍ₙₒᵣₘ₎ₙ = norm(cache.fuₙ)^nₑₓₚ
     end
 
-    if termination_condition(cache.fuₙ, cache.uₙ, cache.uₙ₋₁, cache.abstol, cache.reltol)
-        cache.force_stop = true
-    end
+    check_and_update!(cache, cache.fuₙ, cache.uₙ, cache.uₙ₋₁)
 
     # Update spectral parameter
     @. cache.uₙ₋₁ = cache.uₙ - cache.uₙ₋₁
@@ -215,9 +211,8 @@ function perform_step!(cache::DFSaneCache{true})
 end
 
 function perform_step!(cache::DFSaneCache{false})
-    @unpack alg, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, M, tc_storage = cache
+    @unpack alg, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, M = cache
 
-    termination_condition = cache.termination_condition(tc_storage)
     f = x -> cache.prob.f(x, cache.p)
 
     T = eltype(cache.uₙ)
@@ -260,9 +255,7 @@ function perform_step!(cache::DFSaneCache{false})
         f₍ₙₒᵣₘ₎ₙ = norm(cache.fuₙ)^nₑₓₚ
     end
 
-    if termination_condition(cache.fuₙ, cache.uₙ, cache.uₙ₋₁, cache.abstol, cache.reltol)
-        cache.force_stop = true
-    end
+    check_and_update!(cache, cache.fuₙ, cache.uₙ, cache.uₙ₋₁)
 
     # Update spectral parameter
     cache.uₙ₋₁ = @. cache.uₙ - cache.uₙ₋₁
@@ -290,26 +283,9 @@ function perform_step!(cache::DFSaneCache{false})
     return nothing
 end
 
-function SciMLBase.solve!(cache::DFSaneCache)
-    while !cache.force_stop && cache.stats.nsteps < cache.maxiters
-        cache.stats.nsteps += 1
-        perform_step!(cache)
-    end
-
-    if cache.stats.nsteps == cache.maxiters
-        cache.retcode = ReturnCode.MaxIters
-    else
-        cache.retcode = ReturnCode.Success
-    end
-
-    return SciMLBase.build_solution(cache.prob, cache.alg, cache.uₙ, cache.fuₙ;
-        retcode = cache.retcode, stats = cache.stats)
-end
-
 function SciMLBase.reinit!(cache::DFSaneCache{iip}, u0 = cache.uₙ; p = cache.p,
-        abstol = cache.abstol, reltol = cache.reltol,
-        termination_condition = cache.termination_condition,
-        maxiters = cache.maxiters) where {iip}
+        abstol = cache.abstol, reltol = cache.reltol, maxiters = cache.maxiters,
+        termination_condition = get_termination_mode(cache.tc_cache)) where {iip}
     cache.p = p
     if iip
         recursivecopy!(cache.uₙ, u0)
@@ -330,12 +306,12 @@ function SciMLBase.reinit!(cache::DFSaneCache{iip}, u0 = cache.uₙ; p = cache.p
     T = eltype(cache.uₙ)
     cache.σₙ = T(cache.alg.σ_1)
 
-    termination_condition = _get_reinit_termination_condition(cache, abstol, reltol,
+    abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fuₙ, cache.uₙ,
         termination_condition)
 
     cache.abstol = abstol
     cache.reltol = reltol
-    cache.termination_condition = termination_condition
+    cache.tc_cache = tc_cache
     cache.maxiters = maxiters
     cache.stats.nf = 1
     cache.stats.nsteps = 1