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(<path_to_npz_file>, 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.