Fast and differentiable particle accelerator optics simulation for reinforcement learning and optimisation applications.
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.
Simply install Cheetah from PyPI by running the following command.
pip install cheetah-accelerator
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)
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.
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}
}
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
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.