[rogeriojorge]

Depends on #13 (structured tuple state) and #18. Their commits are included here until they merge; merge #13, then #18, then this.

Problem

Unpreconditioned Newton-GMRES in ImplicitMidpoint needs thousands of Krylov iterations once the timestep exceeds the Hermite streaming/collision time scales, and gives no iteration diagnostics.

Change

• Newton-Krylov solve delegated to solvax.newton_krylov (PyPI; deps only jax + equinox, already pulled by diffrax; solvax>=0.8.0 added to pyproject.toml/requirements.txt, tested with 0.8.0 and 0.26.0). New options linear_rtol, linear_atol, linear_restart, linear_max_restarts, preconditioner.
• collision_diffusion_x_streaming_preconditioner(args, dt): exact inverse of the collision + diffusion + x-streaming Hermite tridiagonal lines (jax.lax.linalg.tridiagonal_solve).
• simulation() returns midpoint_stats (total/max Newton and GMRES iterations) when the solver is ImplicitMidpoint.
• Default solver stays Dopri8 (see below).

Effect

1D Landau, Nx=33, Nn=128, nu=20, t=10, x64, same machine, interleaved runs (min of 3). Error = max |Fk - Fk_ref| / max |Fk_ref| against Dopri8 at tol 1e-11.

solver	dt	wall	error
Newton + jax.scipy GMRES (incremental), no preconditioner	0.5	29.5 s	1.1e-1
Newton + jax.scipy GMRES, line preconditioner	0.5	0.69 s	1.1e-1
this PR (SOLVAX + line preconditioner)	0.5	0.24 s	1.1e-1
Newton + jax.scipy GMRES, line preconditioner	0.02	8.3 s	1.9e-4
this PR	0.02	4.9 s	1.9e-4
Dopri8 adaptive, tol 1e-8 (default)	-	1.3 s	1.2e-7

• The preconditioner is the large win at large dt (about 40x). SOLVAX's Newton-Krylov is a further 1.7-3x over the best equivalent written with jax.scipy.sparse.linalg.gmres at the same preconditioner and iteration counts, and it reports GMRES iteration counts.
• Dopri8 stays the default. At equal accuracy it is 4-10x cheaper than implicit midpoint on the shipped-style cases: 1D Landau, a stiffer Landau (Nn=256, nu=100), 1D two-stream and a 2D (17x9) Landau case.

Where implicit midpoint wins

Dopri8's step is set by stability (step count does not change between tol 1e-4 and 1e-8), so implicit midpoint wins when the dynamics of interest are slow compared with the collision or streaming rates. Runs on 6 pinned cores (taskset), min of 2 runs, compile time excluded (implicit ~9 s, Dopri8 <3 s). Error is max over saved times of |Fk - Fk_ref| / max |Fk_ref|; the reference is Dopri8 at tol 1e-12. Implicit uses the line preconditioner with rtol=1e-8, linear_rtol=1e-6.

case	solver	steps	wall	field error
collisional Landau, Nx=3, Nn=64, nu=2000, t=50	Dopri8, tol 1e-4 to 1e-8	24061-24469	30.6-32.3 s	6e-6 to 3e-7
	implicit, dt=0.05	1000	2.9 s	1.4e-3
	implicit, dt=0.02	2500	2.6 s (12x)	2.2e-4
collisional Landau, Nx=33, Nn=128, nu=1000, t=20	Dopri8, tol 1e-4 to 1e-8	4399-4617	58.7-62.9 s	8e-6 to 3e-7
	implicit, dt=0.05	400	7.2 s (8x)	1.3e-3
	implicit, dt=0.02	1000	22.4 s (2.7x)	2.1e-4
ion acoustic wave, mi/me=100, Ti/Te=0.1, Nx=33, Nn=128, nu=50, t=300	Dopri8, tol 1e-4 to 1e-8	8035-8085	107-113 s	3e-4 to 3e-9
	implicit, dt=0.5	600	23 s (4.7x)	ion density 2.1e-4; field 1.7e-2 for t>=50
	implicit, dt=0.1	3000	89 s (1.2x)	field 5.6e-4 for t>=50, 3.3e-3 incl. the initial electron transient

• The ion acoustic case is a win only when the slow observable (ion density) is what you need. Resolving the field to 1e-3 needs dt=0.1, and then it is only 1.2x faster.
• At dt >= 1.5 in the ion acoustic case Newton fails on the first step, and the solve reports failure instead of returning a result.
• Reverse-mode gradients: neither the old nor this Newton loop is reverse-differentiable (while_loop); use Dopri8 (gradient matches finite differences to 8e-7).

Tests

test_implicit_midpoint_simulation_reports_iterations (preconditioned run, stats present) and test_collision_diffusion_x_streaming_line (matches a dense solve). Full suite passes with and without JAX_ENABLE_X64=1.

---

Review 2026-10-03 (branch not changed; head 9c066cc)

Scope. The only #21-specific commit is 567e0c4. The rest of the 28-commit head is accumulated #9/#12/#13/#18/#34/#46/#48/#50/#51 history.

Restacking needed. #18 was rewritten (now 806c2e4) with an outer-residual and finiteness fix. This branch must be rebuilt onto the new #18 (and #13) before review of the diff itself.

Failure handling. solvax.newton_krylov stops on a nonfinite residual or threshold with converged=False, and decides on the residual of the current iterate. So a NaN is not reported as success here. Two gaps remain:

