simplified changes and speeded it up

Author: zerothi

src/sisl/_indices.pyx

@@ -145,17 +145,18 @@ def indices(ints_st[::1] element, ints_st[::1] test_element, ints_st offset=0,
 @cython.boundscheck(False)
 @cython.wraparound(False)
 @cython.initializedcheck(False)
-def indices_in_cylinder(floats_st[:, ::1] dxyz, R, const floats_st h):
-    """ Indices for all coordinates that are within a cylinde radius `R` and height `h`
+def indices_in_cylinder(floats_st[:, ::1] dxyz, const floats_st R, const floats_st h):
+    """ Indices for all coordinates that are within a cylinder radius `R` and height `h`
 
     Parameters
     ----------
     dxyz :
-       coordinates centered around the cylinder
+       coordinates centered around the cylinder.
+       The last axis is the cylinder height.
     R :
-       radius of cylinder to check
+       radius of cylinder to check.
     h :
-       height of cylinder to check
+       height of cylinder to check.
 
     Returns
     -------
@@ -168,34 +169,22 @@ def indices_in_cylinder(floats_st[:, ::1] dxyz, R, const floats_st h):
     cdef ndarray[int32_t] IDX = np.empty([n], dtype=np.int32)
     cdef int[::1] idx = IDX
 
-    cdef floats_st Rx, Ry
-    cdef floats_st L2
+    cdef floats_st R2
+    cdef floats_st hhalve, L2
     cdef Py_ssize_t i, j, m
 
     # Handle radius input
-    if isinstance(R, (float, int)):
-        Rx = Ry = R
-    else:
-        R = np.asarray(R, dtype=np.float64)
-        if R.shape == (2,):
-            Rx, Ry = R[0], R[1]
-        elif R.shape == (2,3):
-            from sisl._core._dtypes import fnorm
-            Rx = fnorm(R[0])
-            Ry = fnorm(R[1])
-        else:
-            raise ValueError(f"Unsupported radius shape: {R.shape}")
-
-    cdef floats_st Rx2 = Rx * Rx
-    cdef floats_st Ry2 = Ry * Ry
+    R2 = R * R
+    hhalve = h / 2
 
     # Reset number of elements
     m = 0
 
     with nogil:
         for i in range(n):
-            if dxyz[i,nxyz] > 0.5*h or dxyz[i,nxyz] < -0.5*h: continue
-            L2 = (dxyz[i, 0]*dxyz[i, 0])/Rx2 + (dxyz[i, 1]*dxyz[i, 1])/Ry2
+            if dxyz[i,nxyz] > hhalve or dxyz[i,nxyz] < -hhalve: continue
+            # Calculate the distance of the circle
+            L2 = (dxyz[i, 0]*dxyz[i, 0])/R2 + (dxyz[i, 1]*dxyz[i, 1])/R2
             if L2 > 1.0: continue
             idx[m] = <int> i
             m += 1

src/sisl/shape/_cylinder.py

@@ -191,7 +191,7 @@ def within_index(self, other, rtol: float = 1.0e-8):
         # Get indices where we should do the more
         # expensive exact check of being inside shape
         # I.e. this reduces the search space to the box
-        return indices_in_cylinder(tmp, 1.0 + rtol, 1.0 + rtol)
+        return indices_in_cylinder(tmp, 1.0 + rtol)
 
     @deprecation(
         "toSphere is deprecated, use shape.to.Sphere(...) instead.", "0.15", "0.17"

src/sisl/shape/tests/test_cylinder.py

@@ -39,10 +39,24 @@ def test_create_ellipticalcylinder():
 
 def test_ellipticalcylinder_within():
     el = EllipticalCylinder(1.0, 1.0)
-    # center of cylinder
-    assert el.within_index([0, 0, 0])[0] == 0
+    # points in an ellipsis
+    points = [
+        [0, 0, 0],
+        [0, 0, 0.5],
+        [0, 0, -0.5],
+        [1, 0, -0.5],
+        [0, 1, -0.5],
+    ]
+    assert len(el.within_index(points)) == len(points)
+
     # should not be in a circle
-    assert el.within_index([0.2, 0.2, 0.9])[0] == 0
+    points = [
+        [0, 0, 0.6],
+        [0, 0, -0.6],
+        [0.2, 0.2, 0.9],
+        [0.2, 0.2, -0.9],
+    ]
+    assert len(el.within_index(points)) == 0
 
 
 def test_tosphere():