OOP NonlinearSolveAlg: reuse W with SimpleNonlinearSolve solvers

lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl

@@ -20,7 +20,12 @@ using NonlinearSolve: FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg, New
     HomotopySweep, HomotopyPolyAlgorithm, ArcLengthContinuation, step!
 # The operator Jacobian path is implemented in NonlinearSolveBase and needs its own floor.
 import NonlinearSolveBase
-using NonlinearSolveBase: get_linear_cache
+# An algorithm that defines no `__init` (every SimpleNonlinearSolve one) gets this fallback
+# cache, which records the arguments a later `solve!` will forward and nothing else: no
+# `step!`, no `stats`, none of the iteration state `not_terminated` reads. The inner-solver
+# branches key off the cache type rather than the algorithm type so a newly exported
+# algorithm cannot quietly pick up the `step!`-driven path.
+using NonlinearSolveBase: get_linear_cache, NonlinearSolveNoInitCache
 using MuladdMacro: @muladd
 using FastBroadcast: @..
 import FastClosures: @closure
@@ -59,7 +64,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
 

lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl

@@ -41,20 +41,36 @@ 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. `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
         # `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 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
@@ -233,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)
@@ -242,21 +269,37 @@ 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; `_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`
+        # 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)
+        nlcache.u
     end
-    step!(nlcache; recompute_jacobian)
-    nlsolver.ztmp = nlcache.u
+    nlsolver.ztmp = znew
 
     ustep = compute_ustep(tmp, γ, z, method)
     atmp = calculate_residuals(
-        z .- nlcache.u, uprev, ustep, opts.abstol, opts.reltol,
+        z .- znew, uprev, ustep, opts.abstol, opts.reltol,
         opts.internalnorm, t
     )
     ndz = opts.internalnorm(atmp, t)
@@ -799,7 +842,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)

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
 
 """

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,9 +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.
+                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

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