Merge pull request #172 from desy-ml/fix-cavity-tracking-issue

Fix cavity tracking issue introduced in recent updates

Author: jank324

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) (@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) (@jank324, @cr-xu)
 
 ### 🚀 Features
 

cheetah/accelerator/cavity.py

@@ -155,7 +155,8 @@ def _track_beam(self, incoming: Beam) -> Beam:
                     outgoing_energy.unsqueeze(-1) * beta1.unsqueeze(-1)
                 ) * (
                     torch.cos(
-                        incoming.particles[..., 4]
+                        -1
+                        * incoming.particles[..., 4]
                         * beta0.unsqueeze(-1)
                         * k.unsqueeze(-1)
                         + phi.unsqueeze(-1)

tests/test_compare_ocelot.py

@@ -698,13 +698,63 @@ def test_cavity():
 
     # Compare
     assert np.isclose(outgoing_beam.beta_x.cpu().numpy(), derived_twiss.beta_x)
-    assert np.isclose(
-        outgoing_beam.alpha_x.cpu().numpy(), derived_twiss.alpha_x, rtol=2e-5
-    )
+    assert np.isclose(outgoing_beam.alpha_x.cpu().numpy(), derived_twiss.alpha_x)
     assert np.isclose(outgoing_beam.beta_y.cpu().numpy(), derived_twiss.beta_y)
+    assert np.isclose(outgoing_beam.alpha_y.cpu().numpy(), derived_twiss.alpha_y)
     assert np.isclose(
-        outgoing_beam.alpha_y.cpu().numpy(), derived_twiss.alpha_y, rtol=2e-5
+        outgoing_beam.total_charge.cpu().numpy(), np.sum(outgoing_parray.q_array)
+    )
+    assert np.allclose(
+        outgoing_beam.particles[:, :, 5].cpu().numpy(),
+        outgoing_parray.rparticles.transpose()[:, 5],
+    )
+    assert np.allclose(
+        outgoing_beam.particles[:, :, 4].cpu().numpy(),
+        outgoing_parray.rparticles.transpose()[:, 4],
+    )
+
+
+def test_cavity_non_zero_phase():
+    """Compare tracking through a cavity with a phase offset."""
+    # Ocelot
+    tws = ocelot.Twiss()
+    tws.beta_x = 5.91253677
+    tws.alpha_x = 3.55631308
+    tws.beta_y = 5.91253677
+    tws.alpha_y = 3.55631308
+    tws.emit_x = 3.494768647122823e-09
+    tws.emit_y = 3.497810737006068e-09
+    tws.gamma_x = (1 + tws.alpha_x**2) / tws.beta_x
+    tws.gamma_y = (1 + tws.alpha_y**2) / tws.beta_y
+    tws.E = 6e-3
+
+    p_array = ocelot.generate_parray(tws=tws, charge=5e-9)
+
+    cell = [ocelot.Cavity(l=1.0377, v=0.01815975, freq=1.3e9, phi=30.0)]
+    lattice = ocelot.MagneticLattice(cell)
+    navigator = ocelot.Navigator(lattice=lattice)
+
+    _, outgoing_parray = ocelot.track(lattice, deepcopy(p_array), navigator)
+    derived_twiss = ocelot.cpbd.beam.get_envelope(outgoing_parray)
+
+    # Cheetah
+    incoming_beam = cheetah.ParticleBeam.from_ocelot(
+        parray=p_array, dtype=torch.float64
+    )
+    cheetah_cavity = cheetah.Cavity(
+        length=torch.tensor([1.0377]),
+        voltage=torch.tensor([0.01815975e9]),
+        frequency=torch.tensor([1.3e9]),
+        phase=torch.tensor([30.0]),
+        dtype=torch.float64,
     )
+    outgoing_beam = cheetah_cavity.track(incoming_beam)
+
+    # Compare
+    assert np.isclose(outgoing_beam.beta_x.cpu().numpy(), derived_twiss.beta_x)
+    assert np.isclose(outgoing_beam.alpha_x.cpu().numpy(), derived_twiss.alpha_x)
+    assert np.isclose(outgoing_beam.beta_y.cpu().numpy(), derived_twiss.beta_y)
+    assert np.isclose(outgoing_beam.alpha_y.cpu().numpy(), derived_twiss.alpha_y)
     assert np.isclose(
         outgoing_beam.total_charge.cpu().numpy(), np.sum(outgoing_parray.q_array)
     )