[TF24] Prototype of multirate stepper (DO NOT MERGE)#42
Draft
aornugent wants to merge 12 commits into
Draft
Conversation
A model-agnostic multirate-infinitesimal macro step: an explicit Runge-Kutta on the slow block, sub-cycling the fast block over each leg with the existing adaptive solver. A System becomes multirate by declaring its fast/slow split (fast_size, slow_rates, fast_rates, aggregate); such a System still satisfies the plain ODE interface, so the ordinary solver gives a single-rate reference. The fast block is driven through a FastLeg adapter that presents it as a plain System reading the leg's slow-coupling polynomial, so no adaptive RK is reimplemented. Coupling tables: forward-Euler, Heun, Kutta3 (lifted), and MRI-GARK-ERK33a. TwoRateSystem example + gates: outer ERK order, resolved coupling order, consistency under stiffness, and fast-cost independence from the slow-block size. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
A double adaptive pass records the fast sub-cycle schedule (accepted step sizes per leg); a second pass at an active scalar replays it with fixed steps, so the tape is the exact discrete adjoint of the scheme as run. The macro step is now templated on the scalar and drives the fast leg through a single record-or-replay driver (if constexpr guards the double-only adaptive branch), so forward and reverse share one code path. TwoRateSystem gains seedable inputs (the relaxation rate and the initial state). Gate: the adjoint gradient matches a frozen-schedule central finite difference to ~1e-9 across a range of stiffness, and is independent of the finite-difference step. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
The fast block presented as a plain System is the fast subsystem; the per-stage windows it is sub-cycled over are sub-intervals, and the recorded schedule holds one entry per sub-cycle. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
The inner sub-cycle is now a policy: the adaptive black-box RK (default) or an operator split that integrates the model's stiff part exactly (analytic_flow) and steps the gentle remainder with ROS34PW2 (Rang & Angermann), composed Strang-symmetrically. The split is opt-in and only compiles for models that provide the analytic_flow / residual_rhs hooks, so the same model can be run both ways to measure what exposing an analytic partial flow buys. DrainageSystem (a stiff power-law drainage column with a closed-form recession, shaped like TF24 soil) is the A/B subject. At the accuracy the macro step delivers, splitting is ~6x fewer fast steps and the win is flat as the drainage stiffness grows 1000x; the exact recession preserves positivity. ROS34PW2 verified order 3 and L-stable independently. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
ros34pw2_step is templated on the scalar with a passive value-Jacobian and a double-matrix / active-rhs dense solve, so a stage is an active back-substitution with recorded factors -- the discrete adjoint of a constant-matrix solve. subcycle_split templates on the scalar (record double / replay active); its residual is a generic lambda so the coupling is active in the stages but passive in the value-Jacobian. DrainageSystem exposes the drainage rate and initial state as seedable inputs. Gate: the gradient of a functional w.r.t. [c, initial state] through Strang(exact flow, ROS34PW2, exact flow) matches a frozen-schedule finite difference to ~1e-9 across stiffness (exact because the residual is linear in u). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
MriStep matches the Step interface, so SolverInternal dispatches to it exactly as it does for RODAS. One macro step advances the whole state by step_size: an explicit RK on the slow block, the fast block sub-cycled by the existing solver over each leg (mri.hpp). A System opts in by declaring its fast/slow split (is_multirate: fast_size + the partitioned rates), state laid out fast-first; MriStep::supported gates it (double scalar only -- gradients go through the record->replay path). Intended for a fixed macro grid (the forcing-kink grid); adaptivity lives in the fast sub-cycle. The include cycle (mri.hpp sub-cycles through SolverInternal, which now holds an MriStep member) is resolved by defining MriStep::step out of line in ode_step_mri_impl.hpp, included at the bottom of mri.hpp once both are available. Gate: method='mri' through the Solver surface is stable and accurate on a fixed daily grid where a single explicit step of that size blows up -- the multirate value, now on the same 'method' surface the SCM selects. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
A multirate System whose fast rates read a frozen slow context (rather than a linear aggregate g) implements freeze_slow(x); mri_macro_step calls it once per leg before the fast sub-cycle. Detected via if constexpr, a no-op otherwise -- the toys (linear aggregate) are unaffected. This is what the TF24 patch adapter needs: capture the cohorts + light field per leg, then vary only the soil in the fast sub-cycle. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
The macro stepper and MriStep now split state at slow_size (slow first, fast block last), matching Patch's native [cohorts | environment] layout so the SCM's Solver<Patch> needs no re-ordering wrapper. The two demonstrators are flipped to [slow x | fast u] to match; the reverse-mode functional sums the slow block [0, slow_size). This is the layout commitment behind the SCM integration (system-design: candidate C -- one convention, no wrapper, no engine offset). Full suite 275 pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
The MRI sub-cycle and macro step took the System by const reference, which suits the toy models whose partition hooks are pure. A real System such as plant's Patch computes its fast/slow rates by mutating internal state (set soil theta, recompute the environment, rebuild the frozen light field), so its hooks cannot be const. Take the System by non-const reference through FastSubsystem, the sub-cycle strategies and mri_macro_step; the toys' const hooks still bind to a non-const ref, so nothing else changes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
MriStep hardcoded the adaptive black-box inner. Add an mri_wants_split trait (detects a runtime bool mri_split() on the System) and branch: a System that opts in uses the exact-flow + ROS34PW2 split inner (Lever 1), else the adaptive inner. Both instantiate only when mri_split() is present (which implies the analytic_flow / residual_rhs hooks). Systems without it (the two-rate toy) are unaffected -- tests pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
Reuses RodasStep via a Jacobian-policy template param (default unchanged, so RODAS stays bit-identical). BlockFdJacobian finite-differences only the fast (soil) block through the full coupled RHS, capturing d(uptake)/d(theta) where the stiffness lives. Needs no rebind<U>() hook, so it runs on the real coupled Patch which RODAS cannot. Exposed as method='imex', gated on the fast/slow partition trait; toy two-rate demo confirms it is stable at a stiff fixed grid where explicit diverges. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BvyqPT8pVL7bCZvWqmhZHM
aornugent
marked this pull request as draft
July 22, 2026 04:45
aornugent
force-pushed
the
claude/tf24-multirate-engine
branch
from
July 22, 2026 04:48
2f78191 to
b88514d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hi!
I have been tinkering with different solver configurations for TF24 including multirate integration (MRI) and implicit/explicit (IMEX) split operators.
This prototype has a bunch of baggage and is only really a quick spike to see if there were any quick wins. I don't think it constitutes a production quality design, but I'm sharing it in case there are parts you'd like to cherry pick.
Maybe addresses some of: traitecoevo/plant#566
This PR is paired with traitecoevo/plant#569