From 27080f0bfef6d11e3654e060c7a19dc3ab28ab2d Mon Sep 17 00:00:00 2001 From: TANG ZhiXiong Date: Thu, 9 Mar 2023 00:59:33 +0800 Subject: [PATCH] Add nanaflann KdTree, use identical interface to `scipy.spatial.cKDTree` (#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 | 1 - CMakeLists.txt | 4 +- Makefile | 15 +- docs/about/release-notes.md | 4 + fast_crossing/__init__.py | 2 + fast_crossing/cli/__init__.py | 0 fast_crossing/spatial.py | 109 ++++++++++++++ headers | 2 +- setup.py | 9 +- src/densify_polyline.hpp | 44 ------ src/main.cpp | 5 +- src/nanoflann_kdtree.hpp | 194 +++++++++++++++++++++++++ src/point_in_polygon.hpp | 226 ------------------------------ src/pybind11_nanoflann_kdtree.hpp | 73 ++++++++++ tests/test_basic.py | 104 +++++++++++++- 15 files changed, 505 insertions(+), 287 deletions(-) create mode 100644 fast_crossing/__init__.py create mode 100644 fast_crossing/cli/__init__.py create mode 100644 fast_crossing/spatial.py delete mode 100644 src/densify_polyline.hpp create mode 100644 src/nanoflann_kdtree.hpp delete mode 100644 src/point_in_polygon.hpp create mode 100644 src/pybind11_nanoflann_kdtree.hpp diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9feb0ed..8a3e827 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/CMakeLists.txt b/CMakeLists.txt index d7908ed..22fce55 100644 --- a/CMakeLists.txt +++ b/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}) diff --git a/Makefile b/Makefile index 6f57d86..44f44f4 100644 --- a/Makefile +++ b/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 diff --git a/docs/about/release-notes.md b/docs/about/release-notes.md index f89b99e..21fccd9 100644 --- a/docs/about/release-notes.md +++ b/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 diff --git a/fast_crossing/__init__.py b/fast_crossing/__init__.py new file mode 100644 index 0000000..31f4c37 --- /dev/null +++ b/fast_crossing/__init__.py @@ -0,0 +1,2 @@ +from _pybind11_fast_crossing import * # noqa +from _pybind11_fast_crossing import __version__ # noqa diff --git a/fast_crossing/cli/__init__.py b/fast_crossing/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fast_crossing/spatial.py b/fast_crossing/spatial.py new file mode 100644 index 0000000..dd1c8ec --- /dev/null +++ b/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 diff --git a/headers b/headers index 3e9c22d..1c21313 160000 --- a/headers +++ b/headers @@ -1 +1 @@ -Subproject commit 3e9c22d2747cbe5c87ad6433bfdbde95ca70b66b +Subproject commit 1c21313c6112aa597eb843f6c006683168ff1f5c diff --git a/setup.py b/setup.py index 4740242..1ef6f1b 100644 --- a/setup.py +++ b/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="dvorak4tzx@gmail.com", - 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"]}, ) diff --git a/src/densify_polyline.hpp b/src/densify_polyline.hpp deleted file mode 100644 index a321d6e..0000000 --- a/src/densify_polyline.hpp +++ /dev/null @@ -1,44 +0,0 @@ -#ifndef CUBAO_DENSIFY_POLYLINE_HPP -#define CUBAO_DENSIFY_POLYLINE_HPP - -#include - -namespace cubao -{ -inline RowVectors densify_polyline(const Eigen::Ref &coords, - double max_gap) -{ - if (coords.rows() < 2 || max_gap <= 0) { - return coords; - } - std::vector xyzs; - const int N = coords.rows(); - xyzs.reserve(N); - for (int i = 0; i < N - 1; i++) { - const Eigen::RowVector3d &curr = coords.row(i); - const Eigen::RowVector3d &next = coords.row(i + 1); - Eigen::RowVector3d dir = next - curr; - double gap = dir.norm(); - if (gap == 0) { - continue; - } - dir /= gap; - if (gap <= max_gap) { - xyzs.push_back(curr); - continue; - } - RowVectors pos(1 + int(std::ceil(gap / max_gap)), 3); - pos.col(0) = Eigen::ArrayXd::LinSpaced(pos.rows(), curr[0], next[0]); - pos.col(1) = Eigen::ArrayXd::LinSpaced(pos.rows(), curr[1], next[1]); - pos.col(2) = Eigen::ArrayXd::LinSpaced(pos.rows(), curr[2], next[2]); - int M = pos.rows() - 1; - int N = xyzs.size(); - xyzs.resize(N + M); - Eigen::Map(&xyzs[N][0], M, 3) = pos.topRows(M); - } - xyzs.push_back(coords.row(N - 1)); - return Eigen::Map(&xyzs[0][0], xyzs.size(), 3); -} -} // namespace cubao - -#endif diff --git a/src/main.cpp b/src/main.cpp index c3137ed..0059031 100644 --- a/src/main.cpp +++ b/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, // diff --git a/src/nanoflann_kdtree.hpp b/src/nanoflann_kdtree.hpp new file mode 100644 index 0000000..8a2840d --- /dev/null +++ b/src/nanoflann_kdtree.hpp @@ -0,0 +1,194 @@ +#ifndef CUBAO_NANO_KDTREE_HPP +#define CUBAO_NANO_KDTREE_HPP + +#include +#include +#include +#include + +namespace cubao +{ +using RowVectors = Eigen::Matrix; +using RowVectorsNx3 = RowVectors; +using RowVectorsNx2 = Eigen::Matrix; + +// https://github.com/jlblancoc/nanoflann/blob/master/examples/utils.h +struct PointCloud +{ + using Point = Eigen::Vector3d; + using coord_t = double; //!< The type of each coordinate + std::vector pts; + + // Must return the number of data points + inline size_t kdtree_get_point_count() const { return pts.size(); } + + // Returns the dim'th component of the idx'th point in the class: + // Since this is inlined and the "dim" argument is typically an immediate + // value, the + // "if/else's" are actually solved at compile time. + inline double kdtree_get_pt(const size_t idx, const size_t dim) const + { + if (dim == 0) + return pts[idx][0]; + else if (dim == 1) + return pts[idx][1]; + else + return pts[idx][2]; + } + + // Optional bounding-box computation: return false to default to a standard + // bbox computation loop. + // Return true if the BBOX was already computed by the class and returned + // in "bb" so it can be avoided to redo it again. Look at bb.size() to + // find out the expected dimensionality (e.g. 2 or 3 for point clouds) + template bool kdtree_get_bbox(BBOX & /* bb */) const + { + return false; + } +}; + +// https://github.com/jlblancoc/nanoflann/blob/master/examples/pointcloud_example.cpp +using KdTreeIndex = nanoflann::KDTreeSingleIndexAdaptor< + nanoflann::L2_Simple_Adaptor, // + PointCloud, // + 3 /* dim */, // + int /* AccessorType */ + >; + +struct KdTree : PointCloud +{ + KdTree(int leafsize = 10) : leafsize_(leafsize) {} + KdTree(const RowVectors &xyzs) { add(xyzs); } + KdTree(const Eigen::Ref &xys) { add(xys); } + + Eigen::Map points() const + { + return Eigen::Map(&pts[0][0], pts.size(), 3); + } + + void add(const RowVectors &xyzs) + { + int N = pts.size(); + int M = xyzs.rows(); + pts.resize(N + M); + points_(N, M) = xyzs; + index_.reset(); + } + void add(const Eigen::Ref &xys) + { + int N = pts.size(); + int M = xys.rows(); + pts.resize(N + M); + auto points = this->points_(N, M); + points.leftCols(2) = xys; + points.col(2).setZero(); + index_.reset(); + } + + void reset() + { + pts.clear(); + index_.reset(); + } + void reset_index() { index_.reset(); } + void build_index(bool force = false) { index(force); } + + int leafsize() const { return leafsize_; } + void set_leafsize(size_t value) { leafsize_ = value; } + + // https://github.com/kzampog/cilantro/tree/master + std::pair nearest(const Eigen::Vector3d &position, // + bool return_squared_l2 = false) const + { + int ret_index; + double out_dist_sqr; + auto resultSet = nanoflann::KNNResultSet(1); + resultSet.init(&ret_index, &out_dist_sqr); + index().findNeighbors(resultSet, position.data(), + nanoflann::SearchParams()); + return std::make_pair(ret_index, // + return_squared_l2 ? out_dist_sqr + : std::sqrt(out_dist_sqr)); + } + std::pair nearest(int index, // + bool return_squared_l2 = false) const + { + auto ret = nearest(pts[index], 2, false, return_squared_l2); + return ret.first[0] == index + ? std::make_pair(ret.first[1], ret.second[1]) + : std::make_pair(ret.first[0], ret.second[0]); + } + + std::pair + nearest(const Eigen::Vector3d &position, // + int k, // + bool sorted = true, // + bool return_squared_l2 = false) const + { + auto indexes = std::vector(k); + auto distances = std::vector(k); + auto resultSet = nanoflann::KNNResultSet(k); + resultSet.init(&indexes[0], &distances[0]); + auto params = nanoflann::SearchParams(); + params.sorted = sorted; + index().findNeighbors(resultSet, position.data(), params); + const int N = resultSet.size(); + return std::make_pair( + Eigen::VectorXi::Map(&indexes[0], N), + return_squared_l2 + ? Eigen::VectorXd::Map(&distances[0], N) + : Eigen::VectorXd::Map(&distances[0], N).cwiseSqrt().eval()); + } + std::pair + nearest(const Eigen::Vector3d &position, // + double radius, // + bool sorted = true, // + bool return_squared_l2 = false) const + { + auto params = nanoflann::SearchParams(); + params.sorted = sorted; + std::vector> indices_dists; + index().radiusSearch(position.data(), // + radius * radius, // + indices_dists, params); + + const int N = indices_dists.size(); + Eigen::VectorXi indexes(N); + Eigen::VectorXd distances(N); + for (int i = 0; i < N; i++) { + indexes[i] = indices_dists[i].first; + distances[i] = indices_dists[i].second; + } + return std::make_pair(indexes, return_squared_l2 + ? distances + : distances.cwiseSqrt().eval()); + } + + // + private: + Eigen::Map points_(int i, int N) + { + double *data = &pts[i][0]; + return Eigen::Map(data, N, 3); + } + Eigen::Map points_() { return points_(0, pts.size()); } + + // mutable index + // CppCon 2017: Kate Gregory “10 Core Guidelines You Need to Start Using + // Now” https://youtu.be/XkDEzfpdcSg?t=1093 + // Restoring const-correctness + mutable std::unique_ptr index_; // should be noncopyable + KdTreeIndex &index(bool force_rebuild = false) const + { + if (force_rebuild || !index_) { + index_ = std::make_unique( + 3 /*dim*/, (const PointCloud &)*this, + nanoflann::KDTreeSingleIndexAdaptorParams(leafsize_)); + } + return *index_; + } + int leafsize_ = 10; +}; +} // namespace cubao + +#endif diff --git a/src/point_in_polygon.hpp b/src/point_in_polygon.hpp deleted file mode 100644 index b16f193..0000000 --- a/src/point_in_polygon.hpp +++ /dev/null @@ -1,226 +0,0 @@ -#ifndef CUBAO_POINT_IN_POLYGON_HPP -#define CUBAO_POINT_IN_POLYGON_HPP - -#include - -namespace cubao -{ -using RowVectorsNx2 = Eigen::Matrix; -namespace agg -{ -enum path_commands -{ - path_cmd_stop = 0, //----path_cmd_stop - path_cmd_move_to = 1, //----path_cmd_move_to - path_cmd_line_to = 2, //----path_cmd_line_to - path_cmd_curve3 = 3, //----path_cmd_curve3 - path_cmd_curve4 = 4, //----path_cmd_curve4 - path_cmd_curveN = 5, //----path_cmd_curveN - path_cmd_catrom = 6, //----path_cmd_catrom - path_cmd_ubspline = 7, //----path_cmd_ubspline - path_cmd_end_poly = 0x0F, //----path_cmd_end_poly - path_cmd_mask = 0x0F //----path_cmd_mask -}; -} - -// https://github.com/matplotlib/matplotlib/blob/7303d81eb2390208d3bac67dd51d390b585e411f/src/_path.h#L67-L233 -// The following function was found in the Agg 2.3 examples -// (interactive_polygon.cpp). It has been generalized to work on (possibly -// curved) polylines, rather than just polygons. The original comments have -// been kept intact. -// -- Michael Droettboom 2007-10-02 -// -//======= Crossings Multiply algorithm of InsideTest ======================== -// -// By Eric Haines, 3D/Eye Inc, erich@eye.com -// -// This version is usually somewhat faster than the original published in -// Graphics Gems IV; by turning the division for testing the X axis crossing -// into a tricky multiplication test this part of the test became faster, -// which had the additional effect of making the test for "both to left or -// both to right" a bit slower for triangles than simply computing the -// intersection each time. The main increase is in triangle testing speed, -// which was about 15% faster; all other polygon complexities were pretty much -// the same as before. On machines where division is very expensive (not the -// case on the HP 9000 series on which I tested) this test should be much -// faster overall than the old code. Your mileage may (in fact, will) vary, -// depending on the machine and the test data, but in general I believe this -// code is both shorter and faster. This test was inspired by unpublished -// Graphics Gems submitted by Joseph Samosky and Mark Haigh-Hutchinson. -// Related work by Samosky is in: -// -// Samosky, Joseph, "SectionView: A system for interactively specifying and -// visualizing sections through three-dimensional medical image data", -// M.S. Thesis, Department of Electrical Engineering and Computer Science, -// Massachusetts Institute of Technology, 1993. -// -// Shoot a test ray along +X axis. The strategy is to compare vertex Y values -// to the testing point's Y and quickly discard edges which are entirely to one -// side of the test ray. Note that CONVEX and WINDING code can be added as -// for the CrossingsTest() code; it is left out here for clarity. -// -// Input 2D polygon _pgon_ with _numverts_ number of vertices and test point -// _point_, returns 1 if inside, 0 if outside. -inline Eigen::VectorXi -point_in_polygon(const Eigen::Ref &points, - const Eigen::Ref &polygon // or polyline -) -{ - if (!points.rows() || !polygon.rows()) { - throw std::invalid_argument("invalid polygon, or empty points"); - } - - size_t n = points.rows(); - Eigen::VectorXi inside_flag(n); - inside_flag.setZero(); - const int N = polygon.rows(); - - uint8_t yflag1; - double vtx0, vty0, vtx1, vty1; - double tx, ty; - double sx, sy; - double x, y; - bool all_done; - - std::vector yflag0(n); - std::vector subpath_flag(n); - unsigned path_id = 0; - unsigned code = 0; - unsigned vertex; - - do { - if (code != agg::path_cmd_move_to) { - if (path_id >= N) { - x = 0.0; - y = 0.0; - vertex = agg::path_cmd_stop; - } else { - x = polygon(path_id, 0); - y = polygon(path_id, 1); - vertex = path_id == 0 ? agg::path_cmd_move_to - : agg::path_cmd_line_to; - ++path_id; - } - code = vertex; - if (code == agg::path_cmd_stop || - (code & agg::path_cmd_end_poly) == agg::path_cmd_end_poly) { - continue; - } - } - - sx = vtx0 = vtx1 = x; - sy = vty0 = vty1 = y; - - for (size_t i = 0; i < n; ++i) { - ty = points(i, 1); - - if (std::isfinite(ty)) { - // get test bit for above/below X axis - yflag0[i] = (vty0 >= ty); - - subpath_flag[i] = 0; - } - } - - do { - if (path_id >= N) { - x = 0.0; - y = 0.0; - vertex = agg::path_cmd_stop; - } else { - x = polygon(path_id, 0); - y = polygon(path_id, 1); - vertex = path_id == 0 ? agg::path_cmd_move_to - : agg::path_cmd_line_to; - ++path_id; - } - code = vertex; - - // The following cases denote the beginning on a new subpath - if (code == agg::path_cmd_stop || - (code & agg::path_cmd_end_poly) == agg::path_cmd_end_poly) { - x = sx; - y = sy; - } else if (code == agg::path_cmd_move_to) { - break; - } - - for (size_t i = 0; i < n; ++i) { - tx = points(i, 0); - ty = points(i, 1); - - if (!(std::isfinite(tx) && std::isfinite(ty))) { - continue; - } - - yflag1 = (vty1 >= ty); - // Check if endpoints straddle (are on opposite sides) of - // X axis (i.e. the Y's differ); if so, +X ray could - // intersect this edge. The old test also checked whether - // the endpoints are both to the right or to the left of - // the test point. However, given the faster intersection - // point computation used below, this test was found to be - // a break-even proposition for most polygons and a loser - // for triangles (where 50% or more of the edges which - // survive this test will cross quadrants and so have to - // have the X intersection computed anyway). I credit - // Joseph Samosky with inspiring me to try dropping the - // "both left or both right" part of my code. - if (yflag0[i] != yflag1) { - // Check intersection of pgon segment with +X ray. - // Note if >= point's X; if so, the ray hits it. The - // division operation is avoided for the ">=" test by - // checking the sign of the first vertex wrto the test - // point; idea inspired by Joseph Samosky's and Mark - // Haigh-Hutchinson's different polygon inclusion - // tests. - if (((vty1 - ty) * (vtx0 - vtx1) >= - (vtx1 - tx) * (vty0 - vty1)) == yflag1) { - subpath_flag[i] ^= 1; - } - } - - // Move to the next pair of vertices, retaining info as - // possible. - yflag0[i] = yflag1; - } - - vtx0 = vtx1; - vty0 = vty1; - - vtx1 = x; - vty1 = y; - } while (code != agg::path_cmd_stop && - (code & agg::path_cmd_end_poly) != agg::path_cmd_end_poly); - - all_done = true; - for (size_t i = 0; i < n; ++i) { - tx = points(i, 0); - ty = points(i, 1); - - if (!(std::isfinite(tx) && std::isfinite(ty))) { - continue; - } - - yflag1 = (vty1 >= ty); - if (yflag0[i] != yflag1) { - if (((vty1 - ty) * (vtx0 - vtx1) >= - (vtx1 - tx) * (vty0 - vty1)) == yflag1) { - subpath_flag[i] = subpath_flag[i] ^ true; - } - } - inside_flag[i] |= subpath_flag[i]; - if (inside_flag[i] == 0) { - all_done = false; - } - } - - if (all_done) { - break; - } - } while (code != agg::path_cmd_stop); - return inside_flag; -} -} // namespace cubao - -#endif diff --git a/src/pybind11_nanoflann_kdtree.hpp b/src/pybind11_nanoflann_kdtree.hpp new file mode 100644 index 0000000..f47fdc8 --- /dev/null +++ b/src/pybind11_nanoflann_kdtree.hpp @@ -0,0 +1,73 @@ +// should sync +// - +// https://github.com/cubao/fast-crossing/blob/master/src/pybind11_nanoflann_kdtree.hpp +// - +// https://github.com/cubao/headers/tree/main/include/cubao/pybind11_nanoflann_kdtree.hpp + +#pragma once + +#include +#include +#include +#include +#include + +#include "cubao_inline.hpp" +#include "nanoflann_kdtree.hpp" + +namespace cubao +{ +namespace py = pybind11; +using namespace pybind11::literals; +using rvp = py::return_value_policy; + +CUBAO_INLINE void bind_nanoflann_kdtree(py::module &m) +{ + using KdTree = cubao::KdTree; + py::class_(m, "KdTree", py::module_local()) + .def(py::init(), "leafsize"_a = 10) + .def(py::init(), "points"_a) + .def(py::init &>(), "points"_a) + // + .def("points", &KdTree::points, rvp::reference_internal) + // + .def("add", py::overload_cast(&KdTree::add), + "points"_a) + .def("add", + py::overload_cast &>( + &KdTree::add), + "points"_a) + // + .def("reset", &KdTree::reset) + .def("reset_index", &KdTree::reset_index) + .def("build_index", &KdTree::build_index, "force_rebuild"_a = false) + // + .def("leafsize", &KdTree::leafsize) + .def("set_leafsize", &KdTree::set_leafsize, "value"_a) + // + .def("nearest", + py::overload_cast(&KdTree::nearest, + py::const_), + "position"_a, py::kw_only(), "return_squared_l2"_a = false) + .def("nearest", + py::overload_cast(&KdTree::nearest, py::const_), + "index"_a, py::kw_only(), "return_squared_l2"_a = false) + // + .def("nearest", + py::overload_cast( + &KdTree::nearest, py::const_), + "position"_a, py::kw_only(), + "k"_a, // + "sorted"_a = true, // + "return_squared_l2"_a = false) + .def("nearest", + py::overload_cast( + &KdTree::nearest, py::const_), + "position"_a, py::kw_only(), + "radius"_a, // + "sorted"_a = true, // + "return_squared_l2"_a = false) + // + ; +} +} // namespace cubao diff --git a/tests/test_basic.py b/tests/test_basic.py index 2a02628..8f20870 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,6 +1,10 @@ +import os +import sys + import numpy as np import pytest -from fast_crossing import FastCrossing, densify_polyline, point_in_polygon + +from fast_crossing import FastCrossing, KdTree, densify_polyline, point_in_polygon def test_fast_crossing(): @@ -230,3 +234,101 @@ def test_point_in_polygon(): points = [[0.5, 0.5], [10, 0]] mask = point_in_polygon(points=points, polygon=polygon) assert np.all(mask == [1, 0]) + + +def test_kdtree(): + xs = np.linspace(0, 10, 101) + xyzs = np.zeros((len(xs), 3)) + xyzs[:, 0] = xs + tree = KdTree(xyzs) + assert tree.points().shape == (101, 3) + idx, dist = tree.nearest(0) + assert idx == 1, dist == 0.1 + idx, dist = tree.nearest([0.0, 0.0, 0.0]) + assert idx == 0, dist == 0.0 + idx, dist = tree.nearest([0.0, 0.0, 0.0], k=4) + assert np.all(idx == [0, 1, 2, 3]) + np.testing.assert_allclose(dist, [0.0, 0.1, 0.2, 0.3], atol=1e-15) + idx, dist = tree.nearest([0.0, 0.0, 0.0], radius=0.25) + assert np.all(idx == [0, 1, 2]) + np.testing.assert_allclose(dist, [0.0, 0.1, 0.2], atol=1e-15) + + +def _test_cKDTree_query(KDTree): + x, y = np.mgrid[0:5, 2:8] + tree = KDTree(np.c_[x.ravel(), y.ravel()]) + + expected_k1 = [[2.0, 0.2236068], [0, 13]] + expected_k2 = [[[2.0, 2.23606798], [0.2236068, 0.80622577]], [[0, 6], [13, 19]]] + for k, expected in zip([1, 2], [expected_k1, expected_k2]): # noqa + dd, ii = tree.query([[0, 0], [2.2, 2.9]], k=k) + DD, II = expected + np.testing.assert_allclose(dd, DD, atol=1e-6) + np.testing.assert_allclose(ii, II, atol=1e-6) + + expected_k1 = [ + [[2.0], [0.22360679774997916]], + [[0], [13]], + ] + expected_k2 = [ + [[2.23606797749979], [0.8062257748298548]], + [[6], [19]], + ] + expected_k1_k2 = [ + [[2.0, 2.23606797749979], [0.22360679774997916, 0.8062257748298548]], + [[0, 6], [13, 19]], + ] + for k, expected in zip( # noqa + [[1], [2], [1, 2]], [expected_k1, expected_k2, expected_k1_k2] + ): + dd, ii = tree.query([[0, 0], [2.2, 2.9]], k=k) + DD, II = expected + np.testing.assert_allclose(dd, DD, atol=1e-6) + np.testing.assert_allclose(ii, II, atol=1e-6) + + x, y = np.mgrid[0:4, 0:4] + points = np.c_[x.ravel(), y.ravel()] + tree = KDTree(points) + ret = sorted(tree.query_ball_point([2 + 1e-15, 1e-15], 1 + 1e-9)) + assert np.all([4, 8, 9, 12] == np.array(ret)) + + ret = tree.query_ball_point([[2, 0], [3, 0]], 1 + 1e-9) + assert np.all([4, 8, 9, 12] == np.array(sorted(ret[0]))) + assert np.all([8, 12, 13] == np.array(sorted(ret[1]))) + + ret = tree.query_ball_point([[2, 0], [3, 0]], 1 + 1e-9, return_length=True) + assert np.all([4, 3] == np.array(ret)) + + +def test_scipy_cKDTree(): + from scipy.spatial import cKDTree + + _test_cKDTree_query(cKDTree) + + +def test_nanoflann_KDTree(): + from fast_crossing.spatial import KDTree + + _test_cKDTree_query(KDTree) + + +def pytest_main(dir: str, *, test_file: str = None): + + os.chdir(dir) + sys.exit( + pytest.main( + [ + dir, + *(["-k", test_file] if test_file else []), + "--capture", + "tee-sys", + "-vv", + "-x", + ] + ) + ) + + +if __name__ == "__main__": + pwd = os.path.abspath(os.path.dirname(__file__)) + pytest_main(pwd, test_file=os.path.basename(__file__))