IMEX solver: per-block stiff flag with ARK-type additive Runge-Kutta

[milanofthe]

Motivation

Real-world block-diagram systems often have heterogeneous stiffness: a few blocks are locally stiff (fast actuator dynamics, electrical filters, sensor transients) while the bulk is non-stiff. Running a fully implicit solver (ESDIRK43, DIRK3) on the whole system pays Newton cost everywhere. Running an explicit solver requires tiny dt_max to stay stable.

IMEX-ARK schemes let us split the RHS additively: f = f_E + f_I. Non-stiff blocks contribute to f_E, stiff blocks to f_I. Per stage, one explicit pass + one implicit Newton solve over the stiff cluster.

We already have the pieces:

• Per-block AD-Jacobians via JIT (compile_jacobian in src/pybindings/py/helpers.rs)
• DIRK/ESDIRK solvers (src/solvers/)
• Preconditioned Anderson for the implicit stage solve (src/optim/anderson.rs)
• DAG analysis with loop detection (src/utils/graph.rs)

Missing: the IMEX tableau dispatch, per-block stiff flag, and cluster analysis.

Proposed architecture

User-facing API

# Block-level flag (default granularity)
ode = ODE(func, initial_value=x0, stiff=True)

# Also on DynamicalSystem, MassMatrixDAE, SemiExplicitDAE, FullyImplicitDAE

# Subsystem-level flag (flags all internal dynamical blocks)
sub = Subsystem(blocks=[...], connections=[...], stiff=True)

Dispatch rule

• No blocks flagged stiff → existing behavior (user's chosen global solver)
• User chose explicit solver AND ≥1 block flagged stiff → IMEX kicks in (ARK default)
• User chose implicit solver AND blocks flagged → implicit globally; flag is a no-op warning

Stiff cluster analysis

"Stiff cluster" = connected component in the subgraph {stiff blocks} ∪ {connections with both endpoints stiff}. Isolated stiff blocks are 1-cluster, handled exactly like today's per-block ImplicitSolver path.

Each cluster gets one Newton solve per stage with:

• Block-diagonal Jacobian from per-block JIT AD
• Off-diagonal coupling terms from stiff↔stiff connections (trivial for linear connections; FD fallback for nonlinear)

Cross-partition connections (stiff↔non-stiff): handled by existing Simulation FPI between stages. No DAG restructuring needed — we keep the existing topo-sort evaluation order and just split the RHS computation into explicit-pass + implicit-pass per ARK stage.

Tableau choice

Based on SOTA survey (SUNDIALS/ARKode, PETSc/TS, recent literature):

1. Default: ARK3(2)4L[2]SA (Kennedy-Carpenter 2003) — 4 stages, order 3, embedded order 2. L-stable ESDIRK on implicit side, stiffly-accurate, stage-order 2. This is ARKode's default 3rd-order IMEX. Cheap (4 block-graph evaluations per step).

2. Tight tolerances: ARK4(3)6L[2]SAb (Kennedy-Carpenter 2019 — note: the "b" variant from the 2019 paper, not the 2003 original; better nonlinear stability with same L[2]SA properties). 6 stages, order 4, embedded order 3.

Both have embedded error estimators for adaptive timestepping (ARS-style schemes lack these — skip ARS443 despite its ubiquity in older literature).

Skip for v1: ARK5(4)8 (8 stages is expensive), ARS schemes (no embedding), Rosenbrock variants.

Implementation plan

1. ImexTableau in src/solvers/tableaus.rs — struct holding explicit + implicit Butcher coefficients with shared c-vector and embedded weights. Encode ARK3(2)4 and ARK4(3)6b coefficients.
2. stiff: bool kwarg propagated through ODE, DynamicalSystem, MassMatrixDAE, SemiExplicitDAE, FullyImplicitDAE, and Subsystem. Sets a Block.stiff flag in core.
3. Cluster analysis in src/utils/graph.rs: stiff_clusters(blocks, connections) -> Vec<Vec<BlockId>> via connected-components on the stiff subgraph.
4. ImexSolver in src/solvers/: per-stage dispatch (explicit pass → per-cluster Newton). Reuses existing DIRK stage-solve + Anderson acceleration.
5. Simulation step (src/simulation.rs around line 800): dispatch Explicit/Implicit/IMEX based on block partition. Keep existing outer FPI for cross-partition algebraic loops.

Open design questions

• Subsystem boundary: should a cluster search cross Subsystem boundaries? Instinct: no (Subsystems are semantic encapsulation units), but needs validation on real models.
• Tableau default when user has not chosen a solver: ARK3(2)4L[2]SA as sensible default, or require explicit Solver=ARK32 selection?
• Subsystem stiff=True semantics vs. per-block flags: override the latter, or OR them?

References

Tableau construction:

• Kennedy, C. A., & Carpenter, M. H. (2003). Additive Runge-Kutta schemes for convection-diffusion-reaction equations. Appl. Numer. Math. 44, 139-181.
• Kennedy, C. A., & Carpenter, M. H. (2019). Higher-order additive Runge-Kutta schemes for ordinary differential equations. Appl. Numer. Math. 136, 183-205.

Theory + error analysis:

• Boscarino, S., & Russo, G. (2007). Error Analysis of IMEX RK Methods Derived from DAEs. SIAM J. Numer. Anal.
• Biswas et al. (2023). Algebraic Structure of the Weak Stage Order Conditions for RK Methods. SIAM J. Numer. Anal. (arXiv:2204.03603)

Reference implementations:

• Reynolds, D. R., et al. (2023). ARKODE: A flexible IVP solver infrastructure for one-step methods. ACM TOMS. (SUNDIALS/ARKode)
• PETSc/TS TSARKIMEX (PETSc manual, chapter 11).

Future: multirate (deferred)

• Chinomona, R., & Reynolds, D. R. (2021). Implicit-Explicit Multirate Infinitesimal GARK Methods. SIAM J. Sci. Comput. 43(5).
• Fish, A., Reynolds, D. R., & Roberts, S. (2023). Implicit-Explicit Multirate Infinitesimal Stage-Restart Methods. arXiv:2301.00865. — natural evolution of stiff=True into per-block timescale=... once plain IMEX is stable.

Competitive context

None of Simulink, Dymola, OpenModelica use IMEX-ARK as defaults (they use Rosenbrock, BDF, DASSL). Simscape's per-network "local solver" is the closest analog to per-block stiff=True, but uses fixed-step BE/TR — not IMEX. Block-level IMEX-ARK in fastsim would be novel in the control-sim space.