Learning the dynamics of a one-dimensional plasma model with graph
  neural networks
We explore the possibility of fully replacing a plasma physics kinetic
simulator with a graph neural network-based simulator. We focus on this class
of surrogate models given the similarity between their message-passing update
mechanism and the traditional physics solver update, and the possibility of
enforcing known physical priors into the graph construction and update. We show
that our model learns the kinetic plasma dynamics of the one-dimensional plasma
model, a predecessor of contemporary kinetic plasma simulation codes, and
recovers a wide range of well-known kinetic plasma processes, including plasma
thermalization, electrostatic fluctuations about thermal equilibrium, and the
drag on a fast sheet and Landau damping. We compare the performance against the
original plasma model in terms of run-time, conservation laws, and temporal
evolution of key physical quantities. The limitations of the model are
presented and possible directions for higher-dimensional surrogate models for
kinetic plasmas are discussed.
Diogo D Carvalho, Diogo R Ferreira, Luis O Silva
Do Graph Neural Networks Dream of Landau Damping? Insights from Kinetic Simulations of a Plasma Sheet Model

###### Abstract

We explore the possibility of fully replacing a plasma physics kinetic simulator with a graph neural network-based simulator. We focus on this class of surrogate models given the similarity between their message-passing update mechanism and the traditional physics solver update, and the possibility of enforcing known physical priors into the graph construction and update. We show that our model learns the kinetic plasma dynamics of the one-dimensional plasma model, a predecessor of contemporary kinetic plasma simulation codes, and recovers a wide range of well-known kinetic plasma processes, including plasma thermalization, electrostatic fluctuations about thermal equilibrium, and the drag on a fast sheet and Landau damping. We compare the performance against the original plasma model in terms of run-time, conservation laws, and temporal evolution of key physical quantities. The limitations of the model are presented and possible directions for higher-dimensional surrogate models for kinetic plasmas are discussed.

_Keywords: Plasma Physics, Machine Learning, Graph Neural Networks_

## 1 Introduction

Simulating the kinetic behavior of a plasma [1] is a complex and computationally demanding task. Fully relativistic massively-parallelized Particle-in-Cell (PIC) codes are commonly used to model these phenomena and have been shown to correctly recover and predict plasma collective behavior [2, 3, 4].

To obtain computational speed-ups, there have been recent attempts to combine existing PIC codes with machine learning surrogate models. These efforts include approaches to accelerate [5] or fully replace [6] the field solver block, or the integration of surrogate models into advanced physics extensions [7]. However, how to simultaneously obtain a significant computational gain, while enforcing known physics constraints is still an open research question.
Developments in machine learning introduced several physics-inspired surrogate models as an alternative to standard differential equation solvers [8, 9, 10, 11] and \(n\)-body or mesh-based simulators [12, 13, 14, 15, 16, 17].

From the broad set of available surrogate models, one class of algorithms that can be of particular interest for kinetic plasma simulations are graph neural network-based approaches [18, 19], because of their capability to model both particle-particle [12] and particle-mesh interactions [13], as well as the possibility of enforcing known invariances or symmetries into the network architecture [15, 16, 19]. These approaches have been successfully applied to fluid [12, 13], rigid body [13, 20], and charged particle dynamics [15, 16, 21]. However, to the best of our knowledge, their applicability to model kinetic plasma simulations is still to be demonstrated.

In this work we aim to model the predecessor of the PIC loop, the one-dimensional electrostatic sheet plasma model introduced by Dawson [22, 23]. This is an ideal initial testbench since it provides a simpler scenario, in terms of the problem structure and possible computational gains, while capturing a wide range of kinetic plasma physics phenomena, that go beyond "collisionless" physics, including Coulomb collisions, and collisional thermalization [22, 23, 24, 25, 26]. Moreover, recent studies in the fundamental statistical physics processes in plasmas have been using the sheet model and/or direct extensions [27]. We show how to leverage previous work on graph neural network-based simulators by Sanchez-Gonzalez et al. [12] to kinetic plasma simulations and to the one-dimensional sheet model by introducing domain knowledge into the graph construction and simulator update mechanisms which enforce the desired symmetries and invariances. We discuss the advantages and disadvantages of using our surrogate model when compared to the standard physics simulator in terms of accuracy, run-time, energy conservation, and generalization capabilities. Finally, based on our findings, we comment on the expected impact of graph neural network-based simulators for the multi-dimensional PIC scenario.

## 2 Electrostatic Sheet Model

The single-species one-dimensional electrostatic sheet model introduced by Dawson [22, 23] represents a plasma as a group of equally negatively charged sheets, moving freely over a uniformly neutralizing ion background (see figure 1). In a one-dimensional system, this model is exact and describes, within classical physics, the dynamics of an arbitrary non-relativistic plasma.

The system is in equilibrium if the sheets are at rest and equally spaced. In such a scenario, the electric field at a given sheet is zero, and the field profile is represented by a sawtooth function; it varies linearly along the x-axis by a factor of \(4\pi en_{0}\), where \(e\) represents the electron charge and \(n_{0}\) the background ion number density, and jumps by a factor of \(-4\pi en_{0}\delta\) at each sheet, where \(\delta=L/N_{sheets}\) represents the inter-sheet spacing at equilibrium, \(L\) the box length and \(N_{sheets}\) the number of sheets.

If one sheet moves a certain distance \(\xi\) from its equilibrium position, it will
experience an electric field \(E=4\pi en_{0}\xi\), resulting in an equation of motion given by:

\[\ddot{\xi}=-\omega_{p}^{2}\xi \tag{1}\]

where \(m_{e}\) represents the electron mass, and \(\omega_{p}=\sqrt{4\pi n_{0}e^{2}/m_{e}}\) is the plasma frequency. This result implies that, for small displacements, each sheet behaves as an independent harmonic oscillator.

For larger displacements, it is possible that consecutive sheets cross each other, corresponding to a one-dimensional Coulomb collision, meaning their equilibrium positions switch. Alternatively, this interaction can be modeled as an elastic collision, i.e. one can simply switch the velocities of the sheets at the instant of crossing (instead of their equilibrium positions) as it results from the conservation of energy and momentum. An illustration of the difference in the resulting individual trajectories is provided in figure 2.

To simulate this system, two computational directions can be used [23]: a synchronous method, and an asynchronous method. The synchronous method first updates the sheet dynamics according to (1) considering a fixed \(\Delta t\). It then detects crossings by testing the condition \(x_{i}^{t+1}>x_{j}^{t+1}\) for \(j>i\), and proceeds to use an iterative method to estimate the crossing times and correct the motion of the corresponding sheets. This method does not correctly resolve crossings involving more than two sheets in a single time step, which leads to an overall energy loss in the system if the time step is too large compared with the inverse of the typical collision frequency. For this reason, for higher sheet velocities it is necessary to use smaller simulation time steps (the collision frequency increases with increasing thermal velocity).

