Add diffrax/lineax-backed adaptive methods to op_engine's CoreSolver, ported from COVID19_USA's diffrax_engine.py

[jc-macdonald]

Summary

Add diffrax/lineax-backed adaptive solve methods to op_engine's existing CoreSolver, extracted from and informed directly by a working reference implementation: ACCIDDA/COVID19_USA's model_input/plugins/diffrax_engine.py (~2,200 lines) and model_input/plugins/_history.py (~400 lines).

Design note (2026-08-18): this issue originally proposed a structurally separate "JAX engine" package/class alongside CoreSolver. That framing is revised: the target is one CoreSolver, with these capabilities added as new selectable method= values dispatched internally when the input state is a JAX array — mirroring how op_system's compiled eval_fn infers its array namespace at call time rather than exposing a separate compiled artifact per backend. See the comment above for full rationale. The technical extraction inventory below is unchanged; only the "separate engine" framing is corrected.

Why this supersedes the old GPU-backend plan (#2) and the Protocol-injection plan (#19/#21/#22/#23/#25)

Those issues assumed a JAX backend would be a drop-in LinearSolverBackend implementation slotted into the existing SciPy-shaped CoreSolver (same algorithm, different array module). That assumption doesn't hold for the adaptive/differentiable-stepping piece specifically. COVID19_USA already needed JAX-traceable gradients for NUTS/SVI/SMC inference over its op_system-defined ODE RHS, and — finding no engine in op_engine that could do this — built its own engine from scratch. It is not a SciPy-shaped implicit solver ported to JAX; it's a different solve strategy entirely (adaptive explicit RK via diffrax.Tsit5, additive operator splitting for advection/jump terms via a generic finite-volume upwind scheme, diffrax.PIDController/RecursiveCheckpointAdjoint, block-axis spatial vmap, posterior-draw-batch vmap, plus a from-scratch DDE/history ring-buffer convolution system for op_system's convolve_history(...) operator).

Note this is narrower than originally stated: basic implicit/IMEX solving on JAX is not part of this gap — that's covered by #69's Array-API dense fallback (xp.linalg.solve, correct on any conformant backend including JAX, with zero JAX-specific code). What is genuinely JAX-divergent and needs real extraction work is: adaptive/differentiable stepping, operator splitting, vmap batching, and the history/DDE convolution system. Those don't have a SciPy analog to fall back to.

Confirmed via direct code reading that diffrax_engine.py is generic, not COVID-specific — nothing in it references compartment names or COVID domain concepts; it's driven entirely by metadata read off the compiled op_system spec (stepper.option("pytree_stepper_fn", ...), "block_pytree_stepper_fn", "history_stepper_fn", "body_eval_fn", "block_history_stepper_fn", "block_body_eval_fn", "operators"). It is COVID19_USA's private, unofficial, unshared implementation of an additional solve strategy for the same role op_engine.CoreSolver plays for the SciPy case — currently unmaintained outside that one repo, with no shared test coverage, and no protection against drifting out of sync as op_system's compiled-spec contract evolves.

Scope

Extract the generic (non-COVID-specific) machinery into op_engine's CoreSolver as new methods, gated behind an optional jax extra (mirroring op_system's [project.optional-dependencies] jax). Rough inventory from the reference implementation, by conceptual piece (COVID19_USA file:line references for whoever picks this up):

• Solve orchestration: override/cache-key handling (engine_overrides, diffrax_engine.py:202-225), diffrax.ODETerm/MultiTerm assembly, PIDController/ConstantStepSize step-control selection, RecursiveCheckpointAdjoint/DirectAdjoint adjoint selection, diffeqsolve invocation, runner/runner_vmap entry points (diffrax_engine.py:1056-2197). Exposed as new RunConfig-selectable methods (naming TBD, e.g. method="tsit5") rather than a separate solve() entry point.
• Operator-splitting builders: generic upwind finite-volume advection/jump-integral operators over an arbitrary named axis (_build_neighbor_index, _make_advection_term(_pytree), _build_kernel_weights(_pytree), _make_jump_term(_pytree), _build_operator_terms(_pytree), diffrax_engine.py:404-1056) — conceptually overlaps with matrix_ops.py's Laplacian/Crank-Nicolson builders but is a different (explicit, JAX-traced) numerical strategy, not a shared code path.
• Block-axis (spatial) vmap: activated when a spec declares factorize_axes, vmaps a single draw's solve over a spatial axis (diffrax_engine.py:1754-1900).
• Draw-batch vmap: runner_vmap vmaps both y0 and dynamic (traced) params over a posterior-draw axis, with static-arg-signature consistency checks across the batch (diffrax_engine.py:2097-2197).
• DDE/history convolution: fixed-cadence outer ring-buffer provider satisfying op_system's HistoryProvider protocol (query/commit), supporting gamma and protection_response kernels, explicitly incompatible with BacksolveAdjoint (must use RecursiveCheckpointAdjoint) — this is _history.py in full, ~400 lines, self-contained.

Explicitly out of scope for this issue (follow-ups once these methods exist and are proven):

• Migrating COVID19_USA itself onto the new op_engine methods. COVID19_USA's flagship inference config is under active iteration; that migration is real, separate, carefully-sequenced work and shouldn't block or be blocked by landing this in op_engine.
• Solver choice beyond Tsit5 (the reference implementation hardcodes it; op_engine's existing CoreSolver config already supports named-method selection and that pattern extends naturally here — pick additional solver choices up as a fast-follow).
• Sparse-accelerated implicit/stiff solving on JAX at scale via lineax (dense IMEX on JAX is covered by Multi-backend implicit solve: Array-API dense fallback + SciPy/CuPy sparse registry #69; a registered sparse-JAX path, if ever needed, is separate follow-on work, not blocking this issue).

Acceptance criteria

• New optional jax extra on op_engine (and propagated to flepimop2-op_engine, see flepimop2-op_engine: options-surface parity with flepimop2-op_system + jax extra #68) with diffrax/lineax/jax as dependencies.
• New CoreSolver-selectable method(s) (naming TBD during implementation) that: solve an op_system-compiled spec via the same pytree_stepper_fn/block_pytree_stepper_fn/history_stepper_fn/body_eval_fn contract COVID19_USA's engine already consumes; support draw-batch vmap; support block-axis vmap; support convolve_history specs via the ported history provider — all invoked through the same public CoreSolver.run()/RunConfig surface as the existing explicit/IMEX methods, not a separate class or module entry point.
• Test suite ported/adapted from COVID19_USA's tests/test_diffrax_engine.py and tests/test_history_provider.py, run against op_engine's copy.
• just quality and just test pass.
• A short design note (in docs/) on why these are additional methods on CoreSolver rather than a separate engine class, and why sparse/dense dispatch (Multi-backend implicit solve: Array-API dense fallback + SciPy/CuPy sparse registry #69) and adaptive/differentiable stepping (this issue) are two different problems with two different solutions — so the reasoning doesn't get re-litigated later.

Relationship to other issues

Supersedes #2 ("Add GPU compatible version of matrix ops") — closed as folded into this issue, since the real shape of the work turned out to be adaptive-ODE-plus-operator-splitting methods, not GPU-friendly linear-algebra variants of matrix_ops functions. Depends on #69 for the "dense IMEX already works on JAX" baseline (not blocking — can be scoped/designed in parallel, but the full non-goal-loosening claim needs #69 landed). Can start independent of the provider-parity issue (#68), though both will need to land before COVID19_USA migration is realistic. See #26 for the overall tracking issue.

[jc-macdonald]

Design revision (2026-08-18)

Revising this issue's framing after further review. The original title/body proposed a structurally separate JAX/diffrax engine alongside the existing SciPy CoreSolver. That's being walked back in favor of:

• One CoreSolver, one method surface, with diffrax/lineax-backed adaptive methods added as new selectable method= values on the existing class — not a second engine/package with a different public API.
• The "no true IMEX support in the JAX engine" non-goal is loosened: dense IMEX already works on JAX via the Array-API dense fallback being added in Multi-backend implicit solve: Array-API dense fallback + SciPy/CuPy sparse registry #69 (xp.linalg.solve, correct on any conformant backend, zero JAX-specific code). What Multi-backend implicit solve: Array-API dense fallback + SciPy/CuPy sparse registry #69 does not cover is sparse-accelerated IMEX at JAX scale — that's still real follow-on work via lineax, not ruled out architecturally the way the original text implied.
• Rationale for the correction: op_system's and op_engine's whole premise is genuine multi-backend generality — "works with any package meeting the Array API standard" — and the Array API standard's linalg.solve extension is already implemented uniformly across NumPy/CuPy/JAX/PyTorch via array-api-compat. Treating JAX as needing a wholesale separate engine overstated the actual gap, which is specifically about sparse factorization at scale (genuinely JAX-divergent) and differentiable adaptive stepping (the real reason to want diffrax/lineax methods), not about basic implicit solving.

Body below is being updated to match. See #69 for the sparse/dense dispatch layer this depends on, and #26 for the overall revised plan.