[rogeriojorge]

Companion to #44, which makes the same change on main. The Hermite–DG simulation gains three keywords passed to diffrax.diffeqsolve. Their defaults reproduce the current behaviour exactly: max_steps=1000000, adjoint=RecursiveCheckpointAdjoint(), progress_meter=TqdmProgressMeter(1). With them:

• Forward mode works. jax.jacfwd runs through the solver with adjoint=ForwardMode(). Before, it failed with "can't apply forward-mode autodiff (jvp) to a custom_vjp function".
• Reverse-mode memory is bounded with RecursiveCheckpointAdjoint(checkpoints=K).
• Progress meter is optional. Differentiated and batch runs can turn off the host-callback progress meter.

The library change is 7 lines plus one docstring line in spectrax/_simulation.py, with no numerical change.

Tests (tests/test_autodiff.py, 5 tests, about 75 s on CPU)

The case is a two-species 2D Orszag–Tang setup projected with legT, exactly as in Examples/2D_Orszag_Tang.py: 4×4 cells, N_DG=2, 4³ Hermite modes, adaptive Dopri5 at tolerance 1e-8. The objective is the final field energy plus the Hermite free energy. Four Hermite modes are needed because the collision matrix is identically zero with three.

	d/d deltaB	d/d nu
reverse mode	82.55185639040157	−0.23192470646531285
forward mode	82.55185639040175	−0.23192470646531274
centred FD (step 1e-5)	82.55185650405394	−0.23192470379740368
reverse vs forward	2.2e-15	4.8e-16
reverse vs FD	1.4e-9	1.2e-8

• Checkpoint invariance. Gradients with checkpoints=2 and checkpoints=8 match the default adjoint to 1.3e-15.
• Bounded memory. With checkpoints=4, the compiled reverse-pass workspace stays at 4.1 MB as max_steps grows from 256 to 65536. With the default adjoint it grows from 6.0 MB to 56.1 MB.

Verification

• Tests. The branch goes from 4 to 9 passing tests, and the blocking flake8 check reports 0.
• Defaults unchanged. The PI's call pattern simulation(input_parameters, **solver_parameters), run with a reduced copy of input_2D_orszag_tang.toml, returns Ck, Fk and time bitwise identical to the unchanged branch.
• CI. This PR adds Discontinuous-Galerkin to the build_test.yml triggers, so CI now runs here and passes.

Merge notes

• Merge before Keep DG state as a PyTree and shard it over spatial cells #39, which is stacked on this branch (it now includes the former Shard DG state over spatial cells #40).

Differentiable thermal speeds and drifts (20605a8, e75052d)

compute_A_pm_matrices splits the Hermite advection matrices αJ + uI into upwind parts with jnp.linalg.eig. JAX cannot differentiate the complex eigenvectors eig returns, so jax.grad with respect to alpha_s or u_s raised NotImplementedError. These matrices are real symmetric (Jacobi), so jnp.linalg.eigh gives the same split and is differentiable. The change is three calls plus a comment. The constant Maxwell flux matrix keeps eig.

• Same numbers as before. For N = 1, 2, 3, 5 and 8 modes with nonzero drifts, the upwind matrices match the previous eig result to 8.9e-16. A short DG simulation matches to 7.6e-15 relative. A new test compares the split with an independent host LAPACK eigendecomposition.
• Gradients now work. The case is the final field energy on 4×4 cells, N_DG=2, 4³ Hermite modes, at a thermal-speed scale of 1 and an x-drift of 0.02:

	d/d thermal-speed scale	d/d x-drift
reverse mode	0.211158477295316	0.642760648209616
forward mode	0.211158477295316	0.642760648209616
centred FD	0.211158476570716	0.642760647640017
reverse vs FD	3.4e-9	8.9e-10

The drift is tested at a nonzero value because at zero drift the field energy is stationary in it by symmetry.

Found while testing, not changed here

• No energy diagnostics. diagnostics(output) is commented out in the DG simulation, and it expects mi_me and Omega_cs. Energies must therefore be computed from Ck and Fk directly, as the tests do.

Update 2026-10-03 (head c715971)

Base changed to sync/dg-with-main (#54), which is merged in (no history rewrite), so the diff is only this PR's 4 files.

The upwind split no longer calls a differentiable eigh per species and direction. A = alpha*J + u*I shares the eigenvectors of the constant Hermite Jacobi matrix J, so compute_A_pm_matrices diagonalizes J once on the host (per size) and forms V max(alpha*x_k+u, 0) V^T and V min(., 0) V^T.

Evidence (x64 CPU, standalone comparison of the two forms):

• values: max |eigh split - constant-basis split| <= 2.2e-15 for N = 1..40, generic and kink points;
• Jacobians in (alpha, u): agree to <= 1e-14 away from zero eigenvalues; forward = reverse to 1e-15;
• alpha = 0: the eigh Jacobian is NaN (degenerate spectrum), the constant-basis one is finite;
• at a zero eigenvalue (u = 0 for odd N, or u = -alpha*x_k) both forms used (lam +- |lam|)/2, and roundoff in the computed root decided which one-sided derivative autodiff returned (N = 40: Jacobians differ by 0.19). The new code symmetrizes the roots (middle root exactly 0) and uses maximum/minimum, so an exact zero eigenvalue gets the mean of the one-sided derivatives. Near (not exact) ties are still one-sided; that is inherent to the nonsmooth flux;
• cost: jitted reverse gradient of the split, N = 40: 225 us (eigh) vs 75 us (constant basis); N = 8: 41 vs 47 us.

New test test_upwind_split_derivatives_away_from_and_at_zero_speed[N=3,4,7]: finite AD Jacobian, forward = reverse, kink detected exactly where expected, AD = mean of one-sided finite differences, finite at alpha = 0.

Tests: pytest tests 17 passed (Python 3.12, JAX 0.11.2, CPU, x64 set by the package). Local results, not GitHub CI.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
spectrax/_initialization.py	92.48% <100.00%> (ø)
spectrax/_simulation.py	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.