On the other hand, the asynchronous method advances the simulation until the next crossing. The next crossing time can be predicted analytically from (1) by solving for \(x_{i}(t)=x_{j}(t)\) with respect to \(t\) for all pairs of adjacent sheets. This algorithm guarantees energy conservation but implies additional computational effort since a sorted table containing all expected crossing times between neighboring sheets must be pre-computed and updated after each crossing is resolved.

Figure 1: Schematic of the 1D single species electrostatic sheet model. In equilibrium, the negatively charged sheets (red) are equally spaced inside the box. If one sheet is displaced from its equilibrium position, the average electric field the sheet feels is not zero, due to the charge imbalance. Adapted from Dawson [22].
Since the Graph Network Simulator (GNS) is a synchronous model, we use the synchronous sheet model algorithm (illustrated in figure 3) for both data generation and testing purposes to allow for model accuracy comparisons at equivalent simulation time steps.

## 3 Graph Network Simulator

The GNS architecture used here is inspired by the work of Sanchez-Gonzalez et al. [12] while taking into consideration the specifics of the electrostatic sheet model. The main building blocks are presented in figure 3.

Based on the sheet positions \(\mathbf{x}^{t}\), velocities \(\mathbf{v}^{t}\), equilibrium positions \(\mathbf{x}^{t}_{eq}\), and boundary conditions, we generate a graph representation \(\mathcal{G}\) of the plasma. A Graph Neural Network (GNN) predicts each individual sheet acceleration \(\mathbf{a}^{t}\), which will be used to update the positions \(\mathbf{x}^{t}\) and velocities \(\mathbf{v}^{t}\). We enforce the boundary conditions by

Figure 2: Comparison of charged sheet trajectories when considering sheet interactions as crossings (top) _versus_ binary collisions (bottom). The represented system consists of 10 sheets (represented by different colors) moving on a periodic box of length \(L\). Initial velocities were randomly chosen. We will learn to model the dynamics of the first case (crossings) as this is considerably easier, mainly due to the smoothness of the sheet trajectories. More details on the difficulties that arise when attempting to learn collisional dynamics are provided in Appendix A.
re-injecting particles that crossed the boundary, sorting the particles by their position, and updating their equilibrium positions. This process can be repeated to generate longer simulation rollouts.

Our GNS does not treat particle interaction as binary collisions. Instead, we learn to predict the changes in velocity as sheets move through one another. This choice makes it significantly easier for the network to learn the dynamics of the system and also reduces the graph and model complexity. We provide this comparison, as well as the structural changes required to the simulator when considering crossings as collisions, in A.

### Graph Representation

The plasma is represented as a graph by a set of nodes \(\{\mathbf{n}_{i}\}\) representing the negatively charged sheets, and a set of directed edges \(\{\mathbf{r}_{ij}\}\) which connect neighboring sheets. In our case, for simplicity and performance reasons we opted to connect only the first neighbors (additional comments on the impact of higher connectivity are provided throughout the paper). Each node vector \(\mathbf{n}_{i}\) contains the information relative to the corresponding negatively charged sheet, while each edge vector \(\mathbf{r}_{ij}\) contains information about the relative displacement of node \(i\) in relation to node \(j\). They are defined as:

\[\begin{split}\mathbf{n}_{i}^{t}&=\left[\xi_{i}^{t}, v_{i}^{t}\right]\\ \mathbf{r}_{ij}^{t}&=\left[x_{j}^{t}-x_{i}^{t}\right] \end{split} \tag{2}\]

where \(\xi_{i}^{t}\) corresponds to the displacement of the \(i\)th sheet from its equilibrium position \(\left(x_{i}^{t}-x_{eq_{i}}^{t}\right)\), and \(v_{i}^{t}\) the finite difference velocity \(\left(x_{i}^{t}-x_{i}^{t-1}\right)/\Delta t\). To allow the model to generalize to different box sizes and number of sheets, we normalize all distances and velocities by the intersheet distance at equilibrium \(\delta\). This transformation makes

Figure 3: Schematic representations of the synchronous Electrostatic Sheet Model algorithm [23] and the Graph Network Simulator used.
the inputs of the network invariant to the system size (box length and number of particles). This is also the reason why we include in the node vector the displacement from equilibrium instead of the sheet position inside the box, which additionally enforces an invariance to the sheet rank.

When considering reflecting boundary conditions, extra nodes representing mirrored versions of the boundary sheets are added to the graph (figure 4). Ideally, the number of mirror sheets should be large enough so that no boundary sheet crosses all mirror sheets in a single time step (in our case we set it to the number of message-passing steps, which we cover in Section 3.2). On the other hand, for periodic boundaries the graph becomes cyclic and we consider the distance between boundary sheets to be equal to the distance through the simulation box walls. In subsequent sections, we will show that representing boundaries in this way allows the GNN model to be applied to different boundary conditions than the ones it was trained on since it learns interactions between pairs of sheets (not interactions with the wall). However, this comes at the cost of an extra hard-coded step, that re-injects particles that crossed the boundary (procedure to be explained in more detail in later sections).

We experimented with several other possible representations of the system, (e.g. different boundary representations, node and edge parameters, connect \(n\)-nearest neighbors) but these variants either produced worse results (in terms of accuracy or generalization capabilities) or introduced extra complexity and memory requirements that did not result in meaningful accuracy and/or speed-up improvements for the tested scenarios.

### Graph Neural Network Architecture

The GNN module follows an encoder-processor-decoder architecture as introduced by Battaglia et al. [18]. A slight modification is used, which includes the sent edges in the node update mechanism of the processor block (equivalent to the Graph Network implementation available in Jraph [28]). The main building blocks are as follows:

**Encoder**: The encoder transforms each node \(\mathbf{n}_{i}\) and edge \(\mathbf{r}_{ij}\) vectors into a latent space node \(\mathbf{v}_{i}\) and edge \(\mathbf{e}_{ij}\) representation according to:

\[\mathbf{v}_{i}=\varepsilon^{v}\left(\mathbf{n}_{i}\right) \tag{3}\] \[\mathbf{e}_{ij}=\varepsilon^{e}\left(\mathbf{r}_{ij}\right)\]

Figure 4: Graph representation of a four-sheet system for different boundary conditions. The number of mirror sheets used for reflecting boundaries is not necessarily one, and, if necessary, can be larger than the number of sheets inside the box.
where \(\varepsilon^{v}\) and \(\varepsilon^{e}\) are learnable functions.

**Processor**: The processor consists of a series of \(M\) slightly modified Graph Network (GN) blocks [18, 28] which sequentially update the node and edge values according to:

