[rogeriojorge]

Target: main. Head: 806c2e4, two commits on main ab87385. Old head: 60ca682, a 25-commit integration merge carrying #9/#12/#13/#34/#46/#48/#50/#51/#52/#53.

Dependencies: none. The step uses jax.tree.map, so it works with the flat state on main and with #13's (Ck, Fk) tuple state. #21 builds on this solver.

Problem

ImplicitMidpoint on main:

• (a) evaluates the RHS at t1 instead of t_n + dt/2;
• (b) always returns successful;
• (c) breaks on PyTree states.

The previous fix for (b) set not_converged = norm >= 1 and returned its complement. A NaN residual is False for both < 1 and >= 1, so a NaN step ended as success. The flag also described the residual before the last update, not the iterate that was returned.

Change (spectrax/midpoint_solver.py only)

• The RHS is evaluated at t_mid, and all state arithmetic uses jax.tree.map. This is the original commit, kept.
• Convergence is decided on the residual of the returned iterate. That residual must be finite and its scaled norm < 1. Tolerances are unchanged (rtol=1e-6, atol=1e-8, same scaled norm).
• A nonfinite residual or Newton update ends the step as nonlinear_divergence. This also covers a failed inner GMRES solve. An exhausted Newton budget ends as nonlinear_max_steps_reached.
• A failed step returns an infinite error estimate, as Diffrax's own implicit solvers do:
  • with an adaptive controller, the step is rejected and retried smaller;
  • with constant steps, the failure becomes the result of the solve.
• GMRES restarts are max(1, min(20, max_iters // 2)). Before, max_iters=1 gave zero GMRES restarts, so the single allowed update was a no-op.

Tests (local, macOS CPU, Python 3.12, jax 0.11.2, diffrax 0.7.2, at 806c2e4)

python -m pytest -q -n 2 tests: 21 passed. With JAX_ENABLE_X64=1: 21 passed.

New file tests/test_midpoint_solver.py, 10 tests. Run against the previous #18 solver, the 9 failure-path tests all fail.

• NaN / ±Inf RHS: result is nonlinear_divergence and the error estimate is inf.
• NaN reached only after an update.
• Failed inner solve: a custom_jvp with a NaN tangent gives a failure, never success.
• Max iterations: with y'=-y³, one update is not enough, and a converged run satisfies the midpoint equation to 1e-9.
• Last allowed update: linear y'=-y with max_iters=1 succeeds with the exact (1-dt/2)/(1+dt/2).
• Zero iterations: accepts an already exact predictor and rejects an inexact one.
• Whole solves: a constant-step solve ends with the failure, and an adaptive solve rejects failed steps and still finishes accurately.
• Harmonic oscillator (ω=2, dt=0.2, 250 steps, rtol=1e-12):
  • energy is conserved to < 1e-11;
  • phase follows (2/dt) arctan(ω dt/2) to 1e-9;
  • the accumulated lag matches ω³dt²t/12 within the expected O((ω dt)²) correction (2.3% here).

Evidence on a Landau case (not in CI)

Setup: 1D Landau, Nx=4, Nn=40, ν=0, t=20, consistent initial field as in Examples/1D_Landau_damping.py. Reference: Dopri8 at tol 1e-10. Midpoint uses constant steps with rtol=1e-10.

dt	max rel. total-energy change (saves at steps)	rel. E₁ error
0.4	1.5e-8 (saves interpolated)	8.4e-2
0.2	1.7e-12	2.1e-2
0.1	1.7e-12	5.3e-3

• Order: the error falls 4× per halving of dt, so the method is second order.
• Energy: at saves that fall on steps, the energy change equals the Dopri8 reference's own 1.7e-12. At saves between steps, LocalLinearInterpolation does not conserve energy (about 1e-9 at dt=0.1).
• Phase: we compared zero crossings of Re E₁ with the reference (ω = 1.3439 from the reference). The phase lag at t = 17.9 is:

dt	measured lag	predicted by (ω/ω̃ - 1) ω t, ω̃ = (2/dt) arctan(ω dt/2)	ratio
0.4	0.561	0.568	0.988
0.2	0.142	0.144	0.984
0.1	0.0355	0.0361	0.982

Notes for review

• Error estimate is first order. y_error = y1 − (Euler predictor) is an O(dt²) local estimate for a second-order method, so adaptive midpoint steps are conservative. This is unchanged here.
• No reverse-mode AD. The Newton loop is a lax.while_loop, so it is not reverse-differentiable. Use Dopri8 for jax.grad.

GitHub CI: see the checks on this head. The local results above are not CI results.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
spectrax/midpoint_solver.py	98.30% <100.00%> (+64.97%)	⬆️

🚀 New features to boost your workflow:

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

[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.