Linear and efficient neighbour finder

CITATION.cff

@@ -9,6 +9,9 @@ authors:
     given-names: Nick
     orcid: "https://orcid.org/0000-0003-3038-1855"
     alias: zerothi
+  - family-names: Febrer
+    given-names: Pol
+    alias: pfebrer
 title: sisl
 identifiers:
   - description: "Collection DOI for any sisl version."

CMakeLists.txt

@@ -26,7 +26,19 @@ find_package(Python COMPONENTS
   REQUIRED)
 
 # Retrive the cython executable
+message(STATUS "Locating Cython executable")
 find_program(CYTHON_EXECUTABLE "cython" REQUIRED)
+execute_process(COMMAND
+  "${CYTHON_EXECUTABLE}" --version
+  OUTPUT_QUIET
+  OUTPUT_VARIABLE CYTHON_VERSION
+  RESULT_VARIABLE CYTHON_VERSION_ERR
+  OUTPUT_STRIP_TRAILING_WHITESPACE)
+# Print it out
+message(STATUS "Cython version information [${CYTHON_EXECUTABLE} --version]: ${CYTHON_VERSION}")
+if(CYTHON_VERSION_ERR)
+  message(SEND_ERROR "ERROR: Cython version detection failed, this might be problematic")
+endif()
 
 
 # Define global definitions
@@ -287,7 +299,7 @@ function(add_cython_library)
   if(NOT DEFINED _c_SOURCE)
     list(GET _c_UNPARSED_ARGUMENTS 0 _c_SOURCE)
   endif()
-  
+
   # Parse simple flags
   if(NOT DEFINED _c_CYTHON_EXECUTABLE)
     set(_c_CYTHON_EXECUTABLE "${CYTHON_EXECUTABLE}")
@@ -617,7 +629,7 @@ function(add_f2py_library)
     VERBATIM
     COMMENT "Generating module file from signature ${_f_SIGNATURE}"
   )
-    
+
   python_add_library(${_f_LIBRARY} WITH_SOABI
     MODULE "${_module_file}" ${_f_SOURCES} ${_wrappers}
     )

benchmarks/comparisons/neighbors.py

@@ -0,0 +1,90 @@
+import timeit
+from collections import defaultdict
+
+from ase.neighborlist import neighbor_list
+
+from sisl.geom import NeighborFinder, fcc, graphene_nanoribbon
+
+
+def quadratic(geom):
+    neighs = []
+    for at in geom:
+        neighs.append(geom.close(at, R=1.5))
+
+
+what = "fcc"
+times = defaultdict(list)
+na = []
+
+# Iterate over multiple tiles
+for tiles in range(1, 11):
+    print(tiles)
+
+    if what == "fcc":
+        geom = fcc(1.5, "C").tile(3 * tiles, 1).tile(3 * tiles, 2).tile(3 * tiles, 0)
+    elif what == "ribbon":
+        geom = graphene_nanoribbon(9).tile(tiles, 0)
+
+    na.append(geom.na)
+
+    # Compute with ASE
+    ase_time = timeit.timeit(
+        lambda: neighbor_list("ij", geom.to.ase(), 1.5, self_interaction=False),
+        number=1,
+    )
+    times["ASE"].append(ase_time)
+
+    for bs in [4, 8, 12]:
+        print("  bs = ", bs)
+
+        # Compute with this implementation (different variants)
+        my_time = timeit.timeit(
+            lambda: NeighborFinder(geom, R=1.5, bin_size=bs).find_neighbors(
+                None, as_pairs=True
+            ),
+            number=1,
+        )
+        times[f"(PAIRS) [{bs}]"].append(my_time)
+
+        my_time = timeit.timeit(
+            lambda: NeighborFinder(geom, R=1.5, bin_size=bs).find_neighbors(
+                None, as_pairs=False
+            ),
+            number=1,
+        )
+        times[f"(NEIGHS) [{bs}]"].append(my_time)
+
+        my_time = timeit.timeit(
+            lambda: NeighborFinder(geom, R=1.5, bin_size=bs).find_neighbors(
+                None, as_pairs=False
+            ),
+            number=1,
+        )
+        times[f"(NEIGHS) [{bs}]"].append(my_time)
+
+        my_time = timeit.timeit(
+            lambda: NeighborFinder(geom, R=1.5, bin_size=bs).find_all_unique_pairs(),
+            number=1,
+        )
+        times[f"(UNIQUE PAIRS) [{bs}]"].append(my_time)
+
+    # Compute with quadratic search
+    # quadratic_time = timeit.timeit(lambda: quadratic(geom), number=1)
+    # times["QUADRATIC"].append(quadratic_time)
+
+# Plotting
+import plotly.graph_objs as go
+
+fig = go.Figure()
+for key, time in times.items():
+    fig.add_scatter(x=na, y=time, name=key)
+
+fig.update_layout(
+    xaxis_title="Number of atoms",
+    yaxis_title="Time (s)",
+    title=f"Finding neighbours on {what}",
+    yaxis_showgrid=True,
+    xaxis_showgrid=True,
+)
+fig.write_image(f"neighbors_{what}.png")
+fig.show()

ci/doc.yml

@@ -12,7 +12,7 @@ dependencies:
   - pyproject-metadata
   - pyparsing>=1.5.7
   - numpy>=1.13
-  - cython>=0.29.28
+  - cython>=3.0.0
   - scipy>=1.5.0
   - netcdf4
   - xarray>=0.10.0

docs/api/basic.rst

@@ -66,6 +66,6 @@ In particular `oplist` is useful when calculating averages in Brillouin zones (s
 .. autosummary::
    :toctree: generated/
 
+   NeighborFinder
    oplist
    ~sisl.utils.PropertyDict
-

pyproject.toml

@@ -12,7 +12,7 @@
 requires = [
     "setuptools_scm[toml]>=6.2",
     "scikit-build-core[pyproject]",
-    "Cython>=0.29.28",
+    "Cython>=3",
     # see https://github.com/scipy/oldest-supported-numpy/
     # should fix #310
     "oldest-supported-numpy; sys_platform != 'win32'",
@@ -49,7 +49,8 @@ dependencies = [
 ]
 
 authors = [
-    {name = "Nick Papior", email = "[email protected]"}
+    {name = "Nick Papior", email = "[email protected]"},
+    {name = "Pol Febrer"}
 ]
 maintainers = [{name="sisl developers"}]
 
@@ -136,9 +137,9 @@ analysis = [
 ]
 
 viz = [
+    "netCDF4",
     "dill >= 0.3.2",
     "pathos",
-    "netCDF4",
     "scikit-image",
 
     "plotly",
@@ -147,34 +148,34 @@ viz = [
 ]
 
 viz-plotly = [
+    "netCDF4",
     "dill >= 0.3.2",
     "pathos",
-    "netCDF4",
     "scikit-image",
 
     "plotly",
 ]
 
 viz-matplotlib = [
+    "netCDF4",
     "dill >= 0.3.2",
     "pathos",
-    "netCDF4",
     "scikit-image",
 
     "matplotlib",
 ]
 
 viz-blender = [
+    "netCDF4",
     "dill >= 0.3.2",
     "pathos",
-    "netCDF4",
     "scikit-image",
 ]
 
 viz-ase = [
+    "netCDF4",
     "dill >= 0.3.2",
     "pathos",
-    "netCDF4",
     "scikit-image",
 
     "ase",
@@ -186,10 +187,10 @@ test = [
     "pytest-env",
     "pytest-faulthandler",
     "coveralls",
+    "netCDF4",
     "tqdm",
     "dill >= 0.3.2",
     "pathos",
-    "netCDF4",
     "scikit-image",
     "matplotlib",
     "plotly",

src/sisl/CMakeLists.txt

@@ -12,3 +12,4 @@ endforeach()
 add_subdirectory("_core")
 add_subdirectory("io")
 add_subdirectory("physics")
+add_subdirectory("geom")

src/sisl/conftest.py

@@ -227,6 +227,7 @@ def pytest_configure(config):
         "hamiltonian",
         "geometry",
         "geom",
+        "neighbor",
         "shape",
         "state",
         "electron",

src/sisl/geom/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory("_neighbors")

src/sisl/geom/__init__.py

@@ -43,6 +43,7 @@
 
 """
 from ._composite import *
+from ._neighbors import *
 from .basic import *
 from .bilayer import *
 from .category import *

src/sisl/geom/_neighbors/CMakeLists.txt

@@ -0,0 +1,11 @@
+# In this directory we have a set of libraries
+# We will need to link to the Numpy includes
+foreach(source _operations)
+  add_cython_library(
+    SOURCE ${source}.py
+    LIBRARY ${source}
+    OUTPUT ${source}_C
+    )
+  install(TARGETS ${source} LIBRARY
+    DESTINATION ${SKBUILD_PROJECT_NAME}/geom/_neighbors)
+endforeach()

src/sisl/geom/_neighbors/__init__.py

@@ -0,0 +1 @@
+from ._finder import NeighborFinder