def test_rtree_empty(bounds_array, page_size):
# Build rtree
rt = HilbertRtree(bounds_array, page_size=page_size)
n = bounds_array.shape[1] // 2
query_bounds = (-np.inf, ) * n + (np.inf, ) * n
assert set(rt.intersects(query_bounds)) == set()
covers, overlaps = rt.covers_overlaps(query_bounds)
assert set(covers) == set()
assert set(overlaps) == set()
assert rt.total_bounds == (np.nan, ) * (2 * n)
def test_rtree_intersects_none(bounds_array, page_size):
# Build rtree
rt = HilbertRtree(bounds_array, page_size=page_size)
# query with query array
n = bounds_array.shape[1] // 2
b = ([11] * n) + ([100] * n)
intersected = set(rt.intersects(np.array(b)))
# Nothing should be selected
assert intersected == set()
def test_rtree_intersects_different_bounds_2d(bounds_array, query_array,
page_size):
# Build rtree
rt = HilbertRtree(bounds_array, page_size=page_size)
# query with query array
for i in range(query_array.shape[0]):
b = query_array[i, :]
intersected = set(rt.intersects(np.array(b)))
# Compare to brute force implementation
assert intersected == intersects_bruteforce(b, bounds_array)
def test_rtree_intersects_all(bounds_array, page_size):
# Build rtree
rt = HilbertRtree(bounds_array, page_size=page_size)
# query with query array
n = bounds_array.shape[1] // 2
b_mins = bounds_array[:, :n].min(axis=0)
b_maxes = bounds_array[:, n:].max(axis=0)
b = np.concatenate([b_mins, b_maxes])
intersected = set(rt.intersects(np.array(b)))
# All indices should be selected
assert intersected == set(range(bounds_array.shape[0]))
def test_rtree_intersects_input_bounds(bounds_array, page_size):
# Build rtree
rt = HilbertRtree(bounds_array, page_size=page_size)
# query with bounds array, these were the input
for i in range(bounds_array.shape[0]):
b = bounds_array[i, :]
intersected = set(rt.intersects(np.array(b)))
# Should at least select itself
assert i in intersected
# Compare to brute force implementation
assert intersected == intersects_bruteforce(b, bounds_array)
def test_rtree_pickle(bounds_array, page_size):
# Build rtree
rt = HilbertRtree(bounds_array, page_size=page_size)
# Serialize uninitialized rtree
rt2 = pickle.loads(pickle.dumps(rt))
if bounds_array.size == 0:
assert isinstance(rt2, HilbertRtree)
return
rt2_result = set(rt2.intersects(bounds_array[0, :]))
# Call intersection to construct numba rtree
rt_result = set(rt.intersects(bounds_array[0, :]))
# Serialize initialized rtree
rt3 = pickle.loads(pickle.dumps(rt))
rt3_result = set(rt3.intersects(bounds_array[0, :]))
assert rt_result == rt2_result and rt_result == rt3_result
def test_rtree_covers_overlaps_input_bounds(bounds_array, page_size):
# Build rtree
rt = HilbertRtree(bounds_array, page_size=page_size)
# query with bounds array, these were the input
for i in range(bounds_array.shape[0]):
b = bounds_array[i, :]
covers, overlaps = rt.covers_overlaps(np.array(b))
covers = set(covers)
overlaps = set(overlaps)
# Should have nothing in common
assert covers.isdisjoint(overlaps)
# covered should contain all bounding boxes that are fully covered by query
all_covers = covers_bruteforce(b, bounds_array)
if covers != all_covers:
rt.covers_overlaps(np.array(b))
assert covers == all_covers
# covered and overlaps together should contain all intersecting bounding
# boxes (overlaps will typically contain others as well)
intersects = intersects_bruteforce(b, bounds_array)
assert intersects == covers.union(overlaps)
def sindex(self):
if self._sindex is None:
self._sindex = HilbertRtree(self.bounds)
return self._sindex
def partition_sindex(self):
if self._partition_sindex is None:
self._partition_sindex = HilbertRtree(self.partition_bounds.values)
return self._partition_sindex