Add nanaflann KdTree, use identical interface to `scipy.spatial.cKDTr…

…ee` (#8)

* benchmark N points

* sync code

* add nanokdtree

* nanokdtree

* kdtree

* not ready

* kdtree works

* update

* update

* bind search

* test kdtree

* integrated python part

* not ready

* not ready

* same interface

* fix#

* fix

* ready to release

* rename to nanoflann kdtree

* update headers

* add scipy

.pre-commit-config.yaml

@@ -32,7 +32,6 @@ repos:
   - id: mixed-line-ending
   - id: requirements-txt-fixer
   - id: trailing-whitespace
-  - id: end-of-file-fixer
 
 # Black, the code formatter, natively supports pre-commit
 - repo: https://github.com/psf/black

CMakeLists.txt

@@ -27,7 +27,7 @@ set(CMAKE_CXX_STANDARD 17)
 set(PYBIND11_CPP_STANDARD -std=c++17)
 
 add_subdirectory(pybind11)
-pybind11_add_module(fast_crossing src/main.cpp)
+pybind11_add_module(_pybind11_fast_crossing src/main.cpp)
 
-target_compile_definitions(fast_crossing
+target_compile_definitions(_pybind11_fast_crossing
                            PRIVATE VERSION_INFO=${FAST_CROSSING_VERSION_INFO})

Makefile

@@ -85,19 +85,20 @@ tar.gz:
 	tar -cvz --exclude .git -f ../fast_crossing.tar.gz .
 	ls -alh ../fast_crossing.tar.gz
 
+NUM_POINTS = 100000
 benchmark_point_in_polygon:
-	python3 benchmarks/benchmark_point_in_polygon.py generate_test_data -o dist/point_in_polygon
+	python3 benchmarks/benchmark_point_in_polygon.py generate_test_data --num=$(NUM_POINTS) -o dist/point_in_polygon
 	python3 benchmarks/benchmark_point_in_polygon.py shapely \
-		dist/point_in_polygon/random_num_10000__bbox_800.00x600.00__radius_250.00__points.npy \
-		dist/point_in_polygon/random_num_10000__bbox_800.00x600.00__radius_250.00__polygon.npy \
+		dist/point_in_polygon/random_num_$(NUM_POINTS)__bbox_800.00x600.00__radius_250.00__points.npy \
+		dist/point_in_polygon/random_num_$(NUM_POINTS)__bbox_800.00x600.00__radius_250.00__polygon.npy \
 		dist/mask_shapely.npy
 	python3 benchmarks/benchmark_point_in_polygon.py matplotlib \
-		dist/point_in_polygon/random_num_10000__bbox_800.00x600.00__radius_250.00__points.npy \
-		dist/point_in_polygon/random_num_10000__bbox_800.00x600.00__radius_250.00__polygon.npy \
+		dist/point_in_polygon/random_num_$(NUM_POINTS)__bbox_800.00x600.00__radius_250.00__points.npy \
+		dist/point_in_polygon/random_num_$(NUM_POINTS)__bbox_800.00x600.00__radius_250.00__polygon.npy \
 		dist/mask_matplotlib.npy
 	python3 benchmarks/benchmark_point_in_polygon.py cubao \
-		dist/point_in_polygon/random_num_10000__bbox_800.00x600.00__radius_250.00__points.npy \
-		dist/point_in_polygon/random_num_10000__bbox_800.00x600.00__radius_250.00__polygon.npy \
+		dist/point_in_polygon/random_num_$(NUM_POINTS)__bbox_800.00x600.00__radius_250.00__points.npy \
+		dist/point_in_polygon/random_num_$(NUM_POINTS)__bbox_800.00x600.00__radius_250.00__polygon.npy \
 		dist/mask_cubao.npy
 .PHONY: benchmark_point_in_polygon
 

docs/about/release-notes.md

@@ -10,6 +10,10 @@ To upgrade `fast-crossing` to the latest version, use pip:
 pip install -U fast-crossing
 ```
 
+## Version 0.0.5 (2023-03-08)
+
+*   Add nanaflann KdTree, use identical interface to `scipy.spatial.cKDTree`
+
 ## Version 0.0.4 (2023-03-07)
 
 *   Integrate `point_in_polygon` test

fast_crossing/__init__.py

@@ -0,0 +1,2 @@
+from _pybind11_fast_crossing import *  # noqa
+from _pybind11_fast_crossing import __version__  # noqa

fast_crossing/spatial.py

@@ -0,0 +1,109 @@
+import numpy as np
+from _pybind11_fast_crossing import KdTree as _KdTree
+
+
+class KDTree:
+    def __init__(self, data: np.ndarray, leafsize: int = 10, *args, **kwargs):
+        data = np.asarray(data, dtype=np.float64)
+        self.tree: _KdTree = _KdTree(data)
+        self.tree.set_leafsize(leafsize)
+
+    @staticmethod
+    def vec3(arr: np.ndarray):
+        return np.r_[arr, 0.0] if len(arr) == 2 else np.asarray(arr, dtype=np.float64)
+
+    def count_neighbors(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def query(self, x, k=1, *args, **kwargs):
+        x = np.asarray(x, dtype=np.float64)
+        if x.ndim == 1:
+            xyz = self.vec3(x)
+            ii, dd = self.tree.nearest(xyz, k=k)
+            return dd, ii
+        if isinstance(k, (int, np.integer)):
+            ret_ii, ret_dd = [], []
+            for xyz in x:
+                xyz = self.vec3(xyz)
+                if k == 1:
+                    ii, dd = self.tree.nearest(xyz)
+                else:
+                    ii, dd = self.tree.nearest(xyz, k=k)
+                    ii = ii.tolist()
+                    dd = dd.tolist()
+                ret_ii.append(ii)
+                ret_dd.append(dd)
+            return ret_dd, ret_ii
+        K = max(k)
+        ret_ii, ret_dd = [], []
+        for xyz in x:
+            xyz = self.vec3(xyz)
+            ii, dd = self.tree.nearest(xyz, k=K)
+            ii = [ii[kk - 1] for kk in k]
+            dd = [dd[kk - 1] for kk in k]
+            ret_ii.append(ii)
+            ret_dd.append(dd)
+        return ret_dd, ret_ii
+
+    def query_ball_point(
+        self,
+        x,
+        r,
+        p=2.0,
+        eps=0,
+        workers=1,
+        return_sorted=None,
+        return_length=False,
+        *args,
+        **kwargs,
+    ):
+        """
+        https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.cKDTree.query_ball_point.html#scipy.spatial.cKDTree.query_ball_point
+        """
+        x = np.asarray(x, dtype=np.float64)
+        if return_sorted is None:
+            # If None, does not sort single point queries, but does sort
+            # multi-point queries which was the behavior before this option was
+            # added.
+            return_sorted = x.ndim != 1
+        if x.ndim == 1:
+            xyz = self.vec3(x)
+            ii, dd = self.tree.nearest(
+                xyz,
+                radius=r,
+                return_squared_l2=True,
+                sorted=return_sorted,
+            )
+            if return_length:
+                return len(ii)
+            return ii.tolist()
+        if return_sorted is None:
+            return_sorted = True
+        if isinstance(r, (int, float, np.number)):
+            r = [r] * len(x)
+        ret_ii = []
+        for pp, rr in zip(x, r):  # noqa
+            xyz = self.vec3(pp)
+            ii, dd = self.tree.nearest(
+                xyz,
+                radius=rr,
+                return_squared_l2=True,
+                sorted=return_sorted,
+            )
+            ret_ii.append(ii.tolist())
+        if return_length:
+            ret_ii = [len(ii) for ii in ret_ii]
+        return ret_ii
+
+    def query_ball_tree(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def query_pairs(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def query_distance_matrix(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+# create alias
+cKDTree = KDTree

setup.py

@@ -3,7 +3,7 @@
 import subprocess
 import sys
 
-from setuptools import Extension, setup
+from setuptools import Extension, find_packages, setup
 from setuptools.command.build_ext import build_ext
 
 # Convert distutils Windows platform specifiers to CMake -A arguments
@@ -124,16 +124,17 @@ def build_extension(self, ext):
 # logic and declaration, and simpler if you include description/version in a file.
 setup(
     name="fast_crossing",
-    version="0.0.4",
+    version="0.0.5",
     author="tzx",
     author_email="[email protected]",
-    url="https://github.com/cubao/fast-crossing",
+    url="https://fast-crossing.readthedocs.io",
     description="fast crossing",
     long_description=open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
+    packages=find_packages(),
     ext_modules=[CMakeExtension("fast_crossing")],
     cmdclass={"build_ext": CMakeBuild},
     zip_safe=False,
     install_requires=["numpy"],
-    extras_require={"test": ["pytest>=6.0"]},
+    extras_require={"test": ["pytest>=6.0", "scipy"]},
 )

src/main.cpp

@@ -18,6 +18,8 @@
 #include "fast_crossing.hpp"
 #include "pybind11_fast_crossing.hpp"
 #include "pybind11_flatbush.hpp"
+#include "pybind11_nanoflann_kdtree.hpp"
+
 #include "pybind11_polyline_ruler.hpp"
 
 #include "point_in_polygon.hpp"
@@ -29,10 +31,11 @@
 namespace py = pybind11;
 using namespace pybind11::literals;
 
-PYBIND11_MODULE(fast_crossing, m)
+PYBIND11_MODULE(_pybind11_fast_crossing, m)
 {
     cubao::bind_fast_crossing(m);
     cubao::bind_flatbush(m);
+    cubao::bind_nanoflann_kdtree(m);
     cubao::bind_polyline_ruler(m);
 
     m.def("point_in_polygon", &cubao::point_in_polygon, //