Python interface search_box input updated.

Author: Jaybro

examples/python/kd_tree.py

@@ -74,10 +74,20 @@ def tree_creation_and_query_types():
     # A box search returns the same data structure as a radius search.
     # However, instead of containing neighbors it simply contains
     # indices.
-    min = np.array([[0, 0], [2, 2], [0, 0], [6, 6]], dtype=np.float32)
-    max = np.array([[3, 3], [3, 3], [9, 9], [9, 9]], dtype=np.float32)
-    bnns = t.search_box(min, max)
-    t.search_box(min, max, bnns)
+    # An array of input boxes is defined as follows:
+    #   [min_0, max_0, min_1, max_1, ...]
+    boxes = np.array(
+        [[0, 0],
+         [3, 3],
+         [2, 2],
+         [3, 3],
+         [0, 0],
+         [9, 9],
+         [6, 6],
+         [9, 9]],
+        dtype=np.float32)
+    bnns = t.search_box(boxes)
+    t.search_box(boxes, bnns)
     print("Results for the orthogonal box search:")
     for bnn in bnns:
         print(f"{bnn}")

src/pyco_tree/pico_tree/_pyco_tree/darray.hpp

@@ -45,6 +45,16 @@ class DArrayImpl : public DArrayImplBase {
     // It is important that at the binding side of things we ensure that the
     // array is kept alive while the view is alive.
     // NOTE: At the time of writing an undocumented feature.
+
+    // In case the size of a vector equals 0, its data pointer can equal
+    // nullptr. When this happens, the library interface of numpy (as wrapped by
+    // pybind11) will allocate some memory and store the address to that memory
+    // instead of storing the nullptr address of the vector. This means that
+    // each time we create a view for the same empty vector, the memory address
+    // it stores may randomly change. This is not an issue, but good to document
+    // here. See:
+    //  * PyArray_NewFromDescr(...)
+    //  * https://numpy.org/doc/1.13/reference/c-api.array.html
     return pybind11::array_t<T, 0>(
         array_[i].size(), array_[i].data(), pybind11::none());
   }

src/pyco_tree/pico_tree/_pyco_tree/def_kd_tree.cpp

@@ -201,23 +201,19 @@ Search for all neighbors within a radius of each of the input points.
       .def(
           "search_box",
           static_cast<void (KdTree::*)(
-              py::array_t<Scalar, 0> const,
-              py::array_t<Scalar, 0> const,
-              Neighborhoods&) const>(&KdTree::SearchBox),
-          py::arg("min").noconvert().none(false),
-          py::arg("max").noconvert().none(false),
-          py::arg("box").noconvert().none(false),
+              py::array_t<Scalar, 0> const, Neighborhoods&) const>(
+              &KdTree::SearchBox),
+          py::arg("boxes").noconvert().none(false),
+          py::arg("indices").noconvert().none(false),
           R"ptdoc(
 Search for all points within each of the axis aligned input boxes and
 store the result in the specified output.
 )ptdoc")
       .def(
           "search_box",
-          static_cast<Neighborhoods (KdTree::*)(
-              py::array_t<Scalar, 0> const, py::array_t<Scalar, 0> const)
+          static_cast<Neighborhoods (KdTree::*)(py::array_t<Scalar, 0> const)
                           const>(&KdTree::SearchBox),
-          py::arg("min").noconvert().none(false),
-          py::arg("max").noconvert().none(false),
+          py::arg("boxes").noconvert().none(false),
           "Search for all points within each of the axis aligned input boxes.");
 }
 

src/pyco_tree/pico_tree/_pyco_tree/kd_tree.hpp

@@ -112,32 +112,29 @@ class KdTree : public pico_tree::KdTree<Space_, Metric_> {
   }
 
   void SearchBox(
-      py::array_t<ScalarType, 0> const min,
-      py::array_t<ScalarType, 0> const max,
-      DArray& box) const {
-    auto query_min = MakeMap<Dim>(min);
-    auto query_max = MakeMap<Dim>(max);
+      py::array_t<ScalarType, 0> const boxes, DArray& indices) const {
+    auto query = MakeMap<Dim>(boxes);
 
-    if (query_min.size() != query_max.size()) {
+    if (query.size() % 2 != 0) {
       throw std::invalid_argument("Query min and max don't have equal size.");
     }
 
-    auto& box_data = box.data<IndexType>();
-    box_data.resize(query_min.size());
+    std::size_t box_count = query.size() / std::size_t(2);
+    auto& indices_data = indices.data<IndexType>();
+    indices_data.resize(box_count);
 
 #pragma omp parallel for schedule(dynamic, kChunkSize)
     // TODO Reduce the vector resize overhead
-    for (SSize i = 0; i < static_cast<SSize>(query_min.size()); ++i) {
-      Base::SearchBox(query_min[i], query_max[i], box_data[i]);
+    for (SSize i = 0; i < static_cast<SSize>(box_count); ++i) {
+      auto index = static_cast<pico_tree::Size>(i * 2);
+      Base::SearchBox(query[index + 0], query[index + 1], indices_data[i]);
     }
   }
 
-  DArray SearchBox(
-      py::array_t<ScalarType, 0> const min,
-      py::array_t<ScalarType, 0> const max) const {
-    DArray box = DArray(std::vector<std::vector<IndexType>>());
-    SearchBox(min, max, box);
-    return box;
+  DArray SearchBox(py::array_t<ScalarType, 0> const boxes) const {
+    DArray indices = DArray(std::vector<std::vector<IndexType>>());
+    SearchBox(boxes, indices);
+    return indices;
   }
 
   inline ScalarType const* data() const { return points().data(); }

test/pyco_tree/kd_tree_test.py

@@ -107,28 +107,45 @@ def test_search_radius(self):
             self.assertEqual(nns[i][0][0], i)
             self.assertAlmostEqual(nns[i][0][1], 0)
 
-        # Test that the memory is re-used
-        datas = [x.ctypes.data for x in nns]
+        # Test that the memory is re-used by comparing memory addresses.
+        # In case the size of an array equals zero, its memory address is
+        # random. See darray.hpp for more details.
+        def addresses(nns):
+            return [x.ctypes.data if len(x) else 0 for x in nns]
+
+        datas = addresses(nns)
         t.search_radius(a, search_radius, nns)
-        self.assertEqual([x.ctypes.data for x in nns], datas)
+        self.assertEqual(addresses(nns), datas)
         t.search_radius(a, search_radius**2, nns)
-        self.assertNotEqual([x.ctypes.data for x in nns], datas)
+        self.assertNotEqual(addresses(nns), datas)
 
     def test_search_box(self):
         a = np.array([[2, 1], [4, 3], [8, 7]], dtype=np.float32)
         t = pt.KdTree(a, pt.Metric.L2Squared, 10)
-
-        min = np.array([[0, 0], [2, 2], [0, 0], [6, 6]], dtype=np.float32)
-        max = np.array([[3, 3], [3, 3], [9, 9], [9, 9]], dtype=np.float32)
-        nns = t.search_box(min, max)
+        boxes = np.array(
+            [[0, 0],
+             [3, 3],
+             [2, 2],
+             [3, 3],
+             [0, 0],
+             [9, 9],
+             [6, 6],
+             [9, 9]],
+            dtype=np.float32)
+        nns = t.search_box(boxes)
         self.assertEqual(len(nns), 4)
         self.assertEqual(nns.dtype, t.dtype_index)
         self.assertTrue(nns)
 
-        # Test that the memory is re-used
-        nns[0][0] = 42
-        t.search_box(min, max, nns)
-        self.assertEqual(nns[0][0], 0)
+        # Test that the memory is re-used by comparing memory addresses.
+        # In case the size of an array equals zero, its memory address is
+        # random. See darray.hpp for more details.
+        def addresses(nns):
+            return [x.ctypes.data if len(x) else 0 for x in nns]
+
+        datas = addresses(nns)
+        t.search_box(boxes, nns)
+        self.assertEqual(addresses(nns), datas)
 
         # Check the number of indices found.
         sizes = [1, 0, 3, 1]