Use the new Termination Condition API

Author: avik-pal

Manifest.toml

@@ -40,7 +40,7 @@ ADTypes = "0.2"
 ArrayInterface = "6.0.24, 7"
 BandedMatrices = "1"
 ConcreteStructs = "0.2"
-DiffEqBase = "6.130"
+DiffEqBase = "6.136"
 EnumX = "1"
 Enzyme = "0.11"
 FastBroadcast = "0.1.9, 0.2"
@@ -56,7 +56,7 @@ RecursiveArrayTools = "2"
 Reexport = "0.2, 1"
 SciMLBase = "2.4"
 SimpleNonlinearSolve = "0.1.23"
-SparseDiffTools = "2.6"
+SparseDiffTools = "2.9"
 StaticArraysCore = "1.4"
 UnPack = "1.0"
 Zygote = "0.6"

Project.toml

@@ -30,6 +30,7 @@ PrecompileTools.@recompile_invalidations begin
 end
 
 @reexport using ADTypes, LineSearches, SciMLBase, SimpleNonlinearSolve
+import DiffEqBase: AbstractNonlinearTerminationMode
 
 const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences,
     ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode}

src/NonlinearSolve.jl

@@ -52,7 +52,7 @@ end
     reset_check
     prob
     stats::NLStats
-    lscache
+    ls_cache
     termination_condition
     tc_storage
 end
@@ -93,7 +93,7 @@ function perform_step!(cache::GeneralBroydenCache{true})
     T = eltype(u)
 
     mul!(_vec(du), J⁻¹, _vec(fu))
-    α = perform_linesearch!(cache.lscache, u, du)
+    α = perform_linesearch!(cache.ls_cache, u, du)
     _axpy!(-α, du, u)
     f(fu2, u, p)
 
@@ -137,7 +137,7 @@ function perform_step!(cache::GeneralBroydenCache{false})
     T = eltype(cache.u)
 
     cache.du = _restructure(cache.du, cache.J⁻¹ * _vec(cache.fu))
-    α = perform_linesearch!(cache.lscache, cache.u, cache.du)
+    α = perform_linesearch!(cache.ls_cache, cache.u, cache.du)
     cache.u = cache.u .- α * cache.du
     cache.fu2 = f(cache.u, p)
 

src/broyden.jl

@@ -61,7 +61,7 @@ end
     reltol
     prob
     stats::NLStats
-    lscache
+    ls_cache
     termination_condition
     tc_storage
 end
@@ -131,7 +131,7 @@ function perform_step!(cache::GeneralKlementCache{true})
     cache.linsolve = linres.cache
 
     # Line Search
-    α = perform_linesearch!(cache.lscache, u, du)
+    α = perform_linesearch!(cache.ls_cache, u, du)
     _axpy!(-α, du, u)
     f(cache.fu2, u, p)
 
@@ -193,7 +193,7 @@ function perform_step!(cache::GeneralKlementCache{false})
     end
 
     # Line Search
-    α = perform_linesearch!(cache.lscache, cache.u, cache.du)
+    α = perform_linesearch!(cache.ls_cache, cache.u, cache.du)
     cache.u = @. cache.u - α * cache.du  # `u` might not support mutation
     cache.fu2 = f(cache.u, p)
 

src/klement.jl

@@ -59,7 +59,7 @@ end
     reset_check
     prob
     stats::NLStats
-    lscache
+    ls_cache
     termination_condition
     tc_storage
 end
@@ -109,7 +109,7 @@ function perform_step!(cache::LimitedMemoryBroydenCache{true})
 
     termination_condition = cache.termination_condition(tc_storage)
 
-    α = perform_linesearch!(cache.lscache, u, du)
+    α = perform_linesearch!(cache.ls_cache, u, du)
     _axpy!(-α, du, u)
     f(cache.fu2, u, p)
 
@@ -169,7 +169,7 @@ function perform_step!(cache::LimitedMemoryBroydenCache{false})
 
     T = eltype(cache.u)
 
-    α = perform_linesearch!(cache.lscache, cache.u, cache.du)
+    α = perform_linesearch!(cache.ls_cache, cache.u, cache.du)
     cache.u = cache.u .- α * cache.du
     cache.fu2 = f(cache.u, p)
 

src/lbroyden.jl

@@ -70,9 +70,8 @@ end
     reltol
     prob
     stats::NLStats
-    lscache
-    termination_condition
-    tc_storage
+    ls_cache
+    tc_cache
 end
 
 function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphson, args...;
@@ -86,39 +85,34 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphso
     uf, linsolve, J, fu2, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip);
         linsolve_kwargs)
 
