Learning time-dependent and integro-differential collision operators from plasma phase space data using differentiable simulators

TL;DR

The paper tackles learning time-varying collision operators in plasmas from phase-space data using differentiable simulators. It introduces two operator forms—the time-dependent Fokker-Planck operator with $A_\ parallel$, $D_\nparallel$, and $D_\perp$, and a non-local integro-differential form with kernel $\boldsymbol{K}$—and demonstrates their inference from self-consistent PIC data, either via neural-network or discrete-tensor representations. Key findings show that phase-space-based learning yields more accurate long-time dynamics than tracking-based estimates, and a Pareto analysis identifies an advection-diffusion description with kernel size $k=2$ as optimal for the studied regime. The work provides a data-driven framework to infer collision physics in regimes lacking closed-form solutions and sets the stage for incorporating external fields and wave spectra to study non-thermal particle acceleration and electromagnetically dominated plasmas. This approach has potential applications in laboratory and astrophysical plasmas where collisional and stochastic dynamics are intertwined and evolve with the background state.

Abstract

Collisional and stochastic wave-particle dynamics in plasmas far from equilibrium are complex, temporally evolving, stochastic processes which are challenging to model. In this work, we extend previous methods coupling differentiable kinetic simulators and plasma phase space diagnostics to learn collision operators that account for time-varying background distributions. We also introduce a more general integro-differentiable operator formulation to probe relevant terms in the collision operator. To validate the proposed methodology we use data generated by self-consistent electromagnetic Particle-in-Cell simulations. We show that both approaches recover operators that can accurately reproduce the plasma phase space dynamics while being more accurate than estimates based on particle track statistics. These results further demonstrate the potential of using differentiable simulators to infer collision operators for scenarios where no closed form solution exists or deviations from existing theory are expected.

Figures (21)

• Figure 1: Illustration of the methodology used to learn collision operators from 2D PIC simulation data. Similarly to carvalho2026learning we use a differentiable simulator, coupled with phase space diagnostics from sub-populations of particles of the background plasma, to learn the collision operator ($\mathcal{C}$) that best describes the observed long-term phase space dynamics (i.e. that one that minimizes the loss $\mathcal{L}$\ref{['eq:loss']} between predicted $\hat{f}$ and observed $f$ distribution functions). In this work we explore two new description for the collision operator: (a) a Fokker-Planck time-dependent operator as described by \ref{['eq:fp_operator']}, here represented solely by its parallel ($A_\parallel$, $D_\parallel$) and perpendicular ($D_\perp$) components with respect to the direction of propagation of the particle; (b) an integro-differential operator as described in \ref{['eq:K_operator']} which in the discrete limit can be represented by a convolution of the distribution function with a 4-dimensional kernel $K_i(v_x,v_i,l,m)$ where the kernel sizes $k$ controls the non-locality of the operator (in this case we show $k=2$). In both cases, the operators can be approximated either by a discrete tensor or a continuous approximator such as a NN. For the NN case (example in (a)), initial guesses correspond to randomly initialized models, for the discrete case (example in (b)) the initial operator guess is set to zero (for illustrative purposes we instead set random values in this figure).
• Figure 2: (a) Time-dependent advection-diffusion coefficients retrieved using Tracks, and phase space based approaches with a discrete (PS-Tensor) and continuous (PS-NN) function approximators. Main difference to highlight is that the Diffusion predicted from Tracks is lower than the one measured from PS-Tensor and PS-NN. (b) Snapshots of the velocity distribution $f(v) = v \int d\theta f(v\cos{(\theta), v\sin{(\theta))}}$ at different times. All curves are normalized to the same value. Initial velocity distribution function corresponds to an isotropic waterbag ($t=t_{start}$) which by the end of the simulation has relaxed towards a thermal distribution ($t=t_{end}$). Distribution function evolves more rapidly in the initial time-steps of the simulation. This is the reason why the advection-diffusion coefficients vary more noticeable for smaller $t$ values (purple curves).
• Figure 3: Phase space evolution for a subpopulation using operators recovered in Figure \ref{['fig:AD-sim1-comparisons']}. 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_{w}\in [-2,0]$) and the difference to the PIC data on the right ($\hat{f}^{(t)} - f^{(t)}$ for $v_x/v_{w}\in [0,2]$). Values are normalized to the peak of the PIC distribution function at time $t$ ($f^{(t)}_{max}$). The operator estimated from particle tracks fails to reproduce the phase space dynamics. The operators learned from phase space evolution using the differentiable simulator approximate the dynamics relatively well, and overall, the random distribution of errors can be attributed to the granularity of the original distribution function. Examples for other subpopulations are provided in \ref{['app:ad_extra_rollouts']}.
• Figure 4: Pareto-curve on the impact of operator kernel size $k$ on rollout error using the corresponding integro-differential operator (K). The baseline error for a pure advection-diffusion model trained on the simulation (PS-Tensor) is shown in red. The optimal value is $k=2$, which corresponds to an advection-diffusion operator. This is the expected result for small-angle scattering collisions of particles. Further increasing $k$ leads to virtually no improvements and starts degrading the performance due to overfitting/training instabilities.
• Figure 5: Discrete 4-Dimensional kernel operators recovered for different kernel sizes $k$ (correspond to figure \ref{['fig:K-test_l1_avg']}). Only $K_x$ is shown since $K_x(v_x, v_y, l,m) = K_y(v_y, v_x, m, l)$ by construction. It is clear that for $k > 1$ there is a derivative term along $v_x$ that consistently dominates. This corresponds to a diffusion term, which is expected to be relevant for the dynamics of interest. The operator with $k=1$ can not compute this derivative and therefore fails to reproduce the phase space dynamics. Since no second derivative term appears for $k=4$, we can also infer that non-local transport is not relevant.