diff --git a/setup.py b/setup.py index 154011de4..5603d9fd1 100644 --- a/setup.py +++ b/setup.py @@ -214,6 +214,7 @@ def run(self): _ospath("sisl/physics/_matrix_utils.pxd") ] }, + "sisl.geom.neighs.cneighs": {}, } # List of extensions for setup(...) diff --git a/sisl/geom/neighs/cneighs.pyx b/sisl/geom/neighs/cneighs.pyx new file mode 100644 index 000000000..2eb789af1 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/sisl/geom/neighs/neighfinder.py b/sisl/geom/neighs/neighfinder.py index 4881ae717..4f4341309 100644 --- a/sisl/geom/neighs/neighfinder.py +++ b/sisl/geom/neighs/neighfinder.py @@ -3,6 +3,7 @@ from sisl.utils.mathematics import fnorm from . import _fneighs +from . import cneighs class NeighFinder: