Cheetah: Bridging the Gap Between Machine Learning and Particle
  Accelerator Physics with High-Speed, Differentiable Simulations
Jan Kaiser Chenran Xu Annika Eichler Andrea Santamaria Garcia
Machine learning has emerged as a powerful solution to the modern challenges
in accelerator physics. However, the limited availability of beam time, the
computational cost of simulations, and the high-dimensionality of optimisation
problems pose significant challenges in generating the required data for
training state-of-the-art machine learning models. In this work, we introduce
Cheetah, a PyTorch-based high-speed differentiable linear-beam dynamics code.
Cheetah enables the fast collection of large data sets by reducing computation
times by multiple orders of magnitude and facilitates efficient gradient-based
optimisation for accelerator tuning and system identification. This positions
Cheetah as a user-friendly, readily extensible tool that integrates seamlessly
with widely adopted machine learning tools. We showcase the utility of Cheetah
through five examples, including reinforcement learning training,
gradient-based beamline tuning, gradient-based system identification,
physics-informed Bayesian optimisation priors, and modular neural network
surrogate modelling of space charge effects. The use of such a high-speed
differentiable simulation code will simplify the development of machine
learning-based methods for particle accelerators and fast-track their
integration into everyday operations of accelerator facilities.