From b3881a4247874b7f755d86e8c66c9c7abccbe862 Mon Sep 17 00:00:00 2001
From: giadarol <giovanni.iadarola@cern.ch>
Date: Thu, 26 Sep 2019 15:36:25 +0200
Subject: [PATCH] Finalized and tested aperture

---
 examples/aperture/000_test_aperture.py | 42 ++++++++++++++++++++++++++
 pysixtrack/elements.py                 | 25 ++++++++++++++-
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 examples/aperture/000_test_aperture.py

diff --git a/examples/aperture/000_test_aperture.py b/examples/aperture/000_test_aperture.py
new file mode 100644
index 0000000..077aba7
--- /dev/null
+++ b/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()
diff --git a/pysixtrack/elements.py b/pysixtrack/elements.py
index fbd20ad..61ffd80 100644
--- a/pysixtrack/elements.py
+++ b/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 = [