[rogeriojorge]

Problem

The Diffrax state is one flat vector: every RHS call slices/reshapes it and concatenates the derivatives, and the species/Hermite axes are invisible to the compiler (a prerequisite for sharding).

Change (_simulation.py)

The state is the tuple (Ck, Fk) with shapes (Ns, Np, Nm, Nn, Ny, Nx//2+1, Nz) and (6, Ny, Nx//2+1, Nz). Output Ck/Fk layout is unchanged. API note: the public ode_system now takes and returns the tuple instead of a flat vector.

Evidence (tested SHA of the code change d5edbc4, rebased on main ab87385; local, Python 3.12.9, JAX 0.11.2, diffrax 0.7.2)

Same script run in a main worktree and a branch worktree; RHS (flattened) and Ck/Fk trajectories, Tsit5, fixed and adaptive steps, max|diff| = 0 (bitwise identical) for:

• 1D: Nx=9, Nn=6, Ns=2;
• 3D/3V: (Nx,Ny,Nz)=(7,5,4), (Nn,Nm,Np)=(3,2,4), Ns=3, anisotropic shifted widths, nu=0.5, random state.
  Earlier timing on Orszag-Tang (from the original PR): 65x65 / 6^3, 50 steps, 9.74 s -> 9.03 s; 129x129 / 8^3, 10 steps, 22.7 s -> 21.5 s, peak 2.77 -> 2.62 GB (not re-measured here).

Tests

• test_solver_state_keeps_species_and_hermite_axes: integrator y0/RHS shapes (fails on main).
• test_structured_state_maps_flat_hermite_index_to_species_p_m_n: 3 species, unequal orders, flat index s*Nn*Nm*Np + p*Nn*Nm + m*Nn + n -> state[s,p,m,n].
• test_continuation_from_output_matches_single_run: restart from a saved (Ck, Fk) matches the uninterrupted fixed-step run to 1e-13.
  python -m pytest -n 2 -q tests at head 53777eb: 11 passed.

No dependencies. Conflicts trivially with #9/#12/#55 in the simulation output block.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
spectrax/_simulation.py	100.00% <100.00%> (ø)

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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