Cheetah: Bridging the Gap Between Machine Learning and Particle Accelerator Physics with High-Speed, Differentiable Simulations

Abstract

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.

arXiv comment16 pages, 9 figures, 3 tables

References in corpus (6)

• Bayesian optimization of a free-electron laser
• Phase Space Reconstruction from Accelerator Beam Measurements Using Neural Networks and Differentiable Simulations
• High-fidelity Prediction of Megapixel Longitudinal Phase-space Images of Electron Beams using Encoder-Decoder Neural Networks
• Differentiable Preisach Modeling for Characterization and Optimization of Accelerator Systems with Hysteresis
• Differentiable self-consistent space-charge simulation for accelerator design
• Predicting the transverse emittance of space charge dominated beams using the phase advance scan technique and a fully connected neural network

Cited by in corpus (2)

• Optimisation of the Accelerator Control by Reinforcement Learning: A Simulation-Based Approach
• The Memory Scaling of Reverse-Mode Differentiation in Particle Accelerator Simulations with Space Charge