Merge pull request #123 from desy-ml/122-wrong-units-in-docstring-in-…

…multiple-locations

Correct wrong units in docstring in multiple locations

CHANGELOG.md

@@ -27,6 +27,7 @@
 - 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)
 
 ## [v0.6.1](https://github.com/desy-ml/cheetah/releases/tag/v0.6.1) (2023-09-17)
 

cheetah/accelerator.py

@@ -23,8 +23,10 @@
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
 rest_energy = torch.tensor(
-    constants.electron_mass * constants.speed_of_light**2 / constants.elementary_charge
-)  # electron mass
+    constants.electron_mass
+    * constants.speed_of_light**2
+    / constants.elementary_charge  # electron mass
+)
 electron_mass_eV = torch.tensor(
     physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 )
@@ -43,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

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.
         """
@@ -253,6 +254,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
@@ -272,6 +274,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 __repr__(self) -> str:
@@ -504,15 +507,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.
         """
@@ -680,16 +683,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.
@@ -838,14 +841,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.
@@ -951,14 +954,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