[rogeriojorge]

Exact Hermite free-streaming integrating factor for LawsonSolver (stacked on #82)

Stacking. Targets main so CI runs, but the branch contains the #82 commit; review only the last two commits (git log origin/agent/lawson-main..) = spectrax/_lawson_streaming.py, a few lines in _lawson.py / _simulation.py, README section, tests. Merge #82 first (then this PR reduces to the streaming files). The older integration-lineage version is merged into integration/dark-baseline only.

What it is

LawsonSolver(solver, factor_fn=...) generalizes the diagonal rate of #82 to an arbitrary smooth integrating factor Phi(tau), Phi(0) = I. hermite_streaming_factor builds the factor for the free streaming of the Hermite-Fourier system, and simulation(..., integrating_factor="streaming") uses it (default False stays bitwise unchanged; True is #82).

Derivation (details in the spectrax/_lawson_streaming.py docstring)

• Streaming acts along each Hermite axis as G = -i k (a T + u I) (k = 2 pi index / L, a width, u drift, T the symmetric tridiagonal T[j, j+1] = sqrt((j+1)/2)). The three axes act on different tensor factors, so they commute and exp(tau G) is a product. T = V diag(lam) V^T (one real eigh; lam the Gauss-Hermite nodes over sqrt 2, V the nodal transform) gives exp(tau G_ax) = exp(-i tau k u) V diag(exp(-i tau k a lam)) V^T: unitary, exact for any step, V shared by every k, species, width and drift (only the phases depend on them). Cost: two Nn x Nn products per Hermite axis per stage.
• Transformation: w = Phi^{-1} y gives the exact identity w' = Phi^{-1}(f(Phi w) - Phi' w), so the wrapped RK tableau applies to w for any smooth Phi (the choice of Phi only changes what is left in w'). Order statement (corrected): per Lawson (1967, SIAM J. Numer. Anal. 4, 372-380) the order of the underlying RK method is retained for a fixed linear part as h -> 0 (here the rate is the step-start rate, Lawson's re-evaluated-A form). No stiff-order claim is made when r h >> 1: per Hochbruck and Ostermann (2010, Acta Numerica 19, 209-286; Thm 2.17, Ex. 2.32, Table 2.2) the stiff order of Lawson methods drops (to 1 for Lawson-Runge) when ||hA|| is large. Empirical support is measured agreement with plain Dopri8 records (5e-10 in H05 c=64).
• Streaming only (Phi = B): w' = B^{-1} f(B w) - G w, i.e. N = f - G y is what the RK method sees.
• Streaming plus a diagonal rate r (does not commute with G): Strang product Phi = A(h/2) B(h) A(h/2), A = exp(-r .); with z = A w, u = B z, y = A u: w' = A^{-1} B^{-1}[A^{-1} f(y) + (r/2) u] - A^{-1} G z + (r/2) w, one B and one B^{-1} per stage; reduces to Lawson (integrating-factor) wrapper for the diagonal hypercollision damping (opt-in) #82 for G = 0. The symmetric BCH expansion shows the factor equals exp(tau (G - r) + O(tau^3 [G,[G,r]])), so what remains explicit is O((h |G|)^2 |r|). A stiff diagonal rate is therefore removed completely only while h |G| = h k a sqrt(2 Nn) is small. A joint exponential of G - R per k would need a non-Hermitian eig per step (rate changes every step; jnp.linalg.eig is CPU-only): not built.
• Moving basis: hermite_streaming_factor(basis_fn=...) takes (a, u) from the state at the step start (used in the integration lineage for the pump frame); the basis change over the step stays in N.
• The factor is exact in the dense output as well (y(t) = Phi(t - t0) w(t); oscillations are removed from w, which is smooth).

Measured

All on a laptop CPU (2 threads), after warm-up, simulation default two-stream, t_max = 20, 2 species, nu = 3:

Nx, Nn	tolerance	plain steps / wall	streaming steps / wall	rel. diff Ck / Fk vs plain
32, 64	1e-8	161 / 0.18 s	30 / 0.12 s	2.5e-7 / 3.0e-5
32, 64	1e-12	166 / 0.19 s	79 / 0.32 s	5.0e-11 / 3.1e-9
64, 128	1e-8	497 / 1.86 s	28 / 0.36 s	3.8e-7 / 5.4e-5
64, 128	1e-12	507 / 1.90 s	75 / 0.96 s	1.0e-10 / 6.4e-9
128, 128	1e-8	1049 / 7.11 s	29 / 0.58 s	8.2e-7 / 7.3e-5
128, 128	1e-12	1081 / 7.31 s	71 / 1.41 s	1.3e-10 / 8.7e-9

