Refactoring of _pyco_tree:

* Improved compile time using type erasure.
* Added the save_kd_tree and load_kd_tree functions.
* Replaced several std::runtime_error exceptions by std::invalid_argument.
Updated the minimum Python version to 3.10.
Avoid global namespace redeclaration of index error with GNU extensions.

Author: Jaybro

.github/workflows/pip.yml

@@ -18,25 +18,26 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v4
         with:
-          python-version: "3.7"
+          python-version: "3.10"
 
-      # On macos-latest, OpenMP is already installed but not visible to CMake.
-      # Library libomp 15.0+ has been made keg-only.
-      # https://github.com/Homebrew/homebrew-core/issues/112107#issuecomment-1278042927
+      # On macos-latest, OpenMP is not installed and not visible to CMake.
       - name: Setup OpenMP for macOS
         if: runner.os == 'macOS'
-        run: brew link --force libomp
+        run: |
+          brew install libomp 
+          brew link --force libomp
 
       # The global version of pybind11 allows us to find it via CMake.
       - name: Update setup related tools
         run: |
-          python -m pip install --upgrade wheel setuptools
-          pip install ninja
-          pip install scikit-build
-          pip install pybind11-global
+          python -m pip install wheel
+          python -m pip install setuptools
+          python -m pip install ninja
+          python -m pip install scikit-build
+          python -m pip install pybind11-global
 
       - name: Install with pip
-        run: pip install ./ -v
+        run: python -m pip install ./ -v
 
       - name: Test Python bindings
         run: python -m unittest discover -s ./test/pyco_tree -p '*_test.py' -v

README.md

@@ -80,7 +80,7 @@ Optional:
 * [Google Benchmark](https://github.com/google/benchmark) is needed to run any of the benchmarks. The [nanoflann](https://github.com/jlblancoc/nanoflann) and [OpenCV](https://opencv.org/) benchmarks also require their respective libraries to be installed.
 
 Python bindings:
-* [Python](https://www.python.org/). Version 3.7 or higher.
+* [Python](https://www.python.org/). Version 3.10 or higher.
 * [pybind11](https://github.com/pybind/pybind11). Used to ease the creation of Python bindings. Available under the [BSD](https://github.com/pybind/pybind11/blob/master/LICENSE) license and copyright.
 * [OpenMP](https://www.openmp.org/). For parallelization of queries.
 * [numpy](https://numpy.org/). Points and search results are represented by ndarrays.

examples/eigen/eigen.cpp

@@ -84,14 +84,16 @@ void col_major_support() {
   using kd_tree = pico_tree::kd_tree<map>;
   using neighbor = typename kd_tree::neighbor_type;
   using scalar = typename point::Scalar;
-  constexpr int dim = point::RowsAtCompileTime;
+  constexpr Eigen::Index dim = point::RowsAtCompileTime;
 
   auto points = generate_random_eigen_n<point>(num_points, point_area);
   point p = point::Random() * point_area / scalar(2.0);
 
   std::cout << "Eigen RowMajor: " << map::IsRowMajor << std::endl;
 
-  kd_tree tree(map(points.data()->data(), dim, points.size()), max_leaf_count);
+  kd_tree tree(
+      map(points.data()->data(), dim, static_cast<Eigen::Index>(num_points)),
+      max_leaf_count);
 
   std::vector<neighbor> knn;
   pico_tree::scoped_timer t("pico_tree col major", run_count);
@@ -108,14 +110,16 @@ void row_major_support() {
   using kd_tree = pico_tree::kd_tree<map>;
   using neighbor = typename kd_tree::neighbor_type;
   using scalar = typename point::Scalar;
-  constexpr int dim = point::ColsAtCompileTime;
+  constexpr Eigen::Index dim = point::ColsAtCompileTime;
 
   auto points = generate_random_eigen_n<point>(num_points, point_area);
   point p = point::Random() * point_area / scalar(2.0);
 
   std::cout << "Eigen RowMajor: " << point::IsRowMajor << std::endl;
 
-  kd_tree tree(map(points.data()->data(), points.size(), dim), max_leaf_count);
+  kd_tree tree(
+      map(points.data()->data(), static_cast<Eigen::Index>(num_points), dim),
+      max_leaf_count);
 
   std::vector<neighbor> knn;
   pico_tree::scoped_timer t("pico_tree row major", run_count);

examples/kd_tree/kd_tree_search.cpp

@@ -93,7 +93,7 @@ void search_s1() {
   std::cout << "Closest angles (index, distance, value): " << std::endl;
   for (auto const& nn : knn) {
     std::cout << "  " << nn.index << ", " << nn.distance << ", "
-              << tree.space()[nn.index] << std::endl;
+              << tree.space()[static_cast<std::size_t>(nn.index)] << std::endl;
   }
 }
 

examples/opencv/opencv.cpp

@@ -4,11 +4,8 @@
 #include <pico_tree/vector_traits.hpp>
 #include <random>
 
-using index = int;
-using scalar = float;
-
-index const num_points = 1024 * 1024 * 2;
-scalar const point_area = 1000.0;
+std::size_t const num_points = 1024 * 1024 * 2;
+float const point_area = 1000.0;
 std::size_t const run_count = 1024 * 1024;
 
 template <typename Vec_>
@@ -30,6 +27,8 @@ std::vector<Vec_> generate_random_vec_n(
 
 // This example shows to build a kd_tree from a vector of cv::Point3.
 void basic_vector() {
+  using index = int;
+  using scalar = float;
   using point = cv::Vec<scalar, 3>;
   std::vector<point> random =
       generate_random_vec_n<point>(num_points, point_area);
@@ -47,6 +46,9 @@ void basic_vector() {
 
 // This example shows to build a kd_tree using a cv::Mat.
 void basic_matrix() {
+  using index = int;
+  using scalar = float;
+
   // Multiple columns based on the number of coordinates in a point.
   {
     constexpr int dim = 3;
@@ -55,7 +57,7 @@ void basic_matrix() {
 
     pico_tree::kd_tree<pico_tree::opencv_mat_map<scalar, dim>> tree(
         random, pico_tree::max_leaf_size_t(10));
-    pico_tree::point_map<scalar, dim> p = tree.space()[random.rows / 2];
+    pico_tree::point_map<scalar, dim> p = tree.space()[tree.space().size() / 2];
 
     pico_tree::neighbor<index, scalar> nn;
     pico_tree::scoped_timer t("pico_tree cv mat", run_count);

examples/pico_toolshed/pico_toolshed/format/format_bin.hpp

@@ -11,23 +11,24 @@ void write_bin(std::string const& filename, std::vector<T_> const& v) {
     return;
   }
 
-  std::size_t const element_size = sizeof(T_);
   std::fstream stream =
       internal::open_stream(filename, std::ios::out | std::ios::binary);
-  stream.write(reinterpret_cast<char const*>(&v[0]), element_size * v.size());
+  internal::stream_wrapper wrapper(stream);
+  wrapper.write(v.data(), v.size());
 }
 
 template <typename T_>
 void read_bin(std::string const& filename, std::vector<T_>& v) {
   std::fstream stream =
       internal::open_stream(filename, std::ios::in | std::ios::binary);
+  internal::stream_wrapper wrapper(stream);
 
   auto bytes = std::filesystem::file_size(filename);
   std::size_t const element_size = sizeof(T_);
   std::size_t const element_count =
       static_cast<std::size_t>(bytes) / element_size;
   v.resize(element_count);
-  stream.read(reinterpret_cast<char*>(&v[0]), element_size * element_count);
+  wrapper.read(element_count, v.data());
 }
 
 }  // namespace pico_tree

examples/pico_toolshed/pico_toolshed/format/format_mnist.hpp

@@ -54,7 +54,7 @@ inline void read_mnist_images(
     throw std::runtime_error("unexpected mnist image dimensions");
   }
 
-  images.resize(header.image_count);
+  images.resize(static_cast<std::size_t>(header.image_count));
   wrapper.read(images.size(), images.data());
 }
 
@@ -75,7 +75,7 @@ inline void read_mnist_labels(
   wrapper.read(header);
   header.label_count = big_endian<std::int32_t>{header.label_count};
 
-  labels.resize(header.label_count);
+  labels.resize(static_cast<std::size_t>(header.label_count));
   wrapper.read(labels.size(), labels.data());
 }
 

examples/pico_understory/pico_understory/internal/cover_tree_builder.hpp

@@ -40,7 +40,7 @@ class build_cover_tree_impl {
     std::shuffle(indices.begin(), indices.end(), g);
 
     node_type* node = insert_first_two(indices);
-    for (index_type i = 2; i < static_cast<index_type>(npts); ++i) {
+    for (size_t i = 2; i < npts; ++i) {
       node = insert(node, create_node(indices[i]));
     }
 

examples/pico_understory/pico_understory/internal/kd_tree_priority_search.hpp

@@ -68,17 +68,16 @@ class priority_search_nearest_euclidean {
   inline void search_nearest(
       node_type const* const node, scalar_type node_box_distance) {
     if (node->is_leaf()) {
-      for (index_type i = node->data.leaf.begin_idx;
-           i < node->data.leaf.end_idx;
-           ++i) {
-        visitor_(
-            indices_[i],
-            metric_(query_.begin(), query_.end(), space_[indices_[i]]));
+      auto begin = indices_.begin() + node->data.leaf.begin_idx;
+      auto const end = indices_.begin() + node->data.leaf.end_idx;
+      for (; begin < end; ++begin) {
+        visitor_(*begin, metric_(query_.begin(), query_.end(), space_[*begin]));
       }
     } else {
       // Go left or right and then check if we should still go down the other
       // side based on the current minimum distance.
-      scalar_type const v = query_[node->data.branch.split_dim];
+      scalar_type const v =
+          query_[static_cast<size_t>(node->data.branch.split_dim)];
       scalar_type old_offset;
       scalar_type new_offset;
       node_type const* node_1st;

examples/python/kd_tree.py

@@ -1,15 +1,16 @@
 #!/usr/bin/env python3
 
-import pico_tree as pt
+import os
 import numpy as np
+import pico_tree as pt
 from pathlib import Path
 from time import perf_counter
 # from scipy.spatial import KDTree as spKDTree
 # from sklearn.neighbors import KDTree as skKDTree
 # from pykdtree.kdtree import KDTree as pyKDTree
 
 
-def tree_creation_and_query_types():
+def kd_tree_creation_and_query():
     print("*** KdTree Creation And Basic Information ***")
     # An input array must have a dimension of two and it must be
     # contiguous. A C contiguous array contains points in its rows and
@@ -71,9 +72,9 @@ def tree_creation_and_query_types():
     print()
 
     print("*** Box Search ***")
-    # A box search returns the same data structure as a radius search.
-    # However, instead of containing neighbors it simply contains
-    # indices.
+    # A box search returns the same data structure as a radius
+    # search. However, instead of containing neighbors it simply
+    # contains indices.
     # An array of input boxes is defined as follows:
     #   [min_0, max_0, min_1, max_1, ...]
     boxes = np.array(
@@ -106,6 +107,31 @@ def tree_creation_and_query_types():
     print()
 
 
+def kd_tree_file_io():
+    # The save_kd_tree and load_kd_tree functions are considered
+    # convenience functions to save and load a kd_tree on a single
+    # machine. Take the following into account when using them:
+    # * Does not take memory endianness into account.
+    # * Does not check if the stored tree structure is valid for the
+    #   given point set.
+
+    a = np.array([[2, 1], [4, 3], [8, 7]], dtype=np.float64, order='C')
+    t1 = pt.KdTree(a, pt.Metric.L2Squared, 10)
+
+    filename = "tree.bin"
+    # The save_kd_tree *only* saves the tree structure, not the point
+    # set that was used to create the KdTree.
+    pt.save_kd_tree(t1, filename)
+    # Loading the KdTree requires the original point set with which the
+    # KdTree was originally created.
+    t2 = pt.load_kd_tree(a, filename)
+    os.remove(filename)
+
+    print(t1)
+    print(t2)
+    print()
+
+
 def array_initialization():
     print("*** Array Initialization ***")
     p = np.array([[2, 1], [4, 3], [8, 7]], dtype=np.float64)
@@ -126,7 +152,7 @@ def array_initialization():
     print()
 
 
-def performance_test_pico_tree():
+def kd_tree_performance_test():
     print("*** Performance against scans.bin ***")
     # The benchmark documentation, docs/benchmark.md section "Running a
     # new benchmark", explains how to generate a scans.bin file from an
@@ -175,9 +201,10 @@ def performance_test_pico_tree():
 
 
 def main():
-    tree_creation_and_query_types()
+    kd_tree_creation_and_query()
+    kd_tree_file_io()
     array_initialization()
-    performance_test_pico_tree()
+    kd_tree_performance_test()
 
 
 if __name__ == "__main__":

setup.py

@@ -13,6 +13,6 @@
       packages=['pico_tree'],
       package_dir={'': 'src/pyco_tree'},
       cmake_install_dir='src/pyco_tree/pico_tree',
-      python_requires='>=3.7.3',
+      python_requires='>=3.10',
       install_requires=['numpy'],
       )

src/pico_tree/pico_tree/internal/stream_wrapper.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <fstream>
+#include <string>
 #include <vector>
 
 namespace pico_tree::internal {
@@ -46,6 +47,15 @@ class stream_wrapper {
     read(size, values.data());
   }
 
+  //! \brief Reads a string from the stream_wrapper.
+  //! \details Reads the size of the string followed by all its elements.
+  inline void read(std::string& values) {
+    typename std::string::size_type size;
+    read(size);
+    values.resize(size);
+    read(size, values.data());
+  }
+
   //! \brief Reads an array of values from the stream_wrapper.
   //! \tparam T_ Type of a value.
   template <typename T_>
@@ -70,9 +80,14 @@ class stream_wrapper {
   template <typename T_>
   inline void write(std::vector<T_> const& values) {
     write(values.size());
-    stream_.write(
-        reinterpret_cast<char const*>(&values[0]),
-        static_cast<std::streamsize>(sizeof(T_) * values.size()));
+    write(values.data(), values.size());
+  }
+
+  //! \brief Writes a string to the stream_wrapper.
+  //! \details Writes the size of the string followed by all its elements.
+  inline void write(std::string const& values) {
+    write(values.size());
+    write(values.data(), values.size());
   }
 
   //! \brief Writes an array of values to the stream_wrapper.

src/pyco_tree/pico_tree/__init__.py

@@ -1,5 +1,5 @@
-__all__ = ['DArray', 'KdTree', 'Metric']
+__all__ = ['DArray', 'Metric', 'KdTree', 'load_kd_tree', 'save_kd_tree']
 
-from .kd_tree import KdTree
-from .metric import Metric
-from ._pyco_tree import __doc__, DArray
+# TODO Generate .pyi stub files for _pyco_tree.
+
+from ._pyco_tree import __doc__, DArray, Metric, KdTree, load_kd_tree, save_kd_tree

src/pyco_tree/pico_tree/_pyco_tree/core.hpp

@@ -55,10 +55,14 @@ struct array_layout {
         index_inner(row_major ? 1 : 0),
         index_outer(row_major ? 0 : 1) {}
 
+  inline py::ssize_t inner_stride() const { return info.shape[index_inner]; }
+
+  inline py::ssize_t outer_stride() const { return info.shape[index_outer]; }
+
   py::buffer_info info;
   bool row_major;
-  py::ssize_t index_inner;
-  py::ssize_t index_outer;
+  std::size_t index_inner;
+  std::size_t index_outer;
 };
 
 template <pico_tree::size_t Dim_>