[rogeriojorge]

Lawson (integrating-factor) wrapper for the diagonal damping — opt-in, bitwise off

Targets main (bf41830). Self-contained: it needs nothing from the closure PRs (#66 etc.). The older #81 is the same wrapper stacked on the integration lineage (field-scaled closure rate, Dark-SPECTRAX); it is merged into integration/dark-baseline only and is kept as the integration variant. This PR is the clean, main-based version.

What it is

• spectrax.LawsonSolver(solver, rate_fn): wraps any explicit Diffrax Runge-Kutta solver (Dopri8, Heun, ...) so a diagonal damping -r y of the right-hand side is integrated exactly. rate_fn(t, y, args) returns a non-negative pytree like y (zeros for states that must not be damped). It is the hook for any stiff diagonal rate, e.g. the field-scaled closure rate of Field-scaled closure rate for the asymmetric-Hermite top-mode instability #66 (not a dependency).
• simulation(..., integrating_factor=True): uses the hypercollision nu * col and hyper-diffusion D k^2 rates of the existing right-hand side. Default False is bitwise unchanged (tested). Steps are clipped onto the save times (see dense output below).
• hypercollision_rate, simulation_rate helpers. +171 lines of library code, +138 of tests, +20 of simulation, a README section.

Derivation (also in the spectrax/_lawson.py docstring)

Write f(t, y) = -r y + N(t, y), r = r(y0) >= 0 frozen at the step start, N = f + r y0 * y. That is an identity for every r. With w = exp(r (tau - t0)) y:
w' = exp(r (tau - t0)) (f(tau, exp(-r (tau - t0)) w) + r exp(-r (tau - t0)) w); one step of the wrapped RK tableau on w from w(t0) = y0, then y1 = exp(-r h) w1.

• Order is kept. For fixed r this is the classical Lawson-RK method of the wrapped tableau; because the split is an identity for every r, choosing r = r(y0) (a smooth function of the step-start state, fixed before the stages) does not change the order. Verified numerically: Heun-Lawson order 2, Dopri8-Lawson order 8 with a state-dependent rate.
• Stability. The stability function of a Lawson scheme equals that of the underlying RK method applied to N (arXiv:1910.12720, Section 3, which I read). So -r y no longer enters the stability condition h r < ~5.17 of Dopri8.
• Error estimate and dense output are mapped back: y_err = exp(-r h) w_err; y(t) = exp(-r (t - t0)) w(t).
• Frozen rate. The true rate varies over the step; the variation sits inside N and is treated explicitly (this is why r should vary slowly on the step, as the field-scaled closure does through |E|).
• Dense-output caveat (measured). w(t) is a polynomial; y = exp(-r (t - t0)) w is accurate only for r h of order one. Saves on step ends are exact. simulation(integrating_factor=True) therefore wraps the controller in ClipStepSizeController(step_ts=save_times) (users of LawsonSolver directly should do the same when the rate is stiff).
• Overflow guard (found while benchmarking this PR). exp(r h) overflows near r h ~ 700. An adaptive controller will happily propose such a step (the stiff mode is exactly resolved so the error estimate is tiny); the candidate is then inf, Diffrax's scale atol + rtol max(|y0|, |y1|) becomes inf, the scaled error 0, and the garbage step is accepted. Now a non-finite candidate is returned as y0 with a finite huge error (1e30; an infinite one poisons the PID memory and the next step becomes inf), so it is rejected and the step shrinks. Finite steps are bitwise untouched. Regression test included; without the guard the nu = 4000 run below never finishes (20 000 rejected steps).

Evidence (this machine, CPU, x64; simulation(Nx=17, Nn=24, Ns=2, timesteps=5, t_max=2, nu=...), wall after warm-up)

nu	plain steps / wall	Lawson steps / wall	max abs Ck difference at saves
3 (default)	18 / 0.01 s	20 / 0.01 s	agree at solver tolerance (test)
300	106 / 0.04 s	10 / 0.01 s
4 000	1 830 / 0.68 s	21 / 0.01 s
40 000	19 597 / 7.43 s	713 / 0.34 s

(t_max = 5 for the nu = 3 row.) Larger gains are limited by r h < ~700 and by forced-mode accuracy (below). In the Dark-SPECTRAX study (field-scaled closure rate, Nx 64, Nn 128, GPU/CPU records in that repo) the same wrapper gave 5.0x fewer steps at c = 64 to the t = 500 mark and no gain at c = 4 (closure not stiff before ~540); ordinary-dispersion B10 case: 4.6x fewer steps but 1.5x wall (small state, per-step cost 3x), and a ledger defect that is 1e-9 instead of 1e-11 because the plain run over-resolves. Those are in the Dark repo docs/results.md; they are reported, not re-derived here.

When it helps and when it does not

• Helps: a diagonal rate makes the explicit step stability-limited (nu * h, D k^2 h, or a field-scaled closure damping that grows with the field).
• No gain, small cost (one extra RHS evaluation per step, because FSAL's cached stage lives in the other variable): default nu (3), where the step is accuracy-limited. 18 -> 20 steps above.
• Forced stiff modes: w' ~ exp(r tau) N. If the stiff row is driven by a smooth source, the transformed integrand is far from polynomial and the step remains limited by accuracy for r h beyond a few tens. On the toy y' = -4000 y + sin t (rtol 1e-6) it needs 1 143 steps (the plain stability bound alone would be ~1 600; plain not run). This is the known drawback of Lawson relative to phi-function (ETD/Rosenbrock) integrators; see literature.
• It does not fix the asymmetric-weight (AW) Hermite truncation instability. The exponential is exact, so it neither adds nor removes damping. A step-size collapse from a growing instability is not cured.
• Not conservative beyond the order of the scheme; density/momentum/energy rows should get rate 0 (the hypercollision spectrum already vanishes for n <= 2).
• Only ODETerm problems; wrapped solver must be explicit RK with an embedded error estimate (Dopri8, Tsit5, Heun, ...); r h < ~700.

Literature (what I read, and what I did not)

• Crouseilles, Einkemmer, Massot, Exponential methods for solving hyperbolic problems with application to kinetic equations, arXiv:1910.12720 (J. Comput. Phys. 2020): read Sections 1-3. Lawson stability function equals the RK one; for kinetic (hyperbolic) problems Lawson is preferred over ETD/ExpRK for stability, while for parabolic problems Lawson has "all but lost" its appeal because phi-function methods converge better (also: full order with inhomogeneous boundary conditions cannot be achieved by Lawson).
• Kassam, Trefethen, Fourth-order time-stepping for stiff PDEs, SIAM J. Sci. Comput. 26 (2005): read abstract and the comparison. Integrating-factor RK4 is competitive for Burgers, worst for Kuramoto-Sivashinsky, and fails for KdV/Allen-Cahn in their setting; ETD4RK (their contour-integral version of Cox-Matthews) is best. Consistent with the forced-mode caveat above.
• Einkemmer, Ostermann, Convergence analysis of Strang splitting for Vlasov-type equations, arXiv:1207.2090: read abstract/introduction (second order under regularity). Used only for the splitting discussion in the follow-up streaming PR.
• Lawson (1967, SIAM J. Numer. Anal. 4), Cox-Matthews (2002, J. Comput. Phys. 176), Hochbruck-Ostermann (2010, Acta Numerica 19): not read (paywalled; no open copy found by me); cited only as the standard references for integrating-factor and exponential methods. No claim here depends on them beyond what the two open papers above say.
• Delzanno, J. Comput. Phys. 301 (2015): context only (fully implicit, conservative alternative for the same stiffness); not used for any claim.

Tests (tests/test_lawson.py, 10, all pass locally; full CI below)

1. frozen diagonal is exact at r h = 200 (plain Dopri8 blows up at the same step);
2. order 2 (Heun) and 8 (Dopri8) with a state-dependent rate;
3. zero rate is bitwise equal to the wrapped solver, same step count, dense output (SaveAt ts);
4. dense output versus exact decay;
5. pytree state with a rate-0 ledger-like row;
6. simulation default unchanged bitwise; Lawson agrees at solver tolerance; n <= 2 rows agree to 1e-6;
7. gradient (jax.grad) versus central finite difference, rel 1e-5;
8. stiff nu = 4000 simulation matches plain, with at least 2x fewer steps;
9. rate layout equals the flat [Ck, Fk] state, rate zero on fields;
10. overflow step is rejected, not accepted.

Limitations / review notes

• rate_fn is evaluated at the step start; its time variation inside the step is explicit (documented above).
• Complex-state support in Diffrax is flagged "work in progress" by Diffrax itself; the tests compare against the plain solver on the same complex state.
• The simulation option on main covers hypercollision/hyper-diffusion only (main has no field-scaled closure); custom closures use LawsonSolver directly.

Merge recommendation

Safe to merge if the owner wants the option: opt-in, bitwise off, no new dependency, small. Value is concentrated in stiff-closure runs (large nu, strong D, or the #66 field-scaled closure). If SPECTRAX's intended regime is the default nu, there is no benefit and the PR can wait; nothing else depends on it. A stacked follow-up (exact Hermite streaming, #83) is a separate PR on top of this one; its benchmark showed that exact streaming does not help the stiff-closure regime, so #83 is optional and this PR does not depend on it.