From fe0bdf4ba042c95e5ff89e3c8c06343e5ef7051a Mon Sep 17 00:00:00 2001
From: Nick Papior <nickpapior@gmail.com>
Date: Fri, 11 Aug 2023 10:06:09 +0200
Subject: [PATCH] fixed index checks for cylinders

Signed-off-by: Nick Papior <nickpapior@gmail.com>
---
 src/sisl/_indices.pyx                 | 60 +++++++++++++++++++++++++++
 src/sisl/shape/_cylinder.py           |  5 ++-
 src/sisl/shape/tests/test_cylinder.py |  8 ++++
 3 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/src/sisl/_indices.pyx b/src/sisl/_indices.pyx
index 7f12ae959..50f53bd1f 100644
--- a/src/sisl/_indices.pyx
+++ b/src/sisl/_indices.pyx
@@ -173,6 +173,66 @@ cdef void _indices_sorted_arrays(const Py_ssize_t n_element, const int[::1] elem
     for j in range(j, n_test_element):
         idx[j] = -1
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def indices_in_cylinder(np.ndarray[np.float64_t, ndim=2, mode='c'] dxyz, const double R, const double h):
+    """ Indices for all coordinates that are within a cylinde radius `R` and height `h`
+
+    Parameters
+    ----------
+    dxyz : ndarray(np.float64)
+       coordinates centered around the cylinder
+    R : float
+       radius of cylinder to check
+    h : float
+       height of cylinder to check
+
+    Returns
+    -------
+    index : np.ndarray(np.int32)
+       indices of all dxyz coordinates that are within the cylinder
+    """
+    cdef double[:, ::1] dXYZ = dxyz
+    cdef Py_ssize_t n = dXYZ.shape[0]
+    cdef np.ndarray[np.int32_t, ndim=1] idx = np.empty([n], dtype=np.int32)
+    cdef int[::1] IDX = idx
+
+    n = _indices_in_cylinder(dXYZ, R, h, IDX)
+
+    if n == 0:
+        return np.empty([0], dtype=np.int32)
+    return idx[:n].copy()
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+@cython.initializedcheck(False)
+cdef Py_ssize_t _indices_in_cylinder(const double[:, ::1] dxyz, const double R, const double h, int[::1] idx) nogil:
+    cdef Py_ssize_t N = dxyz.shape[0]
+    cdef Py_ssize_t xyz = dxyz.shape[1]
+    cdef double R2 = R * R
+    cdef double L2
+    cdef Py_ssize_t i, j, n
+    cdef int skip
+
+    # Reset number of elements
+    n = 0
+
+    for i in range(N):
+        skip = 0
+        for j in range(xyz-1):
+            skip |= dxyz[i, j] > R
+        if skip or dxyz[i, -1] > h: continue
+
+        L2 = 0.
+        for j in range(xyz-1):
+            L2 += dxyz[i, j] * dxyz[i, j]
+        if L2 > R2: continue
+        idx[n] = i
+        n += 1
+
+    return n
+
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
diff --git a/src/sisl/shape/_cylinder.py b/src/sisl/shape/_cylinder.py
index a4194c124..013ed8472 100644
--- a/src/sisl/shape/_cylinder.py
+++ b/src/sisl/shape/_cylinder.py
@@ -6,7 +6,7 @@
 import numpy as np
 
 import sisl._array as _a
-from sisl._indices import indices_in_sphere
+from sisl._indices import indices_in_cylinder
 from sisl._internal import set_module
 from sisl.messages import warn
 from sisl.utils.mathematics import expand, fnorm, fnorm2, orthogonalize
@@ -152,7 +152,7 @@ def within_index(self, other, tol=1.e-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_sphere(tmp, 1. + tol)
+        return indices_in_cylinder(tmp, 1. + tol, 1. + tol)
 
     @property
     def height(self):
@@ -177,6 +177,7 @@ def height_vector(self):
     def toSphere(self):
         """ Convert to a sphere """
         from .ellipsoid import Sphere
+
         # figure out the distance from the center to the edge (along longest radius)
         h = self.height / 2
         r = self.radius.max()
diff --git a/src/sisl/shape/tests/test_cylinder.py b/src/sisl/shape/tests/test_cylinder.py
index 91ccf76ea..bd2132b96 100644
--- a/src/sisl/shape/tests/test_cylinder.py
+++ b/src/sisl/shape/tests/test_cylinder.py
@@ -36,6 +36,14 @@ def test_create_ellipticalcylinder():
     str(el)
 
 
+def test_ellipticalcylinder_within():
+    el = EllipticalCylinder(1., 1.)
+    # center of cylinder
+    assert el.within_index([0, 0, 0])[0] == 0
+    # should not be in a circle
+    assert el.within_index([0.2, 0.2, 0.9])[0] == 0
+
+
 def test_tosphere():
     el = EllipticalCylinder([1., 1.], 1.)
     el.to.Sphere()