From 2fbd5c4846ef4c69e1650918b1daa5f7f0c7841d Mon Sep 17 00:00:00 2001 From: Chenran Xu Date: Wed, 12 Jun 2024 17:01:45 +0200 Subject: [PATCH 1/2] Fix a missing sign in the cavity tracking method; Remove large rtol in tests --- cheetah/accelerator/cavity.py | 3 +- tests/test_compare_ocelot.py | 58 ++++++++++++++++++++++++++++++++--- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/cheetah/accelerator/cavity.py b/cheetah/accelerator/cavity.py index a897c593..421aeb33 100644 --- a/cheetah/accelerator/cavity.py +++ b/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) diff --git a/tests/test_compare_ocelot.py b/tests/test_compare_ocelot.py index 05ca9af0..4008f9ac 100644 --- a/tests/test_compare_ocelot.py +++ b/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) ) From 1083409a61b5c10c16dd1e1b420cbdb105398cc2 Mon Sep 17 00:00:00 2001 From: Chenran Xu Date: Wed, 12 Jun 2024 17:08:42 +0200 Subject: [PATCH 2/2] Update CHANGELOG.md --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 99e5b4ec..164082d5 100644 --- a/CHANGELOG.md +++ b/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