# A Generative Artificial Intelligence framework for long-time plasma turbulence simulations

B Clavier, D Zarzoso, D Del-Castillo-Negrete, E Frénod

► To cite this version:

B Clavier, D Zarzoso, D Del-Castillo-Negrete, E Frénod. A Generative Artificial Intelligence framework for long-time plasma turbulence simulations. Physics of Plasmas, In press, 10.1063/5.0255386. [hal-05085168](https://hal.science/hal-05085168v1)

HAL Id: hal-05085168

https://hal.science/hal-05085168v1

Submitted on 13 Aug 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# A Generative Artificial Intelligence framework for long-time plasma turbulence simulations

B. Clavier

Aix Marseille Univ, CNRS, Centrale Med, M2P2 UMR 7340, Marseille, France

D. Zarzoso*

Aix Marseille Univ, CNRS, Centrale Med, M2P2 UMR 7340, Marseille, France

david.zarzoso-fernandez@cnrs.fr

D. del-Castillo-Negrete

Institute for Fusion Studies and Department of Physics, University of Texas at Austin, Austin, Texas 78712, USA.

E. Frénod

Université Bretagne Sud, LMBA UMR 6205, Vannes, France

May 26, 2025

## Abstract

Generative deep learning techniques are employed in a novel framework for the construction of surrogate models capturing the spatio-temporal dynamics of 2D plasma turbulence. The proposed Generative Artificial Intelligence Turbulence (GAIT) framework enables the acceleration of turbulence simulations for long-time transport studies. GAIT leverages a convolutional variational auto-encoder and a recurrent neural network to generate new turbulence data from existing simulations, extending the time horizon of transport studies with minimal computational cost. The application of the GAIT framework to plasma turbulence using the Hasegawa-Wakatani (HW) model is presented, evaluating its performance via various analyses. Very good agreement is found between the GAIT and the HW models in the spatio-temporal Fourier and Proper Orthogonal Decomposition spectra, the flow topology characterized by the Okubo-Weiss parameter, and the time autocorrelation function of turbulent fluctuations. Excellent agreement has also been obtained in the probability distribution function of particle displacements and in the effective turbulent diffusivity. In-depth analyses of the latent space of turbulent states, choice of hyper-parameters and alternative deep learning models for the time prediction are presented. Our results highlight the potential of AI-based surrogate models to overcome the computational challenges in turbulence simulation, which can be extended to other situations such as geophysical fluid dynamics.

## Introduction

Turbulence plays a key role in nature, laboratory experiments, and industry. Examples include stars [10], the solar tachocline [48], the solar wind [7], the atmosphere of planets [49], oceanic currents [24] as well as aeronautics and industrial systems [44]. In this paper, we focus on turbulence in magnetized plasmas encountered in controlled nuclear fusion. The strong gradients in temperature, energy and density, along with the intrinsic nonlinearities of the plasma dynamics and the geometric complexity of fusion devices, give rise to instabilities resulting in turbulent regimes. This represents a major challenge in the development of controlled nuclear fusion because turbulent transport can lead to anomalous losses of energy and particles [4].

Magnetized plasma turbulence exhibits non-linear effects and multi-scale interactions, which require numerically demanding simulations. This has led to the development of a variety of numerical tools and physics models of different levels of complexity ranging from fluid [29] to gyro-kinetic descriptions [43, 25, 27, 9, 39, 37]. Simulations of magnetized plasmas are essential to understand and predict the turbulence phenomena and their effect on the performance of a reactor. To study the effect of electromagnetic turbulence on different particles in a fusion device, a standard procedure consists of following the trajectories of passive particles in the presence of the turbulent electromagnetic field. However, to do this one needs to have already computed the turbulence data over the integration time. For this kind of studies to be relevant, passive particle trajectories must be integrated over characteristic times much larger than what is the usual time over which turbulence models are solved, since the
asymptotic nature of transport relies on rare events and intermittency phenomena [17, 11]. Due to the limitations of computational resources, this can limit the accuracy of turbulent transport assessment. Furthermore, the time step and the grid resolution required to solve a turbulence model might need to be finer than those required to solve passive particle trajectories. In this paper, we propose a framework based on Artificial Intelligence (AI) to generate turbulence data at a much lower cost, to overcome the aforementioned numerical limitations of plasma simulations.

In recent years, considerable research efforts in plasma physics and nuclear fusion have been devoted to devel-
oping or using different AI techniques depending on the application. For example, in the framework of prediction
of disruptions in fusion devices, several methods have been employed, such as probabilistic support vector ma-
chines [42], recurrent neural networks [34], random forests [46, 30], deep convolution neural networks [2] and
multiscale convolution layers [53]. In the context of real-time prediction of turbulent fluxes, standard multilayer
perceptron neural networks have been explored [14, 13]. In addition, recent investigations have been conducted to
accelerate plasma turbulence simulations [28]. AI has also been used for the modeling of the hot tail of runaway
electrons in tokamaks. Examples of this application are the use of the physics-informed neural network approach
[41] and the development of pseudo-reversible normalizing flows [50]. To help in the automatic detection of
instabilities in fusion plasmas, segmentation and computer vision techniques have been applied to spectrograms
[8, 16] together with advanced GPU-accelerated Mirnov coil signal processing and unsupervised machine learning
[52]. Among the different types of AI techniques used or developed recently, generative AI has attracted much
interest due to its ability to produce new and original data based on patterns learned from existing data. This
can include text, images, audio, and even video. Unlike traditional AI, which focuses primarily on analyzing and
classifying data, generative AI creates new data by detecting the underlying distributions of the data on which it
has been trained. Notable examples include deep learning architectures such as generative adversarial networks
(GANs) [26], variational autoencoders (VAEs) [35], and transformer-based models such as the Generative Pre-
trained Transformer (GPT) [51]. One of the first and remarkable achievements of generative AI was the design of
new proteins [21]. In nuclear fusion, VAEs have, for example, been applied very recently to predict the impact of
the size of the tokamak on the values of the electron density and temperature at the pedestal [31] or to propagate in
time the electron radial profiles in the AUG tokamak [36]. The anomaly detection capabilities of auto-encoders can
also be used to identify instabilities and help in real-time fusion plasma control [3]. When considering turbulence,
one must take into account the high dimensionality due to the wide range of spatio-temporal scales. It is well
known that architectures such as convolutional networks can process multidimensional data efficiently, whereas
others such as recurrent neural networks are more suitable for processing and predicting time series [40]. There-
fore, intuitively, combining convolutional and recurrent neural networks could be appropriate for the generation
of synthetic turbulence [22]. However, modern generative AI methods (GAN, VAE, diffusion models) [19, 20]
can over-perform standard AI techniques in this kind of tasks, with different fidelity and acceleration capabilities.

The AI framework that we propose is based on our recent work reported in Ref.[15] where the GAIT model was introduced. Here, we delve deeper into the analyses and provide a comprehensive and systematic examination of the implementation of the proposed framework. The GAIT model consists of the coupling of two machine learning models: a convolutional variational autoencoder (CVAE), and a recurrent neural network (RNN). Although the GAIT model can in principle be applied to any turbulence problem, the purpose of this paper is to demonstrate its applicability to 2D plasma turbulence and to give a detailed explanation of the implementation of the model. In addition, we provide a systematic evaluation of its performance by analyzing relevant physical and statistical quantities. We investigate the role of the hyper-parameters and compare different recurrent neural network architectures, showing that our model can generate turbulence data with a great reduction in computational cost.

The paper is structured as follows. In section 1, we introduce the turbulence physical model that will be used. Section 2 provides a description of the CVAE, whereas the RNN is described in section 3. In section 4, we assess the performance of the model to generate turbulence data using both Eulerian and Lagrangian metrics. Section 5 is devoted to the analysis of the reduced representation of turbulence by the CVAE. The influence of the hyper-parameters is studied in section 6. We present and evaluate in section 7 some alternative architectures for the time sequence generation. Conclusions and future directions are given in section 7.

1 The Hasegawa-Wakatani model for drift wave turbulence

The Hasegawa-Wakatani (HW) equations [29] are used to produce a series of snapshots of 2D plasma turbulence that will form the training dataset for our AI framework. The HW equations represent a simplified model of edge magnetized plasma turbulence, specifically drift wave turbulence. We focus on the outer edge of a poloidal cross section, and use a Cartesian slab approximation of the toroidal geometry with the *x*-axis in the radial direction,
<figure><img src="image_2.png" /><figcaption>Figure 1: Results from HW numerical simulation: a snapshot of turbulent Hasegawa-Wakatani potential (left) and time evolution of the energy, showing the growth and saturation phases (right).</figcaption></figure>

the y-axis in the poloidal direction, and the z-axis along the toroidal direction. The magnetic field is assumed to be stationary, uniform and in the toroidal direction, i.e. $\mathbf{B} = B\mathbf{e}_z$. The plasma density decreases towards the edge, in the x-direction, and this gradient drives the drift wave instability. The HW model describes the evolution of the electrostatic potential and the density perturbation:

$$ \partial_t n + [\phi, n] = C(\phi - n) - \kappa \partial_x \phi + \mu \nabla^2 n \quad (1) $$

$$ \partial_t \Omega + [\phi, \Omega] = C(\phi - n) + \mu \nabla^2 \Omega \quad (2) $$

where $\phi$ is the electrostatic potential, $n$ the density perturbation, $\Omega = \nabla^2 \phi$ the vorticity, $C$ the adiabaticity coefficient, $\mu$ the diffusion coefficient, $[A_1, A_2] = \partial_x A \partial_x B - \partial_y A \partial_y B$ the canonical Poisson bracket, and $\kappa = -\partial_x \log n_0$ the drive of the instability, with $n_0$ the background density. Time and space are normalized to the reference cyclotron frequency $\omega_{c0}$ and to the reference Larmor radius $\rho_0$, respectively.

A standard pseudo-spectral method [1] was used for the numerical solution of Eqs.(1)-(2). At each time iteration, the Poisson equation, $\Omega = \nabla^2 \phi$, is solved for the potential in Fourier space, and the HW equations are pushed forward in time using a 4-th order Runge-Kutta scheme. To guarantee convergence and stability, a fine grid and sufficiently small time step was used. However, as explained below, the GAIT model implementation can work with a time step and a grid size coarser than the ones used in the direct numerical solution, and does not need the training density data to generate the potential.

In the present paper the numerical solution of Eqs.(1)-(2) assume $C = 1$, $\kappa = 1$, $\mu = 10^{-3}$. We use a reference wavenumber of $k_0 = 0.15$, meaning the domain size is $2\pi/k_0 \approx 41.9\rho_0$, with a grid resolution of $512 \times 512$ and a timestep of $\Delta t = 2 \cdot 10^{-2}/\omega_{c0}$. These parameters are within the standard range of values, see for example [45, 1], and the high resolution is the minimum required to provide a ground truth for the training of the AI model [28]. The boundary conditions are double periodic in $x$ and $y$.

The simulations are done for $5 \cdot 10^5$ time steps, saving the potential every 50 time steps, i.e. $\Delta t_{saving} = \omega_{c0}^{-1}$, and thus providing training dataset contains 10000 snapshots. The high-resolution $512 \times 512$ grid is down-sampled to a coarse grained $64 \times 64$ grid by averaging neighboring pixels.

The left panel of Fig. 1 shows a snapshot of the turbulent potential $\phi$ obtained after integrating Eqs. 1 and 2 as well as the evolution of the energy $\tilde{E}_p = \langle \frac{1}{2}(\nabla\phi)^2 \rangle_{cell}$, where $\langle \cdot \rangle_{cell}$ indicates the mean over the periodic domain.

To train the GAIT model, we draw a subset of $N_s = 8000$ snapshots from the dataset. As observed on the right panel of Fig. 1, the growth phase of the turbulence ends around $\omega_{c0t} \approx 200$. Beyond this time, the turbulence has reached a saturated state and becomes statistically stationary. To avoid transient effects, the 8000 snapshots used to train the GAIT model correspond to the time window $\omega_{c0t} \in [1000, 9000]$.

## 2 GAIT model - Convolutional variational auto-encoder

As illustrated in Fig. 2, the GAIT model is a combination of a convolutional variational auto-encoder (CVAE) and a recurrent neural network (RNN). The CVAE is designed to create a reduced representation of the turbulent potential, and the RNN extrapolates the evolution of this reduced representation in the latent space. In this section we provide a detailed explanation of the CVAE, together with the description of the training.
<figure><img src="image_3.png" /><figcaption>Figure 2: Schematic representation of the architecture of the proposed generative machine learning framework to construct surrogate models of plasma turbulence.</figcaption></figure>

## 2.1 Convolutional Variational Auto-encoder

Convolutional Variational Autoencoders (CVAEs) are a type of generative model that combines the Variational Autoencoder (VAE) framework with convolutional neural networks, making them especially well suited for handling images [35, 12]. VAEs are designed to learn a compressed and probabilistic representation of data that can be used to generate new samples. Any auto-encoder consists of two parts: the encoder $\mathcal{E}$ and the decoder $\mathcal{D}$. The encoder maps the input data (a turbulent snapshot $\phi$ in our case) into a lower-dimensional space, called the latent space. The decoder attempts to reconstruct the original input from the latent space. Therefore, an auto-encoder is designed and trained to approximate the identity on the training data, i.e. $\mathcal{D}[\mathcal{E}[\phi]] = \hat{\phi} \approx \phi$. For the GAIT model, the dimension of the latent space used is $N = 64$, which means that the images of size $64 \times 64$ representing snapshots of the potential $\phi$ are reduced by a factor of 64, and then reconstructed.

The term *variational* refers to the use of variational inference to approximate complex probability distributions. In essence, a snapshot from the turbulence dataset is not mapped to a fixed point in the latent space, but to the parameters of a distribution. This distribution is afterwards sampled to generate a point in the latent space, which is subsequently mapped by the decoder into the original space to approximate the input. This stochastic process forces the latent space to encode meaningful representations of the turbulence data in a smooth and continuous way, which facilitates interpolation between data points and enables the generation of new samples, which is a requirement for the GAIT model. For the distribution, we use a Gaussian distribution, like most of the VAE implementations: $N(z|\mu, \sigma)$, with $z \in \mathbb{R}^N$, $\mu \in \mathbb{R}^N$ and $\sigma \in \mathbb{R}^N$.

The term *convolutional* indicates that convolutional neural networks (CNNs) are used in the encoder and decoder. Those networks are particularly good at processing grid-like data, such as 2D or 3D images, thanks to spatial convolution operations that can efficiently encode and decode high-dimensional image data while preserving important spatial structures. A CNN is a stack of successive convolution layers, each applying one convolution operation to the data. A convolution layer has three main parameters: the kernel size, the stride and the number of filters. The kernel size refers to the dimensions of the filter used to scan the input, with smaller kernels focusing on local details and larger ones capturing broader patterns. The stride determines how far the filter moves across the input; a larger stride reduces the feature map size and speeds up computation but may miss details, while a smaller stride provides finer detail at a higher computational cost. The number of filters (or the number of kernels) defines how many patterns the network can learn, with more filters allowing the network to capture a wider range of features in the input data. In our case, the convolution operations need to be double periodic, both in the encoder and in the decoder. They need to be periodic in the encoder so that the model can recognize a pattern that overlaps with the boundaries of the domain. They also need to be periodic in the decoder so that all images generated by the decoder respect the periodic boundary conditions. The convolution layers available in standard libraries, such as TensorFlow, do not offer this feature. For this reason, we had to implement customized convolution layers in the GAIT model. To do so, before applying a convolution operation, we pad the data on the boundaries with values from the opposite boundaries, so that the convolution kernel at the border captures the data from the other side. The size of the padding must correspond to the kernel size of the operation. This must be properly implemented so that those operations can be differentiated by the machine learning algorithm during the training process.
The architectures of the encoder $\mathcal{E}$ and the decoder $\mathcal{D}$ are presented in Tables 1 and 2, respectively. Convolution layers allow the architecture of the model to remain relatively small compared to the size of the data. Here, the 8000 images of size 64 × 64 represent approximately 3 × 10⁷ data points, while the CVAE has around 3 × 10⁶ trainable parameters. For the deep convolutional layers, we employ tanh activation functions, whereas linear activation functions are used to calculate $\mu$ and $\log(\sigma)$ so that they are not restricted to a range of values. For the same reason, the last transposed convolution layer also has a linear activation function.

<table><caption>Table 1: Architecture of the encoder</caption><thead><tr><th>Layer - Operation</th><th>Activation function</th><th>Output shape</th></tr></thead><tbody><tr><td>Input</td><td>-</td><td>(64,64)</td></tr><tr><td>Convolution, 32 filters</td><td>tanh</td><td>(32,32,32)</td></tr><tr><td>Convolution, 32 filters</td><td>tanh</td><td>(16,16,32)</td></tr><tr><td>Convolution, 64 filters</td><td>tanh</td><td>(8,8,64)</td></tr><tr><td>Convolution, 128 filters</td><td>tanh</td><td>(4,4,128)</td></tr><tr><td>Convolution, 512 filters</td><td>tanh</td><td>(1,1,512)</td></tr><tr><td>μ Dense layer, 64 units</td><td>linear</td><td>(64)</td></tr><tr><td>log(σ) Dense layer, 64 units</td><td>linear</td><td>(64)</td></tr><tr><td>Sampling layer ~ N(z|μ, σ)</td><td>-</td><td>(64)</td></tr></tbody></table>

<table><caption>Table 2: Architecture of the decoder</caption><thead><tr><th>Layer - Operation</th><th>Activation function</th><th>Output shape</th></tr></thead><tbody><tr><td>Input</td><td>-</td><td>(64)</td></tr><tr><td>Dense layer, 512 units</td><td>tanh</td><td>(1,1,512)</td></tr><tr><td>Transposed convolution, 128 filters</td><td>tanh</td><td>(4,4,128)</td></tr><tr><td>Transposed convolution, 64 filters</td><td>tanh</td><td>(8,8,64)</td></tr><tr><td>Transposed convolution, 32 filters</td><td>tanh</td><td>(16,16,32)</td></tr><tr><td>Transposed convolution, 32 filters</td><td>tanh</td><td>(32,32,32)</td></tr><tr><td>Transposed convolution, 1 filter</td><td>linear</td><td>(64,64,1)</td></tr><tr><td>Data output</td><td>-</td><td>(64,64)</td></tr></tbody></table>

## 2.2 Training of the CVAE

Once the architecture of the CVAE has been decided, the weights used in the convolutions and other operations performed by the CVAE need to be adjusted in order to obtain a good compression in latent space by the encoder $\mathcal{E}$ and satisfactory reconstruction by the decoder $\mathcal{D}$. As in any other machine learning problem, this translates into the minimization of a loss function, which for the GAIT model is

$$L_{\text{CVAE}} = \sum_{i}^{N_s} w_{\phi} \|\phi_i - \hat{\phi}_i\|^2 + w_{\nabla\phi} \|\nabla\phi_i - \nabla\hat{\phi}_i\|^2 \\ + w_{\text{KL}} \text{KL}(\mathcal{N}(\mathbf{z}_i|\mu_i, \sigma_i), \mathcal{N}(\mathbf{z}_i|\mathbf{0}, \mathbf{1})) \quad (3)$$

where $\{\phi_i\}_{i=1}^{N_s}$ and $\{\hat{\phi}_i\}_{i=1}^{N_s}$ are the $N_s$ training and the $N_s$ reconstructed snapshots respectively, $\mathbf{z}_i = \mathcal{E}(\phi_i)$, with $\mathcal{E}$ the encoder, $\text{KL}$ is the Kullback-Leibler divergence [38] between the Gaussian distribution $\mathcal{N}(\mathbf{z}|\mu, \sigma)$ and a standard Gaussian $\mathcal{N}(\mathbf{z}|\mathbf{0}, \mathbf{1})$, and $\hat{\phi} = \mathcal{D}(\mathcal{E}(\phi))$ is the reconstructed potential, with $\mathcal{D}$ the decoder.

The loss function (3) is a weighted sum of three terms. The first two terms, with their respective weights ($w_{\phi}, w_{\nabla\phi}$), penalize reconstruction errors. Together, they form the reconstruction loss. The loss functions usually employed in VAEs include only the direct reconstruction error, here the one associated with $w_{\phi}$. In our case, we added the second term, weighted by $w_{\nabla\phi}$, that compares the gradients of the input and the output of the CVAE. This term is important to make the training more efficient, since it directly targets the reconstruction of the electric field $\mathbf{E} = -\nabla\phi$, which actually plays a more important role in particle transport than the potential itself and emphasizes the conservation of energy through the autoencoder. In the present work, the gradients are computed with a 4-th order method. This applies to the computation of the loss function, the computation of the potential energy, and the computation of the drift velocity to test the GAIT model. The third term in Eq.(3) uses the Kullback-Leibler divergence to quantify the difference between $\mathcal{N}(\mathbf{z}|\mu, \sigma)$ and a standard Gaussian $\mathcal{N}(\mathbf{z}|\mathbf{0}, \mathbf{1})$. To minimize this term, the model needs to compress the information efficiently and distribute similar inputs close together in the
<figure><img src="image_4.png" /><figcaption>Figure 3: Training loss of the convolutional variational auto-encoder, epochs 100 to 5000, ($w_{\phi} = 1, w_{\nabla\phi} = 10, w_{KL} = 0.01$). Left: evolution of the total loss during the training. Right: competition between the three terms of the loss.</figcaption></figure>

latent space. As a result, new points sampled from this space will also produce meaningful outputs, even if they are not part of the training set.

The weights ($w_{\phi}, w_{\nabla\phi}, w_{\text{KL}}$) are free parameters of the GAIT model. The selection of these parameters is based on general physical intuition and does not require fine tuning. As a general rule, the KL divergence loss cannot dominate over the reconstruction loss. Otherwise, the reconstruction would not be good enough. It is therefore natural to try weights satisfying $w_{\phi} < w_{\text{KL}}$ or $w_{\nabla\phi} < w_{\text{KL}}$. Starting from this principle, we adjust the weights to target the features of the turbulence we want to replicate. Increasing $w_{\phi}$ or $w_{\nabla\phi}$ will lead to a good reconstruction of the potential or the electric field, respectively, whereas increasing $w_{\text{KL}}$ means enforcing more regularity in the encoding process, allowing a smoother evolution of the output when moving in the latent space. In this work, we present two study cases: *case 1* with ($w_{\phi} = 1, w_{\nabla\phi} = 10, w_{\text{KL}} = 0.01$), is designed to obtain a good reconstruction of the potential and its gradient. The weight $w_{\text{KL}}$ is kept low to allow the potential to take space in the latent space to encode information. *Case 2*, with ($w_{\phi} = 0, w_{\nabla\phi} = 10, w_{\text{KL}} = 0.1$), focuses on the regularity in latent space and the reconstruction of the velocity field only. As we will see, *case 1* allows a good Eulerian agreement between the training data and the output of the GAIT model, whereas *case 2* helps recover the Lagrangian properties of passive tracers using the velocity field $V = (-\partial_y\phi, \partial_x\phi)$.

Owing to the large number of trainable parameters, we used the standard Adam optimizer. The CVAE was trained on 5000 epochs, with an initial learning rate of $l_r = 10^{-3}$ and a batch size of 1000. The evolution of the loss function during training for *textitcase 1* is shown in Fig. 3. We can see that after 1000 epochs approximately, the different terms of the loss function are in competition. The minimization of the reconstruction error implies an increase of the Kullback-Leibler divergence. This means that in order for the potential to be well reconstructed, the latent vectors need to occupy a larger region of the latent space. After the minimization of the loss function during the training, the reconstruction capability of the CVAE can be tested by looking at snapshots from the potential. Three snapshots and their respective reconstructions are shown in Fig. 4. From left to right are represented: the first snapshot of the training dataset, a snapshot in the middle of the dataset, and the last snapshot of the dataset. We see that the CVAE has learned to reconstruct the whole simulation.

## 3 GAIT model - Recurrent neural network

After the CVAE training, the data is compressed into a latent space containing continuous and meaningful representations of the turbulence states. It is therefore possible to find a time sequence that, once decoded, provides a realistic sequence of 2D images of HW turbulence. For this purpose, the resulting time sequence needs to have the same dynamics as the HW simulation that served as training data. This is the purpose of the second component of the GAIT model, consisting of a neural network that operates in the latent space. Unlike the auto-encoder, here we train the model not to reproduce the input but to predict the evolution of the latent state of the system. For example, the input of the model can be a single latent vector $z_i$, and the model is trained to output $z_{i+1}$. Using such a method means that we assume that knowing the current latent state is enough to predict the next one. In this case, a simple deep neural network with a few hidden dense layers can be trained to do so. But we could also want to use more information to predict the next latent state by looking at the history of its evolution. The input of the model can be a sequence of latent vectors ($z_i, ..., z_{i+1}$), with $l$ representing the input sequence length as a new hyper-parameter of the model. The model can then be trained to either output the next latent state $z_{i+1}$, or the whole sequence advanced by one time step ($z_{i+1}, ..., z_{i+1}$). If we have sequences as inputs, the data becomes too high-dimensional for the use of a basic DNN with dense layers. In this case, we use architectures designed to process sequential data. Recurrent neural networks (RNN) or long-short-term memory networks (LSTM) are
<figure><img src="image_5.png" /><figcaption>Figure 4: Reconstruction tests of the convolutional variational auto-encoder after training, ($w_{\phi} = 1, w_{\nabla\phi} = 10, w_{KL} = 0.01$). Top row: original data. Bottom row: data reconstructed by the auto-encoder.</figcaption></figure>

network architectures that allow to process sequential data without using too many parameters. They are designed to capture time dependencies in the data and can be used to extrapolate time series.

The GAIT model uses a RNN architecture with two hidden recurrent layers and an input/output sequence of length $l = 50$ (see Table 3). The numerical results show that the latent vectors from the training data are all located in a relatively small region of the latent space. Their norms are all around 1.16, and do not exceed 1.2. This is a consequence of the Kullback-Leibler divergence, which penalizes remote vectors in latent space. Moreover, the norm of a latent vector is related to the energy of the potential after decoding. Therefore, to ensure that the electrostatic energy remains stable, the RNN uses bounded tanh activation functions that prevent irregularities during the generation process, whereas a linear activation function is used in the dense output layer. A more in-depth analysis of the latent space is provided in Sec. 5. We also investigate the other networks discussed above, such as LSTM networks, or the most simple approach with just one latent vector as input and one latent vector as output. The performances of those different architectures as well as the influence of the sequence length $l$ are discussed in Sec. 7.

<table><caption>Table 3: Architecture of the recurrent neural network</caption><thead><tr><th>Layer - Operation</th><th>Activation function</th><th>Output shape</th></tr></thead><tbody><tr><td>Inputs</td><td>-</td><td>(50,64)</td></tr><tr><td>Recurrent layer, 128 units</td><td>tanh</td><td>(50,128)</td></tr><tr><td>Recurrent layer, 128 units</td><td>tanh</td><td>(50,128)</td></tr><tr><td>Dense layer, 64 units</td><td>linear</td><td>(50,64)</td></tr></tbody></table>

To train the RNN, the sequence $\{\mathbf{z}_1, \mathbf{z}_2, \dots, \mathbf{z}_N\}$, corresponding to the time evolution of the turbulence snapshots $\{\phi_1(x, y), \phi_2(x, y), \dots, \phi_N(x, y)\}$, is split into training and testing sequences containing 85% and 15% of the data, respectively. The RNN is trained for 500 epochs with an Adam optimizer, an initial learning rate of $l_r = 10^{-3}$, and a batch size of 500. The model is trained to push forward in time whole sequences of latent vectors using the loss function

$$L_{\text{RNN}} = \sum_{i} \| (\mathbf{z}_{i+1}, \dots, \mathbf{z}_{i+l+1}) - \text{RNN}(\mathbf{z}_i, \dots, \mathbf{z}_{i+l}) \|^2 \quad (4)$$

Penalizing the prediction of the whole sequence and not only the next latent state forces the model to assimilate all the temporal dynamics occurring in the time-span of the training sequence $\omega_{c0}\Delta t_{seq} = 50$. This is enough to capture almost all dependencies, since the autocorrelation time of the HW potential is $\tau_{AC,HW} \simeq 4$.

Fig.5 displays two snapshots of turbulent potential: one generated using the GAIT model and one from the direct numerical simulation of the HW turbulence model in Eq. (1)-(2). By visual inspection, it is impossible to
<figure><img src="image_6.png" /><figcaption>Figure 5: Snapshots of the electrostatic potential obtained from the direct numerical integration of the Hasegawa-Wakatani model (left) and using the proposed GAIT (Generative Artificial Intelligence Turbulence) model (right).</figcaption></figure>

tell which one is which. From the physics perspective however, such observations are not sufficient to assess the performance of the GAIT model and more quantitative analysis is required.

# 4 Evaluation of the GAIT model

The assessment of the generation of turbulence by the GAIT model is done through two metrics: the Eulerian metrics provide comparison of quantities derived from the turbulent field, whereas the Lagrangian metrics compare passive particle statistics. Owing to the periodicity of the turbulence, the Eulerian analyses are done over an entire periodic cell. To do those comparisons, we generated a long sequence of turbulent states using the GAIT model. The training dataset of HW turbulence has 8000 frames and we generated 100,000 frames with the GAIT model. For the following analyses, such a long simulation is not necessary, a few thousand frames are well enough to compute the significant metrics of 2D steady-state turbulence, given that the auto-correlation time corresponds to 5 frames only. However, it is to be noted that with such a generative model, there is a risk that it starts deteriorating or reaches a fixed point after long times. The long generation enabled us to verify that the model maintains the same characteristics over much longer time (here, one order of magnitude). Furthermore, having long-time simulations of turbulence is necessary to study of the asymptotic properties of transport. There are systems in which transport is determined by rare events, such as particle trappings and/or long displacements. In plasma turbulence, vortices can induce trapping, and zonal flows can lead to rare long displacements [32, 33]. To have a precise understanding of turbulent transport in those cases, long turbulent simulations are required. In some cases, turbulent transport can transition from super-diffusive to sub-diffusive in the $t \to \infty$ limit [11]. The cheap generation of long turbulent simulations then becomes a useful tool to carry out transport studies.

## 4.1 Eulerian comparison

To test the Eulerian statistics of the GAIT model we consider the Fourier spectra of the electric field fluctuations, both in space and time, the spectrum of the proper orthogonal decomposition (POD) of the turbulent potential, the auto-correlation decay time and the Okubo-Weiss parameter. These metrics were computed over a 5000-frame window for both the DNS and the GAIT-generated potential. For the DNS we took the first 5000 frames out of the 8000-frame sequence. For the GAIT generated data, to verify its consistency over time, we looked at several time windows. As shown in the following analysis, the metrics remain consistent throughout the long generated sequence, showing that the model preserves the turbulent characteristics of Hasegawa-Wakatani system.

Comparing the spatial Fourier spectrum of the two models (the HW model and the surrogate GAIT model) allows to verify if the new samples generated by GAIT model have the same statistical properties in the spatial domain as the HW model. On the other hand, comparison of the time spectrum help verify that the dynamic properties of the turbulence are correctly reproduced. The Fourier comparison is done over the electric field turbulent fluctuations : $\delta E(x,y,t) = E(x,y,t) - \langle E \rangle$, where $\langle f \rangle$ denotes spatial average. The first test compares the Fourier power spectra in space of $\delta E(x,y,t)$. The left panel of Figure 6 displays the comparison of the spatial spectrum $|\hat{\delta E}(k)|$, where $\hat{\delta E}(k,t)$ is the Fourier transform in space of $\delta E(x,y,t) = E(x,y,t) - \langle E \rangle$, $\bar{f}$ denotes time average, $E$ is the magnitude of the electric field, $E = | - \nabla \phi |$, and $k = \sqrt{k_x^2 + k_y^2}$. The Fourier power spectra in time $\langle |\tilde{\delta E}(\omega)| \rangle$, where $\tilde{\delta E}(x,y,\omega)$ is the Fourier transform in time of $\delta E(x,y,t)$, is also displayed on the right panel of Figure 6. The spectra comparisons display a good agreement between the original turbulence and the output of the GAIT model.
<figure><img src="image_7.png" /><figcaption>Figure 6: Comparison of electric field turbulent fluctuations spectra in the HW and GAIT models. Left: Fourier power spectra in space, $|\hat{\delta E}(k)|$. Right: Fourier power spectra in time, $\langle |\tilde{\delta E}(\omega)| \rangle$.</figcaption></figure>

<figure><img src="image_8.png" /><figcaption>Figure 7: Comparison of electrostatic potential fluctuations in the HW and GAIT models. (Left) Proper orthogonal decomposition spectrum. (Right) Eulerian time auto-correlation function.</figcaption></figure>

<figure><img src="image_9.png" /><figcaption>Figure 8: Flow topology regions with the Okubo-Weiss parameter $Q(x, y, t_c)$, in the HW model (left) and the GAIT model (right).</figcaption></figure>

<figure><img src="image_10.png" /><figcaption>Figure 9: Comparison of probability distribution of Okubo-Weiss parameter $Q$ values in space and time.</figcaption></figure>
Although the agreement of Fourier spectra is very good, it does not necessarily account for the formation and preservation of coherent structures, such as vortices, which are key aspects of turbulence [47], and can be responsible for trapping effects that play an important role in particle transport [18]. For this purpose, we use the Proper Orthogonal Decomposition (POD) of the potential. This method allows to identify empirical modes that are proper to the studied data, and has been extensively used for fluid and plasma turbulent flows [5, 6, 23]. The POD consists in decomposing the data into orthogonal modes as follows: $\phi(x,y,t) = \sum_{\ell=1}^{N} w_{\ell} u_{\ell}(x,y)v_{\ell}(t)$, where $u_{\ell}$ are spatial modes and $v_{\ell}$ are temporal modes, forming two families of orthogonal modes. The coefficients $w_{\ell}$ form the POD spectrum. These values indicate the amount of the variance of the data that can be explained with a finite number of spatio-temporal modes, hence providing a hierarchy of structures forming the turbulence. Fig. 7 shows excellent agreement of the POD spectrum values $w_{\ell}$ as a function of the rank $\ell$, especially for case 1, indicating that the HW and the GAIT models have the same hierarchy of spatiotemporal modes.

For the GAIT model to be used as a surrogate model of turbulence, the Eulerian auto-correlation decay time of the potential must be correctly reproduced, for it represents an important feature of the temporal dynamics of turbulence. It is defined as the time it takes for the autocorrelation function $\mathcal{C}(\tau) = \langle\phi(x,t)\phi(x,t+\tau)\rangle$ to decrease under $1/e$. The right panel of Fig. 7 shows the good agreement in this metric especially for case 1. The Okubo-Weiss parameter [33], which is designed to characterize the flow topology of a turbulent flow field, is employed as another Eulerian metric, since it serves as a tool to analyze the dominance of vorticity or deformation. In our case, the turbulent flow is given by the drift velocity $\mathbf{V} = -\nabla\phi \times \mathbf{e}_z = (V_x, V_y) = (-\partial_y\phi, \partial_x\phi)$. The Okubo-Weiss parameter is defined as $Q = s^2 - \Omega^2$, where $\Omega = \nabla^2_{\perp}\phi = \partial_x V_y - \partial_y V_x$ is the vorticity and $s^2 = s_1^2 + s_2^2$ is the deformation, with $s_1 = \partial_x V_x - \partial_y V_y$ and $s_2 = \partial_x V_y + \partial_y V_x$. After computing the threshold value $Q_0 = \sqrt{\langle Q^2 \rangle}$, we can identify three regions. The region $Q \le -Q_0$, vorticity dominated, is strongly elliptic. In this region transport tends to be reduced due to the trapping effect of vortices. On the other hand, the region $Q \ge Q_0$ is deformation dominated, strongly hyperbolic. In this region, mixing and dispersion are enhanced. In between these two regions lies the intermediate region $-Q_0 < Q < Q_0$. We can observe very good agreement on Fig. 8. On the left, we see a snapshot of the flow regions at a fixed time for the HW model and the GAIT model. The threshold is derived from the HW potential. They exhibit similar structures, with large vorticity dominated regions surrounded by deformation dominated regions. For a quantitative comparison, we plot the PDF of all $Q$ values in space on Fig. 9, showing excellent agreement for case 1. Case 2 displays a qualitative agreement, the shape of the PDF being the same, but compressed in the $x$-direction of the graph.

## 4.2 Lagrangian comparison

In this subsection we perform statistical analyses of the trajectories of passive tracers in the turbulent potential and compare the results between HW and GAIT models. To ensure that the transport does not depend on the initial positions where particles are initialized, they are uniformly distributed over the periodic domain $\{(x,y)|0 < x < 2\pi/k_0, 0 < y < 2\pi/k_0\}$. The equation of motion $\frac{dx}{dt} = -\nabla\phi \times \mathbf{e}_z$ is solved for each particle using a 4-th order Runge-Kutta scheme, a timestep of $\Delta t = 0.2\omega_{c0}^{-1}$, linear interpolation in time, and bilinear interpolation in space for the potential. A total of $10^4$ trajectories are integrated in both HW and GAIT turbulent potentials. The time-span is $t_{max,HW} = 8000\omega_{c0}^{-1}$ for the HW turbulence, and the long generated sequence of the GAIT model allowed to extend the integration to $t_{max,GAIT} = 100000\omega_{c0}^{-1}$. The periodicity of the potential, ensured by the implementation of periodic convolution layers in the decoder, allows to integrate the orbits in $\mathbb{R}^2$, and the values of $(x_i(t), y_i(t))$ were not wrapped in the $[0, 2\pi/k_0] \times [0, 2\pi/k_0]$ domain. The main quantities of interest are the probability distribution function of the displacement $P(R)$ and the evolution of the mean squared displacement $\langle R^2(t) \rangle$ where $R = \|\delta\vec{r} - \langle\delta\vec{r}\rangle\|$, $\delta\vec{r}$ is the individual displacement vector and $\langle\cdot\rangle$ is an ensemble average. Fig. 10 shows that both the HW and the GAIT models exhibit the expected diffusive scaling $\langle R^2(t) \rangle \sim t$ and Gaussian profiles of $P(R)$. As seen in the left panel of Fig. 10, with case 1 the GAIT model results in higher diffusivity ($D_{GAIT,1} = 1.57$) than that provided by the HW model ($D_{HW} = 1.311$). On the other hand, case 2 shows a good agreement with $D_{GAIT,2} = 1.243$, and the PDFs of $R$, plotted on the right panel of 10. This can be explained by the larger $w_{KL}$ in the case 2 setting, forcing the regularity in latent space and hence reducing the noise in the dynamics of the generated turbulence.

## 4.3 Computational efficiency

To compare the efficiency of the GAIT model, all computations and training were done with the same computational resources. The numerical integration of the HW model to produce 10000 snapshots of training data takes 50 hours on a V100 GPU node of the French Jean-Zay super computer. Training the CVAE for 5000 epochs requires 1 hour, and the training of the RNN can be done in about five minutes. Once the model is trained, generating 10000 snapshots of turbulent states with the GAIT model takes about 7.5 min, showing that the GAIT model can
<figure><img src="image_11.png" /><figcaption>Figure 10: Comparison of displacements, *R*, of tracers in HW and GAIT models. Left: Mean squared displacement $\langle R^2(t) \rangle$ as function of *t* in log-log scale. Right: Probability distribution function of displacements at different times.</figcaption></figure>

generate turbulent states at a rate 50 hrs/ 7.5 min ~ 400 times faster than the direct numerical integration. In addition, the GAIT model does not generate 512 × 512 snapshots with a time step of $\Delta t = 2 \cdot 10^{-2} \omega_0^{-1}$, i.e. it is able to overcome the resolution of all spatio-temporal scales without numerical stability problems and loss of information in the down-sampling in space and time is minor, since all the important observables happens at larger scales.

# 5 Continuity and energy studies in the latent space

After training the CVAE, it is important to assess the quality of the coding of the HW turbulence. The latent space is 64-dimensional, and it is not straightforward to understand the role of each dimension, but one can first look at how the different latent vectors are organized. On the left panel of Fig 11, it is plotted the PDFs of the latent vector norm $||z||$ for the encoded HW turbulence and for the GAIT-generated latent sequence. The norm we use in the latent space for our calculation is $||z|| = (\frac{1}{64} \sum_{i=1}^{64} z_i^2)^{\frac{1}{2}}$. The results are shown for *case 1* and *case 2*. In both cases, the norm of the latent vectors is very localized, implying that the vectors are located around a hypersphere. For instance, for the HW encoded turbulence with the *case 1* CVAE, the vectors are around a hypersphere of radius *r* = 1.158. In both cases, the mean norm of the vectors generated by the GAIT model is smaller than the mean of the training data, which indicates that the RNN does not capture completely the topology of the encoded HW turbulence. However, the difference remains relatively small: 1.158 vs 1.087 and 1.008 vs 0.919, corresponding to relative errors of 6% and 9%, respectively. Fig.13 shows that this drift has a negligible effect on the energy of the decoded turbulence. We can also observe the effect of the $w_{KL}$ weight on the output of the model. $\langle||z||\rangle_{case2} < \langle||z||\rangle_{case1}$ is a consequence of $w_{KL,case2} > w_{KL,case1}$. Also, as we can see on Fig 11, the PDF of the norm is narrower for *case 2*. Those are effects of the increased regularization of the latent space with this set of hyperparameters.

However, the study of the PDF of the norm fails to provide information about the existence of preferred directions in the set of latent vectors. Theoretically, one could compare the distribution of angles after converting the latent vectors in hyper-spherical coordinates, but in a 64-dimensional space, this would require to compare 63 distributions. With only $N_s = 8000$ samples, those would be too sparse to draw any rigorous conclusion. Another way to detect preferred directions in the latent space is to run a symmetry test. One repeatedly generates a random hyperplane $H_u = \{z|z \cdot u = 0\}$—where *u* is randomly generated on the unit sphere—in the latent space and checks if it splits the set of vectors in two equally populated sets $S_u^+ = \{z_i|z_i \cdot u > 0\}$ and $S_u^- = \{z_i|z_i \cdot u < 0\}$. The population difference is calculated for a large number of random hyperplanes, and statistics of those differences are derived. The box plots on the right panel of Fig. 11 show the average percentage of population difference $\langle(\#S_u^+ - \#S_u^-)/2N_s\rangle_u$ and its standard deviation for the HW encoded data and the GAIT generated data. We observe that the test on the GAIT data is not as satisfactory as the test on the encoded HW data. However, the asymmetry observed remains less than 10% in the worst case, showing that the turbulent states generated by the GAIT model do not get stuck in a specific direction of the latent space, but remain balanced.

In order for the turbulence generation process to be acceptable, we need to make sure that by moving continuously in time, there are no sudden jumps in the latent space. A continuity test can be done by looking at pairs of snapshots from the training dataset, and by comparing their distance in time - the number of timesteps that separate them - with their distance in the latent space. The result is shown in Fig. 12 and exhibits the continuity of the encoding process.

Further understanding of the importance of the norm of the latent vectors can be gained by randomly generating a large number of latent vectors using a Gaussian multivariate sampler. After decoding them, the corresponding
<figure><img src="image_12.png" /><figcaption>Figure 11: Statistical tests of the occupation of the latent space by the encoded data. Left: Distribution of norms. Right: Assymetry tests of the latent vectors</figcaption></figure>

<figure><img src="image_13.png" /><figcaption>Figure 12: Continuity test of the CVAE: distance in latent space as a function of the time difference in the simulation.</figcaption></figure>

energy of the resulting potential can be computed. Fig 13 shows the results of this test for *case 1* and *case 2*. As we can see from both panels, when a point is randomly generated in the latent space and subsequently decoded into a turbulent snapshot, there is a clear dependency of the energy on the norm, which is linear with weak variance at low norms and saturates with increased variance for $||z|| \gtrsim 1$. However, when we look at the points from the encoded data, we see that they are not on the curve obtained with random points, but under the curve for *case 1* and above it for *case 2*. This shows that the organization of the latent space operated by the auto-encoder during the training is complex and cannot be reduced to just a norm and symmetry analysis. The latent vectors of the training dataset cannot be imitated by just producing a random vector of similar norm.

<figure><img src="image_14.png" /><figcaption>Figure 13: Relation between norm in latent space and energy of the electrostatic potential, for the training dataset and for randomly generated vectors. Left: *case 1* - ($w_\phi = 1, w_{\nabla\phi} = 10, w_{KL} = 0.01$). Right: *case 2* - ($w_\phi = 0, w_{\nabla\phi} = 10, w_{KL} = 0.1$)</figcaption></figure>
<figure><img src="image_15.png" /><figcaption>Figure 14: Influence of the training dataset size $N_s$ on the spatial Fourier spectrum of the electric field fluctuations. Left: Fourier spectra comparison. Right: evolution of the error with the number of training samples.</figcaption></figure>

# 6 Discussion on the CVAE design and dependence on the size of the training dataset

The CVAE design is key to the whole model performance, as it processes all the turbulent potential snapshots and constructs a low-dimensional representation of the turbulent system. If the CVAE is too simple, it might not be able to reconstruct essential features of the turbulence. On the other hand, if it is too complex, it might overfit the data and lose its generalization properties. Like any machine learning model, implementing the CVAE, requires choosing the values of the hyper-parameters. For example, the values of $w_\phi$, $w_\nabla\phi$ and $w_{KL}$ can have strong influence on the output of the GAIT model.

Also, the training dataset may play a critical role in the performance of the model and, like any machine learning algorithm, its size $N_s$ is a key element and has been the object of a detailed study in this paper. For this purpose, the same model as presented in sections 2 and 3 with the case 1 set of parameters has been trained with different training dataset sizes, ranging from $N_s = 500$ to $N_s = 8000$, and has been used to generate new sequences of turbulence. For each Eulerian metric introduced in section 4, we compute the error as the percentage of difference in the $L^2$ norm

$$Err_{\overline{|\delta E_i|,i}} = \frac{\left( \int_0^{k_{max}} (|\overline{\delta E_i(k)}| - |\overline{\delta E_{HW}(k)})|^2 dk \right)^{\frac{1}{2}}}{\left( \int_0^{k_{max}} (|\overline{\delta E_{HW}(k)}|^2 dk \right)^{\frac{1}{2}}} \quad (5)$$

where $E_i$ is the electric field of the model that is tested, and $E_{HW}$ is the electric field of the HW simulation. We can see from Figs. 14-18 that there is a global improvement as the number of training samples grows. However, for the Fourier spectra of the electric field fluctuations (Figs 14-15), the error increases again above $N_s = 4000$. The increased error of the Fourier spectra going from $N_s = 4000$ to $N_s = 8000$ raises an interesting point in the elaboration of an AI model that uses dimensionality reduction, such as an auto-encoder in our case. Although it is commonly believed that more training data will improve the model, it is not always the case. For an auto-encoder, increasing the training data is a way to improve the generalization capability, and ensure that new samples will share the same features as the training samples. However, having too much data will start to deteriorate the quality of the reconstruction. Since the latent space is constrained with the KL divergence, there is a finite volume in which all the data has to be encoded. In this case, after a certain threshold, the encoded data will be too concentrated in the latent space, and some features of the samples will be overlooked by the model. This illustrates the complexity of the implementation of an AI model, especially when we want it to be computationally light and efficient: there is often not one optimal design, and a trade-off must be reached to provide a satisfactory output. We also see the need to have a variety of quantitative tools to assess the performance of a model, as limiting ourselves to only one metric could let us disregard potential flaws of the model. Overall, we decided to use as training dataset size $N_s = 8000$, which appears to be the best all around choice for the GAIT model. Note that several parameters can be adjusted in the implementation and training of the GAIT model. Some of these include the size of the training data set, the number and characteristics of convolution layers, the dimension of the latent space, and the loss function hyper-parameters. We have discussed the impact of some of these parameters on the performance of the GAIT model. However, a throughout study of all the dependencies and the optimization of the model's hyper-parameters is beyond the scope of the present paper.
<figure><img src="image_16.png" /></figure>

<figure><img src="image_17.png" /><figcaption>Figure 15: Influence of the training dataset size *N*<sub>s</sub> on the temporal Fourier spectrum of the electric field fluctuations. Left: Fourier spectra comparison. Right: evolution of the error with the number of training samples.</figcaption></figure>

<figure><img src="image_18.png" /></figure>

<figure><img src="image_19.png" /><figcaption>Figure 16: Influence of the training dataset size *N*<sub>s</sub> on the POD spectrum of the potential. Left: Spectra comparison. Right: evolution of the error with the number of training samples.</figcaption></figure>

<figure><img src="image_20.png" /></figure>

<figure><img src="image_21.png" /><figcaption>Figure 17: Influence of the training dataset size *N*<sub>s</sub> on the autocorrelation function of the potential. Left: Autocorrelation comparison. Right: evolution of the error with the number of training samples.</figcaption></figure>

<figure><img src="image_22.png" /></figure>

<figure><img src="image_23.png" /><figcaption>Figure 18: Influence of the training dataset size *N*<sub>s</sub> on the Okubo-Weiss PDF of the potential. Left: Okubo-Weiss PDF comparison. Right: evolution of the error with the number of training samples.</figcaption></figure>
<figure><img src="image_24.png" /><figcaption>Figure 19: Eulerian comparison of the turbulence obtained with different model architectures in latent space. Left: spatial spectrum. Middle: frequency spectrum. Right: auto-correlation functions.</figcaption></figure>

# 7 Variations of the RNN model

If the CVAE has been properly trained, one can expect to obtain good agreement in metrics such as the spatial Fourier spectrum and the Okubo-Weiss parameter, since those are derived from individual snapshots. However, to obtain good agreement in the temporal metrics, such as the temporal Fourier spectrum and the time autocorrelation, a well-trained CVAE will not be enough. This is the role of the RNN model, as it must be trained to replicate the dynamics in the latent space so that the decoded data will have the correct turbulent dynamics. In this section, we discuss two key aspects of the model: (1) the choice of the neural network for the time generation in the latent space and (2) the sequence length *l* used to train the RNN. Regarding the choice of the neural network, several possibilities can be envisioned:

*   **Naive deep neural network (DNN).** It consists of a simple DNN that takes a 64D vector and outputs a 64D vector. We chose three densely connected hidden layers of 128 units each. Such DNN was trained to predict $z_{i+1}$, given $z_i$, therefore its loss function is simply $L_{DNN} = \sum_i \|z_{i+1} - DNN(z_i)\|^2$. This is a relatively light architecture ($\sim 5 \times 10^4$ trainable parameters). However, it does only use one vector as input and cannot therefore capture or replicate properly the dynamics of the potential happening at larger timescales than the time step itself.

*   **Recurrent neural network (RNN).** The aforementioned limitation of the naive DNN architecture led us to the use of a more common choice in time series generation models: having a time sequence as an input the model has to predict the next element. This can technically be done with a usual DNN, but the number of trainable parameters increases dramatically. For example, if we use a similar DNN to process sequences of $l = 10$ 64D vectors, we require more than a million weights to train. This brings an imbalance between the number of parameters of the model and the number of samples in the training dataset, which is $64 \times N_s \approx 5 \times 10^5$. Therefore, the use of specific architectures designed to process sequential data is mandatory. The most common choice is the so-called recurrent neural network (RNN) architecture that we introduced in section 3, it processes series of length $l = 50$, with a limited number of parameters ($\approx 6 \times 10^4$).

*   **Long-short-term memory network (LSTM).** These networks, based on the RNN architecture, are designed to capture long-term dependencies in sequential data. The counterpart is that it requires more trainable parameters than the RNN, i.e. $\approx 2.4 \times 10^5$ for an LSTM performing a similar task than the RNN.

The Eulerian metrics for the three models are given in Fig. 19. On the left panel, we plot the spatial Fourier spectrum, which is not affected by the choice of the model, as expected. However, we see on the middle and right panels that the temporal metrics are modified depending on the model. Except for the lowest frequencies, the RNN appears to be the best at reproducing the spectrum of the HW model. At high frequencies, fluctuations obtained with the LSTM network are too high, whereas those from the DNN are too low. Regarding the autocorrelation function (right panel), the RNN clearly outperforms the other architectures. The initial decay is reproduced in all cases, but the other models fail to reproduce the oscillatory behavior of the training data.

A longer training sequence means that the RNN will be able to capture more temporal dependencies in the training data, but it will make the model more expensive computationally. We tested the models with different training lengths *l* to evaluate the impact of the length sequence. In Fig 20 we plot the temporal Fourier spectra of the electric field fluctuations, along with the error obtained for each training length *l*. In Fig 21 we plot the different autocorrelation functions and the associated errors. We can see on Fig 20 that the error for the temporal spectrum keeps decreasing as we increase *l*, although the evolution is slow: increasing the training length from
<figure><img src="image_25.png" /><figcaption>Figure 20: Comparison of the temporal spectra obtained with different training sequence lengths. Left: Temporal Fourier spectrum. Right: Evolution of the error.</figcaption></figure>

<figure><img src="image_26.png" /><figcaption>Figure 21: Comparison of the autocorrelation function obtained with different training sequence lengths. Left: Autocorrelation functions. Right: Evolution of the error.</figcaption></figure>

$l = 10$ to $l = 100$ gives a 10% decrease of the error. On the other hand, if we look at 21, we can see that the best reproduction of the autocorrelation function is obtained with $l = 50$, which is significantly better than with $l = 10$. A training length of $l = 50$ represents more than 10 autocorrelation times of the potential, whereas with $l = 10$ the RNN has access to the history of the encoded potential over 2 autocorrelation times, which seems to be limiting. Over $l = 50$, the RNN does not improve significantly. Therefore, we kept the RNN architecture with a training length of $l = 50$, since it gave very satisfactory results in a short amount of time.

## Conclusions

We have proposed a new surrogate model for 2D plasma turbulence, called Generative Artificial Intelligence Turbulence (GAIT), applying generative deep-learning techniques to the Hasegawa-Wakatani model. Such model, composed of a Convolutional Variational Auto-Encoder (CVAE) and a Recurrent Neural Network (RNN), has been used to generate synthetic turbulence data without solving all the temporal and spatial scales of the model, while capturing the essential features of turbulence. This has been possible thanks to the combination of the excellent low-dimensional representation of turbulence provided by the CVAE and the capability of the RNN to generate long time sequences in the latent space. A quantitative comparison between the turbulence obtained from direct simulations and that generated artificially by the GAIT model was performed, by looking at various important features of a turbulent flow: the Fourier spectra, the POD spectrum, the Okubo-Weiss parameter, the auto-correlation and the associated Lagrangian transport. At the moment, we have not achieved a model that can reproduce all those features simultaneously, along with an acceleration factor of several orders of magnitude. However, by meticulously designing the model with a focus on the choices that could impact performance, we were able to adapt it to gain precise control over its output. The customization of the hyper-parameters of the loss function allows for easy tuning of the GAIT model to correctly capture either Eulerian or Lagrangian aspects of turbulence. A study of the latent space showed that the topological properties of the encoded turbulence were also reproduced. The choice of different hyper-parameters, architectures, and trainings was discussed to provide deeper insight into the design of the GAIT model. This work shows the applicability of generative AI techniques to tackle the main problem of numerical turbulence, namely its computational cost. The development of surrogate turbulence models will help push some analyses, such as test particle analyses, way further than it was possible
until now. For example, systems in which transport can transition from super-diffusive at intermediate times to sub-diffusive in the $t \to \infty$ limit require very long integration times, and the development of surrogate models opens new possibilities in this direction.

The acceleration obtained with the use of AI is extremely promising and paves the way to its application to more numerically demanding cases, like 3D turbulence or gyro-kinetic turbulence [27].

## Acknowledgements

This work has been carried out within the framework of the EUROfusion Consortium, funded by the European Union via the Euratom Research and Training Programme (Grant Agreement No 101052200 — EUROfusion) Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Commission. Neither the European Union nor the European Commission can be held responsible for them. This work has received financial support from the AIM4EP project (ANR-21-CE30-0018), funded by the French National Research Agency (ANR), and the U.S. Department of Energy under Contracts No. DE-AC05-00OR22725 and DE-FG02-04ER54742. All the simulations and training of neural networks reported here were performed on HPC resources of IDRIS under the allocations 2021-A0100512455, 2022-AD010512455R1 and 2023-A0140514165 made by GENCI.

## Data Availability Statement

The data that support the findings of this study are available from the corresponding author upon reasonable request.

## References

[1] Johan Anderson, Eun-jin Kim, Bogdan Hnat, and Tariq Rafiq. Elucidating plasma dynamics in hasegawa-wakatani turbulence by information geometry. *Physics of Plasmas*, 27(2), 2020.

[2] Enrico Aymerich, Giuliana Sias, Fabio Pisano, Barbara Cannas, Sara Carcangiu, Carlo Sozzi, Chris Stu-art, PJ Carvalho, Alessandra Fanni, and JET Contributors. Disruption prediction at jet through deep convolutional neural networks using spatiotemporal information from plasma profiles. *Nuclear Fusion*, 62(6):066005, 2022.

[3] Zhe Bai, Xishuo Wei, William Tang, Leonid Oliker, Zhihong Lin, and Samuel Williams. Ftl: Transfer learning nonlinear plasma dynamic transitions in low dimensional embeddings via deep neural networks. *arXiv preprint arXiv:2404.17466*, 2024.

[4] R Basu, V Naulin, and J Juul Rasmussen. Particle diffusion in anisotropic turbulence. *Communications in Nonlinear Science and Numerical Simulation*, 8(3-4):477-492, 2003.

[5] Gal Berkooz, Philip Holmes, and John L Lumley. The proper orthogonal decomposition in the analysis of turbulent flows. *Annual review of fluid mechanics*, 25(1):539-575, 1993.

[6] P Beyer, S Benkadda, and X Garbet. Proper orthogonal decomposition and galerkin projection for a three-dimensional plasma dynamical system. *Physical Review E*, 61(1):813, 2000.

[7] Roberto Bruno and Vincenzo Carbone. The solar wind as a turbulence laboratory. *Living Reviews in Solar Physics*, 10:1-208, 2013.

[8] Andrés Bustos, E Ascasíbar, A Cappa, and Rafael Mayo-García. Automatic identification of mhd modes in magnetic fluctuation spectrograms using deep learning techniques. *Plasma Physics and Controlled Fusion*, 63(9):095001, 2021.

[9] Jeff Candy, Emily A Belli, and RV Bravenec. A high-accuracy eulerian gyrokinetic solver for collisional plasmas. *Journal of Computational Physics*, 324:73-93, 2016.

[10] VM Canuto and J Christensen-Dalsgaard. Turbulence in astrophysics: Stars. *Annual review of fluid mechanics*, 30(1):167-198, 1998.
[11] Álvaro Cartea and Diego del Castillo-Negrete. Fluid limit of the continuous-time random walk with general lévy jump distribution functions. *Physical Review E—Statistical, Nonlinear, and Soft Matter Physics*, 76(4):041105, 2007.

[12] M Cheng, F Fang, CC Pain, and IM Navon. An advanced hybrid deep adversarial autoencoder for parameterized nonlinear fluid flow modelling. *Computer Methods in Applied Mechanics and Engineering*, 372:113375, 2020.

[13] J Citrin, P Trochim, T Goerler, D Pfau, KL van de Plassche, and F Jenko. Fast transport simulations with higher-fidelity surrogate models for iter. *Physics of Plasmas*, 30(6), 2023.

[14] Jonathan Citrin, Sarah Breton, Federico Felici, Frederic Imbeaux, T Aniel, JF Artaud, B Baiocchi, C Bourdelle, Y Camenen, and J Garcia. Real-time capable first principle based modelling of tokamak turbulent transport. *Nuclear Fusion*, 55(9):092001, 2015.

[15] B Clavier, D Zarzoso, D del Castillo-Negrete, and E Frénod. Generative-machine-learning surrogate model of plasma turbulence. *Physical Review E*, 111(1):L013202, 2025.

[16] E d D Zapata-Cornejo, D Zarzoso, SD Pinches, SE Sharapov, and M Fitzgerald. Segmentation of mhd modes using fourier transform, wavelets and computer vision algorithms. *Plasma Physics and Controlled Fusion*, 66(9):095016, 2024.

[17] D del Castillo-Negrete. Asymmetric transport and non-gaussian statistics of passive scalars in vortices in shear. *Physics of Fluids*, 10(3):576-594, 1998.

[18] Diego del Castillo-Negrete. Chaotic transport in zonal flows in analogous geophysical and plasma systems. *Physics of Plasmas*, 7(5):1702-1711, 2000.

[19] Claudia Drygala, Edmund Ross, Francesca di Mare, and Hanno Gottschalk. Comparison of generative learning methods for turbulence modeling. *arXiv preprint arXiv:2411.16417*, 2024.

[20] Pan Du, Meet Hemant Parikh, Xiantao Fan, Xin-Yang Liu, and Jian-Xun Wang. Conditional neural field latent diffusion model for generating spatiotemporal turbulence. *Nature Communications*, 15(1):10416, 2024.

[21] Richard Evans, Michael O'Neill, Alexander Pritzel, Natasha Antropova, Andrew Senior, Tim Green, Augustin Žídek, Russ Bates, Sam Blackwell, Jason Yim, et al. Protein complex prediction with alphafold-multimer. *biorxiv*, pages 2021-10, 2021.

[22] Kai Fukami, Yusuke Nabae, Ken Kawai, and Koji Fukagata. Synthetic turbulent inflow generator using machine learning. *Physical Review Fluids*, 4(6):064603, 2019.

[23] S. Futatani, S. Benkadda, and D. del Castillo-Negrete. Spatiotemporal multiscaling analysis of impurity transport in plasma turbulence using proper orthogonal decomposition. *Physics of Plasmas*, 16(4):042506, 04 2009.

[24] Ann E Gargett. Ocean turbulence. *Annual Review of Fluid Mechanics*, 21(1):419-451, 1989.

[25] Tobias Goerler, Xavier Lapillonne, Stephan Brunner, Tilman Dannert, Frank Jenko, Florian Merz, and Daniel Told. The global version of the gyrokinetic turbulence code gene. *Journal of Computational Physics*, 230(18):7053-7071, 2011.

[26] Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. *Advances in neural information processing systems*, 27, 2014.

[27] Virginie Grandgirard, Jérémie Abiteboul, Julien Bigot, Thomas Cartier-Michaud, Nicolas Crouseilles, Guilhem Dif-Pradalier, Ch Ehrlacher, Damien Esteve, Xavier Garbet, Ph Ghendrih, et al. A 5d gyrokinetic full-f global semi-lagrangian code for flux-driven ion turbulence simulations. *Computer physics communications*, 207:35-68, 2016.

[28] Robin Greif, Frank Jenko, and Nils Thuerey. Physics-preserving ai-accelerated simulations of plasma turbulence. *arXiv preprint arXiv:2309.16400*, 2023.

[29] Akira Hasegawa and Masahiro Wakatani. Plasma edge turbulence. *Physical Review Letters*, 50(9):682, 1983.
[30] WH Hu, Cristina Rea, QP Yuan, KG Erickson, DL Chen, Biao Shen, Yao Huang, JY Xiao, JJ Chen, YM Duan, et al. Real-time prediction of high-density east disruptions using random forest. *Nuclear Fusion*, 61(6):066034, 2021.

[31] AE Järvinen, A Kit, YRJ Poels, S Wiesen, V Menkovski, Lorenzo Frassinetti, M Dunne, JET Contributors, ASDEX Upgrade Team, et al. Representation learning algorithms for inferring machine independent latent features in pedestals in jet and aug. *Physics of Plasmas*, 31(3), 2024.

[32] Benjamin Kadoch, Diego del Castillo-Negrete, Wouter JT Bos, and Kai Schneider. Lagrangian statistics and flow topology in forced two-dimensional turbulence. *Physical Review E—Statistical, Nonlinear, and Soft Matter Physics*, 83(3):036314, 2011.

[33] Benjamin Kadoch, Diego del Castillo-Negrete, Wouter JT Bos, and Kai Schneider. Lagrangian conditional statistics and flow topology in edge plasma turbulence. *Physics of Plasmas*, 29(10), 2022.

[34] Julian Kates-Harbeck, Alexey Svyatkovskiy, and William Tang. Predicting disruptive instabilities in controlled fusion plasmas through deep learning. *Nature*, 568(7753):526–531, 2019.

[35] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. *arXiv preprint arXiv:1312.6114*, 2013.

[36] Adam Kit, AE Järvinen, YRJ Poels, S Wiesen, V Menkovski, R Fischer, M Dunne, ASDEX-Upgrade Team, et al. On learning latent dynamics of the aug plasma state. *Physics of Plasmas*, 31(3), 2024.

[37] Ralf Kleiber, Matthias Borchardt, Roman Hatzky, Axel Könies, Henry Leyh, Alexey Mishchenko, Jörg Riemann, Christoph Slaby, JM García-Regaña, Edilberto Sanchez, et al. Euterpe: A global gyrokinetic code for stellarator geometry. *Computer Physics Communications*, 295:109013, 2024.

[38] Solomon Kullback and Richard A Leibler. On information and sufficiency. *The annals of mathematical statistics*, 22(1):79–86, 1951.

[39] Emmanuel Lanti, Noé Ohana, Natalia Tronko, Thomas Hayward-Schneider, Alberto Bottino, Ben F McMillan, Alexey Mishchenko, Aaron Scheinberg, Alessandro Biancalani, P Angelino, et al. Orb5: a global electromagnetic gyrokinetic code using the pic approach in toroidal geometry. *Computer Physics Communications*, 251:107072, 2020.

[40] Romit Maulik, Bethany Lusch, and Prasanna Balaprakash. Reduced-order modeling of advection-dominated systems with recurrent neural networks and convolutional autoencoders. *Physics of Fluids*, 33(3), 2021.

[41] Christopher J McDevitt. A physics-informed deep learning model of the hot tail runaway electron seed. *Physics of Plasmas*, 30(9), 2023.

[42] Andrea Murari, Michele Lungaroni, Emmanuele Peluso, Pasqualino Gaudio, Jesús Vega, Sebastián Dormido-Canto, Matteo Baruzzo, Michela Gelfusa, and JET Contributors. Adaptive predictors based on probabilistic svm for real time disruption mitigation on jet. *Nuclear Fusion*, 58(5):056002, 2018.

[43] AG Peeters, Yann Camenen, F James Casson, WA Hornsby, AP Snodin, D Strintzi, and Gabor Szepesi. The nonlinear gyro-kinetic flux tube code gkw. *Computer physics communications*, 180(12):2650–2672, 2009.

[44] Norbert Peters. Turbulent combustion. *Measurement Science and Technology*, 12(11):2022–2022, 2001.

[45] Andrey Pushkarev. *Self-organization of isotopic and drift-ware turbulence*. PhD thesis, Citeseer, 2013.

[46] Christina Rea, KJ Montes, KG Erickson, RS Granetz, and RA Tinguely. A real-time machine learning-based disruption predictor in diii-d. *Nuclear Fusion*, 59(9):096016, 2019.

[47] Joel Sommeria. Experimental study of the two-dimensional inverse energy cascade in a square box. *Journal of fluid mechanics*, 170:139–168, 1986.

[48] EA Spiegel and J-P Zahn. The solar tachocline. *Astronomy and Astrophysics* (ISSN 0004-6361), vol. 265, no. 1, p. 106-114, 265:106–114, 1992.

[49] John C Wyngaard. Atmospheric turbulence. *Annual Review of Fluid Mechanics*, 24(1):205–234, 1992.
[50] Minglei Yang, Pengjun Wang, Diego del Castillo-Negrete, Yanzhao Cao, and Guannan Zhang. A pseudoreversible normalizing flow for stochastic dynamical systems with various initial distributions. *SIAM Journal on Scientific Computing*, 46(4):C508–C533, 2024.

[51] Gokul Yenduri, M Ramalingam, G Chemmalar Selvi, Y Supriya, Gautam Srivastava, Praveen Kumar Reddy Maddikunta, G Deepti Raj, Rutvij H Jhaveri, B Prabadevi, Weizheng Wang, et al. Gpt (generative pre-trained transformer)-a comprehensive review on enabling technologies, potential applications, emerging challenges, and future directions. *IEEE Access*, 2024.

[52] EdD Zapata-Cornejo, D Zarzoso, SD Pinches, A Bustos, A Cappa, and E Ascasibar. A novel unsupervised machine learning algorithm for automatic alfvénic activity detection in the tj-ii stellarator. *Nuclear Fusion*, 64(12):126057, 2024.

[53] JX Zhu, C Rea, RS Granetz, ES Marmar, R Sweeney, K Montes, and RA Tinguely. Integrated deep learning framework for unstable event identification and disruption prediction of tokamak plasmas. *Nuclear Fusion*, 63(4):046009, 2023.