A Pre-trained Variational Autoencoder for Gyrokinetic Plasma Turbulence Surrogate Modeling

2026-09-29

AI Summary

This study addresses the prohibitive computational cost of direct numerical simulations (DNS) for plasma turbulence, which hinders the rapid generation of large-scale spatiotemporal evolution data. We propose a physics-informed latent-space surrogate model that employs a pretrained variational autoencoder (VAE) to compress turbulent fields and integrates ConvLSTM networks to learn latent dynamical evolution. Notably, we introduce transfer learning using Stable Diffusion pretrained weights and design Fourier spectral loss alongside manifold consistency error metrics, enabling multi-channel prediction without architectural reconstruction. Validated against the GENE code, our approach generates thousands of timesteps within seconds on a single GPU, significantly accelerating simulations while preserving multi-scale spectral accuracy.

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.

Problem

• plasma turbulence
• surrogate modeling
• gyrokinetic simulation
• computational acceleration

Innovation

• Variational Autoencoder
• Surrogate Modeling
• Physics-informed Loss
• Gyrokinetic Turbulence
• Transfer Learning