Cheetah: Fast simulation for linear beam dynamics

Introduction

The successful training of machine learning methods requires large datasets. Acquiring these datasets is often one of the key challenges one must solve when using machine learning. With particle accelerators in particular, beam time for collecting datasets is a sparse and expensive resource, and existing simulation codes, while accurate, can take a long time to compute even on high-performance computers.

Our simulation code [Cheetha](https://github.com/desy-ml/cheetah) solves this problem through intelligent optimisation of computations built on top of the PyTorch Python package, and by trading accuracy for speed. This trade is possible because methods like domain randomisation allow machine learning models to perform well in the real world even when they were trained on a fairly inaccurate model of the latter. In doing so, Cheetah achieves a speed-up of multiple orders of magnitude when compared to existing particle accelerators simulation codes, making large scale dataset collection and reinforcement learning trainings much more feasible.

In addition, being built on top of PyTorch, Cheetah provides a perfect environment to integrate modular learned surrogate models and differentiable particle accelerator simulations. Differentiable simulations open up a plethora of new opportunities, from fast gradient-based tune and lattice optimimsation all the way to system identification.

Purpose/goal/objectives:

  • Provide a fast simulation toolkit for particle accelerators suitable for training machine learning models that require ever larger training datasets.
  • Ease and speed-up the development and deployment of machine learning methods to particle accelerators facilities, with a particular focus on reinforcement learning applications.
  • Provide an easy-to-use and easily extensible differentiable simulation toolkit for gradient-based particle accelerator optimisation and system identification.

References:

  • O. Stein, J. Kaiser, A. Eichler - "Accelerating Linear Beam Dynamics Simulations for Machine Learning Applications" in Proceedings of the 13th International Particle Accelerator Conference, 2022.