# Do Graph Neural Networks Dream of Landau Damping? Insights from Kinetic
Simulations of a Plasma Sheet Model

Diogo D Carvalho1, Diogo R Ferreira2 and Luís O Silva1 1 GoLP/Instituto de
Plasmas e Fusão Nuclear, Instituto Superior Técnico, Universidade de Lisboa,
1049-001 Lisbon, Portugal  2 Instituto de Plasmas e Fusão Nuclear, Instituto
Superior Técnico, Universidade de Lisboa, 1049-001 Lisbon, Portugal
<EMAIL_ADDRESS>
###### 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 [19, 15, 16]. These approaches have been successfully applied to
fluid [12, 13], rigid body [13, 20], and charged particle dynamics [21, 15,
16]. 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, 24, 25, 26, 23]. 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.

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].

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$ (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.

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 A.

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.

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.

Figure 3: Schematic representations of the synchronous Electrostatic Sheet
Model algorithm [23] and the Graph Network Simulator used.

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

###  3.1 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:

$\eqalign{\mathbf{n}^{t}_{i}&=\left[\xi^{t}_{i},v^{t}_{i}\right]\cr\mathbf{r}^{t}_{ij}&=\left[x^{t}_{j}-x^{t}_{i}\right]}$
(2)

where $\xi_{i}^{t}$ corresponds to the displacement of the $i$th sheet from
its equilibrium position $\left(x^{t}_{i}-x^{t}_{eq_{i}}\right)$, and
$v^{t}_{i}$ 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 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).

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.

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.

###  3.2 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:

$\eqalign{\mathbf{v}_{i}=\varepsilon^{v}\left(\mathbf{n}_{i}\right)\cr\mathbf{e}_{ij}=\varepsilon^{e}\left(\mathbf{r}_{ij}\right)}$
(3)

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:

