OOP NonlinearSolveAlg: reuse W with SimpleNonlinearSolve solvers

[singhharsh1708]

The out-of-place NonlinearSolveAlg branch of build_nlsolver builds J, W through build_J_W and then throws both away: the inner NonlinearProblem gets no jac, and NonlinearSolveCache stores nothing for J and W. Every inner Newton iteration recomputes the Jacobian by AD, with no reuse across stages or steps.

This wires the W back in for the one case that is safe today: a StaticWOperator problem with a SimpleNonlinearSolve inner solver. The W lives in a Ref that 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 the way the StaticWOperator branch of calc_W does.

The gate is deliberately narrow, and narrower than the first revision of this PR claimed. A concrete-J WOperator cannot be handed to the inner solver as a plain matrix the same way; it needs the buffer-copy pattern from #3823, where the assembled W is copied into a Jacobian buffer the inner solver owns. Wrapping a matrix-free WOperator so the inner solver applies it by mul! is a possible follow-up as well. Neither is implemented here, so every other W type keeps the existing no-reuse default and this stays additive.

Getting there first needed the out-of-place path to handle SimpleNonlinearSolve at all, which it did not. 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 (force_stop, nsteps, maxiters). So compute_step! drives it with solve!, _rejected_trial_step skips it instead of throwing the moment the displacement hits exactly zero, and the inner cache keeps the solver's own tolerances rather than the zeroed ones the step!-driven path uses, which would send every inner solve to maxiters. All three branch on the cache type, not on a list of algorithm types, so a newly exported SimpleNonlinearSolve algorithm cannot quietly pick up the wrong path. SimpleHomotopySweep, added in 2.13, is exactly what a hand-written list would have missed.

Stats on this path are incomplete, and the code says so instead of pretending otherwise. SimpleNonlinearSolve builds every solution without stats, so sol.stats === nothing and the residual and linear-solve work its solvers do is not attributable to the integrator. What is counted is what this side performs: njacs and nw from the W assemblies. The out-of-place initialize! skips the counter copy outright for these caches rather than reading fields they do not have, and the njacs guard it inherited goes with it, since that guard's comment described a WReuseJac that only exists on the in-place path and W reuse here only ever pairs with a solve!-only cache.

Testing, against master at 0be95df:

• linear_nonlinear_tests.jl 54 pass, up from 44 (10 added)
• nsa_matrixfree_tests.jl 14 pass, unchanged
• newton_tests.jl 6 pass, unchanged
• nsa_jacobian_reuse_tests.jl 17 pass, unchanged

The added case solves a stiff van der Pol in SVector form with TRBDF2 under both SimpleNewtonRaphson and SimpleTrustRegion, checks the result against TRBDF2 with the default nonlinear solver, and asserts njacs == nw with the assembly count strictly below the nonlinear iteration count, which is the reuse actually happening. On that problem SimpleNewtonRaphson goes from 22532 nonlinear iterations without reuse to 14026 with it, over 1343 W assemblies, and the end state agrees with the NLNewton run to 2e-4 either way.

nested_ad_nlsolvealg_tests.jl errors the same way on master (LinearSolve/SciMLOperators version skew) and is untouched here.

AI Disclosure

Claude assisted with this work.

[ChrisRackauckas]

We shouldn't unconditionally convert.

[singhharsh1708]

good catch — split the closure so StaticWOperator returns its .W field directly, only WOperator path converts now

[ChrisRackauckas]

What about non-static opertators

[singhharsh1708]

for non-static WOperator we shouldn't unconditionally convert either.
Updated to mirror the IIP gate from #3753: use_w_reuse now requires either StaticWOperator or a WOperator whose underlying J is concrete (not an AbstractSciMLOperator). The else branch hands NonlinearSolve the cached Ww._concrete_form directly via the closure (no convert call). Matrix-free / non-concrete WOperator cases fall back to the non-w-reuse path.

[ChrisRackauckas]

use the operator...

[singhharsh1708]

fixed — now hands the operator through directly

[singhharsh1708]

