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:

@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.