-    abstol, reltol, termination_condition = _init_termination_elements(abstol,
-        reltol, termination_condition, eltype(u))
-
-    mode = DiffEqBase.get_termination_mode(termination_condition)
-
-    storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() :
-              nothing
+    abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, u,
+        termination_condition)
 
     return NewtonRaphsonCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, uf, linsolve, J,
         jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob,
         NLStats(1, 0, 0, 0, 0),
-        init_linesearch_cache(alg.linesearch, f, u, p, fu1, Val(iip)),
-        termination_condition, storage)
+        init_linesearch_cache(alg.linesearch, f, u, p, fu1, Val(iip)), tc_cache)
 end
 
 function perform_step!(cache::NewtonRaphsonCache{true})
-    @unpack u, u_prev, fu1, f, p, alg, J, linsolve, du, tc_storage = cache
+    @unpack u, u_prev, fu1, f, p, alg, J, linsolve, du = cache
     jacobian!!(J, cache)
 
-    termination_condition = cache.termination_condition(tc_storage)
-
     # u = u - J \ fu
     linres = dolinsolve(alg.precs, linsolve; A = J, b = _vec(fu1), linu = _vec(du),
         p, reltol = cache.abstol)
     cache.linsolve = linres.cache
 
     # Line Search
-    α = perform_linesearch!(cache.lscache, u, du)
+    α = perform_linesearch!(cache.ls_cache, u, du)
     _axpy!(-α, du, u)
     f(cache.fu1, u, p)
 
-    termination_condition(cache.fu1, u, u_prev, cache.abstol, cache.reltol) &&
-        (cache.force_stop = true)
+    if cache.tc_cache(cache.fu1, cache.u, u_prev)
+        # Stores the best solution in cache!
+        cache.tc_cache.u !== nothing && copyto!(cache.u, cache.tc_cache.u)
+        cache.force_stop = true
+    end
 
     @. u_prev = u
     cache.stats.nf += 1
@@ -129,9 +123,7 @@ function perform_step!(cache::NewtonRaphsonCache{true})
 end
 
 function perform_step!(cache::NewtonRaphsonCache{false})
-    @unpack u, u_prev, fu1, f, p, alg, linsolve, tc_storage = cache
-
-    termination_condition = cache.termination_condition(tc_storage)
+    @unpack u, u_prev, fu1, f, p, alg, linsolve = cache
 
     cache.J = jacobian!!(cache.J, cache)
     # u = u - J \ fu
@@ -144,14 +136,17 @@ function perform_step!(cache::NewtonRaphsonCache{false})
     end
 
     # Line Search
-    α = perform_linesearch!(cache.lscache, u, cache.du)
+    α = perform_linesearch!(cache.ls_cache, u, cache.du)
     cache.u = @. u - α * cache.du  # `u` might not support mutation
     cache.fu1 = f(cache.u, p)
 
-    termination_condition(cache.fu1, cache.u, u_prev, cache.abstol, cache.reltol) &&
-        (cache.force_stop = true)
+    if cache.tc_cache(cache.fu1, cache.u, u_prev)
+        # Stores the best solution in cache!
+        cache.tc_cache.u !== nothing && (cache.u = cache.tc_cache.u)
+        cache.force_stop = true
+    end
 
-    cache.u_prev = @. cache.u
+    cache.u_prev = cache.u
     cache.stats.nf += 1
     cache.stats.njacs += 1
     cache.stats.nsolve += 1
@@ -161,7 +156,7 @@ end
 
 function SciMLBase.reinit!(cache::NewtonRaphsonCache{iip}, u0 = cache.u; p = cache.p,
         abstol = cache.abstol, reltol = cache.reltol,
-        termination_condition = cache.termination_condition,
+        termination_condition = get_termination_mode(cache.tc_cache),
         maxiters = cache.maxiters) where {iip}
     cache.p = p
     if iip
