Refactor array-loading methods, add tests

Author: mwtoews

rtree/index.py

@@ -1061,12 +1061,16 @@ def intersection_v(self, mins, maxs):
         """
         import numpy as np
 
-        assert mins.shape == maxs.shape
-        assert mins.strides == maxs.strides
+        # Ensure inputs are 2D float arrays
+        mins = np.atleast_2d(mins).astype(np.float64).copy()
+        maxs = np.atleast_2d(maxs).astype(np.float64).copy()
 
-        # Cast
-        mins = mins.astype(np.float64)
-        maxs = maxs.astype(np.float64)
+        if mins.ndim != 2 or maxs.ndim != 2:
+            raise ValueError("mins/maxs must have 2 dimensions: (n, d)")
+        if mins.shape != maxs.shape:
+            raise ValueError("mins and maxs shapes not equal")
+        if mins.strides != maxs.strides:
+            raise ValueError("mins and maxs strides not equal")
 
         # Extract counts
         n, d = mins.shape
@@ -1109,6 +1113,7 @@ def nearest_v(
         self,
         mins,
         maxs,
+        *,
         num_results=1,
         max_dists=None,
         strict=False,
@@ -1144,12 +1149,16 @@ def nearest_v(
         """
         import numpy as np
 
-        assert mins.shape == maxs.shape
-        assert mins.strides == maxs.strides
+        # Ensure inputs are 2D float arrays
+        mins = np.atleast_2d(mins).astype(np.float64).copy()
+        maxs = np.atleast_2d(maxs).astype(np.float64).copy()
 
-        # Cast
-        mins = mins.astype(np.float64)
-        maxs = maxs.astype(np.float64)
+        if mins.ndim != 2 or maxs.ndim != 2:
+            raise ValueError("mins/maxs must have 2 dimensions: (n, d)")
+        if mins.shape != maxs.shape:
+            raise ValueError("mins and maxs shapes not equal")
+        if mins.strides != maxs.strides:
+            raise ValueError("mins and maxs strides not equal")
 
         # Extract counts
         n, d = mins.shape
