Shard simulations across Hermite modes

[rogeriojorge]

Summary

• shard the first Hermite axis in addition to the existing species axis
• exchange only one neighboring Hermite plane with lax.ppermute
• use one axis-generic global RMS norm for adaptive Diffrax solves
• keep spatial FFTs local and electromagnetic fields replicated

This is stacked on #15 and reuses its species_mesh argument. The Vlasov-Maxwell equations are unchanged.

Communication layout

The distribution uses P("species", None, None, "hermite"). The RHS communicates nearest-neighbor Hermite halos and the low-order current; it does not all-gather the distribution. This follows the Fourier-Hermite layout of Parker and Dellar: retain complete Fourier spectra locally and communicate adjacent Hermite modes.

Validation

• pytest -q: 5 passed
• a randomized four-device adaptive run matches serial to 1e-12
• local fixed-step 4/8-device outputs are bitwise identical to serial
• exact-head one/two-GPU outputs are bitwise identical
• final patch: 133 additions and 33 deletions across four files, reduced from 178/46

Post-warm Orszag-Tang timing on a 14-core Apple M3 Max, 129x129 Fourier, 32x4x4 Hermite, 100 Euler steps:

CPU devices	Warm time	Speedup	Efficiency
1	36.16 s	1.00x	100%
4	12.61 s	2.87x	71.7%
8	10.01 s	3.61x	45.2%

The eight-device process reaches about 840% CPU utilization. A smaller 65x65, 16x4x4 case saturated at 1.88x on four and 1.96x on eight, so useful local work per shard is essential.

Exact-head A4000 timing at 129x129 Fourier, 28x4x4 Hermite, 200 Euler steps:

GPU devices	Warm time	Speedup	Efficiency
1	9.02 s	1.00x	100%
2	6.58 s	1.37x	68.5%

Ck, Fk, and total energy are bitwise identical between the one- and two-GPU runs. Both A4000s were otherwise idle during measurement.

Constraints

• the species mesh dimension must divide Ns
• the Hermite mesh dimension must divide Nn and leave at least two local modes
• saved output is gathered before the existing flattened diagnostics
• mapped runs disable the progress callback

References: https://arxiv.org/abs/1407.1932 and https://docs.jax.dev/en/latest/notebooks/shard_map.html