Skip to content

NonlinearSolveAlg: reuse stored J on dt-only W updates (do_newJW parity)#3820

Merged
ChrisRackauckas merged 2 commits into
SciML:masterfrom
ChrisRackauckas-Claude:nsa-split-j-w-updates
Jul 6, 2026
Merged

NonlinearSolveAlg: reuse stored J on dt-only W updates (do_newJW parity)#3820
ChrisRackauckas merged 2 commits into
SciML:masterfrom
ChrisRackauckas-Claude:nsa-split-j-w-updates

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor

Please ignore until reviewed by @ChrisRackauckas. Draft opened by an agent.

Problem

The W-reuse path recomputes the Jacobian on every W update, including pure dt/γ changes. The legacy do_newJW policy separates the two decisions: a dt change only reassembles W = J - M/γdt from the stored J via jacobian2W! (O(nnz)); a fresh Jacobian evaluation is reserved for first use, always_new, or a diverged iteration.

Measured on the dense-AD 512-unknown brusselator at rtol 1e-8 (master): FBDF+NLNewton evaluates 1 Jacobian for the whole solve (nf = 2,764) while FBDF+NSA evaluates 295 (nf = 152,393) — every dt-driven W refresh pays a full 512-dim AD Jacobian.

Fix

_update_nlsolvealg_W! takes a new_jac flag; initialize! computes new_jac = first_call || always_new || status === Divergence separately from the dt-change test, mirroring do_newJW's jbad/wbad split. dt-only updates reuse cache.J and only rerun jacobian2W!.

The expected trade (same one NLNewton makes) is visible on cheap-J ROBER: user-jac calls drop ~100 → ~4, iteration counts rise to NLNewton's level (1433 vs NLNewton's 1320, from 911 with always-fresh J), wall time unchanged. On expensive-J problems it's decisive — dense brusselator NSA/NLNewton time ratios, master → with this series:

alg rtol 1e-4 1e-6 1e-8
FBDF 5.0× → 0.96× 5.6× → 1.10× 7.5× → 1.29×
QNDF 4.1× → 0.92× 6.8× → 1.16× 7.9× → 0.88×
KenCarp4 17× → 1.20× 39× → 1.79× 66× → 1.72×
TRBDF2 20× → 1.65× 55× → 2.24× 70× → 1.49×

(together with #3819 and SciML/NonlinearSolve.jl#1003; the remaining TRBDF2/KenCarp4 gap is step-count inflation from dt-controller/η coupling, a separate issue.)

Tests

New nsa_jacobian_reuse_tests.jl counts calls to a user-supplied jac on ROBER (the njacs stat can't observe this — it's dominated by the inner analytic-jac closure copies, one per W update): asserts JAC_CALLS < stats.nw / 3. Verified the test fails without this change (JAC_CALLS ≈ nw, 2 failures) and passes with it. Solution accuracy asserted against a 1e-12 reference. Full-stack sparse brusselator + KLU reaches 1.00× of NLNewton.

🤖 Generated with Claude Code

The W-reuse path recomputed the Jacobian on every W update, including pure
dt/gamma changes, so njacs tracked nw one-to-one (~10x the legacy count: on
ROBER at 1e-8, NLNewton evaluates 9 Jacobians vs 94/95 for NSA). The legacy
do_newJW policy separates the two decisions: a dt change only reassembles
W = J - M/γdt from the stored J via jacobian2W! (O(nnz)); a fresh Jacobian is
reserved for first use, always_new, or a diverged iteration.

Apply the same split here: _update_nlsolvealg_W! takes a new_jac flag and
skips the f.jac/jacobian! evaluation when only the γdt scaling changed.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
…e copies)

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor Author

CI triage: the failing lanes here are master-baseline, not from this change — StochasticDiffEq QA/downgrade + GPU lanes fail on all concurrent PRs including the 2-line #3816; the OrdinaryDiffEqNonlinearSolve Core lane failure is the pre-existing Dependent Mass Matrix Tests error (mass_matrix_tests.jl:283, OOP NLNewton + non-concrete operator W, 221 pass/2 error/4 broken — the limitation #3796 documents), reproduced identically on unmodified master locally. The MTK nlstep lane failures are the two order-test failures that #3818 fixes.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants