Dataset and model weights for particle-based plasma simulation using a graph neural network

This record contains the trained graph network simulator models described in the paper "Particle-based plasma simulation using a graph neural network" [1], and the data used to train and evaluate them. The data were generated from simualtions of two counterpropagating beams of electrons in one spatial dimension using the EPOCH particle-in-cell code [2]. 

The provided files are:

models.zip - contains a directory for each of models A, B and C, each containing three files:

hyperparameters.json - contains the model hyperparameters and parameters used for training the model (batch size, learning rate, added noise)

model-*.pt - contains trained model weights

train_state-*.pt - contains the optimizer state, training step and training loss; it can be used to resume training

data.zip - contains the data split into training, validation and test sets, each in its own NPZ file, and JSON files "metadata_long_timestep.json" and "metadata_short_timestep.json" with information about the dataset.

The metadata JSON files contain the following information: 

"bounds": the lower and upper bounds of the simulated space, in metres

"sequence_length" : the number of snapshots in each simulation

"default_connectivity_radius": the value that will be used to construct graphs unless overriden when running the code, in metres

"dim": the number of spatial dimensions

"dt": nominal time step between snapshots, in seconds

"vel_mean" and "vel_std": the mean and standard deviation of velocities computed from displacements between adjacent snapshots in the training data, used for scaling input features

"acc_mean" and "acc_std": the mean and standard deviation of acclerations computed from displacements between adjacent snapshots in the training data, used for scaling the outputs of the model

"E_mean", "E_std", "B_mean", "B_std": the mean and standard deviation of x, y and z components of the electric (E) and magnetic (B) fields in the training data, with x being the single dimension in which particles are constrained to move in the simulation; used for scaling input features

"weight_max" and "weight_min": the maximum and minimum particle weight in the training dataset

Model A was trained with the entire training dataset, using the parameters in metadata_short_timestep.json for scaling. For models B and C, every fourth snapshot from the data was taken and the rest were discarded; the corresponding metadata is contained in metadata_long_timestep.json.

Each NPZ file contains an array named "gns_data", which can be loaded with the following Python code:

import numpy as np
with np.load(, allow_pickle=True) as data_file:
data = data_file['gns_data']

This contains a list of dictionaries, each containing data from one simulation:

'particle_weights': array of shape (number_of_particles,)

'particle_positions': array of shape (number_of_particles, number_of_snapshots, dim), containing the position of each particle in metres

'momenta': array of shape (number_of_particles, number_of_snapshots, dim), containing the momentum of each particle in kg m/s

'grid_positions': array of shape (number_of_grid_points,), containing the positions of grid points in metres

'electric_field': array of shape (number_of_grid_points, number_of_snapshots, 3), containing the x, y and z componenets of the electric field at each grid point from each snapshot in newtons per coulomb

'magnetic_field': array of shape (number_of_grid_points, number_of_snapshots, 3), containing the x, y and z componenets of the electric field at each grid point from each snapshot in teslas

'num_left': number of particles initially in the beam travelling to the left

'num_right': number of particles initially in the beam travelling to the right

'original_index': string, unique ID for the sample

'simulation_time': array of shape (number_of_snapshots,) containing the simulated time in seconds at which each snapshot was recorded

'temp': the initial temperature of the beams in kelvins, as given to EPOCH

'drift_p': the magnitude of the initial drift momentum of the beams in kg m/s, as given to EPOCH

'dens': the initial number of electrons per metre as given to EPOCH

References

[1] M. Mlinarevic, G. K. Holt and A. Agnello, Particle-based plasma simulation using a graph neural network, 2025, arXiv:2503.00274 [physics.plasm-ph]

[2] T. D. Arber et al., Contemporary particle-in-cell approach to laser-plasma modelling, Plasma Phys. Control. Fusion 57 (2015) 113001.