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