"""
Module for binary space partitioning, to facilitate optimal runtime complexity for n-point
problems.
"""
import abc
from rtree import index as rtree
[docs]
class BoxTreeInterface(abc.ABC):
"""
Tree for fast lookup of queries of axis aligned rhomboids.
"""
[docs]
@abc.abstractmethod
def query(self, boxes, unique=False):
"""
Should return a generator that yields lists of intersecting IDs per query box if
``unique=False``. If ``unique=True``, yield a flat list of unique intersecting box
IDs for all queried boxes.
"""
pass
@abc.abstractmethod
def __len__(self):
pass
[docs]
class _BoxRTree(BoxTreeInterface):
"""
Tree for fast lookup of queries of axis aligned rhomboids using the Rtree package.
"""
def __init__(self, boxes):
self._rtree = rtree.Index(properties=rtree.Property(dimension=3))
for id, box in enumerate(boxes):
self._rtree.insert(id, box)
def __len__(self):
return len(self._rtree)
def query(self, boxes, unique=False):
"""
Given an iterable of ``(min_x, min_y, min_z, max_x, max_y, max_z)`` box tuples,
find all the boxes that intersect with them.
:param boxes: Boxes to look for intersections with.
:type boxes: Iterable[Tuple[float, float, float, float, float, float]]
:param unique: If ``True``, return a flat generator of unique results. If ``False``
(default), per box in ``boxes``, return a list of intersecting boxes.
:returns: See ``unique``.
:rtype: Union[Iterator[List[Tuple[float, float, float, float, float, float]]],
Iterator[Tuple[float, float, float, float, float, float]]]
"""
all_ = (list(self._rtree.intersection(box, objects=False)) for box in boxes)
if unique:
seen = set()
# Double for loop over results, skipping those that have been seen before.
yield from (
seen.add(elem) or elem for arr in all_ for elem in arr if elem not in seen
)
else:
yield from all_
# Cheapo provider pattern.
[docs]
class BoxTree(_BoxRTree):
"""
Tree for fast lookup of repeat queries of axis aligned rhomboids.
"""
pass
__all__ = ["BoxTreeInterface", "BoxTree"]