From 358c9993f16dfa149dd55bc178947d8353c65c67 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 13:48:34 +0200 Subject: [PATCH 01/30] Add tests for `Sextupole` implementation --- tests/test_sextupole.py | 51 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 tests/test_sextupole.py diff --git a/tests/test_sextupole.py b/tests/test_sextupole.py new file mode 100644 index 000000000..f23518b40 --- /dev/null +++ b/tests/test_sextupole.py @@ -0,0 +1,51 @@ +from copy import deepcopy + +import ocelot +import torch + +import cheetah + + +def test_compare_sextupole_to_ocelot(): + """Compare the results of tracking through a sextupole in Cheetah and Ocelot.""" + length = 0.34 + k2 = 0.5 + tilt = 0.1 + + # Track through a sextupole in Cheetah + incoming = cheetah.ParticleBeam.from_astra( + "tests/resources/ACHIP_EA1_2021.1351.001" + ) + cheetah_sextupole = cheetah.Sextupole( + length=torch.tensor(length), k2=torch.tensor(k2), tilt=torch.tensor(tilt) + ) + outgoing_cheetah = cheetah_sextupole.track(incoming) + + # Convert to Ocelot sextupole + incoming_p_array = ocelot.astraBeam2particleArray( + "tests/resources/ACHIP_EA1_2021.1351.001" + ) + lattice = ocelot.MagneticLattice([ocelot.Sextupole(l=length, k2=k2, tilt=tilt)]) + navigator = ocelot.Navigator(lattice) + _, outgoing_p_array = ocelot.track(lattice, deepcopy(incoming_p_array), navigator) + outgoing_ocelot = cheetah.ParticleBeam.from_ocelot(outgoing_p_array) + + # Compare the results + assert torch.allclose(outgoing_cheetah.particles, outgoing_ocelot.particles) + + +def test_sextupole_as_drift(): + """Test that a sextupole with k2=0 is equivalent to a drift.""" + incoming = cheetah.ParticleBeam.from_astra( + "tests/resources/ACHIP_EA1_2021.1351.001" + ) + + sextupole = cheetah.Sextupole(length=torch.tensor(0.34), k2=torch.tensor(0.0)) + drift = cheetah.Drift(length=torch.tensor(0.34)) + + # Track through the sextupole and drift + sextupole_outgoing = sextupole.track(incoming) + drift_outgoing = drift.track(incoming) + + # Check that the results are the same + assert torch.allclose(sextupole_outgoing.particles, drift_outgoing.particles) From dc37cff0bb04ffa516117711a78dc7c9d5551045 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 15:04:18 +0200 Subject: [PATCH 02/30] Add sextupole sekelton --- cheetah/__init__.py | 1 + cheetah/accelerator/__init__.py | 1 + cheetah/accelerator/drift.py | 8 +- cheetah/accelerator/sextupole.py | 131 +++++++++++++++++++++++++++++++ 4 files changed, 140 insertions(+), 1 deletion(-) create mode 100644 cheetah/accelerator/sextupole.py diff --git a/cheetah/__init__.py b/cheetah/__init__.py index 5a37aa789..02d936a87 100644 --- a/cheetah/__init__.py +++ b/cheetah/__init__.py @@ -13,6 +13,7 @@ RBend, Screen, Segment, + Sextupole, Solenoid, SpaceChargeKick, TransverseDeflectingCavity, diff --git a/cheetah/accelerator/__init__.py b/cheetah/accelerator/__init__.py index bc783dcee..7a8320e1a 100644 --- a/cheetah/accelerator/__init__.py +++ b/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 diff --git a/cheetah/accelerator/drift.py b/cheetah/accelerator/drift.py index cf72e1dc3..5bf701402 100644 --- a/cheetah/accelerator/drift.py +++ b/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) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py new file mode 100644 index 000000000..2889a7f81 --- /dev/null +++ b/cheetah/accelerator/sextupole.py @@ -0,0 +1,131 @@ +import matplotlib.pyplot as plt +import torch + +from cheetah.accelerator.element import Element +from cheetah.particles import Beam, ParameterBeam, ParticleBeam, Species +from cheetah.utils import compute_relativistic_factors, verify_device_and_dtype + + +class Sextupole(Element): + """ + A sextupole element in a particle accelerator. + + :param length: Length in meters. + :param k2: TODO + :param misalignment: TODO + :param tilt: TODO + :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], 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: + device = self.length.device + dtype = self.length.dtype + + _, igamma2, beta = compute_relativistic_factors(energy, species.mass_eV) + + vector_shape = torch.broadcast_shapes(self.length.shape, igamma2.shape) + + tm = torch.eye(7, device=device, dtype=dtype).repeat((*vector_shape, 1, 1)) + tm[..., 0, 1] = self.length + tm[..., 2, 3] = self.length + tm[..., 4, 5] = -self.length / beta**2 * igamma2 + + return tm + + 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 = torch.eye( + 7, device=self.length.device, dtype=self.length.dtype + ).repeat((*incoming.mu_x.shape, 1, 1)) + second_order_tm[..., 0, 1] = self.length + second_order_tm[..., 2, 3] = self.length + second_order_tm[..., 4, 5] = ( + -self.length + / incoming.relativistic_beta**2 + * incoming.relativistic_gamma**2 + ) + second_order_tm[..., 1, 0] = self.k2 * self.length**3 / 6 + second_order_tm[..., 3, 2] = self.k2 * self.length**3 / 6 + second_order_tm[..., 5, 4] = -self.k2 * self.length**3 / 6 + second_order_tm[..., 1, 2] = self.k2 * self.length**3 / 6 + second_order_tm[..., 3, 4] = self.k2 * self.length**3 / 6 + second_order_tm[..., 5, 0] = -self.k2 * self.length**3 / 6 + second_order_tm[..., 5, 2] = -self.k2 * self.length**3 / 6 + second_order_tm[..., 1, 4] = -self.k2 * self.length**3 / 6 + second_order_tm[..., 3, 0] = -self.k2 * self.length**3 / 6 + second_order_tm[..., 5, 3] = -self.k2 * self.length**3 / 6 + second_order_tm[..., 1, 5] = -self.k2 * self.length**3 / 6 + second_order_tm[..., 3, 2] = -self.k2 * self.length**3 / 6 + second_order_tm[..., 5, 4] = -self.k2 * self.length**3 / 6 + + # Apply the transfer map to the incoming beam + particles = torch.einsum("ijk,ikl->ijl", first_order_tm, incoming.particles) + particles = torch.einsum("ijk,ikl->ijl", second_order_tm, particles) + return ParticleBeam( + particles=particles, + energy=incoming.energy, + particle_charges=incoming.particle_charges, + survival_probabilities=incoming.survival_probabilities, + species=incoming.species, + ) + + @property + def is_skippable(self) -> bool: + return False + + @property + def is_active(self) -> bool: + return torch.any(self.k2 != 0.0) + + 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 + + def defining_features(self) -> list[str]: + return super().defining_features() + ["length", "k2"] + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(length={repr(self.length)}, " + f"k2={repr(self.k2)}, " + f"name={repr(self.name)})" + ) From 75a0f98e2aac729cc8cb0df3b72e812ab4a70cd9 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 15:06:38 +0200 Subject: [PATCH 03/30] Rearrange contents of `track_methods.py` --- cheetah/track_methods.py | 46 ++++++++++++++++++++-------------------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 1c687c4de..5fed5245d 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -6,29 +6,6 @@ from cheetah.utils import compute_relativistic_factors -def rotation_matrix(angle: torch.Tensor) -> torch.Tensor: - """Rotate the transfer map in x-y plane. - - :param angle: Rotation angle in rad, for example `angle = np.pi/2` for vertical = - dipole. - :return: Rotation matrix to be multiplied to the element's transfer matrix. - """ - cs = torch.cos(angle) - sn = torch.sin(angle) - - tm = torch.eye(7, dtype=angle.dtype, device=angle.device).repeat(*angle.shape, 1, 1) - tm[..., 0, 0] = cs - tm[..., 0, 2] = sn - tm[..., 1, 1] = cs - tm[..., 1, 3] = sn - tm[..., 2, 0] = -sn - tm[..., 2, 2] = cs - tm[..., 3, 1] = -sn - tm[..., 3, 3] = cs - - return tm - - def base_rmatrix( length: torch.Tensor, k1: torch.Tensor, @@ -102,6 +79,29 @@ def base_rmatrix( return R +def rotation_matrix(angle: torch.Tensor) -> torch.Tensor: + """Rotate the transfer map in x-y plane. + + :param angle: Rotation angle in rad, for example `angle = np.pi/2` for vertical = + dipole. + :return: Rotation matrix to be multiplied to the element's transfer matrix. + """ + cs = torch.cos(angle) + sn = torch.sin(angle) + + tm = torch.eye(7, dtype=angle.dtype, device=angle.device).repeat(*angle.shape, 1, 1) + tm[..., 0, 0] = cs + tm[..., 0, 2] = sn + tm[..., 1, 1] = cs + tm[..., 1, 3] = sn + tm[..., 2, 0] = -sn + tm[..., 2, 2] = cs + tm[..., 3, 1] = -sn + tm[..., 3, 3] = cs + + return tm + + def misalignment_matrix( misalignment: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: From 7cbe9e04a8c8daf110768b861ae6351d11a33014 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 16:08:56 +0200 Subject: [PATCH 04/30] Implement method to generate second order transfer map --- cheetah/track_methods.py | 208 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 205 insertions(+), 3 deletions(-) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 5fed5245d..5c32874b6 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -15,15 +15,15 @@ def base_rmatrix( energy: torch.Tensor | None = None, ) -> torch.Tensor: """ - Create a universal transfer matrix for a beamline element. + Create a first order universal transfer matrix for a beamline element. :param length: Length of the element in m. :param k1: Quadrupole strength in 1/m**2. - :param hx: Curvature (1/radius) of the element in 1/m**2. + :param hx: Curvature (1/radius) of the element in 1/m**2. TODO: Unit correct? :param species: Particle species of the beam. :param tilt: Roation of the element relative to the longitudinal axis in rad. :param energy: Beam energy in eV. - :return: Transfer matrix for the element. + :return: First order transfer matrix for the element. """ device = length.device dtype = length.dtype @@ -79,6 +79,208 @@ def base_rmatrix( return R +def base_tmatrix( + length: torch.Tensor, + k1: torch.Tensor, + k2: torch.Tensor, + hx: torch.Tensor, + species: Species, + tilt: torch.Tensor | None = None, + energy: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Create a second order universal transfer matrix for a beamline element. + + :param length: Length of the element in m. + :param k1: Quadrupole strength in 1/m**2. + :param k2: Sextupole strength in 1/m**3. + :param hx: Curvature (1/radius) of the element in 1/m**2. + :param species: Particle species of the beam. + :param tilt: Roation of the element relative to the longitudinal axis in rad. + :param energy: Beam energy in eV. + :return: Second order transfer matrix for the element. + """ + device = length.device + dtype = length.dtype + + tilt = tilt if tilt is not None else torch.tensor(0.0, device=device, dtype=dtype) + energy = ( + energy if energy is not None else torch.tensor(0.0, device=device, dtype=dtype) + ) + + _, igamma2, beta = compute_relativistic_factors(energy, species.mass_eV) + + # Avoid division by zero + k1 = k1.clone() + k1[k1 == 0] = 1e-12 + + kx2 = k1 + hx**2 + ky2 = -k1 + kx = torch.sqrt(torch.complex(kx2, torch.tensor(0.0, device=device, dtype=dtype))) + ky = torch.sqrt(torch.complex(ky2, torch.tensor(0.0, device=device, dtype=dtype))) + cx = torch.cos(kx * length).real + cy = torch.cos(ky * length).real + sy = (torch.sin(ky * length) / ky).real + sx = (torch.sin(kx * length) / kx).real + dx = hx / kx2 * (1.0 - cx) + + d2y = 0.5 * sy**2 + s2y = sy * cy + c2y = torch.cos(2 * ky * length).real + fx = torch.where(kx2 != 0, (length - sx) / kx2, length**3 / 6.0) + f2y = torch.where(ky2 != 0, (length - s2y) / ky2, length**3 / 6.0) + + j1 = torch.where(kx2 != 0, (length - sx) / kx2, length**3 / 6.0) + j2 = torch.where( + kx2 != 0, + (3.0 * length - 4.0 * sx + sx * cx) / (2 * kx2**2), + length**5 / 20.0, + ) + j3 = torch.where( + kx2 != 0, + (15.0 * length - 22.5 * sx + 9.0 * sx * cx - 1.5 * sx * cx**2 + kx2 * sx**3) + / (6.0 * kx2**3), + length**7 / 56.0, + ) + j_denominator = kx2 - 4 * ky2 + jc = torch.where(j_denominator != 0, (c2y - cx) / j_denominator, 0.5 * length**2) + js = torch.where( + j_denominator != 0, (cy * sy - sx) / j_denominator, length**3 / 6.0 + ) + jd = torch.where(j_denominator != 0, (d2y - dx) / j_denominator, length**4 / 24.0) + jf = torch.where(j_denominator != 0, (f2y - fx) / j_denominator, length**5 / 120.0) + + khk = k2 + 2 * hx * k1 + + vector_shape = torch.broadcast_shapes( + length.shape, k1.shape, hx.shape, tilt.shape, energy.shape + ) + + T = torch.zeros((7, 7, 7), dtype=dtype, device=device).repeat(*vector_shape, 1, 1) + T[..., 0, 0, 0] = -1 / 6 * khk * (sx**2 + dx) - 0.5 * hx * kx2 * sx**2 + T[..., 0, 0, 1] = 2 * -1 / 6 * khk * sx * dx + 0.5 * hx * sx * cx + T[..., 0, 1, 1] = -1 / 6 * khk * dx**2 + 0.5 * hx * dx * cx + T[..., 0, 0, 5] = ( + 2 * -hx / 12 / beta * khk * (3 * sx * j1 - dx**2) + + 0.5 * hx**2 / beta * sx**2 + + 0.25 / beta * k1 * length * sx + ) + T[..., 0, 1, 5] = ( + 2 * -hx / 12 / beta * khk * (sx * dx**2 - 2 * cx * j2) + + 0.25 * hx**2 / beta * (sx * dx + cx * j1) + - 0.25 / beta * (sx + length * cx) + ) + T[..., 0, 5, 5] = ( + -(hx**2) / 6 / beta**2 * khk * (dx**2 * dx - 2 * sx * j2) + + 0.5 * hx**3 / beta**2 * sx * j1 + - 0.5 * hx / beta**2 * length * sx + - 0.5 * hx / (beta**2) * igamma2 * dx + ) + T[..., 0, 2, 2] = k1 * k2 * jd + 0.5 * (k2 + hx * k1) * dx + T[..., 0, 2, 3] = 2 * 0.5 * k2 * js + T[..., 0, 3, 3] = k2 * jd - 0.5 * hx * dx + T[..., 1, 0, 0] = -1 / 6 * khk * sx * (1 + 2 * cx) + T[..., 1, 0, 1] = 2 * -1 / 6 * khk * dx * (1 + 2 * cx) + T[..., 1, 1, 1] = -1 / 3 * khk * sx * dx - 0.5 * hx * sx + T[..., 1, 0, 5] = 2 * -hx / 12 / beta * khk * ( + 3 * cx * j1 + sx * dx + ) - 0.25 / beta * k1 * (sx - length * cx) + T[..., 1, 1, 5] = ( + 2 * -hx / 12 / beta * khk * (3 * sx * j1 + dx**2) + + 0.25 / beta * k1 * length * sx + ) + T[..., 1, 5, 5] = ( + -(hx**2) / 6 / beta**2 * khk * (sx * dx**2 - 2 * cx * j2) + - 0.5 * hx / beta**2 * k1 * (cx * j1 - sx * dx) + - 0.5 * hx / beta**2 * igamma2 * sx + ) + T[..., 1, 2, 2] = k1 * k2 * js + 0.5 * (k2 + hx * k1) * sx + T[..., 1, 2, 3] = 2 * 0.5 * k2 * jc + T[..., 1, 3, 3] = k2 * js - 0.5 * hx * sx + T[..., 2, 0, 2] = ( + 2 * 0.5 * k2 * (cy * jc - 2 * k1 * sy * js) + 0.5 * hx * k1 * sx * sy + ) + T[..., 2, 0, 3] = 2 * 0.5 * k2 * (sy * jc - 2 * cy * js) + 0.5 * hx * sx * cy + T[..., 2, 1, 2] = ( + 2 * 0.5 * k2 * (cy * js - 2 * k1 * sy * jd) + 0.5 * hx * k1 * dx * sy + ) + T[..., 2, 1, 3] = 2 * 0.5 * k2 * (sy * js - 2 * cy * jd) + 0.5 * hx * dx * cy + T[..., 2, 2, 5] = ( + 2 * 0.5 * hx / beta * k2 * (cy * jd - 2 * k1 * sy * jf) + + 0.5 * hx**2 / beta * k1 * j1 * sy + - 0.25 / beta * k1 * length * sy + ) + T[..., 2, 3, 5] = ( + 2 * 0.5 * hx / beta * k2 * (sy * jd - 2 * cy * jf) + + 0.5 * hx**2 / beta * j1 * cy + - 0.25 / beta * (sy + length * cy) + ) + T[..., 3, 0, 2] = ( + 2 * 0.5 * k1 * k2 * (2 * cy * js - sy * jc) + 0.5 * (k2 + hx * k1) * sx * cy + ) + T[..., 3, 0, 3] = ( + 2 * 0.5 * k2 * (2 * k1 * sy * js - cy * jc) + 0.5 * (k2 + hx * k1) * sx * sy + ) + T[..., 3, 1, 2] = ( + 2 * 0.5 * k1 * k2 * (2 * cy * jd - sy * js) + 0.5 * (k2 + hx * k1) * dx * cy + ) + T[..., 3, 1, 3] = ( + 2 * 0.5 * k2 * (2 * k1 * sy * jd - cy * js) + 0.5 * (k2 + hx * k1) * dx * sy + ) + T[..., 3, 2, 5] = ( + 2 * 0.5 * hx / beta * k1 * k2 * (2 * cy * jf - sy * jd) + + 0.5 * hx / beta * (k2 + hx * k1) * j1 * cy + + 0.25 / beta * k1 * (sy - length * cy) + ) + T[..., 3, 3, 5] = ( + 2 * 0.5 * hx / beta * k2 * (2 * k1 * sy * jf - cy * jd) + + 0.5 * hx / beta * (k2 + hx * k1) * j1 * sy + - 0.25 / beta * k1 * length * sy + ) + T[..., 4, 0, 0] = -1 * hx / 12 / beta * khk * ( + sx * dx + 3 * j1 + ) - 0.25 / beta * k1 * (length - sx * cx) + T[..., 4, 0, 1] = -2 * hx / 12 / beta * khk * dx**2 + 0.25 / beta * k1 * sx**2 + T[..., 4, 1, 1] = ( + -1 * hx / 6 / beta * khk * j2 + - 0.5 / beta * sx + - 0.25 / beta * k1 * (j1 - sx * dx) + ) + T[..., 4, 0, 5] = ( + -2 * hx**2 / 12 / beta**2 * khk * (3 * dx * j1 - 4 * j2) + + 0.25 * hx / beta**2 * k1 * j1 * (1 + cx) + + 0.5 * hx / beta**2 * igamma2 * sx + ) + T[..., 4, 1, 5] = ( + -2 * hx**2 / 12 / beta**2 * khk * (dx * dx**2 - 2 * sx * j2) + + 0.25 * hx / beta**2 * k1 * sx * j1 + + 0.5 * hx / beta**2 * igamma2 * dx + ) + T[..., 4, 5, 5] = ( + -1 * hx**3 / 6 / beta**3 * khk * (3 * j3 - 2 * dx * j2) + + hx**2 / 6 / beta**3 * k1 * (sx * dx**2 - j2 * (1 + 2 * cx)) + + 1.5 / beta**3 * igamma2 * (hx**2 * j1 - length) + ) + T[..., 4, 2, 2] = ( + -1 * -hx / beta * k1 * k2 * jf + - 0.5 * hx / beta * (k2 + hx * k1) * j1 + + 0.25 / beta * k1 * (length - cy * sy) + ) + T[..., 4, 2, 3] = -2 * -0.5 * hx / beta * k2 * jd - 0.25 / beta * k1 * sy**2 + T[..., 4, 3, 3] = ( + -1 * -hx / beta * k2 * jf + + 0.5 * hx**2 / beta * j1 + - 0.25 / beta * (length + cy * sy) + ) + + # Rotate the R matrix for skew / vertical magnets + if torch.any(tilt != 0): + T = torch.einsum( + "...ij,...jk,...kl->...il", rotation_matrix(-tilt), T, rotation_matrix(tilt) + ) + return T + + def rotation_matrix(angle: torch.Tensor) -> torch.Tensor: """Rotate the transfer map in x-y plane. From 747ceb1219480f9a53a200c11263a1ca2d4d7ede Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 17:10:56 +0200 Subject: [PATCH 05/30] Fix sextupole into working state --- cheetah/accelerator/sextupole.py | 82 +++++++++++++++++--------------- cheetah/track_methods.py | 19 ++++---- tests/test_sextupole.py | 13 +++-- 3 files changed, 62 insertions(+), 52 deletions(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index 2889a7f81..0a2df76de 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -3,7 +3,8 @@ from cheetah.accelerator.element import Element from cheetah.particles import Beam, ParameterBeam, ParticleBeam, Species -from cheetah.utils import compute_relativistic_factors, verify_device_and_dtype +from cheetah.track_methods import base_rmatrix, base_tmatrix, misalignment_matrix +from cheetah.utils import verify_device_and_dtype class Sextupole(Element): @@ -48,19 +49,21 @@ def __init__( ) def transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor: - device = self.length.device - dtype = self.length.dtype - - _, igamma2, beta = compute_relativistic_factors(energy, species.mass_eV) - - vector_shape = torch.broadcast_shapes(self.length.shape, igamma2.shape) - - tm = torch.eye(7, device=device, dtype=dtype).repeat((*vector_shape, 1, 1)) - tm[..., 0, 1] = self.length - tm[..., 2, 3] = self.length - tm[..., 4, 5] = -self.length / beta**2 * igamma2 + 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, + ) - return tm + if torch.all(self.misalignment == 0): + return R + else: + R_entry, R_exit = misalignment_matrix(self.misalignment) + R = torch.einsum("...ij,...jk,...kl->...il", R_exit, R, R_entry) + return R def track(self, incoming: Beam) -> Beam: """ @@ -70,36 +73,37 @@ def track(self, incoming: Beam) -> Beam: :return: Beam exiting the element. """ first_order_tm = self.transfer_map(incoming.energy, incoming.species) - - second_order_tm = torch.eye( - 7, device=self.length.device, dtype=self.length.dtype - ).repeat((*incoming.mu_x.shape, 1, 1)) - second_order_tm[..., 0, 1] = self.length - second_order_tm[..., 2, 3] = self.length - second_order_tm[..., 4, 5] = ( - -self.length - / incoming.relativistic_beta**2 - * incoming.relativistic_gamma**2 + second_order_tm = base_tmatrix( + 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, ) - second_order_tm[..., 1, 0] = self.k2 * self.length**3 / 6 - second_order_tm[..., 3, 2] = self.k2 * self.length**3 / 6 - second_order_tm[..., 5, 4] = -self.k2 * self.length**3 / 6 - second_order_tm[..., 1, 2] = self.k2 * self.length**3 / 6 - second_order_tm[..., 3, 4] = self.k2 * self.length**3 / 6 - second_order_tm[..., 5, 0] = -self.k2 * self.length**3 / 6 - second_order_tm[..., 5, 2] = -self.k2 * self.length**3 / 6 - second_order_tm[..., 1, 4] = -self.k2 * self.length**3 / 6 - second_order_tm[..., 3, 0] = -self.k2 * self.length**3 / 6 - second_order_tm[..., 5, 3] = -self.k2 * self.length**3 / 6 - second_order_tm[..., 1, 5] = -self.k2 * self.length**3 / 6 - second_order_tm[..., 3, 2] = -self.k2 * self.length**3 / 6 - second_order_tm[..., 5, 4] = -self.k2 * self.length**3 / 6 # Apply the transfer map to the incoming beam - particles = torch.einsum("ijk,ikl->ijl", first_order_tm, incoming.particles) - particles = torch.einsum("ijk,ikl->ijl", second_order_tm, particles) + first_order_particles = torch.matmul( + incoming.particles, first_order_tm.transpose(-2, -1) + ) + second_order_particles = torch.einsum( + "...ijk,...j,...k->...i", + second_order_tm, + incoming.particles, + incoming.particles, + ) + outgoing_particles = second_order_particles + first_order_particles + from icecream import ic + + ic( + first_order_particles.shape, + second_order_particles.shape, + outgoing_particles.shape, + ) + return ParticleBeam( - particles=particles, + particles=outgoing_particles, energy=incoming.energy, particle_charges=incoming.particle_charges, survival_probabilities=incoming.survival_probabilities, diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 5c32874b6..59b85eb0b 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -19,7 +19,7 @@ def base_rmatrix( :param length: Length of the element in m. :param k1: Quadrupole strength in 1/m**2. - :param hx: Curvature (1/radius) of the element in 1/m**2. TODO: Unit correct? + :param hx: Curvature (1/radius) of the element in 1/m. :param species: Particle species of the beam. :param tilt: Roation of the element relative to the longitudinal axis in rad. :param energy: Beam energy in eV. @@ -94,7 +94,7 @@ def base_tmatrix( :param length: Length of the element in m. :param k1: Quadrupole strength in 1/m**2. :param k2: Sextupole strength in 1/m**3. - :param hx: Curvature (1/radius) of the element in 1/m**2. + :param hx: Curvature (1/radius) of the element in 1/m. :param species: Particle species of the beam. :param tilt: Roation of the element relative to the longitudinal axis in rad. :param energy: Beam energy in eV. @@ -110,19 +110,15 @@ def base_tmatrix( _, igamma2, beta = compute_relativistic_factors(energy, species.mass_eV) - # Avoid division by zero - k1 = k1.clone() - k1[k1 == 0] = 1e-12 - kx2 = k1 + hx**2 ky2 = -k1 kx = torch.sqrt(torch.complex(kx2, torch.tensor(0.0, device=device, dtype=dtype))) ky = torch.sqrt(torch.complex(ky2, torch.tensor(0.0, device=device, dtype=dtype))) cx = torch.cos(kx * length).real cy = torch.cos(ky * length).real - sy = (torch.sin(ky * length) / ky).real - sx = (torch.sin(kx * length) / kx).real - dx = hx / kx2 * (1.0 - cx) + sx = torch.where(kx != 0, (torch.sin(kx * length) / kx).real, length) + sy = torch.where(ky != 0, (torch.sin(ky * length) / ky).real, length) + dx = torch.where(kx != 0, (1.0 - cx) / kx2, length**2 / 2.0) d2y = 0.5 * sy**2 s2y = sy * cy @@ -156,7 +152,9 @@ def base_tmatrix( length.shape, k1.shape, hx.shape, tilt.shape, energy.shape ) - T = torch.zeros((7, 7, 7), dtype=dtype, device=device).repeat(*vector_shape, 1, 1) + T = torch.zeros((7, 7, 7), dtype=dtype, device=device).repeat( + *vector_shape, 1, 1, 1 + ) T[..., 0, 0, 0] = -1 / 6 * khk * (sx**2 + dx) - 0.5 * hx * kx2 * sx**2 T[..., 0, 0, 1] = 2 * -1 / 6 * khk * sx * dx + 0.5 * hx * sx * cx T[..., 0, 1, 1] = -1 / 6 * khk * dx**2 + 0.5 * hx * dx * cx @@ -272,6 +270,7 @@ def base_tmatrix( + 0.5 * hx**2 / beta * j1 - 0.25 / beta * (length + cy * sy) ) + T[..., 6, 6, 6] = 0.0 # Constant term currently handled by first order transfer map # Rotate the R matrix for skew / vertical magnets if torch.any(tilt != 0): diff --git a/tests/test_sextupole.py b/tests/test_sextupole.py index f23518b40..a5d9a17ca 100644 --- a/tests/test_sextupole.py +++ b/tests/test_sextupole.py @@ -25,13 +25,18 @@ def test_compare_sextupole_to_ocelot(): incoming_p_array = ocelot.astraBeam2particleArray( "tests/resources/ACHIP_EA1_2021.1351.001" ) - lattice = ocelot.MagneticLattice([ocelot.Sextupole(l=length, k2=k2, tilt=tilt)]) + lattice = ocelot.MagneticLattice( + [ocelot.Sextupole(l=length, k2=k2, tilt=tilt)], + method={"global": ocelot.SecondTM}, + ) navigator = ocelot.Navigator(lattice) _, outgoing_p_array = ocelot.track(lattice, deepcopy(incoming_p_array), navigator) outgoing_ocelot = cheetah.ParticleBeam.from_ocelot(outgoing_p_array) # Compare the results - assert torch.allclose(outgoing_cheetah.particles, outgoing_ocelot.particles) + assert torch.allclose( + outgoing_cheetah.particles, outgoing_ocelot.particles, atol=1e-5, rtol=1e-6 + ) def test_sextupole_as_drift(): @@ -48,4 +53,6 @@ def test_sextupole_as_drift(): drift_outgoing = drift.track(incoming) # Check that the results are the same - assert torch.allclose(sextupole_outgoing.particles, drift_outgoing.particles) + assert torch.allclose( + sextupole_outgoing.particles, drift_outgoing.particles, atol=1e-5, rtol=1e-6 + ) From f8be117f08f6ab16ffa4f4afe8ae130120a6d4f9 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 17:22:47 +0200 Subject: [PATCH 06/30] Fix test failures resulting from minor oversights --- cheetah/accelerator/sextupole.py | 9 +-------- tests/test_elements.py | 1 + 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index 0a2df76de..3a2dbed94 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -94,13 +94,6 @@ def track(self, incoming: Beam) -> Beam: incoming.particles, ) outgoing_particles = second_order_particles + first_order_particles - from icecream import ic - - ic( - first_order_particles.shape, - second_order_particles.shape, - outgoing_particles.shape, - ) return ParticleBeam( particles=outgoing_particles, @@ -116,7 +109,7 @@ def is_skippable(self) -> bool: @property def is_active(self) -> bool: - return torch.any(self.k2 != 0.0) + return torch.any(self.k2 != 0.0).item() def split(self, resolution: torch.Tensor) -> list[Element]: raise NotImplementedError diff --git a/tests/test_elements.py b/tests/test_elements.py index 6e5fb3992..771a8bb72 100644 --- a/tests/test_elements.py +++ b/tests/test_elements.py @@ -11,6 +11,7 @@ cheetah.HorizontalCorrector: {"length": torch.tensor(1.0)}, cheetah.Quadrupole: {"length": torch.tensor(1.0)}, cheetah.Segment: {"elements": [cheetah.Drift(length=torch.tensor(1.0))]}, + cheetah.Sextupole: {"length": torch.tensor(1.0)}, cheetah.Solenoid: {"length": torch.tensor(1.0)}, cheetah.SpaceChargeKick: {"effect_length": torch.tensor(1.0)}, cheetah.TransverseDeflectingCavity: {"length": torch.tensor(1.0)}, From 9d51d004f1bb919857ef6affbbac31b30c5691e7 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 17:31:03 +0200 Subject: [PATCH 07/30] First draft of `ParameterBeam` implementation for `Sextupole` --- cheetah/accelerator/sextupole.py | 76 ++++++++++++++++++++++++-------- 1 file changed, 57 insertions(+), 19 deletions(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index 3a2dbed94..5bd0ec552 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -83,25 +83,63 @@ def track(self, incoming: Beam) -> Beam: energy=incoming.energy, ) - # Apply the transfer map to the incoming beam - first_order_particles = torch.matmul( - incoming.particles, first_order_tm.transpose(-2, -1) - ) - second_order_particles = torch.einsum( - "...ijk,...j,...k->...i", - second_order_tm, - 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, - ) + if isinstance(incoming, ParameterBeam): + # Apply the transfer map to the incoming particles + first_order_mu = torch.matmul( + first_order_tm, incoming.mu.unsqueeze(-1) + ).squeeze(-1) + second_order_mu = torch.einsum( + "...ijk,...j,...k->...i", + second_order_tm, + incoming.mu, + incoming.mu, + ) + outgoing_mu = first_order_mu + second_order_mu + first_order_cov = torch.matmul( + first_order_tm, + torch.matmul(incoming.cov, first_order_tm.transpose(-2, -1)), + ) + second_order_cov = torch.einsum( + "...ijk,...jkl,...l->...i", + second_order_tm, + incoming.cov, + incoming.mu, + ) + outgoing_cov = first_order_cov + second_order_cov + + return ParameterBeam( + mu=outgoing_mu, + cov=outgoing_cov, + energy=incoming.energy, + total_charge=incoming.total_charge, + survival_probabilities=incoming.survival_probabilities, + species=incoming.species, + ) + elif isinstance(incoming, ParticleBeam): + # Apply the transfer map to the incoming particles + first_order_particles = torch.matmul( + incoming.particles, first_order_tm.transpose(-2, -1) + ) + second_order_particles = torch.einsum( + "...ijk,...j,...k->...i", + second_order_tm, + 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: From 00da6604444cb781d6df8e510c985bd7a45ce949 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 18:12:33 +0200 Subject: [PATCH 08/30] Add a test for `ParameterBeam` tracking through a `Sextupole` --- tests/test_sextupole.py | 46 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/test_sextupole.py b/tests/test_sextupole.py index a5d9a17ca..309413523 100644 --- a/tests/test_sextupole.py +++ b/tests/test_sextupole.py @@ -56,3 +56,49 @@ def test_sextupole_as_drift(): assert torch.allclose( sextupole_outgoing.particles, drift_outgoing.particles, atol=1e-5, rtol=1e-6 ) + + +def test_sextupole_parameter_beam_particle_beam_agreement(): + """ + Test that the results of tracking an `ParameterBeam` and a `ParticleBeam` through a + sextupole agree. + """ + # Create a sextupole + length = 0.34 + k2 = 0.5 + tilt = 0.1 + sextupole = cheetah.Sextupole( + length=torch.tensor(length), k2=torch.tensor(k2), tilt=torch.tensor(tilt) + ) + + # Create an incoming ParticleBeam + incoming_particle_beam = cheetah.ParticleBeam.from_astra( + "tests/resources/ACHIP_EA1_2021.1351.001" + ) + + # Create an incoming ParameterBeam + incoming_parameter_beam = cheetah.ParameterBeam.from_astra( + "tests/resources/ACHIP_EA1_2021.1351.001" + ) + + # Track through the sextupole + outgoing_particle_beam = sextupole.track(incoming_particle_beam) + outgoing_parameter_beam = sextupole.track(incoming_parameter_beam) + + outgoing_particle_beam_as_parameter_beam = ( + outgoing_particle_beam.as_parameter_beam() + ) + + # Check that the results are the same + assert torch.allclose( + outgoing_particle_beam_as_parameter_beam.mu, + outgoing_parameter_beam.mu, + atol=1e-5, + rtol=1e-6, + ) + assert torch.allclose( + outgoing_particle_beam_as_parameter_beam.cov, + outgoing_parameter_beam.cov, + atol=1e-5, + rtol=1e-6, + ) From 96ee1a042271ac041bd39d26f4f803bfa776631b Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Tue, 22 Apr 2025 18:14:55 +0200 Subject: [PATCH 09/30] Add changelog entry --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e5a7f0341..3830aa41d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ### 🚀 Features - Implement `split` method for the `Solenoid` element (see #380) (@cr-xu) +- Add a `Sextupole` element (see #406) (@jank324, @Hespe) ### 🐛 Bug fixes From 2951aa0780bbf757016ef731679fef3c40a25a7f Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Wed, 23 Apr 2025 14:05:39 +0200 Subject: [PATCH 10/30] Implement `ParameterBeam` in `Sextupole` with first order effects only --- cheetah/accelerator/sextupole.py | 33 ++------------------------------ 1 file changed, 2 insertions(+), 31 deletions(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index 5bd0ec552..be35b51ac 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -84,37 +84,8 @@ def track(self, incoming: Beam) -> Beam: ) if isinstance(incoming, ParameterBeam): - # Apply the transfer map to the incoming particles - first_order_mu = torch.matmul( - first_order_tm, incoming.mu.unsqueeze(-1) - ).squeeze(-1) - second_order_mu = torch.einsum( - "...ijk,...j,...k->...i", - second_order_tm, - incoming.mu, - incoming.mu, - ) - outgoing_mu = first_order_mu + second_order_mu - first_order_cov = torch.matmul( - first_order_tm, - torch.matmul(incoming.cov, first_order_tm.transpose(-2, -1)), - ) - second_order_cov = torch.einsum( - "...ijk,...jkl,...l->...i", - second_order_tm, - incoming.cov, - incoming.mu, - ) - outgoing_cov = first_order_cov + second_order_cov - - return ParameterBeam( - mu=outgoing_mu, - cov=outgoing_cov, - energy=incoming.energy, - total_charge=incoming.total_charge, - survival_probabilities=incoming.survival_probabilities, - species=incoming.species, - ) + # 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 = torch.matmul( From 214fc7c692faff2c6c3cad974eaa467dcf647092 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 11:43:07 +0200 Subject: [PATCH 11/30] Rename method for computing T to reflect that T is a tensor --- cheetah/accelerator/sextupole.py | 4 ++-- cheetah/track_methods.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index be35b51ac..fd9949c71 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -3,7 +3,7 @@ from cheetah.accelerator.element import Element from cheetah.particles import Beam, ParameterBeam, ParticleBeam, Species -from cheetah.track_methods import base_rmatrix, base_tmatrix, misalignment_matrix +from cheetah.track_methods import base_rmatrix, base_ttensor, misalignment_matrix from cheetah.utils import verify_device_and_dtype @@ -73,7 +73,7 @@ def track(self, incoming: Beam) -> Beam: :return: Beam exiting the element. """ first_order_tm = self.transfer_map(incoming.energy, incoming.species) - second_order_tm = base_tmatrix( + second_order_tm = base_ttensor( length=self.length, k1=torch.zeros_like(self.length), k2=self.k2, diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 59b85eb0b..f8b18c578 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -79,7 +79,7 @@ def base_rmatrix( return R -def base_tmatrix( +def base_ttensor( length: torch.Tensor, k1: torch.Tensor, k2: torch.Tensor, From 142fca71d5061f2d5941dce4b04ea206fade88ce Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Fri, 25 Apr 2025 12:17:30 +0200 Subject: [PATCH 12/30] Add vectorization and dtype tests for sextupole --- tests/test_device_dtype.py | 2 ++ tests/test_vectorized.py | 1 + 2 files changed, 3 insertions(+) diff --git a/tests/test_device_dtype.py b/tests/test_device_dtype.py index b293b6bc1..ba6868ccf 100644 --- a/tests/test_device_dtype.py +++ b/tests/test_device_dtype.py @@ -53,6 +53,7 @@ def test_move_quadrupole_to_device(target_device: torch.device): cheetah.HorizontalCorrector, cheetah.Quadrupole, cheetah.RBend, + cheetah.Sextupole, cheetah.Solenoid, cheetah.TransverseDeflectingCavity, cheetah.Undulator, @@ -80,6 +81,7 @@ def test_forced_element_dtype(ElementClass): cheetah.HorizontalCorrector, cheetah.Quadrupole, cheetah.RBend, + cheetah.Sextupole, cheetah.Solenoid, cheetah.TransverseDeflectingCavity, cheetah.Undulator, diff --git a/tests/test_vectorized.py b/tests/test_vectorized.py index 1458eae09..3b668ff44 100644 --- a/tests/test_vectorized.py +++ b/tests/test_vectorized.py @@ -342,6 +342,7 @@ def test_vectorized_screen_2d(BeamClass, method): cheetah.HorizontalCorrector, cheetah.Quadrupole, cheetah.RBend, + cheetah.Sextupole, cheetah.Solenoid, cheetah.TransverseDeflectingCavity, cheetah.Undulator, From 2061e241ec4765397c428bdb11087f91fab219d7 Mon Sep 17 00:00:00 2001 From: Christian Hespe Date: Fri, 25 Apr 2025 12:34:36 +0200 Subject: [PATCH 13/30] Add clone test for sextupole --- tests/test_clone.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_clone.py b/tests/test_clone.py index 7d12d0452..1c3382aec 100644 --- a/tests/test_clone.py +++ b/tests/test_clone.py @@ -13,6 +13,7 @@ cheetah.HorizontalCorrector, cheetah.Quadrupole, cheetah.RBend, + cheetah.Sextupole, cheetah.Solenoid, cheetah.TransverseDeflectingCavity, cheetah.Undulator, From 7a1c5fd232aed7c0a627001ebe8380f3038a7ef4 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 12:55:48 +0200 Subject: [PATCH 14/30] Slightly cleaer docstring for `rotation_matrix` method --- cheetah/track_methods.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index f8b18c578..01c2ec7de 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -281,7 +281,8 @@ def base_ttensor( def rotation_matrix(angle: torch.Tensor) -> torch.Tensor: - """Rotate the transfer map in x-y plane. + """ + Rotate the coordinate system in the x-y plane. :param angle: Rotation angle in rad, for example `angle = np.pi/2` for vertical = dipole. From 23a61253653a1b8cea292e17230db23908cc260c Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 13:02:57 +0200 Subject: [PATCH 15/30] Address (mostly) the comments from Copilot review --- cheetah/accelerator/quadrupole.py | 2 +- cheetah/accelerator/sextupole.py | 6 +++--- tests/test_sextupole.py | 6 +++++- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/cheetah/accelerator/quadrupole.py b/cheetah/accelerator/quadrupole.py index 3c58a0557..413a3047e 100644 --- a/cheetah/accelerator/quadrupole.py +++ b/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"`. diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index fd9949c71..9cd1a8f79 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -12,9 +12,9 @@ class Sextupole(Element): A sextupole element in a particle accelerator. :param length: Length in meters. - :param k2: TODO - :param misalignment: TODO - :param tilt: TODO + :param k2: Sextupole strength in TODO What is the unit? + :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. """ diff --git a/tests/test_sextupole.py b/tests/test_sextupole.py index 309413523..1dbf7b016 100644 --- a/tests/test_sextupole.py +++ b/tests/test_sextupole.py @@ -67,8 +67,12 @@ def test_sextupole_parameter_beam_particle_beam_agreement(): length = 0.34 k2 = 0.5 tilt = 0.1 + misalignment = (1e-4, 2e-4) sextupole = cheetah.Sextupole( - length=torch.tensor(length), k2=torch.tensor(k2), tilt=torch.tensor(tilt) + length=torch.tensor(length), + k2=torch.tensor(k2), + tilt=torch.tensor(tilt), + misalignment=torch.tensor(misalignment), ) # Create an incoming ParticleBeam From 9ebf1b7deaa183a9ca1f7992b8a9b0dcdbb4e5a0 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 13:04:42 +0200 Subject: [PATCH 16/30] Add unit to docstring --- cheetah/accelerator/sextupole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index 9cd1a8f79..52c058bc0 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -12,7 +12,7 @@ class Sextupole(Element): A sextupole element in a particle accelerator. :param length: Length in meters. - :param k2: Sextupole strength in TODO What is the unit? + :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. From d27f71069826ff38c3dfcc66b50acc513ee6cce9 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 13:14:18 +0200 Subject: [PATCH 17/30] Fix precedence that I ignored when originally combining the two steps from Ocelot --- cheetah/track_methods.py | 109 +++++++++++++++++++-------------------- 1 file changed, 54 insertions(+), 55 deletions(-) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 01c2ec7de..70b15d7fc 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -156,15 +156,15 @@ def base_ttensor( *vector_shape, 1, 1, 1 ) T[..., 0, 0, 0] = -1 / 6 * khk * (sx**2 + dx) - 0.5 * hx * kx2 * sx**2 - T[..., 0, 0, 1] = 2 * -1 / 6 * khk * sx * dx + 0.5 * hx * sx * cx + T[..., 0, 0, 1] = 2 * (-1 / 6 * khk * sx * dx + 0.5 * hx * sx * cx) T[..., 0, 1, 1] = -1 / 6 * khk * dx**2 + 0.5 * hx * dx * cx - T[..., 0, 0, 5] = ( - 2 * -hx / 12 / beta * khk * (3 * sx * j1 - dx**2) + T[..., 0, 0, 5] = 2 * ( + -hx / 12 / beta * khk * (3 * sx * j1 - dx**2) + 0.5 * hx**2 / beta * sx**2 + 0.25 / beta * k1 * length * sx ) - T[..., 0, 1, 5] = ( - 2 * -hx / 12 / beta * khk * (sx * dx**2 - 2 * cx * j2) + T[..., 0, 1, 5] = 2 * ( + -hx / 12 / beta * khk * (sx * dx**2 - 2 * cx * j2) + 0.25 * hx**2 / beta * (sx * dx + cx * j1) - 0.25 / beta * (sx + length * cx) ) @@ -175,17 +175,17 @@ def base_ttensor( - 0.5 * hx / (beta**2) * igamma2 * dx ) T[..., 0, 2, 2] = k1 * k2 * jd + 0.5 * (k2 + hx * k1) * dx - T[..., 0, 2, 3] = 2 * 0.5 * k2 * js + T[..., 0, 2, 3] = 2 * (0.5 * k2 * js) T[..., 0, 3, 3] = k2 * jd - 0.5 * hx * dx T[..., 1, 0, 0] = -1 / 6 * khk * sx * (1 + 2 * cx) - T[..., 1, 0, 1] = 2 * -1 / 6 * khk * dx * (1 + 2 * cx) + T[..., 1, 0, 1] = 2 * (-1 / 6 * khk * dx * (1 + 2 * cx)) T[..., 1, 1, 1] = -1 / 3 * khk * sx * dx - 0.5 * hx * sx - T[..., 1, 0, 5] = 2 * -hx / 12 / beta * khk * ( - 3 * cx * j1 + sx * dx - ) - 0.25 / beta * k1 * (sx - length * cx) - T[..., 1, 1, 5] = ( - 2 * -hx / 12 / beta * khk * (3 * sx * j1 + dx**2) - + 0.25 / beta * k1 * length * sx + T[..., 1, 0, 5] = 2 * ( + -hx / 12 / beta * khk * (3 * cx * j1 + sx * dx) + - 0.25 / beta * k1 * (sx - length * cx) + ) + T[..., 1, 1, 5] = 2 * ( + -hx / 12 / beta * khk * (3 * sx * j1 + dx**2) + 0.25 / beta * k1 * length * sx ) T[..., 1, 5, 5] = ( -(hx**2) / 6 / beta**2 * khk * (sx * dx**2 - 2 * cx * j2) @@ -193,80 +193,79 @@ def base_ttensor( - 0.5 * hx / beta**2 * igamma2 * sx ) T[..., 1, 2, 2] = k1 * k2 * js + 0.5 * (k2 + hx * k1) * sx - T[..., 1, 2, 3] = 2 * 0.5 * k2 * jc + T[..., 1, 2, 3] = 2 * (0.5 * k2 * jc) T[..., 1, 3, 3] = k2 * js - 0.5 * hx * sx - T[..., 2, 0, 2] = ( - 2 * 0.5 * k2 * (cy * jc - 2 * k1 * sy * js) + 0.5 * hx * k1 * sx * sy + T[..., 2, 0, 2] = 2 * ( + 0.5 * k2 * (cy * jc - 2 * k1 * sy * js) + 0.5 * hx * k1 * sx * sy ) - T[..., 2, 0, 3] = 2 * 0.5 * k2 * (sy * jc - 2 * cy * js) + 0.5 * hx * sx * cy - T[..., 2, 1, 2] = ( - 2 * 0.5 * k2 * (cy * js - 2 * k1 * sy * jd) + 0.5 * hx * k1 * dx * sy + T[..., 2, 0, 3] = 2 * (0.5 * k2 * (sy * jc - 2 * cy * js) + 0.5 * hx * sx * cy) + T[..., 2, 1, 2] = 2 * ( + 0.5 * k2 * (cy * js - 2 * k1 * sy * jd) + 0.5 * hx * k1 * dx * sy ) - T[..., 2, 1, 3] = 2 * 0.5 * k2 * (sy * js - 2 * cy * jd) + 0.5 * hx * dx * cy - T[..., 2, 2, 5] = ( - 2 * 0.5 * hx / beta * k2 * (cy * jd - 2 * k1 * sy * jf) + T[..., 2, 1, 3] = 2 * (0.5 * k2 * (sy * js - 2 * cy * jd) + 0.5 * hx * dx * cy) + T[..., 2, 2, 5] = 2 * ( + 0.5 * hx / beta * k2 * (cy * jd - 2 * k1 * sy * jf) + 0.5 * hx**2 / beta * k1 * j1 * sy - 0.25 / beta * k1 * length * sy ) - T[..., 2, 3, 5] = ( - 2 * 0.5 * hx / beta * k2 * (sy * jd - 2 * cy * jf) + T[..., 2, 3, 5] = 2 * ( + 0.5 * hx / beta * k2 * (sy * jd - 2 * cy * jf) + 0.5 * hx**2 / beta * j1 * cy - 0.25 / beta * (sy + length * cy) ) - T[..., 3, 0, 2] = ( - 2 * 0.5 * k1 * k2 * (2 * cy * js - sy * jc) + 0.5 * (k2 + hx * k1) * sx * cy + T[..., 3, 0, 2] = 2 * ( + 0.5 * k1 * k2 * (2 * cy * js - sy * jc) + 0.5 * (k2 + hx * k1) * sx * cy ) - T[..., 3, 0, 3] = ( - 2 * 0.5 * k2 * (2 * k1 * sy * js - cy * jc) + 0.5 * (k2 + hx * k1) * sx * sy + T[..., 3, 0, 3] = 2 * ( + 0.5 * k2 * (2 * k1 * sy * js - cy * jc) + 0.5 * (k2 + hx * k1) * sx * sy ) - T[..., 3, 1, 2] = ( - 2 * 0.5 * k1 * k2 * (2 * cy * jd - sy * js) + 0.5 * (k2 + hx * k1) * dx * cy + T[..., 3, 1, 2] = 2 * ( + 0.5 * k1 * k2 * (2 * cy * jd - sy * js) + 0.5 * (k2 + hx * k1) * dx * cy ) - T[..., 3, 1, 3] = ( - 2 * 0.5 * k2 * (2 * k1 * sy * jd - cy * js) + 0.5 * (k2 + hx * k1) * dx * sy + T[..., 3, 1, 3] = 2 * ( + 0.5 * k2 * (2 * k1 * sy * jd - cy * js) + 0.5 * (k2 + hx * k1) * dx * sy ) - T[..., 3, 2, 5] = ( - 2 * 0.5 * hx / beta * k1 * k2 * (2 * cy * jf - sy * jd) + T[..., 3, 2, 5] = 2 * ( + 0.5 * hx / beta * k1 * k2 * (2 * cy * jf - sy * jd) + 0.5 * hx / beta * (k2 + hx * k1) * j1 * cy + 0.25 / beta * k1 * (sy - length * cy) ) - T[..., 3, 3, 5] = ( - 2 * 0.5 * hx / beta * k2 * (2 * k1 * sy * jf - cy * jd) + T[..., 3, 3, 5] = 2 * ( + 0.5 * hx / beta * k2 * (2 * k1 * sy * jf - cy * jd) + 0.5 * hx / beta * (k2 + hx * k1) * j1 * sy - 0.25 / beta * k1 * length * sy ) - T[..., 4, 0, 0] = -1 * hx / 12 / beta * khk * ( - sx * dx + 3 * j1 - ) - 0.25 / beta * k1 * (length - sx * cx) - T[..., 4, 0, 1] = -2 * hx / 12 / beta * khk * dx**2 + 0.25 / beta * k1 * sx**2 - T[..., 4, 1, 1] = ( - -1 * hx / 6 / beta * khk * j2 - - 0.5 / beta * sx - - 0.25 / beta * k1 * (j1 - sx * dx) + T[..., 4, 0, 0] = -1 * ( + hx / 12 / beta * khk * (sx * dx + 3 * j1) + - 0.25 / beta * k1 * (length - sx * cx) + ) + T[..., 4, 0, 1] = -2 * (hx / 12 / beta * khk * dx**2 + 0.25 / beta * k1 * sx**2) + T[..., 4, 1, 1] = -1 * ( + hx / 6 / beta * khk * j2 - 0.5 / beta * sx - 0.25 / beta * k1 * (j1 - sx * dx) ) - T[..., 4, 0, 5] = ( - -2 * hx**2 / 12 / beta**2 * khk * (3 * dx * j1 - 4 * j2) + T[..., 4, 0, 5] = -2 * ( + hx**2 / 12 / beta**2 * khk * (3 * dx * j1 - 4 * j2) + 0.25 * hx / beta**2 * k1 * j1 * (1 + cx) + 0.5 * hx / beta**2 * igamma2 * sx ) - T[..., 4, 1, 5] = ( - -2 * hx**2 / 12 / beta**2 * khk * (dx * dx**2 - 2 * sx * j2) + T[..., 4, 1, 5] = -2 * ( + hx**2 / 12 / beta**2 * khk * (dx * dx**2 - 2 * sx * j2) + 0.25 * hx / beta**2 * k1 * sx * j1 + 0.5 * hx / beta**2 * igamma2 * dx ) - T[..., 4, 5, 5] = ( - -1 * hx**3 / 6 / beta**3 * khk * (3 * j3 - 2 * dx * j2) + T[..., 4, 5, 5] = -1 * ( + hx**3 / 6 / beta**3 * khk * (3 * j3 - 2 * dx * j2) + hx**2 / 6 / beta**3 * k1 * (sx * dx**2 - j2 * (1 + 2 * cx)) + 1.5 / beta**3 * igamma2 * (hx**2 * j1 - length) ) - T[..., 4, 2, 2] = ( - -1 * -hx / beta * k1 * k2 * jf + T[..., 4, 2, 2] = -1 * ( + -hx / beta * k1 * k2 * jf - 0.5 * hx / beta * (k2 + hx * k1) * j1 + 0.25 / beta * k1 * (length - cy * sy) ) - T[..., 4, 2, 3] = -2 * -0.5 * hx / beta * k2 * jd - 0.25 / beta * k1 * sy**2 - T[..., 4, 3, 3] = ( - -1 * -hx / beta * k2 * jf + T[..., 4, 2, 3] = -2 * (-0.5 * hx / beta * k2 * jd - 0.25 / beta * k1 * sy**2) + T[..., 4, 3, 3] = -1 * ( + -hx / beta * k2 * jf + 0.5 * hx**2 / beta * j1 - 0.25 / beta * (length + cy * sy) ) From 49980f5892c74458c9494c5d6c752a3aa1d9c40b Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 14:19:54 +0200 Subject: [PATCH 18/30] Use realistic sextupole values in test similar to those in EuXFEL lattice --- tests/test_sextupole.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_sextupole.py b/tests/test_sextupole.py index 1dbf7b016..4eef07665 100644 --- a/tests/test_sextupole.py +++ b/tests/test_sextupole.py @@ -8,9 +8,9 @@ def test_compare_sextupole_to_ocelot(): """Compare the results of tracking through a sextupole in Cheetah and Ocelot.""" - length = 0.34 - k2 = 0.5 - tilt = 0.1 + length = 0.11 + k2 = 87.0 + tilt = torch.pi / 2 # Track through a sextupole in Cheetah incoming = cheetah.ParticleBeam.from_astra( @@ -45,8 +45,8 @@ def test_sextupole_as_drift(): "tests/resources/ACHIP_EA1_2021.1351.001" ) - sextupole = cheetah.Sextupole(length=torch.tensor(0.34), k2=torch.tensor(0.0)) - drift = cheetah.Drift(length=torch.tensor(0.34)) + sextupole = cheetah.Sextupole(length=torch.tensor(0.11), k2=torch.tensor(0.0)) + drift = cheetah.Drift(length=torch.tensor(0.11)) # Track through the sextupole and drift sextupole_outgoing = sextupole.track(incoming) @@ -64,9 +64,9 @@ def test_sextupole_parameter_beam_particle_beam_agreement(): sextupole agree. """ # Create a sextupole - length = 0.34 - k2 = 0.5 - tilt = 0.1 + length = 0.11 + k2 = 87.0 + tilt = torch.pi / 2 misalignment = (1e-4, 2e-4) sextupole = cheetah.Sextupole( length=torch.tensor(length), From 4650997bdb1af0651217b150e331792c1fcdd540 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 14:28:15 +0200 Subject: [PATCH 19/30] Address the fact the `ParameterBeam` and `ParticleBeam` comparison on `Sextupole` might have been phyiscally incorrect --- tests/test_sextupole.py | 92 ++++++++++++++++++----------------------- 1 file changed, 41 insertions(+), 51 deletions(-) diff --git a/tests/test_sextupole.py b/tests/test_sextupole.py index 4eef07665..57dfb6c89 100644 --- a/tests/test_sextupole.py +++ b/tests/test_sextupole.py @@ -6,8 +6,11 @@ import cheetah -def test_compare_sextupole_to_ocelot(): - """Compare the results of tracking through a sextupole in Cheetah and Ocelot.""" +def test_compare_sextupole_to_ocelot_particle(): + """ + Compare the results of tracking through a sextupole in Cheetah and Ocelot. For a + `ParticleBeam` with second order effects in Ocelot. + """ length = 0.11 k2 = 87.0 tilt = torch.pi / 2 @@ -39,70 +42,57 @@ def test_compare_sextupole_to_ocelot(): ) -def test_sextupole_as_drift(): - """Test that a sextupole with k2=0 is equivalent to a drift.""" - incoming = cheetah.ParticleBeam.from_astra( - "tests/resources/ACHIP_EA1_2021.1351.001" - ) - - sextupole = cheetah.Sextupole(length=torch.tensor(0.11), k2=torch.tensor(0.0)) - drift = cheetah.Drift(length=torch.tensor(0.11)) - - # Track through the sextupole and drift - sextupole_outgoing = sextupole.track(incoming) - drift_outgoing = drift.track(incoming) - - # Check that the results are the same - assert torch.allclose( - sextupole_outgoing.particles, drift_outgoing.particles, atol=1e-5, rtol=1e-6 - ) - - -def test_sextupole_parameter_beam_particle_beam_agreement(): +def test_compare_sextupole_to_ocelot_parameter(): """ - Test that the results of tracking an `ParameterBeam` and a `ParticleBeam` through a - sextupole agree. + Compare the results of tracking through a sextupole in Cheetah and Ocelot for a + `ParameterBeam` with only first order effects in Ocelot. """ - # Create a sextupole length = 0.11 k2 = 87.0 tilt = torch.pi / 2 - misalignment = (1e-4, 2e-4) - sextupole = cheetah.Sextupole( - length=torch.tensor(length), - k2=torch.tensor(k2), - tilt=torch.tensor(tilt), - misalignment=torch.tensor(misalignment), - ) - # Create an incoming ParticleBeam - incoming_particle_beam = cheetah.ParticleBeam.from_astra( + # Track through a sextupole in Cheetah + incoming = cheetah.ParameterBeam.from_astra( "tests/resources/ACHIP_EA1_2021.1351.001" ) + cheetah_sextupole = cheetah.Sextupole( + length=torch.tensor(length), k2=torch.tensor(k2), tilt=torch.tensor(tilt) + ) + outgoing_cheetah = cheetah_sextupole.track(incoming) - # Create an incoming ParameterBeam - incoming_parameter_beam = cheetah.ParameterBeam.from_astra( + # Convert to Ocelot sextupole + incoming_p_array = ocelot.astraBeam2particleArray( "tests/resources/ACHIP_EA1_2021.1351.001" ) + lattice = ocelot.MagneticLattice( + [ocelot.Sextupole(l=length, k2=k2, tilt=tilt)], + method={"global": ocelot.TransferMap}, + ) + navigator = ocelot.Navigator(lattice) + _, outgoing_p_array = ocelot.track(lattice, deepcopy(incoming_p_array), navigator) + outgoing_ocelot = cheetah.ParameterBeam.from_ocelot(outgoing_p_array) + + # Compare the results + assert torch.allclose(outgoing_cheetah.mu, outgoing_ocelot.mu, atol=1e-5, rtol=1e-6) + assert torch.allclose( + outgoing_cheetah.cov, outgoing_ocelot.cov, atol=1e-5, rtol=1e-6 + ) - # Track through the sextupole - outgoing_particle_beam = sextupole.track(incoming_particle_beam) - outgoing_parameter_beam = sextupole.track(incoming_parameter_beam) - outgoing_particle_beam_as_parameter_beam = ( - outgoing_particle_beam.as_parameter_beam() +def test_sextupole_as_drift(): + """Test that a sextupole with k2=0 is equivalent to a drift.""" + incoming = cheetah.ParticleBeam.from_astra( + "tests/resources/ACHIP_EA1_2021.1351.001" ) + sextupole = cheetah.Sextupole(length=torch.tensor(0.11), k2=torch.tensor(0.0)) + drift = cheetah.Drift(length=torch.tensor(0.11)) + + # Track through the sextupole and drift + sextupole_outgoing = sextupole.track(incoming) + drift_outgoing = drift.track(incoming) + # Check that the results are the same assert torch.allclose( - outgoing_particle_beam_as_parameter_beam.mu, - outgoing_parameter_beam.mu, - atol=1e-5, - rtol=1e-6, - ) - assert torch.allclose( - outgoing_particle_beam_as_parameter_beam.cov, - outgoing_parameter_beam.cov, - atol=1e-5, - rtol=1e-6, + sextupole_outgoing.particles, drift_outgoing.particles, atol=1e-5, rtol=1e-6 ) From b9e9d48646b8d1cf07397ee64f72ed4e41192050 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 14:30:18 +0200 Subject: [PATCH 20/30] Add docs entry for `Sextupole` --- docs/accelerator.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/accelerator.rst b/docs/accelerator.rst index 8cc0a737d..959728d07 100644 --- a/docs/accelerator.rst +++ b/docs/accelerator.rst @@ -55,6 +55,10 @@ Accelerator :members: :undoc-members: +.. automodule:: accelerator.sextupole + :members: + :undoc-members: + .. automodule:: accelerator.solenoid :members: :undoc-members: From 2fede4305997bd2ec1b1da885fff5fb5cc4d4f43 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 14:33:27 +0200 Subject: [PATCH 21/30] Add vectorised sextupole test --- tests/test_sextupole.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/test_sextupole.py b/tests/test_sextupole.py index 57dfb6c89..d3d07c156 100644 --- a/tests/test_sextupole.py +++ b/tests/test_sextupole.py @@ -42,6 +42,46 @@ def test_compare_sextupole_to_ocelot_particle(): ) +def test_compare_sextupole_to_ocelot_particle_vectorized(): + """ + Compare the results of tracking through a sextupole in Cheetah and Ocelot. For a + `ParticleBeam` with second order effects in Ocelot. + + Vectorised version of the test. + """ + length = 0.11 + k2 = 87.0 + tilt = torch.pi / 2 + + # Track through a sextupole in Cheetah + incoming = cheetah.ParticleBeam.from_astra( + "tests/resources/ACHIP_EA1_2021.1351.001" + ) + cheetah_sextupole = cheetah.Sextupole( + length=torch.tensor(length), + k2=torch.tensor(k2).repeat([3, 2]), + tilt=torch.tensor(tilt), + ) + outgoing_cheetah = cheetah_sextupole.track(incoming) + + # Convert to Ocelot sextupole + incoming_p_array = ocelot.astraBeam2particleArray( + "tests/resources/ACHIP_EA1_2021.1351.001" + ) + lattice = ocelot.MagneticLattice( + [ocelot.Sextupole(l=length, k2=k2, tilt=tilt)], + method={"global": ocelot.SecondTM}, + ) + navigator = ocelot.Navigator(lattice) + _, outgoing_p_array = ocelot.track(lattice, deepcopy(incoming_p_array), navigator) + outgoing_ocelot = cheetah.ParticleBeam.from_ocelot(outgoing_p_array) + + # Compare the results + assert torch.allclose( + outgoing_cheetah.particles, outgoing_ocelot.particles, atol=1e-5, rtol=1e-6 + ) + + def test_compare_sextupole_to_ocelot_parameter(): """ Compare the results of tracking through a sextupole in Cheetah and Ocelot for a From 3645c628100ce25e193e2d13c64217d55aa762ba Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 14:42:12 +0200 Subject: [PATCH 22/30] Fix bug discovered in sextupole vectorisation --- cheetah/accelerator/sextupole.py | 2 +- cheetah/track_methods.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index 52c058bc0..c8588c48c 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -93,7 +93,7 @@ def track(self, incoming: Beam) -> Beam: ) second_order_particles = torch.einsum( "...ijk,...j,...k->...i", - second_order_tm, + second_order_tm.unsqueeze(-4), # Add broadcast dimension for particles incoming.particles, incoming.particles, ) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 70b15d7fc..e45fd85ee 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -149,7 +149,7 @@ def base_ttensor( khk = k2 + 2 * hx * k1 vector_shape = torch.broadcast_shapes( - length.shape, k1.shape, hx.shape, tilt.shape, energy.shape + length.shape, k1.shape, k2.shape, hx.shape, tilt.shape, energy.shape ) T = torch.zeros((7, 7, 7), dtype=dtype, device=device).repeat( From 0b6c7201fef8e98bb480fb9ec219d1373f1b561c Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 15:06:59 +0200 Subject: [PATCH 23/30] Test with vectorisation in first order as well --- cheetah/track_methods.py | 6 +++++- tests/test_sextupole.py | 4 ++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index e45fd85ee..58803ae11 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -274,7 +274,11 @@ def base_ttensor( # Rotate the R matrix for skew / vertical magnets if torch.any(tilt != 0): T = torch.einsum( - "...ij,...jk,...kl->...il", rotation_matrix(-tilt), T, rotation_matrix(tilt) + "...ij,...jkl,...kn,...lm->...inm", + rotation_matrix(-tilt), + T, + rotation_matrix(tilt), + rotation_matrix(tilt), ) return T diff --git a/tests/test_sextupole.py b/tests/test_sextupole.py index d3d07c156..b8fb4d940 100644 --- a/tests/test_sextupole.py +++ b/tests/test_sextupole.py @@ -59,8 +59,8 @@ def test_compare_sextupole_to_ocelot_particle_vectorized(): ) cheetah_sextupole = cheetah.Sextupole( length=torch.tensor(length), - k2=torch.tensor(k2).repeat([3, 2]), - tilt=torch.tensor(tilt), + k2=torch.tensor(k2).repeat([2]), + tilt=torch.tensor(tilt).repeat([3, 1]), ) outgoing_cheetah = cheetah_sextupole.track(incoming) From e1ed559618eb6a3ee1f166a9f64bd9618bf3ccbb Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 15:08:46 +0200 Subject: [PATCH 24/30] Add dependency to speed up einsum operations in `torch` --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index ba02f0522..5b57a13e4 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,6 @@ long_description_content_type="text/markdown", packages=[package for package in find_packages() if package.startswith("cheetah")], python_requires=">=3.10", - install_requires=["matplotlib", "numpy", "scipy", "torch"], + install_requires=["matplotlib", "numpy", "scipy", "torch[opt-einsum]"], extras_require={"openpmd": ["openpmd-beamphysics"]}, ) From f285e3c8d42be2870ab54aa05bde3d20e324d9e8 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 15:11:12 +0200 Subject: [PATCH 25/30] Presumed minor speed up in first order titlt by replacing `einsum` --- cheetah/track_methods.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 58803ae11..99cb03463 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -73,9 +73,8 @@ def base_rmatrix( # Rotate the R matrix for skew / vertical magnets if torch.any(tilt != 0): - R = torch.einsum( - "...ij,...jk,...kl->...il", rotation_matrix(-tilt), R, rotation_matrix(tilt) - ) + R = rotation_matrix(-tilt) @ R @ rotation_matrix(tilt) + return R From 3d3eea7a95248e88fa53b217f5f777c71e5cfc10 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 15:12:02 +0200 Subject: [PATCH 26/30] Fix `Sextupole.defining_features` --- cheetah/accelerator/sextupole.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index c8588c48c..67fddabb6 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -126,8 +126,9 @@ def split(self, resolution: torch.Tensor) -> list[Element]: 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"] + return super().defining_features() + ["length", "k2", "misalignment", "tilt"] def __repr__(self) -> str: return ( From 5971ad0a6fbd95c0ffea4838a23aab78cbd1204c Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 25 Apr 2025 15:20:03 +0200 Subject: [PATCH 27/30] Clean up matrix multiplications across Cheetah replacing `matmul` with `@` --- cheetah/accelerator/cavity.py | 8 +++----- cheetah/accelerator/custom_transfer_map.py | 5 +---- cheetah/accelerator/dipole.py | 8 ++++---- cheetah/accelerator/element.py | 6 +++--- cheetah/accelerator/segment.py | 2 +- cheetah/accelerator/sextupole.py | 6 +++--- cheetah/utils/kde.py | 2 +- 7 files changed, 16 insertions(+), 21 deletions(-) diff --git a/cheetah/accelerator/cavity.py b/cheetah/accelerator/cavity.py index 83e901de2..6e34893c5 100644 --- a/cheetah/accelerator/cavity.py +++ b/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 ) diff --git a/cheetah/accelerator/custom_transfer_map.py b/cheetah/accelerator/custom_transfer_map.py index 355662fd7..a861e273c 100644 --- a/cheetah/accelerator/custom_transfer_map.py +++ b/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) diff --git a/cheetah/accelerator/dipole.py b/cheetah/accelerator/dipole.py index d14af9876..2e89ae483 100644 --- a/cheetah/accelerator/dipole.py +++ b/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: diff --git a/cheetah/accelerator/element.py b/cheetah/accelerator/element.py index 6fb656dc4..e90a094a8 100644 --- a/cheetah/accelerator/element.py +++ b/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, diff --git a/cheetah/accelerator/segment.py b/cheetah/accelerator/segment.py index 916bdb016..7877c0b1d 100644 --- a/cheetah/accelerator/segment.py +++ b/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 diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index 67fddabb6..d78632f2f 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -88,8 +88,8 @@ def track(self, incoming: Beam) -> Beam: return super().track(incoming) elif isinstance(incoming, ParticleBeam): # Apply the transfer map to the incoming particles - first_order_particles = torch.matmul( - incoming.particles, first_order_tm.transpose(-2, -1) + first_order_particles = incoming.particles @ first_order_tm.transpose( + -2, -1 ) second_order_particles = torch.einsum( "...ijk,...j,...k->...i", @@ -128,7 +128,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None: @property def defining_features(self) -> list[str]: - return super().defining_features() + ["length", "k2", "misalignment", "tilt"] + return super().defining_features + ["length", "k2", "misalignment", "tilt"] def __repr__(self) -> str: return ( diff --git a/cheetah/utils/kde.py b/cheetah/utils/kde.py index 7d5c82973..22904785b 100644 --- a/cheetah/utils/kde.py +++ b/cheetah/utils/kde.py @@ -101,7 +101,7 @@ def _kde_joint_pdf_2d( + f"Got {type(kernel_values2)}" ) - joint_kernel_values = torch.matmul(kernel_values1.transpose(-2, -1), kernel_values2) + joint_kernel_values = kernel_values1.transpose(-2, -1) @ kernel_values2 normalization = ( torch.sum(joint_kernel_values, dim=(-2, -1)).unsqueeze(-1).unsqueeze(-1) + epsilon From 894dbd5d44764d22a8da9beee5610fcc626fecda Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Mon, 28 Apr 2025 10:51:55 +0200 Subject: [PATCH 28/30] Add sentence to docstring that MAD convention is used --- cheetah/track_methods.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 99cb03463..82a21ffa0 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -88,7 +88,8 @@ def base_ttensor( energy: torch.Tensor | None = None, ) -> torch.Tensor: """ - Create a second order universal transfer matrix for a beamline element. + Create a second order universal transfer matrix for a beamline element. Uses MAD + convention. :param length: Length of the element in m. :param k1: Quadrupole strength in 1/m**2. From 64ac13ccc69b22eb3da6ae375959217f9d280ebf Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Mon, 28 Apr 2025 11:04:17 +0200 Subject: [PATCH 29/30] Replace `einsum` by matrix multiplications Co-authored-by: Christian Hespe --- cheetah/accelerator/sextupole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index d78632f2f..e8017985b 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -62,7 +62,7 @@ def transfer_map(self, energy: torch.Tensor, species: Species) -> torch.Tensor: return R else: R_entry, R_exit = misalignment_matrix(self.misalignment) - R = torch.einsum("...ij,...jk,...kl->...il", R_exit, R, R_entry) + R = R_exit @ R @ R_entry return R def track(self, incoming: Beam) -> Beam: From e741aca4cf3cd8b131052ea7cb3d890fbd0c2288 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Mon, 28 Apr 2025 11:06:03 +0200 Subject: [PATCH 30/30] Add further suggestions by @Hespe --- cheetah/accelerator/sextupole.py | 6 +++++- cheetah/track_methods.py | 8 ++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py index d78632f2f..fda94e39a 100644 --- a/cheetah/accelerator/sextupole.py +++ b/cheetah/accelerator/sextupole.py @@ -28,7 +28,9 @@ def __init__( device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: - device, dtype = verify_device_and_dtype([length, k2], device, dtype) + device, dtype = verify_device_and_dtype( + [length, k2, misalignment, tilt], device, dtype + ) factory_kwargs = {"device": device, "dtype": dtype} super().__init__(name=name, **factory_kwargs) @@ -134,5 +136,7 @@ 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)})" ) diff --git a/cheetah/track_methods.py b/cheetah/track_methods.py index 82a21ffa0..8956a540e 100644 --- a/cheetah/track_methods.py +++ b/cheetah/track_methods.py @@ -15,7 +15,7 @@ def base_rmatrix( energy: torch.Tensor | None = None, ) -> torch.Tensor: """ - Create a first order universal transfer matrix for a beamline element. + Create a first order universal transfer map for a beamline element. :param length: Length of the element in m. :param k1: Quadrupole strength in 1/m**2. @@ -23,7 +23,7 @@ def base_rmatrix( :param species: Particle species of the beam. :param tilt: Roation of the element relative to the longitudinal axis in rad. :param energy: Beam energy in eV. - :return: First order transfer matrix for the element. + :return: First order transfer map for the element. """ device = length.device dtype = length.dtype @@ -88,7 +88,7 @@ def base_ttensor( energy: torch.Tensor | None = None, ) -> torch.Tensor: """ - Create a second order universal transfer matrix for a beamline element. Uses MAD + Create a second order universal transfer map for a beamline element. Uses MAD convention. :param length: Length of the element in m. @@ -98,7 +98,7 @@ def base_ttensor( :param species: Particle species of the beam. :param tilt: Roation of the element relative to the longitudinal axis in rad. :param energy: Beam energy in eV. - :return: Second order transfer matrix for the element. + :return: Second order transfer map for the element. """ device = length.device dtype = length.dtype