Investigating Automatic Differentiation in Blast/WarpX

Project description

NERSC is seeking an enthusiastic summer-plus intern for a term of up to six months to investigate integrating the Enzyme Automatic Differentiator into the Gordon-Bell award-winning PIC code, WarpX, and its sibling code ImpactX.

The results of this internship will contribute to production-quality, daily-use open source modeling software with a direct impact on next-generation particle accelerators and fusion energy science research. 

Background

When designing novel particle accelerators, such as the beamlines driving the latest Neutrino experiments and compact laser-plasma accelerators for novel light sources, a large number of variables must be considered to produce high-quality, high-charge, high-energy particle beams. To produce these beams, an initial set of parameters is chosen, and a simulation is completed to calculate the relativistic physics and non-linear dynamics in an accelerator. An optimization strategy, such as conjugate gradient or Bayesian Optimization, is then applied to produce a new, improved prediction of the most likely parameter values that yield the desired beam. The simulation is then run again with the new parameters. This is repeated until appropriate parameters are identified that reliably produce the desired particle beam.

The Beam, Plasma & Accelerator Simulation Toolkit (BLAST) team, in close collaboration with NERSC, is investigating the use of automatic differentiation to improve the time-to-solution for WarpX/ImpactX simulations in high-energy physics, with an overarching impact on high-energy physics and fusion-energy science challenges. Automatic differentiation (AD) tools, like Enzyme, instrument specific pieces of a scientific code to automatically track the influence and sensitivity of important simulation parameters. 

Novel beamline accelerator physics present high-dimensional optimization problems. This illustration shows an example beamline setup for this project, including possible variables for optimizing the beam.

Specifically, AD calculates the relationship between changes in input variables and returns the corresponding derivatives. These derivatives can be used by optimization strategies to greatly improve their parameter guesses, thereby reducing the number of simulations needed to find the beam parameters. And, gradient tracking provides the basis for innovative ML workflows, such as training embedded neural networks in high-performance simulations.

Enabling AD in BLAST slows individual simulations (by capturing simulation results and their gradients), but reduces the number of simulations needed, producing an overall savings (time-to-solution/GPU-hours) for high-dimensional optimization problems, as in accelerator physics.

Axel Huebl