Fix typo in cavity tracking; Remove the tolerance in ocelot comparison

desy-ml · Mar 27, 2024 · 87adf12 · 87adf12
1 parent 42dbebf
commit 87adf12
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/cheetah/accelerator.py b/cheetah/accelerator.py
@@ -1007,7 +1007,7 @@ def _track_beam(self, incoming: ParticleBeam) -> ParticleBeam:
                 )
 
             if isinstance(incoming, ParameterBeam):
-                outgoing_mu[4] = incoming._mu[4] + (
+                outgoing_mu[4] = outgoing_mu[4] + (
                     T566 * incoming._mu[5] ** 2
                     + T556 * incoming._mu[4] * incoming._mu[5]
                     + T555 * incoming._mu[4] ** 2
@@ -1024,7 +1024,7 @@ def _track_beam(self, incoming: ParticleBeam) -> ParticleBeam:
                 )
                 outgoing_cov[5, 4] = outgoing_cov[4, 5]
             else:  # ParticleBeam
-                outgoing_particles[:, 4] = incoming.particles[:, 4] + (
+                outgoing_particles[:, 4] = outgoing_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
@@ -661,5 +661,4 @@ def test_cavity():
     assert np.allclose(
         outgoing_beam.particles[:, 4].cpu().numpy(),
         outgoing_parray.rparticles.transpose()[:, 4],
-        atol=1e-4,
     )