Fast and differentiable particle accelerator optics simulation for reinforcement learning and optimisation applications.
Project description
Cheetah
Cheetah is a particle tracking accelerator we built specifically to speed up the training of reinforcement learning models.
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
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.
Cite Cheetah
If you use Cheetah, please cite the following two papers:
@misc{kaiser2024cheetah,
title = {Cheetah: 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},
eprint = {2401.05815},
archiveprefix = {arXiv},
primaryclass = {physics.acc-ph}
}
@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.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for cheetah-accelerator-0.6.2.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | f5a0c83d84e33b24cf332f272c7ae822409af2564e8d6070f305271c8b7f3b54 |
|
MD5 | 506ff201e726c6eb00f8fd47683f04ef |
|
BLAKE2b-256 | d39d04cfa24344f5b705954e573bc1ae27669107319b7cfd51d1bf6b84482213 |
Hashes for cheetah_accelerator-0.6.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | c34d1689d2576a703aebd3392507863b007ed69fb6b626e090dceeb8503ce3c9 |
|
MD5 | d26f6665ad5e09a60118a33ecf9cae0b |
|
BLAKE2b-256 | dbb3e7504d6da4c634e2615d4cee0a203d35be56cddeb087afdbc8916b5d2183 |