Sextupole implementation

CHANGELOG.md

@@ -10,6 +10,7 @@
 
 - Implement `split` method for the `Solenoid` element (see #380) (@cr-xu)
 - Implement a more robust RPN parser, fixing a bug where short strings in an Elegant variable definition would cause parsing to fail. (see #387) (@amylizzle, @Hespe, @jank324)
+- Add a `Sextupole` element (see #406) (@jank324, @Hespe)
 
 ### 🐛 Bug fixes
 

cheetah/__init__.py

@@ -13,6 +13,7 @@
     RBend,
     Screen,
     Segment,
+    Sextupole,
     Solenoid,
     SpaceChargeKick,
     TransverseDeflectingCavity,

cheetah/accelerator/__init__.py

@@ -11,6 +11,7 @@
 from .rbend import RBend  # noqa: F401
 from .screen import Screen  # noqa: F401
 from .segment import Segment  # noqa: F401
+from .sextupole import Sextupole  # noqa: F401
 from .solenoid import Solenoid  # noqa: F401
 from .space_charge_kick import SpaceChargeKick  # noqa: F401
 from .transverse_deflecting_cavity import TransverseDeflectingCavity  # noqa: F401

cheetah/accelerator/cavity.py

@@ -111,12 +111,10 @@ def _track_beam(self, incoming: Beam) -> Beam:
 
         tm = self.transfer_map(incoming.energy, incoming.species)
         if isinstance(incoming, ParameterBeam):
-            outgoing_mu = torch.matmul(tm, incoming.mu.unsqueeze(-1)).squeeze(-1)
-            outgoing_cov = torch.matmul(
-                tm, torch.matmul(incoming.cov, tm.transpose(-2, -1))
-            )
+            outgoing_mu = (tm @ incoming.mu.unsqueeze(-1)).squeeze(-1)
+            outgoing_cov = tm @ incoming.cov @ tm.transpose(-2, -1)
         else:  # ParticleBeam
-            outgoing_particles = torch.matmul(incoming.particles, tm.transpose(-2, -1))
+            outgoing_particles = incoming.particles @ tm.transpose(-2, -1)
         delta_energy = (
             self.voltage * torch.cos(phi) * incoming.species.num_elementary_charges * -1
         )

cheetah/accelerator/custom_transfer_map.py

@@ -69,10 +69,7 @@ def from_merging_elements(
             (*incoming_beam.energy.shape, 1, 1)
         )
         for element in elements:
-            tm = torch.matmul(
-                element.transfer_map(incoming_beam.energy, incoming_beam.species),
-                tm,
-            )
+            tm = element.transfer_map(incoming_beam.energy, incoming_beam.species) @ tm
             incoming_beam = element.track(incoming_beam)
 
         combined_length = sum(element.length for element in elements)

cheetah/accelerator/dipole.py

@@ -407,11 +407,11 @@ def transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor:
             R[..., 2, 3] = self.length
 
         # Apply fringe fields
-        R = torch.matmul(R_exit, torch.matmul(R, R_enter))
+        R = R_exit @ R @ R_enter
+
         # Apply rotation for tilted magnets
-        R = torch.matmul(
-            rotation_matrix(-self.tilt), torch.matmul(R, rotation_matrix(self.tilt))
-        )
+        R = rotation_matrix(-self.tilt) @ R @ rotation_matrix(self.tilt)
+
         return R
 
     def _transfer_map_enter(self) -> torch.Tensor:

cheetah/accelerator/drift.py

@@ -5,7 +5,12 @@
 
 from cheetah.accelerator.element import Element
 from cheetah.particles import Beam, ParticleBeam, Species
-from cheetah.utils import UniqueNameGenerator, bmadx, compute_relativistic_factors
+from cheetah.utils import (
+    UniqueNameGenerator,
+    bmadx,
+    compute_relativistic_factors,
+    verify_device_and_dtype,
+)
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
@@ -30,6 +35,7 @@ def __init__(
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
+        device, dtype = verify_device_and_dtype([length], device, dtype)
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__(name=name, **factory_kwargs)
 

cheetah/accelerator/element.py

@@ -68,8 +68,8 @@ def track(self, incoming: Beam) -> Beam:
         """
         if isinstance(incoming, ParameterBeam):
             tm = self.transfer_map(incoming.energy, incoming.species)
-            mu = torch.matmul(tm, incoming.mu.unsqueeze(-1)).squeeze(-1)
-            cov = torch.matmul(tm, torch.matmul(incoming.cov, tm.transpose(-2, -1)))
+            mu = (tm @ incoming.mu.unsqueeze(-1)).squeeze(-1)
+            cov = tm @ incoming.cov @ tm.transpose(-2, -1)
             return ParameterBeam(
                 mu,
                 cov,
@@ -79,7 +79,7 @@ def track(self, incoming: Beam) -> Beam:
             )
         elif isinstance(incoming, ParticleBeam):
             tm = self.transfer_map(incoming.energy, incoming.species)
-            new_particles = torch.matmul(incoming.particles, tm.transpose(-2, -1))
+            new_particles = incoming.particles @ tm.transpose(-2, -1)
             return ParticleBeam(
                 new_particles,
                 incoming.energy,

cheetah/accelerator/quadrupole.py

@@ -19,7 +19,7 @@ class Quadrupole(Element):
     :param length: Length in meters.
     :param k1: Strength of the quadrupole in 1/m^-2.
     :param misalignment: Misalignment vector of the quadrupole in x- and y-directions.
-    :param tilt: Tilt angle of the quadrupole in x-y plane [rad]. pi/4 for
+    :param tilt: Tilt angle of the quadrupole in x-y plane in radians. pi/4 for
         skew-quadrupole.
     :param num_steps: Number of drift-kick-drift steps to use for tracking through the
         element when tracking method is set to `"bmadx"`.

cheetah/accelerator/segment.py

@@ -399,7 +399,7 @@ def transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor:
         if self.is_skippable:
             tm = torch.eye(7, device=energy.device, dtype=energy.dtype)
             for element in self.elements:
-                tm = torch.matmul(element.transfer_map(energy, species), tm)
+                tm = element.transfer_map(energy, species) @ tm
             return tm
         else:
             return None

cheetah/accelerator/sextupole.py

@@ -0,0 +1,142 @@
+import matplotlib.pyplot as plt
+import torch
+
+from cheetah.accelerator.element import Element
+from cheetah.particles import Beam, ParameterBeam, ParticleBeam, Species
+from cheetah.track_methods import base_rmatrix, base_ttensor, misalignment_matrix
+from cheetah.utils import verify_device_and_dtype
+
+
+class Sextupole(Element):
+    """
+    A sextupole element in a particle accelerator.
+
+    :param length: Length in meters.
+    :param k2: Sextupole strength in 1/m^3.
+    :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 name: Unique identifier of the element.
+    """
+
+    def __init__(
+        self,
+        length: torch.Tensor,
+        k2: torch.Tensor | None = None,
+        misalignment: torch.Tensor | None = None,
+        tilt: torch.Tensor | None = None,
+        name: str | None = None,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        device, dtype = verify_device_and_dtype(
+            [length, k2, misalignment, tilt], device, dtype
+        )
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(name=name, **factory_kwargs)
+
+        self.length = torch.as_tensor(length, **factory_kwargs)
+
+        self.register_buffer_or_parameter(
+            "k2", torch.as_tensor(k2 if k2 is not None else 0.0, **factory_kwargs)
+        )
+        self.register_buffer_or_parameter(
+            "misalignment",
+            torch.as_tensor(
+                misalignment if misalignment is not None else (0.0, 0.0),
+                **factory_kwargs,
+            ),
+        )
+        self.register_buffer_or_parameter(
+            "tilt", torch.as_tensor(tilt if tilt is not None else 0.0, **factory_kwargs)
+        )
+
+    def transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor:
+        R = base_rmatrix(
+            length=self.length,
+            k1=torch.zeros_like(self.length),
+            hx=torch.zeros_like(self.length),
+            species=species,
+            tilt=self.tilt,
+            energy=energy,
+        )
+
+        if torch.all(self.misalignment == 0):
+            return R
+        else:
+            R_entry, R_exit = misalignment_matrix(self.misalignment)
+            R = R_exit @ R @ R_entry
+            return R
+
+    def track(self, incoming: Beam) -> Beam:
+        """
+        Track the beam through the sextupole element.
+
+        :param incoming: Beam entering the element.
+        :return: Beam exiting the element.
+        """
+        first_order_tm = self.transfer_map(incoming.energy, incoming.species)
+        second_order_tm = base_ttensor(
+            length=self.length,
+            k1=torch.zeros_like(self.length),
+            k2=self.k2,
+            hx=torch.zeros_like(self.length),
+            species=incoming.species,
+            tilt=self.tilt,
+            energy=incoming.energy,
+        )
+
+        if isinstance(incoming, ParameterBeam):
+            # For ParameterBeam, only first-order effects are applied
+            return super().track(incoming)
+        elif isinstance(incoming, ParticleBeam):
+            # Apply the transfer map to the incoming particles
+            first_order_particles = incoming.particles @ first_order_tm.transpose(
+                -2, -1
+            )
+            second_order_particles = torch.einsum(
+                "...ijk,...j,...k->...i",
+                second_order_tm.unsqueeze(-4),  # Add broadcast dimension for particles
+                incoming.particles,
+                incoming.particles,
+            )
+            outgoing_particles = second_order_particles + first_order_particles
+
+            return ParticleBeam(
+                particles=outgoing_particles,
+                energy=incoming.energy,
+                particle_charges=incoming.particle_charges,
+                survival_probabilities=incoming.survival_probabilities,
+                species=incoming.species,
+            )
+        else:
+            raise TypeError(
+                f"Unsupported beam type: {type(incoming)}. Expected ParameterBeam or "
+                "ParticleBeam."
+            )
+
+    @property
+    def is_skippable(self) -> bool:
+        return False
+
+    @property
+    def is_active(self) -> bool:
+        return torch.any(self.k2 != 0.0).item()
+
+    def split(self, resolution: torch.Tensor) -> list[Element]:
+        raise NotImplementedError
+
+    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+        raise NotImplementedError
+
+    @property
+    def defining_features(self) -> list[str]:
+        return super().defining_features + ["length", "k2", "misalignment", "tilt"]
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(length={repr(self.length)}, "
+            f"k2={repr(self.k2)}, "
+            f"misalignment={repr(self.misalignment)}, "
+            f"tilt={repr(self.tilt)}, "
+            f"name={repr(self.name)})"
+        )