def test_bbox() -> None:
"""Test the bounding box behaviour against a brute-force loop."""
rand = Random(1234) # Ensure reproducibility.
SIZE = 128.0
# Build a set of points and keys.
points = [(Vec(rand.uniform(-SIZE, SIZE), rand.uniform(-SIZE, SIZE),
rand.uniform(-SIZE, SIZE)),
Vec(rand.uniform(-SIZE, SIZE), rand.uniform(-SIZE, SIZE),
rand.uniform(-SIZE, SIZE)),
rand.getrandbits(64).to_bytes(8, 'little')) for _ in range(200)]
tree = RTree()
for a, b, data in points:
tree.insert(a, b, data)
# Pick a random bounding box.
bb_min, bb_max = Vec.bbox(
Vec(
rand.uniform(-SIZE, SIZE),
rand.uniform(-SIZE, SIZE),
rand.uniform(-SIZE, SIZE),
),
Vec(
rand.uniform(-SIZE, SIZE),
rand.uniform(-SIZE, SIZE),
rand.uniform(-SIZE, SIZE),
))
expected = [
data for a, b, data in points
if Vec.bbox_intersect(*Vec.bbox(a, b), bb_min, bb_max)
]
found = set(tree.find_bbox(bb_min, bb_max))
# Order is irrelevant, but duplicates must all match.
assert sorted(expected) == sorted(found)
def open(self):
""" open the file, get the tree
"""
self.tfile = TFile(self.fname)
self.ttree = self.tfile.Get(self.tname)
self.rtree = RTree(self.ttree)
#print 'TFile Reader ',self.tfile
return
def test_duplicate_insertion() -> None:
"""Test inserting values with the same bbox."""
tree = RTree()
tree.insert(Vec(10, 20, 30), Vec(40, 50, 60), 'value1')
tree.insert(Vec(50, 20, 30), Vec(80, 65, 60), 'another')
tree.insert(Vec(10, 20, 30), Vec(40, 50, 60), 'value2')
assert len(tree) == 3
# Iterating produces both values.
assert {(*mins, *maxes, val)
for mins, maxes, val in tree} == {
(10, 20, 30, 40, 50, 60, 'value1'),
(50, 20, 30, 80, 65, 60, 'another'),
(10, 20, 30, 40, 50, 60, 'value2'),
}
# Check all vecs are unique.
vecs = [vec for mins, maxes, _ in tree for vec in [mins, maxes]]
assert len(set(map(id, vecs))) == 6, list(tree)
# Check we can delete one.
tree.remove(Vec(10, 20, 30), Vec(40, 50, 60), 'value1')
assert {(*mins, *maxes, val)
for mins, maxes, val in tree} == {
(50, 20, 30, 80, 65, 60, 'another'),
(10, 20, 30, 40, 50, 60, 'value2'),
}
class TreeReader(IReader):
""" Read data from a TTree ROOT class inside a ROOT TFile.
data is each entry of the TTree
"""
def __init__(self,fname,tname):
""" Read data from a TTree ROOT class inside a ROOT TFile.
data is each entry of the TTree
"""
IReader.__init__(self,fname)
self.tname = tname
self.tfile = None
self.ttree = None
self.rtree = None
return
def eof(self):
""" end of file
"""
return (self.rtree.ientry >= self.rtree.nentries)
def open(self):
""" open the file, get the tree
"""
self.tfile = TFile(self.fname)
self.ttree = self.tfile.Get(self.tname)
self.rtree = RTree(self.ttree)
#print 'TFile Reader ',self.tfile
return
def read(self):
""" read and returns the next entry of the tree as a dictionary
"""
data = self.rtree.getentry()
return data
def close(self):
""" close the file
"""
self.tfile.Close()
return