\[\begin{split}\mathbf{e}_{ij}^{m+1}&=\phi^{e}\left( \mathbf{e}_{ij}^{m},\mathbf{v}_{i}^{m},\boldsymbol{\cdot}\mathbf{v}_{j}^{m} \right)\\ \overline{\mathbf{e}}_{r_{i}}^{m+1}&=\sum_{j\in \mathcal{N}(i)}\mathbf{e}_{ij}^{m+1}\\ \overline{\mathbf{e}}_{s_{i}}^{m+1}&=\sum_{j\in \mathcal{N}(i)}\mathbf{e}_{ji}^{m+1}\\ \mathbf{v}_{i}^{m+1}&=\phi^{v}\left(\overline{ \mathbf{e}}_{r_{i}}^{m+1},\overline{\mathbf{e}}_{s_{i}}^{m+1},\mathbf{v}_{i}^ {m}\right)\end{split} \tag{4}\]

where the superscript denotes the block number, \(\mathcal{N}(i)\) the set of nodes connected to \(i\), and \(\phi^{v}\), \(\phi^{e}\) are learnable functions. The value of \(M\) is set depending on the training time step and the maximum velocity of the particles present in the training simulations. Ideally, \(M\) should be larger than the maximum number of neighboring sheets that a given sheet crosses in any particular time step, since a graph node will at most receive/send information from/to the \(M^{\text{th}}\) neighbor. This condition can be relaxed if the \(N\)-nearest neighboring nodes are directly connected (with the cost of additional memory requirements), and we have observed similar performance in our tested scenarios for models with an equivalent \(M\times N\) factor. However, for higher crossing frequencies than the ones presented in this work, it might be preferable to increase \(N\) to avoid possible information bottlenecks [29].

**Decoder**: The decoder block transforms the latent node representations of the last layer into the output vector:

\[\mathbf{y}_{i}=\delta^{v}\left(\mathbf{v}_{i}^{M}\right) \tag{5}\]

where \(\delta^{v}\) is a learnable function. In our case, the output vector \(\mathbf{y}_{i}\) is a single real value that corresponds to an estimate of the individual finite difference particle acceleration \(a_{i}^{t}=(v^{t+1}-v^{t})/\Delta t=(x^{t+1}-2x^{t}-x^{t-1})/\Delta t\) normalized to the intersheet spacing \(\delta\).

We parameterize the encoder and decoder functions \((\varepsilon^{e},\varepsilon^{v},\delta^{v})\) as linear transformations. As for the processor functions \((\phi^{e},\phi^{v})\), they are given by a two-layer dense neural network following: Input \(\rightarrow\{\text{LinearLayer}()\rightarrow\text{ReLU}\rightarrow\text{ LinearLayer}()\}\)\(\rightarrow\) Output. In every block (encoder, processor, and decoder), we use a latent space of size 128. A summary of the hyper-parameter tuning experiments that led to these final values is provided in Appendix B.

Although this GNN architecture does not enforce equivariance with respect to reflections (a symmetry present in the sheet model, i.e. if the simulation box is flipped the absolute value of the predicted accelerations should simply switch signs), we observed that the network was nonetheless capable of correctly approximating this symmetry within the training data range. In fact, we developed an alternative architecture
that enforced this equivariance but did not observe relevant gains concerning the required number of training simulations nor improved rollout accuracy or energy conservation capabilities (in fact we observed a deterioration for considerable out-of-training data distribution values). Similarly, not using the sent messages for the node update mechanism led to poorer energy conservation for (considerably) out-of-training data distribution values. More details about these comparisons and the equivariant architecture are provided in C.

### Position and Velocity Update

To update the particle positions and velocities we use a semi-implicit first-order Euler integration scheme as follows:

\[\begin{split}\tilde{v}^{t+1}&=v^{t}+a^{t}\Delta t \\ \tilde{x}^{t+1}&=x^{t}+\tilde{v}^{t+1}\Delta t. \end{split} \tag{6}\]

After this update, we resolve the boundary crossings. When considering reflecting boundary conditions, we flip the position and velocity of the particles that left the simulation box. No change is applied to their equilibrium positions. When considering periodic boundaries, we instead re-insert the particles through the opposite boundary without changing their velocities. Additionally, the equilibrium positions are updated to take into consideration particles that crossed the boundaries (additional information is provided in D).

Finally, we sort particles by their position inside the box. This step is required to correctly attribute equilibrium positions and ensure the necessary relative ordering for graph construction.

## 4 Implementation

For reference, we implemented the synchronous version of the original electrostatic sheet model [22, 23] in Python, using NumPy [30]. This code is used to generate all ground truth training and test data at a high temporal resolution. The GNS was also implemented in Python using JAX [31], Jraph [28], and Haiku [32].

Additionally, from here onwards we will adopt a system of units similar to Dawson [22]. Time will be shown in units of the plasma period \(\omega_{p}^{-1}\) (with \(\omega_{p}\) as defined in (1)), distances will be presented in units of the intersheet spacing in equilibrium \(\delta\), resulting in velocities in units of \(\delta\!\cdot\!\omega_{p}\) and accelerations in units of \(\delta\!\cdot\!\omega_{p}^{2}\). Note that in the adopted units the Debye length \(\lambda_{D}\) is equivalent to the thermal velocity since, by definition, \(v_{th}=\lambda_{D}\omega_{p}\)[1], and the length of the simulation box \(L\) is equivalent to the number of sheets since \(L=N_{sheets}\delta\).
### Generating the ground truth data

Using the electrostatic sheet model, we generate 10,000 simulations of systems consisting of 10 sheets moving inside a periodic box. All simulations are run for a duration of \(t_{max}=10\ \omega_{p}^{-1}\) using a time step of \(\Delta t_{sim}=10^{-4}\ \omega_{p}^{-1}\). The initial displacements from the equilibrium positions and velocities of the sheets are randomly sampled from uniform distributions. The maximum initial displacement equals \(\xi_{max}^{0}=0.2\ \delta\) and the maximum initial velocity is \(v_{max}^{0}=10\ \delta\!\cdot\!\omega_{p}\). In addition, we ensured that the total energy of the system did not vary more than a predefined threshold (\(\Delta\varepsilon/\varepsilon_{0}\!=\!10^{-6}\)) during the full simulation by discarding simulations that did not fulfill this criterion. This guarantees that all crossings are well resolved by the sheet model.

### Data preprocessing and augmentation

Before training the GNN models, we apply the following preprocessing steps. First, we downsample the data to the desired training time step (e.g. \(\Delta t_{train}=10^{-2}\ \omega_{p}^{-1}\)). To take advantage of the system symmetries, we proceed to augment the training dataset by mirroring the simulations along the x and time axis (the latter is not equivalent to simply changing the sign of the velocities since we are using the finite difference velocity). We proceed to generate pairs of input graphs and output target accelerations, where each input graph corresponds to a full simulation rollout (and corresponding augmented versions). More details on the impact of the impact of the training dataset size and data augmentation are provided in Appendix E.

