sketching cython

setup.py

@@ -214,6 +214,7 @@ def run(self):
                     _ospath("sisl/physics/_matrix_utils.pxd")
         ]
     },
+    "sisl.geom.neighs.cneighs": {},
 }
 
 # List of extensions for setup(...)

sisl/geom/neighs/cneighs.pyx

@@ -0,0 +1,57 @@
+import numpy as np
+# This enables Cython enhanced compatibilities
+cimport numpy as np
+
+def build_table(np.int64_t nbins, np.ndarray[np.int64_t, ndim=1, mode='c'] bin_indices):
+    """Builds the table where the location of all atoms is encoded
+    
+    Additionally, it also builds an array with the number of atoms at
+    each bin.
+    
+    Parameters
+    ----------
+    nbins: int
+        Total number of bins
+    bin_indices: array of int
+        An array containing the bin index of each atom.
+        
+    Returns
+    ----------
+    list: 
+        contains the list of atoms, modified to encode all bin
+        locations. Each item in the list contains the index of
+        the next atom that we can find in the same bin. 
+        Indices are 0-indexed (python indices). If an item is
+        -1, that there are no more atoms in the same bin.
+    heads: 
+        For each bin, the index of `list` where we can find the
+        first atom that is contained in it.
+    counts:
+        For each bin, the number of atoms that it contains.
+    """
+    cdef np.int64_t at, bin_index
+
+    cdef np.int64_t Nat = bin_indices.shape[0]
+
+    # Initialize heads and counts arrays. We don't need to initialize the list array
+    # since we are going to modify all its positions.
+    cdef np.int64_t[:] list_array = np.zeros(Nat, dtype=np.int64)
+
+    cdef np.int64_t[:] counts = np.zeros(Nat, dtype=np.int64)
+    cdef np.int64_t[:] heads = np.ones(Nat, dtype=np.int64) * -1
+
+    # Loop through all atoms
+    for at in range(Nat):
+        # Get the index of the bin where this atom is located.
+        bin_index = bin_indices[at]
+
+        # Replace the head of this bin by the current atom index.
+        # Before replacing, store the previous head in this atoms'
+        # position in the list.
+        list_array[at] = heads[bin_index]
+        heads[bin_index] = at
+
+        # Update the count of this bin (increment it by 1).
+        counts[bin_index] = counts[bin_index] + 1
+
+    return list_array, heads, counts

sisl/geom/neighs/neighfinder.py

@@ -3,6 +3,7 @@
 
 from sisl.utils.mathematics import fnorm
 from . import _fneighs
+from . import cneighs
 
 
 class NeighFinder: