5D Neural Surrogates for Nonlinear Gyrokinetic Simulations of Plasma Turbulence

Abstract

Nuclear fusion plays a pivotal role in the quest for reliable and sustainable energy production. A major roadblock to achieving commercially viable fusion power is understanding plasma turbulence, which can significantly degrade plasma confinement. Modelling turbulence is crucial to design performing plasma scenarios for next-generation reactor-class devices and current experimental machines. The nonlinear gyrokinetic equation underpinning turbulence modelling evolves a 5D distribution function over time. Solving this equation numerically is extremely expensive, requiring up to weeks for a single run to converge, making it unfeasible for iterative optimisation and control studies. In this work, we propose a method for training neural surrogates for 5D gyrokinetic simulations. Our method extends a hierarchical vision transformer to five dimensions and is trained on the 5D distribution function for the adiabatic electron approximation. We demonstrate that our model can accurately infer downstream physical quantities such as heat flux time trace and electrostatic potentials for single-step predictions two orders of magnitude faster than numerical codes. Our work paves the way towards neural surrogates for plasma turbulence simulations to accelerate deployment of commercial energy production via nuclear fusion.

Figures (8)

• Figure 1: Overview of the model architecture and n-dimensional attention. Left: Our 5D Swin-UNet with two down/upsampling stages. We indicate shapes at each resolution change. Right: The locality of the n-dimensional shifted window attention for nDWin-MSA (top) and nDSWin-MSA (bottom), illustrated as a 2D plane of 3D window-partitioned tokens. Connected components are highlighted as colored tokens in the 3D blocks, and as lines connecting them across dimensions.
• Figure 2: Left: 5D distribution function with an ITG of $7.9$ in the linear phase ($t=1.9$). Right: 5D distribution function with an ITG of $7.9$ in the saturated phase ($t=50.6$).
• Figure 3: 5D Swin-UNet can accurately predict the distribution functon in five dimensions. We report one-step model predictions of the distribution function, $\delta f_{\text{pred}}$, versus ground truth, $\delta f_{\text{GT}}$, over time. For each axis pair, the corresponding 2D projections are obtained by averaging over the remaining dimensions.
• Figure 4: Visualization and comparison of the electrostatic potentials for model predictions ($\phi_{\text{pred}}$) versus ground truth ($\phi_{GT}$). In accordance with \ref{['fig:pred_df_poten']}, rows correspond to timesteps $[12.7, 55.7, 102.1]$. \ref{['fig:pot_only_zonal']} shows only the zonal component, while \ref{['fig:pot_zonal']} displays all modes of the field. While the vertical waves are well captured, the model overestimates the zonal flow, particularly at the first timestep, corresponding to the linear phase.
• Figure 5: Heat flux time trace ($\int \delta f$) for ground-truth (GT) and single-step prediction of 5D Swin-UNet (pred) for the holdout trajectory (ITG=6.9).