### Training

To train the models, we hold out 100 simulations for validation purposes. We proceed to minimize the mean squared error between the predicted and target accelerations using the Adam optimizer. We use an exponential learning rate scheduler, similarly to Sanchez-Gonzalez et al. [12], for which \(\alpha(j)=\alpha_{final}+(\alpha_{start}-\alpha_{final})\cdot 0.1^{j\cdot 10^{6}}\), where \(j\) represents the gradient update step, and the initial and final learning rates are given by \(\alpha_{start}=10^{-4}\) and \(\alpha_{final}=10^{-6}\). We set the batch size to 1 (one graph corresponds to a full simulation) and compute the validation loss on the full validation set once a full cycle over the training dataset is completed. The training procedure is then run to a maximum of \(1\times 10^{6}\) gradient updates for \(\Delta t_{train}=10^{-1}\ \omega_{p}^{-1}\), and \(1.5\times 10^{6}\) gradient updates for \(\Delta t_{train}=10^{-2}\ \omega_{p}^{-1}\). The final weights of the model are those obtained for the smallest recorded validation loss. The full training procedure lasts approximately 4 hours for \(\Delta t_{train}=10^{-1}\ \omega_{p}^{-1}\) and \(M=5\), and 1 day for \(\Delta t_{train}=10^{-2}\ \omega_{p}^{-1}\) and \(M=3\), using a single Nvidia Titan X GPU. For each value of \(\Delta t_{train}\) we train 5 equivalent models using different random seeds in order to assess the dependence of performance on weight initialization.
## 5 Model Benchmark

In this section, we assess the capability of the GNS to predict individual sheet trajectories. We showcase the generalization capabilities already hinted at in Section 3.1 by evaluating the model accuracy for systems of different sizes and boundary conditions. Additionally, we compare its energy conservation capabilities and run-time against the Sheet Model (SM) and discuss the identified GNS limitations.

### Trajectory Prediction Error

In order to benchmark the rollout accuracy and generalization capabilities of the GNS, we evaluate its accuracy on multiple test sets consisting of systems with different numbers of sheets and boundary conditions. Each test set contains 100 simulations with a similar duration, temporal resolution, maximum initial displacement and velocity as the ones present in the training set. Our evaluation metrics are the rollout mean absolute error (MAE) and the earth mover's distance (EMD) [33] between the predicted and the ground truth sheet trajectories, calculated, for each time step, as

\[\begin{split}\text{MAE}=&\ \frac{1}{N}\sum_{i} \left|x_{i}^{GNS}-x_{i}^{True}\right|\\ \text{EMD}=&\ \min\frac{1}{N}\sum_{i}^{N}\sum_{j}^{N} \left|x_{i}^{GNS}-x_{j}^{True}\right|\cdot c_{ij},\\ &\ \text{subject to}\ \sum_{i}c_{ij}=1\,\ \sum_{j}c_{ij}=1\,\ \text{and}\ c_{ij}\in\{0,1\}\end{split} \tag{7}\]

and then average over the full simulation rollout. In the case of periodic boundaries, we modify both metrics to consider the absolute distance to be the minimum of the distances through the box or through the walls.

The results presented in figure 5 allow us to draw some conclusions. We observe the rollout errors obtained are considerably small (note that they are presented in units of the intersheet spacing \(\delta\)) demonstrating that, despite training solely on single-step acceleration prediction, we achieve a stable rollout accuracy. To provide additional insight into the small scale of the errors, we showcase in figure 6 the worst test simulation rollouts (highest rollout EMD across all models) for different time steps and boundary conditions.

We also observe the simulator accuracy is independent of the number of sheets and boundary conditions, without additional re-training/fine-tuning. These invariances, already hinted at in Section 3.1, are a direct consequence of the chosen graph representation of the system.

Finally, figure 5 illustrates the importance of using both the MAE and the EMD as complementary evaluation metrics. The larger intervals associated with the MAE are produced by a small set of simulations where, due to the accumulation of small prediction errors, two particles switch their expected relative order during a tangential
crossing (i.e. when they are moving in the same direction with very similar velocities). This results in a permutation of their predicted trajectories with respect to their ground truth, leading to larger MAE values (see for example figure 6 for reflecting boundary, around \(t=2\)\(\omega_{p}^{-1}\) the orange and dark blue trajectories permute after reflection). The reason why the error intervals decrease significantly for the EMD case is because this metric is invariant to permutations of sheet trajectories. This invariance (which the MAE does not provide) is an important property for our case study since a simple permutation of sheet trajectories does not change the distribution function of the system (i.e. the systems are equivalent). Therefore, the EMD provides a better assessment of the accuracy of the simulator to model the collective plasma dynamics.

We observed that, overall, equivalent models trained with different random seeds converge to very similar rollout errors (detailed comparisons provided in F). The only exception was one of the models trained for \(\Delta t=10^{-2}\)\(\omega_{p}^{-1}\) which revealed a slightly worse rollout performance. We attribute this larger error to its worse single-step prediction capabilities given the validation loss at train time was approximately \(2\times\) that of equivalent models.

Figure 5: Rollout error metrics for the GNS in the test set simulations. For each value of \(\Delta t\) we compute the metrics for 5 equivalent GNNs trained using different random seeds. The presented mean values are computed by averaging over sheets, time steps, simulations, and GNN models (for a detailed comparison between different models see F). The error bars represent the minimum and maximum rollout error achieved for the corresponding set of test simulations across all models. The results demonstrate that even though the training data contains solely systems consisting of 10 sheets moving over a periodic box, the GNS is capable of generalizing to smaller/larger system sizes and different boundary conditions. Furthermore, the reported errors are considerably small.
Figure 6: Example of simulation rollouts observed for test simulations of 10 sheet systems. These examples correspond to the worst-performing rollouts (largest EMD across all models) for the indicated simulator time step and boundary conditions. The predicted and ground truth trajectories and the MAE/EMD evolution are shown (per time step average over sheets). In both cases, the ground truth trajectories (obtained with the Sheet Model using \(\Delta t=10^{-4}\)\(\omega_{p}^{-1}\)) are downsampled to the same simulation time step as the GNS. We plot the ground truth trajectories with a larger marker size in order to be possible to distinguish them with respect to the prediction. It is clear that the GNS is capable of correctly modeling the sheet trajectories for longer rollouts, even though it was solely trained on single-step prediction.
### Energy Conservation

