Particle-based plasma simulation using a graph neural network

TL;DR

The paper tackles the computational challenge of simulating plasma dynamics by learning a differentiable surrogate for particle-in-cell simulations. It uses a graph network-based simulator (GNS) that encodes particles and (optionally) field grid points as nodes and learns accelerations and electric fields through message passing in an encoder–processor–decoder architecture. The approach reproduces the two-stream instability in 1D counterpropagating electron beams and achieves high accuracy for densities, fields, and phase-space evolution, even with time steps two orders of magnitude larger than conventional PIC; growth rates align with linear theory. This surrogate offers a path toward faster, differentiable forward/inverse plasma solvers and could enable optimization, control, and discovery tasks in plasma physics, albeit with attention to error accumulation and resource requirements.

Abstract

A surrogate model for particle-in-cell plasma simulations based on a graph neural network is presented. The graph is constructed in such a way as to enable the representation of electromagnetic fields on a fixed spatial grid. The model is applied to simulate beams of electrons in one dimension over a wide range of temperatures, drift momenta and densities, and is shown to reproduce two-stream instabilities - a common and fundamental plasma instability. Qualitatively, the characteristic phase-space mixing of counterpropagating electron beams is observed. Quantitatively, the model's performance is evaluated in terms of the accuracy of its predictions of number density distributions, the electric field, and their Fourier decompositions, particularly the growth rate of the fastest-growing unstable mode, as well as particle position, momentum distributions, energy conservation and run time. The model achieves high accuracy with a time step longer than conventional simulation by two orders of magnitude. This work demonstrates that complex plasma dynamics can be learned and shows promise for the development of fast differentiable simulators suitable for solving forward and inverse problems in plasma physics.

Figures (11)

• Figure 1: Validation loss curves for training of three models: two without electric field nodes, with time steps of 0.15ms (model A, blue) and 0.60ms (model B, orange), and one with electric field nodes and a time step of 0.60ms (model C, green). The values of the model hyperparameters are given in \ref{['tab:model_parameters']}. The validation loss is the mean squared error on the predicted position calculated over the full simulation rollout of 0.15s. The faded lines show the loss after every e4 gradient updates, while the solid curves are the result of smoothing by applying a Savitzky--Golay filter SavitzkyGolayFilterSavitzkyGolayFilterErrors with a cubic polynomial and a window size of 100.
• Figure 2: Snapshots from a simulation of two counterpropagating beams of electrons with a density of 24.65 electrons per metre, initial temperature of [scientific-notation=true, round-mode=figures, round-precision=2]20363.209227206735K, and drift momentum of [scientific-notation=true, round-mode=figures, round-precision=2]3.473697853150168e-24kg.m.s^-1 (giving $v_\text{th}\xspace/v_0=0.14$), performed using EPOCH (left column), GNS model A, with a time step of 0.15ms (middle column) and GNS model B, with a time step of 0.60ms (right column). Neither GNS model includes electric field nodes. The values of the model hyperparameters are given in \ref{['tab:model_parameters']}. The first row shows a plot of electron momentum against position after a simulated time of 21ms, the second row shows the state after 42ms, and the third row after 150ms. Each point represents a superparticle, which represents around 7703.09 electrons. The colour of the points shows which of the two beams the particle comes from. The plots show the development of the two-stream instability.
• Figure 3: Snapshots from a simulation of two counterpropagating beams of electrons with a density of 5.1 electrons per metre, initial temperature of [scientific-notation=true, round-mode=figures, round-precision=2]84259.98604037115K, and drift momentum of [scientific-notation=true, round-mode=figures, round-precision=2]7.639175348786281e-25kg.m.s^-1 (giving $v_\text{th}\xspace/v_0=1.35$), performed using EPOCH (left column), GNS model A, with a time step of 0.15ms (middle column) and GNS model B, with a time step of 0.60ms (right column). Neither GNS model includes electric field nodes. The values of the model hyperparameters are given in \ref{['tab:model_parameters']}. The first row shows a plot of electron momentum against position after a simulated time of 21ms, the second row shows the state after 42ms, and the third row after 150ms. Each point represents a superparticle, which represents around 1588.27 electrons. The colour of the points shows which of the two beams the particle comes from.
• Figure 4: Distribution of the number of superparticles in space over time (top row), its Fourier transform (middle row) and the distribution of electron momenta over time (bottom row) for a simulation of two counterpropagating beams of electrons with a density of 25 electrons per metre, initial temperature of 2.0e4K, and drift momentum of 3.5e-24kg.m.s^-1 (giving $v_\text{th}\xspace/v_0=0.14$), performed using EPOCH (leftmost column) and GNS model B (with a time step of 0.60ms, without electric field nodes, in the middle column). The values of the model hyperparameters are given in \ref{['tab:model_parameters']}. The rightmost column shows the difference between the GNS and EPOCH predictions divided by their mean. The Fourier transform is calculated for each time step and its amplitude averaged over sets of five consecutive time steps. The plots clearly show the development of the two-stream instability.
• Figure 5: Distribution of the number of superparticles in space over time (top row), its Fourier transform (middle row) and the distribution of electron momenta over time (bottom row), for the same ground truth simulation (leftmost column) as in \ref{['fig:spectrogram_no_field']}, but performed using GNS model C, which uses field nodes (middle column). The rightmost column shows the difference between the GNS and EPOCH predictions divided by their mean.