[rogeriojorge]

Stacked on #45: merge #45 first; until then this diff also shows #45's commits.

Problem

The Hermite-DG simulation flattened Ck and Fk into one vector, so their layouts could not be controlled, and the old batch mesh in _simulation.py was never used. The DG state could not be distributed over devices.

Change

• Keep the ODE state as the PyTree (Ck, Fk); Hermite_DG_system returns the same structure.
• simulation(..., shard_axis="x"|"y"|"z") shards both arrays over one spatial cell axis across all JAX devices (shard_map). shift_element exchanges the one-cell periodic halo with ppermute at device boundaries, and the adaptive controller uses a global RMS norm (psum) so all devices take the same steps.
• The default shard_axis=None path is unchanged and still jitted as a whole; the output also carries solver_result.

Effect (office, 16 cores, forced host devices, Nx=4096, Nn=20, Bosh3 adaptive, warm)

devices	1	2	4	8	16
time (s)	6.88	5.07	4.61	3.78	2.93

2.35x at 16 devices (machine load ~40, so noisy). Results are identical to serial to 1e-15 relative. The serial path runs in 5.5-6.1 s vs 6.4-6.9 s before (same numbers, no slowdown).

Tests

test_spatial_sharding_matches_serial compares serial and sharded Ck/Fk. With 4 forced host devices it fails when the halo ppermute is removed. Full suite: 10 passed with and without JAX_ENABLE_X64, plus test_simulation_runs.py with 4 host devices.

Update 2026-10-03 (head 99298f1)

Base changed to agent/dg-autodiff (#45); #45 is merged in (no rewrite), so the diff is the PyTree/sharding change only.

Checked with 2 fake CPU devices (XLA_FLAGS=--xla_force_host_platform_device_count=2), 2D, 4x4 cells, N_DG=2, H=3, 2 species, t=2:

• serial vs shard_axis="x" and "y": max |dCk| 4.7e-13 (scale 67), max |dFk| 7.1e-15 (scale 1). Halo exchange and the global RMS step-size norm reproduce the serial step sequence.
• jax.grad through the serial path works; through the sharded path it fails (ValueError: Length of device assignment 1 is not equal to the size of the mesh 2), because _place calls device_get/device_put on traced values.
• The sharded path ignores max_steps, adjoint and progress_meter (hard-coded 1000000, default adjoint, NoProgressMeter, throw=False), and simulation is no longer jit-decorated.
• Tests: pytest tests 18 passed; tests/test_simulation_runs.py with 2 fake devices 4 passed. Local results.

Open: make the sharded path differentiable (or refuse grad explicitly), pass the #45 solver keywords through, and measure cost on real devices before merge.

[codecov]

Codecov Report

❌ Patch coverage is 97.50000% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
spectrax/_simulation.py	96.87%	1 Missing and 1 partial ⚠️

Files with missing lines	Coverage Δ
spectrax/_model.py	100.00% <100.00%> (ø)
spectrax/_simulation.py	97.43% <96.87%> (-2.57%)	⬇️

🚀 New features to boost your workflow:

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