In order to check for energy conservation, we run simulations using two types of initial velocity distributions: thermal - velocities sampled from a normal distribution with standard deviation equal to \(v_{th}\); oscillation - sheets share the same initial velocity \(v_{0}\) (no crossings should occur). For both initial conditions, we perform a scan over the initial thermal/oscillating velocity (one simulation per value). All simulations consider a system of \(10^{3}\) sheets moving over a periodic box for a total of \(t_{max}=5\times 2\pi\)\(\omega_{p}^{-1}\).

While for the sheet model the energy decreases monotonically, we observed this was not the case for the GNS (energy might increase, decrease, or oscillate, examples are provided in F). Furthermore, since the GNS uses the finite difference velocities instead of the instantaneous velocities, there is an oscillation associated with the plasma period (the period is equal to half a plasma period) which is clearly dominant for lower thermal velocities. This oscillation is merely an artifact of the finite different velocities used to compute the energy of the system.

To allow for a fair comparison between the sheet model and GNS we then compute the total energy variation as follows: we skip the first plasma oscillation, apply a moving average with a window size of \(\Delta t=2\pi\)\(\omega_{p}^{-1}\) to the remaining datapoints and retrieve the maximum deviation from the initial energy of the system (all steps are further justified in F). For each time step the energy of the system \(\epsilon\) is computed according to:

\[\epsilon=\frac{1}{2}m_{e}\sum_{i}^{N}\left(v_{i}^{2}+\omega_{p}^{2}\xi_{i}^{2 }\right). \tag{8}\]

The final results containing scans performed for different initial velocities distributions and simulation time steps are presented in figure 7. It is observed that as the initial thermal velocity increases, the GNS showcases a gain of one or two orders of magnitude in terms of energy conservation when compared to the Sheet Model (SM) algorithm running at a similar time step. This happens since the GNS captures more accurately crossings involving \(n>2\) sheets.

The accuracy of the GNS models remains approximately constant for the thermal velocities for which it was trained (\(v_{th}\leq v_{th}^{train}\)). However, for out-of-distribution scenarios the energy variation starts increasing, more noticeable for \(\Delta t=10^{-2}\)\(\omega_{p}^{-1}\). This behavior is not visible for the oscillatory initial conditions, where we observe a very stable energy loss rate across all models for \(v_{0}\leq v_{max}^{train}\). Therefore, the lower performance for higher \(v_{th}\) should be attributed to a failure in correctly modeling all crossing events in such conditions.

A performance degradation is expected since: a) the number of message-passing steps associated with each model (\(M=3\) for \(\Delta t=10^{-2}\)\(\omega_{p}^{-1}\), \(M=5\) for \(\Delta t=10^{-1}\)\(\omega_{p}^{-1}\)) limits the GNS capability to correctly resolve crossings involving a larger number of sheets (which are more likely to occur at larger values of \(v_{th}\)); b) the GNN is only trained on crossings involving \(n\leq n_{max}\), where \(n_{max}\) is the maximum number of sheets involved in any crossing within the training data. However, this does not explain why
outside the training data the performance of the GNS at different time resolutions is equivalent. We attribute this behavior to three main effects.

Firstly, there is one particular GNN model that consistently performed worse across all metrics (validation loss, rollout accuracy, and energy conservation) which biased significantly the average value of the energy variation rate. Removing this model from the average calculation changes considerably the behavior for \(v_{th}^{train}<v_{th}<v_{train}^{max}\) (more details in Appendix F).

Secondly, the percentage of training data points containing crossings at higher temporal resolution is, at least, \(10\times\) smaller than that of the lower resolution scenario. This can bias the training procedure to reduce the prediction error for events that do not involve crossings to the detriment of the (smaller) subset that contains crossings (which becomes problematic at test time for scenarios where crossings dominate the overall dynamics). Additional support for this claim is the fact that the higher temporal resolution models seem to be "overfitting" the purely oscillatory dynamics within the training data range since they fail to generalize to larger oscillation amplitudes (note the steep increase in energy variation for the oscillating initial conditions for \(v_{0}>v_{max}^{train}\)).

Finally, since we are performing an increased number of updates for the same time

Figure 7: Energy variation rate for simulations of \(10^{3}\) sheets moving over a periodic box for \(t_{max}=5\times 2\pi\)\(\omega_{p}^{-1}\). Initial sheet velocities are sampled from a normal distribution (thermal), or all equal to \(v_{0}\) (oscillation). For the Sheet Model (SM) we run a single simulation per set of initial conditions and simulation time step \(\Delta t\). For the GNS we run, per trained \(\Delta t\), the exact same simulations using 5 equivalent GNN models (trained with different random seeds). The mean values across GNN models are presented in full line, and the min/max values are represented by the shaded region (for a detailed comparison between different seeds see Appendix F). No lines are shown for the SM in the oscillation regime since it conserves perfectly the energy when no crossings occur. For dynamics that involve sheet crossings, the GNS either performs significantly better or is on pair with the SM running at an equivalent simulation step.
interval, the prediction errors for the out-of-distribution data might be accumulating faster. The impact of the aforementioned effects could be investigated in future work (and mitigated if necessary) by using alternative data sampling strategies and/or optimizing for multi-step accuracy at train time [17]. However, the main focus should be the improvement of the performance for lower temporal resolutions since this is where larger computational gains are expected.

It would also be important to study, in future work, why some models are capable of achieving better energy conservation capabilities at larger \(v_{th}\) (which might indicate that they learn a more robust crossing resolution algorithm) and how to consistently achieve this level of performance (e.g. by using a regularizing penalty [21]). Additional evaluation metrics/tests should also be devised since neither the validation loss nor the current rollout accuracy tests seem to be good predictors for improved energy conservation capabilities. These new tests could include, for instance, measurements of rollout accuracy for significantly longer simulations and higher thermal velocities.

### Run-time

Using the same setup for thermal initial conditions, we analyzed the run-time of the sheet model _versus_ the GNS. The results obtained are presented in figure 8. It is clear that the sheet model run-time increases with the crossing frequency, while the GNS does not. This difference in behavior coupled with the gain in energy conservation at higher thermal velocities makes a stronger case for the suitability of the GNS when modeling these scenarios. For example, when considering \(v_{th}=10\ \delta\omega_{p}\ (\lambda_{D}=10\ \delta)\), the GNS using \(\Delta t=10^{-1}\ \omega_{p}^{-1}\) obtains an order of magnitude improvement in both energy conservation and run-time when compared to the sheet model using \(\Delta t=10^{-2}\ \omega_{p}^{-1}\).

However, it is important to highlight that the different models are implemented with different packages (NumPy _vs_ JAX) and running on different hardware (CPU _vs_ GPU) which influences their respective run-time. Furthermore, we make no claims that our implementations are optimal, meaning they could both benefit from further speed-ups. Nonetheless, the qualitatively different behaviour of both algorithms, coupled with differences in accuracy of orders of magnitude on both energy conservation and run-time, leads us to believe that there is a computational gain to be expected in using the GNS for larger simulation steps with higher crossing rates.