@@ -173,12 +168,12 @@ function SciMLBase.reinit!(cache::NewtonRaphsonCache{iip}, u0 = cache.u; p = cac
         cache.fu1 = 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.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/raphson.jl

@@ -1,17 +1,4 @@
-
-@inline UNITLESS_ABS2(x) = real(abs2(x))
-@inline DEFAULT_NORM(u::Union{AbstractFloat, Complex}) = @fastmath abs(u)
-@inline function DEFAULT_NORM(u::Array{T}) where {T <: Union{AbstractFloat, Complex}}
-    return sqrt(real(sum(abs2, u)) / length(u))
-end
-@inline function DEFAULT_NORM(u::StaticArray{<:Union{AbstractFloat, Complex}})
-    return sqrt(real(sum(abs2, u)) / length(u))
-end
-@inline function DEFAULT_NORM(u::AbstractVectorOfArray)
-    return sum(sqrt(real(sum(UNITLESS_ABS2, _u)) / length(_u)) for _u in u.u)
-end
-@inline DEFAULT_NORM(u::AbstractArray) = sqrt(real(sum(UNITLESS_ABS2, u)) / length(u))
-@inline DEFAULT_NORM(u) = norm(u)
+const DEFAULT_NORM = DiffEqBase.NONLINEARSOLVE_DEFAULT_NORM
 
 # Ignores NaN
 function __findmin(f, x)
@@ -36,7 +23,7 @@ code.
     `autodiff=<ADTypes>`.
 """
 function default_adargs_to_adtype(; chunk_size = missing, autodiff = nothing,
-    standardtag = missing, diff_type = missing)
+        standardtag = missing, diff_type = missing)
     # If using the new API short circuit
     autodiff === nothing && return nothing
     autodiff isa ADTypes.AbstractADType && return autodiff
@@ -89,8 +76,8 @@ end
 DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, cachedata) = nothing, nothing
 
 function dolinsolve(precs::P, linsolve; A = nothing, linu = nothing, b = nothing,
-    du = nothing, u = nothing, p = nothing, t = nothing, weight = nothing,
-    cachedata = nothing, reltol = nothing) where {P}
+        du = nothing, u = nothing, p = nothing, t = nothing, weight = nothing,
+        cachedata = nothing, reltol = nothing) where {P}
     A !== nothing && (linsolve.A = A)
     b !== nothing && (linsolve.b = b)
     linu !== nothing && (linsolve.u = linu)
@@ -167,7 +154,7 @@ _maybe_mutable(x, _) = x
 
 # Helper function to get value of `f(u, p)`
 function evaluate_f(prob::Union{NonlinearProblem{uType, iip},
-        NonlinearLeastSquaresProblem{uType, iip}}, u) where {uType, iip}
+            NonlinearLeastSquaresProblem{uType, iip}}, u) where {uType, iip}
     @unpack f, u0, p = prob
     if iip
         fu = f.resid_prototype === nothing ? zero(u) : f.resid_prototype
@@ -194,7 +181,7 @@ end
 
 Defaults to `mul!(C, A, B)`. However, for sparse matrices uses `C .= A * B`.
 """
-__matmul!(C, A, B) = mul!(C, A, B)
+__matmul!(C, A, B) = mul!(C, A, B)``
 __matmul!(C::AbstractSparseMatrix, A, B) = C .= A * B
 
 # Concretize Algorithms
@@ -216,11 +203,20 @@ function __get_concrete_algorithm(alg, prob)
         use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff()
     else
         (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)(;
-            tag = NonlinearSolveTag())
+            tag = ForwardDiff.Tag(NonlinearSolveTag(), eltype(prob.u0)))
     end
     return set_ad(alg, ad)
 end
 
+function init_termination_cache(abstol, reltol, u, ::Nothing)
+    return init_termination_cache(abstol, reltol, u, AbsNormTerminationMode())
+end
+function init_termination_cache(abstol, reltol, u, tc::AbstractNonlinearTerminationMode)
+    tc_cache = init(u, tc; abstol, reltol)
+    return DiffEqBase.get_abstol(tc_cache), DiffEqBase.get_reltol(tc_cache), tc_cache
+end
+
+# FIXME: Remove the functions below when we have migrated to the new type stable API
 __cvt_real(::Type{T}, ::Nothing) where {T} = nothing
 __cvt_real(::Type{T}, x) where {T} = real(T(x))
 
@@ -231,7 +227,7 @@ function _get_tolerance(η, tc_η, ::Type{T}) where {T}
 end
 
 function _init_termination_elements(abstol, reltol, termination_condition,
-    ::Type{T}; mode = NLSolveTerminationMode.AbsNorm) where {T}
+        ::Type{T}; mode = NLSolveTerminationMode.AbsNorm) where {T}
     if termination_condition !== nothing
         if abstol !== nothing && abstol != termination_condition.abstol
             error("Incompatible absolute tolerances found. The tolerances supplied as the \
@@ -279,6 +275,7 @@ function _get_reinit_termination_condition(cache, abstol, reltol, termination_co
             termination_condition.safe_termination_options)
     end
 end
+# FIXME: Purge things till here!
 
 __init_identity_jacobian(u::Number, _) = u
 function __init_identity_jacobian(u, fu)