From f9346802df7aa9b75d84e0e851336bba5231c7a6 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Wed, 13 Sep 2023 15:53:50 -0400 Subject: [PATCH 01/12] Start using termination conditions in newton raphson --- src/raphson.jl | 47 +++++++++++++++++++++++++++++++++++++---------- src/utils.jl | 5 +++++ 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/src/raphson.jl b/src/raphson.jl index 642c7ee36..a0bb2eb24 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -30,11 +30,12 @@ for large-scale and numerically-difficult nonlinear systems. which means that no line search is performed. Algorithms from `LineSearches.jl` can be used here directly, and they will be converted to the correct `LineSearch`. """ -@concrete struct NewtonRaphson{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} +@concrete struct NewtonRaphson{CJ, AD, TC <: NLSolveTerminationCondition} <: AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve precs linesearch + termination_condition::TC end function set_ad(alg::NewtonRaphson{CJ}, ad) where {CJ} @@ -42,16 +43,19 @@ function set_ad(alg::NewtonRaphson{CJ}, ad) where {CJ} end function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, - linesearch = LineSearch(), precs = DEFAULT_PRECS, adkwargs...) + linesearch = LineSearch(), precs = DEFAULT_PRECS, termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; + abstol = nothing, + reltol = nothing), adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) - return NewtonRaphson{_unwrap_val(concrete_jac)}(ad, linsolve, precs, linesearch) + return NewtonRaphson{_unwrap_val(concrete_jac)}(ad, linsolve, precs, linesearch, termination_condition) end @concrete mutable struct NewtonRaphsonCache{iip} <: AbstractNonlinearSolveCache{iip} f alg u + uprev fu1 fu2 du @@ -65,9 +69,11 @@ end internalnorm retcode::ReturnCode.T abstol + reltol prob stats::NLStats lscache + tc_storage end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphson, args...; @@ -80,15 +86,28 @@ 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) - return NewtonRaphsonCache{iip}(f, alg, u, fu1, fu2, du, p, uf, linsolve, J, - jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, prob, - NLStats(1, 0, 0, 0, 0), LineSearchCache(alg.linesearch, f, u, p, fu1, Val(iip))) + + tc = alg.termination_condition + mode = DiffEqBase.get_termination_mode(tc) + + atol = _get_tolerance(abstol, tc.abstol, eltype(u)) + rtol = _get_tolerance(reltol, tc.reltol, eltype(u)) + + storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : + nothing + + return NewtonRaphsonCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, uf, linsolve, J, + jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, prob, + NLStats(1, 0, 0, 0, 0), LineSearchCache(alg.linesearch, f, u, p, fu1, Val(iip)), storage) end function perform_step!(cache::NewtonRaphsonCache{true}) - @unpack u, fu1, f, p, alg, J, linsolve, du = cache + @unpack u, uprev, fu1, f, p, alg, J, linsolve, du = cache jacobian!!(J, cache) + tc_storage = cache.tc_storage + termination_condition = cache.alg.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) @@ -99,7 +118,9 @@ function perform_step!(cache::NewtonRaphsonCache{true}) @. u = u - α * du f(cache.fu1, u, p) - cache.internalnorm(fu1) < cache.abstol && (cache.force_stop = true) + termination_condition(cache.fu1, u, uprev, cache.abstol, cache.reltol) && (cache.force_stop = true) + + @. uprev = u cache.stats.nf += 1 cache.stats.njacs += 1 cache.stats.nsolve += 1 @@ -108,7 +129,11 @@ function perform_step!(cache::NewtonRaphsonCache{true}) end function perform_step!(cache::NewtonRaphsonCache{false}) - @unpack u, fu1, f, p, alg, linsolve = cache + @unpack u, uprev, fu1, f, p, alg, linsolve = cache + + tc_storage = cache.tc_storage + termination_condition = cache.alg.termination_condition(tc_storage) + cache.J = jacobian!!(cache.J, cache) # u = u - J \ fu @@ -125,7 +150,9 @@ function perform_step!(cache::NewtonRaphsonCache{false}) cache.u = @. u - α * cache.du # `u` might not support mutation cache.fu1 = f(cache.u, p) - cache.internalnorm(fu1) < cache.abstol && (cache.force_stop = true) + termination_condition(cache.fu1, cache.u, uprev, cache.abstol, cache.reltol) && (cache.force_stop = true) + + cache.uprev = cache.u cache.stats.nf += 1 cache.stats.njacs += 1 cache.stats.nsolve += 1 diff --git a/src/utils.jl b/src/utils.jl index 173c279be..134843055 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -208,3 +208,8 @@ function __get_concrete_algorithm(alg, prob) end return set_ad(alg, ad) end + +function _get_tolerance(η, tc_η, ::Type{T}) where {T} + @show fallback_η + return ifelse(η !== nothing, η, ifelse(tc_η !== nothing, tc_η, fallback_η)) +end From 37b7015f1b3c8d7b26fd74ff73c7e4e9faa4e5a9 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Wed, 13 Sep 2023 16:06:42 -0400 Subject: [PATCH 02/12] typo --- src/utils.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/utils.jl b/src/utils.jl index 134843055..887a92dc6 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -210,6 +210,6 @@ function __get_concrete_algorithm(alg, prob) end function _get_tolerance(η, tc_η, ::Type{T}) where {T} - @show fallback_η + fallback_η = real(oneunit(T)) * (eps(real(one(T))))^(4 // 5) return ifelse(η !== nothing, η, ifelse(tc_η !== nothing, tc_η, fallback_η)) end From efb8c5ef324322ebde0b74fc7af7da9d116e99e0 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Mon, 25 Sep 2023 21:18:59 -0400 Subject: [PATCH 03/12] Update all algorithms to use termination condition --- src/NonlinearSolve.jl | 2 +- src/levenberg.jl | 58 +++++++++++++++++++++++++++++------- src/raphson.jl | 40 +++++++++++++++---------- src/trustRegion.jl | 68 +++++++++++++++++++++++++++++++++++-------- test/basictests.jl | 42 ++++++++++++++++++++++++++ 5 files changed, 170 insertions(+), 40 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 3a38989d4..b0fb28fb1 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -26,7 +26,7 @@ const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences, ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode} abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end -abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end +abstract type AbstractNewtonAlgorithm{CJ, AD, TC} <: AbstractNonlinearSolveAlgorithm end abstract type AbstractNonlinearSolveCache{iip} end diff --git a/src/levenberg.jl b/src/levenberg.jl index f35f35cb2..460377245 100644 --- a/src/levenberg.jl +++ b/src/levenberg.jl @@ -74,7 +74,8 @@ numerically-difficult nonlinear systems. [this paper](https://arxiv.org/abs/1201.5885) to use a minimum value of the elements in `DᵀD` to prevent the damping from being too small. Defaults to `1e-8`. """ -@concrete struct LevenbergMarquardt{CJ, AD, T} <: AbstractNewtonAlgorithm{CJ, AD} +@concrete struct LevenbergMarquardt{CJ, AD, T, TC <: NLSolveTerminationCondition} <: + AbstractNewtonAlgorithm{CJ, AD, TC} ad::AD linsolve precs @@ -85,6 +86,7 @@ numerically-difficult nonlinear systems. α_geodesic::T b_uphill::T min_damping_D::T + termination_condition::TC end function set_ad(alg::LevenbergMarquardt{CJ}, ad) where {CJ} @@ -97,17 +99,22 @@ function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, damping_initial::Real = 1.0, damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, finite_diff_step_geodesic::Real = 0.1, α_geodesic::Real = 0.75, b_uphill::Real = 1.0, min_damping_D::AbstractFloat = 1e-8, + termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; + abstol = nothing, + reltol = nothing), adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) return LevenbergMarquardt{_unwrap_val(concrete_jac)}(ad, linsolve, precs, damping_initial, damping_increase_factor, damping_decrease_factor, - finite_diff_step_geodesic, α_geodesic, b_uphill, min_damping_D) + finite_diff_step_geodesic, α_geodesic, b_uphill, min_damping_D, + termination_condition) end @concrete mutable struct LevenbergMarquardtCache{iip} <: AbstractNonlinearSolveCache{iip} f alg u + u_prev fu1 fu2 du @@ -121,6 +128,7 @@ end internalnorm retcode::ReturnCode.T abstol + reltol prob DᵀD JᵀJ @@ -145,11 +153,13 @@ end Jv mat_tmp stats::NLStats + tc_storage end function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, - NonlinearLeastSquaresProblem{uType, iip}}, alg_::LevenbergMarquardt, - args...; alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + NonlinearLeastSquaresProblem{uType, iip}}, alg_::LevenbergMarquardt, + args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @unpack f, u0, p = prob @@ -184,15 +194,30 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, fu_tmp = zero(fu1) mat_tmp = zero(JᵀJ) - return LevenbergMarquardtCache{iip}(f, alg, u, fu1, fu2, du, p, uf, linsolve, J, - jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, prob, DᵀD, + tc = alg.termination_condition + mode = DiffEqBase.get_termination_mode(tc) + + atol = _get_tolerance(abstol, tc.abstol, eltype(u)) + rtol = _get_tolerance(reltol, tc.reltol, eltype(u)) + + storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : + nothing + + return LevenbergMarquardtCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, uf, linsolve, + J, + jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, prob, DᵀD, JᵀJ, λ, λ_factor, damping_increase_factor, damping_decrease_factor, h, α_geodesic, b_uphill, min_damping_D, v, a, tmp_vec, v_old, loss, δ, loss, make_new_J, fu_tmp, - zero(u), zero(fu1), mat_tmp, NLStats(1, 0, 0, 0, 0)) + zero(u), zero(fu1), mat_tmp, NLStats(1, 0, 0, 0, 0), storage) + end function perform_step!(cache::LevenbergMarquardtCache{true}) @unpack fu1, f, make_new_J = cache + + tc_storage = cache.tc_storage + termination_condition = cache.alg.termination_condition(tc_storage) + if iszero(fu1) cache.force_stop = true return nothing @@ -205,7 +230,7 @@ function perform_step!(cache::LevenbergMarquardtCache{true}) cache.make_new_J = false cache.stats.njacs += 1 end - @unpack u, p, λ, JᵀJ, DᵀD, J, alg, linsolve = cache + @unpack u, u_prev, p, λ, JᵀJ, DᵀD, J, alg, linsolve = cache # Usual Levenberg-Marquardt step ("velocity"). # The following lines do: cache.v = -cache.mat_tmp \ cache.u_tmp @@ -246,7 +271,11 @@ function perform_step!(cache::LevenbergMarquardtCache{true}) if (1 - β)^b_uphill * loss ≤ loss_old # Accept step. cache.u .+= δ - if loss < cache.abstol + if termination_condition(cache.fu_tmp, + cache.u, + u_prev, + cache.abstol, + cache.reltol) cache.force_stop = true return nothing end @@ -258,6 +287,7 @@ function perform_step!(cache::LevenbergMarquardtCache{true}) cache.make_new_J = true end end + @. u_prev = u cache.λ *= cache.λ_factor cache.λ_factor = cache.damping_increase_factor return nothing @@ -265,6 +295,10 @@ end function perform_step!(cache::LevenbergMarquardtCache{false}) @unpack fu1, f, make_new_J = cache + + tc_storage = cache.tc_storage + termination_condition = cache.alg.termination_condition(tc_storage) + if iszero(fu1) cache.force_stop = true return nothing @@ -281,7 +315,8 @@ function perform_step!(cache::LevenbergMarquardtCache{false}) cache.make_new_J = false cache.stats.njacs += 1 end - @unpack u, p, λ, JᵀJ, DᵀD, J, linsolve, alg = cache + + @unpack u, u_prev, p, λ, JᵀJ, DᵀD, J, linsolve, alg = cache cache.mat_tmp = JᵀJ + λ * DᵀD # Usual Levenberg-Marquardt step ("velocity"). @@ -322,7 +357,7 @@ function perform_step!(cache::LevenbergMarquardtCache{false}) if (1 - β)^b_uphill * loss ≤ loss_old # Accept step. cache.u += δ - if loss < cache.abstol + if termination_condition(fu_new, cache.u, u_prev, cache.abstol, cache.reltol) cache.force_stop = true return nothing end @@ -334,6 +369,7 @@ function perform_step!(cache::LevenbergMarquardtCache{false}) cache.make_new_J = true end end + cache.u_prev = @. cache.u cache.λ *= cache.λ_factor cache.λ_factor = cache.damping_increase_factor return nothing diff --git a/src/raphson.jl b/src/raphson.jl index a0bb2eb24..57547baca 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -30,7 +30,8 @@ for large-scale and numerically-difficult nonlinear systems. which means that no line search is performed. Algorithms from `LineSearches.jl` can be used here directly, and they will be converted to the correct `LineSearch`. """ -@concrete struct NewtonRaphson{CJ, AD, TC <: NLSolveTerminationCondition} <: AbstractNewtonAlgorithm{CJ, AD} +@concrete struct NewtonRaphson{CJ, AD, TC <: NLSolveTerminationCondition} <: + AbstractNewtonAlgorithm{CJ, AD, TC} ad::AD linsolve precs @@ -43,19 +44,24 @@ function set_ad(alg::NewtonRaphson{CJ}, ad) where {CJ} end function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, - linesearch = LineSearch(), precs = DEFAULT_PRECS, termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; - abstol = nothing, - reltol = nothing), adkwargs...) + linesearch = LineSearch(), precs = DEFAULT_PRECS, + termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; + abstol = nothing, + reltol = nothing), adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) - return NewtonRaphson{_unwrap_val(concrete_jac)}(ad, linsolve, precs, linesearch, termination_condition) + return NewtonRaphson{_unwrap_val(concrete_jac)}(ad, + linsolve, + precs, + linesearch, + termination_condition) end @concrete mutable struct NewtonRaphsonCache{iip} <: AbstractNonlinearSolveCache{iip} f alg u - uprev + u_prev fu1 fu2 du @@ -77,7 +83,8 @@ end end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphson, args...; - alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @unpack f, u0, p = prob @@ -86,7 +93,6 @@ 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) - tc = alg.termination_condition mode = DiffEqBase.get_termination_mode(tc) @@ -98,11 +104,12 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphso return NewtonRaphsonCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, uf, linsolve, J, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, prob, - NLStats(1, 0, 0, 0, 0), LineSearchCache(alg.linesearch, f, u, p, fu1, Val(iip)), storage) + NLStats(1, 0, 0, 0, 0), LineSearchCache(alg.linesearch, f, u, p, fu1, Val(iip)), + storage) end function perform_step!(cache::NewtonRaphsonCache{true}) - @unpack u, uprev, fu1, f, p, alg, J, linsolve, du = cache + @unpack u, u_prev, fu1, f, p, alg, J, linsolve, du = cache jacobian!!(J, cache) tc_storage = cache.tc_storage @@ -118,9 +125,10 @@ function perform_step!(cache::NewtonRaphsonCache{true}) @. u = u - α * du f(cache.fu1, u, p) - termination_condition(cache.fu1, u, uprev, cache.abstol, cache.reltol) && (cache.force_stop = true) + termination_condition(cache.fu1, u, u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) - @. uprev = u + @. u_prev = u cache.stats.nf += 1 cache.stats.njacs += 1 cache.stats.nsolve += 1 @@ -129,12 +137,11 @@ function perform_step!(cache::NewtonRaphsonCache{true}) end function perform_step!(cache::NewtonRaphsonCache{false}) - @unpack u, uprev, fu1, f, p, alg, linsolve = cache + @unpack u, u_prev, fu1, f, p, alg, linsolve = cache tc_storage = cache.tc_storage termination_condition = cache.alg.termination_condition(tc_storage) - cache.J = jacobian!!(cache.J, cache) # u = u - J \ fu if linsolve === nothing @@ -150,9 +157,10 @@ function perform_step!(cache::NewtonRaphsonCache{false}) cache.u = @. u - α * cache.du # `u` might not support mutation cache.fu1 = f(cache.u, p) - termination_condition(cache.fu1, cache.u, uprev, cache.abstol, cache.reltol) && (cache.force_stop = true) + termination_condition(cache.fu1, cache.u, u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) - cache.uprev = cache.u + cache.u_prev = @. cache.u cache.stats.nf += 1 cache.stats.njacs += 1 cache.stats.nsolve += 1 diff --git a/src/trustRegion.jl b/src/trustRegion.jl index 769f5e75c..bfdc2a557 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -148,7 +148,8 @@ for large-scale and numerically-difficult nonlinear systems. `linsolve` and `precs` are used exclusively for the inplace version of the algorithm. Support for the OOP version is planned! """ -@concrete struct TrustRegion{CJ, AD, MTR} <: AbstractNewtonAlgorithm{CJ, AD} +@concrete struct TrustRegion{CJ, AD, MTR, TC <: NLSolveTerminationCondition} <: + AbstractNewtonAlgorithm{CJ, AD, TC} ad::AD linsolve precs @@ -161,6 +162,7 @@ for large-scale and numerically-difficult nonlinear systems. shrink_factor::MTR expand_factor::MTR max_shrink_times::Int + termination_condition::TC end function set_ad(alg::TrustRegion{CJ}, ad) where {CJ} @@ -175,11 +177,15 @@ function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAU max_trust_radius::Real = 0 // 1, initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, shrink_factor::Real = 1 // 4, - expand_factor::Real = 2 // 1, max_shrink_times::Int = 32, adkwargs...) + expand_factor::Real = 2 // 1, max_shrink_times::Int = 32, + termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; + abstol = nothing, + reltol = nothing), adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) return TrustRegion{_unwrap_val(concrete_jac)}(ad, linsolve, precs, radius_update_scheme, max_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, - expand_threshold, shrink_factor, expand_factor, max_shrink_times) + expand_threshold, shrink_factor, expand_factor, max_shrink_times, + termination_condition) end @concrete mutable struct TrustRegionCache{iip, trustType, floatType} <: @@ -201,6 +207,7 @@ end internalnorm retcode::ReturnCode.T abstol + reltol prob radius_update_scheme::RadiusUpdateSchemes.T trust_r::trustType @@ -228,10 +235,12 @@ end p4::floatType ϵ::floatType stats::NLStats + tc_storage end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::TrustRegion, args...; - alias_u0 = false, maxiters = 1000, abstol = 1e-8, internalnorm = DEFAULT_NORM, + alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @unpack f, u0, p = prob @@ -333,13 +342,22 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::TrustRegion, initial_trust_radius = convert(trustType, 1.0) end + tc = alg.termination_condition + mode = DiffEqBase.get_termination_mode(tc) + + atol = _get_tolerance(abstol, tc.abstol, eltype(u)) + rtol = _get_tolerance(reltol, tc.reltol, eltype(u)) + + storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : + nothing + return TrustRegionCache{iip}(f, alg, u_prev, u, fu_prev, fu1, fu2, p, uf, linsolve, J, - jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, prob, + jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, prob, radius_update_scheme, initial_trust_radius, max_trust_radius, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, loss, loss_new, H, g, shrink_counter, du, u_tmp, u_gauss_newton, u_cauchy, fu_new, make_new_J, r, p1, p2, p3, p4, ϵ, - NLStats(1, 0, 0, 0, 0)) + NLStats(1, 0, 0, 0, 0), storage) end function perform_step!(cache::TrustRegionCache{true}) @@ -416,6 +434,10 @@ end function trust_region_step!(cache::TrustRegionCache) @unpack fu_new, du, g, H, loss, max_trust_r, radius_update_scheme = cache + + tc_storage = cache.tc_storage + termination_condition = cache.alg.termination_condition(tc_storage) + cache.loss_new = get_loss(fu_new) # Compute the ratio of the actual reduction to the predicted reduction. @@ -444,8 +466,11 @@ function trust_region_step!(cache::TrustRegionCache) # No need to make a new J, no step was taken, so we try again with a smaller trust_r cache.make_new_J = false end - - if iszero(cache.fu) || cache.internalnorm(cache.fu) < cache.abstol + if iszero(cache.fu) || termination_condition(cache.fu, + cache.u, + cache.u_prev, + cache.abstol, + cache.reltol) cache.force_stop = true end @@ -513,7 +538,12 @@ function trust_region_step!(cache::TrustRegionCache) cache.trust_r = rfunc(r, shrink_threshold, p1, p3, p4, p2) * cache.internalnorm(du) - if iszero(cache.fu) || cache.internalnorm(cache.fu) < cache.abstol || + if iszero(cache.fu) || + termination_condition(cache.fu, + cache.u, + cache.u_prev, + cache.abstol, + cache.reltol) || cache.internalnorm(g) < cache.ϵ cache.force_stop = true end @@ -538,7 +568,12 @@ function trust_region_step!(cache::TrustRegionCache) @unpack p1 = cache cache.trust_r = p1 * cache.internalnorm(jvp!(cache)) - if iszero(cache.fu) || cache.internalnorm(cache.fu) < cache.abstol || + if iszero(cache.fu) || + termination_condition(cache.fu, + cache.u, + cache.u_prev, + cache.abstol, + cache.reltol) || cache.internalnorm(g) < cache.ϵ cache.force_stop = true end @@ -562,7 +597,12 @@ function trust_region_step!(cache::TrustRegionCache) @unpack p1 = cache cache.trust_r = p1 * (cache.internalnorm(cache.fu)^0.99) - if iszero(cache.fu) || cache.internalnorm(cache.fu) < cache.abstol || + if iszero(cache.fu) || + termination_condition(cache.fu, + cache.u, + cache.u_prev, + cache.abstol, + cache.reltol) || cache.internalnorm(g) < cache.ϵ cache.force_stop = true end @@ -580,7 +620,11 @@ function trust_region_step!(cache::TrustRegionCache) cache.trust_r *= cache.p2 cache.shrink_counter += 1 end - if iszero(cache.fu) || cache.internalnorm(cache.fu) < cache.abstol + if iszero(cache.fu) || termination_condition(cache.fu, + cache.u, + cache.u_prev, + cache.abstol, + cache.reltol) cache.force_stop = true end end diff --git a/test/basictests.jl b/test/basictests.jl index 06cfc103d..250a2e8b4 100644 --- a/test/basictests.jl +++ b/test/basictests.jl @@ -123,6 +123,20 @@ end probN = NonlinearProblem(quadratic_f, u0, 2.0) @test all(solve(probN, NewtonRaphson(; autodiff)).u .≈ sqrt(2.0)) end + + @testset "Termination condition: $(mode) u0: $(_nameof(u0))" for mode in instances(NLSolveTerminationMode.T), + u0 in (1.0, [1.0, 1.0]) + + if mode ∈ + (NLSolveTerminationMode.SteadyStateDefault, NLSolveTerminationMode.RelSafeBest, + NLSolveTerminationMode.AbsSafeBest) + continue + end + termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, + reltol = nothing) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve(probN, NewtonRaphson(; termination_condition)).u .≈ sqrt(2.0)) + end end # --- TrustRegion tests --- @@ -281,6 +295,20 @@ end @test sol_iip.u ≈ sol_oop.u end end + + @testset "Termination condition: $(mode) u0: $(_nameof(u0))" for mode in instances(NLSolveTerminationMode.T), + u0 in (1.0, [1.0, 1.0]) + + if mode ∈ + (NLSolveTerminationMode.SteadyStateDefault, NLSolveTerminationMode.RelSafeBest, + NLSolveTerminationMode.AbsSafeBest) + continue + end + termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, + reltol = nothing) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve(probN, TrustRegion(; termination_condition)).u .≈ sqrt(2.0)) + end end # --- LevenbergMarquardt tests --- @@ -390,6 +418,20 @@ end @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-10) end end + + @testset "Termination condition: $(mode) u0: $(_nameof(u0))" for mode in instances(NLSolveTerminationMode.T), + u0 in (1.0, [1.0, 1.0]) + + if mode ∈ + (NLSolveTerminationMode.SteadyStateDefault, NLSolveTerminationMode.RelSafeBest, + NLSolveTerminationMode.AbsSafeBest) + continue + end + termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, + reltol = nothing) + probN2 = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve(probN, LevenbergMarquardt(; termination_condition)).u .≈ sqrt(2.0)) + end end # --- DFSane tests --- From b58355e024d0102a2208b1ad7ad94df6af18f0fc Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Tue, 26 Sep 2023 00:46:36 -0400 Subject: [PATCH 04/12] fixup! add tests --- Project.toml | 3 ++- test/basictests.jl | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index c3cb802a8..ff58f8560 100644 --- a/Project.toml +++ b/Project.toml @@ -74,6 +74,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" [targets] -test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt"] +test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "DiffEqBase"] diff --git a/test/basictests.jl b/test/basictests.jl index 250a2e8b4..7135b4901 100644 --- a/test/basictests.jl +++ b/test/basictests.jl @@ -1,5 +1,5 @@ using BenchmarkTools, LinearSolve, NonlinearSolve, StaticArrays, Random, LinearAlgebra, - Test, ForwardDiff, Zygote, Enzyme, SparseDiffTools + Test, ForwardDiff, Zygote, Enzyme, SparseDiffTools, DiffEqBase _nameof(x) = applicable(nameof, x) ? nameof(x) : _nameof(typeof(x)) @@ -429,7 +429,7 @@ end end termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, reltol = nothing) - probN2 = NonlinearProblem(quadratic_f, u0, 2.0) + probN = NonlinearProblem(quadratic_f, u0, 2.0) @test all(solve(probN, LevenbergMarquardt(; termination_condition)).u .≈ sqrt(2.0)) end end From 2763d65c008917260a035c0210ee7246e6db7ed9 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Tue, 10 Oct 2023 03:28:34 -0400 Subject: [PATCH 05/12] Add termination condition to gaussnewton and other fixes --- src/gaussnewton.jl | 62 ++++++++++++++++++++++++++++++++++++---------- src/levenberg.jl | 9 ++----- 2 files changed, 51 insertions(+), 20 deletions(-) diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl index 6dcb0f835..73ff5b8dc 100644 --- a/src/gaussnewton.jl +++ b/src/gaussnewton.jl @@ -36,26 +36,34 @@ for large-scale and numerically-difficult nonlinear least squares problems. Jacobian-Free version of `GaussNewton` doesn't work yet, and it forces jacobian construction. This will be fixed in the near future. """ -@concrete struct GaussNewton{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} +@concrete struct GaussNewton{CJ, AD, TC} <: AbstractNewtonAlgorithm{CJ, AD, TC} ad::AD linsolve precs + termination_condition::TC end function set_ad(alg::GaussNewton{CJ}, ad) where {CJ} return GaussNewton{CJ}(ad, alg.linsolve, alg.precs) end -function GaussNewton(; concrete_jac = nothing, linsolve = CholeskyFactorization(), - precs = DEFAULT_PRECS, adkwargs...) +function GaussNewton(; concrete_jac = nothing, linsolve = NormalCholeskyFactorization(), + precs = DEFAULT_PRECS, + termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.AbsNorm; + abstol = nothing, + reltol = nothing), adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) - return GaussNewton{_unwrap_val(concrete_jac)}(ad, linsolve, precs) + return GaussNewton{_unwrap_val(concrete_jac)}(ad, + linsolve, + precs, + termination_condition) end @concrete mutable struct GaussNewtonCache{iip} <: AbstractNonlinearSolveCache{iip} f alg u + u_prev fu1 fu2 fu_new @@ -72,12 +80,15 @@ end internalnorm retcode::ReturnCode.T abstol + reltol prob stats::NLStats + tc_storage end function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_::GaussNewton, - args...; alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + internalnorm = DEFAULT_NORM, kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @unpack f, u0, p = prob @@ -91,17 +102,30 @@ function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_:: uf, linsolve, J, fu2, jac_cache, du, JᵀJ, Jᵀf = jacobian_caches(alg, f, u, p, Val(iip); linsolve_with_JᵀJ = Val(true)) - return GaussNewtonCache{iip}(f, alg, u, fu1, fu2, zero(fu1), du, p, uf, linsolve, J, - JᵀJ, Jᵀf, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, - prob, NLStats(1, 0, 0, 0, 0)) + tc = alg.termination_condition + mode = DiffEqBase.get_termination_mode(tc) + + atol = _get_tolerance(abstol, tc.abstol, eltype(u)) + rtol = _get_tolerance(reltol, tc.reltol, eltype(u)) + + storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : + nothing + + return GaussNewtonCache{iip}(f, alg, u, copy(u), fu1, fu2, zero(fu1), du, p, uf, + linsolve, J, + JᵀJ, Jᵀf, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, + prob, NLStats(1, 0, 0, 0, 0), storage) end function perform_step!(cache::GaussNewtonCache{true}) - @unpack u, fu1, f, p, alg, J, JᵀJ, Jᵀf, linsolve, du = cache + @unpack u, u_prev, fu1, f, p, alg, J, JᵀJ, Jᵀf, linsolve, du = cache jacobian!!(J, cache) __matmul!(JᵀJ, J', J) __matmul!(Jᵀf, J', fu1) + tc_storage = cache.tc_storage + termination_condition = cache.alg.termination_condition(tc_storage) + # u = u - J \ fu linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(JᵀJ), b = _vec(Jᵀf), linu = _vec(du), p, reltol = cache.abstol) @@ -109,9 +133,15 @@ function perform_step!(cache::GaussNewtonCache{true}) @. u = u - du f(cache.fu_new, u, p) - (cache.internalnorm(cache.fu_new .- cache.fu1) < cache.abstol || - cache.internalnorm(cache.fu_new) < cache.abstol) && + (termination_condition(cache.fu_new .- cache.fu1, + cache.u, + u_prev, + cache.abstol, + cache.reltol) || + termination_condition(cache.fu_new, cache.u, u_prev, cache.abstol, cache.reltol)) && (cache.force_stop = true) + + @. u_prev = u cache.fu1 .= cache.fu_new cache.stats.nf += 1 cache.stats.njacs += 1 @@ -121,7 +151,10 @@ function perform_step!(cache::GaussNewtonCache{true}) end function perform_step!(cache::GaussNewtonCache{false}) - @unpack u, fu1, f, p, alg, linsolve = cache + @unpack u, u_prev, fu1, f, p, alg, linsolve = cache + + tc_storage = cache.tc_storage + termination_condition = cache.alg.termination_condition(tc_storage) cache.J = jacobian!!(cache.J, cache) @@ -138,7 +171,10 @@ function perform_step!(cache::GaussNewtonCache{false}) cache.u = @. u - cache.du # `u` might not support mutation cache.fu_new = f(cache.u, p) - (cache.internalnorm(cache.fu_new) < cache.abstol) && (cache.force_stop = true) + termination_condition(cache.fu_new, cache.u, u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) + + cache.u_prev = @. cache.u cache.fu1 = cache.fu_new cache.stats.nf += 1 cache.stats.njacs += 1 diff --git a/src/levenberg.jl b/src/levenberg.jl index 460377245..247b743f4 100644 --- a/src/levenberg.jl +++ b/src/levenberg.jl @@ -99,7 +99,7 @@ function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, damping_initial::Real = 1.0, damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, finite_diff_step_geodesic::Real = 0.1, α_geodesic::Real = 0.75, b_uphill::Real = 1.0, min_damping_D::AbstractFloat = 1e-8, - termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; + termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.AbsNorm; abstol = nothing, reltol = nothing), adkwargs...) @@ -209,7 +209,6 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, JᵀJ, λ, λ_factor, damping_increase_factor, damping_decrease_factor, h, α_geodesic, b_uphill, min_damping_D, v, a, tmp_vec, v_old, loss, δ, loss, make_new_J, fu_tmp, zero(u), zero(fu1), mat_tmp, NLStats(1, 0, 0, 0, 0), storage) - end function perform_step!(cache::LevenbergMarquardtCache{true}) @@ -271,11 +270,7 @@ function perform_step!(cache::LevenbergMarquardtCache{true}) if (1 - β)^b_uphill * loss ≤ loss_old # Accept step. cache.u .+= δ - if termination_condition(cache.fu_tmp, - cache.u, - u_prev, - cache.abstol, - cache.reltol) + if loss < cache.abstol cache.force_stop = true return nothing end From 1c5bb596a8a6d01da75148a221cb96a8532b5620 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Wed, 25 Oct 2023 17:28:28 -0400 Subject: [PATCH 06/12] Start moving termination conditions to solve kwargs --- docs/pages.jl | 3 +- docs/src/tutorials/code_optimization.md | 18 ++++----- docs/src/tutorials/large_systems.md | 10 ++--- docs/src/tutorials/modelingtoolkit.md | 2 +- docs/src/tutorials/small_compile.md | 20 +++++----- src/NonlinearSolve.jl | 2 +- src/gaussnewton.jl | 2 +- src/levenberg.jl | 4 +- src/raphson.jl | 43 +++++++++++--------- src/trustRegion.jl | 2 +- src/utils.jl | 52 +++++++++++++++++++++++++ test/basictests.jl | 2 +- 12 files changed, 109 insertions(+), 51 deletions(-) diff --git a/docs/pages.jl b/docs/pages.jl index 8564261bd..93ad7b391 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -2,8 +2,7 @@ pages = ["index.md", "Getting Started with Nonlinear Rootfinding in Julia" => "tutorials/getting_started.md", - "Tutorials" => Any[ - "Code Optimization for Small Nonlinear Systems" => "tutorials/code_optimization.md", + "Tutorials" => Any["Code Optimization for Small Nonlinear Systems" => "tutorials/code_optimization.md", "Handling Large Ill-Conditioned and Sparse Systems" => "tutorials/large_systems.md", "Symbolic System Definition and Acceleration via ModelingToolkit" => "tutorials/modelingtoolkit.md", "tutorials/small_compile.md", diff --git a/docs/src/tutorials/code_optimization.md b/docs/src/tutorials/code_optimization.md index daefcb855..cd4089d5d 100644 --- a/docs/src/tutorials/code_optimization.md +++ b/docs/src/tutorials/code_optimization.md @@ -34,7 +34,7 @@ Take for example a prototypical small nonlinear solver code in its out-of-place ```@example small_opt using NonlinearSolve -function f(u, p) +function f(u, p) u .* u .- p end u0 = [1.0, 1.0] @@ -54,7 +54,7 @@ using BenchmarkTools Note that this way of writing the function is a shorthand for: ```@example small_opt -function f(u, p) +function f(u, p) [u[1] * u[1] - p, u[2] * u[2] - p] end ``` @@ -69,25 +69,25 @@ In order to avoid this issue, we can use a non-allocating "in-place" approach. W this looks like: ```@example small_opt -function f(du, u, p) +function f(du, u, p) du[1] = u[1] * u[1] - p du[2] = u[2] * u[2] - p nothing end prob = NonlinearProblem(f, u0, p) -@btime sol = solve(prob, NewtonRaphson()) +@btime sol = solve(prob, NewtonRaphson()) ``` Notice how much faster this already runs! We can make this code even simpler by using the `.=` in-place broadcasting. ```@example small_opt -function f(du, u, p) +function f(du, u, p) du .= u .* u .- p end -@btime sol = solve(prob, NewtonRaphson()) +@btime sol = solve(prob, NewtonRaphson()) ``` ## Further Optimizations for Small Nonlinear Solves with Static Arrays and SimpleNonlinearSolve @@ -140,7 +140,7 @@ want to use the out-of-place allocating form, but this time we want to output a static array. Doing it with broadcasting looks like: ```@example small_opt -function f_SA(u, p) +function f_SA(u, p) u .* u .- p end u0 = SA[1.0, 1.0] @@ -153,7 +153,7 @@ Note that only change here is that `u0` is made into a StaticArray! If we needed for a more complex nonlinear case, then we'd simply do the following: ```@example small_opt -function f_SA(u, p) +function f_SA(u, p) SA[u[1] * u[1] - p, u[2] * u[2] - p] end @@ -170,4 +170,4 @@ which are designed for these small-scale static examples. Let's now use `SimpleN @btime solve(prob, SimpleNewtonRaphson()) ``` -And there we go, around 100ns from our starting point of almost 6μs! \ No newline at end of file +And there we go, around 100ns from our starting point of almost 6μs! diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index c1936932e..8e53d193c 100644 --- a/docs/src/tutorials/large_systems.md +++ b/docs/src/tutorials/large_systems.md @@ -1,8 +1,8 @@ # [Efficiently Solving Large Sparse Ill-Conditioned Nonlinear Systems in Julia](@id large_systems) -This tutorial is for getting into the extra features of using NonlinearSolve.jl. Solving ill-conditioned nonlinear systems -requires specializing the linear solver on properties of the Jacobian in order to cut down on the ``\mathcal{O}(n^3)`` -linear solve and the ``\mathcal{O}(n^2)`` back-solves. This tutorial is designed to explain the advanced usage of +This tutorial is for getting into the extra features of using NonlinearSolve.jl. Solving ill-conditioned nonlinear systems +requires specializing the linear solver on properties of the Jacobian in order to cut down on the ``\mathcal{O}(n^3)`` +linear solve and the ``\mathcal{O}(n^2)`` back-solves. This tutorial is designed to explain the advanced usage of NonlinearSolve.jl by solving the steady state stiff Brusselator partial differential equation (BRUSS) using NonlinearSolve.jl. ## Definition of the Brusselator Equation @@ -182,8 +182,8 @@ nothing # hide Notice that this acceleration does not require the definition of a sparsity pattern, and can thus be an easier way to scale for large problems. For more -information on linear solver choices, see the -[linear solver documentation](https://docs.sciml.ai/DiffEqDocs/stable/features/linear_nonlinear/#linear_nonlinear). +information on linear solver choices, see the +[linear solver documentation](https://docs.sciml.ai/DiffEqDocs/stable/features/linear_nonlinear/#linear_nonlinear). `linsolve` choices are any valid [LinearSolve.jl](https://linearsolve.sciml.ai/dev/) solver. !!! note diff --git a/docs/src/tutorials/modelingtoolkit.md b/docs/src/tutorials/modelingtoolkit.md index 7d26cb30c..3ee7d8e75 100644 --- a/docs/src/tutorials/modelingtoolkit.md +++ b/docs/src/tutorials/modelingtoolkit.md @@ -120,4 +120,4 @@ sol[u5] If you're interested in building models in a component or block based form, such as seen in systems like Simulink or Modelica, take a deeper look at [ModelingToolkit.jl's documentation](https://docs.sciml.ai/ModelingToolkit/stable/) which goes into -detail on such features. \ No newline at end of file +detail on such features. diff --git a/docs/src/tutorials/small_compile.md b/docs/src/tutorials/small_compile.md index d7d2c2311..4f4c56b14 100644 --- a/docs/src/tutorials/small_compile.md +++ b/docs/src/tutorials/small_compile.md @@ -19,18 +19,18 @@ sol = solve(prob, SimpleNewtonRaphson()) However, there are a few downsides to SimpleNonlinearSolve's `SimpleX` style algorithms to note: -1. SimpleNonlinearSolve.jl's methods are not hooked into the LinearSolve.jl system, and thus do not have - the ability to specify linear solvers, use sparse matrices, preconditioners, and all of the other features - which are required to scale for very large systems of equations. -2. SimpleNonlinearSolve.jl's methods have less robust error handling and termination conditions, and thus - these methods are missing some flexibility and give worse hints for debugging. -3. SimpleNonlinearSolve.jl's methods are focused on out-of-place support. There is some in-place support, - but it's designed for simple smaller systems and so some optimizations are missing. + 1. SimpleNonlinearSolve.jl's methods are not hooked into the LinearSolve.jl system, and thus do not have + the ability to specify linear solvers, use sparse matrices, preconditioners, and all of the other features + which are required to scale for very large systems of equations. + 2. SimpleNonlinearSolve.jl's methods have less robust error handling and termination conditions, and thus + these methods are missing some flexibility and give worse hints for debugging. + 3. SimpleNonlinearSolve.jl's methods are focused on out-of-place support. There is some in-place support, + but it's designed for simple smaller systems and so some optimizations are missing. However, the major upsides of SimpleNonlinearSolve.jl are: -1. The methods are optimized and non-allocating on StaticArrays -2. The methods are minimal in compilation + 1. The methods are optimized and non-allocating on StaticArrays + 2. The methods are minimal in compilation As such, you can use the code as shown above to have very low startup with good methods, but for more scaling and debuggability we recommend the full NonlinearSolve.jl. But that said, @@ -51,4 +51,4 @@ is not only sufficient but optimal. Julia has tools for building small binaries via static compilation with [StaticCompiler.jl](https://github.com/tshort/StaticCompiler.jl). However, these tools are currently limited to type-stable non-allocating functions. That said, SimpleNonlinearSolve.jl's solvers are -precisely the subset of NonlinearSolve.jl which are compatible with static compilation. \ No newline at end of file +precisely the subset of NonlinearSolve.jl which are compatible with static compilation. diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index b0fb28fb1..3a38989d4 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -26,7 +26,7 @@ const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences, ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode} abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end -abstract type AbstractNewtonAlgorithm{CJ, AD, TC} <: AbstractNonlinearSolveAlgorithm end +abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end abstract type AbstractNonlinearSolveCache{iip} end diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl index 73ff5b8dc..45ba18a5a 100644 --- a/src/gaussnewton.jl +++ b/src/gaussnewton.jl @@ -36,7 +36,7 @@ for large-scale and numerically-difficult nonlinear least squares problems. Jacobian-Free version of `GaussNewton` doesn't work yet, and it forces jacobian construction. This will be fixed in the near future. """ -@concrete struct GaussNewton{CJ, AD, TC} <: AbstractNewtonAlgorithm{CJ, AD, TC} +@concrete struct GaussNewton{CJ, AD, TC} <: AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve precs diff --git a/src/levenberg.jl b/src/levenberg.jl index 247b743f4..07c0ed21d 100644 --- a/src/levenberg.jl +++ b/src/levenberg.jl @@ -75,7 +75,7 @@ numerically-difficult nonlinear systems. `DᵀD` to prevent the damping from being too small. Defaults to `1e-8`. """ @concrete struct LevenbergMarquardt{CJ, AD, T, TC <: NLSolveTerminationCondition} <: - AbstractNewtonAlgorithm{CJ, AD, TC} + AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve precs @@ -157,7 +157,7 @@ end end function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, - NonlinearLeastSquaresProblem{uType, iip}}, alg_::LevenbergMarquardt, + NonlinearLeastSquaresProblem{uType, iip}}, alg_::LevenbergMarquardt, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} diff --git a/src/raphson.jl b/src/raphson.jl index 57547baca..0d3dddb98 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -30,13 +30,12 @@ for large-scale and numerically-difficult nonlinear systems. which means that no line search is performed. Algorithms from `LineSearches.jl` can be used here directly, and they will be converted to the correct `LineSearch`. """ -@concrete struct NewtonRaphson{CJ, AD, TC <: NLSolveTerminationCondition} <: - AbstractNewtonAlgorithm{CJ, AD, TC} +@concrete struct NewtonRaphson{CJ, AD} <: + AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve precs linesearch - termination_condition::TC end function set_ad(alg::NewtonRaphson{CJ}, ad) where {CJ} @@ -44,17 +43,13 @@ function set_ad(alg::NewtonRaphson{CJ}, ad) where {CJ} end function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, - linesearch = LineSearch(), precs = DEFAULT_PRECS, - termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; - abstol = nothing, - reltol = nothing), adkwargs...) + linesearch = LineSearch(), precs = DEFAULT_PRECS, adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) return NewtonRaphson{_unwrap_val(concrete_jac)}(ad, linsolve, precs, - linesearch, - termination_condition) + linesearch) end @concrete mutable struct NewtonRaphsonCache{iip} <: AbstractNonlinearSolveCache{iip} @@ -79,11 +74,13 @@ end prob stats::NLStats lscache + termination_condition tc_storage end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphson, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @@ -93,19 +90,20 @@ 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) - tc = alg.termination_condition - mode = DiffEqBase.get_termination_mode(tc) + abstol, reltol, termination_condition = _init_termination_elements(abstol, + reltol, + termination_condition, + eltype(u)) - atol = _get_tolerance(abstol, tc.abstol, eltype(u)) - rtol = _get_tolerance(reltol, tc.reltol, eltype(u)) + mode = DiffEqBase.get_termination_mode(termination_condition) storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : nothing return NewtonRaphsonCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, uf, linsolve, J, - jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, prob, + jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), LineSearchCache(alg.linesearch, f, u, p, fu1, Val(iip)), - storage) + termination_condition, storage) end function perform_step!(cache::NewtonRaphsonCache{true}) @@ -113,7 +111,7 @@ function perform_step!(cache::NewtonRaphsonCache{true}) jacobian!!(J, cache) tc_storage = cache.tc_storage - termination_condition = cache.alg.termination_condition(tc_storage) + termination_condition = cache.termination_condition(tc_storage) # u = u - J \ fu linres = dolinsolve(alg.precs, linsolve; A = J, b = _vec(fu1), linu = _vec(du), @@ -140,7 +138,7 @@ function perform_step!(cache::NewtonRaphsonCache{false}) @unpack u, u_prev, fu1, f, p, alg, linsolve = cache tc_storage = cache.tc_storage - termination_condition = cache.alg.termination_condition(tc_storage) + termination_condition = cache.termination_condition(tc_storage) cache.J = jacobian!!(cache.J, cache) # u = u - J \ fu @@ -169,7 +167,9 @@ function perform_step!(cache::NewtonRaphsonCache{false}) end function SciMLBase.reinit!(cache::NewtonRaphsonCache{iip}, u0 = cache.u; p = cache.p, - abstol = cache.abstol, maxiters = cache.maxiters) where {iip} + abstol = cache.abstol, reltol = cache.reltol, + termination_condition = cache.termination_condition, + maxiters = cache.maxiters) where {iip} cache.p = p if iip recursivecopy!(cache.u, u0) @@ -179,7 +179,14 @@ function SciMLBase.reinit!(cache::NewtonRaphsonCache{iip}, u0 = cache.u; p = cac cache.u = u0 cache.fu1 = cache.f(cache.u, p) end + + termination_condition = _get_reinit_termination_condition(cache, + abstol, + reltol, + termination_condition) cache.abstol = abstol + cache.reltol = reltol + cache.termination_condition = termination_condition cache.maxiters = maxiters cache.stats.nf = 1 cache.stats.nsteps = 1 diff --git a/src/trustRegion.jl b/src/trustRegion.jl index bfdc2a557..741871f72 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -149,7 +149,7 @@ for large-scale and numerically-difficult nonlinear systems. Support for the OOP version is planned! """ @concrete struct TrustRegion{CJ, AD, MTR, TC <: NLSolveTerminationCondition} <: - AbstractNewtonAlgorithm{CJ, AD, TC} + AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve precs diff --git a/src/utils.jl b/src/utils.jl index 887a92dc6..f51f34f0b 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -213,3 +213,55 @@ function _get_tolerance(η, tc_η, ::Type{T}) where {T} fallback_η = real(oneunit(T)) * (eps(real(one(T))))^(4 // 5) return ifelse(η !== nothing, η, ifelse(tc_η !== nothing, tc_η, fallback_η)) end + +function _init_termination_elements(abstol, + reltol, + termination_condition, + ::Type{T}) where {T} + if termination_condition !== nothing + abstol !== nothing ? + (abstol != termination_condition.abstol ? + error("Incompatible absolute tolerances found. The tolerances supplied as the keyword argument and the one supplied in the termination condition should be same.") : + nothing) : nothing + reltol !== nothing ? + (reltol != termination_condition.abstol ? + error("Incompatible relative tolerances found. The tolerances supplied as the keyword argument and the one supplied in the termination condition should be same.") : + nothing) : nothing + abstol = _get_tolerance(abstol, termination_condition.abstol, T) + reltol = _get_tolerance(reltol, termination_condition.reltol, T) + return abstol, reltol, termination_condition + else + abstol = _get_tolerance(abstol, nothing, T) + reltol = _get_tolerance(reltol, nothing, T) + termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; + abstol, + reltol) + return abstol, reltol, termination_condition + end +end + +function _get_reinit_termination_condition(cache, abstol, reltol, termination_condition) + if termination_condition != cache.termination_condition + if abstol != cache.abstol + if abstol != termination_condition.abstol + error("Incompatible absolute tolerances found") + end + end + + if reltol != cache.reltol + if reltol != termination_condition.reltol + error("Incompatible relative tolerances found") + end + end + termination_condition + else + # Build the termination_condition with new abstol and reltol + return NLSolveTerminationCondition{ + DiffEqBase.get_termination_mode(termination_condition), + eltype(abstol), + typeof(termination_condition.safe_termination_options), + }(abstol, + reltol, + termination_condition.safe_termination_options) + end +end diff --git a/test/basictests.jl b/test/basictests.jl index 7135b4901..0d487bbc9 100644 --- a/test/basictests.jl +++ b/test/basictests.jl @@ -135,7 +135,7 @@ end termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, reltol = nothing) probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, NewtonRaphson(; termination_condition)).u .≈ sqrt(2.0)) + @test all(solve(probN, NewtonRaphson(); termination_condition).u .≈ sqrt(2.0)) end end From eddcfde5b84d55501636dff55496e6785299d791 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Wed, 25 Oct 2023 19:42:14 -0400 Subject: [PATCH 07/12] Complete moving termination condition to all the algorithms --- src/NonlinearSolve.jl | 2 +- src/dfsane.jl | 30 ++++++++++++++++++++++++------ src/gaussnewton.jl | 37 +++++++++++++++++-------------------- src/levenberg.jl | 43 ++++++++++++++++++++++--------------------- src/raphson.jl | 3 +-- src/trustRegion.jl | 31 ++++++++++++++----------------- src/utils.jl | 6 +++--- test/basictests.jl | 20 +++++++++++++++++--- 8 files changed, 99 insertions(+), 73 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 3a38989d4..8046e4ec4 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -14,7 +14,7 @@ import EnumX: @enumx import ForwardDiff: Dual import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A import RecursiveArrayTools: ArrayPartition, - AbstractVectorOfArray, recursivecopy!, recursivefill! + AbstractVectorOfArray, recursivecopy!, recursivefill!, recursive_unitless_bottom_eltype import Reexport: @reexport import SciMLBase: AbstractNonlinearAlgorithm, NLStats, _unwrap_val, has_jac, isinplace import StaticArraysCore: StaticArray, SVector, SArray, MArray diff --git a/src/dfsane.jl b/src/dfsane.jl index 13de5ff6a..669fa7473 100644 --- a/src/dfsane.jl +++ b/src/dfsane.jl @@ -89,12 +89,16 @@ end internalnorm retcode::SciMLBase.ReturnCode.T abstol + reltol prob stats::NLStats + termination_condition + tc_storage end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::DFSane, args...; - alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, kwargs...) where {uType, iip} uₙ = alias_u0 ? prob.u0 : deepcopy(prob.u0) @@ -123,14 +127,27 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::DFSane, args. f₍ₙₒᵣₘ₎₀ = f₍ₙₒᵣₘ₎ₙ₋₁ ℋ = 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 + return DFSaneCache{iip}(alg, uₙ, uₙ₋₁, fuₙ, fuₙ₋₁, 𝒹, ℋ, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, M, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, p, false, maxiters, - internalnorm, ReturnCode.Default, abstol, prob, NLStats(1, 0, 0, 0, 0)) + internalnorm, ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), + termination_condition, storage) end function perform_step!(cache::DFSaneCache{true}) - @unpack alg, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, M = cache + @unpack alg, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, M, tc_storage = cache + termination_condition = cache.termination_condition(tc_storage) f = (dx, x) -> cache.prob.f(dx, x, cache.p) T = eltype(cache.uₙ) @@ -175,7 +192,7 @@ function perform_step!(cache::DFSaneCache{true}) f₍ₙₒᵣₘ₎ₙ = norm(cache.fuₙ)^nₑₓₚ end - if cache.internalnorm(cache.fuₙ) < cache.abstol + if termination_condition(cache.fuₙ, cache.uₙ, cache.uₙ₋₁, cache.abstol, cache.reltol) cache.force_stop = true end @@ -206,8 +223,9 @@ function perform_step!(cache::DFSaneCache{true}) end function perform_step!(cache::DFSaneCache{false}) - @unpack alg, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, M = cache + @unpack alg, f₍ₙₒᵣₘ₎ₙ₋₁, f₍ₙₒᵣₘ₎₀, σₙ, σₘᵢₙ, σₘₐₓ, α₁, γ, τₘᵢₙ, τₘₐₓ, nₑₓₚ, M, tc_storage = cache + termination_condition = cache.termination_condition(tc_storage) f = x -> cache.prob.f(x, cache.p) T = eltype(cache.uₙ) @@ -250,7 +268,7 @@ function perform_step!(cache::DFSaneCache{false}) f₍ₙₒᵣₘ₎ₙ = norm(cache.fuₙ)^nₑₓₚ end - if cache.internalnorm(cache.fuₙ) < cache.abstol + if termination_condition(cache.fuₙ, cache.uₙ, cache.uₙ₋₁, cache.abstol, cache.reltol) cache.force_stop = true end diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl index 45ba18a5a..c2284213f 100644 --- a/src/gaussnewton.jl +++ b/src/gaussnewton.jl @@ -36,11 +36,10 @@ for large-scale and numerically-difficult nonlinear least squares problems. Jacobian-Free version of `GaussNewton` doesn't work yet, and it forces jacobian construction. This will be fixed in the near future. """ -@concrete struct GaussNewton{CJ, AD, TC} <: AbstractNewtonAlgorithm{CJ, AD} +@concrete struct GaussNewton{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve precs - termination_condition::TC end function set_ad(alg::GaussNewton{CJ}, ad) where {CJ} @@ -48,15 +47,11 @@ function set_ad(alg::GaussNewton{CJ}, ad) where {CJ} end function GaussNewton(; concrete_jac = nothing, linsolve = NormalCholeskyFactorization(), - precs = DEFAULT_PRECS, - termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.AbsNorm; - abstol = nothing, - reltol = nothing), adkwargs...) + precs = DEFAULT_PRECS, adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) return GaussNewton{_unwrap_val(concrete_jac)}(ad, linsolve, - precs, - termination_condition) + precs) end @concrete mutable struct GaussNewtonCache{iip} <: AbstractNonlinearSolveCache{iip} @@ -84,10 +79,12 @@ end prob stats::NLStats tc_storage + termination_condition end function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_::GaussNewton, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @@ -102,29 +99,30 @@ function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_:: uf, linsolve, J, fu2, jac_cache, du, JᵀJ, Jᵀf = jacobian_caches(alg, f, u, p, Val(iip); linsolve_with_JᵀJ = Val(true)) - tc = alg.termination_condition - mode = DiffEqBase.get_termination_mode(tc) + abstol, reltol, termination_condition = _init_termination_elements(abstol, + reltol, + termination_condition, + eltype(u); mode = NLSolveTerminationMode.AbsNorm) - atol = _get_tolerance(abstol, tc.abstol, eltype(u)) - rtol = _get_tolerance(reltol, tc.reltol, eltype(u)) + mode = DiffEqBase.get_termination_mode(termination_condition) storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : nothing return GaussNewtonCache{iip}(f, alg, u, copy(u), fu1, fu2, zero(fu1), du, p, uf, linsolve, J, - JᵀJ, Jᵀf, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, - prob, NLStats(1, 0, 0, 0, 0), storage) + JᵀJ, Jᵀf, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, + reltol, + prob, NLStats(1, 0, 0, 0, 0), storage, termination_condition) end function perform_step!(cache::GaussNewtonCache{true}) - @unpack u, u_prev, fu1, f, p, alg, J, JᵀJ, Jᵀf, linsolve, du = cache + @unpack u, u_prev, fu1, f, p, alg, J, JᵀJ, Jᵀf, linsolve, du, tc_storage = cache jacobian!!(J, cache) __matmul!(JᵀJ, J', J) __matmul!(Jᵀf, J', fu1) - tc_storage = cache.tc_storage - termination_condition = cache.alg.termination_condition(tc_storage) + termination_condition = cache.termination_condition(tc_storage) # u = u - J \ fu linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(JᵀJ), b = _vec(Jᵀf), @@ -151,10 +149,9 @@ function perform_step!(cache::GaussNewtonCache{true}) end function perform_step!(cache::GaussNewtonCache{false}) - @unpack u, u_prev, fu1, f, p, alg, linsolve = cache + @unpack u, u_prev, fu1, f, p, alg, linsolve, tc_storage = cache - tc_storage = cache.tc_storage - termination_condition = cache.alg.termination_condition(tc_storage) + termination_condition = cache.termination_condition(tc_storage) cache.J = jacobian!!(cache.J, cache) diff --git a/src/levenberg.jl b/src/levenberg.jl index 07c0ed21d..096cbb4a9 100644 --- a/src/levenberg.jl +++ b/src/levenberg.jl @@ -74,7 +74,7 @@ numerically-difficult nonlinear systems. [this paper](https://arxiv.org/abs/1201.5885) to use a minimum value of the elements in `DᵀD` to prevent the damping from being too small. Defaults to `1e-8`. """ -@concrete struct LevenbergMarquardt{CJ, AD, T, TC <: NLSolveTerminationCondition} <: +@concrete struct LevenbergMarquardt{CJ, AD, T} <: AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve @@ -86,7 +86,6 @@ numerically-difficult nonlinear systems. α_geodesic::T b_uphill::T min_damping_D::T - termination_condition::TC end function set_ad(alg::LevenbergMarquardt{CJ}, ad) where {CJ} @@ -99,15 +98,11 @@ function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, damping_initial::Real = 1.0, damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, finite_diff_step_geodesic::Real = 0.1, α_geodesic::Real = 0.75, b_uphill::Real = 1.0, min_damping_D::AbstractFloat = 1e-8, - termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.AbsNorm; - abstol = nothing, - reltol = nothing), adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) return LevenbergMarquardt{_unwrap_val(concrete_jac)}(ad, linsolve, precs, damping_initial, damping_increase_factor, damping_decrease_factor, - finite_diff_step_geodesic, α_geodesic, b_uphill, min_damping_D, - termination_condition) + finite_diff_step_geodesic, α_geodesic, b_uphill, min_damping_D) end @concrete mutable struct LevenbergMarquardtCache{iip} <: AbstractNonlinearSolveCache{iip} @@ -153,12 +148,14 @@ end Jv mat_tmp stats::NLStats + termination_condition tc_storage end function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, NonlinearLeastSquaresProblem{uType, iip}}, alg_::LevenbergMarquardt, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @@ -168,6 +165,11 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, uf, linsolve, J, fu2, jac_cache, du, JᵀJ, v = jacobian_caches(alg, f, u, p, Val(iip); linsolve_kwargs, linsolve_with_JᵀJ = Val(true)) + abstol, reltol, termination_condition = _init_termination_elements(abstol, + reltol, + termination_condition, + eltype(u); mode = NLSolveTerminationMode.AbsNorm) + λ = convert(eltype(u), alg.damping_initial) λ_factor = convert(eltype(u), alg.damping_increase_factor) damping_increase_factor = convert(eltype(u), alg.damping_increase_factor) @@ -194,28 +196,24 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, fu_tmp = zero(fu1) mat_tmp = zero(JᵀJ) - tc = alg.termination_condition - mode = DiffEqBase.get_termination_mode(tc) - - atol = _get_tolerance(abstol, tc.abstol, eltype(u)) - rtol = _get_tolerance(reltol, tc.reltol, eltype(u)) + mode = DiffEqBase.get_termination_mode(termination_condition) storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : nothing return LevenbergMarquardtCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, uf, linsolve, J, - jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, prob, DᵀD, + jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, + DᵀD, JᵀJ, λ, λ_factor, damping_increase_factor, damping_decrease_factor, h, α_geodesic, b_uphill, min_damping_D, v, a, tmp_vec, v_old, loss, δ, loss, make_new_J, fu_tmp, - zero(u), zero(fu1), mat_tmp, NLStats(1, 0, 0, 0, 0), storage) + zero(u), zero(fu1), mat_tmp, NLStats(1, 0, 0, 0, 0), termination_condition, storage) end function perform_step!(cache::LevenbergMarquardtCache{true}) - @unpack fu1, f, make_new_J = cache + @unpack fu1, f, make_new_J, tc_storage = cache - tc_storage = cache.tc_storage - termination_condition = cache.alg.termination_condition(tc_storage) + termination_condition = cache.termination_condition(tc_storage) if iszero(fu1) cache.force_stop = true @@ -270,7 +268,11 @@ function perform_step!(cache::LevenbergMarquardtCache{true}) if (1 - β)^b_uphill * loss ≤ loss_old # Accept step. cache.u .+= δ - if loss < cache.abstol + if termination_condition(cache.fu_tmp, + cache.u, + u_prev, + cache.abstol, + cache.reltol) cache.force_stop = true return nothing end @@ -289,10 +291,9 @@ function perform_step!(cache::LevenbergMarquardtCache{true}) end function perform_step!(cache::LevenbergMarquardtCache{false}) - @unpack fu1, f, make_new_J = cache + @unpack fu1, f, make_new_J, tc_storage = cache - tc_storage = cache.tc_storage - termination_condition = cache.alg.termination_condition(tc_storage) + termination_condition = cache.termination_condition(tc_storage) if iszero(fu1) cache.force_stop = true diff --git a/src/raphson.jl b/src/raphson.jl index 0d3dddb98..409c36860 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -107,10 +107,9 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphso end function perform_step!(cache::NewtonRaphsonCache{true}) - @unpack u, u_prev, fu1, f, p, alg, J, linsolve, du = cache + @unpack u, u_prev, fu1, f, p, alg, J, linsolve, du, tc_storage = cache jacobian!!(J, cache) - tc_storage = cache.tc_storage termination_condition = cache.termination_condition(tc_storage) # u = u - J \ fu diff --git a/src/trustRegion.jl b/src/trustRegion.jl index 741871f72..5178ce32e 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -148,7 +148,7 @@ for large-scale and numerically-difficult nonlinear systems. `linsolve` and `precs` are used exclusively for the inplace version of the algorithm. Support for the OOP version is planned! """ -@concrete struct TrustRegion{CJ, AD, MTR, TC <: NLSolveTerminationCondition} <: +@concrete struct TrustRegion{CJ, AD, MTR} <: AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve @@ -162,7 +162,6 @@ for large-scale and numerically-difficult nonlinear systems. shrink_factor::MTR expand_factor::MTR max_shrink_times::Int - termination_condition::TC end function set_ad(alg::TrustRegion{CJ}, ad) where {CJ} @@ -177,15 +176,11 @@ function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAU max_trust_radius::Real = 0 // 1, initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, shrink_factor::Real = 1 // 4, - expand_factor::Real = 2 // 1, max_shrink_times::Int = 32, - termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; - abstol = nothing, - reltol = nothing), adkwargs...) + expand_factor::Real = 2 // 1, max_shrink_times::Int = 32, adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) return TrustRegion{_unwrap_val(concrete_jac)}(ad, linsolve, precs, radius_update_scheme, max_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, - expand_threshold, shrink_factor, expand_factor, max_shrink_times, - termination_condition) + expand_threshold, shrink_factor, expand_factor, max_shrink_times) end @concrete mutable struct TrustRegionCache{iip, trustType, floatType} <: @@ -236,10 +231,12 @@ end ϵ::floatType stats::NLStats tc_storage + termination_condition end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::TrustRegion, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @@ -342,22 +339,23 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::TrustRegion, initial_trust_radius = convert(trustType, 1.0) end - tc = alg.termination_condition - mode = DiffEqBase.get_termination_mode(tc) + abstol, reltol, termination_condition = _init_termination_elements(abstol, + reltol, + termination_condition, + eltype(u)) - atol = _get_tolerance(abstol, tc.abstol, eltype(u)) - rtol = _get_tolerance(reltol, tc.reltol, eltype(u)) + mode = DiffEqBase.get_termination_mode(termination_condition) storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : nothing return TrustRegionCache{iip}(f, alg, u_prev, u, fu_prev, fu1, fu2, p, uf, linsolve, J, - jac_cache, false, maxiters, internalnorm, ReturnCode.Default, atol, rtol, prob, + jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, radius_update_scheme, initial_trust_radius, max_trust_radius, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, loss, loss_new, H, g, shrink_counter, du, u_tmp, u_gauss_newton, u_cauchy, fu_new, make_new_J, r, p1, p2, p3, p4, ϵ, - NLStats(1, 0, 0, 0, 0), storage) + NLStats(1, 0, 0, 0, 0), storage, termination_condition) end function perform_step!(cache::TrustRegionCache{true}) @@ -433,10 +431,9 @@ function retrospective_step!(cache::TrustRegionCache) end function trust_region_step!(cache::TrustRegionCache) - @unpack fu_new, du, g, H, loss, max_trust_r, radius_update_scheme = cache + @unpack fu_new, du, g, H, loss, max_trust_r, radius_update_scheme, tc_storage = cache - tc_storage = cache.tc_storage - termination_condition = cache.alg.termination_condition(tc_storage) + termination_condition = cache.termination_condition(tc_storage) cache.loss_new = get_loss(fu_new) diff --git a/src/utils.jl b/src/utils.jl index f51f34f0b..3d621c52f 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -211,13 +211,13 @@ end function _get_tolerance(η, tc_η, ::Type{T}) where {T} fallback_η = real(oneunit(T)) * (eps(real(one(T))))^(4 // 5) - return ifelse(η !== nothing, η, ifelse(tc_η !== nothing, tc_η, fallback_η)) + return T(ifelse(η !== nothing, η, ifelse(tc_η !== nothing, tc_η, fallback_η))) end function _init_termination_elements(abstol, reltol, termination_condition, - ::Type{T}) where {T} + ::Type{T}; mode = NLSolveTerminationMode.NLSolveDefault) where {T} if termination_condition !== nothing abstol !== nothing ? (abstol != termination_condition.abstol ? @@ -233,7 +233,7 @@ function _init_termination_elements(abstol, else abstol = _get_tolerance(abstol, nothing, T) reltol = _get_tolerance(reltol, nothing, T) - termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault; + termination_condition = NLSolveTerminationCondition(mode; abstol, reltol) return abstol, reltol, termination_condition diff --git a/test/basictests.jl b/test/basictests.jl index 0d487bbc9..056ba632c 100644 --- a/test/basictests.jl +++ b/test/basictests.jl @@ -72,7 +72,7 @@ end end @testset "[OOP] [Immutable AD]" begin - for p in 1.0:0.1:100.0 + for p in [1.0, 100.0] @test begin res = benchmark_nlsolve_oop(quadratic_f, @SVector[1.0, 1.0], p) res_true = sqrt(p) @@ -307,7 +307,7 @@ end termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, reltol = nothing) probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, TrustRegion(; termination_condition)).u .≈ sqrt(2.0)) + @test all(solve(probN, TrustRegion(); termination_condition).u .≈ sqrt(2.0)) end end @@ -430,7 +430,7 @@ end termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, reltol = nothing) probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, LevenbergMarquardt(; termination_condition)).u .≈ sqrt(2.0)) + @test all(solve(probN, LevenbergMarquardt(); termination_condition).u .≈ sqrt(2.0)) end end @@ -584,4 +584,18 @@ end @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-10) end end + + @testset "Termination condition: $(mode) u0: $(_nameof(u0))" for mode in instances(NLSolveTerminationMode.T), + u0 in (1.0, [1.0, 1.0]) + + if mode ∈ + (NLSolveTerminationMode.SteadyStateDefault, NLSolveTerminationMode.RelSafeBest, + NLSolveTerminationMode.AbsSafeBest) + continue + end + termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, + reltol = nothing) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve(probN, DFSane(); termination_condition).u .≈ sqrt(2.0)) + end end From 4baf777ff60c610f319952e79e667f9a1404445c Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Wed, 25 Oct 2023 19:52:00 -0400 Subject: [PATCH 08/12] Handle termination conditions in reinit! --- src/dfsane.jl | 10 +++++++++- src/gaussnewton.jl | 11 ++++++++++- src/trustRegion.jl | 10 +++++++++- src/utils.jl | 4 ++-- 4 files changed, 30 insertions(+), 5 deletions(-) diff --git a/src/dfsane.jl b/src/dfsane.jl index 669fa7473..641222990 100644 --- a/src/dfsane.jl +++ b/src/dfsane.jl @@ -315,7 +315,8 @@ function SciMLBase.solve!(cache::DFSaneCache) end function SciMLBase.reinit!(cache::DFSaneCache{iip}, u0 = cache.uₙ; p = cache.p, - abstol = cache.abstol, maxiters = cache.maxiters) where {iip} + abstol = cache.abstol, termination_condition = cache.termination_condition, + maxiters = cache.maxiters) where {iip} cache.p = p if iip recursivecopy!(cache.uₙ, u0) @@ -336,7 +337,14 @@ 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, + termination_condition) + cache.abstol = abstol + cache.reltol = reltol + cache.termination_condition = termination_condition cache.maxiters = maxiters cache.stats.nf = 1 cache.stats.nsteps = 1 diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl index c2284213f..8415b259e 100644 --- a/src/gaussnewton.jl +++ b/src/gaussnewton.jl @@ -181,7 +181,9 @@ function perform_step!(cache::GaussNewtonCache{false}) end function SciMLBase.reinit!(cache::GaussNewtonCache{iip}, u0 = cache.u; p = cache.p, - abstol = cache.abstol, maxiters = cache.maxiters) where {iip} + abstol = cache.abstol, reltol = cache.reltol, + termination_condition = cache.termination_condition, + maxiters = cache.maxiters) where {iip} cache.p = p if iip recursivecopy!(cache.u, u0) @@ -191,7 +193,14 @@ function SciMLBase.reinit!(cache::GaussNewtonCache{iip}, u0 = cache.u; p = cache cache.u = u0 cache.fu1 = cache.f(cache.u, p) end + termination_condition = _get_reinit_termination_condition(cache, + abstol, + reltol, + termination_condition) + cache.abstol = abstol + cache.reltol = reltol + cache.termination_condition = termination_condition cache.maxiters = maxiters cache.stats.nf = 1 cache.stats.nsteps = 1 diff --git a/src/trustRegion.jl b/src/trustRegion.jl index 5178ce32e..08da2b57d 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -724,7 +724,8 @@ end get_fu(cache::TrustRegionCache) = cache.fu function SciMLBase.reinit!(cache::TrustRegionCache{iip}, u0 = cache.u; p = cache.p, - abstol = cache.abstol, maxiters = cache.maxiters) where {iip} + abstol = cache.abstol, termination_condition = cache.termination_condition, + maxiters = cache.maxiters) where {iip} cache.p = p if iip recursivecopy!(cache.u, u0) @@ -734,7 +735,14 @@ function SciMLBase.reinit!(cache::TrustRegionCache{iip}, u0 = cache.u; p = cache cache.u = u0 cache.fu = cache.f(cache.u, p) end + termination_condition = _get_reinit_termination_condition(cache, + abstol, + reltol, + termination_condition) + cache.abstol = abstol + cache.reltol = reltol + cache.termination_condition = termination_condition cache.maxiters = maxiters cache.stats.nf = 1 cache.stats.nsteps = 1 diff --git a/src/utils.jl b/src/utils.jl index 3d621c52f..01fcfece4 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -244,13 +244,13 @@ function _get_reinit_termination_condition(cache, abstol, reltol, termination_co if termination_condition != cache.termination_condition if abstol != cache.abstol if abstol != termination_condition.abstol - error("Incompatible absolute tolerances found") + error("Incompatible absolute tolerances found. The tolerances supplied as the keyword argument and the one supplied in the termination condition should be same.") end end if reltol != cache.reltol if reltol != termination_condition.reltol - error("Incompatible relative tolerances found") + error("Incompatible absolute tolerances found. The tolerances supplied as the keyword argument and the one supplied in the termination condition should be same.") end end termination_condition From 814f704ec0cfd71a764d96365d5f1ce589cc2b5c Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Thu, 26 Oct 2023 01:21:34 -0400 Subject: [PATCH 09/12] Minor fixes after rebasing --- src/dfsane.jl | 3 ++- src/gaussnewton.jl | 2 +- src/levenberg.jl | 13 ++++++++----- src/raphson.jl | 3 ++- src/trustRegion.jl | 3 ++- src/utils.jl | 2 +- test/23_test_problems.jl | 2 +- 7 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/dfsane.jl b/src/dfsane.jl index 573523d0f..d88f18db0 100644 --- a/src/dfsane.jl +++ b/src/dfsane.jl @@ -314,7 +314,8 @@ function SciMLBase.solve!(cache::DFSaneCache) end function SciMLBase.reinit!(cache::DFSaneCache{iip}, u0 = cache.uₙ; p = cache.p, - abstol = cache.abstol, termination_condition = cache.termination_condition, + abstol = cache.abstol, reltol = cache.reltol, + termination_condition = cache.termination_condition, maxiters = cache.maxiters) where {iip} cache.p = p if iip diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl index 43d45e827..42155072a 100644 --- a/src/gaussnewton.jl +++ b/src/gaussnewton.jl @@ -131,7 +131,7 @@ function perform_step!(cache::GaussNewtonCache{true}) jacobian!!(J, cache) termination_condition = cache.termination_condition(tc_storage) - + if JᵀJ !== nothing __matmul!(JᵀJ, J', J) __matmul!(Jᵀf, J', fu1) diff --git a/src/levenberg.jl b/src/levenberg.jl index b9524021b..0aa42f5c4 100644 --- a/src/levenberg.jl +++ b/src/levenberg.jl @@ -219,7 +219,7 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() : nothing - + if _unwrap_val(linsolve_with_JᵀJ) mat_tmp = zero(JᵀJ) rhs_tmp = nothing @@ -232,19 +232,22 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, linsolve = __setup_linsolve(mat_tmp, rhs_tmp, u, p, alg) end - return LevenbergMarquardtCache{iip, !_unwrap_val(linsolve_with_JᵀJ)}(f, alg, u, fu1, + return LevenbergMarquardtCache{iip, !_unwrap_val(linsolve_with_JᵀJ)}(f, alg, u, copy(u), + fu1, fu2, du, p, uf, linsolve, J, - jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, DᵀD, + jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, + DᵀD, JᵀJ, λ, λ_factor, damping_increase_factor, damping_decrease_factor, h, α_geodesic, b_uphill, min_damping_D, v, a, tmp_vec, v_old, loss, δ, loss, make_new_J, fu_tmp, - zero(u), zero(fu1), mat_tmp, rhs_tmp, J², NLStats(1, 0, 0, 0, 0), termination_condition, storage) + zero(u), zero(fu1), mat_tmp, rhs_tmp, J², NLStats(1, 0, 0, 0, 0), + termination_condition, storage) end function perform_step!(cache::LevenbergMarquardtCache{true, fastls}) where {fastls} @unpack fu1, f, make_new_J, tc_storage = cache termination_condition = cache.termination_condition(tc_storage) - + if iszero(fu1) cache.force_stop = true return nothing diff --git a/src/raphson.jl b/src/raphson.jl index 320efe7da..eef64ec7b 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -102,7 +102,8 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphso 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)), + NLStats(1, 0, 0, 0, 0), + init_linesearch_cache(alg.linesearch, f, u, p, fu1, Val(iip)), termination_condition, storage) end diff --git a/src/trustRegion.jl b/src/trustRegion.jl index a6d24f581..3b14d0a38 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -724,7 +724,8 @@ end get_fu(cache::TrustRegionCache) = cache.fu function SciMLBase.reinit!(cache::TrustRegionCache{iip}, u0 = cache.u; p = cache.p, - abstol = cache.abstol, termination_condition = cache.termination_condition, + abstol = cache.abstol, reltol = cache.reltol, + termination_condition = cache.termination_condition, maxiters = cache.maxiters) where {iip} cache.p = p if iip diff --git a/src/utils.jl b/src/utils.jl index 8df6c1733..016d09a4b 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -324,4 +324,4 @@ end _needs_square_A(_, ::Number) = true _needs_square_A(_, ::StaticArray) = true -_needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve) \ No newline at end of file +_needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve) diff --git a/test/23_test_problems.jl b/test/23_test_problems.jl index bd2b932dc..091088ab2 100644 --- a/test/23_test_problems.jl +++ b/test/23_test_problems.jl @@ -75,7 +75,7 @@ end alg_ops = (DFSane(),) broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [1, 2, 3, 5, 6, 8, 12, 13, 14, 21] + broken_tests[alg_ops[1]] = [1, 2, 3, 5, 6, 21] test_on_library(problems, dicts, alg_ops, broken_tests) end From 416e6569e3abc2b4b7f28cd98183285692ddcf27 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Thu, 26 Oct 2023 02:32:31 -0400 Subject: [PATCH 10/12] Add termination condition to Pseudotransient methods --- src/pseudotransient.jl | 59 ++++++++++++++++++++++++++++++++++-------- test/basictests.jl | 16 ++++++++++++ 2 files changed, 64 insertions(+), 11 deletions(-) diff --git a/src/pseudotransient.jl b/src/pseudotransient.jl index a041c258c..c5871b12f 100644 --- a/src/pseudotransient.jl +++ b/src/pseudotransient.jl @@ -3,9 +3,9 @@ precs = DEFAULT_PRECS, alpha_initial = 1e-3, adkwargs...) An implementation of PseudoTransient method that is used to solve steady state problems in an accelerated manner. It uses an adaptive time-stepping to -integrate an initial value of nonlinear problem until sufficient accuracy in the desired steady-state is achieved to switch over to Newton's method and +integrate an initial value of nonlinear problem until sufficient accuracy in the desired steady-state is achieved to switch over to Newton's method and gain a rapid convergence. This implementation specifically uses "switched evolution relaxation" SER method. For detail information about the time-stepping and algorithm, -please see the paper: [Coffey, Todd S. and Kelley, C. T. and Keyes, David E. (2003), Pseudotransient Continuation and Differential-Algebraic Equations, +please see the paper: [Coffey, Todd S. and Kelley, C. T. and Keyes, David E. (2003), Pseudotransient Continuation and Differential-Algebraic Equations, SIAM Journal on Scientific Computing,25, 553-569.](https://doi.org/10.1137/S106482750241044X) ### Keyword Arguments @@ -27,7 +27,7 @@ SIAM Journal on Scientific Computing,25, 553-569.](https://doi.org/10.1137/S1064 preconditioners. For more information on specifying preconditioners for LinearSolve algorithms, consult the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `alpha_initial` : the initial pseudo time step. it defaults to 1e-3. If it is small, + - `alpha_initial` : the initial pseudo time step. it defaults to 1e-3. If it is small, you are going to need more iterations to converge but it can be more stable. """ @concrete struct PseudoTransient{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} @@ -52,6 +52,7 @@ end f alg u + u_prev fu1 fu2 du @@ -67,15 +68,19 @@ end internalnorm retcode::ReturnCode.T abstol + reltol prob stats::NLStats + termination_condition + tc_storage end isinplace(::PseudoTransientCache{iip}) where {iip} = iip function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::PseudoTransient, args...; - alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @@ -93,16 +98,30 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::PseudoTransi alpha = convert(eltype(u), alg.alpha_initial) res_norm = internalnorm(fu1) - return PseudoTransientCache{iip}(f, alg, u, fu1, fu2, du, p, alpha, res_norm, uf, + 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 + + return PseudoTransientCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, alpha, res_norm, + uf, linsolve, J, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, - prob, NLStats(1, 0, 0, 0, 0)) + reltol, + prob, NLStats(1, 0, 0, 0, 0), termination_condition, storage) end function perform_step!(cache::PseudoTransientCache{true}) - @unpack u, fu1, f, p, alg, J, linsolve, du, alpha = cache + @unpack u, u_prev, fu1, f, p, alg, J, linsolve, du, alpha, tc_storage = cache jacobian!!(J, cache) J_new = J - (1 / alpha) * I + termination_condition = cache.termination_condition(tc_storage) + # u = u - J \ fu linres = dolinsolve(alg.precs, linsolve; A = J_new, b = _vec(fu1), linu = _vec(du), p, reltol = cache.abstol) @@ -114,7 +133,10 @@ function perform_step!(cache::PseudoTransientCache{true}) cache.alpha *= cache.res_norm / new_norm cache.res_norm = new_norm - new_norm < cache.abstol && (cache.force_stop = true) + termination_condition(fu1, u, u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) + + @. u_prev = u cache.stats.nf += 1 cache.stats.njacs += 1 cache.stats.nsolve += 1 @@ -123,7 +145,10 @@ function perform_step!(cache::PseudoTransientCache{true}) end function perform_step!(cache::PseudoTransientCache{false}) - @unpack u, fu1, f, p, alg, linsolve, alpha = cache + @unpack u, u_prev, fu1, f, p, alg, linsolve, alpha, tc_storage = cache + + tc_storage = cache.tc_storage + termination_condition = cache.termination_condition(tc_storage) cache.J = jacobian!!(cache.J, cache) # u = u - J \ fu @@ -141,7 +166,9 @@ function perform_step!(cache::PseudoTransientCache{false}) new_norm = cache.internalnorm(fu1) cache.alpha *= cache.res_norm / new_norm cache.res_norm = new_norm - new_norm < cache.abstol && (cache.force_stop = true) + termination_condition(fu1, cache.u, u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) + cache.u_prev = @. cache.u cache.stats.nf += 1 cache.stats.njacs += 1 cache.stats.nsolve += 1 @@ -167,7 +194,9 @@ end function SciMLBase.reinit!(cache::PseudoTransientCache{iip}, u0 = cache.u; p = cache.p, alpha_new, - abstol = cache.abstol, maxiters = cache.maxiters) where {iip} + abstol = cache.abstol, reltol = cache.reltol, + termination_condition = cache.termination_condition, + maxiters = cache.maxiters) where {iip} cache.p = p if iip recursivecopy!(cache.u, u0) @@ -177,9 +206,17 @@ function SciMLBase.reinit!(cache::PseudoTransientCache{iip}, u0 = cache.u; p = c cache.u = u0 cache.fu1 = cache.f(cache.u, p) end + + termination_condition = _get_reinit_termination_condition(cache, + abstol, + reltol, + termination_condition) + cache.alpha = convert(eltype(cache.u), alpha_new) cache.res_norm = cache.internalnorm(cache.fu1) cache.abstol = abstol + cache.reltol = reltol + cache.termination_condition = termination_condition cache.maxiters = maxiters cache.stats.nf = 1 cache.stats.nsteps = 1 diff --git a/test/basictests.jl b/test/basictests.jl index f26c44666..02f100e71 100644 --- a/test/basictests.jl +++ b/test/basictests.jl @@ -720,6 +720,22 @@ end sol = solve(probN, PseudoTransient(alpha_initial = 1.0), abstol = 1e-10) @test all(abs.(newton_fails(sol.u, p)) .< 1e-10) end + + @testset "Termination condition: $(mode) u0: $(_nameof(u0))" for mode in instances(NLSolveTerminationMode.T), + u0 in (1.0, [1.0, 1.0]) + + if mode ∈ + (NLSolveTerminationMode.SteadyStateDefault, NLSolveTerminationMode.RelSafeBest, + NLSolveTerminationMode.AbsSafeBest) + continue + end + termination_condition = NLSolveTerminationCondition(mode; abstol = nothing, + reltol = nothing) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve(probN, + PseudoTransient(; alpha_initial = 10.0); + termination_condition).u .≈ sqrt(2.0)) + end end # --- GeneralBroyden tests --- From f20c9bcc974c7a4d353c2d46ea9fd41e790309cb Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Thu, 26 Oct 2023 10:44:30 -0400 Subject: [PATCH 11/12] Move other algorithms to use termination conditions --- src/broyden.jl | 53 ++++++++++++++++++++++++++++++++++++++++--------- src/klement.jl | 51 ++++++++++++++++++++++++++++++++++++++--------- src/lbroyden.jl | 42 +++++++++++++++++++++++++++++++-------- src/raphson.jl | 3 +-- 4 files changed, 121 insertions(+), 28 deletions(-) diff --git a/src/broyden.jl b/src/broyden.jl index 6be29a77b..22a1a9cc8 100644 --- a/src/broyden.jl +++ b/src/broyden.jl @@ -31,6 +31,7 @@ end f alg u + u_prev du fu fu2 @@ -46,17 +47,21 @@ end internalnorm retcode::ReturnCode.T abstol + reltol reset_tolerance reset_check prob stats::NLStats lscache + termination_condition + tc_storage end get_fu(cache::GeneralBroydenCache) = cache.fu function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralBroyden, args...; - alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, kwargs...) where {uType, iip} @unpack f, u0, p = prob u = alias_u0 ? u0 : deepcopy(u0) @@ -65,15 +70,29 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralBroyde reset_tolerance = alg.reset_tolerance === nothing ? sqrt(eps(eltype(u))) : alg.reset_tolerance reset_check = x -> abs(x) ≤ reset_tolerance - return GeneralBroydenCache{iip}(f, alg, u, _mutable_zero(u), fu, zero(fu), + + 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 + 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, reset_tolerance, + 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))) + init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), termination_condition, + storage) end function perform_step!(cache::GeneralBroydenCache{true}) - @unpack f, p, du, fu, fu2, dfu, u, J⁻¹, J⁻¹df, J⁻¹₂ = cache + @unpack f, p, du, fu, fu2, dfu, u, u_prev, J⁻¹, J⁻¹df, J⁻¹₂, tc_storage = cache + + termination_condition = cache.termination_condition(tc_storage) T = eltype(u) mul!(_vec(du), J⁻¹, -_vec(fu)) @@ -81,7 +100,8 @@ function perform_step!(cache::GeneralBroydenCache{true}) _axpy!(α, du, u) f(fu2, u, p) - cache.internalnorm(fu2) < cache.abstol && (cache.force_stop = true) + termination_condition(fu2, u, u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) cache.stats.nf += 1 cache.force_stop && return nothing @@ -106,12 +126,16 @@ function perform_step!(cache::GeneralBroydenCache{true}) mul!(J⁻¹, _vec(du), J⁻¹₂, 1, 1) end fu .= fu2 + @. u_prev = u return nothing end function perform_step!(cache::GeneralBroydenCache{false}) - @unpack f, p = cache + @unpack f, p, tc_storage = cache + + termination_condition = cache.termination_condition(tc_storage) + T = eltype(cache.u) cache.du = _restructure(cache.du, cache.J⁻¹ * -_vec(cache.fu)) @@ -119,7 +143,8 @@ function perform_step!(cache::GeneralBroydenCache{false}) cache.u = cache.u .+ α * cache.du cache.fu2 = f(cache.u, p) - cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true) + termination_condition(cache.fu2, cache.u, cache.u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) cache.stats.nf += 1 cache.force_stop && return nothing @@ -142,12 +167,15 @@ function perform_step!(cache::GeneralBroydenCache{false}) cache.J⁻¹ = cache.J⁻¹ .+ _vec(cache.du) * cache.J⁻¹₂ end cache.fu = cache.fu2 + cache.u_prev = @. cache.u return nothing end function SciMLBase.reinit!(cache::GeneralBroydenCache{iip}, u0 = cache.u; p = cache.p, - abstol = cache.abstol, maxiters = cache.maxiters) where {iip} + abstol = cache.abstol, reltol = cache.reltol, + termination_condition = cache.termination_condition, + maxiters = cache.maxiters) where {iip} cache.p = p if iip recursivecopy!(cache.u, u0) @@ -157,7 +185,14 @@ function SciMLBase.reinit!(cache::GeneralBroydenCache{iip}, u0 = cache.u; p = ca cache.u = u0 cache.fu = cache.f(cache.u, p) end + termination_condition = _get_reinit_termination_condition(cache, + abstol, + reltol, + termination_condition) + cache.abstol = abstol + cache.reltol = reltol + cache.termination_condition = termination_condition cache.maxiters = maxiters cache.stats.nf = 1 cache.stats.nsteps = 1 diff --git a/src/klement.jl b/src/klement.jl index a16ed2873..32aec15d5 100644 --- a/src/klement.jl +++ b/src/klement.jl @@ -41,6 +41,7 @@ end f alg u + u_prev fu fu2 du @@ -65,7 +66,8 @@ end get_fu(cache::GeneralKlementCache) = cache.fu function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralKlement, args...; - alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} @unpack f, u0, p = prob u = alias_u0 ? u0 : deepcopy(u0) @@ -84,16 +86,30 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralKleme linsolve = __setup_linsolve(J, _vec(fu), _vec(du), p, alg) end - return GeneralKlementCache{iip}(f, alg, u, fu, zero(fu), du, p, linsolve, + 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 + + return GeneralKlementCache{iip}(f, alg, u, zero(u), fu, zero(fu), du, p, linsolve, J, zero(J), zero(J), _vec(zero(fu)), _vec(zero(fu)), 0, false, - maxiters, internalnorm, ReturnCode.Default, abstol, prob, NLStats(1, 0, 0, 0, 0), - init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip))) + maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, + NLStats(1, 0, 0, 0, 0), + init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), termination_condition, + storage) end function perform_step!(cache::GeneralKlementCache{true}) - @unpack u, fu, f, p, alg, J, linsolve, du = cache + @unpack u, u_prev, fu, f, p, alg, J, linsolve, du, tc_storage = cache T = eltype(J) + termination_condition = cache.termination_condition(tc_storage) + singular, fact_done = _try_factorize_and_check_singular!(linsolve, J) if singular @@ -118,7 +134,8 @@ function perform_step!(cache::GeneralKlementCache{true}) _axpy!(α, du, u) f(cache.fu2, u, p) - cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true) + termination_condition(cache.fu2, u, u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) cache.stats.nf += 1 cache.stats.nsolve += 1 cache.stats.nfactors += 1 @@ -138,13 +155,17 @@ function perform_step!(cache::GeneralKlementCache{true}) mul!(cache.J_cache2, cache.J_cache, J) J .+= cache.J_cache2 + @. u_prev = u cache.fu .= cache.fu2 return nothing end function perform_step!(cache::GeneralKlementCache{false}) - @unpack fu, f, p, alg, J, linsolve = cache + @unpack fu, f, p, alg, J, linsolve, tc_storage = cache + + termination_condition = cache.termination_condition(tc_storage) + T = eltype(J) singular, fact_done = _try_factorize_and_check_singular!(linsolve, J) @@ -174,7 +195,10 @@ function perform_step!(cache::GeneralKlementCache{false}) cache.u = @. cache.u + α * cache.du # `u` might not support mutation cache.fu2 = f(cache.u, p) - cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true) + termination_condition(cache.fu2, cache.u, cache.u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) + + cache.u_prev = @. cache.u cache.stats.nf += 1 cache.stats.nsolve += 1 cache.stats.nfactors += 1 @@ -198,7 +222,9 @@ function perform_step!(cache::GeneralKlementCache{false}) end function SciMLBase.reinit!(cache::GeneralKlementCache{iip}, u0 = cache.u; p = cache.p, - abstol = cache.abstol, maxiters = cache.maxiters) where {iip} + abstol = cache.abstol, reltol = cache.reltol, + termination_condition = cache.termination_condition, + maxiters = cache.maxiters) where {iip} cache.p = p if iip recursivecopy!(cache.u, u0) @@ -208,7 +234,14 @@ function SciMLBase.reinit!(cache::GeneralKlementCache{iip}, u0 = cache.u; p = ca cache.u = u0 cache.fu = cache.f(cache.u, p) end + + termination_condition = _get_reinit_termination_condition(cache, + abstol, + reltol, + termination_condition) cache.abstol = abstol + cache.reltol = reltol + cache.termination_condition = termination_condition cache.maxiters = maxiters cache.stats.nf = 1 cache.stats.nsteps = 1 diff --git a/src/lbroyden.jl b/src/lbroyden.jl index d045d0b20..db4353b41 100644 --- a/src/lbroyden.jl +++ b/src/lbroyden.jl @@ -34,6 +34,7 @@ end f alg u + u_prev du fu fu2 @@ -53,17 +54,21 @@ end internalnorm retcode::ReturnCode.T abstol + reltol reset_tolerance reset_check prob stats::NLStats lscache + termination_condition + tc_storage end get_fu(cache::LimitedMemoryBroydenCache) = cache.fu function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::LimitedMemoryBroyden, - args...; alias_u0 = false, maxiters = 1000, abstol = 1e-6, internalnorm = DEFAULT_NORM, + args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, kwargs...) where {uType, iip} @unpack f, u0, p = prob u = alias_u0 ? u0 : deepcopy(u0) @@ -80,23 +85,38 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::LimitedMemory reset_tolerance = alg.reset_tolerance === nothing ? sqrt(eps(eltype(u))) : alg.reset_tolerance reset_check = x -> abs(x) ≤ reset_tolerance - return LimitedMemoryBroydenCache{iip}(f, alg, u, du, fu, zero(fu), + + 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 + + return LimitedMemoryBroydenCache{iip}(f, alg, u, zero(u), du, fu, zero(fu), zero(fu), p, U, Vᵀ, similar(u, threshold), similar(u, 1, threshold), zero(u), zero(u), false, 0, 0, alg.max_resets, maxiters, internalnorm, - ReturnCode.Default, abstol, reset_tolerance, reset_check, prob, + 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))) + init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), termination_condition, + storage) end function perform_step!(cache::LimitedMemoryBroydenCache{true}) - @unpack f, p, du, u = cache + @unpack f, p, du, u, tc_storage = cache T = eltype(u) + termination_condition = cache.termination_condition(tc_storage) + α = perform_linesearch!(cache.lscache, u, du) _axpy!(α, du, u) f(cache.fu2, u, p) - cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true) + termination_condition(cache.fu2, cache.u, cache.u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) cache.stats.nf += 1 cache.force_stop && return nothing @@ -138,20 +158,25 @@ function perform_step!(cache::LimitedMemoryBroydenCache{true}) cache.iterations_since_reset += 1 end + cache.u_prev .= cache.u cache.fu .= cache.fu2 return nothing end function perform_step!(cache::LimitedMemoryBroydenCache{false}) - @unpack f, p = cache + @unpack f, p, tc_storage = cache + + termination_condition = cache.termination_condition(tc_storage) + T = eltype(cache.u) α = perform_linesearch!(cache.lscache, cache.u, cache.du) cache.u = cache.u .+ α * cache.du cache.fu2 = f(cache.u, p) - cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true) + termination_condition(cache.fu2, cache.u, cache.u_prev, cache.abstol, cache.reltol) && + (cache.force_stop = true) cache.stats.nf += 1 cache.force_stop && return nothing @@ -194,6 +219,7 @@ function perform_step!(cache::LimitedMemoryBroydenCache{false}) cache.iterations_since_reset += 1 end + cache.u_prev = @. cache.u cache.fu = cache.fu2 return nothing diff --git a/src/raphson.jl b/src/raphson.jl index eef64ec7b..a34d860ce 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -135,9 +135,8 @@ function perform_step!(cache::NewtonRaphsonCache{true}) end function perform_step!(cache::NewtonRaphsonCache{false}) - @unpack u, u_prev, fu1, f, p, alg, linsolve = cache + @unpack u, u_prev, fu1, f, p, alg, linsolve, tc_storage = cache - tc_storage = cache.tc_storage termination_condition = cache.termination_condition(tc_storage) cache.J = jacobian!!(cache.J, cache) From 350fac5976644f6d77c6245092ea0fced7e2f989 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Thu, 26 Oct 2023 14:53:44 -0400 Subject: [PATCH 12/12] fixup! cache structs in Klement --- src/klement.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/klement.jl b/src/klement.jl index 32aec15d5..8fc44be59 100644 --- a/src/klement.jl +++ b/src/klement.jl @@ -58,9 +58,12 @@ end internalnorm retcode::ReturnCode.T abstol + reltol prob stats::NLStats lscache + termination_condition + tc_storage end get_fu(cache::GeneralKlementCache) = cache.fu