diff --git a/.vscode/ltex.dictionary.en-GB.txt b/.vscode/ltex.dictionary.en-GB.txt index a9db879df..1caf72051 100644 --- a/.vscode/ltex.dictionary.en-GB.txt +++ b/.vscode/ltex.dictionary.en-GB.txt @@ -52,3 +52,4 @@ Shenoy EuXFEL dtype DOOCS +subclassing diff --git a/CHANGELOG.md b/CHANGELOG.md index fcfb1d474..7208e9cdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,7 @@ ### 🐆 Other - Non-deterministic torch features now raise a warning in tests instead of a failure. This also fixes a test failure introduced by changes on the MPS backend in PyTorch 2.13. (see #655) (@Hespe, @jank324) +- Add example notebook on implementing custom elements to documentation (see #522) (@jank324) ### 🌟 First Time Contributors diff --git a/cheetah/accelerator/aperture.py b/cheetah/accelerator/aperture.py index 33999e920..0586598c8 100644 --- a/cheetah/accelerator/aperture.py +++ b/cheetah/accelerator/aperture.py @@ -32,6 +32,8 @@ class Aperture(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/bpm.py b/cheetah/accelerator/bpm.py index 14441a9e9..f40d02375 100644 --- a/cheetah/accelerator/bpm.py +++ b/cheetah/accelerator/bpm.py @@ -25,6 +25,8 @@ class BPM(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/cavity.py b/cheetah/accelerator/cavity.py index 48b2ec0a7..3229dd2a5 100644 --- a/cheetah/accelerator/cavity.py +++ b/cheetah/accelerator/cavity.py @@ -44,6 +44,8 @@ class Cavity(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/combined_corrector.py b/cheetah/accelerator/combined_corrector.py index 82d296839..944c47f88 100644 --- a/cheetah/accelerator/combined_corrector.py +++ b/cheetah/accelerator/combined_corrector.py @@ -32,6 +32,8 @@ class CombinedCorrector(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear"] diff --git a/cheetah/accelerator/custom_transfer_map.py b/cheetah/accelerator/custom_transfer_map.py index 7931d3be7..9122c031b 100644 --- a/cheetah/accelerator/custom_transfer_map.py +++ b/cheetah/accelerator/custom_transfer_map.py @@ -24,6 +24,8 @@ class CustomTransferMap(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear"] diff --git a/cheetah/accelerator/dipole.py b/cheetah/accelerator/dipole.py index 85d9941ee..9f549661f 100644 --- a/cheetah/accelerator/dipole.py +++ b/cheetah/accelerator/dipole.py @@ -49,6 +49,8 @@ class Dipole(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear", "second_order", "drift_kick_drift"] diff --git a/cheetah/accelerator/drift.py b/cheetah/accelerator/drift.py index ff8f54629..1cf1aa847 100644 --- a/cheetah/accelerator/drift.py +++ b/cheetah/accelerator/drift.py @@ -26,6 +26,8 @@ class Drift(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear", "second_order", "drift_kick_drift"] diff --git a/cheetah/accelerator/horizontal_corrector.py b/cheetah/accelerator/horizontal_corrector.py index 7b3789a64..82a856818 100644 --- a/cheetah/accelerator/horizontal_corrector.py +++ b/cheetah/accelerator/horizontal_corrector.py @@ -30,6 +30,8 @@ class HorizontalCorrector(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear"] diff --git a/cheetah/accelerator/marker.py b/cheetah/accelerator/marker.py index 15c9d86b1..6a71ea347 100644 --- a/cheetah/accelerator/marker.py +++ b/cheetah/accelerator/marker.py @@ -21,6 +21,8 @@ class Marker(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/quadrupole.py b/cheetah/accelerator/quadrupole.py index 4e1fbf4b5..b895d8d39 100644 --- a/cheetah/accelerator/quadrupole.py +++ b/cheetah/accelerator/quadrupole.py @@ -42,6 +42,8 @@ class Quadrupole(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear", "second_order", "drift_kick_drift"] diff --git a/cheetah/accelerator/rbend.py b/cheetah/accelerator/rbend.py index 0479ce37d..8d72c8e0a 100644 --- a/cheetah/accelerator/rbend.py +++ b/cheetah/accelerator/rbend.py @@ -42,6 +42,8 @@ class RBend(Dipole): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/screen.py b/cheetah/accelerator/screen.py index e4ce3c918..2d5fc8e11 100644 --- a/cheetah/accelerator/screen.py +++ b/cheetah/accelerator/screen.py @@ -56,6 +56,8 @@ class Screen(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/segment.py b/cheetah/accelerator/segment.py index b51b1f94a..99a4a743c 100644 --- a/cheetah/accelerator/segment.py +++ b/cheetah/accelerator/segment.py @@ -34,6 +34,8 @@ class Segment(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index 164b477a5..013cd8141 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -34,6 +34,8 @@ class Sextupole(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear", "second_order"] diff --git a/cheetah/accelerator/solenoid.py b/cheetah/accelerator/solenoid.py index 07ab98230..200fe7d0b 100644 --- a/cheetah/accelerator/solenoid.py +++ b/cheetah/accelerator/solenoid.py @@ -34,6 +34,8 @@ class Solenoid(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear"] diff --git a/cheetah/accelerator/space_charge_kick.py b/cheetah/accelerator/space_charge_kick.py index 6dcc209ac..e1c837c29 100644 --- a/cheetah/accelerator/space_charge_kick.py +++ b/cheetah/accelerator/space_charge_kick.py @@ -47,6 +47,8 @@ class SpaceChargeKick(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/superimposed.py b/cheetah/accelerator/superimposed.py index 655bc2a33..985c19455 100644 --- a/cheetah/accelerator/superimposed.py +++ b/cheetah/accelerator/superimposed.py @@ -32,6 +32,8 @@ class Superimposed(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ def __init__( diff --git a/cheetah/accelerator/transverse_deflecting_cavity.py b/cheetah/accelerator/transverse_deflecting_cavity.py index 8d0f47ea6..6cb886926 100644 --- a/cheetah/accelerator/transverse_deflecting_cavity.py +++ b/cheetah/accelerator/transverse_deflecting_cavity.py @@ -37,6 +37,8 @@ class TransverseDeflectingCavity(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["drift_kick_drift"] diff --git a/cheetah/accelerator/undulator.py b/cheetah/accelerator/undulator.py index adb8614d2..bee10b845 100644 --- a/cheetah/accelerator/undulator.py +++ b/cheetah/accelerator/undulator.py @@ -31,6 +31,8 @@ class Undulator(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear"] diff --git a/cheetah/accelerator/vertical_corrector.py b/cheetah/accelerator/vertical_corrector.py index cbbf3364e..dad0ac137 100644 --- a/cheetah/accelerator/vertical_corrector.py +++ b/cheetah/accelerator/vertical_corrector.py @@ -30,6 +30,8 @@ class VerticalCorrector(Element): element (e.g. control-system addresses or PVs). This information is *not* used in simulation and may contain any extra data the user wants to store along with the lattice. See :doc:`/examples/including_metadata` for more information. + :param device: Device on which to create the element's tensors. + :param dtype: Data type of the element's tensors. """ supported_tracking_methods = ["linear"] diff --git a/docs/examples/custom_elements.ipynb b/docs/examples/custom_elements.ipynb new file mode 100644 index 000000000..5a142fb64 --- /dev/null +++ b/docs/examples/custom_elements.ipynb @@ -0,0 +1,311 @@ +{ + "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 +} diff --git a/docs/index.rst b/docs/index.rst index a6ffbfd65..f86ece9fd 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -39,6 +39,7 @@ We provide some examples to demonstrate some features of *Cheetah* and show how examples/optimize_speed examples/gradientbased examples/including_metadata + examples/custom_elements Getting Started