Add particle species in Beam classes and update tracking methods

.github/workflows/pytest.yaml

@@ -15,21 +15,28 @@ jobs:
         include:
           - os: macos-latest
             python-version: "3.12"
-          - os: windows-latest
-            python-version: "3.12"
     runs-on: ${{ matrix.os }}
 
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+      - uses: conda-incubator/setup-miniconda@v3
         with:
+          miniconda-version: 'latest'
+          conda-solver: 'classic'
           python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
+          activate-environment: test_environment
+      - name: Install Bmad and Pytao
+        shell: bash -el {0}
+        run: |
+          conda install -c conda-forge bmad
+          conda install -c conda-forge pytao
+      - name: Install pip dependencies
+        shell: bash -el {0}
         run: |
           python -m pip install --upgrade pip
           pip install -e .
           pip install -r test_requirements.txt
       - name: Test with pytest
+        shell: bash -el {0}
         run: |
           pytest

.github/workflows/pytest_windows.yaml

@@ -0,0 +1,30 @@
+name: pytest_windows
+
+on:
+  push:
+    branches: [master]
+  pull_request:
+    branches: [master]
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        os: [windows-latest]
+        python-version: ["3.12"]
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+          pip install -r test_requirements.txt
+      - name: Test with pytest
+        run: |
+          pytest --ignore=tests/bmad_benchmarks

cheetah/accelerator/aperture.py

@@ -61,7 +61,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: float
+    ) -> torch.Tensor:
         device = self.x_max.device
         dtype = self.x_max.dtype
 

cheetah/accelerator/bpm.py

@@ -31,7 +31,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: float
+    ) -> 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 torch import nn
 
 from cheetah.accelerator.element import Element
@@ -14,8 +13,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-electron_mass_eV = physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-
 
 class Cavity(Element):
     """
@@ -75,14 +72,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: float
+    ) -> 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,
             ),
@@ -113,7 +113,7 @@ def _track_beam(self, incoming: Beam) -> Beam:
 
         phi = torch.deg2rad(self.phase)
 
-        tm = self.transfer_map(incoming.energy)
+        tm = self.transfer_map(incoming.energy, incoming.mass_eV)
         if isinstance(incoming, ParameterBeam):
             outgoing_mu = torch.matmul(tm, incoming._mu.unsqueeze(-1)).squeeze(-1)
             outgoing_cov = torch.matmul(
@@ -161,7 +161,7 @@ def _track_beam(self, incoming: Beam) -> Beam:
                     - torch.cos(phi).unsqueeze(-1)
                 )
 
-            dgamma = self.voltage / electron_mass_eV
+            dgamma = self.voltage / incoming.mass_eV
             if torch.any(delta_energy > 0):
                 T566 = (
                     self.length
@@ -246,7 +246,9 @@ 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: float
+    ) -> torch.Tensor:
         """Produces an R-matrix for a cavity when it is on, i.e. voltage > 0.0."""
         device = self.length.device
         dtype = self.length.dtype
@@ -255,8 +257,8 @@ def _cavity_rmatrix(self, energy: torch.Tensor) -> torch.Tensor:
         delta_energy = self.voltage * torch.cos(phi)
         # Comment from Ocelot: Pure pi-standing-wave case
         eta = torch.tensor(1.0, device=device, dtype=dtype)
-        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"
 
@@ -306,14 +308,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

@@ -64,14 +64,20 @@ 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.mass_eV).device
+        )
+        dtype = (
+            elements[0].transfer_map(incoming_beam.energy, incoming_beam.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.mass_eV), tm
+            )
             incoming_beam = element.track(incoming_beam)
 
         combined_length = sum(element.length for element in elements)
@@ -82,7 +88,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: float
+    ) -> torch.Tensor:
         return self._transfer_map
 
     @property

cheetah/accelerator/dipole.py

@@ -4,7 +4,6 @@
 import numpy as np
 import torch
 from matplotlib.patches import Rectangle
-from scipy.constants import physical_constants
 from torch import nn
 
 from cheetah.accelerator.element import Element
@@ -14,8 +13,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-electron_mass_eV = physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-
 
 class Dipole(Element):
     """
@@ -196,10 +193,9 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
         py = incoming.py
         tau = incoming.tau
         delta = incoming.p
+        mc2 = incoming.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(
@@ -208,9 +204,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)
 
@@ -220,9 +214,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)
@@ -235,6 +227,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             particle_charges=incoming.particle_charges,
             device=incoming.particles.device,
             dtype=incoming.particles.dtype,
+            species=incoming.species,
         )
         return outgoing_beam
 
@@ -368,7 +361,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: float
+    ) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
@@ -380,6 +375,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,7 +2,6 @@
 
 import matplotlib.pyplot as plt
 import torch
-from scipy.constants import physical_constants
 from torch import nn
 
 from cheetah.accelerator.element import Element
@@ -11,8 +10,6 @@
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
-electron_mass_eV = physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
-
 
 class Drift(Element):
     """
@@ -40,11 +37,13 @@ def __init__(
         self.register_buffer("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: float
+    ) -> 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)
 
@@ -92,18 +91,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.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.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.mass_eV
         )
 
         # Broadcast to align their shapes so that they can be stacked
@@ -117,6 +116,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             particle_charges=incoming.particle_charges,
             device=incoming.particles.device,
             dtype=incoming.particles.dtype,
+            species=incoming.species,
         )
         return outgoing_beam
 

cheetah/accelerator/element.py

@@ -24,7 +24,9 @@ def __init__(self, name: Optional[str] = None) -> None:
         self.name = name if name is not None else generate_unique_name()
         self.register_buffer("length", torch.tensor(0.0))
 
-    def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
+    def transfer_map(
+        self, energy: torch.Tensor, particle_mass_eV: float
+    ) -> torch.Tensor:
         r"""
         Generates the element's transfer map that describes how the beam and its
         particles are transformed when traveling through the element.
@@ -63,7 +65,7 @@ def track(self, incoming: Beam) -> Beam:
         if incoming is Beam.empty:
             return incoming
         elif isinstance(incoming, ParameterBeam):
-            tm = self.transfer_map(incoming.energy)
+            tm = self.transfer_map(incoming.energy, incoming.mass_eV)
             mu = torch.matmul(tm, incoming._mu.unsqueeze(-1)).squeeze(-1)
             cov = torch.matmul(tm, torch.matmul(incoming._cov, tm.transpose(-2, -1)))
             return ParameterBeam(
@@ -75,7 +77,7 @@ def track(self, incoming: Beam) -> Beam:
                 dtype=mu.dtype,
             )
         elif isinstance(incoming, ParticleBeam):
-            tm = self.transfer_map(incoming.energy)
+            tm = self.transfer_map(incoming.energy, incoming.mass_eV)
             new_particles = torch.matmul(incoming.particles, tm.transpose(-2, -1))
             return ParticleBeam(
                 new_particles,