Cost per Runge-Kutta step: plain 3.5 ms, integrating_factor=True 4.6 ms, "streaming" about 12 ms (3.5x plain) at Nx 64, Nn 128. The streaming factor wins when the explicit step is stability-limited by streaming (h k a sqrt(2 Nn) < ~3.7) and the dynamics is smooth; it loses on small or accuracy-limited problems.

Decisive stiff-closure benchmark (Dark-SPECTRAX, gates declared in a committed script before any run; numbers in its docs/results.md): H05 v_q/v_te = 0.1, Nx 64, Nn 128, field-scaled closure (#66) with streaming plus the closure rate inside the factor, against the same-session closure-only Lawson (#82) and the plain record:

• c = 64 to t = 500: 9606 steps (closure-only Lawson) vs 9652 (streaming): streaming was not what limited the step; wall 52 s vs 130 s (2.5x slower; per-step cost 2.7x). Accuracy vs the plain CPU record: dK_e 3.4e-10, W_ext 2.7e-11; ledger 3.7e-12.
• c = 4: identical to t = 400; then the run ends at t = 456.5 on the step budget (201 927 steps in one chunk, closure-only: 5 444 for the same chunk): the Strang leftover O((h|G|)^2 |r|) is stiff in w.
• Ordinary B10 (nu = 100): closure-only 25 845 steps; streaming plus rate hits the 400 000-step budget.
• Gates S1 (steps ratio >= 1.5 or wall ratio >= 1) failed for both cases, S3 failed for c = 4; S2 (accuracy) passed. Verdict as declared: streaming does not earn its place in this regime. Consequently simulation exposes "streaming" only (the rate-inside-the-factor combination remains available through hermite_streaming_factor(rate_fn=...) and is documented as harmful at large h|G| and |r|).

Literature (read / not read)

• Crouseilles-Einkemmer-Massot, arXiv:1910.12720: read Sections 1-3 (Lawson for hyperbolic/kinetic problems, stability function of the Lawson scheme equals the RK one; the CFL for Lawson b dt < C dv, with b the advection speed; exact exponential for Fourier/semi-Lagrangian streaming). This is the precedent for exponentiating streaming; their numerical tests are Vlasov-Poisson and drift-kinetic, not Hermite.
• Einkemmer-Ostermann, arXiv:1207.2090: read abstract/introduction (Strang splitting for Vlasov, second order under regularity). Used to frame the Strang product above; the factor here is inside a Runge-Kutta step, not an operator split of the solution.
• Kassam-Trefethen (SIAM J. Sci. Comput. 2005): read the comparison; integrating-factor schemes mixed results. Same caution applies.
• Lawson (1967), Cox-Matthews (2002), Hochbruck-Ostermann (Acta Numerica 2010), Delzanno (JCP 2015, fully implicit alternative; local PDF consulted only for context): not used for any claim; the first three could not be fetched open-access.

Tests (tests/test_lawson_streaming.py, 7, all pass; full suite 25 passed locally)

1. G equals the streaming part of the real RHS (Fk = 0, nu = 0): 3-D, drifts, unequal widths, Nm, Np > 1;
2. B unitary, group property, equals scipy.linalg.expm, V diag(lam) V^T = T;
3. analytic Hermite front c_n = (-i x)^n / sqrt(n!) exp(-x^2/2) to 1e-12 in one step at 13x the Dopri8 imaginary-axis step (plain solver gives |c| > 50 at that step);
4. order (Heun 2, Dopri8 pre-asymptotic like the plain solver) with streaming alone and with a rate inside the factor;
5. simulation agrees with plain at the solver tolerance, k = 0 density preserved to round-off, default bitwise unchanged;
6. gradient w.r.t. the thermal width vs finite difference (1e-5);
7. basis_fn overrides the argument basis.

Limitations

• Cost per step 3.5x plain; benefit only for streaming-limited smooth runs. Not a cure for the stiff closure nor the asymmetric-weight (AW) Hermite truncation truncation instability (the exponential is exact).
• Joint streaming+closure exponential per k not implemented (non-Hermitian eig).
• Complex-state warning from Diffrax applies.

Merge recommendation

Do not merge unless you want an opt-in exact-streaming integrator for large Nx x Nn smooth runs. It is correct, tested, and bitwise off, but the measured benefit is narrow and it does not help the stiff-closure regime that motivated it. #82 is the one with demonstrated value. If merged, merge after #82.