• y_error is not set to inf on failure, so an adaptive controller may accept a failed step's estimate before the failure ends the solve. Corrected Fix implicit midpoint time, convergence reporting and pytree states #18 sets it to inf.
• The stopping test ‖F‖ ≤ max(atol, rtol‖F(y_Euler)‖) is unscaled and relative to the predictor's residual. That is a different contract from the scaled per-component test in Fix implicit midpoint time, convergence reporting and pytree states #18.

Coupling. The preconditioner reads parameters by negative positions in args (args[-20:-18], …). That is fragile against any change of the args tuple.

AD. No reverse-mode AD: the Newton loop is an early-exit while_loop. SOLVAX's fixed_work=True scan path is not used. Forward mode was not tested.

Matched-error profile

• Case: collisional 1D Landau (input_1D_landau_damping.toml with Nx=3, Nn=64, ν=2000, t=20), consistent initial field.
• Reference: Dopri8 at tol 1e-12.
• Error: max over saves of |Fk − Fk_ref| / max|Fk_ref|.
• Timing: time of the second call, so compile is excluded. Single run per row, on a laptop CPU.
• Midpoint settings: rtol=1e-8.
  • Use SOLVAX Newton-Krylov with a Hermite line preconditioner in implicit midpoint #21 uses linear_rtol=1e-6 with the line preconditioner.
  • Corrected Fix implicit midpoint time, convergence reporting and pytree states #18 uses atol=1e-10.

solver	dt / tol	steps	Newton / GMRES iterations	run time	error
Dopri8	tol 1e-4	8955	–	2.35 s	7.7e-6
Dopri8	tol 1e-8	9362	–	2.43 s	4.5e-8
#21 SOLVAX + line preconditioner	dt 0.004	5000	5048 / 9533	1.12 s	7.7e-6
corrected #18 (jax.scipy GMRES, no preconditioner)	dt 0.004	5000	–	4.18 s	8.8e-6
#21 SOLVAX + line preconditioner	dt 0.02	1000	1029 / 2144	0.23 s	2.2e-4
#21 SOLVAX, no preconditioner	dt 0.02	1000	2000 / 15459	0.90 s	2.2e-4
corrected #18	dt 0.02	1000	–	0.89 s	2.2e-4

• At matched error (about 8e-6), Use SOLVAX Newton-Krylov with a Hermite line preconditioner in implicit midpoint #21 is 2.1× faster than Dopri8 and 3.7× faster than plain midpoint.
• At 2e-4, Use SOLVAX Newton-Krylov with a Hermite line preconditioner in implicit midpoint #21 is about 4× faster than plain midpoint. Dopri8 cannot go that coarse: it is stability-limited, about 9000 steps at any tolerance.
• The gain comes from the line preconditioner. SOLVAX without it is no faster than corrected Fix implicit midpoint time, convergence reporting and pytree states #18.
• This is one stiff, collisional case. The PR's own tables show Dopri8 is 4–10× cheaper on the shipped low-collision examples.

Recommendation

Hold for the owner's numerical and performance review. Restack it on the corrected #18, then:

• set y_error = inf on failure;
• replace the positional args reads with named parameters;
• consider loading SOLVAX as an optional dependency.

[codecov]

Codecov Report

❌ Patch coverage is 98.19820% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
spectrax/_diagnostics.py	93.93%	1 Missing and 1 partial ⚠️

Files with missing lines	Coverage Δ
spectrax/_initialize_maxwellian.py	100.00% <100.00%> (+62.96%)	⬆️
spectrax/_simulation.py	100.00% <100.00%> (ø)
spectrax/midpoint_solver.py	98.59% <100.00%> (+65.25%)	⬆️
spectrax/_diagnostics.py	82.90% <93.93%> (+9.04%)	⬆️

🚀 New features to boost your workflow:

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

[rogeriojorge]

Production validation completed on the exact SPECTRAX/SOLVAX heads combined with species sharding: 129x129 x 8^3, 500 midpoint steps, 182.56 s on one A4000 and 98.31 s on two (1.86x, 93% efficiency). Both outputs are finite, have 1.09e-15 relative energy change, and agree to 9.5e-24 relative state norm.

[rogeriojorge]

Rebased as a stack so the PRs merge in sequence without conflicts: merge (with a merge commit) #48, #46, #50, #51, #9, #12, #13, #34, #52, #53, then #18, #21, #44. Until the earlier ones land this diff also shows their commits. The conflicts with #13/#50 (same simulation() signature, diffeqsolve call and test lines) are resolved here; the combined branch passes all 36 tests with and without JAX_ENABLE_X64.

[rogeriojorge]

Where implicit midpoint beats Dopri8 at matched accuracy: when Dopri8's step is stability-limited (its step count does not change with tolerance) and the physics is slow.

• Collisional Landau, nu=2000, Nx=3, Nn=64: 24k Dopri8 steps in 31 s vs 2.6 s implicit (dt=0.02, field error 2e-4), 12x.
• nu=1000, Nx=33, Nn=128: 59 s vs 7.2 s (dt=0.05, 1.3e-3), 8x; 22 s at 2e-4.
• Ion acoustic, mi/me=100, t=300: 4.7x for ion density at 2e-4. For the field at 1e-3, only 1.2x.

Dopri8 still wins on 1D Landau, two-stream and 2D Landau at shipped-like parameters, so it stays the default. Full table in the description.