Vectorised simulations

CHANGELOG.md

@@ -4,11 +4,12 @@
 
 ### 🚨 Breaking Changes
 
-- Cheetah is not 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. (see #116) (@jank324)
+- 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. (see #116) (@jank324)
 
 ### 🚀 Features
 
 - `CustomTransferMap` elements created by combining multiple other elements will now reflect that in their `name` attribute (see #100) (@jank324)
+- Now all `Element` have a default length of `torch.zeros((1))`, i.e. also for `Marker`, `BPM`, `Screen`, and `Aperture`.
 
 ### 🐛 Bug fixes
 

cheetah/accelerator.py

@@ -39,6 +39,8 @@ class Element(ABC, nn.Module):
     :param name: Unique identifier of the element.
     """
 
+    length: torch.Tensor = torch.zeros((1))
+
     def __init__(self, name: Optional[str] = None) -> None:
         super().__init__()
 
@@ -216,9 +218,7 @@ def from_merging_elements(
             tm = torch.matmul(element.transfer_map(incoming_beam.energy), tm)
             incoming_beam = element.track(incoming_beam)
 
-        combined_length = sum(
-            element.length for element in elements if hasattr(element, "length")
-        )
+        combined_length = sum(element.length for element in elements)
 
         combined_name = "combined_" + "_".join(element.name for element in elements)
 
@@ -1275,7 +1275,9 @@ def track(self, incoming: Beam) -> Beam:
         return deepcopy(incoming)
 
     def broadcast(self, shape: Size) -> Element:
-        return self.__class__(is_active=self.is_active, name=self.name)
+        new_bpm = self.__class__(is_active=self.is_active, name=self.name)
+        new_bpm.length = self.length.repeat(shape)
+        return new_bpm
 
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
@@ -1316,7 +1318,9 @@ def track(self, incoming: Beam) -> Beam:
         return incoming
 
     def broadcast(self, shape: Size) -> Element:
-        return self.__class__(name=self.name)
+        new_marker = self.__class__(name=self.name)
+        new_marker.length = self.length.repeat(shape)
+        return new_marker
 
     @property
     def is_skippable(self) -> bool:
@@ -1546,14 +1550,16 @@ def set_read_beam(self, value: Beam) -> None:
         self.cached_reading = None
 
     def broadcast(self, shape: Size) -> Element:
-        return self.__class__(
+        new_screen = self.__class__(
             resolution=self.resolution,
             pixel_size=self.pixel_size,
             binning=self.binning,
             misalignment=self.misalignment.repeat((*shape, 1)),
             is_active=self.is_active,
             name=self.name,
         )
+        new_screen.length = self.length.repeat(shape)
+        return new_screen
 
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
@@ -1678,13 +1684,15 @@ def track(self, incoming: Beam) -> Beam:
         )
 
     def broadcast(self, shape: Size) -> Element:
-        return self.__class__(
+        new_aperture = self.__class__(
             x_max=self.x_max.repeat(shape),
             y_max=self.y_max.repeat(shape),
             shape=self.shape,
             is_active=self.is_active,
             name=self.name,
         )
+        new_aperture.length = self.length.repeat(shape)
+        return new_aperture
 
     def split(self, resolution: torch.Tensor) -> list[Element]:
         # TODO: Implement splitting for aperture properly, for now just return self
@@ -2075,7 +2083,7 @@ def without_inactive_zero_length_elements(
             elements=[
                 element
                 for element in self.elements
-                if (hasattr(element, "length") and element.length > 0.0)
+                if element.length > 0.0
                 or (hasattr(element, "is_active") and element.is_active)
                 or element.name in except_for
             ],
@@ -2104,7 +2112,7 @@ def inactive_elements_as_drifts(
                 (
                     element
                     if (hasattr(element, "is_active") and element.is_active)
-                    or not hasattr(element, "length")
+                    or element.length == 0.0
                     or element.name in except_for
                     else Drift(element.length)
                 )
@@ -2217,7 +2225,7 @@ def is_skippable(self) -> bool:
     @property
     def length(self) -> torch.Tensor:
         lengths = torch.stack(
-            [element.length for element in self.elements if hasattr(element, "length")],
+            [element.length for element in self.elements],
             dim=1,
         )
         return torch.sum(lengths, dim=1)
@@ -2265,10 +2273,7 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
         ]
 
     def plot(self, ax: matplotlib.axes.Axes, s: float) -> None:
-        element_lengths = [
-            element.length[0] if hasattr(element, "length") else 0.0
-            for element in self.elements
-        ]
+        element_lengths = [element.length[0] for element in self.elements]
         element_ss = [0] + [
             sum(element_lengths[: i + 1]) for i, _ in enumerate(element_lengths)
         ]
@@ -2305,9 +2310,7 @@ def plot_reference_particle_traces(
         reference_segment = deepcopy(self)
         splits = reference_segment.split(resolution=torch.tensor(resolution))
 
-        split_lengths = [
-            split.length[0] if hasattr(split, "length") else 0.0 for split in splits
-        ]
+        split_lengths = [split.length[0] for split in splits]
         ss = [0] + [sum(split_lengths[: i + 1]) for i, _ in enumerate(split_lengths)]
 
         references = []
@@ -2393,7 +2396,7 @@ def plot_twiss(self, beam: Beam, ax: Optional[Any] = None) -> None:
         longitudinal_beams = [beam]
         s_positions = [0.0]
         for element in self.elements:
-            if not hasattr(element, "length") or element.length == 0:
+            if element.length == 0:
                 continue
 
             outgoing = element.track(longitudinal_beams[-1])

tests/test_tracking_lengthless_elements.py

@@ -0,0 +1,31 @@
+import torch
+
+import cheetah
+
+beam_in = cheetah.ParticleBeam.from_parameters(num_particles=100)
+
+
+# Only Marker
+def test_tracking_marker_only():
+    segment = cheetah.Segment([cheetah.Marker(name="start")])
+
+    beam_out = segment.track(beam_in)
+
+    assert torch.allclose(beam_out.particles, beam_in.particles)
+
+
+# Only length-less elements between non-skippable elements
+def test_tracking_lengthless_elements():
+    segment = cheetah.Segment(
+        [
+            cheetah.Cavity(
+                length=torch.tensor([0.1]), voltage=torch.tensor([1e6]), name="C2"
+            ),
+            cheetah.Marker(name="start"),
+            cheetah.Cavity(
+                length=torch.tensor([0.1]), voltage=torch.tensor([1e6]), name="C1"
+            ),
+        ]
+    )
+
+    _ = segment.track(beam_in)