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