diff --git a/Project.toml b/Project.toml index 37513ae88..d8719ddfb 100644 --- a/Project.toml +++ b/Project.toml @@ -12,10 +12,12 @@ EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" @@ -43,22 +45,24 @@ ADTypes = "0.2" ArrayInterface = "6.0.24, 7" BandedMatrices = "1" ConcreteStructs = "0.2" -DiffEqBase = "6.136" +DiffEqBase = "6.141" EnumX = "1" Enzyme = "0.11" FastBroadcast = "0.1.9, 0.2" FastLevenbergMarquardt = "0.1" FiniteDiff = "2" ForwardDiff = "0.10.3" +LazyArrays = "1.8" LeastSquaresOptim = "0.8" LineSearches = "7" LinearAlgebra = "<0.0.1, 1" LinearSolve = "2.12" NonlinearProblemLibrary = "0.1" PrecompileTools = "1" +Printf = "<0.0.1, 1" RecursiveArrayTools = "2" Reexport = "0.2, 1" -SciMLBase = "2.8.2" +SciMLBase = "2.9" SciMLOperators = "0.3" SimpleNonlinearSolve = "0.1.23" SparseArrays = "<0.0.1, 1" diff --git a/docs/pages.jl b/docs/pages.jl index a68b0f8e2..bea06e3fa 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -12,6 +12,7 @@ pages = ["index.md", "basics/solve.md", "basics/NonlinearSolution.md", "basics/TerminationCondition.md", + "basics/Logging.md", "basics/FAQ.md"], "Solver Summaries and Recommendations" => Any["solvers/NonlinearSystemSolvers.md", "solvers/BracketingSolvers.md", diff --git a/docs/src/basics/Logging.md b/docs/src/basics/Logging.md new file mode 100644 index 000000000..30f579538 --- /dev/null +++ b/docs/src/basics/Logging.md @@ -0,0 +1,63 @@ +# 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`: + + 1. `show_trace`: Must be `Val(true)` or `Val(false)`. This controls whether the trace is + displayed to the console. (Defaults to `Val(false)`) + 2. `trace_level`: Needs to be one of Trace Objects: [`TraceMinimal`](@ref), + [`TraceWithJacobianConditionNumber`](@ref), or [`TraceAll`](@ref). This controls the + level of detail of the trace. (Defaults to `TraceMinimal()`) + 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)`) + +## Example Usage + +```@example tracing +using ModelingToolkit, NonlinearSolve + +@variables x y z +@parameters σ ρ β + +# Define a nonlinear system +eqs = [0 ~ σ * (y - x), + 0 ~ x * (ρ - z) - y, + 0 ~ x * y - β * z] +@named ns = NonlinearSystem(eqs, [x, y, z], [σ, ρ, β]) + +u0 = [x => 1.0, y => 0.0, z => 0.0] + +ps = [σ => 10.0 ρ => 26.0 β => 8 / 3] + +prob = NonlinearProblem(ns, u0, ps) + +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 +solve(prob; show_trace = Val(true), trace_level = TraceAll(10)) +``` + +You can also store the trace in the solution object: + +```@example tracing +sol = solve(prob; trace_level = TraceAll(), store_trace = Val(true)); + +sol.trace +``` + +!!! note + + 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/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 369de3669..c591eb4ee 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -8,7 +8,9 @@ import Reexport: @reexport import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_workload @recompile_invalidations begin - using DiffEqBase, LinearAlgebra, LinearSolve, SparseArrays, SparseDiffTools + using DiffEqBase, + LazyArrays, LinearAlgebra, LinearSolve, Printf, SparseArrays, + SparseDiffTools using FastBroadcast: @.. import ArrayInterface: restructure @@ -50,6 +52,32 @@ abstract type AbstractNonlinearSolveCache{iip} end isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip +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 + 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 + function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, alg::AbstractNonlinearSolveAlgorithm, args...; kwargs...) cache = init(prob, alg, args...; kwargs...) @@ -79,11 +107,18 @@ 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 + return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode, cache.stats) + cache.retcode, cache.stats, trace) end include("utils.jl") +include("trace.jl") include("extension_algs.jl") include("linesearch.jl") include("raphson.jl") @@ -162,4 +197,7 @@ export SteadyStateDiffEqTerminationMode, SimpleNonlinearSolveTerminationMode, AbsNormTerminationMode, RelSafeTerminationMode, AbsSafeTerminationMode, RelSafeBestTerminationMode, AbsSafeBestTerminationMode +# Tracing Functionality +export TraceAll, TraceMinimal, TraceWithJacobianConditionNumber + end # module diff --git a/src/broyden.jl b/src/broyden.jl index 5c43e9cd8..4db44d308 100644 --- a/src/broyden.jl +++ b/src/broyden.jl @@ -1,6 +1,6 @@ # Sadly `Broyden` is taken up by SimpleNonlinearSolve.jl """ - GeneralBroyden(; max_resets = 3, linesearch = LineSearch(), reset_tolerance = nothing) + GeneralBroyden(; max_resets = 3, linesearch = nothing, reset_tolerance = nothing) An implementation of `Broyden` with reseting and line search. @@ -21,7 +21,7 @@ An implementation of `Broyden` with reseting and line search. linesearch end -function GeneralBroyden(; max_resets = 3, linesearch = LineSearch(), +function GeneralBroyden(; max_resets = 3, linesearch = nothing, reset_tolerance = nothing) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) return GeneralBroyden(max_resets, reset_tolerance, linesearch) @@ -54,6 +54,7 @@ end stats::NLStats ls_cache tc_cache + trace end get_fu(cache::GeneralBroydenCache) = cache.fu @@ -66,6 +67,7 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralBroyde @unpack f, u0, p = prob u = alias_u0 ? u0 : deepcopy(u0) fu = evaluate_f(prob, u) + du = _mutable_zero(u) J⁻¹ = __init_identity_jacobian(u, fu) reset_tolerance = alg.reset_tolerance === nothing ? sqrt(eps(real(eltype(u)))) : alg.reset_tolerance @@ -73,12 +75,14 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralBroyde abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u, termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu, J⁻¹, du; uses_jac_inverse = Val(true), + kwargs...) - return GeneralBroydenCache{iip}(f, alg, u, zero(u), _mutable_zero(u), fu, zero(fu), + return GeneralBroydenCache{iip}(f, alg, u, zero(u), du, fu, zero(fu), zero(fu), p, J⁻¹, zero(_reshape(fu, 1, :)), _mutable_zero(u), false, 0, alg.max_resets, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, reset_tolerance, reset_check, prob, NLStats(1, 0, 0, 0, 0), - init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache) + init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache, trace) end function perform_step!(cache::GeneralBroydenCache{true}) @@ -90,6 +94,9 @@ function perform_step!(cache::GeneralBroydenCache{true}) _axpy!(-α, du, u) f(fu2, u, p) + update_trace_with_invJ!(cache.trace, cache.stats.nsteps + 1, get_u(cache), + get_fu(cache), J⁻¹, du, α) + check_and_update!(cache, fu2, u, u_prev) cache.stats.nf += 1 @@ -131,6 +138,9 @@ function perform_step!(cache::GeneralBroydenCache{false}) cache.u = cache.u .- α * cache.du cache.fu2 = f(cache.u, p) + update_trace_with_invJ!(cache.trace, cache.stats.nsteps + 1, get_u(cache), + get_fu(cache), cache.J⁻¹, cache.du, α) + check_and_update!(cache, cache.fu2, cache.u, cache.u_prev) cache.stats.nf += 1 @@ -173,6 +183,7 @@ function SciMLBase.reinit!(cache::GeneralBroydenCache{iip}, u0 = cache.u; p = ca cache.fu = cache.f(cache.u, p) end + reset!(cache.trace) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fu, cache.u, termination_condition) diff --git a/src/default.jl b/src/default.jl index 96c1d906b..04198a1e0 100644 --- a/src/default.jl +++ b/src/default.jl @@ -82,7 +82,7 @@ end 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])) + original = $(sol_syms[i]), trace = $(sol_syms[i]).trace) end cache.current = $(i + 1) end @@ -103,7 +103,7 @@ end u = cache.caches[idx].u return SciMLBase.build_solution(cache.caches[idx].prob, cache.alg, u, - fus[idx]; retcode, stats) + fus[idx]; retcode, stats, cache.caches[idx].trace) end) return Expr(:block, calls...) @@ -125,7 +125,7 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb 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)) + original = $(cur_sol), trace = $(cur_sol).trace) end end) end @@ -147,7 +147,7 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb 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]).stats, $(sol_syms[i]).trace) end end) end diff --git a/src/dfsane.jl b/src/dfsane.jl index 5648f8af9..4b31ff9f3 100644 --- a/src/dfsane.jl +++ b/src/dfsane.jl @@ -90,6 +90,7 @@ end prob stats::NLStats tc_cache + trace end get_fu(cache::DFSaneCache) = cache.fu @@ -113,11 +114,12 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::DFSane, args. abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, uprev, termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu, nothing, du; kwargs...) return DFSaneCache{iip}(alg, u, uprev, fu, fuprev, 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) + abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache, trace) end function perform_step!(cache::DFSaneCache{true}) @@ -164,6 +166,9 @@ function perform_step!(cache::DFSaneCache{true}) f_norm = cache.internalnorm(cache.fu)^n_exp end + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), nothing, + cache.du, α₊) + check_and_update!(cache, cache.fu, cache.u, cache.uprev) # Update spectral parameter @@ -236,6 +241,9 @@ function perform_step!(cache::DFSaneCache{false}) f_norm = cache.internalnorm(cache.fu)^n_exp end + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), nothing, + cache.du, α₊) + check_and_update!(cache, cache.fu, cache.u, cache.uprev) # Update spectral parameter @@ -288,6 +296,7 @@ function SciMLBase.reinit!(cache::DFSaneCache{iip}, u0 = cache.u; p = cache.p, T = eltype(cache.u) cache.σ_n = T(cache.alg.σ_1) + reset!(cache.trace) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fu, cache.u, termination_condition) diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl index 07b1adaec..ea1855e68 100644 --- a/src/gaussnewton.jl +++ b/src/gaussnewton.jl @@ -1,5 +1,5 @@ """ - GaussNewton(; concrete_jac = nothing, linsolve = nothing, linesearch = LineSearch(), + GaussNewton(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, precs = DEFAULT_PRECS, adkwargs...) An advanced GaussNewton implementation with support for efficient handling of sparse @@ -47,7 +47,7 @@ function set_ad(alg::GaussNewton{CJ}, ad) where {CJ} end function GaussNewton(; concrete_jac = nothing, linsolve = nothing, - linesearch = LineSearch(), precs = DEFAULT_PRECS, vjp_autodiff = nothing, + linesearch = nothing, precs = DEFAULT_PRECS, vjp_autodiff = nothing, adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) @@ -82,6 +82,7 @@ end tc_cache_1 tc_cache_2 ls_cache + trace end function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_::GaussNewton, @@ -108,11 +109,12 @@ function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_:: abstol, reltol, tc_cache_1 = init_termination_cache(abstol, reltol, fu1, u, termination_condition) _, _, tc_cache_2 = init_termination_cache(abstol, reltol, fu1, u, termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu1, ApplyArray(__zero, J), du; kwargs...) return GaussNewtonCache{iip}(f, alg, u, copy(u), fu1, fu2, zero(fu1), du, p, uf, linsolve, J, JᵀJ, Jᵀf, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache_1, tc_cache_2, - init_linesearch_cache(alg.linesearch, f, u, p, fu1, Val(iip))) + init_linesearch_cache(alg.linesearch, f, u, p, fu1, Val(iip)), trace) end function perform_step!(cache::GaussNewtonCache{true}) @@ -137,6 +139,9 @@ function perform_step!(cache::GaussNewtonCache{true}) _axpy!(-α, du, u) f(cache.fu_new, u, p) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), J, + cache.du, α) + check_and_update!(cache.tc_cache_1, cache, cache.fu_new, cache.u, cache.u_prev) if !cache.force_stop cache.fu1 .= cache.fu_new .- cache.fu1 @@ -179,6 +184,9 @@ function perform_step!(cache::GaussNewtonCache{false}) cache.u = @. u - α * cache.du # `u` might not support mutation cache.fu_new = f(cache.u, p) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), cache.J, + cache.du, α) + check_and_update!(cache.tc_cache_1, cache, cache.fu_new, cache.u, cache.u_prev) if !cache.force_stop cache.fu1 = cache.fu_new .- cache.fu1 @@ -207,6 +215,7 @@ function SciMLBase.reinit!(cache::GaussNewtonCache{iip}, u0 = cache.u; p = cache cache.fu1 = cache.f(cache.u, p) end + reset!(cache.trace) abstol, reltol, tc_cache_1 = init_termination_cache(abstol, reltol, cache.fu1, cache.u, termination_condition) _, _, tc_cache_2 = init_termination_cache(abstol, reltol, cache.fu1, cache.u, diff --git a/src/klement.jl b/src/klement.jl index d5031ae24..ec32dc6b8 100644 --- a/src/klement.jl +++ b/src/klement.jl @@ -1,6 +1,6 @@ """ GeneralKlement(; max_resets = 5, linsolve = nothing, - linesearch = LineSearch(), precs = DEFAULT_PRECS) + linesearch = nothing, precs = DEFAULT_PRECS) An implementation of `Klement` with line search, preconditioning and customizable linear solves. @@ -32,7 +32,7 @@ function set_linsolve(alg::GeneralKlement, linsolve) end function GeneralKlement(; max_resets::Int = 5, linsolve = nothing, - linesearch = LineSearch(), precs = DEFAULT_PRECS) + linesearch = nothing, precs = DEFAULT_PRECS) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) return GeneralKlement(max_resets, linsolve, precs, linesearch) end @@ -63,6 +63,7 @@ end stats::NLStats ls_cache tc_cache + trace end get_fu(cache::GeneralKlementCache) = cache.fu @@ -91,12 +92,13 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralKleme abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u, termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu, J, du; kwargs...) return GeneralKlementCache{iip}(f, alg, u, zero(u), fu, zero(fu), du, p, linsolve, J, zero(J), zero(J), _vec(zero(fu)), _vec(zero(fu)), 0, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), - init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache) + init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache, trace) end function perform_step!(cache::GeneralKlementCache{true}) @@ -127,6 +129,9 @@ function perform_step!(cache::GeneralKlementCache{true}) _axpy!(-α, du, u) f(cache.fu2, u, p) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), cache.fu2, J, + cache.du, α) + check_and_update!(cache, cache.fu2, cache.u, cache.u_prev) cache.stats.nf += 1 cache.stats.nsolve += 1 @@ -186,6 +191,9 @@ function perform_step!(cache::GeneralKlementCache{false}) cache.u = @. cache.u - α * cache.du # `u` might not support mutation cache.fu2 = f(cache.u, p) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), cache.fu2, J, + cache.du, α) + check_and_update!(cache, cache.fu2, cache.u, cache.u_prev) cache.u_prev = cache.u cache.stats.nf += 1 @@ -224,6 +232,7 @@ function SciMLBase.reinit!(cache::GeneralKlementCache{iip}, u0 = cache.u; p = ca cache.fu = cache.f(cache.u, p) end + reset!(cache.trace) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fu, cache.u, termination_condition) diff --git a/src/lbroyden.jl b/src/lbroyden.jl index 42000db4e..ecd098f0e 100644 --- a/src/lbroyden.jl +++ b/src/lbroyden.jl @@ -1,5 +1,5 @@ """ - LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = LineSearch(), + LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = nothing, threshold::Int = 10, reset_tolerance = nothing) An implementation of `LimitedMemoryBroyden` with reseting and line search. @@ -24,7 +24,7 @@ An implementation of `LimitedMemoryBroyden` with reseting and line search. reset_tolerance end -function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = LineSearch(), +function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = nothing, threshold::Int = 27, reset_tolerance = nothing) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) return LimitedMemoryBroyden(max_resets, threshold, linesearch, reset_tolerance) @@ -61,6 +61,7 @@ end stats::NLStats ls_cache tc_cache + trace end get_fu(cache::LimitedMemoryBroydenCache) = cache.fu @@ -88,13 +89,17 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::LimitedMemory 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...) return LimitedMemoryBroydenCache{iip}(f, alg, u, zero(u), du, fu, zero(fu), zero(fu), p, U, Vᵀ, similar(u, threshold), similar(u, 1, threshold), zero(u), zero(u), false, 0, 0, alg.max_resets, maxiters, internalnorm, ReturnCode.Default, abstol, 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) + init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache, trace) end function perform_step!(cache::LimitedMemoryBroydenCache{true}) @@ -105,6 +110,12 @@ function perform_step!(cache::LimitedMemoryBroydenCache{true}) _axpy!(-α, du, u) f(cache.fu2, u, p) + idx = min(cache.iterations_since_reset, size(cache.U, 1)) + U_part = selectdim(cache.U, 1, 1:idx) + Vᵀ_part = selectdim(cache.Vᵀ, 2, 1:idx) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), cache.fu2, + ApplyArray(*, Vᵀ_part, U_part), du, α) + check_and_update!(cache, cache.fu2, cache.u, cache.u_prev) cache.stats.nf += 1 @@ -162,6 +173,12 @@ function perform_step!(cache::LimitedMemoryBroydenCache{false}) cache.u = cache.u .- α * cache.du cache.fu2 = f(cache.u, p) + idx = min(cache.iterations_since_reset, size(cache.U, 1)) + U_part = selectdim(cache.U, 1, 1:idx) + Vᵀ_part = selectdim(cache.Vᵀ, 2, 1:idx) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), cache.fu2, + ApplyArray(*, Vᵀ_part, U_part), cache.du, α) + check_and_update!(cache, cache.fu2, cache.u, cache.u_prev) cache.stats.nf += 1 @@ -225,6 +242,7 @@ function SciMLBase.reinit!(cache::LimitedMemoryBroydenCache{iip}, u0 = cache.u; cache.fu = cache.f(cache.u, p) end + reset!(cache.trace) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fu, cache.u, termination_condition) diff --git a/src/levenberg.jl b/src/levenberg.jl index 83729a634..dcc07d85e 100644 --- a/src/levenberg.jl +++ b/src/levenberg.jl @@ -160,6 +160,7 @@ end stats::NLStats tc_cache_1 tc_cache_2 + trace end function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, @@ -220,6 +221,8 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, tc_cache_2 = nothing end + trace = init_nonlinearsolve_trace(alg, u, fu1, ApplyArray(__zero, J), du; kwargs...) + if _unwrap_val(linsolve_with_JᵀJ) mat_tmp = zero(JᵀJ) rhs_tmp = nothing @@ -237,7 +240,8 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip}, ReturnCode.Default, abstol, reltol, prob, DᵀD, JᵀJ, λ, λ_factor, damping_increase_factor, damping_decrease_factor, h, α_geodesic, b_uphill, min_damping_D, v, a, tmp_vec, v_old, loss, δ, loss, make_new_J, fu_tmp, zero(u), - zero(fu1), mat_tmp, rhs_tmp, J², NLStats(1, 0, 0, 0, 0), tc_cache_1, tc_cache_2) + zero(fu1), mat_tmp, rhs_tmp, J², NLStats(1, 0, 0, 0, 0), tc_cache_1, tc_cache_2, + trace) end function perform_step!(cache::LevenbergMarquardtCache{true, fastls}) where {fastls} @@ -276,6 +280,9 @@ function perform_step!(cache::LevenbergMarquardtCache{true, fastls}) where {fast _vec(cache.v) .= -_vec(cache.du) end + 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). @unpack v, α_geodesic, h = cache cache.u_tmp .= _restructure(cache.u_tmp, _vec(u) .+ h .* _vec(v)) @@ -373,6 +380,9 @@ function perform_step!(cache::LevenbergMarquardtCache{false, fastls}) where {fas end end + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), cache.J, + cache.v) + @unpack v, h, α_geodesic = cache # Geodesic acceleration (step_size = v + a / 2). rhs_term = _vec(((2 / h) .* ((_vec(f(u .+ h .* _restructure(u, v), p)) .- diff --git a/src/pseudotransient.jl b/src/pseudotransient.jl index b343138de..3873202c4 100644 --- a/src/pseudotransient.jl +++ b/src/pseudotransient.jl @@ -76,6 +76,7 @@ end prob stats::NLStats tc_cache + trace end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::PseudoTransient, @@ -94,10 +95,11 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::PseudoTransi abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu1, u, termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu1, ApplyArray(__zero, J), du; kwargs...) return PseudoTransientCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, alpha, res_norm, uf, linsolve, J, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, - abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache) + abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache, trace) end function perform_step!(cache::PseudoTransientCache{true}) @@ -125,6 +127,9 @@ function perform_step!(cache::PseudoTransientCache{true}) @. u = u - du f(fu1, u, p) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), J, + cache.du) + new_norm = cache.internalnorm(fu1) cache.alpha *= cache.res_norm / new_norm cache.res_norm = new_norm @@ -157,6 +162,9 @@ function perform_step!(cache::PseudoTransientCache{false}) cache.u = @. u - cache.du # `u` might not support mutation cache.fu1 = f(cache.u, p) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), cache.J, + cache.du) + new_norm = cache.internalnorm(fu1) cache.alpha *= cache.res_norm / new_norm cache.res_norm = new_norm @@ -185,6 +193,7 @@ function SciMLBase.reinit!(cache::PseudoTransientCache{iip}, u0 = cache.u; p = c cache.fu1 = cache.f(cache.u, p) end + reset!(cache.trace) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fu1, cache.u, termination_condition) diff --git a/src/raphson.jl b/src/raphson.jl index a28dec699..594b893e5 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -1,5 +1,5 @@ """ - NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = LineSearch(), + NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, precs = DEFAULT_PRECS, adkwargs...) An advanced NewtonRaphson implementation with support for efficient handling of sparse @@ -30,8 +30,7 @@ for large-scale and numerically-difficult nonlinear systems. which means that no line search is performed. Algorithms from `LineSearches.jl` can be used here directly, and they will be converted to the correct `LineSearch`. """ -@concrete struct NewtonRaphson{CJ, AD} <: - AbstractNewtonAlgorithm{CJ, AD} +@concrete struct NewtonRaphson{CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD} ad::AD linsolve precs @@ -42,8 +41,8 @@ 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 = LineSearch(), precs = DEFAULT_PRECS, adkwargs...) +function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, + precs = DEFAULT_PRECS, adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) return NewtonRaphson{_unwrap_val(concrete_jac)}(ad, linsolve, precs, linesearch) @@ -72,12 +71,13 @@ end 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} + termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), + kwargs...) where {uType, iip} alg = get_concrete_algorithm(alg_, prob) @unpack f, u0, p = prob u = alias_u0 ? u0 : deepcopy(u0) @@ -88,10 +88,12 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::NewtonRaphso abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu1, u, termination_condition) + ls_cache = init_linesearch_cache(alg.linesearch, f, u, p, fu1, Val(iip)) + trace = init_nonlinearsolve_trace(alg, u, fu1, ApplyArray(__zero, J), du; kwargs...) + return NewtonRaphsonCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, uf, linsolve, J, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, - NLStats(1, 0, 0, 0, 0), - init_linesearch_cache(alg.linesearch, f, u, p, fu1, Val(iip)), tc_cache) + NLStats(1, 0, 0, 0, 0), ls_cache, tc_cache, trace) end function perform_step!(cache::NewtonRaphsonCache{true}) @@ -108,6 +110,9 @@ function perform_step!(cache::NewtonRaphsonCache{true}) _axpy!(-α, du, u) f(cache.fu1, u, p) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), J, + cache.du, α) + check_and_update!(cache, cache.fu1, cache.u, cache.u_prev) @. u_prev = u @@ -136,6 +141,9 @@ function perform_step!(cache::NewtonRaphsonCache{false}) cache.u = @. u - α * cache.du # `u` might not support mutation cache.fu1 = f(cache.u, p) + update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), cache.J, + cache.du, α) + check_and_update!(cache, cache.fu1, cache.u, cache.u_prev) cache.u_prev = cache.u @@ -159,6 +167,7 @@ function SciMLBase.reinit!(cache::NewtonRaphsonCache{iip}, u0 = cache.u; p = cac cache.fu1 = cache.f(cache.u, p) end + reset!(cache.trace) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fu1, cache.u, termination_condition) diff --git a/src/trace.jl b/src/trace.jl new file mode 100644 index 000000000..c458c7d07 --- /dev/null +++ b/src/trace.jl @@ -0,0 +1,240 @@ +abstract type AbstractNonlinearSolveTraceLevel end + +""" + TraceMinimal(freq) + TraceMinimal(; print_frequency = 1, store_frequency::Int = 1) + +Trace Minimal Information + + 1. Iteration Number + 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)`. +""" +@kwdef struct TraceMinimal <: AbstractNonlinearSolveTraceLevel + print_frequency::Int = 1 + store_frequency::Int = 1 +end + +""" + TraceWithJacobianConditionNumber(freq) + TraceWithJacobianConditionNumber(; print_frequency = 1, store_frequency::Int = 1) + +`TraceMinimal` + Print the Condition Number of the Jacobian. + +## 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)`. +""" +@kwdef struct TraceWithJacobianConditionNumber <: AbstractNonlinearSolveTraceLevel + print_frequency::Int = 1 + store_frequency::Int = 1 +end + +""" + TraceAll(freq) + TraceAll(; print_frequency = 1, store_frequency::Int = 1) + +`TraceWithJacobianConditionNumber` + 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)`. +""" +@kwdef struct TraceAll <: AbstractNonlinearSolveTraceLevel + print_frequency::Int = 1 + store_frequency::Int = 1 +end + +for Tr in (:TraceMinimal, :TraceWithJacobianConditionNumber, :TraceAll) + @eval begin + $(Tr)(freq) = $(Tr)(; print_frequency = freq, store_frequency = freq) + end +end + +# NonlinearSolve Tracing Utilities +@concrete struct NonlinearSolveTraceEntry + iteration::Int + fnorm + stepnorm + condJ + J + u + fu + δu +end + +function __show_top_level(io::IO, entry::NonlinearSolveTraceEntry) + if entry.condJ === nothing + @printf io "%-8s %-20s %-20s\n" "----" "-------------" "-----------" + @printf io "%-8s %-20s %-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm" + @printf io "%-8s %-20s %-20s\n" "----" "-------------" "-----------" + else + @printf io "%-8s %-20s %-20s %-20s\n" "----" "-------------" "-----------" "-------" + @printf io "%-8s %-20s %-20s %-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm" "cond(J)" + @printf io "%-8s %-20s %-20s %-20s\n" "----" "-------------" "-----------" "-------" + end +end + +function Base.show(io::IO, entry::NonlinearSolveTraceEntry) + entry.iteration == 0 && __show_top_level(io, entry) + if entry.iteration < 0 + # Special case for final entry + @printf io "%-8s %-20.8e\n" "Final" entry.fnorm + @printf io "%-28s\n" "----------------------" + elseif entry.condJ === nothing + @printf io "%-8d %-20.8e %-20.8e\n" entry.iteration entry.fnorm entry.stepnorm + else + @printf io "%-8d %-20.8e %-20.8e %-20.8e\n" entry.iteration entry.fnorm entry.stepnorm entry.condJ + end + return nothing +end + +function NonlinearSolveTraceEntry(iteration, fu, δu) + return NonlinearSolveTraceEntry(iteration, norm(fu, Inf), norm(δu, 2), nothing, + nothing, nothing, nothing, nothing) +end + +function NonlinearSolveTraceEntry(iteration, fu, δu, J) + return NonlinearSolveTraceEntry(iteration, norm(fu, Inf), norm(δu, 2), __cond(J), + nothing, nothing, nothing, nothing) +end + +function NonlinearSolveTraceEntry(iteration, fu, δu, J, u) + return NonlinearSolveTraceEntry(iteration, norm(fu, Inf), norm(δu, 2), __cond(J), + __copy(J), __copy(u), __copy(fu), __copy(δu)) +end + +__cond(J::AbstractMatrix) = cond(J) +__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 + trace_level::Tr +end + +function reset!(trace::NonlinearSolveTrace) + (trace.history !== nothing && resize!(trace.history, 0)) +end + +function Base.show(io::IO, trace::NonlinearSolveTrace) + for entry in trace.history + show(io, entry) + end + return nothing +end + +function init_nonlinearsolve_trace(alg, u, fu, J, δu; show_trace::Val = Val(false), + trace_level::AbstractNonlinearSolveTraceLevel = TraceMinimal(), + store_trace::Val = Val(false), uses_jac_inverse = Val(false), kwargs...) + return init_nonlinearsolve_trace(alg, show_trace, trace_level, store_trace, u, fu, J, + δu, uses_jac_inverse) +end + +function init_nonlinearsolve_trace(alg, ::Val{show_trace}, + trace_level::AbstractNonlinearSolveTraceLevel, ::Val{store_trace}, u, fu, J, + δu, ::Val{uses_jac_inverse}) where {show_trace, store_trace, uses_jac_inverse} + if show_trace + print("\nAlgorithm: ") + Base.printstyled(alg, "\n\n"; color = :green, bold = true) + end + J_ = uses_jac_inverse ? (trace_level isa TraceMinimal ? J : inv(J)) : J + history = __init_trace_history(Val{show_trace}(), trace_level, Val{store_trace}(), u, + fu, J_, δu) + return NonlinearSolveTrace{show_trace, store_trace}(history, trace_level) +end + +function __init_trace_history(::Val{show_trace}, trace_level, ::Val{store_trace}, u, fu, J, + δu) where {show_trace, store_trace} + !store_trace && !show_trace && return nothing + entry = __trace_entry(trace_level, 0, u, fu, J, δu) + show_trace && show(entry) + store_trace && return [entry] + return nothing +end + +function __trace_entry(::TraceMinimal, iter, u, fu, J, δu, α = 1) + return NonlinearSolveTraceEntry(iter, fu, δu .* α) +end +function __trace_entry(::TraceWithJacobianConditionNumber, iter, u, fu, J, δu, α = 1) + return NonlinearSolveTraceEntry(iter, fu, δu .* α, J) +end +function __trace_entry(::TraceAll, iter, u, fu, J, δu, α = 1) + return NonlinearSolveTraceEntry(iter, fu, δu .* α, J, u) +end + +function update_trace!(trace::NonlinearSolveTrace{ShT, StT}, iter, u, fu, J, δu, + α = 1; last::Val{L} = Val(false)) where {ShT, StT, L} + !StT && !ShT && return nothing + + if L + entry = NonlinearSolveTraceEntry(-1, norm(fu, Inf), NaN32, nothing, nothing, + nothing, nothing, nothing) + ShT && show(entry) + return trace + end + + show_now = ShT && (iter % trace.trace_level.print_frequency == 1) + store_now = StT && (iter % trace.trace_level.store_frequency == 1) + (show_now || store_now) && (entry = __trace_entry(trace.trace_level, iter, u, fu, J, + δu, α)) + store_now && push!(trace.history, entry) + show_now && show(entry) + return trace +end + +# Needed for Algorithms which directly use `inv(J)` instead of `J` +function update_trace_with_invJ!(trace::NonlinearSolveTrace{ShT, StT}, iter, u, fu, J, δu, + α = 1; last::Val{L} = Val(false)) where {ShT, StT, L} + !StT && !ShT && return nothing + + if L + entry = NonlinearSolveTraceEntry(-1, norm(fu, Inf), NaN32, nothing, nothing, + nothing, nothing, nothing) + show(entry) + return trace + end + + show_now = ShT && (iter % trace.trace_level.print_frequency == 1) + store_now = StT && (iter % trace.trace_level.store_frequency == 1) + if show_now || store_now + J_ = trace.trace_level isa TraceMinimal ? J : inv(J) + entry = __trace_entry(trace.trace_level, iter, u, fu, J_, δu, α) + end + store_now && push!(trace.history, entry) + show_now && show(entry) + return trace +end diff --git a/src/trustRegion.jl b/src/trustRegion.jl index d0149c31e..8b4041b75 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -230,6 +230,7 @@ end ϵ::floatType stats::NLStats tc_cache + trace end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::TrustRegion, args...; @@ -339,13 +340,14 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::TrustRegion, abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu1, u, termination_condition) + trace = init_nonlinearsolve_trace(alg, u, fu1, ApplyArray(__zero, J), du; kwargs...) return TrustRegionCache{iip}(f, alg, u_prev, u, fu_prev, fu1, fu2, p, uf, linsolve, J, jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob, radius_update_scheme, initial_trust_radius, max_trust_radius, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, loss, loss_new, H, g, shrink_counter, du, u_tmp, u_gauss_newton, u_cauchy, fu_new, make_new_J, r, - p1, p2, p3, p4, ϵ, NLStats(1, 0, 0, 0, 0), tc_cache) + p1, p2, p3, p4, ϵ, NLStats(1, 0, 0, 0, 0), tc_cache, trace) end function perform_step!(cache::TrustRegionCache{true}) @@ -462,6 +464,8 @@ function trust_region_step!(cache::TrustRegionCache) # No need to make a new J, no step was taken, so we try again with a smaller trust_r cache.make_new_J = false end + update_trace!(cache.trace, cache.stats.nsteps + 1, cache.u, cache.fu, cache.J, + @~(cache.u.-cache.u_prev)) check_and_update!(cache, cache.fu, cache.u, cache.u_prev) elseif radius_update_scheme === RadiusUpdateSchemes.NLsolve @@ -483,6 +487,8 @@ function trust_region_step!(cache::TrustRegionCache) cache.trust_r = max(cache.trust_r, 2 * norm(cache.du)) # cache.expand_factor * norm(cache.du)) end + update_trace!(cache.trace, cache.stats.nsteps + 1, cache.u, cache.fu, cache.J, + @~(cache.u.-cache.u_prev)) # convergence test check_and_update!(cache, cache.fu, cache.u, cache.u_prev) @@ -502,6 +508,8 @@ function trust_region_step!(cache::TrustRegionCache) cache.trust_r = min(2 * cache.trust_r, cache.max_trust_r) end + update_trace!(cache.trace, cache.stats.nsteps + 1, cache.u, cache.fu, cache.J, + @~(cache.u.-cache.u_prev)) # convergence test check_and_update!(cache, cache.fu, cache.u, cache.u_prev) @@ -524,6 +532,8 @@ function trust_region_step!(cache::TrustRegionCache) cache.trust_r = rfunc(r, shrink_threshold, p1, p3, p4, p2) * cache.internalnorm(du) + update_trace!(cache.trace, cache.stats.nsteps + 1, cache.u, cache.fu, cache.J, + @~(cache.u.-cache.u_prev)) check_and_update!(cache, cache.fu, cache.u, cache.u_prev) cache.internalnorm(g) < cache.ϵ && (cache.force_stop = true) @@ -547,6 +557,9 @@ function trust_region_step!(cache::TrustRegionCache) @unpack p1 = cache cache.trust_r = p1 * cache.internalnorm(jvp!(cache)) + + update_trace!(cache.trace, cache.stats.nsteps + 1, cache.u, cache.fu, cache.J, + @~(cache.u.-cache.u_prev)) check_and_update!(cache, cache.fu, cache.u, cache.u_prev) cache.internalnorm(g) < cache.ϵ && (cache.force_stop = true) #Fan's update scheme @@ -569,6 +582,9 @@ function trust_region_step!(cache::TrustRegionCache) @unpack p1 = cache cache.trust_r = p1 * (cache.internalnorm(cache.fu)^0.99) + + update_trace!(cache.trace, cache.stats.nsteps + 1, cache.u, cache.fu, cache.J, + @~(cache.u.-cache.u_prev)) check_and_update!(cache, cache.fu, cache.u, cache.u_prev) cache.internalnorm(g) < cache.ϵ && (cache.force_stop = true) elseif radius_update_scheme === RadiusUpdateSchemes.Bastin @@ -585,6 +601,9 @@ function trust_region_step!(cache::TrustRegionCache) cache.trust_r *= cache.p2 cache.shrink_counter += 1 end + + update_trace!(cache.trace, cache.stats.nsteps + 1, cache.u, cache.fu, cache.J, + @~(cache.u.-cache.u_prev)) check_and_update!(cache, cache.fu, cache.u, cache.u_prev) end end @@ -707,6 +726,7 @@ function SciMLBase.reinit!(cache::TrustRegionCache{iip}, u0 = cache.u; p = cache cache.fu = cache.f(cache.u, p) end + reset!(cache.trace) abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, cache.fu, cache.u, termination_condition) diff --git a/src/utils.jl b/src/utils.jl index c5161df7c..bf6d1152f 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -354,3 +354,10 @@ end # Define special concatenation for certain Array combinations @inline _vcat(x, y) = vcat(x, y) + +__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) diff --git a/test/basictests.jl b/test/basictests.jl index e3928a7bc..b26db9d5b 100644 --- a/test/basictests.jl +++ b/test/basictests.jl @@ -28,13 +28,13 @@ const TERMINATION_CONDITIONS = [ # --- NewtonRaphson tests --- @testset "NewtonRaphson" begin - function benchmark_nlsolve_oop(f, u0, p = 2.0; linesearch = LineSearch()) + function benchmark_nlsolve_oop(f, u0, p = 2.0; linesearch = nothing) prob = NonlinearProblem{false}(f, u0, p) return solve(prob, NewtonRaphson(; linesearch), abstol = 1e-9) end function benchmark_nlsolve_iip(f, u0, p = 2.0; linsolve, precs, - linesearch = LineSearch()) + linesearch = nothing) prob = NonlinearProblem{true}(f, u0, p) return solve(prob, NewtonRaphson(; linsolve, precs, linesearch), abstol = 1e-9) end @@ -692,12 +692,12 @@ end # --- GeneralBroyden tests --- @testset "GeneralBroyden" begin - function benchmark_nlsolve_oop(f, u0, p = 2.0; linesearch = LineSearch()) + function benchmark_nlsolve_oop(f, u0, p = 2.0; linesearch = nothing) prob = NonlinearProblem{false}(f, u0, p) return solve(prob, GeneralBroyden(; linesearch), abstol = 1e-9) end - function benchmark_nlsolve_iip(f, u0, p = 2.0; linesearch = LineSearch()) + function benchmark_nlsolve_iip(f, u0, p = 2.0; linesearch = nothing) prob = NonlinearProblem{true}(f, u0, p) return solve(prob, GeneralBroyden(; linesearch), abstol = 1e-9) end @@ -789,12 +789,12 @@ end # --- GeneralKlement tests --- @testset "GeneralKlement" begin - function benchmark_nlsolve_oop(f, u0, p = 2.0; linesearch = LineSearch()) + function benchmark_nlsolve_oop(f, u0, p = 2.0; linesearch = nothing) prob = NonlinearProblem{false}(f, u0, p) return solve(prob, GeneralKlement(; linesearch), abstol = 1e-9) end - function benchmark_nlsolve_iip(f, u0, p = 2.0; linesearch = LineSearch()) + function benchmark_nlsolve_iip(f, u0, p = 2.0; linesearch = nothing) prob = NonlinearProblem{true}(f, u0, p) return solve(prob, GeneralKlement(; linesearch), abstol = 1e-9) end @@ -886,14 +886,14 @@ end # --- LimitedMemoryBroyden tests --- @testset "LimitedMemoryBroyden" begin - function benchmark_nlsolve_oop(f, u0, p = 2.0; linesearch = LineSearch(), + function benchmark_nlsolve_oop(f, u0, p = 2.0; linesearch = nothing, termination_condition = AbsNormTerminationMode()) prob = NonlinearProblem{false}(f, u0, p) return solve(prob, LimitedMemoryBroyden(; linesearch); abstol = 1e-9, termination_condition) end - function benchmark_nlsolve_iip(f, u0, p = 2.0; linesearch = LineSearch(), + function benchmark_nlsolve_iip(f, u0, p = 2.0; linesearch = nothing, termination_condition = AbsNormTerminationMode()) prob = NonlinearProblem{true}(f, u0, p) return solve(prob, LimitedMemoryBroyden(; linesearch); abstol = 1e-9,