[rogeriojorge]

Summary

• guard the Newton GMRES update with lax.cond
• return the current iterate immediately when its scaled nonlinear residual is already below tolerance

This is stacked on #19. It changes one source file with a net two-line increase. The Newton criterion, Eisenstat–Walker forcing term, GMRES configuration, maximum iterations, and converged solution are unchanged.

Motivation

The Newton loop previously called GMRES before testing norm < 1. Every successful midpoint step therefore ended with a redundant zero-iteration linear solve. A zero-iteration Krylov solve still applies the Jacobian-vector operator to form true residuals, which is expensive for the Fourier–Hermite RHS.

A 33x33, 4^3 Orszag–Tang one-step profile over dt=1e-3 ... 1e-1 found that the useful Newton solves normally required one FGMRES iteration (at most two), followed in every case by a converged check whose GMRES result was discarded.

Performance

Five fixed midpoint steps, 33x33, 4^3, CPU, median of five cached runs:

• JAX GMRES: 1.116 s -> 0.736 s (34% reduction)
• SOLVAX pytree FGMRES from Support pytree operands in GMRES SOLVAX#10: 0.324 s -> 0.168 s (48% reduction)

Relative total-energy drift remained exactly 1.211254265768263e-16 in all comparisons.

Validation

• pytest -q: 5 passed
• ruff check spectrax/midpoint_solver.py tests/test_simulation_runs.py: passed
• git diff --check: passed

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

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

🚀 New features to boost your workflow:

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

[rogeriojorge]

Squashed into #18.