Learning collision operators from plasma phase space data using differentiable simulators

TL;DR

The paper tackles the challenge of inferring collision operators for plasmas when analytical forms are unknown by learning advection and diffusion tensors directly from phase-space data using a differentiable FP solver embedded in a kinetic simulator. It contrasts learning from particle tracks with learning from phase-space evolution of multiple subpopulations, employing temporal unrolling and symmetry constraints to combat non-uniqueness and improve generalization. Results show that phase-space-based operators, including NN-parametrized models, match or surpass track-based estimates in long-term predictive accuracy and align closely with theoretical PIC collision operators in the non-relativistic regime. The approach offers memory-efficient, scalable diagnostics and enables data-driven operator discovery in electromagnetically dominated and potentially strongly coupled plasma regimes, with clear pathways to time-varying backgrounds and extended operator forms.

Abstract

We propose a methodology to infer collision operators from phase space data of plasma dynamics. Our approach combines a differentiable kinetic simulator, whose core component in this work is a differentiable Fokker-Planck solver, with a gradient-based optimisation method to learn the collisional operators that best describe the phase space dynamics. We test our method using data from two-dimensional Particle-in-Cell simulations of spatially uniform thermal plasmas, and learn the collision operator that captures the self-consistent electromagnetic interaction between finite-size charged particles over a wide variety of simulation parameters. We demonstrate that the learned operators are more accurate than alternative estimates based on particle tracks, while making no prior assumptions about the relevant time-scales of the processes and significantly reducing memory requirements. We find that the retrieved operators, obtained in the non-relativistic regime, are in excellent agreement with theoretical predictions derived for electrostatic scenarios. Our results show that differentiable simulators offer a powerful and computational efficient approach to infer novel operators for a wide rage of problems, such as electromagnetically dominated collisional dynamics and stochastic wave-particle interactions.

Figures (33)

• Figure 1: Illustration of how advection and diffusion coefficients can be inferred from 2D particle tracks. (a) By following a group of particles with similar initial velocities over time, we observe that advection leads to an average velocity drift ($<\Delta v_i>$), while diffusion leads to an increased spread of the distribution ($<\Delta v_i\Delta v_j>$). Drift is visible by noting that the average particle velocity (red line) is changing when compared to the initial velocity (dashed black line). Diffusion is clearly illustrated by the increased distribution spread (highlighted via the phase space histograms). (b) In practice, the advection-diffusion values are estimated by measuring the rate of change of $<\Delta v_i>$ and $<\Delta v_i\Delta v_j>$ with respect to a period of time where the evolution is linear. To obtain an accurate estimate of the coefficients it is crucial that statistics are computed during the linear evolution phase.
• Figure 2: (a) Illustration of how advection and diffusion coefficients are inferred from the 2D phase space evolution of $N_s$ subpopulations using a differentiable Fokker-Planck (FP) solver. The evolution of the phase space and the changes in the operator are exaggerated for visualization purposes. (b) Using the FP solver and the current operator state, we advance the phase space over $N_u$ time steps and compare against the PIC results via a scalar error metric $\mathcal{L}$. This operation is performed over all training subpopulations $s\in[1, N_s]$ and initial time-steps $t\in[0,N_t - N_u]$. (c) We then update the operator using the gradient-based optimizer Adam kingma2014adam, to minimize the unrolled error across all subpopulations and time-steps. The two operations are performed sequentially in a loop until the results have converged.
• Figure 3: Advection and diffusion coefficients retrieved with different approaches for simulation $index=0$ ($N_{ppc}=4$, $m=1$, $\Delta_x/\lambda_D=1, v_{th}=0.01c)$. Track operator is computed from statistics over all macroparticles, PS operators are computed from the phase space evolution of 9 sub-populations using a differentiable FP solver and a discrete (PS-Tensor) or continuous (PS-NN, PS-NN-Multi) approximator. Note that all these operators perform similarly when predicting the phase space evolution of most tested subpopulations. More examples for other simulation parameters are provided in Supplementary Material \ref{['app:example_single_simulation']}.
• Figure 4: Phase space evolution for a ring subpopulation using operators recovered in Figure \ref{['fig:AD_comparison_index_0']}. This subpopulation was not used during the training of PS-Tensor, PS-NN, and PS-NN-Multi models. The top row corresponds to the observed dynamics in the PIC simulation ($f^{(t)}$). The remaining rows represent the predicted phase space evolution on the left ($\hat{f}^{(t)}$ for $v_x/v_{th}\in [-5,0]$) and the difference to the PIC data on the right ($\hat{f}^{(t)} - f^{(t)}$ for $v_x/v_{th}\in [0,5]$). Values are normalized to the peak of the PIC distribution function at time $t$ ($f^{(t)}_{max}$). All operators approximate the dynamics relatively well, and overall, the random distribution of errors can be attributed to the granularity of the original distribution function. However, both Tracks and PS-NN-Multi seem to systematically slightly underestimate advection (central blue error region and outer red halo). Examples for other subpopulations are provided in Supplementary Material \ref{['app:example_single_simulation']}.
• Figure 5: Distribution of rollout errors for advection and diffusion models obtained from particle tracks (Tracks) or phase space evolution of subpopulations (PS-Tensor, PS-NN, PS-NN-Multi). Boxplots represent statistics over the full dataset of PIC simulations (see table \ref{['tab:simulation_parameters']} for more details) averaged over initial subpopulations (Train - 9 subpopulations; Test - 19 subpopulations, more information in Supplementary Material \ref{['app:subpopulation_sampling']}). The mean error is shown with a dashed line, the median with a full line. The filled area corresponds to values between the first and third quartiles ($Q_1$ and $Q_3$). Whiskers represent the lowest values up to $Q_1 - 1.5(Q_3 - Q_1)$ and $Q_3 + 1.5(Q_3 - Q_1)$. Dots represent values outside this range. The performance of the different methods is, on average, equivalent on the test data. These results demonstrate that estimating the operators from phase space information using a differentiable simulator is a viable alternative approach to particle tracks.