GyroSwin: 5D Surrogates for Gyrokinetics Plasma Turbulence Simulations
Gianluca Galletti, Fabian Paischer, William Hornsby, Paul Setinek, Lorenzo Zanisi, Naomi Carey, Stanislas Pamela, Johannes Brandstetter
Nuclear fusion plays a pivotal role in the quest for reliable and sustainable energy production. A major roadblock to viable fusion power is understanding plasma turbulence, which significantly impairs plasma confinement, and is crucial for modelling plasma scenarios for next-generation reactor design. Plasma turbulence is governed by the nonlinear gyrokinetic equation, which tracks a 5D distribution function over time, but its high computational cost leads to the use of reduced models for turbulent transport. However, these models omit nonlinear effects unique to the full 5D dynamics. To this end, we introduce GyroSwin -- the first 5D neural surrogate for modelling nonlinear gyrokinetic simulations that can capture physical phenomena neglected by reduced models. GyroSwin (i) extends a hierarchical Vision Transformer to 5D and is trained on the 5D distribution function for the adiabatic electron approximation, (ii) introduces latent cross-attention and integration modules for 3D-5D interactions between potential fields and the distribution function, and (iii) performs channelwise mode separation based on nonlinear physics. Experimentally, we show that GyroSwin captures nonlinear physics such as zonal flows and significantly improves flux predictions over reduced numerical models in most cases. GyroSwin has the potential to drastically reduce the complexity of plasma turbulence simulation in 5D.