Merge branch '100-batched-execution' of github.com:desy-ml/cheetah into 100-batched-execution

Author: jank324

.flake8

@@ -7,4 +7,4 @@ exclude =
     .ipynb_checkpoints
 max-line-length = 88
 extend-ignore = E203, W503
-ignore = BLK100
+ignore = BLK100

CHANGELOG.md

@@ -1,6 +1,16 @@
 # Changelog
 
-## v0.6.2 [🚧 Work in Progress]
+## v0.7.0 [🚧 Work in Progress]
+
+### 🚨 Breaking Changes
+
+### 🚀 Features
+
+### 🐛 Bug fixes
+
+### 🐆 Other
+
+## [v0.6.2](https://github.com/desy-ml/cheetah/releases/tag/v0.6.2) (2024-02-13)
 
 ### 🚨 Breaking Changes
 
@@ -19,12 +29,16 @@
 - Fix error raised when tracking a `ParameterBeam` through an active `BPM` (see #101) (@jank324)
 - Fix error in ASTRA beam import where the energy was set to `float64` instead of `float32` (see #111) (@jank324)
 - Fix missing passing of `total_charge` in `ParameterBeam.transformed_to` (see #112) (@jank324)
+- Fix `Cavitiy.__repr__` printing `voltage` value for `phase` property (see #121) (@jank324)
 
 ### 🐆 Other
 
 - Add pull request template (see #97) (@jank324)
 - Add _Acknowledgements_ section to README (see #103) (@jank324)
 - `benchmark` directory was moved to `desy-ml/cheetah-demos` repository (see #108) (@jank324)
+- Update citations to new arXiv preprint (see #117) (@jank324)
+- Improve the docstring with proper units for the phase space dimensions (see #122) (@cr-xu)
+- Link to the new paper (on arXiv) in the documentation (see #125) (@jank324)
 
 ## [v0.6.1](https://github.com/desy-ml/cheetah/releases/tag/v0.6.1) (2023-09-17)
 

CITATION.bib

@@ -0,0 +1,14 @@
+@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}
+}

README.md

@@ -85,15 +85,22 @@ where the optional keyword argument `beam` is the incoming beam represented by t
 
 ## Cite Cheetah
 
-To cite Cheetah in publications:
+If you use Cheetah, please cite the following two papers:
 
 ```bibtex
+@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},
-    url          = {https://github.com/desy-ml/cheetah}
+  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}
 }
 ```
 

cheetah/accelerator.py

@@ -25,8 +25,8 @@
 rest_energy = torch.tensor(
     constants.electron_mass
     * constants.speed_of_light**2
-    / constants.elementary_charge
-)  # electron mass
+    / constants.elementary_charge  # electron mass
+)
 electron_mass_eV = torch.tensor(
     physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 )
@@ -45,21 +45,25 @@ def __init__(self, name: Optional[str] = None) -> None:
         self.name = name if name is not None else generate_unique_name()
 
     def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
-        """
+        r"""
         Generates the element's transfer map that describes how the beam and its
-        particles are transformed when traveling through the element. The state vector
-        consists of 6 values with a physical meaning:
+        particles are transformed when traveling through the element.
+        The state vector consists of 6 values with a physical meaning:
+        (in the trace space notation)
+
         - x: Position in x direction
-        - xp: Momentum in x direction
+        - xp: Angle in x direction
         - y: Position in y direction
-        - yp: Momentum in y direction
-        - s: Position in z direction, the zero value is set to the middle of the pulse
-        - sp: Momentum in s direction
+        - yp: Angle in y direction
+        - s: Position in longitudinal direction, the zero value is set to the
+        reference position (usually the center of the pulse)
+        - p: Relative energy deviation from the reference particle
+           :math:`p = \frac{\Delta E}{p_0 C}`
         As well as a seventh value used to add constants to some of the prior values if
         necessary. Through this seventh state, the addition of constants can be
         represented using a matrix multiplication.
 
-        :param energy: Energy of the Beam. Read from the fed-in Cheetah Beam.
+        :param energy: Reference energy of the Beam. Read from the fed-in Cheetah Beam.
         :return: A 7x7 Matrix for further calculations.
         """
         raise NotImplementedError
@@ -1237,7 +1241,7 @@ def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}(length={repr(self.length)}, "
             + f"voltage={repr(self.voltage)}, "
-            + f"phase={repr(self.voltage)}, "
+            + f"phase={repr(self.phase)}, "
             + f"frequency={repr(self.frequency)}, "
             + f"name={repr(self.name)})"
         )

cheetah/converters/dontbmad.py

@@ -426,7 +426,7 @@ def validate_understood_properties(understood: list[str], properties: dict) -> N
         assert property in understood, (
             f"Property {property} with value {properties[property]} for element type"
             f" {properties['element_type']} is currently not understood. Other values"
-            f" in properties are {properties.keys()}."
+            f" in properties are {properties.keys()}."  # noqa: B038
         )
 
 

cheetah/particles.py

@@ -69,11 +69,11 @@ def from_twiss(
         Create a beam from twiss parameters.
 
         :param beta_x: Beta function in x direction in meters.
-        :param alpha_x: Alpha function in x direction in meters.
-        :param emittance_x: Emittance in x direction.
+        :param alpha_x: Alpha function in x direction in rad.
+        :param emittance_x: Emittance in x direction in m*rad.
         :param beta_y: Beta function in y direction in meters.
-        :param alpha_y: Alpha function in y direction in meters.
-        :param emittance_y: Emittance in y direction.
+        :param alpha_y: Alpha function in y direction in rad.
+        :param emittance_y: Emittance in y direction in m*rad.
         :param energy: Energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
         """
@@ -110,15 +110,16 @@ def transformed_to(
         Create version of this beam that is transformed to new beam parameters.
 
         :param mu_x: Center of the particle distribution on x in meters.
-        :param mu_xp: Center of the particle distribution on px in meters.
+        :param mu_xp: Center of the particle distribution on x' in rad.
         :param mu_y: Center of the particle distribution on y in meters.
-        :param mu_yp: Center of the particle distribution on py in meters.
+        :param mu_yp: Center of the particle distribution on y' in rad.
         :param sigma_x: Sigma of the particle distribution in x direction in meters.
-        :param sigma_xp: Sigma of the particle distribution in px direction in meters.
+        :param sigma_xp: Sigma of the particle distribution in x' direction in rad.
         :param sigma_y: Sigma of the particle distribution in y direction in meters.
-        :param sigma_yp: Sigma of the particle distribution in py direction in meters.
+        :param sigma_yp: Sigma of the particle distribution in y' direction in rad.
         :param sigma_s: Sigma of the particle distribution in s direction in meters.
-        :param sigma_p: Sigma of the particle distribution in p direction in meters.
+        :param sigma_p: Sigma of the particle distribution in p direction,
+        dimensionless.
         :param energy: Energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
         """
@@ -283,6 +284,7 @@ def beta_x(self) -> torch.Tensor:
 
     @property
     def alpha_x(self) -> torch.Tensor:
+        """Alpha function in x direction in rad."""
         return -self.sigma_xxp / self.emittance_x
 
     @property
@@ -307,6 +309,7 @@ def beta_y(self) -> torch.Tensor:
 
     @property
     def alpha_y(self) -> torch.Tensor:
+        """Alpha function in y direction in rad."""
         return -self.sigma_yyp / self.emittance_y
 
     def broadcast(self, shape: torch.Size) -> "Beam":
@@ -611,15 +614,15 @@ def transformed_to(
 
         :param n: Number of particles to generate.
         :param mu_x: Center of the particle distribution on x in meters.
-        :param mu_xp: Center of the particle distribution on px in meters.
+        :param mu_xp: Center of the particle distribution on x' in rad.
         :param mu_y: Center of the particle distribution on y in meters.
-        :param mu_yp: Center of the particle distribution on py in meters.
+        :param mu_yp: Center of the particle distribution on y' in rad.
         :param sigma_x: Sigma of the particle distribution in x direction in meters.
-        :param sigma_xp: Sigma of the particle distribution in px direction in meters.
+        :param sigma_xp: Sigma of the particle distribution in x' direction in rad.
         :param sigma_y: Sigma of the particle distribution in y direction in meters.
-        :param sigma_yp: Sigma of the particle distribution in py direction in meters.
+        :param sigma_yp: Sigma of the particle distribution in y' direction in rad.
         :param sigma_s: Sigma of the particle distribution in s direction in meters.
-        :param sigma_p: Sigma of the particle distribution in p direction in meters.
+        :param sigma_p: Sigma of the particle distribution in p, dimensionless.
         :param energy: Energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
         """
@@ -820,16 +823,16 @@ def from_parameters(
         :param num_particles: Number of particles to generate.
         :param mu_x: Center of the particle distribution on x in meters.
         :param mu_y: Center of the particle distribution on y in meters.
-        :param mu_xp: Center of the particle distribution on px in meters.
-        :param mu_yp: Center of the particle distribution on py in meters.
+        :param mu_xp: Center of the particle distribution on x' in rad.
+        :param mu_yp: Center of the particle distribution on y' in metraders.
         :param sigma_x: Sigma of the particle distribution in x direction in meters.
         :param sigma_y: Sigma of the particle distribution in y direction in meters.
-        :param sigma_xp: Sigma of the particle distribution in px direction in meters.
-        :param sigma_yp: Sigma of the particle distribution in py direction in meters.
+        :param sigma_xp: Sigma of the particle distribution in x' direction in rad.
+        :param sigma_yp: Sigma of the particle distribution in y' direction in rad.
         :param sigma_s: Sigma of the particle distribution in s direction in meters.
-        :param sigma_p: Sigma of the particle distribution in p direction in meters.
-        :param cor_x: Correlation between x and xp.
-        :param cor_y: Correlation between y and yp.
+        :param sigma_p: Sigma of the particle distribution in p, dimensionless.
+        :param cor_x: Correlation between x and x'.
+        :param cor_y: Correlation between y and y'.
         :param cor_s: Correlation between s and p.
         :param energy: Energy of the beam in eV.
         :total_charge: Total charge of the beam in C.
@@ -1041,14 +1044,14 @@ def make_linspaced(
         :param n: Number of particles to generate.
         :param mu_x: Center of the particle distribution on x in meters.
         :param mu_y: Center of the particle distribution on y in meters.
-        :param mu_xp: Center of the particle distribution on px in meters.
-        :param mu_yp: Center of the particle distribution on py in meters.
+        :param mu_xp: Center of the particle distribution on x' in rad.
+        :param mu_yp: Center of the particle distribution on y' in rad.
         :param sigma_x: Sigma of the particle distribution in x direction in meters.
         :param sigma_y: Sigma of the particle distribution in y direction in meters.
-        :param sigma_xp: Sigma of the particle distribution in px direction in meters.
-        :param sigma_yp: Sigma of the particle distribution in py direction in meters.
+        :param sigma_xp: Sigma of the particle distribution in x' direction in rad.
+        :param sigma_yp: Sigma of the particle distribution in y' direction in rad.
         :param sigma_s: Sigma of the particle distribution in s direction in meters.
-        :param sigma_p: Sigma of the particle distribution in p direction in meters.
+        :param sigma_p: Sigma of the particle distribution in p, dimensionless.
         :param energy: Energy of the beam in eV.
         :param device: Device to move the beam's particle array to. If set to `"auto"` a
             CUDA GPU is selected if available. The CPU is used otherwise.
@@ -1232,14 +1235,14 @@ def transformed_to(
         :param n: Number of particles to generate.
         :param mu_x: Center of the particle distribution on x in meters.
         :param mu_y: Center of the particle distribution on y in meters.
-        :param mu_xp: Center of the particle distribution on px in meters.
-        :param mu_yp: Center of the particle distribution on py in meters.
+        :param mu_xp: Center of the particle distribution on x' in rad.
+        :param mu_yp: Center of the particle distribution on y' in rad.
         :param sigma_x: Sigma of the particle distribution in x direction in meters.
         :param sigma_y: Sigma of the particle distribution in y direction in meters.
-        :param sigma_xp: Sigma of the particle distribution in px direction in meters.
-        :param sigma_yp: Sigma of the particle distribution in py direction in meters.
+        :param sigma_xp: Sigma of the particle distribution in x' direction in rad.
+        :param sigma_yp: Sigma of the particle distribution in y' direction in rad.
         :param sigma_s: Sigma of the particle distribution in s direction in meters.
-        :param sigma_p: Sigma of the particle distribution in p direction in meters.
+        :param sigma_p: Sigma of the particle distribution in p, dimensionless.
         :param energy: Energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
         :param device: Device to move the beam's particle array to. If set to `"auto"` a

cheetah/track_methods.py

@@ -6,9 +6,7 @@
 from scipy import constants
 
 REST_ENERGY = torch.tensor(
-    constants.electron_mass
-    * constants.speed_of_light**2
-    / constants.elementary_charge
+    constants.electron_mass * constants.speed_of_light**2 / constants.elementary_charge
 )  # Electron mass
 
 

docs/index.rst

@@ -10,7 +10,7 @@ Welcome to Cheetah's documentation!
 
 GitHub repository: https://github.com/desy-ml/cheetah
 
-Paper: https://accelconf.web.cern.ch/ipac2022/papers/wepoms036.pdf
+Paper: https://arxiv.org/abs/2401.05815
 
 
 Installation
@@ -61,18 +61,23 @@ For more advanced usage, please refer to the in-depth documentation.
 Cite Cheetah
 ------------
 
-To cite Cheetah in publications:
+If you use Cheetah, please cite the following two papers:
 
 .. code-block:: bibtex
 
+    @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,
-        author = {Stein, Oliver and
-                Kaiser, Jan and
-                Eichler, Annika},
-        title = {Accelerating Linear Beam Dynamics Simulations for Machine Learning Applications},
-        booktitle = {Proceedings of the 13th International Particle Accelerator Conference},
-        year = {2022},
-        url = {https://github.com/desy-ml/cheetah},
+      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

tests/test_compare_beam_type.py

@@ -1,6 +1,7 @@
 """
 Tests that ensure that both beam types produce (roughly) the same results.
 """
+
 import torch
 
 import cheetah