Skip to main content

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

Project description

format pytest Documentation Status codestyle License: GPL v3

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)

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:

@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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

cheetah-accelerator-0.6.3.tar.gz (58.7 kB view hashes)

Uploaded Source

Built Distribution

cheetah_accelerator-0.6.3-py3-none-any.whl (49.2 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page