Skip to content

Commit 689d0b7

Browse files
ChrisRackauckas-ClaudeChrisRackauckasclaude
authored
Add Jacobian reuse for Rosenbrock-W methods (#3075)
* Add Jacobian reuse for Rosenbrock-W methods (rebased onto master) Implements CVODE-inspired Jacobian reuse for Rosenbrock-W methods. W-methods guarantee correctness with a stale Jacobian, so we skip expensive J recomputations when conditions allow: - Reuse J but always rebuild W (cheap LU vs expensive AD/finite-diff) - Recompute J on: first iter, step rejection, callback, resize, gamma ratio change >30%, every 20 accepted steps, algorithm switch - Disabled for: strict Rosenbrock, DAEs, linear problems, CompositeAlgorithm Squashed and rebased from PR #3075 (7 commits) onto current master after substantial upstream restructuring (cache consolidation, generic_rosenbrock.jl deletion, RodasTableau unification). Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Fix AD dual propagation and OOP numerical consistency in Jacobian reuse - Strip ForwardDiff.Dual from dtgamma before storing in JacReuseState (these values are heuristic-only and don't need to carry derivatives) - When new_jac=true in OOP path, delegate to standard calc_W instead of custom W construction to ensure numerical consistency with IIP path (fixes regression in special_interps test for Rosenbrock23) Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Disable Jacobian reuse for non-adaptive solves Non-adaptive solves (adaptive=false) with prescribed timesteps don't benefit meaningfully from J reuse, and it causes IIP/OOP inconsistency: adaptive solves have step rejections (EEst > 1) that trigger fresh J recomputation and reset reuse state, while non-adaptive solves following the same timesteps never reject and thus evolve different reuse state. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Relax special_interps tolerance for Rosenbrock-W methods Rosenbrock-W methods with Jacobian reuse take slightly different adaptive steps than the non-adaptive OOP solve (which uses fresh J every step). This causes interpolation differences at ~1e-8, well below the solver tolerance of 1e-14. Relax the test bound to 1e-7 for W-methods. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Remove Julia 1.10 LTS fallback for calc_rosenbrock_differentiation Import unconditionally; accept LTS failures. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Remove do_newJW delegation; keep all Rosenbrock J/W logic in one place _rosenbrock_jac_reuse_decision now always returns (new_jac, new_W) directly instead of returning nothing to delegate to do_newJW. For Rosenbrock methods (no nlsolver), do_newJW just returned (true, true) anyway — the indirection split logic across two functions for no benefit. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Relax lazy-W utility test tolerance for Rosenbrock-W J reuse The lazy-W vs explicit-W comparison for Rosenbrock23 diverges at ~3e-4 with Jacobian reuse active. Relax atol from 2e-5 to 5e-4. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Disable J reuse for operator-based W (Krylov solvers) WOperator and AbstractSciMLOperator wrap J internally for Krylov solvers. A stale J degrades Krylov convergence, causing order loss (e.g. Rodas5P+Enzyme+KrylovJL dropping from order 5 to 1.7). Always recompute J for these W types. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Fix linear-problem fast path and DelayDiffEq @test_broken _rosenbrock_jac_reuse_decision was returning (true, true) for linear operator ODEs with a WOperator-wrapped W because the WOperator check fired before the linear-function check. That made Rosenbrock23 rebuild J and W every step on ODEFunction(::MatrixOperator) problems (seen as nw = 628 / 454 in lib/OrdinaryDiffEqNonlinearSolve linear_solver_tests expecting nw == 1), instead of building W once and reusing it. Reorder the decision so the linear-function branch fires immediately after the iter<=1 check, matching the pre-reuse do_newJW behavior: - iter <= 1 -> (true, true) [first step must build W] - islin -> (false, false) [reuse afterwards, regardless of W type] - non-adaptive / WOperator / mass-matrix / composite checks follow Also flip DelayDiffEq jacobian.jl:57 from @test_broken to @test — the nWfact_ts[] == njacs[] assertion now passes with the Rosenbrock J/W accounting from this PR. Verified locally: Rosenbrock23 on ODEFunction(MatrixOperator, mass_matrix=...): nw=1 Rodas5P+KrylovJL convergence test: L2 order 5.004 (tight reltol) Rosenbrock23 convergence on prob_ode_2Dlinear: order 1.996 Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * DelayDiffEq jacobian test: only assert Wfact_t==jac count for Rodas5P The `nWfact_ts[] == njacs[]` assertion holds for Rosenbrock methods (one Wfact_t / jac call per step) but not for TRBDF2 — SDIRK methods call Wfact_t per stage, so the SDIRK Wfact_t count is much larger than the jac count from the jac-based solve (observed 282 vs 3). The earlier 'Unexpected Pass' was for Rodas5P only, so flip @test_broken to @test just for that alg and keep TRBDF2 broken. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Expose max_jac_age and jac_reuse_gamma_tol as Rosenbrock kwargs Both thresholds used in _rosenbrock_jac_reuse_decision were hardcoded: - Gamma ratio tolerance (|dtgamma/last_dtgamma - 1|) at 0.3 - Max accepted-step age between J recomputations at 20 Expose them as constructor kwargs on all Rosenbrock algorithm structs (both macro-generated blocks, RosenbrockW6S4OS, and HybridExplicitImplicitRK): Rosenbrock23(; max_jac_age = 20, jac_reuse_gamma_tol = 0.3) Threading: - Fields added to each alg struct; constructors take matching kwargs - alg_cache passes alg.max_jac_age into JacReuseState at construction - Decision function reads alg.jac_reuse_gamma_tol (hasproperty guarded for robustness against any alg struct that skips the field) - JacReuseState(dtgamma, max_jac_age = 20) gets an optional second arg Set max_jac_age = 1 (or any small value) to effectively disable reuse for debugging / comparison. Non-W-methods accept the kwargs harmlessly since reuse is gated on isWmethod(alg) in the decision function. Verified: Rosenbrock23 on prob_ode_2Dlinear with jac_reuse_gamma_tol=0.01 doubles njacs (6 → 12) as expected while preserving order 1.996. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Restore rejected-step J reuse for strict Rosenbrock methods Regression in calc_rosenbrock_differentiation! for non-W-methods: the IIP path was passing newJW = _rosenbrock_jac_reuse_decision(...) to calc_W! for all algorithms, including strict Rosenbrock. For strict methods _rosenbrock_jac_reuse_decision returns (true, true) via the !isWmethod early return, which overrode calc_W!'s do_newJW errorfail branch that reuses J across step rejections (since retries land at the same (uprev, t) with only a smaller dt). Effect on master: Rodas5P on Brusselator 1D reported njacs = naccept (125), rejected steps did not add to njacs. Regression on the PR: njacs = naccept + nreject (139), one extra J per rejection. Fix: branch on isWmethod inside the IIP wrapper. W-methods use the reuse decision + newJW; strict methods call calc_W! without newJW, letting do_newJW handle the errorfail/reject case as on master. Verified on Brusselator 1D at reltol=1e-6: Rodas5P: njacs 139 → 125 (naccept=125, nreject=14), 23.7 → 23.1 ms Rodas4: njacs 212 → 194 (naccept=194, nreject=18), 34.8 → 33.7 ms Rosenbrock23 / ROS34PW2 / ROS34PW3 / Rodas23W unchanged. Reported by @gstein3m. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Revert DelayDiffEq jacobian test back to @test_broken for nWfact_ts==njacs My earlier 47ecc05 flipped @test_broken → @test based on a CI "Unexpected Pass" report, but that pass was a symptom of the strict-Rosenbrock jac-reuse regression I later fixed in 2db475d. On genuine master, njacs counts accepted steps (do_newJW's errorfail branch reuses J on retries) while nWfact_ts counts every step attempt (calc_W! calls Wfact_t unconditionally). The test's invariant nWfact_ts == njacs only holds when there are zero rejections, which isn't true for Rodas5P on this DDE (57 Wfact_t calls vs 54 jac calls = 3 rejected steps). Keep the assertion as @test_broken and document why — this is the correct master behavior, not a bug in the counting. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Lower jac_reuse_gamma_tol default from 0.3 to 0.03 A full benchmark sweep on ODEProblemLibrary stiff problems shows the original 0.3 gamma tolerance was tuned only for PDE-like dynamics (Brusselator) where dt barely changes. For chemistry and relaxation oscillators (ROBER, Van der Pol stiff), 30% dt changes are normal, so J goes stale fast and the solver pays for it with heavy step rejections. Benchmarking Rosenbrock23 across Brusselator1D, ROBER, VdP stiff, HIRES, Pollution (geometric mean wall-time ratio vs master): γ=0.30 +19.8% slower on average (old default) γ=0.20 +18.6% γ=0.15 +12.4% γ=0.10 +8.5% γ=0.05 +1.4% γ=0.03 +1.4% <-- new default — dominates γ=0.30 on every class γ=0.02 +1.6% γ=0.01 +9.9% γ=0.03 is strictly better than γ=0.30 on every problem class: Problem class γ=0.30 γ=0.03 Bruss1D -31.6% -34.7% (bigger PDE win) vdp_stiff +131.8% +43.4% (3x less cost) rober +17.4% +5.7% (3x less cost) hires +12.3% -0.5% (≈ master) pollution +17.9% +8.7% Intuition: at γ=0.3 we "reuse J until dt changes by 30%", which is way too loose — chemistry phase transitions and relaxation oscillator fast/slow switches move dt by 2–10× and the reuse keeps old J across the transition. At γ=0.03 ("reuse until dt changes by 3%"), PDE problems with near-constant dt still hit the reuse path (their dt typically changes by <1% per step), but chemistry problems catch their phase transitions early and recompute before rejections start cascading. Strict Rosenbrock methods are unaffected (verified via trace fingerprint comparison across 9 strict algorithms × 3 problems × 4 tolerances — byte-identical t-sequences and u endpoints, 0 diff lines vs master). Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Rosenbrock23/32: default max_jac_age=1 (no reuse), opt-in via kwarg Rosenbrock23 and Rosenbrock32 are low-order Rosenbrock-W methods typically used on small problems, as a fallback in auto-switching algorithms, or for stiff-detection — workloads where Jacobian evaluation is cheap and per-step overhead dominates. On these problems the general W-method reuse defaults (max_jac_age=20, jac_reuse_gamma_tol=0.03) are a net loss: on Van der Pol stiff the reuse causes step rejections that cost 2.3× master wall time even at optimal γ. Change: Rosenbrock23 and Rosenbrock32 default `max_jac_age = 1`. Combined with a new cache-side optimization that skips JacReuseState allocation when max_jac_age ≤ 1, the decision function short-circuits through do_newJW (master-compatible path) and the entire reuse code path is bypassed. Higher-order W-methods (Rodas23W, Rodas4P2, Rodas5P/Pe/Pr, Rodas6P, RosenbrockW6S4OS, ROS34PW1a/b/2/3, ROS34PRw, ROS3PRL/2, ROK4a) keep the default `max_jac_age = 20`, so the PDE reuse wins are preserved. Supporting changes: - _make_jac_reuse_state(dtgamma, max_jac_age) helper returns `nothing` when age ≤ 1, so Rosenbrock23/32 caches allocate no jac_reuse state. - calc_rosenbrock_differentiation! now dispatches on `isWmethod(alg) && jac_reuse !== nothing` so W-methods with reuse disabled take the same master-compatible path as strict Rosenbrock. - get_jac_reuse uses hasfield(typeof(cache), :jac_reuse) instead of hasproperty — compile-time constant-foldable, removes runtime reflection cost from the hot path on caches that opt out. - gamma_tol fallback in the decision function also uses hasfield. Benchmarks (wall time vs master, Rosenbrock23, best-of-3): Problem/tol master PR Δ% Bruss1D/1e-6 71.99 ms 55.61 ms -22.8% Bruss1D/1e-8 258.79 ms 251.76 ms -2.7% rober/1e-6 0.28 ms 0.29 ms +3.6% rober/1e-8 1.02 ms 1.03 ms +1.0% vdp_stiff/1e-6 3.16 ms 3.30 ms +4.4% vdp_stiff/1e-8 19.73 ms 21.01 ms +6.5% hires/1e-6 0.71 ms 0.73 ms +2.8% hires/1e-8 3.26 ms 3.33 ms +2.1% pollution/1e-6 0.53 ms 0.52 ms -1.9% pollution/1e-8 2.08 ms 2.02 ms -2.9% ------------------------------------ Geomean -1.3% All 10 configurations produce byte-identical (njacs, naccept, nreject) sequences to master for Rosenbrock23. The remaining wall- time deltas are measurement noise on sub-millisecond solves. Previous Rosenbrock23 defaults on the same problem set: γ=0.30 (original): +19.8% slower (Van der Pol +131.8%) γ=0.03 (last tune): +1.4% slower (Van der Pol +43.4%) max_jac_age=1 (this commit): -1.3% (Van der Pol +5.4%) Strict Rosenbrock methods are unaffected — 144-config trace fingerprint test still byte-identical to master. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 8d5ed6a commit 689d0b7

9 files changed

Lines changed: 384 additions & 43 deletions

File tree

lib/DelayDiffEq/test/interface/jacobian.jl

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,11 @@ using Test
5454

5555
# check number of function evaluations
5656
@test !iszero(nWfact_ts[])
57+
# Wfact_t is called on every step attempt (accepted + rejected), while
58+
# the user-supplied `jac` is only called on accepted steps because
59+
# do_newJW's errorfail branch reuses J across rejected retries. So
60+
# nWfact_ts[] == njacs[] only when there are no rejections, which
61+
# isn't true for Rodas5P / TRBDF2 on this DDE.
5762
@test_broken nWfact_ts[] == njacs[]
5863
@test iszero(sol_Wfact_t.stats.njacs)
5964
@test_broken nWfact_ts[] == sol_Wfact_t.stats.nw

lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ using OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqAdaptiveImplici
4444
isnewton, _unwrap_val,
4545
set_new_W!, set_W_γdt!, alg_difftype, unwrap_cache, diffdir,
4646
get_W, isfirstcall, isfirststage, isJcurrent,
47-
get_new_W_γdt_cutoff,
47+
get_new_W_γdt_cutoff, isWmethod,
4848
TryAgain, DIRK, COEFFICIENT_MULTISTEP, NORDSIECK_MULTISTEP, GLM,
4949
FastConvergence, Convergence, SlowConvergence,
5050
VerySlowConvergence, Divergence, NLStatus, MethodType, constvalue, @SciMLMessage

lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl

Lines changed: 258 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,174 @@
11
using SciMLOperators: StaticWOperator, WOperator
22

3+
"""
4+
get_jac_reuse(cache)
5+
6+
Duck-typed accessor for the `jac_reuse` field. Returns `nothing` if the cache
7+
does not have a `jac_reuse` field.
8+
"""
9+
@inline function get_jac_reuse(cache)
10+
# hasfield on typeof(cache) is compile-time constant-foldable, unlike
11+
# hasproperty which walks propertynames at runtime. This makes the
12+
# non-reuse fast path free on caches without the field.
13+
if hasfield(typeof(cache), :jac_reuse)
14+
return cache.jac_reuse
15+
else
16+
return nothing
17+
end
18+
end
19+
20+
# Strip ForwardDiff.Dual to plain value for heuristic storage in JacReuseState.
21+
# JacReuseState fields are Float64 (or similar) and don't need to carry AD derivatives.
22+
_jac_reuse_value(x) = ForwardDiff.value(x)
23+
24+
"""
25+
_rosenbrock_jac_reuse_decision(integrator, cache, dtgamma) -> NTuple{2,Bool}
26+
27+
Decide whether to recompute the Jacobian and/or W matrix for Rosenbrock methods.
28+
All Rosenbrock/W-method J/W logic lives here — no delegation to `do_newJW`.
29+
30+
For W-methods on adaptive, non-DAE, non-composite, non-linear ODE problems,
31+
implements CVODE-inspired Jacobian reuse:
32+
- Always recompute on first iteration
33+
- Recompute after step rejection (EEst > 1), callback, resize, algorithm switch
34+
- Recompute when gamma ratio changes too much:
35+
`|dtgamma/last_dtgamma - 1| > alg.jac_reuse_gamma_tol` (default 0.03)
36+
- Recompute every `alg.max_jac_age` accepted steps (default 20)
37+
- Otherwise reuse J but rebuild W from the cached J and current dtgamma.
38+
The Jacobian evaluation (finite-diff / AD) is the expensive part; W
39+
construction and LU factorization are comparatively cheap.
40+
41+
Both thresholds are tunable via Rosenbrock algorithm constructor kwargs
42+
`max_jac_age` and `jac_reuse_gamma_tol`, e.g.
43+
`Rosenbrock23(max_jac_age = 50, jac_reuse_gamma_tol = 0.1)`.
44+
Set `max_jac_age = 1` to effectively disable reuse (recompute every step).
45+
46+
Returns `(true, true)` (fresh J and W) for all non-reusable cases:
47+
strict Rosenbrock methods, linear problems, mass-matrix (DAE) problems,
48+
CompositeAlgorithm, non-adaptive solves, and first iteration.
49+
"""
50+
function _rosenbrock_jac_reuse_decision(integrator, cache, dtgamma)
51+
alg = OrdinaryDiffEqCore.unwrap_alg(integrator, true)
52+
53+
# Non-W-methods always recompute
54+
if !isWmethod(alg)
55+
return (true, true)
56+
end
57+
58+
jac_reuse = get_jac_reuse(cache)
59+
# No reuse state available — always recompute
60+
if jac_reuse === nothing
61+
return (true, true)
62+
end
63+
64+
# First iteration: always compute J and W.
65+
if integrator.iter <= 1
66+
return (true, true)
67+
end
68+
69+
# Linear problems: J is the constant linear operator and W wrapping it
70+
# doesn't need rebuilding after the first step — SciMLOperators update
71+
# W's internal gamma in place. This must come before the non-adaptive,
72+
# WOperator, mass-matrix, and composite checks so linear operator ODEs
73+
# (e.g. ODEFunction(::MatrixOperator), including with a mass matrix) get
74+
# the fast path — matches the pre-reuse `do_newJW` behavior.
75+
islin, _ = islinearfunction(integrator)
76+
if islin
77+
return (false, false)
78+
end
79+
80+
# Non-adaptive solves always recompute.
81+
# J reuse provides negligible benefit with prescribed timesteps and causes
82+
# IIP/OOP inconsistency (adaptive solves have step rejections that reset
83+
# reuse state, while non-adaptive solves following the same timesteps don't).
84+
if !integrator.opts.adaptive
85+
return (true, true)
86+
end
87+
88+
# Operator-based W (WOperator for Krylov solvers, AbstractSciMLOperator):
89+
# always recompute. Krylov convergence depends on W quality, and stale J
90+
# in the operator causes convergence degradation.
91+
if cache.W isa WOperator || cache.W isa AbstractSciMLOperator
92+
return (true, true)
93+
end
94+
95+
# Mass matrix (DAE) problems always recompute.
96+
# Stale Jacobians cause order reduction for DAEs because the algebraic
97+
# constraint derivatives must remain accurate. See Steinebach (2024).
98+
if integrator.f.mass_matrix !== I
99+
return (true, true)
100+
end
101+
102+
# CompositeAlgorithm always recomputes.
103+
# Rapid stiff↔nonstiff transitions make reuse counterproductive.
104+
if integrator.alg isa CompositeAlgorithm
105+
return (true, true)
106+
end
107+
108+
# Commit pending_dtgamma from previous step if it was accepted.
109+
# This ensures rejected steps don't pollute last_dtgamma.
110+
naccept = integrator.stats.naccept
111+
if naccept > jac_reuse.last_naccept
112+
jac_reuse.last_dtgamma = jac_reuse.pending_dtgamma
113+
jac_reuse.last_naccept = naccept
114+
end
115+
116+
# Fresh cache (e.g., algorithm switch where iter > 1 but the Rosenbrock
117+
# cache is freshly created with cached_J = nothing).
118+
if iszero(jac_reuse.last_dtgamma)
119+
return (true, true)
120+
end
121+
122+
# Detect algorithm switch in CompositeAlgorithm: if integrator.iter jumped
123+
# by more than 1 since our last Rosenbrock step, another algorithm ran in
124+
# between and the cached Jacobian is evaluated at a stale u.
125+
if jac_reuse.last_step_iter != 0 && integrator.iter > jac_reuse.last_step_iter + 1
126+
return (true, true)
127+
end
128+
129+
# Callback modification: recompute
130+
if integrator.u_modified
131+
return (true, true)
132+
end
133+
134+
# Resize detection: if u changed length since last J computation,
135+
# the cached LU factorization has wrong dimensions.
136+
# (u_modified is already cleared by reeval_internals_due_to_modification!
137+
# before perform_step! runs, so we need this explicit check.)
138+
if length(integrator.u) != jac_reuse.last_u_length && jac_reuse.last_u_length != 0
139+
return (true, true)
140+
end
141+
142+
# Previous step was rejected (EEst > 1): the old W wasn't good enough.
143+
# Recompute everything since we're retrying with a different dt anyway.
144+
if integrator.EEst > 1
145+
return (true, true)
146+
end
147+
148+
# Gamma ratio check (uses only accepted-step dtgamma).
149+
# Tunable via `alg.jac_reuse_gamma_tol` (default 0.03).
150+
last_dtg = jac_reuse.last_dtgamma
151+
gamma_tol = if hasfield(typeof(alg), :jac_reuse_gamma_tol)
152+
alg.jac_reuse_gamma_tol
153+
else
154+
0.03
155+
end
156+
if !iszero(last_dtg) && abs(dtgamma / last_dtg - 1) > gamma_tol
157+
return (true, true)
158+
end
159+
160+
# Age check: recompute J after max_jac_age accepted steps.
161+
if (naccept - jac_reuse.last_naccept) >= jac_reuse.max_jac_age
162+
return (true, true)
163+
end
164+
165+
# Reuse J but rebuild W with the current dtgamma. Following CVODE's
166+
# approach: the Jacobian evaluation (finite-diff / AD) is expensive,
167+
# while reconstructing W = J − M/(dt·γ) and its LU is comparatively
168+
# cheap and keeps the step controller accurate.
169+
return (false, true)
170+
end
171+
3172
function calc_tderivative!(integrator, cache, dtd1, repeat_step)
4173
return @inbounds begin
5174
(; t, dt, uprev, u, f, p) = integrator
@@ -686,18 +855,104 @@ end
686855

687856
function calc_rosenbrock_differentiation!(integrator, cache, dtd1, dtgamma, repeat_step)
688857
nlsolver = nothing
858+
alg = OrdinaryDiffEqCore.unwrap_alg(integrator, true)
689859
# we need to skip calculating `J` and `W` when a step is repeated
690860
new_jac = new_W = false
691861
if !repeat_step
692-
new_jac, new_W = calc_W!(
693-
cache.W, integrator, nlsolver, cache, dtgamma, repeat_step
694-
)
862+
jac_reuse = get_jac_reuse(cache)
863+
if isWmethod(alg) && jac_reuse !== nothing
864+
# W-methods with reuse enabled: use CVODE-inspired reuse logic
865+
newJW = _rosenbrock_jac_reuse_decision(integrator, cache, dtgamma)
866+
new_jac, new_W = calc_W!(
867+
cache.W, integrator, nlsolver, cache, dtgamma, repeat_step, newJW
868+
)
869+
jac_reuse.last_step_iter = integrator.iter
870+
if new_jac
871+
jac_reuse.pending_dtgamma = _jac_reuse_value(dtgamma)
872+
jac_reuse.last_u_length = length(integrator.u)
873+
end
874+
else
875+
# Strict Rosenbrock, or W-method with reuse disabled (jac_reuse ===
876+
# nothing). Defer to do_newJW inside calc_W! so that the errorfail
877+
# branch reuses J across step rejections (matches master).
878+
new_jac, new_W = calc_W!(
879+
cache.W, integrator, nlsolver, cache, dtgamma, repeat_step
880+
)
881+
end
695882
end
696883
# If the Jacobian is not updated, we won't have to update ∂/∂t either.
697884
calc_tderivative!(integrator, cache, dtd1, repeat_step || !new_jac)
698885
return new_W
699886
end
700887

888+
"""
889+
calc_rosenbrock_differentiation(integrator, cache, dtgamma, repeat_step)
890+
891+
Non-mutating (OOP) version of `calc_rosenbrock_differentiation!`.
892+
Returns `(dT, W)` where `dT` is the time derivative and `W` is the factorized
893+
system matrix. Supports Jacobian reuse for W-methods via `jac_reuse` in the cache.
894+
"""
895+
function calc_rosenbrock_differentiation(integrator, cache, dtgamma, repeat_step)
896+
jac_reuse = get_jac_reuse(cache)
897+
898+
# If no reuse support or repeat step, use standard path
899+
if repeat_step || jac_reuse === nothing
900+
dT = calc_tderivative(integrator, cache)
901+
W = calc_W(integrator, cache, dtgamma, repeat_step)
902+
return dT, W
903+
end
904+
905+
new_jac, new_W = _rosenbrock_jac_reuse_decision(integrator, cache, dtgamma)
906+
907+
# Track iteration for algorithm-switch detection in CompositeAlgorithm
908+
jac_reuse.last_step_iter = integrator.iter
909+
910+
if new_jac
911+
# Fresh Jacobian needed — use standard calc_tderivative + calc_W
912+
# for numerical consistency with the IIP path (calc_W handles
913+
# StaticWOperator, WOperator, AbstractSciMLOperator, etc.).
914+
dT = calc_tderivative(integrator, cache)
915+
W = calc_W(integrator, cache, dtgamma, repeat_step)
916+
917+
# Cache J for future reuse (calc_W computes J internally but
918+
# doesn't expose it, so we recompute — cheap relative to W).
919+
jac_reuse.cached_J = calc_J(integrator, cache)
920+
jac_reuse.cached_dT = dT
921+
jac_reuse.cached_W = W
922+
jac_reuse.pending_dtgamma = _jac_reuse_value(dtgamma)
923+
jac_reuse.last_u_length = length(integrator.u)
924+
925+
return dT, W
926+
end
927+
928+
# Reusing cached J — build W from it directly.
929+
# Safety: if cached_J is nothing (e.g. first use after algorithm switch),
930+
# fall back to standard path.
931+
if jac_reuse.cached_J === nothing
932+
dT = calc_tderivative(integrator, cache)
933+
W = calc_W(integrator, cache, dtgamma, repeat_step)
934+
return dT, W
935+
end
936+
937+
J = jac_reuse.cached_J
938+
dT = jac_reuse.cached_dT
939+
940+
mass_matrix = integrator.f.mass_matrix
941+
update_coefficients!(mass_matrix, integrator.uprev, integrator.p, integrator.t)
942+
943+
# Rebuild W from cached J and current dtgamma.
944+
# The Jacobian evaluation is the expensive part; W = J − M/(dt·γ) and
945+
# its factorization are comparatively cheap and keep step control accurate.
946+
W = J - mass_matrix * inv(dtgamma)
947+
if !isa(W, Number)
948+
W = DiffEqBase.default_factorize(W)
949+
end
950+
integrator.stats.nw += 1
951+
jac_reuse.cached_W = W
952+
953+
return dT, W
954+
end
955+
701956
# update W matrix (only used in Newton method)
702957
function update_W!(integrator, cache, dtgamma, repeat_step, newJW = nothing)
703958
return update_W!(cache.nlsolver, integrator, cache, dtgamma, repeat_step, newJW)

lib/OrdinaryDiffEqRosenbrock/src/OrdinaryDiffEqRosenbrock.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ using OrdinaryDiffEqDifferentiation: TimeDerivativeWrapper, TimeGradientWrapper,
3535
calc_W, calc_rosenbrock_differentiation!, build_J_W,
3636
UJacobianWrapper, dolinsolve, WOperator, resize_J_W!
3737

38+
using OrdinaryDiffEqDifferentiation: calc_rosenbrock_differentiation
39+
3840
using Reexport
3941
@reexport using SciMLBase
4042

0 commit comments

Comments
 (0)