### Limitations

The main limitations that we identified for the GNS are the requirement to use a fixed simulation step (equal to the training simulation step) and the performance degradation on out-of-training data distribution scenarios (as demonstrated in figure 7).

The first constraint is solely due to the fact that the network has to learn to predict sheet crossings, which implicitly forces it to know what is the simulation step. If one wishes to train a single model for different simulation steps, it would be necessary to provide \(\Delta t\) as an input to the network (in the GNS the time step only appears
explicitly in the ODE integrator). Alternatively, as we have shown, different models can be trained, one per \(\Delta t\). As long as enough training data is provided and the model architecture is scaled accordingly (e.g. by increasing the number of message passing steps for larger \(\Delta t\) or connecting the \(n^{\text{th}}\)-closest neighbors) there is no limitation on the time step which can be used to train the model.

Regarding the second constraint, note that the GNS still performs better than the sheet model at equivalent simulation steps for out-of-training distribution \(v_{th}\) values. Additionally, high-fidelity simulations for larger values of \(v_{th}\) can be generated in order to fine-tune or retrain models for a broader dynamic range.

## 6 Recovering Known Kinetic Plasma Processes

In order to provide stronger evidence of the GNS generalization capabilities, we now showcase a broad range of known kinetic plasma processes that the simulator is able to recover. These examples, present in both the original sheet model benchmarks [22, 23, 24, 25, 26] and other kinetic codes benchmarks [2, 3, 27, 34], aim to demonstrate the capability of the GNS to simulate collective behavior in accordance with known kinetic theory. An important point to stress is that the surrogate simulator was not explicitly trained to

Figure 8: Run-time of the Sheet Model (SM) _vs._ the GNS for systems of \(10^{3}\) sheets moving on a periodic box over the period of a single plasma oscillation (\(t_{max}=2\pi\ \omega_{p}^{-1}\)). Values shown are averages over 5 simulations. For a fairer comparison among different time resolutions, the rollout data is only saved to a buffer every \(\Delta t=10^{-1}\ \omega_{p}^{-1}\). The just-in-time compilation time for the GNS is not included since it is a fixed cost that does not change for longer simulations (where it is considerably diluted). It amounts to \(t_{JIT}=1.2\) s for \(\Delta t_{GNS}=10^{-2}\ \omega_{p}^{-1}\) and \(t_{JIT}=2.3\) s for \(\Delta t_{GNS}=10^{-1}\ \omega_{p}^{-1}\). The results demonstrate that the SM run-time increases with \(v_{th}\) (higher number of crossings) while the GNS run-time is constant. By comparing these results with those of figure 7 it is observed that the GNS is faster than the SM at equivalent energy variation rates.
reproduce these effects. The GNN only learned to (correctly) model single-step updates over a reduced system size (10 sheets). However, when we apply it to larger system sizes and longer time durations, we observe the emergence of the expected kinetic collective plasma dynamics.

The results presented hereafter are produced using the GNN trained for \(\Delta t=10^{-1}\)\(\omega_{p}^{-1}\) which showcased the best energy conservation capabilities (Model #4 in Appendix F). The same collective plasma dynamics are recovered for the equivalent GNN models trained using different random seeds, and those trained using the larger time step \(\Delta t=10^{-2}\)\(\omega_{p}^{-1}\).

### Plasma Thermalization

In [22, 25], Dawson demonstrated that, independently of the initial velocity distribution of the sheets, it is expected that over time the system will move towards thermal equilibrium, and that this happens due to crossings/collisions involving more than 2 sheets [25] (cf. Section 2 for a discussion on the physics of \(n=2\) crossing/elastic collisions and why they do not modify the distribution function). The distribution function of the sheet velocities is expected to converge to a normal distribution whose standard deviation corresponds to the thermal velocity of the plasma.

We demonstrate this behavior by performing 50 simulations of systems consisting of \(10^{3}\) sheets with initial velocities randomly sampled from a uniform distribution (\(v\ \in\ [-5,5]\ \delta\cdot\omega_{p}\)). We provide snapshots of the evolution of the distribution function (averaged over simulations) in figure 9.

It is clear that the system does indeed thermalize, and that the measured thermal velocity \(v_{th}=2.671\ \delta\cdot\omega_{p}\) is in accordance with the expected value \(v_{th}=2.679\ \delta\cdot\omega_{p}\). The latter is computed according to \(v_{th}^{2}=1/3\ v_{max}^{2}r_{kin}\)[23], where \(v_{max}\) corresponds to the initial uniform distribution maximum value, and \(r_{kin}\) represents the ratio of the available kinetic energy with respect to the total energy of the system (estimated by averaging over time steps and simulations) since a percentage of the initial kinetic energy is deposited in the fields.

Additionally, using a diagnostic similar to the one introduced by Liang et al. [35], figure 9 demonstrates that there is a steep increase in the entropy (\(S\)) of the system until \(t\approx 1.25\ \omega_{p}^{-1}\). This increase is associated with the establishment of correlations between sheets as crossings start to occur, and the length of this time interval is actually independent of the initial velocity range [25].

### Debye Shielding

Another fundamental property of plasmas is their quasi-neutrality [1], i.e. on a macroscopic scale the overall charge density of positive and negative particles will cancel out. However, within local regions of characteristic length \(\lambda_{D}=v_{th}/\omega_{p}\) (referred to as the Debye length) the local electric fields generated by a charged particle will not be fully screened by the oppositely charged particles. We expect to observe the same behavior
Figure 9: Evolution of the velocity and displacement from the equilibrium position density functions, and the entropy (\(S\)) of the system. Histograms represent the density function ensemble averaged over 50 simulations at the corresponding time. The entropy variation of the different phase-space components (\(\xi\), \(v\)) is obtained using a diagnostic similar to the one implemented by Liang et al. [35]. For the calculation of the distribution functions, we discretized the (\(\xi\), \(v\)) phase-space for the range \(\xi\in[-6.5,6.5]\,\,\delta\) and \(v\in[-12.6,12.6]\,\,\delta\cdot\omega_{p}\) using 51 bins along each axis. These results demonstrate that the GNS is capable of correctly modeling the process of plasma thermalization from a non-equilibrium state, with the thermal velocity of the system in equilibrium \(v_{th}=2.671\,\,\delta\cdot\omega_{p}\) (measured by fitting a Gaussian to the final distribution) in excellent agreement with the theoretical prediction \(v_{th}=2.679\,\,\delta\cdot\omega_{p}\).
for the sheet model. More precisely, the density of sheets at a certain distance from a test position is expected to follow [22]:

\[n(x)=n_{0}\left(1-\frac{\delta}{2\lambda_{D}}e^{-|x|/\lambda_{D}}\right). \tag{9}\]

To test the GNS, we initialize systems of \(10^{4}\) sheets following different initial thermal distributions (\(v_{th}=[1.5,\ 2.5,\ 5.0]\ \delta\cdot\omega_{p}\)). The simulations are run for \(t_{max}=80\tau_{relax}\), where \(\tau_{relax}=\sqrt{2\pi}\lambda_{D}/\delta\ \omega_{P}^{-1}\) is an estimate of the relaxation time of the system [22], i.e. the time it takes for the system to forget its current state.

To compute the sheet density profiles shown in figure 10, we follow a similar procedure as Dawson [22]. We choose a set of equally spaced test sheets and, for each of them, measure the number of neighboring sheets within a pre-defined range of increasing distances. In our case, we compute the number of sheets within a distance \(d\in\,]0.2i,0.2(i+1)]\lambda_{D}\) up to \(3\lambda_{D}\) (\(i=15\)). We repeat this procedure for every \((3\lambda_{D}/\delta)\)-th sheet, over multiple independent time steps (\(t_{j}=j\cdot\tau_{relax}\) for \(j>0\)). The counts are then averaged over the number of test sheets, and time steps.

