Discovering Physics and Improving Simulations using Neural Networks and Differentiable Kinetic Simulations

Modern libraries with implementations of Automatic Differentiation (AD) for machine learning can be used as general-purpose scientific computing libraries. We use two such libraries, JAX and PyTorch, to build differentiable kinetic plasma physics solvers i.e. solvers where each term and iteration in the solver is differentiable end-to-end. We leverage the ability to acquire fast and exact (to machine precision) gradients to 1/ discover novel physics using Vlasov-Fokker-Planck simulations and to 2/ develop improved numerical techniques for Particle-In-Cell modeling. In each, we judiciously place neural networks to learn functions which fulfill a holistic objective inline with the kinetic simulation. We also discuss the performance as well as the validation process for differentiable simulations.