Pre-trained AI model speeds up plasma turbulence simulations by thousands of times
A Pre-trained Variational Autoencoder for Gyrokinetic Plasma Turbulence Surrogate Modeling
Machine Learning
Summary
Simulating plasma turbulence, which is important for understanding fusion energy, takes a lot of computer power and time. The authors created a smart computer model that learns from images and physics rules to simplify the complex plasma data and predict how it changes over time. This method makes running many simulation steps much faster while keeping important details accurate. Their approach works on different kinds of plasma turbulence data without needing big changes to the model.
What this means in practice
- •For fusion simulation engineers: Run large-scale plasma turbulence simulations faster on GPUs by using this surrogate model to approximate detailed physics accurately.
- •For computational fluid dynamics teams: Accelerate turbulence modeling for related fluid systems by adapting the pre-trained variational autoencoder framework to different turbulence datasets.
Authors
Minglei Yang, Marshall Nicholson, Diego Del-Castillo-Negrete, David Hatch, Guannan Zhang
Abstract
Machine learning surrogate models offer a promising path toward accelerating plasma turbulence simulations. We present PreVAE-Turb, a surrogate modeling framework that leverages pre-trained variational autoencoders (VAEs) from the Stable Diffusion image generation model for efficient spatial compression of turbulence fields. The pre-trained VAE is fine-tuned on turbulence data using a physics-informed loss function that includes a spectral loss operating in Fourier space to enforce spectral accuracy across scales. The VAE is combined with convolutional long short-term memory (ConvLSTM) networks to learn temporal dynamics in latent space, with a manifold consistency error metric that monitors encode--decode consistency during autoregressive rollouts. We validate the framework on two-dimensional Hasegawa-Wakatani drift-wave turbulence and extend it to gyrokinetic turbulence from the GENE code, where a four-channel adaptation simultaneously predicts electrostatic potential, density, and parallel/perpendicular temperature fluctuations without requiring architecture redesign. Once trained, inference generates thousands of time steps in seconds on a single GPU, providing substantial computational acceleration compared to direct numerical simulation. The pre-trained approach offers a transferable methodology broadly applicable to various turbulence simulation codes.