{ "cells": [ { "cell_type": "markdown", "id": "2ff3b814", "metadata": {}, "source": [ "# Implementing custom elements in Cheetah\n", "\n", "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.\n", "\n", "**If you implement a new element and think it could be useful for others, please consider contributing it directly to Cheetah!**\n" ] }, { "cell_type": "markdown", "id": "a863a3d2", "metadata": {}, "source": [ "To make any contributions to Cheetah, please open a pull request on the [Cheetah GitHub repository](https://github.com/desy-ml/cheetah). 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.\n" ] }, { "cell_type": "markdown", "id": "0ebe23f1", "metadata": {}, "source": [ "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.\n", "\n", "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`.\n", "\n", "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.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "567c7a26", "metadata": {}, "outputs": [], "source": [ "from typing import Literal\n", "\n", "import matplotlib.pyplot as plt\n", "import torch\n", "\n", "from cheetah.accelerator.element import Element\n", "from cheetah.particles import Beam, ParameterBeam, ParticleBeam, Species\n", "from cheetah.track_methods import base_rmatrix, base_ttensor, misalignment_matrix\n", "from cheetah.utils import cache_transfer_map\n", "\n", "\n", "class MyCustomElement(Element):\n", " \"\"\"\n", " TODO: Describe the purpose of your custom element here. For example: A quadrupole\n", " magnet that focuses the beam in the horizontal direction and defocuses it in the \n", " vertical direction.\n", "\n", " :param length: Length in meters. TODO: If your element has a physical length, keep \n", " this line (but delete the TODO comment). If it does not, like for example a \n", " `Marker`, delete this line.\n", " :param a_tensor_parameter: TODO: Specify the parameters of your element here, \n", " following the format `:param parameter_name: Description of the parameter.`.\n", " :param not_a_tensor_parameter: TODO: Specify the parameters of your element here, \n", " following the format `:param parameter_name: Description of the parameter.`.\n", " :param name: Unique identifier of the element.\n", " :param sanitize_name: Whether to sanitise the name to be a valid Python variable\n", " name.\n", " :param metadata: Dictionary of arbitrary, serialisable annotations attached to the\n", " element.\n", " :param device: Device on which to create the element's tensors.\n", " :param dtype: Data type of the element's tensors.\n", " \"\"\"\n", "\n", " def __init__(\n", " self,\n", " length: torch.Tensor,\n", " # TODO: Replace `a_tensor_parameter` with the actual parameters of your element. \n", " # Please remember to add correct type annotations, and to update the docstring \n", " # above accordingly. Note that parameters with default values must not specify \n", " # these values here in the signature of `__init__`, but rather have them set \n", " # in the body of `__init__` (see below).\n", " a_tensor_parameter: torch.Tensor | None = None,\n", " not_a_tensor_parameter: Literal[\n", " \"for_example_a_mode\", \"another_mode\"\n", " ] = \"for_example_a_mode\",\n", " name: str | None = None,\n", " sanitize_name: bool = False,\n", " metadata: dict | None = None,\n", " device: torch.device | None = None,\n", " dtype: torch.dtype | None = None,\n", " ) -> None:\n", " factory_kwargs = {\"device\": device, \"dtype\": dtype}\n", " super().__init__(\n", " name=name, sanitize_name=sanitize_name, metadata=metadata, **factory_kwargs\n", " )\n", "\n", " # TODO: Delete the following line if your element does not have a physical \n", " # length.\n", " self.length = length\n", "\n", " # TODO: Assign with `self.parameter_name = ...` only the parameters that are not \n", " # `torch.Tensor`s.\n", " self.not_a_tensor_parameter = not_a_tensor_parameter\n", "\n", " # TODO: Register all your torch.Tensor parameters like the following. If they\n", " # are `torch.Tensor`s, do not simply assign them with \n", " # `self.my_parameter = ...`! Note that `length` is already registered by the \n", " # call to `super().__init__()`, so the registration of `length` is not needed, \n", " # and the assignment `self.length = ...` is allowed.\n", " self.register_buffer_or_parameter(\n", " \"a_tensor_parameter\",\n", " (\n", " a_tensor_parameter\n", " if a_tensor_parameter is not None\n", " else torch.tensor(0.0, **factory_kwargs)\n", " ),\n", " )\n", "\n", " # TODO: Add other initialisation code you may need here.\n", "\n", " @cache_transfer_map\n", " def first_order_transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor:\n", " # TODO: If your element only has first-order effects, you should delete the \n", " # `track` method below and only implement this method, which computes the \n", " # first-order transfer map. Cheetah will then automatically handle the \n", " # tracking for you. \n", " # The transfer map needs to be a tensor of shape (..., 7, 7), where all but \n", " # the last two dimensions are vector dimensions, the upper left 6x6 block is \n", " # the first-order transfer map you probably know from accelerator physics. The \n", " # last row should be zeros, except for the last element, which should be 1.0. \n", " # The last column can be used to add zero-th order effects.\n", " # There is no need to write a docstring for this method, as the superclass \n", " # `Element` already provides one.\n", "\n", " # Example placeholder: identity transfer map (no effect).\n", " R = torch.eye(7, device=energy.device, dtype=energy.dtype)\n", "\n", " return R\n", "\n", " @cache_transfer_map\n", " def second_order_transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor:\n", " # TODO: If your element has second-order effects, you can implement this method,\n", " # which computes the second-order transfer map. The transfer map needs to be a \n", " # tensor of shape (..., 7, 7, 7).\n", " # There is no need to write a docstring for this method, as the superclass \n", " # `Element` already provides one.\n", "\n", " # Example placeholder: zero second-order tensor with identity first-order map.\n", " T = torch.zeros((7, 7, 7), device=energy.device, dtype=energy.dtype)\n", " T[:, 6, :] = torch.eye(7, device=energy.device, dtype=energy.dtype)\n", "\n", " return T\n", "\n", " def track(self, incoming: Beam) -> Beam:\n", " # TODO: This is the main method where you implement your physics. It transforms \n", " # a beam entering your element into a beam exiting your element.\n", " # If all you need is linear first- or second-order tracking, you can delete \n", " # this method and only implement `first_order_transfer_map` and/or \n", " # `second_order_transfer_map` above. Cheetah will then automatically handle \n", " # the tracking for you.\n", "\n", " # TODO: You will probably have to distinguish between the two different types of \n", " # beams in Cheetah, `ParameterBeam` and `ParticleBeam`. See the documentation \n", " # for more information on them.\n", " if isinstance(incoming, ParameterBeam):\n", " # TODO: Implement your logic for `ParameterBeam`s here. If your element has \n", " # complex physics, especially if those include interactions between \n", " # particles, it is often legitimate to fall back to the first-order \n", " # transfer map for `ParameterBeam`s. You can do so by calling \n", " # `super().track(incoming)`.\n", " return super().track(incoming)\n", " elif isinstance(incoming, ParticleBeam):\n", " # TODO: Implement your logic for `ParticleBeam`s here. Return a new \n", " # `ParticleBeam` object with the transformed particles.\n", "\n", " outgoing_particles = torch.rand_like(\n", " incoming.particles\n", " ) # Placeholder for transformed particles\n", "\n", " return ParticleBeam(\n", " particles=outgoing_particles,\n", " energy=incoming.energy.clone(),\n", " particle_charges=incoming.particle_charges.clone(),\n", " survival_probabilities=incoming.survival_probabilities.clone(),\n", " s=incoming.s + self.length,\n", " species=incoming.species.clone(),\n", " )\n", " else:\n", " raise TypeError(\n", " f\"Unsupported beam type: {type(incoming)}. Expected ParameterBeam or \"\n", " \"ParticleBeam.\"\n", " )\n", "\n", " @property\n", " def is_skippable(self) -> bool:\n", " # TODO: If your element follows strictly linear beam dynamics, Cheetah can \n", " # perform speed optimisations. To see whether it can do so, it checks whether \n", " # the element `is_skippable` property returns `True`. If your element has only \n", " # first-order effects, you can return `True` here. If it has second-order \n", " # effects, you should return `False`. In some cases it makes sense to \n", " # determine this dynamically.\n", " return False\n", "\n", " @property\n", " def is_active(self) -> bool:\n", " # TODO: Return `True` if your element is active, i.e. \"on\", and `False` if it is \n", " # inactive, i.e. \"off\". If your element does not support being \n", " # active/inactive, you can delete this property.\n", " return True\n", "\n", " def split(self, resolution: torch.Tensor) -> list[Element]:\n", " # TODO: Implement your logic for splitting your element longitudinally here. \n", " # This is used to provide beams at multiple positions along the element, \n", " # specifically no longer than `resolution` meters apart. \n", " # If your element cannot be split, or is a lengthless element, you can simply \n", " # omit this method and the base class implementation will return `[self]`.\n", " num_splits = (self.length.abs().max() / resolution).ceil().int()\n", " return [\n", " MyCustomElement(\n", " self.length / num_splits,\n", " a_tensor_parameter=self.a_tensor_parameter,\n", " not_a_tensor_parameter=self.not_a_tensor_parameter,\n", " name=f\"{self.name}_split_{i}\",\n", " sanitize_name=self.is_name_sanitized(),\n", " dtype=self.length.dtype,\n", " device=self.length.device,\n", " )\n", " for i in range(num_splits)\n", " ]\n", "\n", " def plot(self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None) -> plt.Axes:\n", " # TODO: Cheetah can plot using in Matplotlib a straight representation of the \n", " # lattice. You can implement here how your element should be represented.\n", " raise NotImplementedError\n", "\n", " @property\n", " def defining_features(self) -> list[str]:\n", " # TODO: Add to this list all the properties that define your element, i.e. those \n", " # that should be saved and loaded, when the element is serialised. Please make \n", " # sure to add the list of your element's defining features to the end of the \n", " # list returned by `super().defining_features`.\n", " return super().defining_features + [\n", " \"length\", \"a_tensor_parameter\", \"not_a_tensor_parameter\"\n", " ]\n" ] }, { "cell_type": "markdown", "id": "5285532e", "metadata": {}, "source": [ "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.\n", "\n", "- Add your element to the `cheetah/accelerator/__init__.py` file, so that it can be imported from the `cheetah.accelerator` module.\n", "- Add your element to the `cheetah/__init__.py` file, so that it can be imported directly from the `cheetah` module as `cheetah.MyCustomElement`.\n", "- If you want to add a 3D visual representation for your new element, contribute a `.glb` mesh file named after your class in snake_case (e.g. `my_custom_element.glb` for `MyCustomElement`) to the [3D assets repository](https://github.com/desy-ml/3d-assets). Cheetah's base class `Element` will then automatically resolve and fetch this mesh for 3D visualisations, so you do not need to implement `to_mesh` yourself.\n" ] }, { "cell_type": "markdown", "id": "29857de4", "metadata": {}, "source": [ "If you further want to contribute your element to the main Cheetah repository, you will also need to take care of the following:\n", "\n", "- In `CHANGELOG.md` ...\n", " - Create an entry to advertise your new element under \"Features\" for the currently under development version.\n", " - Add your name and GitHub handle at the bottom of the list of \"First Time Contributors\".\n", "- In `docs/accelerator.rst`, add an entry like the following for your element in the correct alphabetical place:\n", " ```rst\n", " .. automodule:: accelerator.my_custom_element\n", " :members:\n", " :undoc-members:\n", " ```\n", "- Write a test for your element in a new appropriately named file in the `tests` directory. The file name should follow the pattern `test_.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.\n", "- Add a minimal dictionary of arguments to your element to the `ELEMENT_SUBCLASSES_ARGS` dictionary in `tests/conftest.py` so that common element functionality (such as clone, parameter types, device/dtype conversions, and vectorised tracking) is automatically tested for your new element.\n", "- Support serialisation and conversion to/from other formats:\n", " - In `cheetah/latticejson.py` (LatticeJSON): Ensure your custom element is exported in `cheetah/__init__.py` and constructor parameters align with `defining_features` to automatically support LatticeJSON.\n", " - In `cheetah/converters/ocelot.py` (Ocelot): Add a check to `convert_element` using `elif isinstance(element, ocelot.YourElementClass):` to map Ocelot parameters to your Cheetah class.\n", " - In `cheetah/converters/elegant.py` (Elegant): Add a check to `convert_element` using `elif parsed[\"element_type\"] == \"your_elegant_type_code\":` inside the parsed dictionary block to instantiate your class, and list understood properties under `validate_understood_properties`.\n", " - In `cheetah/converters/bmad.py` (Bmad): Add a check to `convert_element` using `elif bmad_parsed[\"element_type\"] == \"your_bmad_type_code\":` to map the element properties to your class.\n", "- In both `README.md` and `docs/index.rst` ...\n", " - ... add your name and GitHub handle at the bottom of the list of contributors under \"Author Contributions\".\n", " - ... if not there yet, add the logo of your institution.\n", " - ... if needed, add to the funding string/acknowledgements section.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5cf9dedf", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "cheetah-dev", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.14.0" } }, "nbformat": 4, "nbformat_minor": 5 }