From 9cef5937032d6f8dcb74e319d51735d0d5d34121 Mon Sep 17 00:00:00 2001 From: singhharsh1708 Date: Fri, 31 Jul 2026 15:22:27 +0530 Subject: [PATCH 1/2] OOP NonlinearSolveAlg: support solve!-only inner solvers SimpleNonlinearSolve defines no `__init`, so all of its algorithms land in `NonlinearSolveBase.NonlinearSolveNoInitCache`: a cache that records the arguments a later `solve!` will forward and nothing else. It has no `step!`, no `stats` field, and none of the iteration state `not_terminated` reads, so the out-of-place path threw on the first `step!`. Drive those caches with `solve!`, skip the counter copy in `initialize!`, and keep the solver's own tolerances instead of the zeroed ones the `step!`-driven path uses, which would send every inner solve to `maxiters`. `_rejected_trial_step` skips them as well: a whole inner solve ran to its own termination, so a zero displacement there is convergence, and `compute_step!` already turned an unsuccessful return into `Inf`. Each branch keys off the cache type rather than a list of algorithm types, so a newly exported solver cannot quietly pick up the wrong path. --- .../src/OrdinaryDiffEqNonlinearSolve.jl | 6 ++- .../src/newton.jl | 51 +++++++++++-------- .../src/nlsolve.jl | 8 ++- lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl | 11 +++- 4 files changed, 53 insertions(+), 23 deletions(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl index 332c7492bb..afc801db69 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl @@ -22,7 +22,11 @@ using NonlinearSolve: FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg, New import NonlinearSolveBase # `get_u`/`get_fu` are the only inner-state reads that hold for every inner cache type: # polyalgorithm caches keep `u`/`fu` on the active branch, not as top-level fields. -using NonlinearSolveBase: get_linear_cache, get_u, get_fu +# `NonlinearSolveNoInitCache` is the fallback cache for algorithms that define no `__init` +# (every SimpleNonlinearSolve one): it records the arguments a later `solve!` will forward +# and nothing else, so the inner-solver branches key off the cache type rather than the +# algorithm type and a newly exported algorithm cannot quietly pick up the `step!` path. +using NonlinearSolveBase: get_linear_cache, get_u, get_fu, NonlinearSolveNoInitCache using MuladdMacro: @muladd using FastBroadcast: @.. import FastClosures: @closure diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl index c1093f07ab..955cedc3ae 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl @@ -41,18 +41,15 @@ function initialize!( cache.tstep = integrator.t + nlsolver.c * dt (; ustep, tstep, k, invγdt) = cache - if SciMLBase.has_stats(integrator) + # A `solve!`-only cache keeps no counters, and `reinit!` on it only remakes the problem + # rather than re-evaluating the residual, so neither term below applies. The inner + # solver on the remaining path owns its own Jacobian, so its `njacs` count stands. + if SciMLBase.has_stats(integrator) && !(cache.cache isa NonlinearSolveNoInitCache) # The `reinit!` below evaluates the residual at the new `z` and *then* zeroes the # inner cache's counters, so that evaluation is never visible in # `cache.cache.stats.nf`. Left uncounted it loses one `f` call per stage. integrator.stats.nf += cache.cache.stats.nf + 1 - # Under `W` reuse the inner solver's "Jacobian" is `WReuseJac`, which copies the - # `W` assembled here rather than evaluating anything, so its `njacs` counts copies - # (it tracks `nw`, not Jacobian evaluations). `_update_nlsolvealg_W!` counts the - # real ones. Without reuse the inner solver owns the Jacobian and its count stands. - if cache.W === nothing - integrator.stats.njacs += cache.cache.stats.njacs - end + integrator.stats.njacs += cache.cache.stats.njacs integrator.stats.nsolve += cache.cache.stats.nsolve end if f isa DAEFunction @@ -242,22 +239,36 @@ end nlcache = nlsolver.cache.cache recompute_jacobian = nlsolver.iter == 1 && (cache.W === nothing || cache.new_W) - # A terminated inner cache makes `step!` a no-op, which would leave `ndz == 0` and read to - # the outer convergence test as a perfect solve. Terminating *successfully* is the normal - # path (the inner solve converged before the integrator's own test was satisfied); any - # other terminal state is a genuine failure and must reach the step controller, so report - # it the way `NLNewton` reports a failed linear solve. - if !NonlinearSolveBase.not_terminated(nlcache) && - !SciMLBase.successful_retcode(nlcache.retcode) - return convert(eltype(z), Inf) + znew = if nlcache isa NonlinearSolveNoInitCache + # No `step!` here, so each outer iteration runs a whole inner solve. Those solvers + # also report no counters at all (SimpleNonlinearSolve builds every solution + # without `stats`), so the residual and linear-solve work they do is not + # attributable to the integrator. + sol = SciMLBase.solve!(nlcache) + # A `solve!` runs to its own termination every outer iteration, so an unsuccessful + # return is a genuine failure and must reach the step controller the way `NLNewton` + # reports a failed linear solve. Without this a maxiters solve that never moved + # would hand back `ndz == 0` and read as convergence. + SciMLBase.successful_retcode(sol.retcode) || return convert(eltype(z), Inf) + sol.u + else + # A terminated inner cache makes `step!` a no-op, which would leave `ndz == 0` and + # read to the outer convergence test as a perfect solve. Terminating *successfully* + # is the normal path (the inner solve converged before the integrator's own test was + # satisfied); any other terminal state is a genuine failure and must reach the step + # controller, so report it the way `NLNewton` reports a failed linear solve. + if !NonlinearSolveBase.not_terminated(nlcache) && + !SciMLBase.successful_retcode(nlcache.retcode) + return convert(eltype(z), Inf) + end + step!(nlcache; recompute_jacobian) + get_u(nlcache) end - step!(nlcache; recompute_jacobian) - active_u = get_u(nlcache) - nlsolver.ztmp = active_u + nlsolver.ztmp = znew ustep = compute_ustep(tmp, γ, z, method) atmp = calculate_residuals( - z .- active_u, uprev, ustep, opts.abstol, opts.reltol, + z .- znew, uprev, ustep, opts.abstol, opts.reltol, opts.internalnorm, t ) ndz = opts.internalnorm(atmp, t) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/nlsolve.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/nlsolve.jl index 1736e3f981..a9c21c2c9d 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/nlsolve.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/nlsolve.jl @@ -9,9 +9,15 @@ # either converged exactly (which genuinely is convergence) or failed, and failures were # already turned into `Inf` by `compute_step!`, so termination is the discriminator between # "already at the root" and "has not decided anything yet". +# A `solve!`-driven inner solver has no trial step to reject: each outer iteration runs a +# whole inner solve that ran to its own termination, and `compute_step!` already turned an +# unsuccessful one into `Inf`, so a zero displacement there is convergence. Its cache also +# carries none of the iteration state `not_terminated` reads. _rejected_trial_step(nlsolver, ndz) = false function _rejected_trial_step(nlsolver::NLSolver{<:NonlinearSolveAlg}, ndz) - return iszero(ndz) && NonlinearSolveBase.not_terminated(nlsolver.cache.cache) + nlcache = nlsolver.cache.cache + nlcache isa NonlinearSolveNoInitCache && return false + return iszero(ndz) && NonlinearSolveBase.not_terminated(nlcache) end """ diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl index c558958a7f..ec5f9653d0 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl @@ -649,11 +649,11 @@ function build_nlsolver( else (tmp, γ, α, tstep, invγdt, DIRK, p, dt, f) end + inner_alg = _nlalg_with_linsolve(nlalg.alg, alg.linsolve) prob = NonlinearProblem( NonlinearFunction{false, SciMLBase.FullSpecialize}(nlf), copy(ztmp), nlp_params ) - inner_alg = _nlalg_with_linsolve(nlalg.alg, alg.linsolve) # Zero tolerances: the integrator owns convergence (see the in-place branch above). cache = init( prob, inner_alg; verbose = verbose.nonlinear_verbosity, @@ -663,6 +663,15 @@ function build_nlsolver( abstol = _inner_lintol(uTolType), reltol = _inner_lintol(uTolType), ) ) + if cache isa NonlinearSolveNoInitCache + # This cache only records the arguments a `solve!` will forward, so + # building one is the cheapest way to ask whether the inner solver + # iterates, and the answer stays right as SimpleNonlinearSolve grows. + # Each outer iteration is then a whole inner solve that has to terminate + # on its own, so it keeps the solver's own tolerances: the zeroed ones + # above would send every one of them to `maxiters`. + cache = init(prob, inner_alg; verbose = verbose.nonlinear_verbosity) + end nlcache = NonlinearSolveCache( nothing, tstep, nothing, nothing, invγdt, prob, cache, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, From 2b2920d74b81b6691a20a1a2e5365c98e84e4fe5 Mon Sep 17 00:00:00 2001 From: singhharsh1708 Date: Fri, 31 Jul 2026 15:26:34 +0530 Subject: [PATCH 2/2] OOP NonlinearSolveAlg: reuse a static W across stages The out-of-place branch of `build_nlsolver` builds `J, W` through `build_J_W` and then throws both away: the inner `NonlinearProblem` gets no `jac` and the cache stores `nothing`, so every inner Newton iteration recomputes the Jacobian by AD with no reuse across stages or steps. Wire the `W` back in for the case that is safe today, a `StaticWOperator` problem with a `solve!`-only inner solver. The `W` lives in a `Ref` the inner `NonlinearFunction`'s analytic `jac` reads, and `initialize!` refreshes it under the same `always_new` / divergence / `new_W_dt_cutoff` policy `NLNewton` uses, assembling it through `calc_J` like the `StaticWOperator` branch of `calc_W`. A concrete-J `WOperator` needs the buffer-copy pattern from #3823 instead, so it stays on the existing no-reuse default. `njacs` and `nw` now count the assemblies, which is the only inner work the integrator can attribute on this path. --- .../src/OrdinaryDiffEqNonlinearSolve.jl | 2 +- .../src/newton.jl | 40 +++++++++++++++++-- lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl | 23 +++++++++-- .../test/linear_nonlinear_tests.jl | 31 +++++++++++++- 4 files changed, 87 insertions(+), 9 deletions(-) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl index afc801db69..9f1ee0ede3 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl @@ -65,7 +65,7 @@ import OrdinaryDiffEqCore: _initialize_dae!, import OrdinaryDiffEqDifferentiation: update_W!, is_always_new, build_uf, build_J_W, WOperator, StaticWOperator, wrapprecs, default_krylov_warm_start, build_jac_config, dolinsolve, - resize_jac_config!, jacobian2W!, jacobian! + resize_jac_config!, jacobian2W!, jacobian!, calc_J import StaticArraysCore: StaticArray diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl index 955cedc3ae..9370ca9cc7 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl @@ -42,8 +42,8 @@ function initialize!( (; ustep, tstep, k, invγdt) = cache # A `solve!`-only cache keeps no counters, and `reinit!` on it only remakes the problem - # rather than re-evaluating the residual, so neither term below applies. The inner - # solver on the remaining path owns its own Jacobian, so its `njacs` count stands. + # rather than re-evaluating the residual, so neither term below applies. `W` reuse is + # gated to those caches, which leaves the inner solver here owning its own Jacobian. if SciMLBase.has_stats(integrator) && !(cache.cache isa NonlinearSolveNoInitCache) # The `reinit!` below evaluates the residual at the new `z` and *then* zeroes the # inner cache's counters, so that evaluation is never visible in @@ -52,6 +52,25 @@ function initialize!( integrator.stats.njacs += cache.cache.stats.njacs integrator.stats.nsolve += cache.cache.stats.nsolve end + + if cache.W !== nothing + dtgamma = method === DIRK ? γ * dt : γ * dt / α + W_γdt = cache.W_γdt + first_call = iszero(W_γdt) + # No stored `J` on this path (`W` is rebuilt from a fresh Jacobian), so a + # dt/gamma drift past the cutoff costs a Jacobian evaluation, not just a + # reassembly as in the in-place `new_jac`/`new_w` split. + should_update = first_call || alg.always_new || + nlsolver.status === Divergence || + abs(inv(dtgamma) / inv(W_γdt) - 1) > oftype(dtgamma, alg.new_W_dt_cutoff) + if should_update + _update_nlsolvealg_W_oop!(cache, integrator, dtgamma) + cache.new_W = true + else + cache.new_W = false + end + end + if f isa DAEFunction nlp_params = (tmp, α, tstep, invγdt, p, dt, uprev, f) else @@ -230,6 +249,17 @@ function rescale_stale_W_rhs!(nlcache, nlsolver, integrator, nlstep_data) return nothing end +function _update_nlsolvealg_W_oop!(nlcache, integrator, dtgamma) + # Same construction as the `StaticWOperator` branch of `calc_W`: `calc_J` picks + # user `jac` vs differentiation through `nlcache.uf` and records the evaluation + # in `integrator.stats.njacs` — the only njacs the reused-`W` path can count. + J_new = calc_J(integrator, nlcache) + nlcache.W[] = J_new - integrator.f.mass_matrix * inv(dtgamma) + nlcache.W_γdt = dtgamma + integrator.stats.nw += 1 + return nothing +end + ## compute_step! @muladd function compute_step!(nlsolver::NLSolver{<:NonlinearSolveAlg, false}, integrator) @@ -243,7 +273,8 @@ end # No `step!` here, so each outer iteration runs a whole inner solve. Those solvers # also report no counters at all (SimpleNonlinearSolve builds every solution # without `stats`), so the residual and linear-solve work they do is not - # attributable to the integrator. + # attributable to the integrator; `_update_nlsolvealg_W_oop!` counts the Jacobian + # and `W` assemblies, which are the part performed on this side. sol = SciMLBase.solve!(nlcache) # A `solve!` runs to its own termination every outer iteration, so an unsuccessful # return is a genuine failure and must reach the step controller the way `NLNewton` @@ -815,7 +846,8 @@ function Base.resize!(nlcache::NonlinearSolveCache, ::AbstractNLSolver, integrat nlcache.weight === nothing || resize!(nlcache.weight, i) nlcache.dz === nothing || resize!(nlcache.dz, i) nlcache.jac_config === nothing || resize_jac_config!(nlcache, integrator) - nlcache.W === nothing || resize_J_W!(nlcache, integrator, i) + # The reused out-of-place `W` is a `Ref` over a static matrix, which has no `resize!`. + nlcache.W === nothing || nlcache.W isa Ref || resize_J_W!(nlcache, integrator, i) # `nlcache.linsolve` is re-pointed by `initialize!` when it rebuilds the inner cache on the # next length mismatch, before any estimator solve — nothing to rebuild here. nlcache.W_γdt = zero(nlcache.W_γdt) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl index ec5f9653d0..58456eaee4 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl @@ -663,18 +663,35 @@ function build_nlsolver( abstol = _inner_lintol(uTolType), reltol = _inner_lintol(uTolType), ) ) + W_ref = nothing if cache isa NonlinearSolveNoInitCache # This cache only records the arguments a `solve!` will forward, so # building one is the cheapest way to ask whether the inner solver # iterates, and the answer stays right as SimpleNonlinearSolve grows. - # Each outer iteration is then a whole inner solve that has to terminate - # on its own, so it keeps the solver's own tolerances: the zeroed ones + if !isdae && f.nlstep_data === nothing && W isa StaticWOperator + # StaticWOperator stores inv(W) in its W field for n <= 7, so the + # Ref holds the plain W matrix; initialize! rewrites it before the + # first solve (W_γdt starts at zero). + W_ref = Ref(W.W) + nlf_jac = let W_ref = W_ref + (z, p) -> W_ref[] + end + prob = NonlinearProblem( + NonlinearFunction{false, SciMLBase.FullSpecialize}( + nlf; jac = nlf_jac, jac_prototype = W.W + ), + copy(ztmp), nlp_params + ) + end + # Each outer iteration is a whole inner solve that has to terminate on + # its own, so it keeps the solver's own tolerances: the zeroed ones # above would send every one of them to `maxiters`. cache = init(prob, inner_alg; verbose = verbose.nonlinear_verbosity) end nlcache = NonlinearSolveCache( nothing, tstep, nothing, nothing, invγdt, prob, cache, - nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, + nothing, W_ref, W_ref === nothing ? nothing : uf, + nothing, nothing, nothing, nothing, nothing, zero(tstep), true ) else diff --git a/lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl b/lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl index 049ee7b034..b46552f840 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl @@ -1,7 +1,8 @@ using OrdinaryDiffEqSDIRK, OrdinaryDiffEqBDF, OrdinaryDiffEqRosenbrock, Test, Random, - LinearAlgebra, LinearSolve, ADTypes + LinearAlgebra, LinearSolve, ADTypes, SciMLBase using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg using NonlinearSolve: NewtonRaphson +using SimpleNonlinearSolve: SimpleNewtonRaphson, SimpleTrustRegion Random.seed!(123) A = 0.01 * rand(3, 3) @@ -208,3 +209,31 @@ let integ = init( step!(integ) @test !iszero(integ.cache.nlsolver.cache.weight) end + +using StaticArrays +vdp_static(u, p, t) = SVector(u[2], p[1] * ((1 - u[1]^2) * u[2] - u[1])) +let + prob = ODEProblem( + vdp_static, SVector(2.0, 0.0), (0.0, 1.0), SVector(1.0e3) + ) + solref = solve(prob, TRBDF2(); reltol = 1.0e-8, abstol = 1.0e-10) + for inner in ( + SimpleNewtonRaphson(; autodiff = AutoForwardDiff()), + SimpleTrustRegion(; autodiff = AutoForwardDiff()), + ) + integ = init( + prob, + TRBDF2(nlsolve = NonlinearSolveAlg(inner), concrete_jac = true); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + @test integ.cache.nlsolver.cache.W isa Base.RefValue + sol = solve!(integ) + @test SciMLBase.successful_retcode(sol.retcode) + @test maximum(abs.(sol.u[end] .- solref.u[end])) < 1.0e-3 + # These inner solvers report no counters of their own, so every Jacobian the + # integrator sees comes from the reused `W`: one evaluation per assembly, and + # fewer assemblies than nonlinear iterations because the `W` is held across them. + @test sol.stats.njacs == sol.stats.nw + @test 0 < sol.stats.nw < sol.stats.nnonliniter + end +end