Fix twiss plot

CHANGELOG.md

@@ -4,7 +4,7 @@
 
 ### 🚨 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, #215, #218) (@jank324, @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, #213, #215, #218) (@jank324, @cr-xu)
 
 ### 🚀 Features
 

cheetah/accelerator/segment.py

@@ -511,10 +511,10 @@ def plot_twiss(self, beam: Beam, ax: Optional[Any] = None) -> None:
             outgoing = element.track(longitudinal_beams[-1])
 
             longitudinal_beams.append(outgoing)
-            s_positions.append(s_positions[-1] + element.length)
+            s_positions.append(s_positions[-1] + element.length[0])
 
-        beta_x = [beam.beta_x for beam in longitudinal_beams]
-        beta_y = [beam.beta_y for beam in longitudinal_beams]
+        beta_x = [beam.beta_x[0] for beam in longitudinal_beams]
+        beta_y = [beam.beta_y[0] for beam in longitudinal_beams]
 
         if ax is None:
             fig = plt.figure()

tests/test_plotting.py

@@ -0,0 +1,26 @@
+import cheetah
+
+from .resources import ARESlatticeStage3v1_9 as ares
+
+
+def test_twiss_plot():
+    """
+    Test that the Twiss plot does not raise an exception using the ARES EA as an
+    example.
+    """
+    cell = cheetah.converters.ocelot.subcell_of_ocelot(
+        ares.cell, "AREASOLA1", "AREABSCR1"
+    )
+    ares.areamqzm1.k1 = 5.0
+    ares.areamqzm2.k1 = -5.0
+    ares.areamcvm1.k1 = 1e-3
+    ares.areamqzm3.k1 = 5.0
+    ares.areamchm1.k1 = -2e-3
+
+    incoming_beam = cheetah.ParticleBeam.from_astra(
+        "tests/resources/ACHIP_EA1_2021.1351.001"
+    )
+    segment = cheetah.Segment.from_ocelot(cell)
+
+    # Do the actual plotting
+    segment.plot_twiss(incoming_beam)