It is clear from the results presented in figure 10, that the expected behavior is recovered.

### Electrostatic fluctuations

Although the plasma is in thermal equilibrium, there are constant exchanges of energy between the sheets and the (electrostatic) waves propagating inside the plasma. This leads to the appearance of electrostatic fluctuations, with an average power spectrum that follows [2]:

\[\frac{\langle E^{2}(k)\rangle}{8\pi}\ =\frac{k_{B}T}{2L\left(1+k^{2}\lambda_{D}^ {2}\right)} \tag{10}\]

Figure 10: Example of Debye shielding for systems with different Debye lengths. The GNS can correctly recover the expected density profiles for all the tested scenarios.
where \(k\) represents the wave vector, \(k_{B}\) the Boltzmann constant, \(T\) the plasma temperature (\(k_{B}T=mv_{th}^{2}\)), and \(\langle\cdot\rangle\) the time average.

In figure 11 we recover this spectrum for a system of \(10^{3}\) sheets with \(\lambda_{D}=5\ \delta\). We make use of the ergodic theorem [36] to compute the statistical average of the power spectrum by averaging over independent time steps (separated by \(\Delta t=\tau_{relax}\)).

### Drag on a Fast Sheet

A fast sheet (\(v_{sheet}\gg v_{th}\)) moving through the plasma is expected to feel a constant drag given by [22]:

\[\frac{dv}{dt}=-\frac{\omega_{p}^{2}\delta}{2}. \tag{11}\]

This drag is independent of the velocity of the sheet and is caused by the excitation of a electrostatic wake on the rear of the fast sheet, i.e. the sheet transfers energy to the electrostatic wake.

In figure 12 we demonstrate this behavior. The results were obtained by performing simulations of periodic systems of 100 sheets with \(\lambda_{D}=5\ \delta\) (\(v_{th}=5\ \delta\cdot\omega_{p}\)) over a period of \(t_{max}=5\ \omega_{p}^{-1}\). For each simulation, we set the initial velocity of the first sheet to \(v_{0}=\pm\alpha v_{th}\) and track its evolution over time. We then average over simulations, 1000 for each initial sheet velocity (accounting for the sign change).

Figure 11: Electric field power spectrum for a system in thermal equilibrium. The power spectrum for the last time step and the temporal average (computed over relaxation periods \(\Delta t=\tau_{relax}\approx 13\ \omega_{p}^{-1}\)) are shown. The time-averaged power spectrum retrieved from the GNS simulation matches the theoretical curve, thus demonstrating that it correctly models the electrostatic fluctuations around thermal equilibrium.
### Landau Damping

While fast sheets are able to excite an electrostatic wake in their rear, the resulting electrostatic wake can also accelerate sheets moving slightly slower than its phase velocity [1, 22]. Electrostatic modes are therefore self-consistently generated by particles moving close to its phase velocity, while being damped by particles moving slightly slower. However, since a plasma in thermal equilibrium follows a Maxwellian distribution in velocity space, there exist on average more particles moving faster than the wave, than those moving slower. Therefore, on average, the modes will be damped. This mechanism is known as Landau damping [1] and is an inherently collisionless kinetic process that the sheet model has been shown to recover [22].

For a given mode \(m\), with a wavelength \(\lambda_{m}=2L/m\) and wave vector \(k_{m}=2\pi/\lambda_{m}\), we can compute its wave frequency and damping time by finding numerically the roots of the dispersion relation for \(k=k_{m}\)[37]:

\[1=\frac{\omega_{p}^{2}}{k^{2}}\int_{-\infty}^{\infty}\frac{\partial\hat{f}_{0} /\partial v}{v-(\omega/k)}dv \tag{12}\]

where \(\hat{f}_{0}\) corresponds to the distribution function in velocity space. The solution will have both a real and imaginary part. The real part corresponds to the wave angular frequency, while the imaginary part corresponds to the inverse of the damping time.

To reproduce the expected damping behavior we follow a similar procedure to that of Dawson [22]. We produce 50 simulations, each with a duration of \(t_{max}=500\)\(\omega_{p}^{-1}\), of thermal plasmas consisting of \(10^{3}\) sheets for \(\lambda_{D}=5\)\(\delta\) and reflecting boundaries. For each simulation time step we compute the mode "amplitude" \(A_{m}\) and its rate of change

Figure 12: Average drag on fast sheets of different initial velocities (\(v_{0}\gg v_{th}=5\)\(\delta\) \(\cdot\)\(\omega_{p}\)). The GNS recovers the expected drag felt by the fast sheets independently of their initial velocity and propagation direction.
\(\dot{A}_{m}\) using the cross-correlation:

\[\begin{split} A_{m}^{t}&=\frac{2}{N}\sum_{i=0}^{N-1} \left(x_{i}^{t}-x_{eq_{i}}^{t}\right)\ \sin\left(\frac{m\pi}{N}\left(i+\frac{1}{2}\right)\right)\\ \dot{A}_{m}^{t}&=\frac{A_{m}^{t}-A_{m}^{t-1}}{\Delta t }\end{split} \tag{13}\]

where the index \(i\) indicates the relative ordering of the sheets in the box.

We then collect trajectories of equal time length every time the mode crosses the region of phase-space \((A_{m},\dot{A}_{m})\) defined by a ring of radius \(R\) and thickness \(dR\) (\(dR\ll R\)). Finally, we rotate the trajectories so that they all start on the same position in phase-space, and compute their average.

In figure 13 we showcase the results obtained for several modes. It is possible to see that, although for some trajectories the mode is still growing, on average it decreases according to the expected damping time. This demonstrates that the GNS is capable of correctly modeling Landau damping, an inherently kinetic mechanism associated with the collective collisionless dynamics of a plasma.

