Implementing custom elements in Cheetah¶
Cheetah is designed to be extensible, making it easy to add custom elements when the existing elements do not meet your needs. The following guide will illustrate with examples how to implement elements following linear beam dynamics as well as elements following more complex physics.
If you implement a new element and think it could be useful for others, please consider contributing it directly to Cheetah!
To make any contributions to Cheetah, please open a pull request on the Cheetah GitHub repository. To do so, fork the repository, implement your new element and then open a pull request. One of the Cheetah maintainers will then review your code, maybe ask for some improvements and then merge it onto the master branch.
The most important part of implementing a new element is to create a subclass of the cheetah.accelerator.Element class. The latter is an abstract base class that defines the interface for all elements in Cheetah. By subclassing it, you can define the specific behaviour and properties of your custom element.
Start by creating a new Python file for your custom element in the cheetah/accelerator/ directory of the Cheetah repository. For example, if you want to create a custom element called MyCustomElement, you would create a file named my_custom_element.py.
The following code snippet shows how to implement a simple custom element. You can paste it into your new file and then follow the instructions in the TODO comments to complete the implementation of the element itself.
[1]:
from typing import Literal
import matplotlib.pyplot as plt
import torch
from cheetah.accelerator.element import Element
from cheetah.particles import Beam, ParameterBeam, ParticleBeam, Species
from cheetah.track_methods import base_rmatrix, base_ttensor, misalignment_matrix
from cheetah.utils import cache_transfer_map
class MyCustomElement(Element):
"""
TODO: Describe the purpose of your custom element here. For example: A quadrupole
magnet that focuses the beam in the horizontal direction and defocuses it in the
vertical direction.
:param length: Length in meters. TODO: If your element has a physical length, keep
this line (but delete the TODO comment). If it does not, like for example a
`Marker`, delete this line.
:param a_tensor_parameter: TODO: Specify the parameters of your element here,
following the format `:param parameter_name: Description of the parameter.`.
:param not_a_tensor_parameter: TODO: Specify the parameters of your element here,
following the format `:param parameter_name: Description of the parameter.`.
:param name: Unique identifier of the element.
:param sanitize_name: Whether to sanitise the name to be a valid Python variable
name.
:param metadata: Dictionary of arbitrary, serialisable annotations attached to the
element.
:param device: Device on which to create the element's tensors.
:param dtype: Data type of the element's tensors.
"""
def __init__(
self,
length: torch.Tensor,
# TODO: Replace `a_tensor_parameter` with the actual parameters of your element.
# Please remember to add correct type annotations, and to update the docstring
# above accordingly. Note that parameters with default values must not specify
# these values here in the signature of `__init__`, but rather have them set
# in the body of `__init__` (see below).
a_tensor_parameter: torch.Tensor | None = None,
not_a_tensor_parameter: Literal[
"for_example_a_mode", "another_mode"
] = "for_example_a_mode",
name: str | None = None,
sanitize_name: bool = False,
metadata: dict | None = None,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
) -> None:
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__(
name=name, sanitize_name=sanitize_name, metadata=metadata, **factory_kwargs
)
# TODO: Delete the following line if your element does not have a physical
# length.
self.length = length
# TODO: Assign with `self.parameter_name = ...` only the parameters that are not
# `torch.Tensor`s.
self.not_a_tensor_parameter = not_a_tensor_parameter
# TODO: Register all your torch.Tensor parameters like the following. If they
# are `torch.Tensor`s, do not simply assign them with
# `self.my_parameter = ...`! Note that `length` is already registered by the
# call to `super().__init__()`, so the registration of `length` is not needed,
# and the assignment `self.length = ...` is allowed.
self.register_buffer_or_parameter(
"a_tensor_parameter",
(
a_tensor_parameter
if a_tensor_parameter is not None
else torch.tensor(0.0, **factory_kwargs)
),
)
# TODO: Add other initialisation code you may need here.
@cache_transfer_map
def first_order_transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor:
# TODO: If your element only has first-order effects, you should delete the
# `track` method below and only implement this method, which computes the
# first-order transfer map. Cheetah will then automatically handle the
# tracking for you.
# The transfer map needs to be a tensor of shape (..., 7, 7), where all but
# the last two dimensions are vector dimensions, the upper left 6x6 block is
# the first-order transfer map you probably know from accelerator physics. The
# last row should be zeros, except for the last element, which should be 1.0.
# The last column can be used to add zero-th order effects.
# There is no need to write a docstring for this method, as the superclass
# `Element` already provides one.
# Example placeholder: identity transfer map (no effect).
R = torch.eye(7, device=energy.device, dtype=energy.dtype)
return R
@cache_transfer_map
def second_order_transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor:
# TODO: If your element has second-order effects, you can implement this method,
# which computes the second-order transfer map. The transfer map needs to be a
# tensor of shape (..., 7, 7, 7).
# There is no need to write a docstring for this method, as the superclass
# `Element` already provides one.
# Example placeholder: zero second-order tensor with identity first-order map.
T = torch.zeros((7, 7, 7), device=energy.device, dtype=energy.dtype)
T[:, 6, :] = torch.eye(7, device=energy.device, dtype=energy.dtype)
return T
def track(self, incoming: Beam) -> Beam:
# TODO: This is the main method where you implement your physics. It transforms
# a beam entering your element into a beam exiting your element.
# If all you need is linear first- or second-order tracking, you can delete
# this method and only implement `first_order_transfer_map` and/or
# `second_order_transfer_map` above. Cheetah will then automatically handle
# the tracking for you.
# TODO: You will probably have to distinguish between the two different types of
# beams in Cheetah, `ParameterBeam` and `ParticleBeam`. See the documentation
# for more information on them.
if isinstance(incoming, ParameterBeam):
# TODO: Implement your logic for `ParameterBeam`s here. If your element has
# complex physics, especially if those include interactions between
# particles, it is often legitimate to fall back to the first-order
# transfer map for `ParameterBeam`s. You can do so by calling
# `super().track(incoming)`.
return super().track(incoming)
elif isinstance(incoming, ParticleBeam):
# TODO: Implement your logic for `ParticleBeam`s here. Return a new
# `ParticleBeam` object with the transformed particles.
outgoing_particles = torch.rand_like(
incoming.particles
) # Placeholder for transformed particles
return ParticleBeam(
particles=outgoing_particles,
energy=incoming.energy.clone(),
particle_charges=incoming.particle_charges.clone(),
survival_probabilities=incoming.survival_probabilities.clone(),
s=incoming.s + self.length,
species=incoming.species.clone(),
)
else:
raise TypeError(
f"Unsupported beam type: {type(incoming)}. Expected ParameterBeam or "
"ParticleBeam."
)
@property
def is_skippable(self) -> bool:
# TODO: If your element follows strictly linear beam dynamics, Cheetah can
# perform speed optimisations. To see whether it can do so, it checks whether
# the element `is_skippable` property returns `True`. If your element has only
# first-order effects, you can return `True` here. If it has second-order
# effects, you should return `False`. In some cases it makes sense to
# determine this dynamically.
return False
@property
def is_active(self) -> bool:
# TODO: Return `True` if your element is active, i.e. "on", and `False` if it is
# inactive, i.e. "off". If your element does not support being
# active/inactive, you can delete this property.
return True
def split(self, resolution: torch.Tensor) -> list[Element]:
# TODO: Implement your logic for splitting your element longitudinally here.
# This is used to provide beams at multiple positions along the element,
# specifically no longer than `resolution` meters apart.
# If your element cannot be split, or is a lengthless element, you can simply
# omit this method and the base class implementation will return `[self]`.
num_splits = (self.length.abs().max() / resolution).ceil().int()
return [
MyCustomElement(
self.length / num_splits,
a_tensor_parameter=self.a_tensor_parameter,
not_a_tensor_parameter=self.not_a_tensor_parameter,
name=f"{self.name}_split_{i}",
sanitize_name=self.is_name_sanitized(),
dtype=self.length.dtype,
device=self.length.device,
)
for i in range(num_splits)
]
def plot(self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None) -> plt.Axes:
# TODO: Cheetah can plot using in Matplotlib a straight representation of the
# lattice. You can implement here how your element should be represented.
raise NotImplementedError
@property
def defining_features(self) -> list[str]:
# TODO: Add to this list all the properties that define your element, i.e. those
# that should be saved and loaded, when the element is serialised. Please make
# sure to add the list of your element's defining features to the end of the
# list returned by `super().defining_features`.
return super().defining_features + [
"length", "a_tensor_parameter", "not_a_tensor_parameter"
]
After you have finished your main implementation, you will need to do some housekeeping to make sure that your new element is properly usable in Cheetah.
Add your element to the
cheetah/accelerator/__init__.pyfile, so that it can be imported from thecheetah.acceleratormodule.Add your element to the
cheetah/__init__.pyfile, so that it can be imported directly from thecheetahmodule ascheetah.MyCustomElement.If you want to add a 3D visual representation for your new element, contribute a
.glbmesh file named after your class in snake_case (e.g.my_custom_element.glbforMyCustomElement) to the 3D assets repository. Cheetah’s base classElementwill then automatically resolve and fetch this mesh for 3D visualisations, so you do not need to implementto_meshyourself.
If you further want to contribute your element to the main Cheetah repository, you will also need to take care of the following:
In
CHANGELOG.md…Create an entry to advertise your new element under “Features” for the currently under development version.
Add your name and GitHub handle at the bottom of the list of “First Time Contributors”.
In
docs/accelerator.rst, add an entry like the following for your element in the correct alphabetical place:.. automodule:: accelerator.my_custom_element :members: :undoc-members:
Write a test for your element in a new appropriately named file in the
testsdirectory. The file name should follow the patterntest_<element_name>.py. This test should test the physics of your element, for example by checking something that should always hold true about those physics, or by checking an example case against results from another code or from a paper.Add a minimal dictionary of arguments to your element to the
ELEMENT_SUBCLASSES_ARGSdictionary intests/conftest.pyso that common element functionality (such as clone, parameter types, device/dtype conversions, and vectorised tracking) is automatically tested for your new element.Support serialisation and conversion to/from other formats:
In
cheetah/latticejson.py(LatticeJSON): Ensure your custom element is exported incheetah/__init__.pyand constructor parameters align withdefining_featuresto automatically support LatticeJSON.In
cheetah/converters/ocelot.py(Ocelot): Add a check toconvert_elementusingelif isinstance(element, ocelot.YourElementClass):to map Ocelot parameters to your Cheetah class.In
cheetah/converters/elegant.py(Elegant): Add a check toconvert_elementusingelif parsed["element_type"] == "your_elegant_type_code":inside the parsed dictionary block to instantiate your class, and list understood properties undervalidate_understood_properties.In
cheetah/converters/bmad.py(Bmad): Add a check toconvert_elementusingelif bmad_parsed["element_type"] == "your_bmad_type_code":to map the element properties to your class.
In both
README.mdanddocs/index.rst…… add your name and GitHub handle at the bottom of the list of contributors under “Author Contributions”.
… if not there yet, add the logo of your institution.
… if needed, add to the funding string/acknowledgements section.
[ ]: