NonlinearSolveAlg OOP: extend W-reuse to WOperator/StaticWOperator

[singhharsh1708]

The OOP NonlinearSolveAlg branch of build_nlsolver builds J, W via build_J_W and then discards both — the NonlinearProblem is built with no jac, and NonlinearSolveCache stores nothing for J/W. Every Newton iter inside the OOP path recomputes the Jacobian via AD/FD with no cross-step reuse.

This is the OOP mirror of #3753 (which fixed the same gap on the IIP path for WOperator with a concrete underlying J). The architecture is identical: detect a reusable W at build time, wire a closure jac into the inner NonlinearProblem, store J/W on the cache so compute_step!'s existing recompute_jacobian flow picks them up.

Scope

Enables W-reuse when W isa Union{WOperator, StaticWOperator} (the case produced for OOP problems with concrete_jac=true, a matrix-free linsolve like KrylovJL_GMRES, or a SciMLOperator jac_prototype). Other W types (plain Matrix, Factorization, generic AbstractSciMLOperator) fall through to the existing no-reuse default — so this is purely additive.

Changes

1. utils.jl OOP build_nlsolver for NonlinearSolveAlg: detect use_w_reuse, build a closure (z, p) -> convert(AbstractMatrix, W), pass it as jac to NonlinearFunction, store J, W on NonlinearSolveCache.
2. newton.jl OOP initialize! for NonlinearSolveAlg: when cache.W !== nothing, run the same first-call / new_W_dt_cutoff / divergence logic the IIP path uses, calling update_coefficients!(cache.W, uprev, p, tstep; gamma = dtgamma) to refresh the WOperator in place and toggling cache.new_W for compute_step!.

Local validation

Robertson, OOP, ImplicitEuler(linsolve = KrylovJL_GMRES(), concrete_jac = true, ...):

	retcode	steps	nf	njacs	nw	norm(u_end vs NLNewton)
NLNewton	Success	1302	2560	641	641	—
NonlinearSolveAlg	Success	1302	2386	0	641	2.3e-13

njacs = 0 on the NS side reflects that the inner solver no longer computes J itself — the WOperator's update_coefficients! does it once per step instead. nw = 641 over 1302 steps means W is reused across ~2 steps on average, matching the NLNewton modified-Newton cadence.

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.
• NonlinearSolveAlg OOP: extend W-reuse to WOperator/StaticWOperator #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.