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.

Author: singhharsh1708

lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl

@@ -64,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

@@ -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`
@@ -811,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/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

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