Fix plotting in 0-dimensional case

cheetah/accelerator/cavity.py

@@ -346,7 +346,7 @@ def plot(self, ax: plt.Axes, s: float) -> None:
         height = 0.4
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="gold", alpha=alpha, zorder=2
+            (s, 0), self.length, height, color="gold", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/custom_transfer_map.py

@@ -105,5 +105,5 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
     def plot(self, ax: plt.Axes, s: float) -> None:
         height = 0.4
 
-        patch = Rectangle((s, 0), self.length[0], height, color="tab:olive", zorder=2)
+        patch = Rectangle((s, 0), self.length, height, color="tab:olive", zorder=2)
         ax.add_patch(patch)

cheetah/accelerator/dipole.py

@@ -482,9 +482,9 @@ def defining_features(self) -> list[str]:
 
     def plot(self, ax: plt.Axes, s: float) -> None:
         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(self.angle) if self.is_active else 1)
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:green", alpha=alpha, zorder=2
+            (s, 0), self.length, height, color="tab:green", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)

cheetah/accelerator/horizontal_corrector.py

@@ -82,10 +82,10 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
 
     def plot(self, ax: plt.Axes, s: float) -> None:
         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(self.angle) if self.is_active else 1)
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:blue", alpha=alpha, zorder=2
+            (s, 0), self.length, height, color="tab:blue", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/quadrupole.py

@@ -219,9 +219,9 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
 
     def plot(self, ax: plt.Axes, s: float) -> None:
         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(self.k1) if self.is_active else 1)
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:red", alpha=alpha, zorder=2
+            (s, 0), self.length, height, color="tab:red", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/segment.py

@@ -386,7 +386,7 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
         ]
 
     def plot(self, ax: plt.Axes, s: float) -> None:
-        element_lengths = [element.length[0] for element in self.elements]
+        element_lengths = [element.length for element in self.elements]
         element_ss = [0] + [
             sum(element_lengths[: i + 1]) for i, _ in enumerate(element_lengths)
         ]
@@ -423,7 +423,7 @@ def plot_reference_particle_traces(
         reference_segment = deepcopy(self)
         splits = reference_segment.split(resolution=torch.tensor(resolution))
 
-        split_lengths = [split.length[0] for split in splits]
+        split_lengths = [split.length for split in splits]
         ss = [0] + [sum(split_lengths[: i + 1]) for i, _ in enumerate(split_lengths)]
 
         references = []
@@ -464,7 +464,7 @@ def plot_reference_particle_traces(
 
         for particle_index in range(num_particles):
             xs = [
-                float(reference_beam.x[0, particle_index].cpu())
+                reference_beam.x[particle_index]
                 for reference_beam in references
                 if reference_beam is not Beam.empty
             ]
@@ -475,7 +475,7 @@ def plot_reference_particle_traces(
 
         for particle_index in range(num_particles):
             ys = [
-                float(reference_beam.ys[0, particle_index].cpu())
+                reference_beam.y[particle_index]
                 for reference_beam in references
                 if reference_beam is not Beam.empty
             ]

cheetah/accelerator/solenoid.py

@@ -121,7 +121,7 @@ def plot(self, ax: plt.Axes, s: float) -> None:
         height = 0.8
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:orange", alpha=alpha, zorder=2
+            (s, 0), self.length, height, color="tab:orange", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/transverse_deflecting_cavity.py

@@ -223,7 +223,7 @@ def plot(self, ax: plt.Axes, s: float) -> None:
         height = 0.4
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="olive", alpha=alpha, zorder=2
+            (s, 0), self.length, height, color="olive", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/undulator.py

@@ -69,7 +69,7 @@ def plot(self, ax: plt.Axes, s: float) -> None:
         height = 0.4
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:purple", alpha=alpha, zorder=2
+            (s, 0), self.length, height, color="tab:purple", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)
 

cheetah/accelerator/vertical_corrector.py

@@ -85,10 +85,10 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
 
     def plot(self, ax: plt.Axes, s: float) -> None:
         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(self.angle) if self.is_active else 1)
 
         patch = Rectangle(
-            (s, 0), self.length[0], height, color="tab:cyan", alpha=alpha, zorder=2
+            (s, 0), self.length, height, color="tab:cyan", alpha=alpha, zorder=2
         )
         ax.add_patch(patch)