Fix broken references to mass_eV

Author: jank324

cheetah/accelerator/cavity.py

@@ -105,7 +105,7 @@ def _track_beam(self, incoming: Beam) -> Beam:
 
         phi = torch.deg2rad(self.phase)
 
-        tm = self.transfer_map(incoming.energy, incoming.mass_eV)
+        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(
@@ -153,7 +153,7 @@ def _track_beam(self, incoming: Beam) -> Beam:
                     - torch.cos(phi).unsqueeze(-1)
                 )
 
-            dgamma = self.voltage / incoming.mass_eV
+            dgamma = self.voltage / incoming.species.mass_eV
             if torch.any(delta_energy > 0):
                 T566 = (
                     self.length

cheetah/accelerator/custom_transfer_map.py

@@ -63,18 +63,25 @@ def from_merging_elements(
         )
 
         device = (
-            elements[0].transfer_map(incoming_beam.energy, incoming_beam.mass_eV).device
+            elements[0]
+            .transfer_map(incoming_beam.energy, incoming_beam.species.mass_eV)
+            .device
         )
         dtype = (
-            elements[0].transfer_map(incoming_beam.energy, incoming_beam.mass_eV).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, incoming_beam.mass_eV), tm
+                element.transfer_map(
+                    incoming_beam.energy, incoming_beam.species.mass_eV
+                ),
+                tm,
             )
             incoming_beam = element.track(incoming_beam)
 

cheetah/accelerator/dipole.py

@@ -186,7 +186,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
         py = incoming.py
         tau = incoming.tau
         delta = incoming.p
-        mc2 = incoming.mass_eV
+        mc2 = incoming.species.mass_eV
 
         z, pz, p0c = bmadx.cheetah_to_bmad_z_pz(tau, delta, incoming.energy, mc2)
 

cheetah/accelerator/drift.py

@@ -90,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, incoming.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, incoming.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, incoming.mass_eV
+            z, pz, p0c, incoming.species.mass_eV
         )
 
         # Broadcast to align their shapes so that they can be stacked

cheetah/accelerator/element.py

@@ -64,7 +64,7 @@ def track(self, incoming: Beam) -> Beam:
         :return: Beam of particles exiting the element.
         """
         if isinstance(incoming, ParameterBeam):
-            tm = self.transfer_map(incoming.energy, incoming.mass_eV)
+            tm = self.transfer_map(incoming.energy, incoming.species.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(
@@ -76,7 +76,7 @@ def track(self, incoming: Beam) -> Beam:
                 dtype=mu.dtype,
             )
         elif isinstance(incoming, ParticleBeam):
-            tm = self.transfer_map(incoming.energy, incoming.mass_eV)
+            tm = self.transfer_map(incoming.energy, incoming.species.mass_eV)
             new_particles = torch.matmul(incoming.particles, tm.transpose(-2, -1))
             return ParticleBeam(
                 new_particles,

cheetah/accelerator/quadrupole.py

@@ -114,7 +114,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
         py = incoming.py
         tau = incoming.tau
         delta = incoming.p
-        mc2 = incoming.mass_eV
+        mc2 = incoming.species.mass_eV
 
         z, pz, p0c = bmadx.cheetah_to_bmad_z_pz(tau, delta, incoming.energy, mc2)
 

cheetah/accelerator/transverse_deflecting_cavity.py

@@ -118,7 +118,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
         py = incoming.py
         tau = incoming.tau
         delta = incoming.p
-        mc2 = incoming.mass_eV
+        mc2 = incoming.species.mass_eV
 
         z, pz, p0c = bmadx.cheetah_to_bmad_z_pz(tau, delta, incoming.energy, mc2)
 

cheetah/particles/particle_beam.py

@@ -1547,7 +1547,7 @@ def energies(self) -> torch.Tensor:
     @property
     def momenta(self) -> torch.Tensor:
         """Momenta of the individual particles."""
-        return torch.sqrt(self.energies**2 - self.mass_eV**2)
+        return torch.sqrt(self.energies**2 - self.species.mass_eV**2)
 
     def clone(self) -> "ParticleBeam":
         return ParticleBeam(

tests/test_cavity.py

@@ -72,7 +72,9 @@ def test_vectorized_cavity_zero_voltage(voltage):
 
     outgoing = cavity.track(incoming)
 
-    assert not torch.isnan(cavity.transfer_map(incoming.energy, incoming.mass_eV)).any()
+    assert not torch.isnan(
+        cavity.transfer_map(incoming.energy, incoming.species.mass_eV)
+    ).any()
 
     assert not torch.isnan(outgoing.sigma_x).any()
     assert not torch.isnan(outgoing.sigma_y).any()

tests/test_speed_optimizations.py

@@ -232,10 +232,10 @@ def test_skippable_elements_reset():
     )
 
     original_tm = original_segment.elements[2].transfer_map(
-        energy=incoming_beam.energy, particle_mass_eV=incoming_beam.mass_eV
+        energy=incoming_beam.energy, particle_mass_eV=incoming_beam.species.mass_eV
     )
     merged_tm = merged_segment.elements[2].transfer_map(
-        energy=incoming_beam.energy, particle_mass_eV=incoming_beam.mass_eV
+        energy=incoming_beam.energy, particle_mass_eV=incoming_beam.species.mass_eV
     )
 
     assert torch.allclose(original_tm, merged_tm)