Keep solver state axes structured

[rogeriojorge]

Summary

• keep distribution and field coefficients as separate Diffrax PyTree leaves
• preserve (species, p, m, n, y, kx, z) axes during time integration
• remove full-state slicing and concatenation from every RHS evaluation
• preserve the existing flattened Ck and Fk output layout

This is the smallest state-layout prerequisite for later JAX sharding: the distribution's species and Hermite axes now remain visible to the compiler throughout the solve.

Numerical validation

• pytest -q (3 passed)
• fixed-step Ck and Fk outputs are bitwise identical to origin/main
• adaptive-step Ck and Fk outputs are bitwise identical to origin/main

CPU benchmark

Fresh-process Orszag-Tang runs with Tsit5 and fixed dt=0.002:

• 65x65 Fourier, 6x6x6 Hermite, 50 steps: 9.736 s -> 9.027 s mean compile-plus-run time across three trials (7.3%)
• 129x129 Fourier, 8x8x8 Hermite, 10 steps: 22.741 s -> 21.500 s stepping time (5.5%)
• 129x129 peak footprint: 2.77 GB -> 2.62 GB (5.1%)

Lowering plus compilation was slightly slower in the separated 65x65 measurement (0.911 s versus 0.829 s); the measured gain is in stepping, not compilation.

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