Add example notebook on implementing custom elements to documentation #522
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR adds comprehensive documentation for implementing custom elements in the Cheetah accelerator physics library. The documentation consists of a Jupyter notebook that provides detailed guidance on extending Cheetah with user-defined elements.
- Adds a complete tutorial notebook showing how to subclass the
Elementbase class - Provides extensive code examples with TODO comments explaining each implementation step
- Includes guidelines for contributing new elements back to the main repository
Reviewed Changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| docs/index.rst | Adds the new custom elements example to the documentation table of contents |
| docs/examples/custom_elements.ipynb | Complete tutorial notebook with detailed implementation guidance and code templates |
| CHANGELOG.md | Documents the addition of the new example notebook |
| .vscode/ltex.dictionary.en-GB.txt | Adds "subclassing" to the spell-check dictionary |
| " def transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor:\n", | ||
| " # 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\n", | ||
| " # for you.\n", | ||
| " # The transfer map needs to be a tensor of shape (..., 7, 7), where all but the last two dimensions are vector dimensions, the 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.\n", |
There was a problem hiding this comment.
Duplicate word 'the' in the comment. Should be 'where all but the last two dimensions are vector dimensions, the upper left 6x6 block'.
| " def to_mesh(\n", | ||
| " self, cuteness: float | dict = 1.0, show_download_progress: bool = True\n", | ||
| " ) -> \"tuple[trimesh.Trimesh | None, np.ndarray]\": # noqa: F821 # type: ignore\n", | ||
| " # TODO: Cheetah can create 3D mesh representations of elements and lattices. In most cases you don't have to implement this method. If you feel creative, you can open a PR with a mesh file on the `desy-ml/3d-assets` repository with a file name of your class name in snake case, and then your mesh will automatically be used by Cheetah. In some cases, placing the mesh or computing the 3D trasnformation to place its beam output position is more complex. In those cases, you might need to implement this method. (See `Dipole` for an example of such a case.)\n", |
There was a problem hiding this comment.
Typo in the word 'trasnformation'. Should be 'transformation'.
| " # TODO: Cheetah can create 3D mesh representations of elements and lattices. In most cases you don't have to implement this method. If you feel creative, you can open a PR with a mesh file on the `desy-ml/3d-assets` repository with a file name of your class name in snake case, and then your mesh will automatically be used by Cheetah. In some cases, placing the mesh or computing the 3D trasnformation to place its beam output position is more complex. In those cases, you might need to implement this method. (See `Dipole` for an example of such a case.)\n", | |
| " # TODO: Cheetah can create 3D mesh representations of elements and lattices. In most cases you don't have to implement this method. If you feel creative, you can open a PR with a mesh file on the `desy-ml/3d-assets` repository with a file name of your class name in snake case, and then your mesh will automatically be used by Cheetah. In some cases, placing the mesh or computing the 3D transformation to place its beam output position is more complex. In those cases, you might need to implement this method. (See `Dipole` for an example of such a case.)\n", |
| " ]\n", | ||
| "\n", | ||
| " def __repr__(self) -> str:\n", | ||
| " # TODO: Implement a string representation for your element. This is usally the class name, followed by the values of the defining features in such a way that the result could be pasted into a Python interpreter to create an equivalent object.\n", |
There was a problem hiding this comment.
Typo in the word 'usally'. Should be 'usually'.
|
@Hespe, when you review this, can you check the result in the built documentation. Can you also check if this matches the current version of Cheetah? I modelled this after when we did in #406. I'm not sure if we should go ahead and merge this PR when it's ready, or wait for #476, because that will change some things. Anyways, @frankmayet can have a look at this already. |
|
Notebooks on readthedocs don't seem to use the full screen width. Of course you can scroll horizontally, but this makes the docstrings anoying to read: 
|
| " magnet that focuses the beam in the horizontal direction and defocuses it in the vertical direction.\n", | ||
| "\n", | ||
| " :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.\n", | ||
| " :param my_second_parameter: TODO: Specify the parameters of your element here, following the format `:param parameter_name: Description of the parameter.`.\n", |
There was a problem hiding this comment.
my_second_parameter does not exist, you probably renamed it.
| " def __init__(\n", | ||
| " self,\n", | ||
| " length: torch.Tensor,\n", | ||
| " # TODO: Replace `my_second_parameter` with the actual parameters of your element. Note that the first parameter must always be `length`, and the last two parameters must always be `name`, `device` and `dtype`. 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).\n", |
There was a problem hiding this comment.
my_second_parameter again, and it should be the last three parameters, not two.
Also, for some parameters it would be fine to give the default directly. Do you want to keep the remark this general for simplicity without going into the required detail?
| " :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.\n", | ||
| " :param my_second_parameter: TODO: Specify the parameters of your element here, following the format `:param parameter_name: Description of the parameter.`.\n", | ||
| " :param name: Unique identifier of the element.\n", | ||
| " \"\"\"\n", |
There was a problem hiding this comment.
No docstring for device and dtype?
| "\n", | ||
| " def split(self, resolution: torch.Tensor) -> list[Element]:\n", | ||
| " # TODO: Implement your logic for splitting your element longitudinally here. This is used to provide beams at multiple positions along the element, specifically `resolution` meters apart. Not all elements can easily be split, which is why Cheetah does not guarantee to the user that `resolution` will always be respected. If your element cannot be split, you may return a list containing only `self.clone()`.\n", | ||
| " return [self.clone()]\n", |
There was a problem hiding this comment.
We don't typically clone() in the existing implementations of split. Maybe we should?
| " return ParticleBeam(\n", | ||
| " particles=outgoing_particles,\n", | ||
| " energy=incoming.energy,\n", | ||
| " particle_charges=incoming.particle_charges,\n", | ||
| " survival_probabilities=incoming.survival_probabilities,\n", | ||
| " species=incoming.species,\n", | ||
| " )\n", |
There was a problem hiding this comment.
Should we add clone() here to the passed tensors?
| " # TODO: Implement a string representation for your element. This is usally the class name, followed by the values of the defining features in such a way that the result could be pasted into a Python interpreter to create an equivalent object.\n", | ||
| " return (\n", | ||
| " f\"{self.__class__.__name__}(length={repr(self.length)}, \"\n", | ||
| " f\"my_second_parameter={repr(self.my_second_parameter)}, \"\n", |
There was a problem hiding this comment.
| " f\"my_second_parameter={repr(self.my_second_parameter)}, \"\n", | |
| " f\"a_tensor_parameter={repr(self.a_tensor_parameter)}, \"\n", |
| " # 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`.\n", | ||
| " return super().defining_features + [\n", | ||
| " \"length\",\n", | ||
| " \"my_second_parameter\",\n", |
There was a problem hiding this comment.
| " \"my_second_parameter\",\n", | |
| " \"a_tensor_parameter\",\n", |
| " # TODO: Register all your torch.Tensor parameters like the following. If they\n", | ||
| " # 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.\n", | ||
| " self.register_buffer_or_parameter(\n", | ||
| " \"my_second_parameter\",\n", |
There was a problem hiding this comment.
| " \"my_second_parameter\",\n", | |
| " \"a_tensor_parameter\",\n", |
| " :undoc-members:\n", | ||
| " ```\n", | ||
| "- Write a test for your element in a new appropriately named in the `tests` directory. The file name should follow the pattern `test_<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.\n", | ||
| "- Add your new element to the parameterisation of the `test_element_buffer_contents_and_location` test in `tests/test_clone.py`.\n", |
There was a problem hiding this comment.
These tests are parametrized with @pytest.mark.for_every_element, no need to adjust them.
| "- Add your new element to the parameterisation of the `test_element_buffer_contents_and_location` test in `tests/test_clone.py`.\n", | ||
| "- Add your new element to the parametrisations of the `test_forced_element_dtype` and `test_infer_element_dtype` tests in `tests/test_device_dtype.py`.\n", | ||
| "- Add a minimal dictionary of arguments to your element to the `ELEMENT_CLASSES_REQUIRING_ARGS` dictionary in `tests/test_elements.py`.\n", | ||
| "- Add your element to the parametrisation of the `test_drift_broadcasting_two_different_inputs` test in `tests/test_vectorized.py`.\n", |
I would be in favour of merging #476 first. Otherwise, I think its likely that we will not update this example soon enough. If we merge this PR first, #476 should in best case also update the example. |
Seconded! |
Description
Add a notebook that guides people along the implementation of new elements in Cheetah.
Motivation and Context
Closes #396.
Types of changes
Checklist
flake8(required).pytesttests pass (required).pyteston a machine with a CUDA GPU and made sure all tests pass (required).Note: We are using a maximum length of 88 characters per line.