Generative diffusion models for learning stochastic flow maps in particle-based kinetic simulation
Guannan Zhang
Oak Ridge National Laboratory
Computer Science and Mathematics Division
We report recent progress on the use of generative artificial intelligence (AI) methods to accelerate particle-based plasma kinetic computations of interest to magnetically confined fusion plasmas. These computations are time-consuming due to multiscale dynamics, boundary conditions and the need to follow large ensembles of particles to avoid statistical sampling errors. The models of interest are Fokker-Planck (FP) equations for the particle distribution function in phase space including drifts and collisions. The AI methods include Normalizing Flows (NF) and Diffusion Models (DM) that learn the probability distribution function of the final state conditioned to the initial state, such that the model only needs to be trained once and then used to handle arbitrary initial conditions. Going beyond Ref.[1], where we proposed the use of NF to accelerate the computation of hot-tail generation of RE, we present new results based on the use of DM that allow the quantification of confinement losses in bounded domains. We present a unified hybrid data-driven approach that combines a conditional diffusion model with an exit prediction neural network to capture both interior stochastic dynamics and boundary exit phenomena [2]. We present applications to the computation of RE generation and confinement.
[1] M. Yang et al, A pseudo-reversible normalizing flow for stochastic dynamical systems with various initial conditions, SIAM Journal of Scientific Computing, 46, (4) C508-C533 (2024).
[2] M. Yang et al, Generative AI models for learning flow maps of stochastic dynamical systems in bounded domains, J. Comp. Phys. 544, 114434, 2026.