Cheetah

Fast and differentiable particle accelerator optics simulation for reinforcement learning and optimisation applications.

7
mentions
11
contributors

What Cheetah can do for you

format pytest Documentation Status codestyle License: GPL v3

Cheetah

Cheetah is a high-speed differentiable beam dynamics code specifically design to support machine learning applications for particle accelerators.

Its speed helps generate data fast, for example for extremely data-hungry reinforcement learning algorithms, while its differentiability can be used for a variety of applications, including accelerator tuning, system identification and physics-informed prior means for Bayesian optimisation. Its native integration with machine learning toolchains around PyTorch also makes Cheetah an ideal candidate for coupling of physics-based and neural network beam dynamics models that remain fast and differentiable.

To learn more about what Cheetah can do, we recommend reading our PRAB paper. To learn how to use Cheetah, we refer to the example notebooks in the Cheetah documentation.

Installation

Simply install Cheetah from PyPI by running the following command.

pip install cheetah-accelerator

How To Use

A sequence of accelerator elements (or a lattice) is called a Segment in Cheetah. You can create a Segment as follows

import torch
from cheetah import BPM, Drift, HorizontalCorrector, Segment, VerticalCorrector

segment = Segment(
    elements=[
        BPM(name="BPM1SMATCH"),
        Drift(length=torch.tensor(1.0)),
        BPM(name="BPM6SMATCH"),
        Drift(length=torch.tensor(1.0)),
        VerticalCorrector(length=torch.tensor(0.3), name="V7SMATCH"),
        Drift(length=torch.tensor(0.2)),
        HorizontalCorrector(length=torch.tensor(0.3), name="H10SMATCH"),
        Drift(length=torch.tensor(7.0)),
        HorizontalCorrector(length=torch.tensor(0.3), name="H12SMATCH"),
        Drift(length=torch.tensor(0.05)),
        BPM(name="BPM13SMATCH"),
    ]
)

Alternatively you can create a segment from an Ocelot cell by running

segment = Segment.from_ocelot(cell)

All elements can be accesses as a property of the segment via their name. The strength of a quadrupole named AREAMQZM2 for example, may be set by running

segment.AREAMQZM2.k1 = torch.tensor(4.2)

In order to track a beam through the segment, simply call the segment like so

outgoing_beam = segment.track(incoming_beam)

You can choose to track either a beam defined by its parameters (fast) or by its particles (precise). Cheetah defines two different beam classes for this purpose and beams may be created by

beam1 = ParameterBeam.from_parameters()
beam2 = ParticleBeam.from_parameters()

It is also possible to load beams from Ocelot ParticleArray or Astra particle distribution files for both types of beam

ocelot_beam = ParticleBeam.from_ocelot(parray)
astra_beam = ParticleBeam.from_astra(filepath)

You may plot a segment with reference particle traces bay calling

segment.plot_overview(beam=beam)

Overview Plot

where the optional keyword argument beam is the incoming beam represented by the reference particles. Cheetah will use a default incoming beam, if no beam is passed.

Cite Cheetah

If you use Cheetah, please cite the following two papers:

@article{kaiser2024cheetah,
    title        = {Bridging the gap between machine learning and particle accelerator physics with high-speed, differentiable simulations},
    author       = {Kaiser, Jan and Xu, Chenran and Eichler, Annika and Santamaria Garcia, Andrea},
    year         = 2024,
    month        = {May},
    journal      = {Phys. Rev. Accel. Beams},
    publisher    = {American Physical Society},
    volume       = 27,
    pages        = {054601},
    doi          = {10.1103/PhysRevAccelBeams.27.054601},
    url          = {https://link.aps.org/doi/10.1103/PhysRevAccelBeams.27.054601},
    issue        = 5,
    numpages     = 17
}
@inproceedings{stein2022accelerating,
    title        = {Accelerating Linear Beam Dynamics Simulations for Machine Learning Applications},
    author       = {Stein, Oliver and Kaiser, Jan and Eichler, Annika},
    year         = 2022,
    booktitle    = {Proceedings of the 13th International Particle Accelerator Conference}
}

For Developers

Activate your virtual environment. (Optional)

Install the cheetah package as editable

pip install -e .

We suggest installing pre-commit hooks to automatically conform with the code formatting in commits:

pip install pre-commit
pre-commit install

Acknowledgements

We acknowledge the contributions of the following people to the development of Cheetah: Jan Kaiser, Chenran Xu, Oliver Stein, Annika Eichler, Andrea Santamaria Garcia and others.

The work to develop Cheetah has in part been funded by the IVF project InternLabs-0011 (HIR3X) and the Initiative and Networking Fund by the Helmholtz Association (Autonomous Accelerator, ZT-I-PF-5-6). In addition, we acknowledge support from DESY (Hamburg, Germany) and KIT (Karlsruhe, Germany), members of the Helmholtz Association HGF.

Participating organisations

Deutsches Elektronen-Synchrotron DESY
Karlsruhe Institute of Technology (KIT)
Law
SLA
Uni

Reference papers

Mentions

Contributors

JK
Author/Developer/Maintainer
Deutsches Elektronen-Synchrotron DESY
CX
Chenran Xu
AE
Annika Eichler
Deutsches Elektronen-Synchrotron DESY
ASG
Andrea Santamaria Garcia
CH
Christian Hespe
Deutsches Elektronen-Synchrotron DESY
GC
Grégoire Charleux
Lawrence Berkeley National Laboratory
RL
Remi Lehe
Lawrence Berkeley National Laboratory
AH
Axel Huebl
Lawrence Berkeley National Laboratory
JG
Juan Pablo Gonzalez-Aguilera
RR
Ryan Roussel
SLAC National Accelerator Laboratory
AE
Auralee Edelen
SLAC National Accelerator Laboratory

Helmholtz Program-oriented Funding IV

Research Field
Research Program
PoF Topic
6 Matter
6.2 Matter and Technologies
6.2.1 Accelerator Research and Development
6.2.3 Data Management and Analysis
  • 6 Matter
    • 6.2 Matter and Technologies
      • 6.2.1 Accelerator Research and Development
      • 6.2.3 Data Management and Analysis