Description
Constructing a physically consistent equilibrium particle distribution that satisfies the steady-state gyrokinetic Vlasov equation is essential for high-fidelity gyrokinetic simulations, as it suppresses nonphysical initial transients and reduces particle noise. Such a consistent distribution is also crucial for long-timescale transport studies, where it enables accurate characterization of phase-space zonal structures (PSZS). However, obtaining this distribution requires solving a high-dimensional optimization problem, which can be computationally expensive especially at high resolution.
In this work, we present a novel high-performance implementation of this extraction method utilizing Python and the JAX library. Unlike traditional implementations, our approach leverages JAX’s automatic differentiation and Just-In-Time (JIT) compilation to optimize the iterative natural gradient descent scheme. A significant advantage of this new framework is its capability for seamless deployment on Graphics Processing Units (GPUs), allowing for massive parallelization of the phase space calculations.
We verify the method by comparing it with the conventional orbit average operator approach. The results demonstrate that our JAX-based solver not only maintains high numerical accuracy—effectively resolving the numerical errors near the trapped-passing boundary caused by diverging orbit periods—but also achieves a substantial improvement in computational speed. The successful deployment on GPUs drastically reduces the time-to-solution, making this algorithm highly suitable for large-scale kinetic simulations and rapid data analysis in fusion research.