rstar-python 0.1.0__tar.gz → 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {rstar_python-0.1.0 → rstar_python-0.2.0}/CHANGELOG.md +12 -0
- {rstar_python-0.1.0 → rstar_python-0.2.0}/Cargo.lock +1 -32
- {rstar_python-0.1.0 → rstar_python-0.2.0}/Cargo.toml +2 -2
- {rstar_python-0.1.0 → rstar_python-0.2.0}/PKG-INFO +42 -12
- {rstar_python-0.1.0 → rstar_python-0.2.0}/README.md +41 -11
- {rstar_python-0.1.0 → rstar_python-0.2.0}/pyproject.toml +1 -1
- rstar_python-0.2.0/python/benchmarks/benchmark_bbox_rtree.py +102 -0
- rstar_python-0.2.0/python/rstar_python/__init__.py +3 -0
- {rstar_python-0.1.0 → rstar_python-0.2.0}/python/rstar_python/rstar_python.pyi +68 -0
- rstar_python-0.2.0/python/tests/test_bbox_rtree.py +162 -0
- {rstar_python-0.1.0 → rstar_python-0.2.0}/python/tests/test_rtree.py +36 -0
- rstar_python-0.2.0/src/lib.rs +898 -0
- rstar_python-0.1.0/python/rstar_python/__init__.py +0 -3
- rstar_python-0.1.0/src/lib.rs +0 -373
- {rstar_python-0.1.0 → rstar_python-0.2.0}/.github/workflows/ci.yml +0 -0
- {rstar_python-0.1.0 → rstar_python-0.2.0}/.github/workflows/release.yml +0 -0
- {rstar_python-0.1.0 → rstar_python-0.2.0}/.gitignore +0 -0
- {rstar_python-0.1.0 → rstar_python-0.2.0}/LICENSE +0 -0
- {rstar_python-0.1.0 → rstar_python-0.2.0}/python/rstar_python/py.typed +0 -0
|
@@ -7,6 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|
|
7
7
|
|
|
8
8
|
## [Unreleased]
|
|
9
9
|
|
|
10
|
+
## [0.2.0] - 2026-08-21
|
|
11
|
+
|
|
12
|
+
### Added
|
|
13
|
+
- `PyRTree.locate_in_envelope_ids()` for stored ids within a closed bounding box.
|
|
14
|
+
- `PyRTree.remove_item()` for exact point-and-id removal when points coincide.
|
|
15
|
+
- `PyBBoxRTree`, a 2–8 dimensional index for stored axis-aligned bounding boxes,
|
|
16
|
+
with insertion, optimized NumPy bulk loading, closed-boundary intersection,
|
|
17
|
+
and exact removal.
|
|
18
|
+
- `PyBBoxRTree.intersection_batch()` for vectorised queries with CSR-style output.
|
|
19
|
+
- Bounding-box bulk loading and batched intersections release the GIL while
|
|
20
|
+
running native tree operations.
|
|
21
|
+
|
|
10
22
|
## [0.1.0] - 2026-06-07
|
|
11
23
|
|
|
12
24
|
### Added
|
|
@@ -231,13 +231,12 @@ checksum = "5912b862fa5ffb462607bfd1e35036c458c537921f508c8235a83d5f3987edfe"
|
|
|
231
231
|
dependencies = [
|
|
232
232
|
"heapless",
|
|
233
233
|
"num-traits",
|
|
234
|
-
"serde",
|
|
235
234
|
"smallvec",
|
|
236
235
|
]
|
|
237
236
|
|
|
238
237
|
[[package]]
|
|
239
238
|
name = "rstar-python"
|
|
240
|
-
version = "0.
|
|
239
|
+
version = "0.2.0"
|
|
241
240
|
dependencies = [
|
|
242
241
|
"numpy",
|
|
243
242
|
"pyo3",
|
|
@@ -250,36 +249,6 @@ version = "2.1.2"
|
|
|
250
249
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
251
250
|
checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe"
|
|
252
251
|
|
|
253
|
-
[[package]]
|
|
254
|
-
name = "serde"
|
|
255
|
-
version = "1.0.228"
|
|
256
|
-
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
257
|
-
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
|
|
258
|
-
dependencies = [
|
|
259
|
-
"serde_core",
|
|
260
|
-
"serde_derive",
|
|
261
|
-
]
|
|
262
|
-
|
|
263
|
-
[[package]]
|
|
264
|
-
name = "serde_core"
|
|
265
|
-
version = "1.0.228"
|
|
266
|
-
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
267
|
-
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
|
|
268
|
-
dependencies = [
|
|
269
|
-
"serde_derive",
|
|
270
|
-
]
|
|
271
|
-
|
|
272
|
-
[[package]]
|
|
273
|
-
name = "serde_derive"
|
|
274
|
-
version = "1.0.228"
|
|
275
|
-
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
276
|
-
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
|
|
277
|
-
dependencies = [
|
|
278
|
-
"proc-macro2",
|
|
279
|
-
"quote",
|
|
280
|
-
"syn",
|
|
281
|
-
]
|
|
282
|
-
|
|
283
252
|
[[package]]
|
|
284
253
|
name = "smallvec"
|
|
285
254
|
version = "1.15.1"
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[package]
|
|
2
2
|
name = "rstar-python"
|
|
3
|
-
version = "0.
|
|
3
|
+
version = "0.2.0"
|
|
4
4
|
edition = "2021"
|
|
5
5
|
license = "MIT"
|
|
6
6
|
description = "Python bindings for the rstar R*-tree spatial index"
|
|
@@ -15,4 +15,4 @@ crate-type = ["cdylib"]
|
|
|
15
15
|
[dependencies]
|
|
16
16
|
pyo3 = { version = "0.28", features = ["extension-module", "abi3-py310"] }
|
|
17
17
|
numpy = "0.28"
|
|
18
|
-
rstar =
|
|
18
|
+
rstar = "0.13"
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: rstar-python
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
Classifier: Development Status :: 3 - Alpha
|
|
5
5
|
Classifier: Intended Audience :: Developers
|
|
6
6
|
Classifier: License :: OSI Approved :: MIT License
|
|
@@ -31,10 +31,10 @@ Project-URL: Repository, https://github.com/kephale/rstar-python
|
|
|
31
31
|
|
|
32
32
|
Python bindings for the [rstar](https://github.com/georust/rstar) R*-tree spatial index library.
|
|
33
33
|
|
|
34
|
-
A fast, **dynamic** spatial index for points
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
return references to *your* data
|
|
34
|
+
A fast, **dynamic** spatial index for points and stored axis-aligned bounding
|
|
35
|
+
boxes in 2–8 dimensions. Insert and remove items after construction, run
|
|
36
|
+
nearest-neighbour and radius searches, and query spatial regions. Each item can
|
|
37
|
+
carry an integer id so queries can return references to *your* data.
|
|
38
38
|
|
|
39
39
|
## Installation
|
|
40
40
|
|
|
@@ -64,7 +64,7 @@ easy prebuilt wheels and no C/C++ system dependency.
|
|
|
64
64
|
|
|
65
65
|
```python
|
|
66
66
|
import numpy as np
|
|
67
|
-
from rstar_python import PyRTree
|
|
67
|
+
from rstar_python import PyBBoxRTree, PyRTree
|
|
68
68
|
|
|
69
69
|
# Create a 3D R-tree
|
|
70
70
|
tree = PyRTree(dims=3)
|
|
@@ -85,6 +85,10 @@ tree.nearest_neighbor([1.1, 2.1, 3.1]) # -> [1.0, 2.0, 3.0]
|
|
|
85
85
|
tree.k_nearest_neighbors([1.1, 2.1, 3.1], k=2) # -> [[...], [...]]
|
|
86
86
|
tree.neighbors_within_radius([1.0, 2.0, 3.0], radius=1.0)
|
|
87
87
|
tree.locate_in_envelope(min_corner=[0, 0, 0], max_corner=[2, 2, 2])
|
|
88
|
+
tree.locate_in_envelope_ids([0, 0, 0], [2, 2, 2]) # stored ids, closed AABB
|
|
89
|
+
|
|
90
|
+
# Remove an exact point+id when coordinates are shared by multiple items.
|
|
91
|
+
tree.remove_item([1.0, 2.0, 3.0], data=10)
|
|
88
92
|
|
|
89
93
|
# --- Vectorised, id-returning queries (scipy/sklearn style) ---
|
|
90
94
|
query_pts = np.array([[1.1, 2.1, 3.1], [7.0, 8.0, 9.0]], dtype=np.float64)
|
|
@@ -99,18 +103,41 @@ len(tree) # number of points
|
|
|
99
103
|
tree.dims # 3
|
|
100
104
|
[1.0, 2.0, 3.0] in tree # membership test
|
|
101
105
|
tree.remove([1.0, 2.0, 3.0])
|
|
106
|
+
|
|
107
|
+
# --- Stored bounding boxes ---
|
|
108
|
+
boxes = PyBBoxRTree(dims=2)
|
|
109
|
+
boxes.insert([0.0, 0.0], [2.0, 2.0], data=100)
|
|
110
|
+
boxes.insert([2.0, 1.0], [3.0, 3.0])
|
|
111
|
+
# Returns both ids: intersection uses closed bounds, so touching counts.
|
|
112
|
+
boxes.intersection([1.0, 1.0], [2.0, 2.0])
|
|
113
|
+
|
|
114
|
+
# Query many boxes in one native call. Results use CSR-style row offsets:
|
|
115
|
+
# the matches for row i are ids[offsets[i]:offsets[i + 1]].
|
|
116
|
+
query_mins = np.array([[0.0, 0.0], [5.0, 5.0]], dtype=np.float64)
|
|
117
|
+
query_maxs = np.array([[2.0, 2.0], [6.0, 6.0]], dtype=np.float64)
|
|
118
|
+
ids, offsets = boxes.intersection_batch(query_mins, query_maxs)
|
|
119
|
+
|
|
120
|
+
boxes.remove_item([0.0, 0.0], [2.0, 2.0], data=100)
|
|
121
|
+
|
|
122
|
+
# Bulk loading replaces the stored boxes.
|
|
123
|
+
boxes.bulk_load(
|
|
124
|
+
min_corners=[[0.0, 0.0], [5.0, 5.0]],
|
|
125
|
+
max_corners=[[1.0, 1.0], [6.0, 6.0]],
|
|
126
|
+
data=[10, 20],
|
|
127
|
+
)
|
|
102
128
|
```
|
|
103
129
|
|
|
104
130
|
## Features
|
|
105
131
|
|
|
106
|
-
- Points in 2–8 dimensions
|
|
107
|
-
- Dynamic
|
|
108
|
-
- Per-
|
|
132
|
+
- Points and stored axis-aligned bounding boxes in 2–8 dimensions
|
|
133
|
+
- Dynamic insertion and exact item removal
|
|
134
|
+
- Per-item integer ids (auto-assigned or supplied)
|
|
109
135
|
- Nearest-neighbour and k-nearest-neighbour queries
|
|
110
136
|
- Radius search and axis-aligned bounding-box (envelope) queries
|
|
111
|
-
- Vectorised `query` / `query_radius`
|
|
112
|
-
|
|
113
|
-
-
|
|
137
|
+
- Vectorised point `query` / `query_radius` and bounding-box
|
|
138
|
+
`intersection_batch` queries over NumPy arrays
|
|
139
|
+
- NumPy-native batched intersection results without Python integer boxing
|
|
140
|
+
- Bulk loading (lists or NumPy arrays) for fast construction
|
|
114
141
|
- Type stubs (`py.typed`) for IDE and `mypy` support
|
|
115
142
|
- Built on the fast Rust [rstar](https://github.com/georust/rstar) library
|
|
116
143
|
|
|
@@ -135,6 +162,9 @@ maturin develop --release
|
|
|
135
162
|
|
|
136
163
|
# Run tests
|
|
137
164
|
pytest python/tests -v
|
|
165
|
+
|
|
166
|
+
# Run the bounding-box performance benchmark
|
|
167
|
+
python python/benchmarks/benchmark_bbox_rtree.py
|
|
138
168
|
```
|
|
139
169
|
|
|
140
170
|
## License
|
|
@@ -2,10 +2,10 @@
|
|
|
2
2
|
|
|
3
3
|
Python bindings for the [rstar](https://github.com/georust/rstar) R*-tree spatial index library.
|
|
4
4
|
|
|
5
|
-
A fast, **dynamic** spatial index for points
|
|
6
|
-
|
|
7
|
-
|
|
8
|
-
return references to *your* data
|
|
5
|
+
A fast, **dynamic** spatial index for points and stored axis-aligned bounding
|
|
6
|
+
boxes in 2–8 dimensions. Insert and remove items after construction, run
|
|
7
|
+
nearest-neighbour and radius searches, and query spatial regions. Each item can
|
|
8
|
+
carry an integer id so queries can return references to *your* data.
|
|
9
9
|
|
|
10
10
|
## Installation
|
|
11
11
|
|
|
@@ -35,7 +35,7 @@ easy prebuilt wheels and no C/C++ system dependency.
|
|
|
35
35
|
|
|
36
36
|
```python
|
|
37
37
|
import numpy as np
|
|
38
|
-
from rstar_python import PyRTree
|
|
38
|
+
from rstar_python import PyBBoxRTree, PyRTree
|
|
39
39
|
|
|
40
40
|
# Create a 3D R-tree
|
|
41
41
|
tree = PyRTree(dims=3)
|
|
@@ -56,6 +56,10 @@ tree.nearest_neighbor([1.1, 2.1, 3.1]) # -> [1.0, 2.0, 3.0]
|
|
|
56
56
|
tree.k_nearest_neighbors([1.1, 2.1, 3.1], k=2) # -> [[...], [...]]
|
|
57
57
|
tree.neighbors_within_radius([1.0, 2.0, 3.0], radius=1.0)
|
|
58
58
|
tree.locate_in_envelope(min_corner=[0, 0, 0], max_corner=[2, 2, 2])
|
|
59
|
+
tree.locate_in_envelope_ids([0, 0, 0], [2, 2, 2]) # stored ids, closed AABB
|
|
60
|
+
|
|
61
|
+
# Remove an exact point+id when coordinates are shared by multiple items.
|
|
62
|
+
tree.remove_item([1.0, 2.0, 3.0], data=10)
|
|
59
63
|
|
|
60
64
|
# --- Vectorised, id-returning queries (scipy/sklearn style) ---
|
|
61
65
|
query_pts = np.array([[1.1, 2.1, 3.1], [7.0, 8.0, 9.0]], dtype=np.float64)
|
|
@@ -70,18 +74,41 @@ len(tree) # number of points
|
|
|
70
74
|
tree.dims # 3
|
|
71
75
|
[1.0, 2.0, 3.0] in tree # membership test
|
|
72
76
|
tree.remove([1.0, 2.0, 3.0])
|
|
77
|
+
|
|
78
|
+
# --- Stored bounding boxes ---
|
|
79
|
+
boxes = PyBBoxRTree(dims=2)
|
|
80
|
+
boxes.insert([0.0, 0.0], [2.0, 2.0], data=100)
|
|
81
|
+
boxes.insert([2.0, 1.0], [3.0, 3.0])
|
|
82
|
+
# Returns both ids: intersection uses closed bounds, so touching counts.
|
|
83
|
+
boxes.intersection([1.0, 1.0], [2.0, 2.0])
|
|
84
|
+
|
|
85
|
+
# Query many boxes in one native call. Results use CSR-style row offsets:
|
|
86
|
+
# the matches for row i are ids[offsets[i]:offsets[i + 1]].
|
|
87
|
+
query_mins = np.array([[0.0, 0.0], [5.0, 5.0]], dtype=np.float64)
|
|
88
|
+
query_maxs = np.array([[2.0, 2.0], [6.0, 6.0]], dtype=np.float64)
|
|
89
|
+
ids, offsets = boxes.intersection_batch(query_mins, query_maxs)
|
|
90
|
+
|
|
91
|
+
boxes.remove_item([0.0, 0.0], [2.0, 2.0], data=100)
|
|
92
|
+
|
|
93
|
+
# Bulk loading replaces the stored boxes.
|
|
94
|
+
boxes.bulk_load(
|
|
95
|
+
min_corners=[[0.0, 0.0], [5.0, 5.0]],
|
|
96
|
+
max_corners=[[1.0, 1.0], [6.0, 6.0]],
|
|
97
|
+
data=[10, 20],
|
|
98
|
+
)
|
|
73
99
|
```
|
|
74
100
|
|
|
75
101
|
## Features
|
|
76
102
|
|
|
77
|
-
- Points in 2–8 dimensions
|
|
78
|
-
- Dynamic
|
|
79
|
-
- Per-
|
|
103
|
+
- Points and stored axis-aligned bounding boxes in 2–8 dimensions
|
|
104
|
+
- Dynamic insertion and exact item removal
|
|
105
|
+
- Per-item integer ids (auto-assigned or supplied)
|
|
80
106
|
- Nearest-neighbour and k-nearest-neighbour queries
|
|
81
107
|
- Radius search and axis-aligned bounding-box (envelope) queries
|
|
82
|
-
- Vectorised `query` / `query_radius`
|
|
83
|
-
|
|
84
|
-
-
|
|
108
|
+
- Vectorised point `query` / `query_radius` and bounding-box
|
|
109
|
+
`intersection_batch` queries over NumPy arrays
|
|
110
|
+
- NumPy-native batched intersection results without Python integer boxing
|
|
111
|
+
- Bulk loading (lists or NumPy arrays) for fast construction
|
|
85
112
|
- Type stubs (`py.typed`) for IDE and `mypy` support
|
|
86
113
|
- Built on the fast Rust [rstar](https://github.com/georust/rstar) library
|
|
87
114
|
|
|
@@ -106,6 +133,9 @@ maturin develop --release
|
|
|
106
133
|
|
|
107
134
|
# Run tests
|
|
108
135
|
pytest python/tests -v
|
|
136
|
+
|
|
137
|
+
# Run the bounding-box performance benchmark
|
|
138
|
+
python python/benchmarks/benchmark_bbox_rtree.py
|
|
109
139
|
```
|
|
110
140
|
|
|
111
141
|
## License
|
|
@@ -0,0 +1,102 @@
|
|
|
1
|
+
"""Reproducible microbenchmark for bounding-box bulk loads and intersections."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import gc
|
|
7
|
+
import json
|
|
8
|
+
import statistics
|
|
9
|
+
import time
|
|
10
|
+
from collections.abc import Callable
|
|
11
|
+
from typing import Any
|
|
12
|
+
|
|
13
|
+
import numpy as np
|
|
14
|
+
import numpy.typing as npt
|
|
15
|
+
|
|
16
|
+
from rstar_python import PyBBoxRTree
|
|
17
|
+
|
|
18
|
+
FloatArray = npt.NDArray[np.float64]
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def make_boxes(
|
|
22
|
+
rng: np.random.Generator, rows: int, dims: int, size: float
|
|
23
|
+
) -> tuple[FloatArray, FloatArray]:
|
|
24
|
+
mins = rng.random((rows, dims), dtype=np.float64)
|
|
25
|
+
maxs = mins + rng.random((rows, dims), dtype=np.float64) * size
|
|
26
|
+
return mins, maxs
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def median_seconds(function: Callable[[], Any], repeats: int) -> float:
|
|
30
|
+
function() # Warm native code and allocations before measuring.
|
|
31
|
+
samples: list[float] = []
|
|
32
|
+
for _ in range(repeats):
|
|
33
|
+
gc.disable()
|
|
34
|
+
start = time.perf_counter()
|
|
35
|
+
result = function()
|
|
36
|
+
samples.append(time.perf_counter() - start)
|
|
37
|
+
gc.enable()
|
|
38
|
+
del result
|
|
39
|
+
return statistics.median(samples)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def run_benchmark(
|
|
43
|
+
num_boxes: int, num_queries: int, dims: int, repeats: int
|
|
44
|
+
) -> dict[str, float | int]:
|
|
45
|
+
rng = np.random.default_rng(42)
|
|
46
|
+
mins, maxs = make_boxes(rng, num_boxes, dims, 0.005)
|
|
47
|
+
query_mins, query_maxs = make_boxes(rng, num_queries, dims, 0.05)
|
|
48
|
+
ids = np.arange(num_boxes, dtype=np.int64)
|
|
49
|
+
|
|
50
|
+
def build() -> PyBBoxRTree:
|
|
51
|
+
tree = PyBBoxRTree(dims)
|
|
52
|
+
tree.bulk_load(mins, maxs, ids)
|
|
53
|
+
return tree
|
|
54
|
+
|
|
55
|
+
bulk_seconds = median_seconds(build, repeats)
|
|
56
|
+
tree = build()
|
|
57
|
+
|
|
58
|
+
def scalar_list() -> int:
|
|
59
|
+
return sum(
|
|
60
|
+
len(tree.intersection(lower, upper))
|
|
61
|
+
for lower, upper in zip(query_mins, query_maxs, strict=True)
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
def batch() -> int:
|
|
65
|
+
result_ids, offsets = tree.intersection_batch(query_mins, query_maxs)
|
|
66
|
+
assert offsets[-1] == len(result_ids)
|
|
67
|
+
return len(result_ids)
|
|
68
|
+
|
|
69
|
+
hits = scalar_list()
|
|
70
|
+
assert batch() == hits
|
|
71
|
+
scalar_list_seconds = median_seconds(scalar_list, repeats)
|
|
72
|
+
batch_seconds = median_seconds(batch, repeats)
|
|
73
|
+
|
|
74
|
+
return {
|
|
75
|
+
"boxes": num_boxes,
|
|
76
|
+
"queries": num_queries,
|
|
77
|
+
"dims": dims,
|
|
78
|
+
"hits": hits,
|
|
79
|
+
"bulk_seconds": bulk_seconds,
|
|
80
|
+
"scalar_list_seconds": scalar_list_seconds,
|
|
81
|
+
"batch_seconds": batch_seconds,
|
|
82
|
+
"batch_speedup_over_list": scalar_list_seconds / batch_seconds,
|
|
83
|
+
}
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def main() -> None:
|
|
87
|
+
parser = argparse.ArgumentParser(description=__doc__)
|
|
88
|
+
parser.add_argument("--boxes", type=int, default=200_000)
|
|
89
|
+
parser.add_argument("--queries", type=int, default=20_000)
|
|
90
|
+
parser.add_argument("--dims", type=int, default=3)
|
|
91
|
+
parser.add_argument("--repeats", type=int, default=7)
|
|
92
|
+
args = parser.parse_args()
|
|
93
|
+
print(
|
|
94
|
+
json.dumps(
|
|
95
|
+
run_benchmark(args.boxes, args.queries, args.dims, args.repeats),
|
|
96
|
+
indent=2,
|
|
97
|
+
)
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
if __name__ == "__main__":
|
|
102
|
+
main()
|
|
@@ -52,6 +52,11 @@ class PyRTree:
|
|
|
52
52
|
) -> list[list[float]]:
|
|
53
53
|
"""Coordinates of all points inside the axis-aligned box."""
|
|
54
54
|
|
|
55
|
+
def locate_in_envelope_ids(
|
|
56
|
+
self, min_corner: Sequence[float], max_corner: Sequence[float]
|
|
57
|
+
) -> list[int]:
|
|
58
|
+
"""Ids of all points inside the closed axis-aligned box."""
|
|
59
|
+
|
|
55
60
|
def query(
|
|
56
61
|
self, points: npt.NDArray[np.float64], k: int = ...
|
|
57
62
|
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.int64]]:
|
|
@@ -69,9 +74,72 @@ class PyRTree:
|
|
|
69
74
|
def remove(self, point: Sequence[float]) -> bool:
|
|
70
75
|
"""Remove a point by coordinates. Returns ``True`` if one was removed."""
|
|
71
76
|
|
|
77
|
+
def remove_item(self, point: Sequence[float], data: int) -> bool:
|
|
78
|
+
"""Remove the point with exactly matching coordinates and id."""
|
|
79
|
+
|
|
72
80
|
def size(self) -> int:
|
|
73
81
|
"""Number of points in the tree."""
|
|
74
82
|
|
|
75
83
|
def __len__(self) -> int: ...
|
|
76
84
|
def __contains__(self, point: Sequence[float]) -> bool: ...
|
|
77
85
|
def __repr__(self) -> str: ...
|
|
86
|
+
|
|
87
|
+
class PyBBoxRTree:
|
|
88
|
+
"""An R*-tree spatial index over axis-aligned bounding boxes.
|
|
89
|
+
|
|
90
|
+
Corner arguments are normalized per axis, so either corner order is
|
|
91
|
+
accepted for inserts and queries. Bulk loading and batched intersections
|
|
92
|
+
release the GIL during native tree operations. Synchronize mutation of the
|
|
93
|
+
same instance while a batch is in flight to avoid an ``Already borrowed``
|
|
94
|
+
runtime error.
|
|
95
|
+
"""
|
|
96
|
+
|
|
97
|
+
def __init__(self, dims: int) -> None: ...
|
|
98
|
+
@property
|
|
99
|
+
def dims(self) -> int:
|
|
100
|
+
"""Number of dimensions this tree indexes."""
|
|
101
|
+
|
|
102
|
+
def insert(
|
|
103
|
+
self,
|
|
104
|
+
min_corner: Sequence[float],
|
|
105
|
+
max_corner: Sequence[float],
|
|
106
|
+
data: int | None = ...,
|
|
107
|
+
) -> int:
|
|
108
|
+
"""Insert a bounding box, returning its id."""
|
|
109
|
+
|
|
110
|
+
def bulk_load(
|
|
111
|
+
self,
|
|
112
|
+
min_corners: Sequence[Sequence[float]] | npt.NDArray[np.float64],
|
|
113
|
+
max_corners: Sequence[Sequence[float]] | npt.NDArray[np.float64],
|
|
114
|
+
data: Sequence[int] | npt.NDArray[np.int64] | None = ...,
|
|
115
|
+
) -> None:
|
|
116
|
+
"""Build from bounding boxes at once, replacing existing contents."""
|
|
117
|
+
|
|
118
|
+
def intersection(
|
|
119
|
+
self, min_corner: Sequence[float], max_corner: Sequence[float]
|
|
120
|
+
) -> list[int]:
|
|
121
|
+
"""Ids of boxes whose closed bounds overlap or touch the query."""
|
|
122
|
+
|
|
123
|
+
def intersection_batch(
|
|
124
|
+
self,
|
|
125
|
+
min_corners: npt.NDArray[np.float64],
|
|
126
|
+
max_corners: npt.NDArray[np.float64],
|
|
127
|
+
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
|
|
128
|
+
"""Query many boxes, returning ids and CSR row offsets.
|
|
129
|
+
|
|
130
|
+
Results for row ``i`` are ``ids[offsets[i]:offsets[i + 1]]``.
|
|
131
|
+
"""
|
|
132
|
+
|
|
133
|
+
def remove_item(
|
|
134
|
+
self,
|
|
135
|
+
min_corner: Sequence[float],
|
|
136
|
+
max_corner: Sequence[float],
|
|
137
|
+
data: int,
|
|
138
|
+
) -> bool:
|
|
139
|
+
"""Remove the box with exactly matching normalized geometry and id."""
|
|
140
|
+
|
|
141
|
+
def size(self) -> int:
|
|
142
|
+
"""Number of bounding boxes in the tree."""
|
|
143
|
+
|
|
144
|
+
def __len__(self) -> int: ...
|
|
145
|
+
def __repr__(self) -> str: ...
|
|
@@ -0,0 +1,162 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import pytest
|
|
3
|
+
|
|
4
|
+
import rstar_python
|
|
5
|
+
from rstar_python import PyBBoxRTree
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_public_export_and_protocols():
|
|
9
|
+
assert rstar_python.PyBBoxRTree is PyBBoxRTree
|
|
10
|
+
assert "PyBBoxRTree" in rstar_python.__all__
|
|
11
|
+
|
|
12
|
+
tree = PyBBoxRTree(dims=2)
|
|
13
|
+
assert tree.dims == 2
|
|
14
|
+
assert tree.size() == len(tree) == 0
|
|
15
|
+
assert repr(tree) == "PyBBoxRTree(dims=2, size=0)"
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def test_insert_ids_and_intersection_semantics():
|
|
19
|
+
tree = PyBBoxRTree(dims=2)
|
|
20
|
+
assert tree.insert([0.0, 0.0], [4.0, 4.0]) == 0
|
|
21
|
+
assert tree.insert([3.0, 3.0], [6.0, 6.0], data=10) == 10
|
|
22
|
+
assert tree.insert([8.0, 8.0], [9.0, 9.0]) == 11
|
|
23
|
+
|
|
24
|
+
# A query contained by the first box still intersects it; the second only touches.
|
|
25
|
+
assert sorted(tree.intersection([1.0, 1.0], [3.0, 3.0])) == [0, 10]
|
|
26
|
+
assert sorted(tree.intersection([6.0, 6.0], [8.0, 8.0])) == [10, 11]
|
|
27
|
+
assert tree.intersection([6.1, 6.1], [7.9, 7.9]) == []
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def test_bulk_load_without_data_assigns_ids_and_advances_counter():
|
|
31
|
+
tree = PyBBoxRTree(dims=2)
|
|
32
|
+
mins = np.array([[0.0, 0.0], [5.0, 5.0]], dtype=np.float64)
|
|
33
|
+
maxs = mins + 1.0
|
|
34
|
+
tree.bulk_load(mins, maxs)
|
|
35
|
+
|
|
36
|
+
assert tree.intersection([0.5, 0.5], [0.5, 0.5]) == [0]
|
|
37
|
+
assert tree.intersection([5.5, 5.5], [5.5, 5.5]) == [1]
|
|
38
|
+
assert tree.insert([10.0, 10.0], [11.0, 11.0]) == 2
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def test_ids_saturate_instead_of_wrapping():
|
|
42
|
+
max_id = 2**63 - 1
|
|
43
|
+
|
|
44
|
+
inserted = PyBBoxRTree(dims=2)
|
|
45
|
+
assert inserted.insert([0.0, 0.0], [1.0, 1.0], data=max_id) == max_id
|
|
46
|
+
assert inserted.insert([2.0, 2.0], [3.0, 3.0]) == max_id
|
|
47
|
+
|
|
48
|
+
bulk_loaded = PyBBoxRTree(dims=2)
|
|
49
|
+
bulk_loaded.bulk_load(
|
|
50
|
+
np.array([[0.0, 0.0]], dtype=np.float64),
|
|
51
|
+
np.array([[1.0, 1.0]], dtype=np.float64),
|
|
52
|
+
data=np.array([max_id], dtype=np.int64),
|
|
53
|
+
)
|
|
54
|
+
assert bulk_loaded.insert([2.0, 2.0], [3.0, 3.0]) == max_id
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def test_bulk_load_numpy_4d_replaces_contents():
|
|
58
|
+
tree = PyBBoxRTree(dims=4)
|
|
59
|
+
tree.insert([0.0] * 4, [1.0] * 4, data=99)
|
|
60
|
+
# Fortran order exercises the strided ndarray fast path.
|
|
61
|
+
mins = np.asfortranarray([[0.0] * 4, [5.0] * 4], dtype=np.float64)
|
|
62
|
+
maxs = np.asfortranarray([[2.0] * 4, [6.0] * 4], dtype=np.float64)
|
|
63
|
+
tree.bulk_load(mins, maxs, data=[20, 30])
|
|
64
|
+
|
|
65
|
+
assert len(tree) == 2
|
|
66
|
+
assert tree.intersection([1.0] * 4, [1.0] * 4) == [20]
|
|
67
|
+
assert tree.intersection([5.5] * 4, [5.5] * 4) == [30]
|
|
68
|
+
assert tree.insert([10.0] * 4, [11.0] * 4) == 100
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def test_numpy_and_batched_intersection_results():
|
|
72
|
+
tree = PyBBoxRTree(dims=2)
|
|
73
|
+
mins = np.array([[0.0, 0.0], [2.0, 2.0], [8.0, 8.0]], dtype=np.float64)
|
|
74
|
+
maxs = np.array([[4.0, 4.0], [6.0, 6.0], [9.0, 9.0]], dtype=np.float64)
|
|
75
|
+
strided_ids = np.array([10, -1, 20, -1, 30], dtype=np.int64)[::2]
|
|
76
|
+
tree.bulk_load(mins, maxs, data=strided_ids)
|
|
77
|
+
|
|
78
|
+
query_mins = np.asfortranarray([[3.0, 3.0], [6.1, 6.1]], dtype=np.float64)
|
|
79
|
+
query_maxs = np.asfortranarray([[8.0, 8.0], [7.9, 7.9]], dtype=np.float64)
|
|
80
|
+
ids, offsets = tree.intersection_batch(query_mins, query_maxs)
|
|
81
|
+
assert ids.dtype == offsets.dtype == np.int64
|
|
82
|
+
assert offsets.tolist() == [0, 3, 3]
|
|
83
|
+
assert sorted(ids[offsets[0] : offsets[1]].tolist()) == [10, 20, 30]
|
|
84
|
+
|
|
85
|
+
empty_ids, empty_offsets = tree.intersection_batch(
|
|
86
|
+
np.empty((0, 2), dtype=np.float64),
|
|
87
|
+
np.empty((0, 2), dtype=np.float64),
|
|
88
|
+
)
|
|
89
|
+
assert empty_ids.tolist() == []
|
|
90
|
+
assert empty_offsets.tolist() == [0]
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
def test_numpy_inputs_reject_wrong_dtypes():
|
|
94
|
+
tree = PyBBoxRTree(dims=2)
|
|
95
|
+
float32_corners = np.zeros((2, 2), dtype=np.float32)
|
|
96
|
+
float64_corners = np.zeros((2, 2), dtype=np.float64)
|
|
97
|
+
|
|
98
|
+
with pytest.raises(
|
|
99
|
+
ValueError,
|
|
100
|
+
match=r"min_corners is a numpy array of dtype float32; expected float64",
|
|
101
|
+
):
|
|
102
|
+
tree.bulk_load(float32_corners, float64_corners)
|
|
103
|
+
with pytest.raises(
|
|
104
|
+
ValueError,
|
|
105
|
+
match=r"data is a numpy array of dtype int32; expected int64",
|
|
106
|
+
):
|
|
107
|
+
tree.bulk_load(
|
|
108
|
+
float64_corners,
|
|
109
|
+
float64_corners,
|
|
110
|
+
data=np.arange(2, dtype=np.int32),
|
|
111
|
+
)
|
|
112
|
+
with pytest.raises(
|
|
113
|
+
ValueError,
|
|
114
|
+
match=r"min_corners is a numpy array of dtype float32; expected float64",
|
|
115
|
+
):
|
|
116
|
+
tree.intersection_batch(float32_corners, float64_corners)
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def test_remove_item_distinguishes_coincident_boxes():
|
|
120
|
+
tree = PyBBoxRTree(dims=2)
|
|
121
|
+
tree.insert([0.0, 0.0], [1.0, 1.0], data=10)
|
|
122
|
+
tree.insert([0.0, 0.0], [1.0, 1.0], data=20)
|
|
123
|
+
|
|
124
|
+
assert tree.remove_item([0.0, 0.0], [1.0, 1.0], 10) is True
|
|
125
|
+
assert tree.intersection([0.5, 0.5], [0.5, 0.5]) == [20]
|
|
126
|
+
assert tree.remove_item([0.0, 0.0], [1.0, 1.0], 10) is False
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def test_empty_bulk_load_resets_contents():
|
|
130
|
+
tree = PyBBoxRTree(dims=2)
|
|
131
|
+
tree.insert([0.0, 0.0], [1.0, 1.0])
|
|
132
|
+
tree.bulk_load([], [])
|
|
133
|
+
assert len(tree) == 0
|
|
134
|
+
assert tree.intersection([0.0, 0.0], [1.0, 1.0]) == []
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
@pytest.mark.parametrize("dims", [1, 9])
|
|
138
|
+
def test_invalid_tree_dimensions(dims):
|
|
139
|
+
with pytest.raises(ValueError, match="Dimensions must be between 2 and 8"):
|
|
140
|
+
PyBBoxRTree(dims=dims)
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def test_dimension_row_and_id_mismatch_errors():
|
|
144
|
+
tree = PyBBoxRTree(dims=2)
|
|
145
|
+
with pytest.raises(ValueError, match="Expected 2 dimensions, got 3"):
|
|
146
|
+
tree.insert([0.0, 0.0, 0.0], [1.0, 1.0])
|
|
147
|
+
with pytest.raises(ValueError, match="same number of rows"):
|
|
148
|
+
tree.bulk_load([[0.0, 0.0]], [])
|
|
149
|
+
with pytest.raises(ValueError, match="Expected 1 ids.*got 2"):
|
|
150
|
+
tree.bulk_load([[0.0, 0.0]], [[1.0, 1.0]], data=[1, 2])
|
|
151
|
+
with pytest.raises(ValueError, match="max_corners row 0 must have 2 dimensions"):
|
|
152
|
+
tree.bulk_load([[0.0, 0.0]], [[1.0, 1.0, 1.0]])
|
|
153
|
+
with pytest.raises(ValueError, match="min_corners must have 2 columns, got 4"):
|
|
154
|
+
tree.bulk_load(
|
|
155
|
+
np.empty((0, 4), dtype=np.float64),
|
|
156
|
+
np.empty((0, 2), dtype=np.float64),
|
|
157
|
+
)
|
|
158
|
+
with pytest.raises(ValueError, match="same number of rows"):
|
|
159
|
+
tree.intersection_batch(
|
|
160
|
+
np.empty((1, 2), dtype=np.float64),
|
|
161
|
+
np.empty((2, 2), dtype=np.float64),
|
|
162
|
+
)
|