Fix Aperture vectorisation issue

CHANGELOG.md

@@ -9,6 +9,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des
 - Cheetah is now vectorised. This means that you can run multiple simulations in parallel by passing a batch of beams and settings, resulting a number of interfaces being changed. For Cheetah developers this means that you now have to account for an arbitrary-dimensional tensor of most of the properties of you element, rather than a single value, vector or whatever else a property was before. (see #116, #157, #170, #172, #173, #198, #208, #213, #215, #218, #229, #233, #258, #265) (@jank324, @cr-xu, @hespe, @roussel-ryan)
 - The fifth particle coordinate `s` is renamed to `tau`. Now Cheetah uses the canonical variables in phase space $(x,px=\frac{P_x}{p_0},y,py, \tau=c\Delta t, \delta=\Delta E/{p_0 c})$. In addition, the trailing "s" was removed from some beam property names (e.g. `beam.xs` becomes `beam.x`). (see #163) (@cr-xu)
 - `Screen` no longer blocks the beam (by default). To return to old behaviour, set `Screen.is_blocking = True`. (see #208) (@jank324, @roussel-ryan)
+- Rework the `Aperture` element. Now `ParticleBeam` has a `particle_survival` attribute that keeps track of the lost particles. The statistical beam parameters are calculated only w.r.t. surviving particles. Note that the `Aperture` breaks differentiability if activated. (see #268) (@cr-xu, @jank324)
 
 ### 🚀 Features
 

cheetah/accelerator/aperture.py

@@ -15,6 +15,8 @@
 class Aperture(Element):
     """
     Physical aperture.
+    The aperture is only considered if tracking a `ParticleBeam` and the aperture is
+    active.
 
     :param x_max: half size horizontal offset in [m]
     :param y_max: half size vertical offset in [m]
@@ -72,41 +74,50 @@ def track(self, incoming: Beam) -> Beam:
         if not (isinstance(incoming, ParticleBeam) and self.is_active):
             return incoming
 
-        assert self.x_max >= 0 and self.y_max >= 0
+        assert torch.all(self.x_max >= 0) and torch.all(self.y_max >= 0)
         assert self.shape in [
             "rectangular",
             "elliptical",
         ], f"Unknown aperture shape {self.shape}"
 
+        # broadcast x_max and y_max and the ParticleBeam to the same shape
+        vector_shape = torch.broadcast_shapes(
+            self.x_max.shape,
+            self.y_max.shape,
+            incoming.x.shape[:-1],
+            incoming.energy.shape,
+        )
+        x_max = self.x_max.expand(vector_shape).unsqueeze(-1)
+        y_max = self.y_max.expand(vector_shape).unsqueeze(-1)
+        outgoing_particles = incoming.particles.expand(vector_shape + (-1, 7))
+
         if self.shape == "rectangular":
             survived_mask = torch.logical_and(
-                torch.logical_and(incoming.x > -self.x_max, incoming.x < self.x_max),
-                torch.logical_and(incoming.y > -self.y_max, incoming.y < self.y_max),
+                torch.logical_and(incoming.x > -x_max, incoming.x < x_max),
+                torch.logical_and(incoming.y > -y_max, incoming.y < y_max),
             )
         elif self.shape == "elliptical":
-            survived_mask = (
-                incoming.x**2 / self.x_max**2 + incoming.y**2 / self.y_max**2
-            ) <= 1.0
-        outgoing_particles = incoming.particles[survived_mask]
+            survived_mask = (incoming.x**2 / x_max**2 + incoming.y**2 / y_max**2) <= 1.0
 
-        outgoing_particle_charges = incoming.particle_charges[survived_mask]
+        outgoing_survival = incoming.particle_survival * survived_mask
 
-        self.lost_particles = incoming.particles[torch.logical_not(survived_mask)]
+        # outgoing_particles = incoming.particles[survived_mask]
 
-        self.lost_particle_charges = incoming.particle_charges[
-            torch.logical_not(survived_mask)
-        ]
+        # outgoing_particle_charges = incoming.particle_charges[survived_mask]
 
-        return (
-            ParticleBeam(
-                outgoing_particles,
-                incoming.energy,
-                particle_charges=outgoing_particle_charges,
-                device=outgoing_particles.device,
-                dtype=outgoing_particles.dtype,
-            )
-            if outgoing_particles.shape[0] > 0
-            else ParticleBeam.empty
+        # self.lost_particles = incoming.particles[torch.logical_not(survived_mask)]
+
+        # self.lost_particle_charges = incoming.particle_charges[
+        #     torch.logical_not(survived_mask)
+        # ]
+
+        return ParticleBeam(
+            outgoing_particles,
+            incoming.energy,
+            particle_charges=incoming.particle_charges,
+            device=incoming.particles.device,
+            dtype=incoming.particles.dtype,
+            particle_survival=outgoing_survival,
         )
 
     def split(self, resolution: torch.Tensor) -> list[Element]:

cheetah/accelerator/cavity.py

@@ -243,6 +243,7 @@ def _track_beam(self, incoming: Beam) -> Beam:
                 particle_charges=incoming.particle_charges,
                 device=outgoing_particles.device,
                 dtype=outgoing_particles.dtype,
+                particle_survival=incoming.particle_survival,
             )
             return outgoing
 

cheetah/accelerator/dipole.py

@@ -235,6 +235,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             particle_charges=incoming.particle_charges,
             device=incoming.particles.device,
             dtype=incoming.particles.dtype,
+            particle_survival=incoming.particle_survival,
         )
         return outgoing_beam
 

cheetah/accelerator/drift.py

@@ -117,6 +117,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             particle_charges=incoming.particle_charges,
             device=incoming.particles.device,
             dtype=incoming.particles.dtype,
+            particle_survival=incoming.particle_survival,
         )
         return outgoing_beam
 

cheetah/accelerator/element.py

@@ -83,6 +83,7 @@ def track(self, incoming: Beam) -> Beam:
                 particle_charges=incoming.particle_charges,
                 device=new_particles.device,
                 dtype=new_particles.dtype,
+                particle_survival=incoming.particle_survival,
             )
         else:
             raise TypeError(f"Parameter incoming is of invalid type {type(incoming)}")

cheetah/accelerator/quadrupole.py

@@ -190,6 +190,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             particle_charges=incoming.particle_charges,
             device=incoming.particles.device,
             dtype=incoming.particles.dtype,
+            particle_survival=incoming.particle_survival,
         )
         return outgoing_beam
 

cheetah/accelerator/space_charge_kick.py

@@ -151,7 +151,8 @@ def _deposit_charge_on_grid(
         )
 
         # Accumulate the charge contributions
-        repeated_charges = beam.particle_charges.repeat_interleave(
+        survived_particle_charges = beam.particle_charges * beam.particle_survival
+        repeated_charges = survived_particle_charges.repeat_interleave(
             repeats=8, dim=-1
         )  # Shape:(..., 8 * num_particles)
         values = (cell_weights.flatten(start_dim=-2) * repeated_charges)[valid_mask]
@@ -563,6 +564,7 @@ def track(self, incoming: ParticleBeam) -> ParticleBeam:
                     particle_charges=incoming.particle_charges.unsqueeze(0),
                     device=incoming.particles.device,
                     dtype=incoming.particles.dtype,
+                    particle_survival=incoming.particle_survival.unsqueeze(0),
                 )
             else:
                 is_incoming_vectorized = True
@@ -577,6 +579,9 @@ def track(self, incoming: ParticleBeam) -> ParticleBeam:
                 ),
                 device=vectorized_incoming.particles.device,
                 dtype=vectorized_incoming.particles.dtype,
+                particle_survival=vectorized_incoming.particle_survival.flatten(
+                    end_dim=-2
+                ),
             )
             flattened_length_effect = self.effect_length.flatten(end_dim=-1)
 
@@ -614,9 +619,10 @@ def track(self, incoming: ParticleBeam) -> ParticleBeam:
                 outgoing = ParticleBeam.from_xyz_pxpypz(
                     xp_coordinates.squeeze(0),
                     vectorized_incoming.energy.squeeze(0),
-                    vectorized_incoming.particle_charges.squeeze(0),
-                    vectorized_incoming.particles.device,
-                    vectorized_incoming.particles.dtype,
+                    particle_charges=vectorized_incoming.particle_charges.squeeze(0),
+                    device=vectorized_incoming.particles.device,
+                    dtype=vectorized_incoming.particles.dtype,
+                    particle_survival=vectorized_incoming.particle_survival.squeeze(0),
                 )
             else:
                 # Reverse the flattening of the vector dimensions
@@ -625,9 +631,10 @@ def track(self, incoming: ParticleBeam) -> ParticleBeam:
                         dim=0, sizes=vectorized_incoming.particles.shape[:-2]
                     ),
                     vectorized_incoming.energy,
-                    vectorized_incoming.particle_charges,
-                    vectorized_incoming.particles.device,
-                    vectorized_incoming.particles.dtype,
+                    particle_charges=vectorized_incoming.particle_charges,
+                    particle_survival=vectorized_incoming.particle_survival,
+                    device=vectorized_incoming.particles.device,
+                    dtype=vectorized_incoming.particles.dtype,
                 )
             return outgoing
         else:

cheetah/accelerator/transverse_deflecting_cavity.py

@@ -207,6 +207,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             torch.stack((x, px, y, py, tau, delta, torch.ones_like(x)), dim=-1),
             ref_energy,
             particle_charges=incoming.particle_charges,
+            particle_survival=incoming.particle_survival,
             device=incoming.particles.device,
             dtype=incoming.particles.dtype,
         )