Hi @ChrisRackauckas — quick nudge on this stack when you have a moment. All three PRs (#3795, #3796, #3798) have addressed the earlier review comments and are rebased on the latest master:

• NonlinearSolveAlg: pipe wrapprecs into inner alg and update weight per step #3795 — Added the explicit KrylovJL_GMRES() linsolve test you asked for, plus a weight-buffer check that verifies it's allocated and filled after one step.
• NLNewton OOP: handle AbstractSciMLOperator W instead of erroring #3796 — Reverted the operator convert; kept the informative error you suggested for the OOP path.
• OOP NonlinearSolveAlg: reuse W with SimpleNonlinearSolve solvers #3798 — Mirrored the IIP gate from NonlinearSolveAlg: extend W-reuse to WOperator with concrete underlying Jacobian #3753 (use_w_reuse requires StaticWOperator or a WOperator with a concrete underlying J), and now hands the operator through directly instead of Ww._concrete_form (which was stale between Newton iterations since update_coefficients! doesn't refresh _concrete_form).

Once the next SciMLOperators release tags, I'll open a follow-up here to route OOP + SciMLOperator mass matrix through WOperator and drop #3796's informative error for that path.

No rush — just wanted to make sure these weren't waiting on me. Thanks!

[ChrisRackauckas-Claude]

Ran this branch against an OOP nonlinear problem and hit two defects in the W-reuse path — posting reproductions so they can be addressed:

1. The WOperator jac closure crashes on first Jacobian use (plain Vector problems).

using OrdinaryDiffEqSDIRK, OrdinaryDiffEqNonlinearSolve, NonlinearSolve, ADTypes
using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
vdp(u, p, t) = [u[2], p[1] * ((1 - u[1]^2) * u[2] - u[1])]
prob = ODEProblem(vdp, [2.0, 0.0], (0.0, 1.0), [1.0e3])
integ = init(prob, TRBDF2(nlsolve = NonlinearSolveAlg(NewtonRaphson(; autodiff = AutoForwardDiff())), concrete_jac = true); reltol = 1e-8, abstol = 1e-10)
step!(integ)
# ERROR: MethodError: Cannot `convert` an object of type WOperator{false, ...} to an object of type Matrix{Float64}

The (z, p) -> Ww closure hands the raw WOperator to NonlinearSolve, whose jacobian cache holds a dense Matrix (no jac_prototype on the inner NonlinearFunction), so the assignment tries convert(Matrix, WOperator) and throws. init succeeds — it dies on the first step! that evaluates the jac, which is also why this wasn't caught: the PR currently has no tests. A test exercising exactly the TRBDF2 + concrete_jac=true + Vector u0 path above (plus a StaticArrays one for StaticWOperator) would have flagged it.

2. Even with the convert fixed, the Jacobian values are never refreshed — J is frozen at its build-time state for the entire integration.

The update in initialize! is update_coefficients!(cache.W, uprev, p, tstep; gamma = dtgamma). For WOperator that forwards to update_coefficients!(W.J, u, p, t) (SciMLOperators woperator.jl), which is a no-op when W.J is a plain matrix — and the use_w_reuse gate specifically selects for a concrete non-operator W.J. So only gamma ever updates; the J entries stay at their build_nlsolver-time values. (For StaticWOperator it's worse: Ws.W is a fixed matrix, so not even gamma changes.) Newton silently degrades to a chord iteration with J(u₀, t₀) forever — no failure on mildly nonlinear problems, just quietly worse convergence, and on strongly state-dependent J (e.g. the VdP case above) divergence→step-rejection loops.

The IIP counterpart (#3753) avoids this because _update_nlsolvealg_W! explicitly recomputes J via f.jac/jacobian! before jacobian2W!. The OOP path needs the same: an explicit J recompute (the OOP analog is calc_J/jacobian(uf, uprev, integrator)) and a W rebuild — for StaticWOperator, reconstructing StaticWOperator(J - mass_matrix * inv(dtgamma)); for WOperator, refreshing the concrete J and gamma together — inside the should_update branch.

Also worth folding in while touching this: pass jac_prototype matching W's structure to the inner NonlinearFunction — without it NonlinearSolve allocates a dense J, which is the direct cause of crash 1 and also silently densifies sparse problems (I've opened a separate PR fixing that on the IIP path).

[singhharsh1708]

Thanks for the review @ChrisRackauckas-Claude — both defects reproduced and fixed. New commit bd5af110:

Defect 1 (crash on convert): the OOP closure now returns the concrete matrix directly — W_ref[].W for StaticWOperator, convert(AbstractMatrix, W_ref[]) for WOperator — and I also pass jac_prototype matching that shape to the inner NonlinearFunction. NonlinearSolve no longer tries to convert(Matrix, WOperator).

Defect 2 (frozen J): added _update_nlsolvealg_W_oop! (mirrors the IIP _update_nlsolvealg_W!) — inside should_update we now recompute J via f.jac(uprev, p, tstep) or jacobian(uf, uprev, integrator), then either copyto!(W.J, J_new); W.gamma = dtgamma for WOperator or rebuild StaticWOperator(J_new - mass_matrix * inv(dtgamma)) and swap the Ref for StaticWOperator. uf is now stored in the nlcache so jacobian(uf, ...) works on the OOP path.

Ref indirection: the closure captures W_ref instead of W directly, so the StaticWOperator swap in initialize! is visible to NonlinearSolve without rebuilding the outer NonlinearProblem. For WOperator this is just a passthrough since it's already mutable.

Tests: added two regressions at the end of linear_nonlinear_tests.jl covering exactly the bot's two failure modes:

• Plain Vector VdP + TRBDF2 + concrete_jac = true — used to crash; now succeeds with njacs > 0.
• SVector VdP variant — same, exercising the StaticWOperator path.

Locally: both new tests pass; existing linear_nonlinear_tests suite still passes.

[ChrisRackauckas]

no it's still concretizing.

[ChrisRackauckas]

OOP is purposefully not mutable. We want to make things fold away when static. This will violate that. I think it's simply not possible to do similar W-reuse in the OOP path without violating the principle of making this path non-mutating.

[singhharsh1708]

For StaticWOperator the payload is an immutable SMatrix — reassigning nlcache.W[] re-points the Ref cell, the SMatrix itself is never mutated, so fold-away on the value side still holds. The mutation you're objecting to only shows up on the WOperator branch (copyto!(W.J, …)).
And picking NonlinearSolveAlg already forces a Matrix{Float64} jac cache inside construct_jacobian_cache — the OOP path isn't fold-away in that mode regardless of what this PR does. If Static-only reuse is still off the table given that, I'll close.

[ChrisRackauckas]

I guess it's fine if there's a simplediffeq path that fully non-allocates with simplenonlinearsolve. But the conversion stuff is never okay.

[ChrisRackauckas]

On "still concretizing": NonlinearSolveBase.construct_jacobian_cache calls safe_similar(jac_prototype) and allocates a Matrix{Float64} regardless of what the closure returns — so the OOP jac already flows through a mutable buffer once NonlinearSolveAlg is picked. If that's the real blocker, the fix belongs upstream (similar/copyto! for WOperator in SciMLOperators), not here.

That should get fixed upstream. Making everything dense is not a valid answer.

[singhharsh1708]

Dropped WOperator branch. W-reuse only fires on StaticWOperator + AbstractSimpleNonlinearSolveAlgorithm — SMatrix rebuilt, Ref re-pointed, no convert.

[ChrisRackauckas]

Why no WOperator?

[singhharsh1708]

@ChrisRackauckas Rebased on master and fixed two bugs the Static+Simple gate exposed. Both surfaced because that path introduces execution paths that weren't previously exercised.

1. Inner cache doesn't support step!

SimpleNonlinearSolve algorithms init to a NonlinearSolveNoInitCache, which only implements solve!. compute_step! called step! on it:

FieldError: type NonlinearSolveNoInitCache has no field `force_stop`

compute_step! now branches on the cache type and calls solve! for that case, pulling stats off the returned solution.

2. StaticWOperator.W holds inv(W) for n ≤ 7

StaticWOperator's constructor sets isinv = n <= 7 and then stores inv(W) in the W field. My jac closure returned W_ref[].W, so the inner Newton was handed the inverse Jacobian. This didn't crash — VdP just ran with wrong math, dt collapsed to 5e-324, and it bailed with Unstable and njacs = 0.

The Ref now holds the plain J - M/γdt SMatrix and the closure returns that directly. Verified against a ForwardDiff-computed Jacobian:

closure W  = [-7.032484918065255e9 1.0; -1000.0 -7.032487918065255e9]
J - M/γdt  = [-7.032484918065255e9 1.0; -1000.0 -7.032487918065255e9]
match: true

Test change

sol.stats.njacs > 0 doesn't work as a reuse assertion here — SimpleNonlinearSolve returns stats = nothing, so the counter never moves regardless. Replaced with a check that the W Ref is live on the cache plus an accuracy comparison against TRBDF2() with NLNewton (end-state diff 2e-4).

Validation

All 11 OrdinaryDiffEqNonlinearSolve Core suites pass locally on this branch (Julia 1.12), including Mass Matrix, W-Operator Prototype, and NSA Jacobian Reuse. Runic clean.

Remaining CI reds are unrelated to this branch — constructorof/TimeDerivativeWrapper on InterfaceI/Downstream, and the nlstep order failures that #3818 addresses.