Finalized and tested aperture

examples/aperture/000_test_aperture.py

@@ -0,0 +1,42 @@
+import numpy as np
+import pysixtrack
+
+p = pysixtrack.Particles()
+
+rect_aperture = pysixtrack.elements.LimitRect( min_x=-1e-2, max_x=2e-2,
+        min_y=-0.5e-2, max_y=2.5e-2)
+
+ellip_aperture = pysixtrack.elements.LimitEllipse(a=1.5e-2, b=.3e-2)
+
+# Test scalar
+assert(rect_aperture.track(p) is None)
+assert(ellip_aperture.track(p) is None)
+
+p.x= 100.
+assert(rect_aperture.track(p)==0)
+assert(ellip_aperture.track(p)==0)
+
+# Test vector
+from numpy.random import uniform
+
+N_part = 10000
+
+p.x = uniform(low=-3e-2, high=3e-2, size=N_part)
+p.y = uniform(low=-3e-2, high=3e-2, size=N_part)
+
+p.state = np.ones_like(p.x, dtype=np.int)
+
+rect_aperture.track(p)
+ellip_aperture.track(p)
+
+import matplotlib.pyplot as plt
+
+plt.close('all')
+fig1 = plt.figure(1)
+ax = fig1.add_subplot(111)
+
+ax.plot(p.x, p.y, '.b')
+ax.plot(p.lost_particles[0].x, p.lost_particles[0].y, 'r.')
+ax.plot(p.lost_particles[1].x, p.lost_particles[1].y, 'g.')
+
+plt.show()

pysixtrack/elements.py

@@ -270,7 +270,7 @@ def track(self, particle):
         x=particle.x
         y=particle.y
 
-        if not hasattr(particle, "__iter__"):
+        if not hasattr(particle.state, "__iter__"):
             particle.state = int(
                 x >= self.min_x
                 and x <= self.max_x
@@ -290,6 +290,29 @@ def track(self, particle):
             if len(particle.state == 0):
                 return -1
 
+class LimitEllipse(Element):
+    _description = [
+        ("a", "m^2", "Horizontal semiaxis", 1.0),
+        ("b", "m^2", "Vertical semiaxis", 1.0),
+    ]
+
+    def track(self, particle):
+
+        x=particle.x
+        y=particle.y
+
+        if not hasattr(particle.state, "__iter__"):
+            particle.state = int(
+                x*x/(self.a*self.a) + y*y/(self.b*self.b) <= 1.)
+            if particle.state != 1:
+                return particle.state
+        else:
+            particle.state = np.int_(
+                x*x/(self.a*self.a) + y*y/(self.b*self.b) <= 1.)
+            particle.remove_lost_particles()
+            if len(particle.state == 0):
+                return -1
+
 
 class BeamMonitor(Element):
     _description = [