### Two-stream Instability

As a final example, we show the two-stream instability in the cold beam regime [1, 24]. For this scenario, two counter-propagating beams with velocities \(\pm v_{0}\) and no energy spread excite a wave that grows exponentially until all particles are trapped inside the electric field, at which point the instability saturates.

Figure 13: Damping of different modes. The initial mode amplitudes are normalized to \(R=0.3\ \delta\). Individual trajectories length is \(\Delta t_{traj}=50\ \omega_{p}^{-1}\). Damping times are respectively \(\tau_{28}=10.25\ \omega_{p}^{-1}\), \(\tau_{30}=7.97\ \omega_{p}^{-1}\), and \(\tau_{32}=6.41\ \omega_{p}^{-1}\). The agreement between the theoretical curves and the average mode trajectories demonstrates that the GNS is capable of correctly modeling Landau damping, an inherently kinetic mechanism.
From linear theory [1, 24], we expect that for two-cold beams with density \(n_{beams}=n_{0}/2\), the fastest growing mode will correspond to \(k=\sqrt{3/8}\cdot\omega_{p}/v_{0}\) with a corresponding growth rate of \(\gamma=\omega_{p}/\sqrt{8}\). Therefore, to excite mode \(m\) we need to set \(v_{0_{m}}=\sqrt{3/8}\cdot\omega_{p}\delta/m\pi\cdot N_{sheets}\). Furthermore, the number of sheets per wavelength of this mode is given by \(N_{\lambda_{m}}=\lambda_{m}/\delta=2/m\cdot N_{sheets}\). Note that both \(v_{0_{m}}\) and \(N_{\lambda_{m}}\) are proportional to the number of sheets used. Therefore, to excite a mode whose wavelength can be resolved by a significant amount of sheets, we need to increase \(v_{0}\) proportionally.

For a system of \(10^{4}\) sheets and \(m=4\), we obtain \(v_{0_{m}}=486\)\(\delta\cdot\omega_{p}\) and \(N_{\lambda_{m}}=5\times 10^{3}\). The chosen velocity is considerably out of the training data range, therefore we expect the energy loss of the GNS to be slightly higher than the one observed for the training scenarios (as shown in figure 7). Nonetheless, we will see that the GNS is still able to capture the macrophysics. Additionally, these (significantly) out-of-training distribution values allow us to once again highlight the potential of the GNS to achieve better energy conservation than the sheet model algorithm at comparable simulation time steps.

In figure 14 we provide a comparison between the evolution of the phase-space and the potential energy for the sheet model at high (\(\Delta t=10^{-3}\)\(\omega_{p}^{-1}\)) and low (\(\Delta t=10^{-1}\)\(\omega_{p}^{-1}\)) temporal resolution and the GNS at the same low temporal resolution. The energy variation during the full simulation was approximately 0.1% for the sheet model at \(\Delta t=10^{-3}\)\(\omega_{p}^{-1}\) and 73% at \(\Delta t=10^{-1}\)\(\omega_{p}^{-1}\), while for the GNS it was approximately 2%. These values are in accordance with what was measured in figure 7.

It is observed that the GNS recovers similar macrophysics when compared to the higher temporal resolution simulation (which we consider as a good approximation of the ground truth) during the linear phase and up to the saturation time. An extra diffusion in phase-space is observed for the GNS, which is associated with the aforementioned higher energy variation. This is expected since the GNN is not capable of correctly resolving crossings involving more than \(2M+1=11\) sheets, and, on average, a sheet moving with \(v_{0_{m}}=486\)\(\delta\cdot\omega_{p}\) should cross \(\approx 49\) neighbors just in the first timestep (since when the simulation starts all sheets are equally spaced by \(\delta\) and, on average, half of the nearest neighbors are counter-propagating with the same absolute velocity). Nonetheless, the overall phase-space structure and growth rates are similar, which provides further support for the generalization capabilities of the model. Additionally, when using the sheet model at a comparable time resolution the results are strikingly different, which reinforces the conclusions derived in Section 5.2 that the GNS is capable of significantly improving simulation results at larger time steps.

Finally, we provide in Appendix G further comparisons for higher order modes (lower \(v_{0_{m}}\)), larger simulation steps (\(\Delta t=10^{-2}\)\(\omega_{p}^{-1}\)), as well as results for the remaining GNNs trained with different seeds (with worse energy conservation capabilities at high \(v_{th}\)). These results further support the claim that the GNS is consistently able to correctly model the overall dynamics of the instability.
## 7 Conclusions

In this work, we demonstrated that graph neural network-based simulators are capable of fully replacing a one dimensional kinetic plasma physics simulator. By introducing domain knowledge into the graph representation and the overall structure of the simulator, we showed that the GNS is capable of generalizing to a broad range of scenarios and is limited only by its training data distribution and a fixed simulation step. Furthermore, we showed that the GNS, when trained on high-fidelity data, conserves energy better than the synchronous algorithm of the sheet model, for larger simulation steps.

In future work, the accuracy of the simulator for higher thermal velocity values

Figure 14: Comparison of the phase-space and potential energy evolution for the Sheet Model (SM) and the GNS in the two counter-propagating cold beams scenario. The GNS is able to recover the same macrophysics as the sheet model at a significantly higher temporal resolution and conserves the system energy significantly better than the SM at an equivalent time resolution.
can be improved by generating additional ground truth data, ideally using the asynchronous version of the sheet model, which is guaranteed to perfectly resolve crossings. Additionally, although not explored, the developed simulator is fully differentiable, which opens the way to explore gradient-based optimization strategies for the discovery of new physics of interest [38].

This work also indicates that it would be possible to accurately model the standard PIC loop. However, we do not believe that this would result in meaningful computational gains since the standard PIC loop implemented in modern architectures is extremely optimized and does not suffer from the issues that are present in the sheet model scenario (e.g. time-consuming serial routines). The usage of graph neural networks can be relevant to model multi-scale plasma dynamics or extra physics modules. This will be explored in future publications.

## Data availability statement

The code used for this project is made available through a git repository: (to be disclosed). The training and test data, as well as a set of the final model weights are available in a zenodo repository: (to be disclosed)

## Conflict of interest

The authors declare no conflict of interest.

## Acknowledgments

The authors would like to thank P. J. Bilbao, T. Grismayer, A. Joglekar, and E. P. Alves for helpful discussions. This work was supported by the FCT (Portugal) under the Project No 2022.02230.PTDC (X-MASER) and PhD Fellowship Grant 2022.13261.BD and has received funding from the European Union's H2020 programme through the project IMPULSE (grant agreement No 871161). The graphics processing units (GPUs) used in this work were donated by NVIDIA Corporation.
