Differentiable Programming for Plasma Physics: From Diagnostics to Discovery and Design

Abstract

Differentiable programming, enabled by automatic differentiation (AD), provides a robust framework for gradient-based optimization in computational plasma physics. While optimization is often only used towards design, we demonstrate that it can also be used for discovery and bridging the gap towards multi-scale modeling. We discuss four applications: (1) discovering novel nonlinear plasma phenomena, including a previously unknown superadditive wavepacket interaction regime, by optimizing differentiable kinetic simulations; (2) learning hidden variables that capture spatiotemporally non-local kinetic effects in fluid simulations, enabling hydrodynamic models to reproduce large Knudsen number physics typically requiring kinetic solvers; (3) accelerating Thomson scattering analysis by over $140\times$ while enabling extraction of velocity distribution functions with $\mathcal{O}(10^3)$ parameters; and (4) inverse design of spatiotemporal laser pulses that achieve target far-field behavior where full space-time coupling improves performance by $15\times$ over spatial or temporal optimization alone. These examples illustrate that differentiable programming not only accelerates existing design and inference workflows but enables qualitatively new capabilities, from algorithmic physics discovery to high-dimensional inference and design previously considered intractable.

Figures (6)

• Figure 1: Computational workflows enabled by differentiable programming. (a) Manual iteration. (b) Parameter scan with cost scaling $\mathcal{O}(k^N)$. (c) Gradient-based optimization using reverse-mode AD. (d) Function learning through a neural network embedded in a differentiable solver. The four applications in this article span stages (c) and (d).
• Figure 2: Superadditive wavepacket interactions discovered through optimization. Electric field amplitude $E^2(x)$ at t = 400, 900, and 1300 $\omega_p^{-1}$ for (a) first wavepacket only, (b) second wavepacket only, and (c) both with optimized excitation parameters. Late in time, the isolated second wavepacket has damped away while the combined system retains substantial energy—the whole exceeds the sum of its parts. From Joglekar and Thomas joglekar_unsupervised_2022.
• Figure 3: Generalization to finite-length wavepackets in a domain 100$\times$ larger than the training geometry. (a) Vlasov simulation shows non-uniform damping with the rear eroding faster than the front. (b) Local closure produces uniform damping. (c) Learned hidden-variable closure reproduces the kinetic behavior without retraining. (d) The hidden variable $\delta$ grows at the wavepacket location and advects forward at the phase velocity, encoding spatial memory. From Joglekar and Thomas joglekar_machine_2023.
• Figure 4: Computation time versus number of fitting parameters (lineouts) for Thomson-scattering analysis using finite differencing (blue) and automatic differentiation (red), on CPU (solid) and GPU (dashed). Reverse-mode AD combined with GPU acceleration yields over two orders of magnitude speedup, with the advantage increasing for larger parameter counts.
• Figure 5: Plasma conditions extracted from temporally resolved Thomson-scattering data: (a) electron density, (b) electron temperature, and (c) super-Gaussian order of the electron distribution. Blue curves show fitted values at every pixel; shaded regions indicate $3\sigma$ uncertainty from the Hessian (blue) and from the standard deviation within a 5-pixel moving window matching the diagnostic resolution (red). The two uncertainty estimates are in good agreement. From Milder et al.milder_qualitative_2024.