Expose midpoint iteration diagnostics

[rogeriojorge]

Summary

• accumulate accepted-step Newton and GMRES iteration totals in Diffrax solver state
• retain the maximum Newton and GMRES work observed in any accepted time step
• expose the compact state as midpoint_stats only for ImplicitMidpoint

Why

Steady-state timing alone cannot distinguish a costly spectral residual/JVP from
poor nonlinear convergence. These four counters provide that distinction without
logging inside JIT-compiled loops or retaining per-step arrays. Diffrax rolls
solver state back when an adaptive step is rejected, so the totals describe the
accepted trajectory.

The explicit-solver output is unchanged: SPECTRAX does not request or expose
Diffrax solver state for other solvers, whose state can contain arrays as large
as the physical state.

Validation

• 7 passed in the complete SPECTRAX test suite
• Ruff passes for changed code (pre-existing _simulation.py import-order and unused-import findings excluded)
• sdist and wheel build successfully
• 65² Fourier, 28×4×4 Hermite, 20 steps: 40 Newton / 60 GMRES iterations, 17.75 s
• 129² Fourier, 40×4×4 Hermite, 50 steps: 100 Newton / 150 GMRES iterations, 326.07 s, finite, relative energy change 1.21e-16
• matched 128², 40×4×4, 50-step control: identical iteration counts and energy change
• matched 129²/40-mode one-A4000 run: identical counters, 45.08 s (restart 20) and 35.08 s (restart 4)
• 193² Fourier, 48×4×4 Hermite, 100 steps on one A4000: 200 Newton / 300 GMRES iterations, 313.48 s, finite
• 257 Fourier × 128 Hermite Landau, 100 steps: counters identify restart-16 and linear-rtol 1e-3 as the measured work minimum, while ruling out restart 4

The larger run preserves exactly 2 Newton and 3 GMRES iterations per step. Its
superlinear runtime growth therefore lies in residual/JVP execution rather than
nonlinear convergence.

Stack

This draft intentionally targets agent/midpoint-collision-diffusion-preconditioner
(#28), so the PR contains only the diagnostics change.