Merge pull request #284 from Hespe/279-beam-properties

Improve `Beam` docstrings and type hints
desy-ml · Nov 20, 2024 · de3c956 · de3c956
2 parents 01dbf8e + a56ca57
commit de3c956
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 26 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,8 +6,8 @@ This is a major release with significant upgrades under the hood of Cheetah. Des
 
 ### 🚨 Breaking Changes
 
-- Cheetah is now vectorised. This means that you can run multiple simulations in parallel by passing a batch of beams and settings, resulting a number of interfaces being changed. For Cheetah developers this means that you now have to account for an arbitrary-dimensional tensor of most of the properties of you element, rather than a single value, vector or whatever else a property was before. (see #116, #157, #170, #172, #173, #198, #208, #213, #215, #218, #229, #233, #258, #265) (@jank324, @cr-xu, @hespe, @roussel-ryan)
-- The fifth particle coordinate `s` is renamed to `tau`. Now Cheetah uses the canonical variables in phase space $(x,px=\frac{P_x}{p_0},y,py, \tau=c\Delta t, \delta=\Delta E/{p_0 c})$. In addition, the trailing "s" was removed from some beam property names (e.g. `beam.xs` becomes `beam.x`). (see #163) (@cr-xu)
+- Cheetah is now vectorised. This means that you can run multiple simulations in parallel by passing a batch of beams and settings, resulting a number of interfaces being changed. For Cheetah developers this means that you now have to account for an arbitrary-dimensional tensor of most of the properties of you element, rather than a single value, vector or whatever else a property was before. (see #116, #157, #170, #172, #173, #198, #208, #213, #215, #218, #229, #233, #258, #265, #284) (@jank324, @cr-xu, @hespe, @roussel-ryan)
+- The fifth particle coordinate `s` is renamed to `tau`. Now Cheetah uses the canonical variables in phase space $(x,px=\frac{P_x}{p_0},y,py, \tau=c\Delta t, \delta=\Delta E/{p_0 c})$. In addition, the trailing "s" was removed from some beam property names (e.g. `beam.xs` becomes `beam.x`). (see #163, #284) (@cr-xu, @hespe)
 - `Screen` no longer blocks the beam (by default). To return to old behaviour, set `Screen.is_blocking = True`. (see #208) (@jank324, @roussel-ryan)
 
 ### 🚀 Features
@@ -51,6 +51,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des
 - Add CI runs for macOS (arm64) and Windows (see #226) (@cr-xu, @jank324, @hespe)
 - Clean up CI pipelines (see #243, #244) (@jank324)
 - Fix logo display in README (see #252) (@jank324)
+- Made `Beam` an abstract class (see #284) (@hespe)
 
 ## [v0.6.3](https://github.com/desy-ml/cheetah/releases/tag/v0.6.3) (2024-03-28)
 

diff --git a/cheetah/particles/beam.py b/cheetah/particles/beam.py
@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from typing import Optional
 
 import torch
@@ -7,7 +8,7 @@
 electron_mass_eV = physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
 
 
-class Beam(nn.Module):
+class Beam(ABC, nn.Module):
     r"""
     Parent class to represent a beam of particles. You should not instantiate this
     class directly, but use one of the subclasses.
@@ -35,6 +36,7 @@ class directly, but use one of the subclasses.
     empty = "I'm an empty beam!"
 
     @classmethod
+    @abstractmethod
     def from_parameters(
         cls,
         mu_x: Optional[torch.Tensor] = None,
@@ -52,6 +54,8 @@ def from_parameters(
         cor_tau: Optional[torch.Tensor] = None,
         energy: Optional[torch.Tensor] = None,
         total_charge: Optional[torch.Tensor] = None,
+        device=None,
+        dtype=torch.float32,
     ) -> "Beam":
         """
         Create beam that with given beam parameters.
@@ -75,10 +79,14 @@ def from_parameters(
         :param cor_tau: Correlation between tau and p.
         :param energy: Reference energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
+        :param device: Device to create the beam on. If set to `"auto"` a CUDA GPU is
+            selected if available. The CPU is used otherwise.
+        :param dtype: Data type of the beam.
         """
         raise NotImplementedError
 
     @classmethod
+    @abstractmethod
     def from_twiss(
         cls,
         beta_x: Optional[torch.Tensor] = None,
@@ -87,9 +95,9 @@ def from_twiss(
         beta_y: Optional[torch.Tensor] = None,
         alpha_y: Optional[torch.Tensor] = None,
         emittance_y: Optional[torch.Tensor] = None,
-        sigma_s: Optional[torch.Tensor] = None,
+        sigma_tau: Optional[torch.Tensor] = None,
         sigma_p: Optional[torch.Tensor] = None,
-        cor_s: Optional[torch.Tensor] = None,
+        cor_tau: Optional[torch.Tensor] = None,
         energy: Optional[torch.Tensor] = None,
         total_charge: Optional[torch.Tensor] = None,
         device=None,
@@ -104,24 +112,29 @@ def from_twiss(
         :param beta_y: Beta function in y direction in meters.
         :param alpha_y: Alpha function in y direction in rad.
         :param emittance_y: Emittance in y direction in m*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_s: Correlation of the particle distribution in s direction.
+        :param sigma_tau: Sigma of the particle distribution in longitudinal direction,
+            in meters.
+        :param sigma_p: Sigma of the particle distribution in p direction,
+            dimensionless.
+        :param cor_tau: Correlation between tau and p.
         :param energy: Energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
-        :param device: Device to create the beam on.
+        :param device: Device to create the beam on. If set to `"auto"` a CUDA GPU is
+            selected if available. The CPU is used otherwise.
         :param dtype: Data type of the beam.
         """
         raise NotImplementedError
 
     @classmethod
+    @abstractmethod
     def from_ocelot(cls, parray) -> "Beam":
         """
         Convert an Ocelot ParticleArray `parray` to a Cheetah Beam.
         """
         raise NotImplementedError
 
     @classmethod
+    @abstractmethod
     def from_astra(cls, path: str, **kwargs) -> "Beam":
         """Load an Astra particle distribution as a Cheetah Beam."""
         raise NotImplementedError
@@ -140,6 +153,8 @@ def transformed_to(
         sigma_p: Optional[torch.Tensor] = None,
         energy: Optional[torch.Tensor] = None,
         total_charge: Optional[torch.Tensor] = None,
+        device=None,
+        dtype=torch.float32,
     ) -> "Beam":
         """
         Create version of this beam that is transformed to new beam parameters.
@@ -160,6 +175,9 @@ def transformed_to(
             dimensionless.
         :param energy: Reference energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
+        :param device: Device to create the transformed beam on. If set to `"auto"` a
+            CUDA GPU is selected if available. The CPU is used otherwise.
+        :param dtype: Data type of the transformed beam.
         """
         # Figure out vector dimensions of the original beam and check that passed
         # arguments have the same vector dimensions.
@@ -213,6 +231,8 @@ def transformed_to(
             sigma_p=sigma_p,
             energy=energy,
             total_charge=total_charge,
+            device=device,
+            dtype=dtype,
         )
 
     @property
@@ -232,50 +252,62 @@ def parameters(self) -> dict:
         }
 
     @property
+    @abstractmethod
     def mu_x(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def sigma_x(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def mu_px(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def sigma_px(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def mu_y(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def sigma_y(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def mu_py(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def sigma_py(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
-    def mu_s(self) -> torch.Tensor:
+    @abstractmethod
+    def mu_tau(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def sigma_tau(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def mu_p(self) -> torch.Tensor:
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def sigma_p(self) -> torch.Tensor:
         raise NotImplementedError
 
@@ -299,11 +331,13 @@ def p0c(self) -> torch.Tensor:
         return self.relativistic_beta * self.relativistic_gamma * electron_mass_eV
 
     @property
+    @abstractmethod
     def sigma_xpx(self) -> torch.Tensor:
         # The covariance of (x,px) ~ $\sigma_{xpx}$
         raise NotImplementedError
 
     @property
+    @abstractmethod
     def sigma_ypy(self) -> torch.Tensor:
         raise NotImplementedError
 

diff --git a/cheetah/particles/parameter_beam.py b/cheetah/particles/parameter_beam.py
@@ -307,6 +307,9 @@ def transformed_to(
         :param sigma_p: Sigma of the particle distribution in p, dimensionless.
         :param energy: Reference energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
+        :param device: Device to create the transformed beam on. If set to `"auto"` a
+            CUDA GPU is selected if available. The CPU is used otherwise.
+        :param dtype: Data type of the transformed beam.
         """
         device = device if device is not None else self.mu_x.device
         dtype = dtype if dtype is not None else self.mu_x.dtype

diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py
@@ -24,6 +24,7 @@ class ParticleBeam(Beam):
     :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
         CUDA GPU is selected if available. The CPU is used otherwise.
+    :param dtype: Data type of the generated particles.
     """
 
     def __init__(
@@ -55,7 +56,7 @@ def __init__(
     @classmethod
     def from_parameters(
         cls,
-        num_particles: Optional[torch.Tensor] = None,
+        num_particles: int = 100_000,
         mu_x: Optional[torch.Tensor] = None,
         mu_y: Optional[torch.Tensor] = None,
         mu_px: Optional[torch.Tensor] = None,
@@ -99,12 +100,10 @@ def from_parameters(
         :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
             CUDA GPU is selected if available. The CPU is used otherwise.
+        :param dtype: Data type of the generated particles.
         """
 
         # Set default values without function call in function signature
-        num_particles = (
-            num_particles if num_particles is not None else torch.tensor(100_000)
-        )
         mu_x = mu_x if mu_x is not None else torch.tensor(0.0)
         mu_px = mu_px if mu_px is not None else torch.tensor(0.0)
         mu_y = mu_y if mu_y is not None else torch.tensor(0.0)
@@ -188,7 +187,7 @@ def from_parameters(
     @classmethod
     def from_twiss(
         cls,
-        num_particles: Optional[torch.Tensor] = None,
+        num_particles: int = 1_000_000,
         beta_x: Optional[torch.Tensor] = None,
         alpha_x: Optional[torch.Tensor] = None,
         emittance_x: Optional[torch.Tensor] = None,
@@ -228,9 +227,6 @@ def from_twiss(
             ), "Arguments must have the same shape."
 
         # Set default values without function call in function signature
-        num_particles = (
-            num_particles if num_particles is not None else torch.tensor(1_000_000)
-        )
         beta_x = beta_x if beta_x is not None else torch.full(shape, 0.0)
         alpha_x = alpha_x if alpha_x is not None else torch.full(shape, 0.0)
         emittance_x = emittance_x if emittance_x is not None else torch.full(shape, 0.0)
@@ -276,7 +272,7 @@ def from_twiss(
     @classmethod
     def uniform_3d_ellipsoid(
         cls,
-        num_particles: Optional[torch.Tensor] = None,
+        num_particles: int = 1_000_000,
         radius_x: Optional[torch.Tensor] = None,
         radius_y: Optional[torch.Tensor] = None,
         radius_tau: Optional[torch.Tensor] = None,
@@ -310,7 +306,8 @@ def uniform_3d_ellipsoid(
         :param sigma_p: Sigma of the particle distribution in p, dimensionless.
         :param energy: Reference 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.
+        :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.
         :param dtype: Data type of the generated particles.
 
         :return: ParticleBeam with uniformly distributed particles inside an ellipsoid.
@@ -343,9 +340,6 @@ def uniform_3d_ellipsoid(
         # Set default values without function call in function signature
         # NOTE that this does not need to be done for values that are passed to the
         # Gaussian beam generation.
-        num_particles = (
-            num_particles if num_particles is not None else torch.tensor(1_000_000)
-        )
         radius_x = (
             radius_x.expand(vector_shape)
             if radius_x is not None
@@ -414,7 +408,7 @@ def uniform_3d_ellipsoid(
     @classmethod
     def make_linspaced(
         cls,
-        num_particles: Optional[torch.Tensor] = None,
+        num_particles: int = 10,
         mu_x: Optional[torch.Tensor] = None,
         mu_y: Optional[torch.Tensor] = None,
         mu_px: Optional[torch.Tensor] = None,
@@ -450,10 +444,10 @@ def make_linspaced(
         :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.
+        :param dtype: Data type of the generated particles.
         """
 
         # Set default values without function call in function signature
-        num_particles = num_particles if num_particles is not None else torch.tensor(10)
         mu_x = mu_x if mu_x is not None else torch.tensor(0.0)
         mu_px = mu_px if mu_px is not None else torch.tensor(0.0)
         mu_y = mu_y if mu_y is not None else torch.tensor(0.0)
@@ -564,7 +558,6 @@ def transformed_to(
         """
         Create version of this beam that is transformed to new beam parameters.
 
-        :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_px: Center of the particle distribution on px, dimensionless.
@@ -582,6 +575,7 @@ def transformed_to(
         :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
             CUDA GPU is selected if available. The CPU is used otherwise.
+        :param dtype: Data type of the transformed particles.
         """
         device = device if device is not None else self.mu_x.device
         dtype = dtype if dtype is not None else self.mu_x.dtype