[rogeriojorge]

Summary

• expose Diffrax's max_steps through simulation
• retain the existing default of 1,000,000
• add a regression that deliberately reaches the limit and verifies RESULTS.max_steps_reached

This is stacked on #14 so callers using throw=False can inspect the failure result. Diffrax uses max_steps for the bounded solver loop and checkpointing machinery, so production fixed-step runs should set it close to ceil(t_max / dt) with a small safety margin.

Validation

• pytest -q (4 passed)
• a completed fixed-step run with max_steps=11 is bitwise identical to the existing 1,000,000-step bound
• a one-step bound reports RESULTS.max_steps_reached

Two-A4000 benchmark

Orszag-Tang, 129x129 Fourier, 8x8x8 Hermite, Tsit5, 100 fixed steps:

• max_steps=1_000_000: 23.996 s
• max_steps=101: 13.941 s
• one A4000 with max_steps=101: 30.776 s
• resulting 1-to-2 GPU speedup: 2.21x

Both bounds conserve total energy to machine precision. Compilation remains approximately 1.9 s.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
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.

[rogeriojorge]

Additional production-duration GPU evidence (2026-07-12):

• Orszag–Tang, Nx=Ny=129, Nn=Nm=Np=12, two NVIDIA A4000 GPUs with species sharding
• 200 fixed steps (dt=1e-4, t=0.02), two saved states, callback-free execution
• compile + one-step warm-up: 25.587 s
• synchronized production integration with the current loose max_steps=1_000_000: 375.674 s
• relative total-energy error: 1.566e-16; field energy 0.13 -> 0.13000007998864274

This six-minute run reinforces the earlier controlled H8 result: setting max_steps close to the known fixed-step count is important at production sizes, independent of JIT compilation time.

[rogeriojorge]

Superseded by #44, which adds max_steps to simulation with the same default of 1000000. The RESULTS.max_steps_reached regression test here needs throw and solver_result from #14, so it can come back together with #14. Closing without merging; the branch is kept.