Spatio-temporal forecasting of plasma turbulence using deep learning

Recent years have seen massive growth and improvement in machine learning models in fusion plasma physics in models that mitigate plasma disruption (posters by Jalavand et al., Kim et al., Shousha et al.), reconstruct equilibria (EFIT AI, poster by Jang et al.), and perform turbulence closures (Shukla et al., poster by Barbour et al.). The number of trainable parameters in these models is in the 10^3-10^5 range. However, creating a robust model would require relatively many more trainable parameters ~10^8 for spatiotemporal plasma turbulence prediction. One must distribute it over multiple nodes on a computing cluster to ensure fast training with such a large model.

In this talk, we present a parallelized, data-driven, deep-learning model for performing closures in plasma turbulence: the Fourier GRU (Gated Recurrent Unit). After distributing the weights of the model over multiple GPUs, we take the ground truth data from the fast flux tube gyrokinetic solver GX and train these networks to predict the heat flux. This effectively provides a purely data-driven driven closure for any flux tube simulation. Then we test our models for different equilibrium configurations, such as a Z-pinch, tokamak, and stellarators, and present our results. We also use the neural network to learn the dynamical properties of the gyrokinetic model, such as the Lyapunov exponents and the chaotic attractor dimension.