Merge branch 'master' into 241-vectorised-aperture-tracking-is-broken

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)
 - Rework the `Aperture` element. Now `ParticleBeam` has a `particle_survival` attribute that keeps track of the lost particles. The statistical beam parameters are calculated only w.r.t. surviving particles. Note that the `Aperture` breaks differentiability if activated. (see #268) (@cr-xu, @jank324)
 
@@ -38,6 +38,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des
 - Fix issue in Bmad import where collimators had no length by interpreting them as `Drift` + `Aperture` (see #249) (@jank324)
 - Fix NumPy 2 compatibility issues with PyTorch on Windows (see #220, #242) (@hespe)
 - Fix issue with Dipole hgap conversion in Bmad import (see #261) (@cr-xu)
+- Fix plotting for segments that contain tensors with `require_grad=True` (see #288) (@hespe)
 
 ### 🐆 Other
 
@@ -52,6 +53,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)
 

cheetah/accelerator/segment.py

@@ -400,7 +400,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> No
             dimension_reordered_ss[vector_idx]
             if stacked_ss.dim() > 1
             else dimension_reordered_ss
-        )
+        ).detach()
 
         ax.plot([0, plot_ss[-1]], [0, 0], "--", color="black")
 
@@ -465,17 +465,17 @@ def plot_reference_particle_traces(
             dimensions_reordered_ss[vector_idx]
             if stacked_ss.dim() > 1
             else dimensions_reordered_ss
-        )
+        ).detach()
         plot_xs = (
             dimension_reordered_xs[vector_idx]
             if stacked_xs.dim() > 2
             else dimension_reordered_xs
-        )
+        ).detach()
         plot_ys = (
             dimension_reordered_ys[vector_idx]
             if stacked_ys.dim() > 2
             else dimension_reordered_ys
-        )
+        ).detach()
 
         for particle_idx in range(num_particles):
             axx.plot(plot_ss, plot_xs[particle_idx])
@@ -558,7 +558,7 @@ def plot_twiss(
             dimension_reordered_s_positions[vector_idx]
             if stacked_s_positions.dim() > 1
             else dimension_reordered_s_positions
-        )
+        ).detach()
 
         broadcast_beta_x = torch.broadcast_tensors(*beta_x)
         stacked_beta_x = torch.stack(broadcast_beta_x)
@@ -567,7 +567,7 @@ def plot_twiss(
             dimension_reordered_beta_x[vector_idx]
             if stacked_beta_x.dim() > 2
             else dimension_reordered_beta_x
-        )
+        ).detach()
 
         broadcast_beta_y = torch.broadcast_tensors(*beta_y)
         stacked_beta_y = torch.stack(broadcast_beta_y)
@@ -576,7 +576,7 @@ def plot_twiss(
             dimension_reordered_beta_y[vector_idx]
             if stacked_beta_y.dim() > 2
             else dimension_reordered_beta_y
-        )
+        ).detach()
 
         if ax is None:
             fig = plt.figure()

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
 

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