@@ -1164,9 +1173,11 @@ def nearest_v(
         offn, offi = 0, 0
 
         if max_dists is not None:
-            assert len(max_dists) == n
-
-            dists = max_dists.astype(np.float64).copy()
+            dists = np.atleast_1d(max_dists).astype(np.float64).copy()
+            if dists.ndim != 1:
+                raise ValueError("max_dists must have 1 dimension")
+            if len(dists) != n:
+                raise ValueError(f"max_dists must have length {n}")
         elif return_max_dists:
             dists = np.zeros(n)
         else:
@@ -1189,7 +1200,7 @@ def nearest_v(
                 ctypes.byref(nr),
             )
 
-            # If we got the expected nuber of results then return
+            # If we got the expected number of results then return
             if nr.value == n - offn:
                 if return_max_dists:
                     return ids[: counts.sum()], counts, dists

tests/common.py

@@ -1,6 +1,10 @@
 """Common test functions."""
 
+import pytest
+
 from rtree.core import rt
 
 sidx_version_string = rt.SIDX_Version().decode()
 sidx_version = tuple(map(int, sidx_version_string.split(".", maxsplit=3)[:3]))
+
+skip_sidx_lt_210 = pytest.mark.skipif(sidx_version < (2, 1, 0), reason="SIDX < 2.1.0")

tests/test_index.py

@@ -14,6 +14,8 @@
 from rtree import core, index
 from rtree.exceptions import RTreeError
 
+from .common import skip_sidx_lt_210
+
 
 class IndexTestCase(unittest.TestCase):
     def setUp(self) -> None:
@@ -268,6 +270,26 @@ def test_double_insertion(self) -> None:
 
         self.assertEqual([1, 1], list(idx.intersection((0, 0, 5, 5))))
 
+    @skip_sidx_lt_210
+    def test_intersection_v(self) -> None:
+        mins = np.array([[0, 1]] * 2).T
+        maxs = np.array([[60, 50]] * 2).T
+        ret = self.idx.intersection_v(mins, maxs)
+        assert type(ret) is tuple
+        ids, counts = ret
+        assert ids.dtype == np.int64
+        ids0 = [0, 4, 16, 27, 35, 40, 47, 50, 76, 80]
+        ids1 = [0, 16, 27, 35, 47, 76]
+        assert ids.tolist() == ids0 + ids1
+        assert counts.dtype == np.uint64
+        assert counts.tolist() == [len(ids0), len(ids1)]
+
+        # errors
+        with pytest.raises(ValueError, match="must have 2 dimensions"):
+            self.idx.intersection_v(np.ones((2, 3, 4)), 4)
+        with pytest.raises(ValueError, match="shapes not equal"):
+            self.idx.intersection_v([0], [10, 12])
+
 
 class TestIndexIntersectionUnion:
     @pytest.fixture(scope="class")
@@ -314,6 +336,17 @@ def test_intersection_interleaved(
             else:
                 assert False
 
+    @skip_sidx_lt_210
+    def test_intersection_v_interleaved(
+        self, index_a_interleaved: index.Index, index_b_interleaved: index.Index
+    ) -> None:
+        index_c_interleaved = index_a_interleaved & index_b_interleaved
+        mins = index_c_interleaved.bounds[0:2]
+        maxs = index_c_interleaved.bounds[2:4]
+        idxs, counts = index_c_interleaved.intersection_v(mins, maxs)
+        assert idxs.tolist() == [0, 1]
+        assert counts.tolist() == [2]
+
     def test_intersection_uninterleaved(
         self, index_a_uninterleaved: index.Index, index_b_uninterleaved: index.Index
     ) -> None:
@@ -330,6 +363,17 @@ def test_intersection_uninterleaved(
             else:
                 assert False
 
+    @skip_sidx_lt_210
+    def test_intersection_v_uninterleaved(
+        self, index_a_uninterleaved: index.Index, index_b_uninterleaved: index.Index
+    ) -> None:
+        index_c_uninterleaved = index_a_uninterleaved & index_b_uninterleaved
+        mins = index_c_uninterleaved.bounds[0::2]
+        maxs = index_c_uninterleaved.bounds[1::2]
+        idxs, counts = index_c_uninterleaved.intersection_v(mins, maxs)
+        assert idxs.tolist() == [0, 1]
+        assert counts.tolist() == [2]
+
     def test_intersection_mismatch(
         self, index_a_interleaved: index.Index, index_b_uninterleaved: index.Index
     ) -> None:
@@ -617,6 +661,46 @@ def test_nearest_basic(self) -> None:
         hits = sorted(idx.nearest((13, 0, 20, 2), 3))
         self.assertEqual(hits, [3, 4, 5])
 
+    @skip_sidx_lt_210
+    def test_nearest_v_basic(self) -> None:
+        mins = np.array([[0, 5]] * 2).T
+        maxs = np.array([[10, 15]] * 2).T
+        ret = self.idx.nearest_v(mins, maxs, num_results=3)
+        assert type(ret) is tuple
+        ids, counts = ret
+        assert ids.dtype == np.int64
+        ids0 = [76, 48, 19]
+        ids1 = [76, 47, 48]
+        assert ids.tolist() == ids0 + ids1
+        assert counts.dtype == np.uint64
+        assert counts.tolist() == [3, 3]
+
+        ret = self.idx.nearest_v(mins, maxs, num_results=3, return_max_dists=True)
+        assert type(ret) is tuple
+        ids, counts, max_dists = ret
+        assert ids.tolist() == ids0 + ids1
+        assert counts.tolist() == [3, 3]
+        assert max_dists.dtype == np.float64
+        np.testing.assert_allclose(max_dists, [7.54938045, 11.05686397])
+
+        ret = self.idx.nearest_v(
+            mins, maxs, num_results=3, max_dists=[10, 10], return_max_dists=True
+        )
+        ids, counts, max_dists = ret
+        assert ids.tolist() == ids0 + ids1[:2]
+        assert counts.tolist() == [3, 2]
+        np.testing.assert_allclose(max_dists, [7.54938045, 3.92672575])
+
+        # errors
+        with pytest.raises(ValueError, match="must have 2 dimensions"):
+            self.idx.nearest_v(np.ones((2, 3, 4)), 4)
+        with pytest.raises(ValueError, match="shapes not equal"):
+            self.idx.nearest_v([0], [10, 12])
+        with pytest.raises(ValueError, match="max_dists must have 1 dimension"):
+            self.idx.nearest_v(maxs, mins, max_dists=[[10]])
+        with pytest.raises(ValueError, match="max_dists must have length 2"):
+            self.idx.nearest_v(maxs, mins, max_dists=[10])
+
     def test_nearest_equidistant(self) -> None:
         """Test that if records are equidistant, both are returned."""
         point = (0, 0)
@@ -677,25 +761,47 @@ def test_deletion(self) -> None:
         self.assertEqual(hits, [])
 
 
-class IndexMoreDimensions(IndexTestCase):
-    def test_3d(self) -> None:
-        """Test we make and query a 3D index"""
+class Index3d(IndexTestCase):
+    """Test we make and query a 3D index"""
+
+    def setUp(self) -> None:
         p = index.Property()
         p.dimension = 3
-        idx = index.Index(properties=p, interleaved=False)
-        idx.insert(1, (0, 0, 60, 60, 22, 22.0))
-        hits = idx.intersection((-1, 1, 58, 62, 22, 24))
+        self.idx = index.Index(properties=p, interleaved=False)
+        self.idx.insert(1, (0, 0, 60, 60, 22, 22.0))
+        self.coords = (-1, 1, 58, 62, 22, 24)
+
+    def test_intersection(self) -> None:
+        hits = self.idx.intersection(self.coords)
         self.assertEqual(list(hits), [1])
 
-    def test_4d(self) -> None:
-        """Test we make and query a 4D index"""
+    @skip_sidx_lt_210
+    def test_intersection_v(self) -> None:
+        idxs, counts = self.idx.intersection_v(self.coords[0::2], self.coords[1::2])
+        assert idxs.tolist() == [1]
+        assert counts.tolist() == [1]
+
+
+class Index4d(IndexTestCase):
+    """Test we make and query a 4D index"""
+
+    def setUp(self) -> None:
         p = index.Property()
         p.dimension = 4
-        idx = index.Index(properties=p, interleaved=False)
-        idx.insert(1, (0, 0, 60, 60, 22, 22.0, 128, 142))
-        hits = idx.intersection((-1, 1, 58, 62, 22, 24, 120, 150))
+        self.idx = index.Index(properties=p, interleaved=False)
+        self.idx.insert(1, (0, 0, 60, 60, 22, 22.0, 128, 142))
+        self.coords = (-1, 1, 58, 62, 22, 24, 120, 150)
+
+    def test_intersection(self) -> None:
+        hits = self.idx.intersection(self.coords)
         self.assertEqual(list(hits), [1])
 
+    @skip_sidx_lt_210
+    def test_intersection_v(self) -> None:
+        idxs, counts = self.idx.intersection_v(self.coords[0::2], self.coords[1::2])
+        assert idxs.tolist() == [1]
+        assert counts.tolist() == [1]
+
 
 class IndexStream(IndexTestCase):
     def test_stream_input(self) -> None: