Unsupervised discovery of nonlinear plasma physics using differentiable kinetic simulations

Abstract

Plasma supports collective modes and particle–wave interactions that lead to complex behaviour in, for example, inertial fusion energy applications. While plasma can sometimes be modelled as a charged fluid, a kinetic description is often crucial for studying nonlinear effects in the higher-dimensional momentum–position phase space that describes the full complexity of the plasma dynamics. We create a differentiable solver for the three-dimensional partial-differential equation describing the plasma kinetics and introduce a domain-specific objective function. Using this framework, we perform gradient-based optimization of neural networks that provide forcing function parameters to the differentiable solver given a set of initial conditions. We apply this to an inertial-fusion-relevant configuration and find that the optimization process exploits a novel physical effect.

1. Introduction

Kinetic plasma physics is described, typically, using various formulations or realizations of Boltzmann-like transport equations for the electron and ion species. A special case that is frequently studied is that of the dynamics restricted to one spatial dimension and one momentum dimension (‘1D-1V’) for a single species of particles, given by

(1.1)\begin{equation} \frac{\partial \tilde f}{\partial \tilde t} + \tilde{v}\frac{\partial \tilde f}{\partial \tilde{x}} - \tilde E \frac{\partial \tilde f}{\partial \tilde{v}} = \left[\frac{\delta \tilde f}{\delta \tilde{t}}\right]_{{\rm coll}} , \end{equation}

where the normalized electric field is given by solving $\boldsymbol {\nabla } \boldsymbol {\cdot } \tilde {E} = 1 - \int \tilde {f}\,{\rm d}\tilde {v}$, and the normalized particle distribution function is $\tilde {f} = f(t,x,v) / n_e v_{{\rm th}}^3$. The quantities are normalized to electrostatic units where $\tilde {v} = v / v_{{\rm th}}, \tilde {x} = x/\lambda _D, \tilde {t} = \omega _p t, \tilde {m} = m/m_e$ and $\tilde {E} = e E / m_e v_{{\rm th}} \omega _p$. Here, $v_{{\rm th}}$ is the thermal velocity, $\omega _p$ is the plasma frequency, $\lambda_D$ is the Debye length, and $m_e$, and $e$, are the electron mass, and charge, respectively.

The left-hand side of this equation describes the evolution of the particle distribution in ‘macroscopic’ field $E$ that arises self-consistently from the one-particle distribution $f$. The right-hand side provides a description of two-particle and higher correlations (i.e. collisions). Equation (1.1), along with Gauss's law and with a particular form of the collision operator for two-particle interactions, is often termed the Vlasov–Poisson–Fokker–Planck (VPFP) equation set. Solving the VPFP set is often analytically intractable, even in 1D-1V. This is because the left-hand side has a stiff linear transport term, has a nonlinear term in $E \partial f/\partial v$ and can sustain wave propagation and other hyperbolic partial-differential equation (PDE) behaviour. Additionally, the right-hand side is typically represented by a hyperbolic advection–diffusion PDE. Making progress on kinetic plasma physics often requires computational simulation tools.

Numerical solutions to the 1D-1V VPFP equation set have been applied in research on laser–plasma interactions in the context of inertial fusion, for example in plasma-based accelerators (Krall, Joyce & Esarey 1991; Thomas 2016), space physics (Chen, Klein & Howes 2019), fundamental plasma physics (Pezzi et al. 2019) and inertial fusion (Strozzi et al. 2007; Fahlen et al. 2009; Banks et al. 2011). Such numerical simulations may be used to explore initial conditions and forcing functions to understand the behaviour of a physical effect in response to input parameters. Multi-dimensional ‘brute-force’ scans are, however, inefficient and costly and it is therefore beneficial to seek a more guided approach by leveraging optimization techniques.

Optimization techniques can be grouped into gradient-based or gradient-free methods. Gradient-free methods are typically useful when it is not practical to obtain the gradient of the desired objective function (Nocedal & Wright 2006). For example, when the objective function is the outcome of a plasma-based accelerator experiment. In this case, one can use classical approaches like the Nelder–Mead algorithm (Shalloo et al. 2020) or more modern alternatives like Bayesian (Shalloo et al. 2020; Jalas et al. 2021) or evolutionary (He et al. 2015; Smith et al. 2020) algorithms. These algorithms have been shown to be effective in the small $N_p$ regime, where $N_p$ is the number of parameters to optimize. While there has been some success in using gradient-free approaches in large $N_p$ regimes e.g. in reinforcement learning scenarios (Salimans et al. 2017), it has been when the gradient information is unreliable. In general, as $N_p \gg 1$, utilizing gradient-free methods becomes computationally intractable or unstable. In the work here, we learn functions reparametrized via a neural network, and therefore we seek the optimal neural network weights and biases. In this work, $\mathcal {O}(10) \leq N_p \leq \mathcal {O}(10^3)$, and automatic-differentiation-driven, gradient-based methods offer significant advantages.

While not focused on the nonlinear kinetic plasma dynamics, recent works have applied gradient-based optimization to fusion device design. Analytic approaches have resulted in the development of adjoint methods for shape derivatives of functions that depend on magnetohydrodynamic equilibria (Antonsen, Paul & Landreman 2019; Paul et al. 2020). These methods have been used to perform optimization of stellarator design (Paul, Landreman & Antonsen 2021). Other work uses gradients obtained from analytic (Zhu et al. 2017) and automatic differentiation¹ (AD) (McGreivy, Hudson & Zhu 2021; Conlin et al. 2022; Dudt et al. 2022; Panici et al. 2022) towards similar device optimization goals. Here, we propose the application of gradient-based optimization towards understanding the nonlinear, kinetic plasma dynamics using differentiable kinetic simulations.

Differentiable simulations have been used in a variety of contexts for such guided searches, for example, learning parameters for molecular dynamics (Schoenholz, Cubuk & Jax 2019), learning differencing stencils in PDEs (Bar-Sinai et al. 2019; Zhuang et al. 2020; Kochkov et al. 2021) and controlling PDEs (Holl, Koltun & Thuerey 2020). Here, we apply AD towards learning physical relationships and discovering novel phenomena in the VPFP dynamical system by training neural networks through differentiable simulations that solve (1.1).

In the rest of the manuscript, we first develop a conceptual framework for how one might implement gradient-descent-based optimization of parameters and learning of functions via differentiable programs. We then apply this framework towards kinetic plasma physics. Specifically, we extend the findings in Fahlen et al. (2009) of the etching of nonlinear plasma wavepackets to study the effect of the hot electrons from one wavepacket on downstream wavepackets using a gradient-based approach. To do this, we train a neural network that provides control parameters to the PDE solver. By choosing physical parameters as inputs and control parameters as outputs of the neural network, we enable the neural network to learn a function that describes the physical relationship between the plasma parameters and the forcing function parameters e.g. the resonance frequency. We train the neural network in an unsupervised fashion using a cost function based on minimizing the free energy and maximizing the non-Maxwellian-ness of the plasma distribution function. This enables us to create self-learning plasma physics simulations, where the optimization process provides a physically interpretable function that can enable physics discovery.

2. Physics Discovery using Differentiable Simulations

In this section, we provide a step-by-step description of how a traditional simulation-based computational physics workflow may be modified to perform closed-loop optimization.

2.1. Open loop: manual workflow

Figure 1(a) depicts a typical simulation workflow represented as a cyclic graph. The user defines the parametric inputs that create the state vector $\boldsymbol {x}$. This can contain any parameters that are used to define the simulation e.g. the grid size, the number of solver steps, etc. For didactic purposes, the physical parameters to the simulation may be separated from $\boldsymbol {x}$ into a different vector of inputs $\boldsymbol {p_d}$ e.g. the forcing function parameters, the viscosity coefficient etc.

Figure 1. (a) A typical workflow where the user provides the initial conditions and forcing function parameters to a PDE solve. The output of the solve is stored as the final state $\boldsymbol {x_f}$. The final state is analysed using domain-specific postprocessing algorithms. (b) A cost function and a parameter scan are introduced which enables= a closed-loop search. (c) A gradient-descent-based optimization algorithm replaces the parameter scan to provide a more efficient search mechanism. This requires the components in the purple background to be written in an auto-differentiable framework. (d) We add a neural network that generates the forcing function parameters as a function of other parameters. This generalizes the learnings from (c) and enables a continuous representation within the learned parameters.

Each of $\boldsymbol {x}$ and $\boldsymbol {p_d}$ is passed to the algorithm that solves the PDE which is represented by the function, $\mathcal {V}$. The output of these simulations is stored in the final state vector $\boldsymbol {x_f}$. The final state is postprocessed using a domain-specific set of algorithms devised by the user or otherwise. The results of the postprocessing are interpreted by the user who then determines the next set of inputs and parameters.

2.2. Closed loop: brute-force parameter scan

Figure 1(b) shows a more automated workflow. We replace the grey-box postprocessing step with the calculation of a scalar quantity $S$ using a cost function $\mathcal {C}$ on the final state $\boldsymbol {x_f}$. This reduces the complexity of the interpretation of the postprocessing and enables a more rapid search in parameter space. The decrease in required human effort for completing one cycle enables the user to execute this loop as a brute-force parameter scan over a pre-defined parameter space. At the end, the user can look up the minimum/maximum of the scalar cost function, and find the parameters which provide that minimum.

The parameter scan approach scales with the number of different unique parameters and the number of values of each parameter. e.g. a two-dimensional search in $x$ and $y$ requires $N_x \times N_y$ calculations. Therefore, the parameter scan approach quickly becomes inefficient when there are many parameters to scan, or when the required resolution in parameter space is very high. To search this parameter space efficiently, and to escape the linear scaling with each parameter, we can use gradient descent.

2.3. Gradient-descent-based parameter learning

Figure 1(c) includes two modifications. The user/parameter search grey box has been replaced with a gradient-descent-based optimization algorithm. This algorithm provides the updated parameters, e.g. $\omega _G$, a guess for the resonant frequency of the system, for the next iteration of the loop. The gradient-descent algorithm requires the calculation of an accurate gradient.

Symbolic differentiation is out of the question here as we do not have an analytical form for our system. In the Appendix, we compare the performance of finite differencing to acquire the gradient and confirm that AD is a superior method for this purpose. Therefore, by writing our PDE solver $\mathcal {V}$ and the cost function $\mathcal {C}$ using a numerical framework that supports automatic differentiation, we are able to perform gradient descent. Since

(2.1)\begin{equation} \mathcal{S} = \mathcal{C}(\boldsymbol{x_f}) = \mathcal{C}(\mathcal{V}(\boldsymbol{x}, \boldsymbol{p_d})), \end{equation}

the gradient for the update step is given by

(2.2)\begin{equation} \frac{\partial \mathcal{S}}{\partial \boldsymbol{p_d}} = \frac{\partial \mathcal{C}(\mathcal{V}(\boldsymbol{x}, \boldsymbol{p_d}))}{\partial \boldsymbol{p_d}} = \frac{\partial \mathcal{S}}{\partial \mathcal{V}} \frac{\partial \mathcal{V}}{\partial \boldsymbol{p_d}}. \end{equation}

For example, if we wish to learn the resonant frequency, $\omega$, that optimizes for the scalar, $\mathcal {S}$, we compute

(2.3)\begin{equation} \frac{\partial \mathcal{S}}{\partial \boldsymbol{\omega}} = \frac{\partial \mathcal{S}}{\partial \mathcal{V}} \frac{\partial \mathcal{V}}{\partial \boldsymbol{\omega}}. \end{equation}

Assuming a well-behaved solution manifold, performing gradient descent tends to reduce the number of iterations required to find the minimum in comparison with an evenly spaced parameter scan, especially when the required resolution is unknown (Nocedal & Wright 2006). While gradient-free methods such as the Bayesian framework, genetic algorithms and others are also often used to perform parametric optimization, we do not perform a more exhaustive comparison because our primary goal is to learn functions, rather than parameters. We discuss this in the following section.

2.4. Gradient-descent-based function learning

In the final step, we can replace the lookup-like capability of the parameter optimization and choose to learn a black-box function that can do the same. Through that process, we acquire a continuous representation of the function rather than the discrete version acquired in § 2.3

Here, we choose to use neural networks, with a parameter vector $\theta$, representing the black-box function and providing the fitting function with a large amount of flexibility. This allows us to extend the gradient-descent-based methodology and leverage existing numerical software to implement this differentiable programming loop. Now,

(2.4)\begin{equation} \mathcal{S} = \mathcal{C}(\boldsymbol{x_f}) = \mathcal{C}(\mathcal{V}(\boldsymbol{x}, \boldsymbol{p_d})) = \mathcal{C}(\mathcal{V}(\boldsymbol{x}, \mathcal{G}(\boldsymbol{x}, \boldsymbol{p_d}; \theta)), \end{equation}

where $\mathcal {G}$ is a function that generates the desired forcing function parameter given a parameter vector $\boldsymbol {\theta }$. To extend the example from § 2.3, $\omega$ is now a function given by $\omega = \mathcal {G}(\boldsymbol {x}, \boldsymbol {p_d}; \theta )$.

We compute the same gradient as in (2.3) and add a correction factor that arises because the parameter (vector) is now $\theta$, rather than $\omega$. The necessary gradient for the gradient update is now given by

(2.5)\begin{equation} \frac{\partial \mathcal{S}}{\partial \boldsymbol{\theta}} = \left[\frac{\partial \mathcal{S}}{\partial \mathcal{V}} \frac{\partial \mathcal{V}}{\partial \mathcal{G}} \right] \frac{\partial \mathcal{G}(\boldsymbol{x}, \boldsymbol{p_d}; \boldsymbol{\theta})}{\partial \boldsymbol{\theta}}. \end{equation}

Since neural networks typically have $\gg \mathcal {O}(10^{2})$ parameters, training these via finite-difference or gradient-free methods is typically avoided except for the most extreme cases (Zhang, Clune & Stanley 2017). It is for this reason that, while it is possible to use finite-difference gradients or even a gradient-free method to perform parameter optimization, these methods are less useful when training neural networks.

3. Discovery of long-lived nonlinear plasma wavepackets

When electrostatic waves are driven to large amplitude, electrons can become trapped in the large potential (Bernstein, Greene & Kruskal 1957; O'Neil 1965). These nonlinear electrostatic wavepackets are dynamically evolving, finite-length analogues of the well-known, time-independent, periodic Bernstein-Greene-Kruskal (BGK) modes described in Bernstein et al. (1957). Simulations of stimulated Raman scattering (SRS) in inertial confinement fusion scenarios show that similar large-amplitude waves, but of finite extent, are generated in the laser–plasma interaction, and that particle trapping is correlated with the transition to the high-reflectivity burst regime of SRS (Strozzi et al. 2007; Ellis et al. 2012).

Simulating wavepackets, similar to those generated in SRS, but in individual electrostatic simulations, has isolated their kinetic dynamics. Fahlen et al. (2009) showed in these isolated simulations that resonant electrons transit through the slower moving wavepacket. This is because the resonant electrons have velocity $v \approx v_{ph} = \omega / k$, where $\omega$ is the frequency and $k$ is the wavenumber of the wavepacket. On the other hand, the group velocity of the wavepacket is $v_g = \partial \omega / \partial k$. Approximating $\omega = \sqrt {1 + 3 k^2}$ gives $v_{ph} / v_{g} = (3k^2+1)/3k^2$. For these wavepackets, this ratio is roughly 3–5. The transit of the resonant electrons from the back of the wavepacket to the front results in the resumption of Landau damping at the back. The wavepacket is then damped away, as seen in Fahlen et al. (2009).

Winjum et al. (2019) modelled the interaction of multiple speckles with a magnetic field acting as a control parameter. Since the effect of the magnetic field is to rotate the distribution in velocity space, the field strength serves as a parameter by which the authors control scattered particle propagation. Using this, along with carefully placed laser speckles, they show that scattered light and particles can serve as the trigger for SRS.

Here, we ask: What happens when a nonlinear electron plasma wavepacket is driven on top of another? To answer this question, we reframe it as an optimization problem and ask: What is the best way to excite a wavepacket that interacts with a pre-existing wavepacket?

We start with a large-amplitude, finite-length electrostatic wavepacket driven by a forcing function with parameters given by

(3.1)\begin{equation} \boldsymbol{p_0} = \left[x_{0}, \omega_{0},t_{0},k_{0}\right], \end{equation}

where $x_{i}$ is the location of excitation, $\omega _i$ is the frequency, $t_i$ is the time of excitation and $k_i$ is the wavenumber, of the $i$th wavepacket.

Since we seek to excite a second wavepacket that can interact with the detrapped electrons, we stipulate that the phase velocity of the resonant electrons from both wavepackets is roughly the same. For this reason, we set the wavenumber $k_1 = k_0$. We reparameterize the resonant frequency, $\omega _1$, with a frequency shift, $\Delta \omega _1$ and the linear resonant frequency $\omega _0$ such that $\omega _1 = \omega _0 + \Delta \omega _1$.

We use the time of excitation of the second wavepacket, $t_1$, as an independent variable along with $k_0$. For each $t_1$ and $k_0$, we seek to learn functions that produce $x_1$ and $\Delta \omega _1$ i.e. we seek to learn $x_1(t_1, k_0, \omega _0)$ and $\Delta \omega _1( t_1, k_0)$. The entire parameter vector for the second wavepacket is given by

(3.2)\begin{equation} \boldsymbol{p_1} = \left[ x_{1}(t_1, k_0), \Delta \omega_{1}(t_1, k_0),t_{1}, k_{0} \right]. \end{equation}

This framing is also illustrated in figure 2 where, given $k_0$ and $t_1$, we seek functions for $\omega _1$ and $x_1$.

Figure 2. Given a first wavepacket with wavenumber $k_0$ and a desired time of second wavepacket excitation $t_1$, the task is to learn functions that give the optimal frequency $\omega _1$ and spatial location $x_1$ of the second wavepacket.

We reparameterize $\Delta \omega _1$ and $x_1$ with a neural network with a parameter vector, $\theta ^*$, that maximizes the electrostatic energy (minimizes the free energy) and maximizes the difference in the kinetic entropy. These relationships are given by

(3.3)\begin{gather} x_1 = x_1(t_1,k_0; \theta^*), \end{gather}

(3.4)\begin{gather}\Delta \omega_1 = \Delta \omega_1(t_1, k_0; \theta^*) \end{gather}

where²,

(3.5)\begin{equation} \theta^* = {\rm argmin} \left[U_\text{es}(\boldsymbol{p}; \theta) -\Delta \mathcal{KE}(\boldsymbol{p}; \theta)\right], \end{equation}

and

(3.6)\begin{gather} U_\text{es} = \sum_{t_{i}}^{t_f} \Delta t \sum_x \Delta x E^2 \end{gather}

(3.7)\begin{gather}\Delta \mathcal{KE} = \sum_{t_{i}}^{t_f} \Delta t \sum_x \Delta x \sum_v \Delta v (f \log(f) - f_\text{MX} \log(f_\text{MX})), \end{gather}

are the electrostatic energy and difference in the kinetic entropy terms in the loss function, respectively. Also, $f_\text {MX} = f_\text {MX}(n(t, x), T(t, x))$, $n(t, x) = \int f(t, x, v)\,{\rm d} v$, $T = \int f(t, x, v) v^2\,{\rm d} v /n$, where $f_\text {MX}$ is the local Maxwell–Boltzmann distribution. Equation (3.7) describes the deviation of the local distribution function from the equivalent Maxwell–Boltzmann distribution function. It has been used as a measure of the non-Maxwellian-ness by Kaufmann & Paterson (2009). The reason we maximize this deviation is because we seek to enhance the likelihood of nonlinear kinetic effects.

We vary the independent variables such that

(3.8)\begin{gather} k_0 \in [0.26, 0.27,\ldots, 0.32], \end{gather}

(3.9)\begin{gather}t_1 \in [400, 500,\ldots, 800], \end{gather}

giving an input space of 35 samples from which we seek to learn these functions.

We are able to reproduce similar model training results using a few different neural network configurations and learning parameters. To summarize that study, we find vanilla multi-layer perceptrons 8 nodes wide and 2 layers deep to be effective, especially when ‘activated’ using the Leaky-Rectified-Linear-Unit function. The specifics behind the neural network, optimizer, and data normalization parameters are provided in the Appendix.

The training simulations are performed with the following parameters. The grid is discretized using $N_x = 6656, N_v = 160, N_t = 1200, t_{\max } = 1100 \omega _p^{-1}, x_{\max } = 6600 \lambda _D$. Assuming inertial fusion conditions where $n_e = 10^{20}\,\text {cm}^{-3}, T_e = 2$ keV gives a non-zero collision frequency $\nu _{ee}/\omega _p = 5 \times 10^{-5}$ (Huba 2011) for the simulations shown here. Details on the solvers and forcing function parameters are provided in the Appendix.

Figure 3 shows that the loss value is reduced over the duration of the training process. Each simulation takes approximately 3 minutes to run on a NVIDIA T4 GPU. We train for 60 epochs. The convergence in the loss metric suggests that we were able to train a overparametrized neural network with 35 samples of data in 60 epochs. While Metz et al. (2022) discuss that differentiating through the evolution of chaotic dynamical systems can be ill posed, ablation studies indicate that the objective function posed here has a robust minimum. One reason is because the chaotic particle dynamics is subject to dispersion, Landau damping and collisions. This may result in a smoothing of the loss surface. Second, since the objective functions used here are space, time and velocity integrals over the chaotic trajectories, the gradients calculated here may benefit from the smoothness associated with time- and phase-space averaging of a well-behaved ergodic system.

Figure 3. The loss, a sum of (3.6) and (3.7), is plotted as a function of time. Each cross represents an epoch, and batch-wise fluctuations are also displayed. The training converges after roughly 150 GPU hours on a NVIDIA T4 and 2100 simulations, which amounts to 60 epochs.

Figure 4 shows the electric field profile for three different simulations. In figure 4(a), only the first wavepacket is excited, and in figure 4(b), only the second wavepacket is excited. In figure 4(c), both wavepackets are excited. Early in time, $t=400 \omega _p^{-1}$ (green), when only the first wavepacket has been excited figures 4(a) and 4(c) agree perfectly. The second wavepacket is excited at $t=500 \omega _p^{-1}$. At $t=900 \omega _p^{-1}$, once some time has passed after the excitation of the second wavepacket, the first wavepacket has not fully damped away. It is visible as small bumps in figures 4(a) and 4(c). The second wavepacket is also present at this time and easily seen in figure 4(b). A larger-amplitude wavepacket is seen in figure 4(c). Late in time, the difference in amplitude between the second wavepacket in figures 4(b) and 4(c) is obvious. The second wavepacket has nearly damped away in figure 4(b). In figure 4(c), the second wavepacket persists, at nearly the same energy as it was at $t=900 \omega _p^{-1}$. We observe this superadditive behaviour, where $f(x) + f(y) \leq f(x+y)$, for all wavenumbers we model.

Figure 4. Early in time (green), (c) is the same as (a). Later in time, $t=900 \omega _p^{-1}$ (blue), (a,c) show very similar magnitudes for the first wavepacket near $x=1500 \lambda _D$ but the second wavepacket excitation is larger in (c) than (b). At $t=1300 \omega _p^{-1}$ (red), it is clear that (c) is not a superposition of (a,b) because (b) has damped away, while (c) retains electrostatic energy, suggesting the involvement of a superadditive process. (a) First wavepacket only; (b) second wavepacket only; (c) both wavepackets.

To determine the mechanism behind this phenomenon, we turn to the phase-space dynamics. In figure 5, (i) is a space–time plot of the electric field. The two dashed-dot red lines at the front and back of the wavepacket are parallel and indicate the velocity of the wavefront. In figure 5(a)(i), the rear of the wavepacket propagates at a seemingly faster rate than the front. Fahlen et al. (2009) describe this as etching of the wavepacket. In figure 5(b), the wave survives for a much longer time, as was also illustrated in figure 4. Figure 5(ii, iii) are phase-space plots with their centre indicated by the intersection of the horizontal timestamp line and the dashed-dot line at the rear wavepacket. The insets (iv) and (v) correspond to the intersection with the dashed-dot line at the front. Insets (ii) and (iv) show the phase space within a window in $\hat {x}$, while (iii) and (v) are the spatially averaged distribution function. Insets (iii) and (v) serve as a proxy for approximating the propensity of Landau damping in that region. In figure 5(a)(ii, iii), we see that the rear of the wavepacket is Maxwellian. Fahlen et al. (2009) showed this is why the rear of the wavepacket damps faster than the front, as in figure 5(a)(i).

Figure 5. Left – space–time plot of the electrostatic energy shows the long-lived wavepacket in (b) where the field in (b) survives for a longer duration than in (a). The horizontal line indicates the timestamp of the snapshots in the middle and right. The diagonal dashed-dot lines indicate the spatial location of the snapshots in the middle and right. Middle phase-space plots at the back (top) and front (bottom) of the wavepacket. In (b), the phase space shows significant activity at the back of the wavepacket while in (a), the distribution function is nearly undisturbed. Right – T = the spatially averaged distribution function. This confirms the fact that the distribution function has returned to a Maxwell–Boltzmann at the back of the wavepacket in (a), while in (b), the distribution function remains flat at the phase velocity of the wave. This is the reason behind the loss of damping. (a) Second wavepacket only; (b) both wavepackets.

In the simulations described here, the plots in figure 5(b) show that the distribution function at the back of the wavepacket has trapped particle activity (figure 5(b)(ii)) and there is a near zero slope at the phase velocity of the wave (figure 5(b)(iii)). Both plots show that the slope is negligible because of the arrival of streaming detrapped particles from the first wavepacket. Due to this effect, the re-emergence of Landau damping that occurs due to the loss of trapped particles in isolated wavepackets no longer occurs here. This results in a reduction of the etching and the wavepacket propagates freely for some time while the particles from the first wavepacket propagate and arrive at the rear of the second wavepacket.

Figure 6 shows the results of the optimization process for the resonant spatial location. For $t_1 < 500 \omega _p^{-1}$, the resonant location decreases as a function of wavenumber. From analysing phase space, we have determined that the spatial location is related to the resonant electron transport i.e. $x_1 \approx v_{ph} t_1$. Waves with larger wavenumbers have smaller phase velocities. Because of this, the resonant spatial location decreases as a function of wavenumber. For a fixed wavenumber, figure 6(b) shows a locus of points in space–time where long-lived wavepackets can be excited in the presence of a pre-existing wavepacket. This suggests the possibility of a critical space–time radius within which collisional relaxation has yet to occur, and nonlinear effects can be exploited.

Figure 6. (a) Learned function for the resonant spatial location as a function of wavenumber of the first wavepacket and time of excitation of the second. (b) The locus in space–time (in red) where long-lived wavepackets can be excited for $k=0.28$: (a) $x_1(k_0, t_1)$; (b) space–timelocus.

Likewise, we also learn the dependence of the optimum frequency shift as a function of $k_0$ and $t_c$, as shown in figure 7. The learned frequency shift, a few per cent here, is similar in magnitude as to that observed in previous work related to SRS (Ellis et al. 2012). Furthermore, the frequency shift increases in magnitude as a function of wavenumber. As before, waves with larger wavenumbers have smaller phase velocities, and therefore, interact with more particles. Because of this, waves with larger wavenumbers have a larger nonlinear frequency shift associated with them, as we see here (Manheimer & Flynn 1971; Dewar 1972; Morales & O'Neil 1972; Berger et al. 2013).

Figure 7. The learned function for the frequency shift $\Delta \omega _1(k_0, t_1)$ as a function of wavenumber of the first wavepacket, $k_0$, and time of excitation, $t_1$, of the second.

4. Conclusion

We show how one may be able to discover novel physics using differentiable simulations by posing a physical question as an optimization problem. This required domain expertise in determining which functional dependencies to learn using neural networks.

In figure 1, we show how one may adapt an existing computational science workflow to the autodidactic process described here. In the work performed here, this process enabled the discovery of functional relationships between optimal parameters in a four-dimensional search space with known bounds but an unknown resolution requirement. We trained the model over a coarse grid in $k_0$ and $t_1$, and learned functions for $x_1$ and $\omega _1$. Using gradient descent here allows an escape from the curse of dimensionality and reduces the problem from a four-dimensional search to a two-dimensional search $+$ two-dimensional gradient descent.

This discovery process is not limited to differentiable simulations. While in figure 1, $\mathcal {V}$ represents a PDE solve, it only needs to be a AD-enabled function that is a model for a physical system. For example, rather than a PDE solve, $\mathcal {V}$ could represent a pre-trained neural-network-based emulator for experimental data. In such a scenario, one may be able to learn forcing function parameters for an experiment using the proposed workflow.

In the neural network literature, the gradient required for the update is $\partial \mathcal {S} / \partial \theta = \partial \mathcal {S} / \partial \mathcal {G} \times \partial \mathcal {G} / \partial \theta$. We see that this is the same as (2.5) after the addition of one more node in the computational graph for $\mathcal {V}$, the function that models the physical system. This allows the neural network training process to become unsupervised and data efficient.

It remains to be determined how the physical mechanism discussed here behaves over a range of amplitudes and collision frequencies in addition to the wavenumber scan we perform here. Reduced models of the wavepacket dynamics in SRS remain useful for the development of inertial confinement fusion schemes where laser–plasma instabilities occur.

Differentiating through an entire simulation, in this case, required the intermediate storage of all $N_t \times N_x \times N_v$ arrays representing the distribution function in time and phase space. To enable this on a single NVIDIA T4 GPU, we use gradient checkpointing (Griewank & Walther 2000; Chen et al. 2016), as implemented in JAX, at every timestep. That is, we require that the Jacobian of the function representing a single timestep recompute its internal linearization points rather than save them to memory during the forward pass. Another technique by which we can save on the memory cost is by improving the efficiency of the accumulation of the Jacobian e.g. as proposed in Naumann (2004) and similar works. An example application towards a PDE is provided by Oktay et al. (2021). A final route to memory savings we wish to highlight here is by solving a backwards-in-time ordinary-differential equation for the adjoint (Pontryagin 1987; Chen et al. 2018). For modelling systems of higher-dimensionality, larger time scales and other constraints that require larger in-memory calculations, it will be crucial to leverage methods that help circumvent the memory burden encountered with differentiable simulations.