from __future__ import print_function
import copy
__all__ = ['AABB', 'AABBTree']
__author__ = 'Kenneth (Kip) Hart'
[docs]class AABB(object):
"""Axis aligned bounding box (AABB)
The AABB is a d-dimensional box.
Args:
limits (iterable, optional): The limits of the box. These should be
specified in the following manner::
limits = [(xmin, xmax), (ymin, ymax), (zmin, zmax), ...]
The default value is None.
"""
def __init__(self, limits=None):
if limits is not None:
for lims in limits:
assert len(lims) == 2
assert lims[0] <= lims[1]
self.limits = limits
def __str__(self):
return str(self.limits)
def __repr__(self):
return 'AABB(' + repr(self.limits) + ')'
def __iter__(self):
self._i = 0
return self
def __next__(self):
if self._i < len(self):
val = self.limits[self._i]
self._i += 1
return val
raise StopIteration
def next(self): # pragma: no cover
"""___next__ for Python 2"""
return self.__next__()
def __getitem__(self, key):
return self.limits[key]
def __len__(self):
return len(self.limits)
def __eq__(self, aabb):
if not isinstance(aabb, AABB):
return False
if (self.limits is None) and (aabb.limits is None):
return True
if self.limits is None:
return False
if aabb.limits is None:
return False
if len(self) != len(aabb):
return False
for i in range(len(self)):
lims1 = self[i]
lims2 = aabb[i]
if (lims1[0] != lims2[0]) or (lims1[1] != lims2[1]):
return False
return True
def __ne__(self, aabb):
return not self.__eq__(aabb)
[docs] @classmethod
def merge(cls, aabb1, aabb2):
"""Merge AABB
Find the AABB of the union of AABBs.
Args:
aabb1 (AABB): An AABB
aabb2 (AABB): An AABB
Returns:
AABB: An AABB that contains both of the inputs
"""
if (aabb1.limits is None) and (aabb2.limits is None):
return cls(None)
elif aabb1.limits is None:
return cls(aabb2.limits)
elif aabb2.limits is None:
return cls(aabb1.limits)
else:
merged_limits = []
for lims1, lims2 in zip(aabb1, aabb2):
lb = min(lims1[0], lims2[0])
ub = max(lims1[1], lims2[1])
merged_limits.append((lb, ub))
return cls(merged_limits)
@property
def perimeter(self):
r"""Perimeter of AABB
The perimeter :math:`p_n` of an AABB with side lengths
:math:`l_1, ..., l_n` is:
.. math::
\begin{align}
p_1 &= 0 \\
p_2 &= 2 (l_1 + l_2) \\
p_3 &= 2 (l_1 l_2 + l_2 l_3 + l_1 l_3) \\
p_n &= 2 \sum_{i=1}^n \prod_{j=1\neq i}^n l_j
\end{align}
"""
if len(self) == 1:
return 0
p = 0
side_lens = [ub - lb for lb, ub in self]
for i in range(len(side_lens)):
p_edge = 1
for j in range(len(side_lens)):
if j != i:
p_edge *= side_lens[j]
p += p_edge
return 2 * p
[docs] def overlaps(self, aabb):
"""Determine if two AABBs overlap
Args:
aabb (AABB): The AABB to check for overlap
Returns:
bool: Flag set to true if the two AABBs overlap
"""
for lims1, lims2 in zip(self, aabb):
min1, max1 = lims1
min2, max2 = lims2
overlaps = (max1 >= min2) and (min1 <= max2)
if not overlaps:
return False
return True
[docs]class AABBTree(object):
"""Python Implementation of the AABB Tree
This is a pure Python implementation of the static d-dimensional AABB tree.
It is heavily based on
`Introductory Guide to AABB Tree Collision Detection`_
from *Azure From The Trenches*.
Args:
aabb (AABB): An AABB
value: The value associated with the AABB
left (AABBTree, optional): The left branch of the tree
right (AABBTree, optional): The right branch of the tree
.. _`Introductory Guide to AABB Tree Collision Detection` : https://www.azurefromthetrenches.com/introductory-guide-to-aabb-tree-collision-detection/
""" # NOQA: E501
def __init__(self, aabb=AABB(), value=None, left=None, right=None):
self.aabb = aabb
self.value = value
self.left = left
self.right = right
def __repr__(self):
inp_strs = []
if self.aabb != AABB():
inp_strs.append('aabb=' + repr(self.aabb))
if self.value is not None:
inp_strs.append('value=' + repr(self.value))
if self.left is not None:
inp_strs.append('left=' + repr(self.left))
if self.right is not None:
inp_strs.append('right=' + repr(self.right))
return 'AABBTree(' + ', '.join(inp_strs) + ')'
def __str__(self, n=0):
strs = []
pre = n * ' '
aabb_str = pre + 'AABB: '
if self.aabb == AABB():
aabb_str += 'None'
else:
aabb_str += str(self.aabb)
value_str = pre + 'Value: ' + str(self.value)
left_str = pre + 'Left:'
if self.left is None:
left_str += ' None'
else:
left_str += '\n' + self.left.__str__(n + 1)
right_str = pre + 'Right:'
if self.right is None:
right_str += ' None'
else:
right_str += '\n' + self.right.__str__(n + 1)
return '\n'.join([aabb_str, value_str, left_str, right_str])
def __eq__(self, aabbtree):
if not isinstance(aabbtree, AABBTree):
return False
if self.aabb != aabbtree.aabb:
return False
if self.is_leaf != aabbtree.is_leaf:
return False
return (self.left == aabbtree.left) and (self.right == aabbtree.right)
def __ne__(self, aabbtree):
return not self.__eq__(aabbtree)
@property
def is_leaf(self):
"""bool: returns True if is leaf node"""
return (self.left is None) and (self.right is None)
[docs] def add(self, aabb, value=None):
"""Add node to tree
This function inserts a node into the AABB tree.
Args:
aabb (AABB): The AABB to add.
value: The value associated with the AABB. Defaults to None.
"""
if self.aabb == AABB():
self.aabb = aabb
self.value = value
elif self.is_leaf:
self.left = copy.deepcopy(self)
self.right = AABBTree(aabb, value)
self.aabb = AABB.merge(self.aabb, aabb)
self.value = None
else:
tree_p = self.aabb.perimeter
tree_merge_p = AABB.merge(self.aabb, aabb).perimeter
new_parent_cost = 2 * tree_merge_p
min_pushdown_cost = 2 * (tree_merge_p - tree_p)
left_merge_p = AABB.merge(self.left.aabb, aabb).perimeter
cost_left = left_merge_p + min_pushdown_cost
if not self.left.is_leaf:
cost_left -= self.left.aabb.perimeter
right_merge_p = AABB.merge(self.right.aabb, aabb).perimeter
cost_right = right_merge_p + min_pushdown_cost
if not self.right.is_leaf:
cost_right -= self.right.aabb.perimeter
if new_parent_cost < min(cost_left, cost_right):
self.left = copy.deepcopy(self)
self.right = AABBTree(aabb, value)
self.value = None
elif cost_left < cost_right:
self.left.add(aabb, value)
else:
self.right.add(aabb, value)
self.aabb = AABB.merge(self.left.aabb, self.right.aabb)
[docs] def does_overlap(self, aabb):
"""Check for overlap
This function checks if the limits overlap any leaf nodes in the tree.
It returns true if there is an overlap.
Args:
aabb (AABB): The AABB to check.
Returns:
bool: True if overlaps with a leaf node of tree.
"""
if self.is_leaf:
return self.aabb.overlaps(aabb)
else:
left_aabb_over = self.left.aabb.overlaps(aabb)
right_aabb_over = self.right.aabb.overlaps(aabb)
if left_aabb_over and self.left.does_overlap(aabb):
return True
if right_aabb_over and self.right.does_overlap(aabb):
return True
return False
[docs] def overlap_values(self, aabb):
"""Get values of overlapping AABBs
This function gets the value field of each overlapping AABB.
Args:
aabb (AABB): The AABB to check.
Returns:
list: Value fields of each node that overlaps.
"""
values = []
if self.is_leaf and self.does_overlap(aabb):
values.append(self.value)
elif self.is_leaf:
pass
else:
if self.left.aabb.overlaps(aabb):
values.extend(self.left.overlap_values(aabb))
if self.right.aabb.overlaps(aabb):
values.extend(self.right.overlap_values(aabb))
return values