Year: 2026
Type: preprint
Abstract: We present the latest version of iGENE, a differentiable implementation of the local (flux-tube) gyrokinetic code GENE, now fully rewritten in JAX. iGENE features a multispecies, collisional, fully electromagnetic nonlinear gyrokinetic model while preserving end-to-end differentiability. The code is verified against GENE for linear ITG, TEM and KBM physics and for nonlinear turbulent fluxes with adiabatic and kinetic electrons. Forward-mode derivatives of the linear growth rate agree with finite differences for all five parameters tested, resolving the 30--50\% discrepancy in $\hat{s}$ and $ε$ reported for the earlier TensorFlow version. All collectives are native JAX operations, and the code strong-scales from 2 to 32 H100 GPUs with 86\% parallel efficiency. We then assess three routes to a shorter time to solution. Custom CUDA FFT kernels operating in bfloat16 are up to $3.97\times$ faster than fp32 cuFFT on the transforms, and on a single H100 iGENE is $2.63\times$ faster per time step than single-precision GENE. Multi-rate sub-cycling of the parallel dynamics yields a further $1.89\times$ speedup with kinetic electrons. Finally, a Hermite-Laguerre velocity-space representation is implemented, offering tunable fidelity.
Source: arXiv (Cornell University)
Institution: Max Planck Institute for Plasma Physics
Language: English
Subfield: Nuclear and High Energy Physics
Field: Physics and Astronomy
Domain: Physical Sciences
Open Access status: green
Funder: National Science Foundation
OpenAlex ID: W7223948686