diff --git a/cheetah/accelerator/transverse_deflecting_cavity.py b/cheetah/accelerator/transverse_deflecting_cavity.py index 9380d2405..2a227ef0d 100644 --- a/cheetah/accelerator/transverse_deflecting_cavity.py +++ b/cheetah/accelerator/transverse_deflecting_cavity.py @@ -117,9 +117,9 @@ def _track_drift_kick_drift(self, incoming: ParticleBeam) -> ParticleBeam: `ParticleBeam`. :return: Beam exiting the element. """ - assert isinstance( - incoming, ParticleBeam - ), "Drift-kick-drift tracking is currently only supported for `ParticleBeam`." + assert isinstance(incoming, ParticleBeam), ( + "Drift-kick-drift tracking is currently only supported for `ParticleBeam`." + ) # Compute Bmad coordinates and p0c x = incoming.x @@ -140,43 +140,56 @@ def _track_drift_kick_drift(self, incoming: ParticleBeam) -> ParticleBeam: x_offset, y_offset, self.tilt, x, px, y, py ) - x, y, z = bmadx.track_a_drift(self.length / 2, x, px, y, py, z, pz, p0c, mc2) + # compute the size of drift steps + dl = self.length / self.num_steps + + # start by tracking half a dl + x, y, z = bmadx.track_a_drift(dl / 2, x, px, y, py, z, pz, p0c, mc2) - voltage = self.voltage * -1 * incoming.species.num_elementary_charges / p0c + voltage = ( + self.voltage * -1 * incoming.species.num_elementary_charges / p0c + ) / self.num_steps k_rf = 2 * torch.pi * self.frequency / speed_of_light - # Phase that the particle sees - phase = ( - 2 - * torch.pi - * ( - self.phase.unsqueeze(-1) - - ( - bmadx.particle_rf_time(z, pz, p0c, mc2) - * self.frequency.unsqueeze(-1) + + for i in range(1, self.num_steps + 1): + # Phase that the particle sees + phase = ( + 2 + * torch.pi + * ( + self.phase.unsqueeze(-1) + - self.frequency.unsqueeze(-1) + * (bmadx.particle_rf_time(z, pz, p0c, mc2)) ) ) - ) - # TODO: Assigning px to px is really bad practice and should be separated into - # two separate variables - px = px + voltage.unsqueeze(-1) * phase.sin() + # TODO: Assigning px to px is really bad practice and should be separated into + # two separate variables + px = px + voltage.unsqueeze(-1) * phase.sin() - beta_old = ( - (1 + pz) - * p0c.unsqueeze(-1) - / (((1 + pz) * p0c.unsqueeze(-1)).square() + mc2.square()).sqrt() - ) - E_old = (1 + pz) * p0c.unsqueeze(-1) / beta_old - E_new = E_old + voltage.unsqueeze(-1) * phase.cos() * k_rf.unsqueeze( - -1 - ) * x * p0c.unsqueeze(-1) - pc = (E_new.square() - mc2.square()).sqrt() - beta = pc / E_new + beta_old = ( + (1 + pz) + * p0c.unsqueeze(-1) + / (((1 + pz) * p0c.unsqueeze(-1)).square() + mc2.square()).sqrt() + ) + E_old = (1 + pz) * p0c.unsqueeze(-1) / beta_old + E_new = E_old + voltage.unsqueeze(-1) * phase.cos() * k_rf.unsqueeze( + -1 + ) * x * p0c.unsqueeze(-1) + pc = (E_new.square() - mc2.square()).sqrt() + beta = pc / E_new + + pz = (pc - p0c.unsqueeze(-1)) / p0c.unsqueeze(-1) + z = z * beta / beta_old + + # if this is the last slice, skip tracking a full dl + if i == self.num_steps: + break - pz = (pc - p0c.unsqueeze(-1)) / p0c.unsqueeze(-1) - z = z * beta / beta_old + x, y, z = bmadx.track_a_drift(dl, x, px, y, py, z, pz, p0c, mc2) - x, y, z = bmadx.track_a_drift(self.length / 2, x, px, y, py, z, pz, p0c, mc2) + # Final half drift + x, y, z = bmadx.track_a_drift(dl / 2, x, px, y, py, z, pz, p0c, mc2) x, px, y, py = bmadx.offset_particle_unset( x_offset, y_offset, self.tilt, x, px, y, py @@ -198,6 +211,39 @@ def _track_drift_kick_drift(self, incoming: ParticleBeam) -> ParticleBeam: ) return outgoing_beam + def split(self, resolution: torch.Tensor) -> list[Element]: + # require that the num_steps is odd and greater than 1 + if self.num_steps % 2 != 1 or self.num_steps < 3: + raise ValueError( + "The number of steps for the transverse deflecting cavity must be" + " odd and greater than 1 to split the element." + ) + + # only allow splitting into 2 elements for now + if int(self.length / resolution) != 2: + raise ValueError( + "Transverse deflecting cavity can only be split into 2 elements for" + " now." + ) + num_splits = 2 + split_steps = int((self.num_steps - 1) / 2) + + return [ + TransverseDeflectingCavity( + self.length / num_splits, + self.voltage / num_splits, + self.phase, + self.frequency, + misalignment=self.misalignment, + tilt=self.tilt, + num_steps=split_steps, + tracking_method=self.tracking_method, + dtype=self.length.dtype, + device=self.length.device, + ) + for i in range(num_splits) + ] + def plot( self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None ) -> plt.Axes: diff --git a/tests/test_transverse_deflecting_cavity.py b/tests/test_transverse_deflecting_cavity.py index 886f422e3..bd8e96406 100644 --- a/tests/test_transverse_deflecting_cavity.py +++ b/tests/test_transverse_deflecting_cavity.py @@ -135,3 +135,76 @@ def test_transverse_deflecting_cavity_all_parameters_vectorization(): outgoing_beam = tdc.track(incoming_beam) assert outgoing_beam.particles.shape[:-2] == torch.Size([4, 3, 2, 2]) + + +def test_tdc_benchmark(): + cavity = cheetah.TransverseDeflectingCavity( + length=torch.tensor(0.2), + voltage=torch.tensor(1.0e6), + phase=torch.tensor(0.0), + frequency=torch.tensor(1.0e9), + ) + test_beam = cheetah.ParticleBeam( + torch.tensor([2e-3, 3e-3, -3e-3, -1e-3, -2e-3, 2e-3, 1.0]).unsqueeze(0), + energy=torch.tensor(4.0e7), + ) + assert torch.allclose( + cavity.track(test_beam).particles.flatten()[:-1], + torch.tensor( + [ + 2.705670627614420e-03, + 4.047421479988640e-03, + -3.200281391645270e-03, + -1.000000000000000e-03, + 1.998582178711370e-03, + -7.955332028185950e-04, + ], + ), + atol=1e-2, + ) + + +def test_transverse_deflecting_cavity_split(): + """ + Test that splitting a TDC into smaller segments works as expected. + """ + tdc = cheetah.TransverseDeflectingCavity( + length=torch.tensor(1.0), + voltage=torch.tensor(1e7), + phase=torch.tensor(0.4), + frequency=torch.tensor(1e9), + tracking_method="drift_kick_drift", + num_steps=11, + ) + + segments = tdc.split(resolution=torch.tensor(0.5)) + + assert len(segments) == 2 + for segment in segments: + assert isinstance(segment, cheetah.TransverseDeflectingCavity) + assert torch.isclose(segment.length, torch.tensor(0.5), rtol=1e-5) + # assert torch.equal(segment.voltage, tdc.voltage) + assert torch.equal(segment.phase, tdc.phase) + assert torch.equal(segment.frequency, tdc.frequency) + assert segment.num_steps == 5 + assert segment.tracking_method == tdc.tracking_method + + # test to make sure that tracking through the split segments gives the same result + # as tracking through the original segment + incoming_beam = cheetah.ParticleBeam.from_parameters( + num_particles=10, + sigma_px=torch.tensor(2e-7), + sigma_py=torch.tensor(2e-7), + energy=torch.tensor(50e6), + ) + + outgoing_beam_full = tdc.track(incoming_beam) + outgoing_beam_split = incoming_beam + for segment in segments: + outgoing_beam_split = segment.track(outgoing_beam_split) + + assert torch.allclose( + outgoing_beam_full.particles, + outgoing_beam_split.particles, + rtol=1e-2, + )