def test_choose_leaf_uses_least_area_enlargement_for_higher_levels(
self, least_area_enlargement_mock, least_overlap_enlargement_mock):
"""
Ensure that the choose subtree strategy uses the least area enlargement strategy when picking a subtree at
levels higher than the one just above the leaf level.
"""
# Arrange
tree = RStarTree()
leaf = RTreeNode(tree, is_leaf=True)
intermediate = RTreeNode(
tree,
is_leaf=False,
entries=[RTreeEntry(Rect(0, 0, 0, 0), child=leaf)])
intermediate_entry = RTreeEntry(Rect(0, 0, 0, 0), child=intermediate)
root = RTreeNode(tree, is_leaf=False, entries=[intermediate_entry])
tree.root = root
e = RTreeEntry(Rect(0, 0, 0, 0))
least_area_enlargement_mock.return_value = intermediate_entry
# Act
rstar_choose_leaf(tree, e)
# Assert
least_area_enlargement_mock.assert_called_once_with(
root.entries, e.rect)
least_overlap_enlargement_mock.assert_called_once_with(
intermediate.entries, e.rect)
def test_get_rstar_stat_same_distribution_for_all_4_sort_types(self):
"""
Tests get_rstar_stat when all 4 sort types (min_x, max_x, min_y, and max_y) result in the same distribution
of entries.
"""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 0, 1, 1))
b = RTreeEntry(data='b', rect=Rect(1, 1, 2, 2))
c = RTreeEntry(data='c', rect=Rect(2, 2, 3, 3))
d = RTreeEntry(data='d', rect=Rect(3, 3, 4, 4))
# Act
stat = get_rstar_stat([a, b, c, d], 1, 3)
# Assert
unique_distributions = [
EntryDistribution(([a], [b, c, d])),
EntryDistribution(([a, b], [c, d])),
EntryDistribution(([a, b, c], [d]))
]
self.assertCountEqual(unique_distributions, stat.unique_distributions)
self.assertCountEqual(unique_distributions,
stat.get_axis_unique_distributions('x'))
self.assertCountEqual(unique_distributions,
stat.get_axis_unique_distributions('y'))
self.assertEqual(96, stat.get_axis_perimeter('x'))
self.assertEqual(96, stat.get_axis_perimeter('y'))
def test_rstar_insert_empty(self):
"""Tests inserting into an empty tree"""
# Arrange
tree = RStarTree(min_entries=1, max_entries=3)
# Act
tree.insert('a', Rect(0, 0, 5, 5))
# Assert
# Ensure root entry has the correct data and bounding box
self.assertEqual(1, len(tree.root.entries))
e = tree.root.entries[0]
self.assertEqual('a', e.data)
self.assertEqual(Rect(0, 0, 5, 5), e.rect)
self.assertIsNone(e.child)
# Ensure root node has correct structure
node = tree.root
self.assertTrue(node.is_root)
self.assertTrue(node.is_leaf)
self.assertEqual(Rect(0, 0, 5, 5), tree.root.get_bounding_rect())
# Ensure root entry has correct structure
self.assertTrue(e.is_leaf)
self.assertIsNone(e.child)
# Ensure there is only 1 level and 1 node in the tree
self.assertEqual(1, len(tree.get_levels()))
self.assertEqual(1, len(list(tree.get_nodes())))
def test_get_rstar_stat_sorts_entries_by_both_min_and_max(self):
"""
List of possible divisions should be based on entries sorted by both the minimum as well as maximum coordinate.
In the example below, when the entries are sorted by either minx, miny, or maxy, the sort order is always
(a,b,c), but when sorted by maxx, the order is (b,a,c). This ordering enables the [(b), (a,c)] division (which
turns out to be optimal).
"""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 0, 7, 2))
b = RTreeEntry(data='b', rect=Rect(1, 1, 2, 3))
c = RTreeEntry(data='c', rect=Rect(2, 2, 8, 4))
# Act
stat = get_rstar_stat([a, b, c], 1, 2)
# Assert
self.assertCountEqual([
EntryDistribution(([a], [b, c])),
EntryDistribution(([a, b], [c])),
EntryDistribution(([b], [a, c]))
], stat.unique_distributions)
self.assertCountEqual([
EntryDistribution(([a], [b, c])),
EntryDistribution(([a, b], [c])),
EntryDistribution(([b], [a, c]))
], stat.get_axis_unique_distributions('x'))
self.assertCountEqual([
EntryDistribution(([a], [b, c])),
EntryDistribution(([a, b], [c]))
], stat.get_axis_unique_distributions('y'))
self.assertEqual(140, stat.get_axis_perimeter('x'))
self.assertEqual(148, stat.get_axis_perimeter('y'))
def test_rstar_insert_no_split(self):
"""Tests multiple inserts which do not require a node split"""
# Arrange
tree = RStarTree(min_entries=1, max_entries=2)
# Act
tree.insert('a', Rect(0, 0, 5, 2))
tree.insert('b', Rect(2, 3, 4, 7))
# Assert
# Root node
self.assertTrue(tree.root.is_root)
self.assertTrue(tree.root.is_leaf)
self.assertEqual(2, len(tree.root.entries))
self.assertEqual(Rect(0, 0, 5, 7), tree.root.get_bounding_rect())
# Entry 'a'
entry_a = next((e for e in tree.root.entries if e.data == 'a'))
self.assertEqual(Rect(0, 0, 5, 2), entry_a.rect)
self.assertTrue(entry_a.is_leaf)
self.assertIsNone(entry_a.child)
# Entry 'b'
entry_b = next((e for e in tree.root.entries if e.data == 'b'))
self.assertEqual(Rect(2, 3, 4, 7), entry_b.rect)
self.assertTrue(entry_b.is_leaf)
self.assertIsNone(entry_b.child)
def test_union_disjoint(self):
"""Tests union of two disjoint (non-intersecting) rectanges"""
# Arrange
r1 = Rect(min_x=0, min_y=-1, max_x=3, max_y=3)
r2 = Rect(min_x=2, min_y=5, max_x=6, max_y=8)
# Act
rect = r1.union(r2)
# Assert
self.assertEqual(Rect(min_x=0, min_y=-1, max_x=6, max_y=8), rect)
def test_union_intersecting(self):
"""Tests union of two intersecting rectanges"""
# Arrange
r1 = Rect(min_x=0, min_y=0, max_x=3, max_y=3)
r2 = Rect(min_x=-2, min_y=-2, max_x=2, max_y=2)
# Act
rect = r1.union(r2)
# Assert
self.assertEqual(Rect(min_x=-2, min_y=-2, max_x=3, max_y=3), rect)
def test_rect_centroid(self):
"""Tests calculating the centroid of a rectangle"""
# Arrange
r = Rect(2, 2, 6, 5)
# Act
centroid = r.centroid()
# Assert
self.assertTrue(isclose(4, centroid[0], rel_tol=EPSILON))
self.assertTrue(isclose(3.5, centroid[1], rel_tol=EPSILON))
def test_intersection_area(self):
"""Tests getting the intersection area of two intersecting rectangles"""
# Arrange
r1 = Rect(min_x=0, min_y=0, max_x=4, max_y=4)
r2 = Rect(min_x=2, min_y=2, max_x=5, max_y=5)
# Act
area = r1.get_intersection_area(r2)
# Assert
self.assertEqual(4, area)
def test_union_contains(self):
"""Tests union of two rectanges where one contains the other"""
# Arrange
r1 = Rect(min_x=1, min_y=1, max_x=3, max_y=3)
r2 = Rect(min_x=-2, min_y=-2, max_x=5, max_y=5)
# Act
rect = r1.union(r2)
# Assert
self.assertEqual(Rect(min_x=-2, min_y=-2, max_x=5, max_y=5), rect)
def test_insert_creates_entry(self):
"""Basic test ensuring an insert creates an entry."""
# Arrange
t = RTree()
# Act
e = t.insert('foo', Rect(0, 0, 1, 1))
# Assert
self.assertEqual('foo', e.data)
self.assertEqual(Rect(0, 0, 1, 1), e.rect)
def test_intersects_disjoint(self):
"""Ensures intersects returns False when two rectangles are completely disjoint."""
# Arrange
r1 = Rect(min_x=0, min_y=0, max_x=5, max_y=2)
r2 = Rect(min_x=1, min_y=5, max_x=3, max_y=9)
# Act
result = r1.intersects(r2)
# Assert
self.assertFalse(result)
def test_intersects(self):
"""Ensures intersects returns True when two rectangles intersect."""
# Arrange
r1 = Rect(min_x=0, min_y=0, max_x=5, max_y=2)
r2 = Rect(min_x=2, min_y=1, max_x=4, max_y=3)
# Act
result = r1.intersects(r2)
# Assert
self.assertTrue(result)
def test_intersection_area_disjoint(self):
"""
Tests getting the intersection area of two completely disjoint (non-intersecting) rectangles with no overlap in
either dimension.
"""
# Arrange
r1 = Rect(min_x=0, min_y=0, max_x=3, max_y=3)
r2 = Rect(min_x=5, min_y=5, max_x=8, max_y=8)
# Act
area = r1.get_intersection_area(r2)
# Assert
self.assertEqual(0, area)
def test_intersection_area_disjoint_y_overlap(self):
"""
Tests getting the intersection area of two disjoint rectangles where there is overlap in the y-dimension but not
in the x-dimension.
"""
# Arrange
r1 = Rect(min_x=0, min_y=0, max_x=2, max_y=4)
r2 = Rect(min_x=2, min_y=2, max_x=3, max_y=3)
# Act
area = r1.get_intersection_area(r2)
# Assert
self.assertEqual(0, area)
def test_least_overlap_enlargement_tie(self):
"""Ensure least area enlargement is used as a tie-breaker when overlap enlargements are equal."""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 0, 4, 5))
b = RTreeEntry(data='b', rect=Rect(3, 4, 5, 6))
rect = Rect(2, 5, 3, 6)
# Act
entry = least_overlap_enlargement([a, b], rect)
# Assert
self.assertEqual(b, entry)
def test_intersects_touches(self):
"""
Ensures intersects returns False when two rectangles merely touch along a border but do not have any interior
intersection area.
"""
# Arrange
r1 = Rect(min_x=0, min_y=0, max_x=5, max_y=2)
r2 = Rect(min_x=5, min_y=0, max_x=7, max_y=2)
# Act
result = r1.intersects(r2)
# Assert
self.assertFalse(result)
def test_least_overlap_enlargement(self):
"""
Basic test of least overlap enlargement helper method. This test demonstrates a scenario where least area
enlargement would favor one entry, but least overlap enlargement favors another.
"""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 0, 4, 5))
b = RTreeEntry(data='b', rect=Rect(2, 4, 5, 6))
rect = Rect(4, 3, 5, 4)
# Act
entry = least_overlap_enlargement([a, b], rect)
# Assert
self.assertEqual(a, entry)
def test_least_area_enlargement(self):
"""
Ensure the node whose bounding box needs least enlargement is chosen for a new entry in the case where there is
a clear winner.
"""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 0, 3, 3))
b = RTreeEntry(data='b', rect=Rect(9, 9, 10, 10))
rect = Rect(2, 2, 4, 4)
# Act
entry = least_area_enlargement([a, b], rect)
# Assert
self.assertEqual(a, entry)
def test_bounding_rect_multiple_inserts_without_split(self):
"""
Ensure root note bounding rect encompasses the bounding rect of all entries after multiple inserts (without
forcing a split)
"""
# Arrange
t = RTree(max_entries=5)
# Act
t.insert('a', Rect(0, 0, 5, 5))
t.insert('b', Rect(1, 1, 3, 3))
t.insert('c', Rect(4, 4, 6, 6))
# Assert
rect = t.root.get_bounding_rect()
self.assertEqual(Rect(0, 0, 6, 6), rect)
def test_least_area_enlargement_tie(self):
"""
When two nodes need to be enlarged by the same amount, the strategy should pick the node having the smallest
area as a tie-breaker.
"""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 0, 4, 2))
b = RTreeEntry(data='b', rect=Rect(5, 1, 7, 3))
c = RTreeEntry(data='c', rect=Rect(0, 4, 1, 5))
rect = Rect(4, 1, 5, 2)
# Act
entry = least_area_enlargement([a, b, c], rect)
# Assert
self.assertEqual(b, entry)
def test_init(self):
"""Basic test ensuring a Rect can be instantiated"""
r = Rect(0, 1, 5, 9)
self.assertEqual(0, r.min_x)
self.assertEqual(1, r.min_y)
self.assertEqual(5, r.max_x)
self.assertEqual(9, r.max_y)
def _add_electoral_wards(dataset_id):
areas_to_add = []
for feature in geojson.loads(wd20_filepath.read_text())["features"]:
ward_code = feature["properties"]["wd20cd"]
ward_name = feature["properties"]["wd20nm"]
ward_id = "wd20-" + ward_code
print() # noqa: T001
print(ward_name) # noqa: T001
try:
la_id = "lad20-" + ward_code_to_la_id_mapping[ward_code]
feature, simple_feature = (polygons_and_simplified_polygons(
feature["geometry"]))
if feature:
rtree_index.insert(ward_id, Rect(*Polygons(feature).bounds))
areas_to_add.append([
ward_id,
ward_name,
dataset_id,
la_id,
feature,
simple_feature,
estimate_number_of_smartphones_in_area(ward_code),
])
except KeyError:
print("Skipping", ward_code, ward_name) # noqa: T001
rtree_index_path.open('wb').write(pickle.dumps(rtree_index))
repo.insert_broadcast_areas(areas_to_add, keep_old_polygons)
def overlapping_areas(self):
if not self.polygons:
return []
return broadcast_area_libraries.get_areas([
overlap.data
for overlap in rtree_index.query(Rect(*self.polygons.bounds))
])
def test_multiple_inserts_without_split(self):
"""
Ensure multiple inserts work (all original entries are returned) without a split (fewer entries than
max_entries)
"""
# Arrange
t = RTree(max_entries=5)
# Act
t.insert('a', Rect(0, 0, 5, 5))
t.insert('b', Rect(1, 1, 3, 3))
t.insert('c', Rect(4, 4, 6, 6))
# Assert
entries = list(t.get_leaf_entries())
self.assertCountEqual(['a', 'b', 'c'],
[entry.data for entry in entries])
def test_choose_split_axis(self):
"""
Ensure split axis is chosen based on smallest overall perimeter of all possible divisions of a list of entries.
In the below scenario, there is a clear winner with the best division being ([a, b, c], [d]).
"""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 0, 1, 1))
b = RTreeEntry(data='b', rect=Rect(1, 0, 2, 1))
c = RTreeEntry(data='c', rect=Rect(2, 0, 3, 1))
d = RTreeEntry(data='d', rect=Rect(1, 7, 2, 8))
stat = get_rstar_stat([a, b, c, d], 1, 3)
# Act
result = choose_split_axis(stat)
# Assert
self.assertEqual('y', result)
def test_quadratic_split(self):
"""Ensures that a split results in the smallest total area."""
# Arrange
t = RTreeGuttman(max_entries=4)
t.insert('a', Rect(2, 8, 5, 9))
t.insert('b', Rect(4, 0, 5, 10))
t.insert('c', Rect(5, 0, 6, 10))
t.insert('d', Rect(5, 7, 8, 8))
# Act
split_node = quadratic_split(t, t.root)
# Assert
group1 = [e.data for e in t.root.entries]
group2 = [e.data for e in split_node.entries]
self.assertCountEqual(['a', 'd'], group1)
self.assertCountEqual(['b', 'c'], group2)
def test_choose_split_index(self):
"""Ensures best split index is chosen based on minimum overlap."""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 1, 4, 5))
b = RTreeEntry(data='b', rect=Rect(3, 5, 6, 8))
c = RTreeEntry(data='c', rect=Rect(7, 0, 9, 4))
d = RTreeEntry(data='d', rect=Rect(8, 7, 10, 9))
distributions = [
EntryDistribution(([a], [b, c, d])),
EntryDistribution(([a, b], [c, d])),
EntryDistribution(([a, b, c], [d]))
]
# Act
i = choose_split_index(distributions)
# Assert
self.assertEqual(1, i)
def test_choose_split_index_tie(self):
"""When multiple divisions have the same overlap, ensure split index is chosen based on minimum area."""
# Arrange
a = RTreeEntry(data='a', rect=Rect(0, 0, 2, 1))
b = RTreeEntry(data='b', rect=Rect(1, 0, 3, 2))
c = RTreeEntry(data='c', rect=Rect(2, 2, 4, 3))
d = RTreeEntry(data='d', rect=Rect(9, 9, 10, 10))
distributions = [
EntryDistribution(([a], [b, c, d])),
EntryDistribution(([a, b], [c, d])),
EntryDistribution(([a, b, c], [d]))
]
# Act
i = choose_split_index(distributions)
# Assert
self.assertEqual(2, i)
def test_rstar_overflow_split_root(self):
"""
When the root node overflows, the root node should be split and the tree should grow a level. Forced reinsert
should not occur at the root level.
"""
# Arrange
t = RStarTree(max_entries=3)
r1 = Rect(0, 0, 3, 2)
r2 = Rect(7, 7, 10, 9)
r3 = Rect(2, 1, 5, 3)
entry_a = RTreeEntry(r1, data='a')
entry_b = RTreeEntry(r2, data='b')
entry_c = RTreeEntry(r3, data='c')
t.root.entries = [entry_a, entry_b, entry_c]
# Arrange entry being inserted. Since the root node is at max capacity, this entry should cause the root
# to overflow.
r4 = Rect(6, 6, 8, 8)
# Act
entry_d = t.insert('d', r4)
# Assert
# Root node should no longer be a leaf node (but should still be root)
self.assertFalse(t.root.is_leaf)
self.assertTrue(t.root.is_root)
# Root node bounding box should encompass all entries
self.assertEqual(Rect(0, 0, 10, 9), t.root.get_bounding_rect())
# Root node should have 2 child entries
self.assertEqual(2, len(t.root.entries))
e1 = t.root.entries[0]
e2 = t.root.entries[1]
# e1 bounding box should encompass entries [a, c]
self.assertEqual(Rect(0, 0, 5, 3), e1.rect)
# e2 bounding box should encompass entries [b ,d]
self.assertEqual(Rect(6, 6, 10, 9), e2.rect)
# Ensure children nodes of e1 and e2 are leaf nodes
leaf_node_1 = e1.child
leaf_node_2 = e2.child
self.assertIsNotNone(leaf_node_1)
self.assertIsNotNone(leaf_node_2)
self.assertTrue(leaf_node_1.is_leaf)
self.assertTrue(leaf_node_2.is_leaf)
# Leaf node 1 should contain entries [a, c]
self.assertEqual(Rect(0, 0, 5, 3), leaf_node_1.get_bounding_rect())
self.assertCountEqual([entry_a, entry_c], leaf_node_1.entries)
# Leaf node 2 should contain entries [b, d]
self.assertEqual(Rect(6, 6, 10, 9), leaf_node_2.get_bounding_rect())
self.assertCountEqual([entry_b, entry_d], leaf_node_2.entries)