$\eqalign{\mathbf{e}_{ij}^{m+1}&=\phi^{e}\left(\mathbf{e}^{m}_{ij},\mathbf{v}^{m}_{i},\cdot\mathbf{v}^{m}_{j}\right)\cr\overline{\mathbf{e}}_{r_{i}}^{m+1}&=\sum_{j\in\mathcal{N}(i)}\mathbf{e}^{m+1}_{ij}\cr\overline{\mathbf{e}}_{s_{i}}^{m+1}&=\sum_{j\in\mathcal{N}(i)}\mathbf{e}^{m+1}_{ji}\cr\mathbf{v}_{i}^{m+1}&=\phi^{v}\left(\overline{\mathbf{e}}^{m+1}_{r_{i}},\overline{\mathbf{e}}^{m+1}_{s_{i}},\mathbf{v}^{m}_{i}\right)}$
(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)$ (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^{t}_{i}=(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$ {LinearLayer()
$\rightarrow$ ReLU $\rightarrow$ 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 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.

###  3.3 Position and Velocity Update

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

$\eqalign{\tilde{v}^{t+1}&=v^{t}+a^{t}\Delta
t\cr\tilde{x}^{t+1}&=x^{t}+\tilde{v}^{t+1}\Delta t\ .}$ (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$.

###  4.1 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
$\left(\Delta\varepsilon/\varepsilon_{0}=10^{-6}\right)$ 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.

###  4.2 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 E.

###  4.3 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.

###  5.1 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

$\eqalign{\mathrm{MAE}=&\
\frac{1}{N}\sum_{i}^{N}\left|x_{i}^{GNS}-x_{i}^{True}\right|\cr\mathrm{EMD}=&\
\min\frac{1}{N}\sum_{i}^{N}\sum_{j}^{N}\left|x_{i}^{GNS}-x_{j}^{True}\right|\cdot
c_{ij},\cr&\ \mathrm{subject\ to}\ \sum_{i}c_{ij}=1\ ,\ \sum_{j}c_{ij}=1\ ,\
\mathrm{and}\ c_{ij}\in\\{0,1\\}}$ (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.

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.

Periodic Boundary - $\Delta t_{GNS}=10^{-2}~{}\omega_{p}^{-1}$
Reflecting Boundary - $\Delta t_{GNS}=10^{-1}~{}\omega_{p}^{-1}$

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.

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.

###  5.2 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).$
(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.

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 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.

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^{max}_{train}$ (more
details in 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 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.

###  5.3 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\\!\cdot\\!\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.

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.

###  5.4 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, 24, 25, 26, 23] and other kinetic codes benchmarks
[2, 34, 3, 27], 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
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 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}$.

###  6.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.

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\left[-6.5,6.5\right]~{}\delta$ and
$v\in\left[-12.6,12.6\right]~{}\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}$.

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].

###  6.2 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 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).$
(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.

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.

###  6.3 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)}$ (10)

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}$).

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.

###  6.4 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}.$ (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 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.

###  6.5 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$ (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 $\dot{A}_{m}$ using the cross-correlation:

$\eqalign{A^{t}_{m}=\frac{2}{N}\sum_{i=0}^{N-1}\left(x^{t}_{i}-x^{t}_{eq_{i}}\right)\
\sin\left(\frac{m\pi}{N}\left(i+\frac{1}{2}\right)\right)\cr\dot{A}_{m}^{t}=\frac{A_{m}^{t}-A_{m}^{t-1}}{\Delta
t}}$ (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.

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.

###  6.6 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.

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.

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.

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 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 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.

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.

## References

## References

  * [1] Chen F F et al. 1984 Introduction to plasma physics and controlled fusion vol 1 (Springer)
  * [2] Dawson J M 1983 Reviews of modern physics 55 403
  * [3] Birdsall C K and Langdon A B 2004 Plasma physics via computer simulation (CRC press)
  * [4] Fonseca R A, Silva L O, Tsung F S, Decyk V K, Lu W, Ren C, Mori W B, Deng S, Lee S, Katsouleas T et al. 2002 Osiris: A three-dimensional, fully relativistic particle in cell code for modeling plasma based accelerators International Conference on Computational Science (Springer) pp 342–351
  * [5] Kube R, Churchill R and Sturdevant B 2021 arXiv preprint arXiv:2110.12444
  * [6] Aguilar X and Markidis S 2021 A deep learning-based particle-in-cell method for plasma simulations 2021 IEEE International Conference on Cluster Computing (CLUSTER) (IEEE) pp 692–697
  * [7] Badiali C, Bilbao P J, Cruz F and Silva L O 2022 Journal of Plasma Physics 88 895880602
  * [8] Chen R T, Rubanova Y, Bettencourt J and Duvenaud D K 2018 Advances in neural information processing systems 31
  * [9] Li Z, Kovachki N B, Azizzadenesheli K, Bhattacharya K, Stuart A and Anandkumar A 2021 Fourier neural operator for parametric partial differential equations International Conference on Learning Representations
  * [10] Karniadakis G E, Kevrekidis I G, Lu L, Perdikaris P, Wang S and Yang L 2021 Nature Reviews Physics 3 422–440
  * [11] Brandstetter J, Worrall D E and Welling M 2022 Message passing neural PDE solvers International Conference on Learning Representations
  * [12] Sanchez-Gonzalez A, Godwin J, Pfaff T, Ying R, Leskovec J and Battaglia P 2020 Learning to simulate complex physics with graph networks International Conference on Machine Learning (PMLR) pp 8459–8468
  * [13] Pfaff T, Fortunato M, Sanchez-Gonzalez A and Battaglia P 2020 Learning mesh-based simulation with graph networks International Conference on Learning Representations
  * [14] Zhong Y D, Dey B and Chakraborty A 2021 Benchmarking energy-conserving neural networks for learning dynamics from data Learning for dynamics and control (PMLR) pp 1218–1229
  * [15] Satorras V G, Hoogeboom E and Welling M 2021 E (n) equivariant graph neural networks International conference on machine learning (PMLR) pp 9323–9332
  * [16] Brandstetter J, Hesselink R, van der Pol E, Bekkers E J and Welling M 2022 Geometric and physical quantities improve e(3) equivariant message passing International Conference on Learning Representations
  * [17] Prantl L, Ummenhofer B, Koltun V and Thuerey N 2022 Guaranteed conservation of momentum for learning particle-based fluid dynamics Advances in Neural Information Processing Systems vol 35 ed Koyejo S, Mohamed S, Agarwal A, Belgrave D, Cho K and Oh A (Curran Associates, Inc.) pp 6901–6913
  * [18] Battaglia P W, Hamrick J B, Bapst V, Sanchez-Gonzalez A, Zambaldi V, Malinowski M, Tacchetti A, Raposo D, Santoro A, Faulkner R et al. 2018 arXiv preprint arXiv:1806.01261
  * [19] Bronstein M M, Bruna J, Cohen T and Veličković P 2021 arXiv preprint arXiv:2104.13478
  * [20] Allen K R, Rubanova Y, Lopez-Guevara T, Whitney W F, Sanchez-Gonzalez A, Battaglia P and Pfaff T 2023 Learning rigid dynamics with face interaction graph networks The Eleventh International Conference on Learning Representations
  * [21] Cranmer M, Sanchez Gonzalez A, Battaglia P, Xu R, Cranmer K, Spergel D and Ho S 2020 Advances in Neural Information Processing Systems 33 17429–17442
  * [22] Dawson J 1962 The Physics of Fluids 5 445–459
  * [23] Dawson J 1970 Methods Computational Physics 9 1–28
  * [24] Dawson J M 1962 Some Investigations of Plasma Instabilitiies in One-dimensional Plasmas vol 152 (US Atomic Energy Commission, Division of Technical Information)
  * [25] Dawson J M 1964 The Physics of Fluids 7 419–425
  * [26] Dawson J M and Shanny R 1968 The Physics of Fluids 11 1506–1523
  * [27] Gravier E, Drouot T, Lesur M, Guillevic A, Lo-Cascio G, Moritz J, Escande D and Manfredi G 2023 Physics of Plasmas 30
  * [28] Godwin J, Keck T, Battaglia P, Bapst V, Kipf T, Li Y, Stachenfeld K, Veličković P and Sanchez-Gonzalez A 2020 Jraph: A library for graph neural networks in jax. URL http://github.com/deepmind/jraph
  * [29] Topping J, Giovanni F D, Chamberlain B P, Dong X and Bronstein M M 2022 Understanding over-squashing and bottlenecks on graphs via curvature International Conference on Learning Representations
  * [30] Harris C R, Millman K J, van der Walt S J, Gommers R, Virtanen P, Cournapeau D, Wieser E, Taylor J, Berg S, Smith N J, Kern R, Picus M, Hoyer S, van Kerkwijk M H, Brett M, Haldane A, del Río J F, Wiebe M, Peterson P, Gérard-Marchant P, Sheppard K, Reddy T, Weckesser W, Abbasi H, Gohlke C and Oliphant T E 2020 Nature 585 357–362 URL https://doi.org/10.1038/s41586-020-2649-2
  * [31] Bradbury J, Frostig R, Hawkins P, Johnson M J, Leary C, Maclaurin D, Necula G, Paszke A, VanderPlas J, Wanderman-Milne S and Zhang Q 2018 JAX: composable transformations of Python+NumPy programs URL http://github.com/google/jax
  * [32] Hennigan T, Cai T, Norman T and Babuschkin I 2020 Haiku: Sonnet for JAX URL http://github.com/deepmind/dm-haiku
  * [33] Villani C et al. 2009 Optimal transport: old and new vol 338 (Springer)
  * [34] Opher M, Silva L O, Dauger D E, Decyk V K and Dawson J M 2001 Physics of Plasmas 8 2454–2460
  * [35] Liang H, Cassak P A, Servidio S, Shay M A, Drake J F, Swisdak M, Argall M R, Dorelli J C, Scime E E, Matthaeus W H et al. 2019 Physics of Plasmas 26
  * [36] Huang K 2008 Statistical mechanics (John Wiley & Sons)
  * [37] Jackson J 1960 Journal of Nuclear Energy. Part C, Plasma Physics, Accelerators, Thermonuclear Research 1 171
  * [38] Joglekar A S and Thomas A G 2022 Journal of Plasma Physics 88 905880608

##  Appendix A Collisions vs. Crossings

It is possible to model the sheet crossings as binary collisions (as
previously illustrated in figure 2). This was the strategy initially explored
in this work since the Graph Network simulator proposed by Sanchez-Gonzalez et
al. [12] was introduced for 2D/3D simulations where particle collisions occur.
However, the accuracy of the models trained in this fashion was significantly
worse. In this subsection, we explain why this is happening.

When trying to predict dynamics considering binary collisions, the range of
target accelerations changes significantly (see figure 15) since at the moment
of collision, particles feel a force that is orders of magnitude larger than
the one felt during their oscillatory motion. This broadening makes it harder
for the network to accurately model the full dynamic range of accelerations.

Figure 15: Target acceleration distributions in the training set data when
considering collisions or crossings. The range of target accelerations is
considerably smaller for the crossing scenario. This makes the training
procedure significantly easier.

Additionally, in the time step where sheets collide, their respective finite
difference velocities significantly decrease, and information is lost about
the initial momentum of the particles (especially for larger $v$ and small
$\Delta t$). Therefore, to allow for correct momentum transfer, it is
necessary to include the previous $C$ step velocities in the node
representation
$\textbf{n}^{t}_{i}=\left[\xi^{t}_{i},v^{t}_{i},...,v^{t-C+1}\right]$ as in
Sanchez-Gonzalez et al. [12]. This introduces a new hyperparameter that
requires tuning and whose value depends on the frequency of collisions.

For larger collision frequencies, it then becomes hard to correctly model
momentum transfer since the model: a) needs to predict a larger dynamics range
of accelerations; b) needs to take into account dynamics across multiple time
steps and larger neighboring regions. In our case, this resulted in
considerably larger neural networks architectures, whose accuracy was still
significantly worse than models trained to predict crossing dynamics.

To illustrate this, we showcase in figure 16 a comparison between the rollout
accuracy difference presented in the main body of the paper for $\Delta
t=10^{-1}~{}\omega_{p}^{-1}$, and an equivalent one trained for collisional
dynamics using $C=5$ (same neural network architecture and number of message
passing steps). It is clear that the accuracy of the simulator degrades
significantly when attempting to predict collisions.

Figure 16: Accuracy of equivalent models trained to predict crossing or
collisional dynamics. Both models use $\Delta t=10^{-1}~{}\omega_{p}^{-1}$.
The rollout error of the GNS model trained for crossing dynamics is
significantly lower. This is the main reason why we opted to model sheet
dynamics considering crossings in the main body of the paper.

It is possible that for larger simulation steps, the problems stated above
become less pronounced since the effect of collisions is smoothed out over
larger time steps (leading to a smaller dynamic range of target accelerations
as observed in figure 15, when moving from $\Delta
t=10^{-2}~{}\omega_{p}^{-1}$ to $\Delta t=10^{-1}~{}\omega_{p}^{-1}$).
However, in the 1D sheet model scenario, our results indicate that modeling
sheet interactions as crossings is the best and most natural option in terms
of model/graph size and accuracy, while also allowing the tracking of
particular sheets (useful, for example, for the fast sheet diagnostic).

Finally, it is important to highlight that when using reflecting boundaries
the sheet trajectories contain “collisions” with the wall even when
considering crossing dynamics. This is the reason why we first train our
models on periodic boundaries and only at test time demonstrate that it is
possible to simulate reflecting boundaries (and not the opposite). It would be
possible to bypass this issue, while maintaining the same graph structure if,
for example, we introduced an extra preprocessing step during training that
masks out time steps for which collisions with the wall happen (or correct the
target accelerations). However, this removes any interaction with the wall
from the training data (defeating the purpose of having reflecting boundaries
in the first place) and would require us to store extra information during the
generation of the training simulations.

##  Appendix B Parameter Scans

The main goal of this work was to understand if we can build a surrogate
simulator capable of recovering kinetic plasma processes with acceptable run-
time. We did not perform extensive hyperparameter tuning and instead opted for
a coarse scan. For the models showcased in the main body of the text, we
prioritized simplicity in favor of smaller gains in performance. Additionally,
we decided to use the same architecture for both training simulation steps,
except the number of message-passing steps, as this parameter is related to
the crossing frequency.

The rollout accuracy on the validation set for models trained with different
hyper-parameters for $\Delta t=10^{-1}~{}\omega_{p}^{-1}$ is shown in figure
17. The most impactful parameter is the number of message-passing steps.
Increasing the number of layers and hidden dimensions size of the MLP networks
does not affect performance significantly and might lead the model to overfit
the training data earlier in the training process. Similarly, introducing
extra complexity by using MLPs for the encoder/decoder instead of linear
transformation does not result in significant performance changes.

We also checked if the introduction of residual connections between message-
passing blocks would improve performance for larger numbers of message-passing
steps, but it did not provide significant improvements.

Figure 17: Parameter scan results for GNS trained with $\Delta
t=10^{-1}~{}\omega_{p}^{-1}$. The rollout EMD values were computed on the
validation set. The final chosen hyperparameter values are highlighted in red.
The most important parameter to tune is the number of message-passing steps,
whose optimal value depends on the expected sheet crossing frequency.

##  Appendix C Equivariant vs Non-Equivariant Architectures

The Graph Network block presented in Section 3.2 can be altered to use solely
the sent edges in the node update mechanism, i.e.:

$\mathbf{v}_{i}^{m+1}=\phi^{v}\left(\overline{\mathbf{e}}^{m+1}_{r_{i}},\mathbf{v}^{m}_{i}\right)$
(14)

which corresponds to the update mechanism presented in the original work by
Battaglia et al. [18].

Additionally, one can make the full architecture (1D) rotationally equivariant
by modifying the architecture as follows. The encoder should now be given by:

$\eqalign{\mathbf{v}_{i}=\varepsilon^{v}\left(\mathbf{n}_{i}\cdot\mathrm{sgn}\left(\xi_{i}\right)\right)\cr\mathbf{e}_{ij}=\varepsilon^{e}\left(\mathbf{r}_{ij}\cdot\mathrm{sgn}\left(\xi_{i}\right)\right)}$
(15)

where $\mathrm{sgn}(\cdot)$ represents the sign function. In this case the
inputs become rotationally invariant, and , therefore, the latent
representations $\mathbf{v}_{i}$ and $\mathbf{e}_{ij}$ are also guaranteed to
be invariant. The multiplication by $\mathrm{sgn}\left(\xi_{i}\right)$ can be
interpreted as a change of coordinate system (or a projection since in 1D the
$\mathrm{sgn}(x)\equiv x/||x||$) to a frame where the $i$th sheet is on the
right-hand side of its equilibrium position (i.e. $\xi>0$). For the case where
$\xi_{i}=0$ (which would result in $\mathrm{sgn(0)}=0$), we use instead the
velocity (i.e. $\mathrm{sgn}(v_{i})$). If the latter is also zero, we set the
value to 1. These extra steps are required to ensure that the architecture
maintains its equivariance properties.

Similarly, the processor is also made rotationally invariant with a slight
modification:

$\eqalign{\mathbf{e}_{ij}^{m+1}&=\phi^{e}\left(\mathbf{e}^{m}_{ij},\mathbf{v}^{m}_{i},\mathrm{sgn}\left(\xi_{i}\xi_{j}\right)\cdot\mathbf{v}^{m}_{j}\right)\cr\overline{\mathbf{e}}_{r_{i}}^{m+1}&=\sum_{j\in\mathcal{N}(i)}\mathbf{e}^{m+1}_{ij}\cr\mathbf{v}_{i}^{m+1}&=\phi^{v}\left(\overline{\mathbf{e}}^{m+1}_{r_{i}},\mathbf{v}^{m}_{i}\right)}$
(16)

since $\mathbf{e}^{0}_{ij}$, $\mathbf{v}^{0}_{i,j}$, and
$\mathrm{sgn}\left(\xi_{i}\xi_{j}\right)$ are invariant to this
transformation, making all subsequent updates invariant as well. In this case,
the multiplication of the sender node latent vector $\mathbf{v}^{m}_{j}$ by
$\mathrm{sgn}\left(\xi_{i}\xi_{j}\right)$ is required to break the symmetry
introduced by the encoder (mirrored sheets around the equilibrium position
have the same latent representation) and can be interpreted as a
transformation to the frame of the receiver latent node $\mathbf{v}^{m}_{i}$
(i.e. where $\xi_{i}>0$).

Finally, the modified decoder enforces equivariance as:

$\mathbf{y}_{i}=\delta^{v}\left(\mathbf{v}_{i}^{M}\right)\cdot\mathrm{sgn}\left(\xi_{i}\right)$
(17)

since $\mathbf{v}_{i}^{M}$ is guaranteed to be invariant under reflections.

Comparisons between the rollout accuracy and energy conservation capabilities
of the architecture used in the main body of the paper (“Non-Equivariant
Default”), and the variants that only use the sent edge information (“Non-
Equivariant Only Sent” and “Equivariant Only Sent”) are presented in figure
18, figure 19, and figure 20.

Figure 18: Comparison of rollout accuracy of different architectures. We use
the test sets described in Section 5.1 that consider periodic boundary
conditions. Mean values for each architecture are computed over 5 equivalent
models (different random seeds), and error bars represent minimum/maximum
values observed over all models. All models are trained considering a time
step $\Delta t=10^{-1}~{}\omega_{p}^{-1}$ using the same hyperparameters
described in Section 3.2 and the training procedure described in Section 4.1.
Non-Equivarient models are trained with data augmentation along both axes ($x$
and $t$) while equivariant models are only provided with data augmented along
the $t$-axis. No significant differences in rollout accuracy are observed
across the different architectures. Figure 19: Comparison of energy
conservation capabilities of different architectures for different initial
conditions. We use the same energy conservation diagnostic described in
Section 5.2. All sheets are initialized at their equilibrium position and
their initial velocities are sampled from a normal distribution (thermal),
chosen randomly from $\pm v_{beam}$ (cold beams), or all equal to $v_{0}$
(oscillation). Mean values are computed over 5 equivalent models (different
random seeds). All models are trained considering a time step $\Delta
t=10^{-1}~{}\omega_{p}^{-1}$ using the same hyperparameters described in
Section 3.2 and the training procedure described in Section 4.1. Non-
Equivariant models are trained with data augmentation along both axes ($x$ and
$t$) while equivariant models are only provided with data augmented along the
$t$-axis. Figure 20: Rollout accuracy as a function of the number of training
simulations. We consider architectures that only use the sent edges for the
latent node update. A single model is trained for each dataset size. Training
is slightly modified compared to previous results. We use a larger number of
maximum gradient updates ($2\times 10^{6}$) and early stopping with a patience
of 100 epochs. Augmented datasets include reflections along both axes ($x$ and
$t$). Overall, it is observed that there are no significant differences in
performance between the architectures.

It is observed that the rollout accuracy is similar across all architectures
(figure 18). Similarly, the models demonstrate equivalent energy conservation
within the training data thermal velocity range (figure 19). However, as one
further increases the thermal velocity of the plasma, the non-equivariant
architectures demonstrate better energy conservation capabilities, in
particular the one used in the main body of the paper. In terms of data
requirements, we do not observe relevant performance gains for a varying
number of training simulations (figure 20). There is a slight improvement by
using the equivariant architecture in the case where no data augmentation is
used. This effect disappears once data augmentation is used at train time.

There can be several reasons for the lack of performance improvement when
considering the equivariant architecture. Firstly, we are dealing with a
relatively simple one-dimensional system, and the symmetry that we are trying
to enforce might be easily learned by the network. This can explain why we do
not observe significant performance improvements in the low data regime when
compared to other works which learn to model charged-particle dynamics in 3D
using equivariant architectures [16, 15]. Further support for this thesis is
the fact that other experiments that included augmented data solely along the
$x$-axis did not cause significant changes in the performance of the non-
equivariant architecture (only using both $x$ and $t$ augmentation produced
improvements).

Secondly, it can be argued that the way we enforce the equivariance is not
ideal, since we are constraining significantly the latent node representations
with the introduction of the multiplication term
$\textrm{sgn}(\xi_{i}\xi_{j})$ in the edge update. This was the option we
found to work better. Furthermore, the performance of the equivariant and the
equivalent non-equivariant architecture (the one which uses solely the sent
edges for the node update) is basically the same across all metrics (rollout
accuracy, energy conservation, etc.) except for considerably large thermal
velocities ($v_{th}\gg v^{train}_{th}$). Therefore, we conclude that the
imposed constraints are not affecting significantly the network learning
capabilities, and the lack of improvements is mostly due to the simplicity of
the setup and symmetry.

Finally, note that we do not propose an equivariant architecture that uses the
sent edges for the node update mechanism. This is due to the fact that, in the
proposed equivariant scheme, the sent edges are computed in the frame of the
receiving node. Therefore, one would have to add an operation that transforms
them back to the sender node frame, or duplicate the amount of operations in
the network in order to store and update the (latent) edges in both frames.

At this point, it is not clear why using the sent edges in the node update of
the non-equivariant architecture improves the energy conservation capabilities
at larger thermal velocities. We do not observe any relevant difference in the
validation loss between the different non-equivariant architectures, which
leads us to believe that this effect is not due simply to a larger number of
degrees of freedom of the network which allows for a better fit. Instead, the
network seems to learn a structurally different update mechanism, which is
more resilient to higher crossing rates.

Lastly, it is important to mention that all architectures are capable of
reproducing the kinetic plasma process presented in Section 6 with the
exception of the Two-Stream Instability. The latter is only recovered using
the non-equivariant architecture presented in the main body of the paper, due
to the aforementioned improved energy conservation capabilities for larger
sheet velocities.

##  Appendix D Boundary crossings re-injection

For reflecting boundary conditions, the new positions, velocities, and
equilibrium positions of the sheets are updated according to:

$\eqalign{x^{t+1}&=\begin{cases}-\tilde{x}^{t+1}&,\textrm{ if}\
\tilde{x}^{t+1}<0\\\ \tilde{x}^{t+1}-L&,\textrm{ if}\ \tilde{x}^{t+1}\geq L\\\
\tilde{x}^{t+1}&,\textrm{ elsewhere}\end{cases}\cr
v^{t+1}&=\begin{cases}-\tilde{v}^{t+1}&,\textrm{ if}\ \tilde{x}^{t+1}<0\ \lor\
\tilde{x}^{t+1}\geq L\\\ \tilde{v}^{t+1}&,\textrm{
elsewhere}\end{cases}\cr\mathbf{x}_{eq}^{t+1}&=\mathbf{x}_{eq}^{t}}$ (18)

where $\tilde{x}$, $\tilde{v}$ represent the positions and velocities after
the ODE integration step.

For periodic boundary conditions, the update rule is:

$\eqalign{x^{t+1}&=\begin{cases}L+\tilde{x}^{t+1}&,\textrm{ if}\
\tilde{x}^{t+1}<0\\\ \tilde{x}^{t+1}-L&,\textrm{ if}\ \tilde{x}^{t+1}\geq L\\\
\tilde{x}^{t+1}&,\textrm{ elsewhere}\end{cases}\cr
v^{t+1}&=\tilde{v}^{t+1}\cr\mathbf{x}_{eq}^{t+1}&=\mathbf{x}_{eq}^{t}+(n_{left}-n_{right})\delta}$
(19)

where $n_{left}$, $n_{right}$ represent the number of particle that crossed
the left and right boundary. For both boundary conditions, $\mathbf{x}^{t+1}$
is sorted after the update (and $\mathbf{v}^{t+1}$ accordingly) so that its
indices match the equilibrium positions array $\mathbf{x}^{t+1}_{eq}$ (i.e.
the correct equilibrium position is attributed).

##  Appendix E Impact of training dataset size and data augmentation

In figure 21, we show the accuracy of equivalent models to those presented in
the main body of the paper, when trained on datasets of different sizes. For
$\Delta t=10^{-2}~{}\omega_{p}^{-1}$ the accuracy plateaus before the maximum
amount of training simulations used. The same does not happen for a $\Delta
t=10^{-1}~{}\omega_{p}^{-1}$, which exhibits a consistent power law scaling,
meaning that increasing the dataset size could still lead to performance
gains. We believe the difference in behavior occurs not only because more
training time steps are available per simulation for smaller $\Delta t$, but
also because it is easier to model the dynamics in smaller time steps since
fewer crossings per simulation step occur.

Figure 21: Model accuracy as a function of the number of training simulations.
The test set contained 100 simulations of 10 sheets moving on a periodic box
(same as in figure 5). It is observed that rollout EMD follows approximately a
power law for $\Delta t=10^{-1}~{}\omega_{p}^{-1}$ meaning that further
improvements can be achieved by increasing the dataset size. For $\Delta
t=10^{-2}~{}\omega_{p}^{-1}$ the performance has already plateaued. Overall,
the proposed data augmentation strategy is proven to be beneficial.

Additionally, the results demonstrate the impact of using data augmentation,
especially at smaller training dataset sizes. Further investigation revealed
that this improvement in performance is mainly due to the time reversal
augmentation since introducing the data reflected along the $x$-axis did not
produce significant differences. We investigated if the data requirements
could be further lowered by enforcing the symmetries expressed by the data
augmentation procedure in the network architecture. However, due to the finite
difference nature of the velocities, solely the spatial rotational symmetry
could be enforced. No significant improvements were observed for a GNN
architecture that enforced the latter, reinforcing the notion that this
symmetry is easily learned/approximated by the network. More details are
provided in C.

##  Appendix F Model Benchmark – Extra

### Performance across different seeds

In figure 22 and figure 23 we show the same metrics presented in Section 5 but
now presented for all the individual GNN models trained (instead of the
average value and min/max range over equivalent models). The results clearly
identify the existence of a consistently worse performing model for $\Delta
t=10^{-2}~{}\omega_{p}^{-1}$, both in terms of rollout accuracy and energy
conservation. We attribute this difference in performance to a worse single-
step prediction capability, which was already clear in its training
performance ($2\times$ larger validation loss than equivalent models). When
removing this model from the set considered for the energy conservation
diagnostic, the behavior for different time resolutions ($\Delta
t=10^{-1}~{}\omega_{p}^{-1}$ and $\Delta t=10^{-1}~{}\omega_{p}^{-2}$) is
significantly more consistent than what was observed in figure 7.

Figure 22: Rollout error metrics for all models used in the benchmark tests
presented in Section 5 (5 models for each $\Delta t$, trained in a similar
fashion using different random seeds). Mean values for each model are computed
by averaging over sheets, time steps, and simulations. The error bars
represent the minimum and maximum rollout error achieved for the corresponding
set of test simulations. We observe that equivalent models perform similarly,
with the exception of model #4 for $\Delta t=10^{-2}\ \omega_{p}^{-1}$. This
difference in performance was already expected after completing the training
procedure since the validation loss was approximately two times that of
equivalent models (i.e. achieved worse single-step acceleration prediction).
Figure 23: Energy conservation metrics for all models used in the benchmark
tests presented in Section 5 (5 models for each $\Delta t$, trained in a
similar fashion using different random seeds). Similarly to the rollout errors
metrics results presented in figure 22, the energy conservation does not vary
significantly with the random seed. The exception is once again model #4 for
$\Delta t=10^{-2}~{}\omega_{p}^{-1}$, which causes a significant increase in
the average energy variation in figure 7 (see “Thermal”). Figure 24: Same
diagnostic presented in figure 7 without considering the contribution of model
$\\#4$ for $\Delta t=10^{-2}~{}\omega_{p}^{-1}$, which was shown to converge
to a significantly worse validation loss, rollout accuracy, and energy
conservation capabilities. The behavior of the GNS curves becomes more
consistent across the different time resolutions.

### Energy Conservation Diagnostic

In figure 25 we provide examples of the observed energy variation during a
simulation rollout for the sheet model and the GNS. The moving average
($\Delta t=2\pi~{}\omega_{p}^{-1}$) as well as its maximum value (what we
consider for the results in figure 7) are highlighted. The results illustrate
the importance of using the moving average for a fair comparison between
algorithms since the usage of finite difference velocities for the GNS
introduces an oscillation in the energy calculation (period equal to half a
plasma oscillation). It also showcases why we do not use the first plasma
period, since before crossings occur (i.e. first time-steps) the initial
“energy increase” introduced by the finite difference velocities is not
compensated in the remaining of the plasma period (which would lead to an
incorrect estimation of the energy using the moving average strategy).

Figure 25: Illustration of the energy conservation diagnostic calculation. All
systems consist of $10^{3}$ sheets inside a periodic box. Both the sheet model
and the GNS use a time step of $\Delta t=10^{-1}~{}\omega_{p}^{-1}$. The
results illustrate why it is important to use the moving average of the energy
over a plasma oscillation ($\Delta t=2\pi~{}\omega_{p}^{-1}$) for a fair
comparison between algorithms (since the GNS uses finite difference
velocities). Additionally, it shows the reason for not considering the first
plasma oscillation, since the finite difference velocities introduce a large
oscillation in the estimation of the energy before crossings occur (i.e.
initial time-steps).

It is important to highlight that the initial energy of the system, which is
considered for both the sheet model algorithm and the GNS, is computed using
the initial instantaneous velocities (i.e. the input used for the sheet model
simulation) which are different from the initial finite difference velocities
used for the GNS. We compute the initial finite difference velocities by
performing a simulation backward in time until $t=-\Delta t_{GNS}$ using the
sheet model at a very high temporal resolution ($\Delta
t_{SM}=10^{-4}~{}\omega_{p}^{-1}$). This ensures that the initial conditions
of the different simulators are equivalent.

##  Appendix G Two-stream Instability – Extra

### Different Modes

In figure 26 and figure 27 we showcase a similar setup as the one shown in
Section 6.6 with the difference that the main modes excited are now $m=8$
($v_{0}=244~{}\delta\omega_{p}^{-1}$) and $m=12$
($v_{0}=162~{}\delta\omega_{p}^{-1}$). We observe once again that the GNS is
capable of modeling the same macrophysics as the sheet model at a considerably
higher time resolution. The total energy variation of the system (for $m=8$
and $m=12$) is similar to the one presented for $m=4$ in Section 6.6 ($\approx
2\%$).

Figure 26: Two stream instability simulation which excites the $m=8$ mode
($v_{0}=244~{}\delta\omega_{p}$). We use the same model and setup as described
in Section 6.6. The GNS is once again capable of recovering the same dynamics
as the higher temporal resolution sheet model. Figure 27: Two stream
instability simulation which excites the $m=12$ mode
($v_{0}=162~{}\delta\omega_{p}$). We use the same model and setup as described
in Section 6.6 and figure 26. Results provide further demonstration of the GNS
capability to recover the instability dynamics.

### Different Model Seeds

In figure 28 we reproduce the two-stream simulation from Section 6.6 now using
equivalent GNNs trained with different random seeds. As previously mentioned,
these models showcase (overall) worse energy conservation capabilities for
out-of-training distribution velocities (figure 23). In particular, for the
two-stream instability simulation presented in figure 28, we observe relative
energy variations between $10-16\%$ in comparison with the $2\%$ reported in
Section 6.6. Nonetheless, the macrophysics obtained for these models is
consistent with the results shown in the main body of the paper (figure 14).
The only exception is Model $\\#1$, whose phase space looks considerably
different (the potential energy growth rate still matches that of other
models). Once again we would like to highlight the significant difference
between the results obtained with the sheet model at a comparable time
resolution (figure 14).

Figure 28: Two-stream instability results ($m=4$) for GNNs initialized with
different random seeds. We recover similar dynamics to those presented in the
main body of the paper (figure 14) while using models with worse energy
conservation capabilities.

### Different Time Step

In figure 29, we provide an additional example of the two-stream instability
test performed in Section 6.6, now performed at $\Delta
t=10^{-2}~{}\omega_{p}^{-1}$. The relative energy variation is approximately
$2\%$ for the GNS (Model $\\#0$) and $9\%$ for the sheet model at $\Delta
t=10^{-2}~{}\omega_{p}^{-1}$. We observe that both the GNS and the sheet model
are able to correctly recover the dynamics using this time step. The reasons
why we do not observe improved energy conservation capabilities for the GNS at
$\Delta t=10^{-2}~{}\omega_{p}^{-1}$ when compared to models trained at
$\Delta t=10^{-1}~{}\omega_{p}^{-1}$ have already been discussed in Section
5.2.

The results obtained across GNN models trained with different seeds for this
higher temporal resolution are also consistent (figure 30), with the exception
of model $\\#4$ ($51\%$ energy loss compared to $2-7\%$ for the remaining
seeds) whose problems have already been identified and discussed in F.

Figure 29: Comparison of the phase space and potential energy evolution of the
two-stream scenario (with $m=4$) using a smaller time step ($\Delta
t=10^{-2}~{}\omega_{p}^{-1}$). We observe similar results to those obtained
for the GNS at lower time resolutions (figure 14). Figure 30: Two-stream
instability ($m=4$) results obtained for the remaining GNN models trained with
a time step $\Delta t=10^{-2}~{}\omega_{p}^{-1}$. Apart from Model $\\#4$
(whose problems were already identified in F), all others retrieve similar
dynamics than those observed in figure 29.