Add particle species in Beam classes and update tracking methods

.github/workflows/pytest.yaml

@@ -18,18 +18,26 @@ jobs:
           - os: windows-latest
             python-version: "3.12"
     runs-on: ${{ matrix.os }}
+    defaults:
+      run:
+        shell: bash -el {0}
 
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
+          miniforge-version: "latest"
+      - name: Install Bmad and PyTao
+        run: |
+          if [[ ! ${{ runner.os }} == 'Windows' ]]; then
+            conda install bmad pytao
+          fi
+      - name: Install pip dependencies
         run: |
           python -m pip install --upgrade pip
           pip install -e .[openpmd]
           pip install -r test_requirements.txt
       - name: Test with pytest
         run: |
-          pytest --seed 42
+          pytest

CHANGELOG.md

@@ -9,6 +9,7 @@
 ### 🚀 Features
 
 - `ParticleBeam` now supports importing from and exporting to [openPMD-beamphysics](https://github.com/ChristopherMayes/openPMD-beamphysics) HDF5 files and `ParticleGroup` objects. This allows for easy conversion to and from other file formats supported by openPMD-beamphysics. (see #305, #320) (@cr-xu, @Hespe)
+- Add support for particle species through a new `Species` class (see #276) (@cr-xu, @jank324)
 
 ### 🐛 Bug fixes
 

cheetah/__init__.py

@@ -18,4 +18,4 @@
     Undulator,
     VerticalCorrector,
 )
-from .particles import Beam, ParameterBeam, ParticleBeam  # noqa: F401
+from .particles import Beam, ParameterBeam, ParticleBeam, Species  # noqa: F401

cheetah/accelerator/aperture.py

@@ -56,7 +56,9 @@ def __init__(
     def is_skippable(self) -> bool:
         return not self.is_active
 
-    def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
+    def transfer_map(
+        self, energy: torch.Tensor, particle_mass_eV: torch.Tensor
+    ) -> torch.Tensor:
         device = self.x_max.device
         dtype = self.x_max.dtype
 

cheetah/accelerator/bpm.py

@@ -30,7 +30,9 @@ def __init__(self, is_active: bool = False, name: Optional[str] = None) -> None:
     def is_skippable(self) -> bool:
         return not self.is_active
 
-    def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
+    def transfer_map(
+        self, energy: torch.Tensor, particle_mass_eV: torch.Tensor
+    ) -> torch.Tensor:
         return torch.eye(7, device=energy.device, dtype=energy.dtype).repeat(
             (*energy.shape, 1, 1)
         )

cheetah/accelerator/cavity.py

@@ -4,7 +4,6 @@
 import torch
 from matplotlib.patches import Rectangle
 from scipy import constants
-from scipy.constants import physical_constants
 
 from cheetah.accelerator.element import Element
 from cheetah.particles import Beam, ParameterBeam, ParticleBeam
@@ -17,15 +16,15 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-electron_mass_eV = physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-
 
 class Cavity(Element):
     """
     Accelerating cavity in a particle accelerator.
 
     :param length: Length in meters.
-    :param voltage: Voltage of the cavity in volts.
+    :param voltage: Voltage of the cavity in volts. NOTE: This assumes the effective
+        voltage. If the particle does not have unit charge, the voltage needs to be
+        scaled properly.
     :param phase: Phase of the cavity in degrees.
     :param frequency: Frequency of the cavity in Hz.
     :param name: Unique identifier of the element.
@@ -67,14 +66,17 @@ def is_active(self) -> bool:
     def is_skippable(self) -> bool:
         return not self.is_active
 
-    def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
+    def transfer_map(
+        self, energy: torch.Tensor, particle_mass_eV: torch.Tensor
+    ) -> torch.Tensor:
         return torch.where(
             (self.voltage != 0).unsqueeze(-1).unsqueeze(-1),
-            self._cavity_rmatrix(energy),
+            self._cavity_rmatrix(energy, particle_mass_eV),
             base_rmatrix(
                 length=self.length,
                 k1=torch.zeros_like(self.length),
                 hx=torch.zeros_like(self.length),
+                particle_mass_eV=particle_mass_eV,
                 tilt=torch.zeros_like(self.length),
                 energy=energy,
             ),
@@ -99,11 +101,13 @@ def _track_beam(self, incoming: Beam) -> Beam:
         Track particles through the cavity. The input can be a `ParameterBeam` or a
         `ParticleBeam`.
         """
-        gamma0, igamma2, beta0 = compute_relativistic_factors(incoming.energy)
+        gamma0, igamma2, beta0 = compute_relativistic_factors(
+            incoming.energy, incoming.species.mass_eV
+        )
 
         phi = torch.deg2rad(self.phase)
 
-        tm = self.transfer_map(incoming.energy)
+        tm = self.transfer_map(incoming.energy, incoming.species.mass_eV)
         if isinstance(incoming, ParameterBeam):
             outgoing_mu = torch.matmul(tm, incoming._mu.unsqueeze(-1)).squeeze(-1)
             outgoing_cov = torch.matmul(
@@ -120,7 +124,9 @@ def _track_beam(self, incoming: Beam) -> Beam:
         if torch.any(incoming.energy + delta_energy > 0):
             k = 2 * torch.pi * self.frequency / constants.speed_of_light
             outgoing_energy = incoming.energy + delta_energy
-            gamma1, _, beta1 = compute_relativistic_factors(outgoing_energy)
+            gamma1, _, beta1 = compute_relativistic_factors(
+                outgoing_energy, incoming.species.mass_eV
+            )
 
             if isinstance(incoming, ParameterBeam):
                 outgoing_mu[..., 5] = incoming._mu[..., 5] * incoming.energy * beta0 / (
@@ -151,7 +157,7 @@ def _track_beam(self, incoming: Beam) -> Beam:
                     - torch.cos(phi).unsqueeze(-1)
                 )
 
-            dgamma = self.voltage / electron_mass_eV
+            dgamma = self.voltage / incoming.species.mass_eV
             if torch.any(delta_energy > 0):
                 T566 = (
                     self.length
@@ -237,16 +243,18 @@ def _track_beam(self, incoming: Beam) -> Beam:
             )
             return outgoing
 
-    def _cavity_rmatrix(self, energy: torch.Tensor) -> torch.Tensor:
+    def _cavity_rmatrix(
+        self, energy: torch.Tensor, particle_mass_eV: torch.Tensor
+    ) -> torch.Tensor:
         """Produces an R-matrix for a cavity when it is on, i.e. voltage > 0.0."""
         factory_kwargs = {"device": self.length.device, "dtype": self.length.dtype}
 
         phi = torch.deg2rad(self.phase)
         delta_energy = self.voltage * torch.cos(phi)
         # Comment from Ocelot: Pure pi-standing-wave case
         eta = torch.tensor(1.0, **factory_kwargs)
-        Ei = energy / electron_mass_eV
-        Ef = (energy + delta_energy) / electron_mass_eV
+        Ei = energy / particle_mass_eV
+        Ef = (energy + delta_energy) / particle_mass_eV
         Ep = (Ef - Ei) / self.length  # Derivative of the energy
         assert torch.all(Ei > 0), "Initial energy must be larger than 0"
 
@@ -296,14 +304,14 @@ def _cavity_rmatrix(self, energy: torch.Tensor) -> torch.Tensor:
                 * self.length
                 * beta0
                 * self.voltage
-                / electron_mass_eV
+                / particle_mass_eV
                 * torch.sin(phi)
                 * (g0 * g1 * (beta0 * beta1 - 1) + 1)
                 / (beta1 * g1 * (g0 - g1) ** 2)
             )
 
         r66 = Ei / Ef * beta0 / beta1
-        r65 = k * torch.sin(phi) * self.voltage / (Ef * beta1 * electron_mass_eV)
+        r65 = k * torch.sin(phi) * self.voltage / (Ef * beta1 * particle_mass_eV)
 
         # Make sure that all matrix elements have the same shape
         r11, r12, r21, r22, r55_cor, r56, r65, r66 = torch.broadcast_tensors(

cheetah/accelerator/custom_transfer_map.py

@@ -62,14 +62,27 @@ def from_merging_elements(
             " incorrect tracking results."
         )
 
-        device = elements[0].transfer_map(incoming_beam.energy).device
-        dtype = elements[0].transfer_map(incoming_beam.energy).dtype
+        device = (
+            elements[0]
+            .transfer_map(incoming_beam.energy, incoming_beam.species.mass_eV)
+            .device
+        )
+        dtype = (
+            elements[0]
+            .transfer_map(incoming_beam.energy, incoming_beam.species.mass_eV)
+            .dtype
+        )
 
         tm = torch.eye(7, device=device, dtype=dtype).repeat(
             (*incoming_beam.energy.shape, 1, 1)
         )
         for element in elements:
-            tm = torch.matmul(element.transfer_map(incoming_beam.energy), tm)
+            tm = torch.matmul(
+                element.transfer_map(
+                    incoming_beam.energy, incoming_beam.species.mass_eV
+                ),
+                tm,
+            )
             incoming_beam = element.track(incoming_beam)
 
         combined_length = sum(element.length for element in elements)
@@ -80,7 +93,9 @@ def from_merging_elements(
             tm, length=combined_length, device=device, dtype=dtype, name=combined_name
         )
 
-    def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
+    def transfer_map(
+        self, energy: torch.Tensor, particle_mass_eV: torch.Tensor
+    ) -> torch.Tensor:
         return self.predefined_transfer_map
 
     @property

cheetah/accelerator/dipole.py

@@ -3,7 +3,6 @@
 import matplotlib.pyplot as plt
 import torch
 from matplotlib.patches import Rectangle
-from scipy.constants import physical_constants
 
 from cheetah.accelerator.element import Element
 from cheetah.particles import Beam, ParticleBeam
@@ -12,8 +11,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-electron_mass_eV = physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-
 
 class Dipole(Element):
     """
@@ -189,10 +186,9 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
         py = incoming.py
         tau = incoming.tau
         delta = incoming.p
+        mc2 = incoming.species.mass_eV
 
-        z, pz, p0c = bmadx.cheetah_to_bmad_z_pz(
-            tau, delta, incoming.energy, electron_mass_eV
-        )
+        z, pz, p0c = bmadx.cheetah_to_bmad_z_pz(tau, delta, incoming.energy, mc2)
 
         # Begin Bmad-X tracking
         x, px, y, py = bmadx.offset_particle_set(
@@ -207,9 +203,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
 
         if self.fringe_at == "entrance" or self.fringe_at == "both":
             px, py = self._bmadx_fringe_linear("entrance", x, px, y, py)
-        x, px, y, py, z, pz = self._bmadx_body(
-            x, px, y, py, z, pz, p0c, electron_mass_eV
-        )
+        x, px, y, py, z, pz = self._bmadx_body(x, px, y, py, z, pz, p0c, mc2)
         if self.fringe_at == "exit" or self.fringe_at == "both":
             px, py = self._bmadx_fringe_linear("exit", x, px, y, py)
 
@@ -225,9 +219,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
         # End of Bmad-X tracking
 
         # Convert back to Cheetah coordinates
-        tau, delta, ref_energy = bmadx.bmad_to_cheetah_z_pz(
-            z, pz, p0c, electron_mass_eV
-        )
+        tau, delta, ref_energy = bmadx.bmad_to_cheetah_z_pz(z, pz, p0c, mc2)
 
         # Broadcast to align their shapes so that they can be stacked
         x, px, y, py, tau, delta = torch.broadcast_tensors(x, px, y, py, tau, delta)
@@ -241,6 +233,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             survival_probabilities=incoming.survival_probabilities,
             device=incoming.particles.device,
             dtype=incoming.particles.dtype,
+            species=incoming.species,
         )
         return outgoing_beam
 
@@ -374,7 +367,9 @@ def _bmadx_fringe_linear(
 
         return px_f, py_f
 
-    def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
+    def transfer_map(
+        self, energy: torch.Tensor, particle_mass_eV: torch.Tensor
+    ) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
@@ -386,6 +381,7 @@ def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
                 length=self.length,
                 k1=self.k1,
                 hx=self.hx,
+                particle_mass_eV=particle_mass_eV,
                 tilt=torch.zeros_like(self.length),
                 energy=energy,
             )  # Tilt is applied after adding edges

cheetah/accelerator/drift.py

@@ -2,16 +2,13 @@
 
 import matplotlib.pyplot as plt
 import torch
-from scipy.constants import physical_constants
 
 from cheetah.accelerator.element import Element
 from cheetah.particles import Beam, ParticleBeam
 from cheetah.utils import UniqueNameGenerator, bmadx, compute_relativistic_factors
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-electron_mass_eV = physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-
 
 class Drift(Element):
     """
@@ -39,11 +36,13 @@ def __init__(
         self.length = torch.as_tensor(length, **factory_kwargs)
         self.tracking_method = tracking_method
 
-    def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
+    def transfer_map(
+        self, energy: torch.Tensor, particle_mass_eV: torch.Tensor
+    ) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
-        _, igamma2, beta = compute_relativistic_factors(energy)
+        _, igamma2, beta = compute_relativistic_factors(energy, particle_mass_eV)
 
         vector_shape = torch.broadcast_shapes(self.length.shape, igamma2.shape)
 
@@ -91,18 +90,18 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
         delta = incoming.p
 
         z, pz, p0c = bmadx.cheetah_to_bmad_z_pz(
-            tau, delta, incoming.energy, electron_mass_eV
+            tau, delta, incoming.energy, incoming.species.mass_eV
         )
 
         # Begin Bmad-X tracking
         x, y, z = bmadx.track_a_drift(
-            self.length, x, px, y, py, z, pz, p0c, electron_mass_eV
+            self.length, x, px, y, py, z, pz, p0c, incoming.species.mass_eV
         )
         # End of Bmad-X tracking
 
         # Convert back to Cheetah coordinates
         tau, delta, ref_energy = bmadx.bmad_to_cheetah_z_pz(
-            z, pz, p0c, electron_mass_eV
+            z, pz, p0c, incoming.species.mass_eV
         )
 
         # Broadcast to align their shapes so that they can be stacked
@@ -117,6 +116,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             survival_probabilities=incoming.survival_probabilities,
             device=incoming.particles.device,
             dtype=incoming.particles.dtype,
+            species=incoming.species,
         )
         return outgoing_beam