def setUp(self):
self.big_rect = rect2.Rect2(1000, 1000)
self.big_rect_sub_1 = rect2.Rect2(500, 500)
self.big_rect_sub_2 = rect2.Rect2(500, 500, vector2.Vector2(500, 0))
self.big_rect_sub_3 = rect2.Rect2(500, 500, vector2.Vector2(500, 500))
self.big_rect_sub_4 = rect2.Rect2(500, 500, vector2.Vector2(0, 500))
random.seed()
def test_think(self):
ent1 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(15, 15)))
ent2 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(20, 20)))
ent3 = quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(0,
0)))
ent4 = quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(5,
0)))
ent5 = quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(0,
5)))
_tree = quadtree.QuadTree(2,
2,
self.big_rect,
entities=[ent1, ent2, ent3, ent4, ent5])
_tree.think(True)
self.assertIsNotNone(_tree.children) # depth 1
self.assertIsNotNone(_tree.children[0].children) # depth 2
self.assertIsNone(_tree.children[0].children[0].children
) # depth 3 shouldn't happen because
self.assertEqual(5, len(
_tree.children[0].children[0].entities)) # max_depth reached
_tree2 = quadtree.QuadTree(2, 2, self.big_rect, entities=[ent1, ent2])
_tree2.think(True)
self.assertIsNone(_tree2.children)
def split(self):
"""
Split this quadtree.
.. caution::
A call to split will always split the tree or raise an error. Use
:py:meth:`.think` if you want to ensure the quadtree is operating
efficiently.
.. caution::
This function will not respect :py:attr:`.bucket_size` or
:py:attr:`.max_depth`.
:raises ValueError: if :py:attr:`.children` is not empty
"""
if self.children:
raise ValueError("cannot split twice")
_cls = type(self)
def _cstr(r):
return _cls(self.bucket_size, self.max_depth, r, self.depth + 1)
_halfwidth = self.location.width / 2
_halfheight = self.location.height / 2
_x = self.location.mincorner.x
_y = self.location.mincorner.y
self.children = [
_cstr(rect2.Rect2(_halfwidth, _halfheight, vector2.Vector2(_x,
_y))),
_cstr(
rect2.Rect2(_halfwidth, _halfheight,
vector2.Vector2(_x + _halfwidth, _y))),
_cstr(
rect2.Rect2(_halfwidth, _halfheight,
vector2.Vector2(_x + _halfwidth,
_y + _halfheight))),
_cstr(
rect2.Rect2(_halfwidth, _halfheight,
vector2.Vector2(_x, _y + _halfheight)))
]
_newents = []
for ent in self.entities:
quad = self.get_quadrant(ent)
if quad < 0:
_newents.append(ent)
else:
self.children[quad].entities.append(ent)
self.entities = _newents
def test_split_entities(self):
ent1 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(50, 50)))
ent2 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(550, 75)))
ent3 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(565, 585)))
ent4 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(95, 900)))
ent5 = quadtree.QuadTreeEntity(
rect2.Rect2(10, 10, vector2.Vector2(495, 167)))
_tree = quadtree.QuadTree(64,
5,
self.big_rect,
entities=[ent1, ent2, ent3, ent4, ent5])
_tree.split()
self.assertEqual(1, len(_tree.children[0].entities))
self.assertEqual(50, _tree.children[0].entities[0].aabb.mincorner.x)
self.assertEqual(50, _tree.children[0].entities[0].aabb.mincorner.y)
self.assertEqual(1, len(_tree.children[1].entities))
self.assertEqual(550, _tree.children[1].entities[0].aabb.mincorner.x)
self.assertEqual(75, _tree.children[1].entities[0].aabb.mincorner.y)
self.assertEqual(1, len(_tree.children[2].entities))
self.assertEqual(565, _tree.children[2].entities[0].aabb.mincorner.x)
self.assertEqual(585, _tree.children[2].entities[0].aabb.mincorner.y)
self.assertEqual(1, len(_tree.children[3].entities))
self.assertEqual(95, _tree.children[3].entities[0].aabb.mincorner.x)
self.assertEqual(900, _tree.children[3].entities[0].aabb.mincorner.y)
self.assertEqual(1, len(_tree.entities))
self.assertEqual(495, _tree.entities[0].aabb.mincorner.x)
self.assertEqual(167, _tree.entities[0].aabb.mincorner.y)
_tree2 = _tree.children[3]
_tree2.split()
for i in range(3):
self.assertEqual(0,
len(_tree2.children[i].entities),
msg="i={}".format(i))
self.assertEqual(1, len(_tree2.children[3].entities))
self.assertEqual(95, _tree2.children[3].entities[0].aabb.mincorner.x)
self.assertEqual(900, _tree2.children[3].entities[0].aabb.mincorner.y)
def test_repr(self):
_tree = quadtree.QuadTree(1, 5, rect2.Rect2(100, 100))
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(2, 2, vector2.Vector2(5, 5))))
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(2, 2, vector2.Vector2(95, 5))))
_olddiff = self.maxDiff
def cleanup(self2=self):
self2.maxDiff = _olddiff
self.addCleanup(cleanup)
self.maxDiff = None
self.assertEqual(
"quadtree(bucket_size=1, max_depth=5, location=rect2(width=100, height=100, mincorner=vector2(x=0, y=0)), depth=0, entities=[], children=[quadtree(bucket_size=1, max_depth=5, location=rect2(width=50.0, height=50.0, mincorner=vector2(x=0, y=0)), depth=1, entities=[quadtreeentity(aabb=rect2(width=2, height=2, mincorner=vector2(x=5, y=5)))], children=None), quadtree(bucket_size=1, max_depth=5, location=rect2(width=50.0, height=50.0, mincorner=vector2(x=50.0, y=0)), depth=1, entities=[quadtreeentity(aabb=rect2(width=2, height=2, mincorner=vector2(x=95, y=5)))], children=None), quadtree(bucket_size=1, max_depth=5, location=rect2(width=50.0, height=50.0, mincorner=vector2(x=50.0, y=50.0)), depth=1, entities=[], children=None), quadtree(bucket_size=1, max_depth=5, location=rect2(width=50.0, height=50.0, mincorner=vector2(x=0, y=50.0)), depth=1, entities=[], children=None)])",
repr(_tree))
def test_nodes_per_depth(self):
_tree = quadtree.QuadTree(1, 5, self.big_rect)
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(50,
50))))
self.assertDictEqual({0: 1}, _tree.find_nodes_per_depth())
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(450, 450))))
self.assertDictEqual({0: 1, 1: 4, 2: 4}, _tree.find_nodes_per_depth())
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(550, 550))))
self.assertDictEqual({0: 1, 1: 4, 2: 4}, _tree.find_nodes_per_depth())
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(850, 550))))
self.assertDictEqual({0: 1, 1: 4, 2: 8}, _tree.find_nodes_per_depth())
def test_insert(self):
_tree = quadtree.QuadTree(2, 2, self.big_rect)
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(15,
15))))
self.assertIsNone(_tree.children)
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(20,
20))))
self.assertIsNone(_tree.children)
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(0, 0))))
self.assertIsNotNone(_tree.children) # depth 1
self.assertIsNotNone(_tree.children[0].children) # depth 2
self.assertIsNone(_tree.children[0].children[0].children
) # depth 3 shouldn't happen because
self.assertEqual(3, len(
_tree.children[0].children[0].entities)) # max_depth reached
def test_get_quadrant_shifted(self):
_tree = quadtree.QuadTree(64, 5, self.big_rect_sub_3)
ent1 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(515, 600)))
self.assertEqual(0, _tree.get_quadrant(ent1))
ent2 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(800, 550)))
self.assertEqual(1, _tree.get_quadrant(ent2))
ent3 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(950, 850)))
self.assertEqual(2, _tree.get_quadrant(ent3))
ent4 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(515, 751)))
self.assertEqual(3, _tree.get_quadrant(ent4))
def test_str(self):
_tree = quadtree.QuadTree(1, 5, rect2.Rect2(100, 100))
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(2, 2, vector2.Vector2(5, 5))))
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(2, 2, vector2.Vector2(95, 5))))
_olddiff = self.maxDiff
def cleanup(self2=self):
self2.maxDiff = _olddiff
self.addCleanup(cleanup)
self.maxDiff = None
self.assertEqual(
"quadtree(at rect(100x100 at <0, 0>) with 0 entities here (2 in total); (nodes, entities) per depth: [ 0: (1, 0), 1: (4, 2) ] (allowed max depth: 5, actual: 1), avg ent/leaf: 0.5 (target 1), misplaced weight 0.0 (0 best, >1 bad)",
str(_tree))
def test_misplaced_ents(self):
_tree = quadtree.QuadTree(3, 5, self.big_rect)
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(75,
35))))
self.assertEqual(
0, _tree.calculate_weight_misplaced_ents()) # 0 misplaced, 1 total
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(300, 499))))
self.assertEqual(
0, _tree.calculate_weight_misplaced_ents()) # 0 misplaced, 2 total
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(800, 600))))
self.assertEqual(
0, _tree.calculate_weight_misplaced_ents()) # 0 misplaced 3 total
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(550, 700))))
self.assertAlmostEqual(1, _tree.calculate_weight_misplaced_ents()
) # 1 misplaced (1 deep), 4 total
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(900, 900))))
self.assertAlmostEqual(4 / 5, _tree.calculate_weight_misplaced_ents()
) # 1 misplaced (1 deep), 5 total
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(950, 950))))
self.assertAlmostEqual(8 / 6, _tree.calculate_weight_misplaced_ents()
) # 1 misplaced (2 deep), 6 total
def test_avg_ents_per_leaf(self):
_tree = quadtree.QuadTree(3, 5, self.big_rect)
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(75,
35))))
self.assertEqual(
1, _tree.calculate_avg_ents_per_leaf()) # 1 ent on 1 leaf
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(300, 499))))
self.assertEqual(2,
_tree.calculate_avg_ents_per_leaf()) # 2 ents 1 leaf
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(800, 600))))
self.assertEqual(3,
_tree.calculate_avg_ents_per_leaf()) # 3 ents 1 leaf
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(450, 300))))
self.assertEqual(0.75, _tree.calculate_avg_ents_per_leaf()
) # 3 ents 4 leafs (1 misplaced)
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(150, 100))))
self.assertEqual(1, _tree.calculate_avg_ents_per_leaf()
) # 4 ents 4 leafs (1 misplaced)
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(450, 450))))
self.assertAlmostEqual(5 / 7, _tree.calculate_avg_ents_per_leaf()
) # 5 ents 7 leafs (1 misplaced)
def test_ents_per_depth(self):
_tree = quadtree.QuadTree(3, 5, self.big_rect)
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(75,
35))))
self.assertDictEqual({0: 1}, _tree.find_entities_per_depth())
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(300, 499))))
self.assertDictEqual({0: 2}, _tree.find_entities_per_depth())
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(800, 600))))
self.assertDictEqual({0: 3}, _tree.find_entities_per_depth())
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(450, 300))))
self.assertDictEqual({0: 1, 1: 3}, _tree.find_entities_per_depth())
_tree.insert_and_think(
quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(150, 100))))
self.assertDictEqual({0: 1, 1: 4}, _tree.find_entities_per_depth())
_tree.insert_and_think(
quadtree.QuadTreeEntity(rect2.Rect2(5, 5, vector2.Vector2(80,
40))))
self.assertDictEqual({
0: 1,
1: 1,
2: 4
}, _tree.find_entities_per_depth())
def test_get_quadrant_none(self):
_tree = quadtree.QuadTree(64, 5, self.big_rect)
ent1 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(497, 150)))
self.assertEqual(-1, _tree.get_quadrant(ent1))
ent2 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(800, 499)))
self.assertEqual(-1, _tree.get_quadrant(ent2))
ent3 = quadtree.QuadTreeEntity(
rect2.Rect2(15, 15, vector2.Vector2(486, 505)))
self.assertEqual(-1, _tree.get_quadrant(ent3))
ent4 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 20, vector2.Vector2(15, 490)))
self.assertEqual(-1, _tree.get_quadrant(ent4))
ent5 = quadtree.QuadTreeEntity(
rect2.Rect2(17, 34, vector2.Vector2(485, 470)))
self.assertEqual(-1, _tree.get_quadrant(ent5))
def test_get_quadrant(self):
_tree = quadtree.QuadTree(64, 5, self.big_rect)
ent1 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(320, 175)))
quad1 = _tree.get_quadrant(ent1)
self.assertEqual(0, quad1)
ent2 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(600, 450)))
quad2 = _tree.get_quadrant(ent2)
self.assertEqual(1, quad2)
ent3 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(700, 950)))
quad3 = _tree.get_quadrant(ent3)
self.assertEqual(2, quad3)
ent4 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(0, 505)))
quad4 = _tree.get_quadrant(ent4)
self.assertEqual(3, quad4)
def test_sum_ents(self):
# it shouldn't matter where we put entities in, adding entities
# to a quadtree should increment this number by 1. So lets fuzz!
_tree = quadtree.QuadTree(64, 5, self.big_rect)
for i in range(1000):
w = random.randrange(1, 10)
h = random.randrange(1, 10)
x = random.uniform(0, 1000 - w)
y = random.uniform(0, 1000 - h)
ent = quadtree.QuadTreeEntity(
rect2.Rect2(w, h, vector2.Vector2(x, y)))
_tree.insert_and_think(ent)
# avoid calculating sum every loop which would take way too long.
# on average, try to sum about 50 times total (5% of the time),
# evenly split between both ways of summing
rnd = random.random()
if rnd > 0.95 and rnd <= 0.975:
_sum = _tree.sum_entities()
self.assertEqual(i + 1, _sum)
elif rnd > 0.975:
_sum = _tree.sum_entities(_tree.find_entities_per_depth())
self.assertEqual(i + 1, _sum)
def test_get_quadrant_0_shifted(self):
_tree = quadtree.QuadTree(
64, 5, rect2.Rect2(500, 800, vector2.Vector2(200, 200)))
ent1 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 10, vector2.Vector2(445, 224)))
self.assertEqual(-1, _tree.get_quadrant(ent1))
ent2 = quadtree.QuadTreeEntity(
rect2.Rect2(11, 17, vector2.Vector2(515, 585)))
self.assertEqual(-1, _tree.get_quadrant(ent2))
ent3 = quadtree.QuadTreeEntity(
rect2.Rect2(20, 20, vector2.Vector2(440, 700)))
self.assertEqual(-1, _tree.get_quadrant(ent3))
ent4 = quadtree.QuadTreeEntity(
rect2.Rect2(15, 15, vector2.Vector2(215, 590)))
self.assertEqual(-1, _tree.get_quadrant(ent4))
ent5 = quadtree.QuadTreeEntity(
rect2.Rect2(7, 12, vector2.Vector2(449, 589)))
self.assertEqual(-1, _tree.get_quadrant(ent5))
def test_retrieve(self):
_tree = quadtree.QuadTree(2, 2, self.big_rect)
ent1 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(25, 25)))
_tree.insert_and_think(ent1)
retr = _tree.retrieve_collidables(ent1)
self.assertIsNotNone(retr)
self.assertEqual(1, len(retr))
self.assertEqual(25, retr[0].aabb.mincorner.x)
self.assertEqual(25, retr[0].aabb.mincorner.y)
# note this is not nicely in a quadrant
ent2 = quadtree.QuadTreeEntity(
rect2.Rect2(20, 10, vector2.Vector2(490, 300)))
_tree.insert_and_think(ent2)
retr = _tree.retrieve_collidables(ent1)
self.assertIsNotNone(retr)
self.assertEqual(
2,
len(retr)) # both ent1 and ent2 are "collidable" in this quad tree
# this should cause a split (bucket_size)
ent3 = quadtree.QuadTreeEntity(
rect2.Rect2(15, 10, vector2.Vector2(700, 450)))
_tree.insert_and_think(ent3)
ent4 = quadtree.QuadTreeEntity(
rect2.Rect2(5, 5, vector2.Vector2(900, 900)))
_tree.insert_and_think(ent4)
# ent1 should collide with ent1 or ent2
# ent2 with ent1 or ent2, or ent3
# ent3 with ent2 or ent3
# ent4 with ent2 or ent4
retr = _tree.retrieve_collidables(ent1)
self.assertIsNotNone(retr)
self.assertEqual(2, len(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 25), None),
str(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 490), None),
str(retr))
retr = _tree.retrieve_collidables(ent2)
self.assertEqual(3, len(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 25), None),
str(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 490), None),
str(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 700), None),
str(retr))
retr = _tree.retrieve_collidables(ent3)
self.assertEqual(2, len(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 490), None),
str(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 700), None),
str(retr))
retr = _tree.retrieve_collidables(ent4)
self.assertEqual(2, len(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 900), None),
str(retr))
self.assertIsNotNone(
next((e for e in retr if e.aabb.mincorner.x == 490), None),
str(retr))
def setUp(self):
self.rect1 = rect2.Rect2(1, 1, vector2.Vector2(2, 2))