Add LongitudinalAperture and fix other apperture classes

tracking/__init__.py

@@ -14,7 +14,9 @@ from mbtrack2.tracking.optics import Optics, PhyisicalModel
 from mbtrack2.tracking.element import (Element, LongitudinalMap, 
                                        TransverseMap, SynchrotronRadiation,
                                        SkewQuadrupole)
-from mbtrack2.tracking.aperture import (CircularAperture, ElipticalAperture,
-                                        RectangularAperture)
+from mbtrack2.tracking.aperture import (CircularAperture, 
+                                        ElipticalAperture,
+                                        RectangularAperture, 
+                                        LongitudinalAperture)
 from mbtrack2.tracking.wakepotential import WakePotential
                                        

tracking/aperture.py

@@ -16,13 +16,11 @@ class CircularAperture(Element):
     
     Parameters
     ----------
-    ring : Synchrotron object
     radius : float
         radius of the circle in [m]
     """
     
-    def __init__(self, ring, radius):
-        self.ring = ring
+    def __init__(self, radius):
         self.radius = radius
         self.radius_squared = radius**2
     
@@ -37,8 +35,8 @@ class CircularAperture(Element):
         ----------
         bunch : Bunch or Beam object
         """
-        alive = bunch.particles["x"]**2 + bunch.particles["y"]**2 > self.radius_squared
-        bunch.alive[alive] = False
+        alive = bunch.particles["x"]**2 + bunch.particles["y"]**2 < self.radius_squared
+        bunch.alive[~alive] = False
         
 class ElipticalAperture(Element):
     """
@@ -47,15 +45,13 @@ class ElipticalAperture(Element):
     
     Parameters
     ----------
-    ring : Synchrotron object
     X_radius : float
         horizontal radius of the elipse in [m]
     Y_radius : float
         vertical radius of the elipse in [m]
     """
     
-    def __init__(self, ring, X_radius, Y_radius):
-        self.ring = ring
+    def __init__(self, X_radius, Y_radius):
         self.X_radius = X_radius
         self.X_radius_squared = X_radius**2
         self.Y_radius = Y_radius
@@ -73,8 +69,8 @@ class ElipticalAperture(Element):
         bunch : Bunch or Beam object
         """
         alive = (bunch.particles["x"]**2/self.X_radius_squared + 
-                 bunch.particles["y"]**2/self.Y_radius_squared > 1)
-        bunch.alive[alive] = False
+                 bunch.particles["y"]**2/self.Y_radius_squared < 1)
+        bunch.alive[~alive] = False
 
 class RectangularAperture(Element):
     """
@@ -83,7 +79,6 @@ class RectangularAperture(Element):
     
     Parameters
     ----------
-    ring : Synchrotron object
     X_right : float
         right horizontal aperture of the rectangle in [m]
     Y_top : float
@@ -94,8 +89,7 @@ class RectangularAperture(Element):
         bottom vertical aperture of the rectangle in [m]
     """
     
-    def __init__(self, ring, X_right, Y_top, X_left=None, Y_bottom=None):
-        self.ring = ring
+    def __init__(self, X_right, Y_top, X_left=None, Y_bottom=None):
         self.X_right = X_right
         self.X_left = X_left
         self.Y_top = Y_top
@@ -127,4 +121,42 @@ class RectangularAperture(Element):
                      (bunch.particles["y"] > self.Y_bottom))
 
         alive = alive_X & alive_Y
-        bunch.alive[alive] = False
\ No newline at end of file
+        bunch.alive[~alive] = False
+        
+class LongitudinalAperture(Element):
+    """
+    Longitudinal aperture element. The particles which are outside of the 
+    longitudinal bounds are 'lost' and not used in the tracking any more.
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    tau_up : float
+        Upper longitudinal bound in [s].
+    tau_low : float, optional
+        Lower longitudinal bound in [s].
+    """
+    
+    def __init__(self, tau_up, tau_low=None):
+        self.tau_up = tau_up
+        if tau_low is None:
+            self.tau_low = tau_up*-1
+        else:
+            self.tau_low = tau_low
+    
+    @Element.parallel    
+    def track(self, bunch):
+        """
+        Tracking method for the element.
+        No bunch to bunch interaction, so written for Bunch objects and
+        @Element.parallel is used to handle Beam objects.
+        
+        Parameters
+        ----------
+        bunch : Bunch or Beam object
+        """
+        
+        alive = ((bunch.particles["tau"] < self.tau_up) &
+                 (bunch.particles["tau"] > self.tau_low))
+
+        bunch.alive[~alive] = False
\ No newline at end of file