A multi-language auto-differentiation module and its application to a parallel particle-in-cell code on distributed computers

TL;DR

This work tackles the need for efficient, transparent auto-differentiation in large-scale accelerator simulations by developing a fast forward-mode TPSA-based AD module that stores first derivatives in a compact array and supports multi-language implementations. The module is integrated into a two-dimensional MPI-based parallel PIC code to enable sensitivity analysis and gradient-based optimization of lattice parameters under space-charge effects, with demonstrated accuracy against finite-difference methods and substantial speedups over traditional TPSA approaches. The approach facilitates in-situ computation of derivatives during simulation, enabling practical optimization workflows for accelerator design in distributed computing environments. Overall, the paper provides a portable, low-overhead tool for differentiable simulations that can enhance design, control, and optimization in beam dynamics and related fields.

Abstract

The auto differentiable simulation is a type of simulation that outputs of the simulation include not only the simulation result itself, but also their derivatives with respect to various input parameters. It provides an efficient method to study sensitivity of the simulation results with respect to the input parameters. Furthermore, it can be used in gradient based optimization methods for rapidly optimizing design parameters. In this paper, we present the development of a fast and transparent auto-differentiation module/class designed for easy integration into numerous simulation codes. As an application, this auto-differentiation module is integrated into a parallel particle-in-cell code with message passing interface (MPI) on distributed memory computers.

Figures (12)

• Figure 1: Layout of a FODO lattice (left) and derivatives of final horizontal and vertical emittances with respect to seven lattice parameters from the single differentiable self-consistent space-charge simulation and from the finite difference approximation to the first derivative qiang2023 (right).
• Figure 2: A section of the auto differentiation module class definition in C++.
• Figure 3: A C++ code showing the usage of the AD class module.
• Figure 4: Relative Twiss parameter changes with respect to 7 lattice parameters from a single auto differentiable simulation and from central finite difference approximation with 14 simulations.
• Figure 5: Schematic plot of two periodic FODO lattices and a section of matching section between them.