tcav splitting

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:

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,
+    )