SPECTRAX: A JAX-based Spectral Kinetic Plasma Solver

We present SPECTRAX, a JAX-based spectral kinetic plasma solver, where the use of just-in-time compilation enables a performance comparable to C++ while retaining the simplicity of Python. Spectral solvers address the multiscale problem of weakly collisional plasmas by performing a spectral decomposition of the one-particle probability density function and the fields, and keeping only enough modes to model the physics of interest. Spectral plasma solvers have been developed before (e.g., SPS [1], GX [2]), but our implementation has the advantage of JAX’s auto-differentiation, allowing the computation of the exact gradients of the output parameters with respect to the input parameters with a single simulation. This potentially makes SPECTRAX a powerful tool for stellarator optimization, a problem of significant importance in fusion research. We demonstrate SPECTRAX's accuracy and performance through benchmark plasma physics simulations, including Landau damping and the two-stream instability.

[1] - V. Roytershteyn, G. L. Delzanno, Spectral Approach to Plasma Kinetic Simulations Based on Hermite Decomposition in the Velocity Space, 2018, Front. Astron. Space Sci. 5:27.

[2] - N.R. Mandell et al., GX: a GPU-native gyrokinetic turbulence code for tokamak and stellarator design, 2024, J. Plasma Phys., vol. 90, 905900402.