Fix twiss plot

CHANGELOG.md

@@ -6,7 +6,7 @@ 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, #215, #218, #229, #233, #258) (@jank324, @cr-xu, @hespe, @roussel-ryan)
+- 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) (@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)
 - `Screen` no longer blocks the beam (by default). To return to old behaviour, set `Screen.is_blocking = True`. (see #208) (@jank324, @roussel-ryan)
 

cheetah/accelerator/aperture.py

@@ -113,14 +113,16 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
         # TODO: Implement splitting for aperture properly, for now just return self
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
+        plot_s = s[vector_idx] if s.dim() > 0 else s
+
         alpha = 1 if self.is_active else 0.2
         height = 0.4
 
         dummy_length = 0.0
 
         patch = Rectangle(
-            (s, 0), dummy_length, height, color="tab:pink", alpha=alpha, zorder=2
+            (plot_s, 0), dummy_length, height, color="tab:pink", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/bpm.py

@@ -51,10 +51,12 @@ def track(self, incoming: Beam) -> Beam:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
+        plot_s = s[vector_idx] if s.dim() > 0 else s
+
         alpha = 1 if self.is_active else 0.2
         patch = Rectangle(
-            (s, -0.3), 0, 0.3 * 2, color="darkkhaki", alpha=alpha, zorder=2
+            (plot_s, -0.3), 0, 0.3 * 2, color="darkkhaki", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/cavity.py

@@ -341,12 +341,15 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
         # element itself
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
+        plot_s = s[vector_idx] if s.dim() > 0 else s
+        plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
+
         alpha = 1 if self.is_active else 0.2
         height = 0.4
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="gold", alpha=alpha, zorder=2
+            (plot_s, 0), plot_length, height, color="gold", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/custom_transfer_map.py

@@ -102,8 +102,11 @@ def defining_features(self) -> list[str]:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
+        plot_s = s[vector_idx] if s.dim() > 0 else s
+        plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
+
         height = 0.4
 
-        patch = Rectangle((s, 0), self.length[0], height, color="tab:olive", zorder=2)
+        patch = Rectangle((plot_s, 0), plot_length, height, color="tab:olive", zorder=2)
         ax.add_patch(patch)

cheetah/accelerator/dipole.py

@@ -480,11 +480,15 @@ def defining_features(self) -> list[str]:
             "tracking_method",
         ]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
+        plot_s = s[vector_idx] if s.dim() > 0 else s
+        plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
+        plot_angle = self.angle[vector_idx] if self.angle.dim() > 0 else self.angle
+
         alpha = 1 if self.is_active else 0.2
-        height = 0.8 * (np.sign(self.angle[0]) if self.is_active else 1)
+        height = 0.8 * (np.sign(plot_angle) if self.is_active else 1)
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:green", alpha=alpha, zorder=2
+            (plot_s, 0), plot_length, height, color="tab:green", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)

cheetah/accelerator/drift.py

@@ -136,7 +136,7 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for i in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
         pass
 
     @property

cheetah/accelerator/element.py

@@ -126,12 +126,16 @@ def split(self, resolution: torch.Tensor) -> list["Element"]:
         raise NotImplementedError
 
     @abstractmethod
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
         """
         Plot a representation of this element into a `matplotlib` Axes at position `s`.
 
         :param ax: Axes to plot the representation into.
         :param s: Position of the object along s in meters.
+        :param vector_idx: Index of the vector dimension to plot. If the model has more
+            than one vector dimension, this can be used to select a specific one. In the
+            case of present vector dimension but no index provided, the first one is
+            used by default.
         """
         raise NotImplementedError
 

cheetah/accelerator/horizontal_corrector.py

@@ -80,12 +80,16 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for i in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
+        plot_s = s[vector_idx] if s.dim() > 0 else s
+        plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
+        plot_angle = self.angle[vector_idx] if self.angle.dim() > 0 else self.angle
+
         alpha = 1 if self.is_active else 0.2
-        height = 0.8 * (np.sign(self.angle[0]) if self.is_active else 1)
+        height = 0.8 * (np.sign(plot_angle) if self.is_active else 1)
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:blue", alpha=alpha, zorder=2
+            (plot_s, 0), plot_length, height, color="tab:blue", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/marker.py

@@ -37,7 +37,7 @@ def is_skippable(self) -> bool:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
         # Do nothing on purpose. Maybe later we decide markers should be shown, but for
         # now they are invisible.
         pass

cheetah/accelerator/quadrupole.py

@@ -217,11 +217,15 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for i in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
+        plot_k1 = self.k1[vector_idx] if self.k1.dim() > 0 else self.k1
+        plot_s = s[vector_idx] if s.dim() > 0 else s
+        plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
+
         alpha = 1 if self.is_active else 0.2
-        height = 0.8 * (np.sign(self.k1[0]) if self.is_active else 1)
+        height = 0.8 * (np.sign(plot_k1) if self.is_active else 1)
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:red", alpha=alpha, zorder=2
+            (plot_s, 0), plot_length, height, color="tab:red", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/screen.py

@@ -297,10 +297,13 @@ def set_read_beam(self, value: Beam) -> None:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float) -> None:
+    def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> None:
+        plot_s = s[vector_idx] if s.dim() > 0 else s
+
         alpha = 1 if self.is_active else 0.2
+
         patch = Rectangle(
-            (s, -0.6), 0, 0.6 * 2, color="tab:green", alpha=alpha, zorder=2
+            (plot_s, -0.6), 0, 0.6 * 2, color="tab:green", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)