Fix transfer map in cavity;

Update CHANGELOG.md
desy-ml · Feb 14, 2024 · 7a0eadb · 7a0eadb
1 parent 4006982
commit 7a0eadb
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 29 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,8 @@
 
 ### 🐛 Bug fixes
 
+- Fix bug in `Cavity` transfer map bug. (see #128) (@cr-xu)
+
 ### 🐆 Other
 
 ## [v0.6.2](https://github.com/desy-ml/cheetah/releases/tag/v0.6.2) (2024-02-13)

diff --git a/cheetah/accelerator.py b/cheetah/accelerator.py
@@ -940,41 +940,31 @@ def _track_beam(self, incoming: ParticleBeam) -> ParticleBeam:
         T566 = 1.5 * self.length * igamma2 / beta0**3
         T556 = 0
         T555 = 0
+
         if incoming.energy + delta_energy > 0:
             k = 2 * torch.pi * self.frequency / constants.speed_of_light
             outgoing_energy = incoming.energy + delta_energy
             g1 = outgoing_energy / electron_mass_eV
             beta1 = torch.sqrt(1 - 1 / g1**2)
 
             if isinstance(incoming, ParameterBeam):
-                outgoing_mu[5] = (
-                    incoming._mu[5]
-                    + incoming.energy * beta0 / (outgoing_energy * beta1)
-                    + self.voltage
-                    * beta0
-                    / (outgoing_energy * beta1)
-                    * (torch.cos(incoming._mu[4] * beta0 * k + phi) - torch.cos(phi))
+                outgoing_mu[5] = incoming._mu[5] * incoming.energy * beta0 / (
+                    outgoing_energy * beta1
+                ) + self.voltage * beta0 / (outgoing_energy * beta1) * (
+                    torch.cos(-incoming._mu[4] * beta0 * k + phi) - torch.cos(phi)
                 )
                 outgoing_cov[5, 5] = incoming._cov[5, 5]
-                # outgoing_cov[5, 5] = (
-                #     incoming._cov[5, 5]
-                #     + incoming.energy * beta0 / (outgoing_energy * beta1)
-                #     + self.voltage
-                #     * beta0
-                #     / (outgoing_energy * beta1)
-                #     * (torch.cos(incoming._mu[4] * beta0 * k + phi) - torch.cos(phi))
-                # )
+
             else:  # ParticleBeam
-                outgoing_particles[:, 5] = (
-                    incoming.particles[:, 5]
-                    + incoming.energy * beta0 / (outgoing_energy * beta1)
-                    + self.voltage
-                    * beta0
-                    / (outgoing_energy * beta1)
-                    * (
-                        torch.cos(incoming.particles[:, 4] * beta0 * k + phi)
-                        - torch.cos(phi)
-                    )
+                outgoing_particles[:, 5] = incoming.particles[
+                    :, 5
+                ] * incoming.energy * beta0 / (
+                    outgoing_energy * beta1
+                ) + self.voltage * beta0 / (
+                    outgoing_energy * beta1
+                ) * (
+                    torch.cos(-incoming.particles[:, 4] * beta0 * k + phi)
+                    - torch.cos(phi)
                 )
 
             dgamma = self.voltage / electron_mass_eV
@@ -1017,7 +1007,7 @@ def _track_beam(self, incoming: ParticleBeam) -> ParticleBeam:
                 )
 
             if isinstance(incoming, ParameterBeam):
-                outgoing_mu[4] = (
+                outgoing_mu[4] = incoming._mu[4] + (
                     T566 * incoming._mu[5] ** 2
                     + T556 * incoming._mu[4] * incoming._mu[5]
                     + T555 * incoming._mu[4] ** 2
@@ -1034,7 +1024,7 @@ def _track_beam(self, incoming: ParticleBeam) -> ParticleBeam:
                 )
                 outgoing_cov[5, 4] = outgoing_cov[4, 5]
             else:  # ParticleBeam
-                outgoing_particles[:, 4] = (
+                outgoing_particles[:, 4] = incoming.particles[:, 4] + (
                     T566 * incoming.particles[:, 5] ** 2
                     + T556 * incoming.particles[:, 4] * incoming.particles[:, 5]
                     + T555 * incoming.particles[:, 4] ** 2

diff --git a/tests/test_compare_ocelot.py b/tests/test_compare_ocelot.py
@@ -645,12 +645,21 @@ 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=1e-4
+        outgoing_beam.alpha_x.cpu().numpy(), derived_twiss.alpha_x, rtol=2e-5
     )
     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, rtol=1e-4
+        outgoing_beam.alpha_y.cpu().numpy(), derived_twiss.alpha_y, rtol=2e-5
     )
     assert np.isclose(
         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],
+        atol=1e-4,
+    )