def test_branch(self): inputFile = 'sample_01.root' treeName = 'tree' inputPath = os.path.join(os.path.dirname(os.path.realpath(__file__)), inputFile) inputFile = ROOT.TFile.Open(inputPath) tree = inputFile.Get(treeName) events = Events(tree) jet_pt = events.jet_pt met_pt = events.met_pt self.assertIsInstance(jet_pt, Branch) self.assertIsInstance(met_pt, Branch) self.assertEqual(0, len(jet_pt)) self.assertEqual(1, len(met_pt)) self.assertEqual(0.0, met_pt[0]) tree.GetEntry(0) self.assertEqual(2, len(jet_pt)) self.assertEqual(1, len(met_pt)) self.assertEqual(124.55626678466797, jet_pt[0]) self.assertEqual(86.90544128417969, jet_pt[1]) self.assertAlmostEqual(43.783382415771484, met_pt[0]) tree.GetEntry(1) self.assertEqual(3, len(jet_pt)) self.assertEqual(1, len(met_pt)) self.assertEqual(112.48554992675781, jet_pt[0]) self.assertEqual(52.32780075073242, jet_pt[1]) self.assertEqual(48.861289978027344, jet_pt[2]) self.assertAlmostEqual(20.483951568603516, met_pt[0])
def test_nEvents(self): tree = MockTree(Entries=100) events = Events(tree) self.assertEqual(100, events.nEvents) # default the same as entries events = Events(tree, -1) self.assertEqual(100, events.nEvents) # the same as entries events = Events(tree, 50) self.assertEqual(50, events.nEvents) events = Events(tree, 120) self.assertEqual(100, events.nEvents) events = Events(tree, 100) self.assertEqual(100, events.nEvents)
def test_iter_getattr(self): tree = MockTree(Entries=5) events = Events(tree) it = iter(events) event = next(it) self.assertEqual(1111, event.branch1) tree.branch1 = 2222 self.assertEqual(2222, event.branch1)
def test_nEvents_start(self): tree = MockTree(Entries=100) events = Events(tree, maxEvents=-1, start=1) self.assertEqual(99, events.nEvents) events = Events(tree, maxEvents=10, start=1) self.assertEqual(10, events.nEvents) events = Events(tree, maxEvents=-1, start=99) self.assertEqual(1, events.nEvents) events = Events(tree, maxEvents=20, start=99) self.assertEqual(1, events.nEvents) events = Events(tree, maxEvents=-1, start=100) self.assertEqual(0, events.nEvents) events = Events(tree, maxEvents=-1, start=110) self.assertEqual(0, events.nEvents) events = Events(tree, maxEvents=10, start=110) self.assertEqual(0, events.nEvents) self.assertRaises(ValueError, Events, tree, maxEvents=-1, start=-10)
def use_Events(): inputFile = ROOT.TFile.Open(inputPath) tree = inputFile.Get(treeName) tree.SetBranchStatus("*", 0) tree.SetBranchStatus("jet_pt", 1) events = Events(tree) for event in events: jet_pt = event.jet_pt for i in range(len(jet_pt)): jet_pt[i]
def test_getitem_start(self): tree = MockTree(Entries=4) events = Events(tree, start=2) self.assertEqual(-1, events.iEvent) event = events[0] self.assertEqual(0, event.iEvent) self.assertEqual(2, tree.iEvent) event = events[1] self.assertEqual(1, event.iEvent) self.assertEqual(3, tree.iEvent) self.assertRaises(IndexError, events.__getitem__, 4) self.assertEqual(-1, events.iEvent)
def test_iter_iEvent_start(self): tree = MockTree(Entries=4) events = Events(tree, start=2) self.assertEqual(-1, events.iEvent) it = iter(events) event = next(it) self.assertEqual(0, event.iEvent) self.assertEqual(2, tree.iEvent) event = next(it) self.assertEqual(1, event.iEvent) self.assertEqual(3, tree.iEvent) self.assertRaises(StopIteration, next, it) self.assertEqual(-1, event.iEvent)
def test_iter_maxEvents(self): tree = MockTree(Entries=40) events = Events(tree, maxEvents=4) self.assertEqual(-1, events.iEvent) it = iter(events) event = next(it) self.assertEqual(0, event.iEvent) event = next(it) self.assertEqual(1, event.iEvent) event = next(it) self.assertEqual(2, event.iEvent) event = next(it) self.assertEqual(3, event.iEvent) self.assertRaises(StopIteration, next, it) self.assertEqual(-1, event.iEvent)
def test_vector(self): inputFile = 'sample_02.root' treeName = 'tree' inputPath = os.path.join(os.path.dirname(os.path.realpath(__file__)), inputFile) inputFile = ROOT.TFile.Open(inputPath) tree = inputFile.Get(treeName) events = Events(tree) trigger_path = events.trigger_path trigger_decision = events.trigger_decision self.assertGreater(tree.GetEntry(0), 0) self.assertEqual(449, len(trigger_path)) self.assertEqual('AlCa_EcalEtaEBonly', trigger_path[0]) self.assertEqual('DST_Physics', trigger_path[12]) self.assertEqual('HLT_SingleForJet25', trigger_path[438]) self.assertEqual(449, len(trigger_decision)) self.assertEqual(0, trigger_decision[0]) self.assertEqual(0, trigger_decision[13]) self.assertEqual(0, trigger_decision[438]) self.assertGreater(tree.GetEntry(1), 0) self.assertEqual(438, len(trigger_path)) self.assertEqual('AlCa_EcalEtaEBonly', trigger_path[0]) self.assertEqual('DST_Ele8_CaloIdL_CaloIsoVL_TrkIdVL_TrkIsoVL_HT250', trigger_path[12]) self.assertRaises(IndexError, trigger_path.__getitem__, 438) self.assertEqual(438, len(trigger_decision)) self.assertEqual(0, trigger_decision[0]) self.assertEqual(1, trigger_decision[13]) self.assertRaises(IndexError, trigger_decision.__getitem__, 438) # This sample file has only two entries. When the 3rd entry is # tried to be accessed, GetEntry(2) returns 0, but the vectors # won't be cleared. These have the previous contents. self.assertEqual(tree.GetEntry(2), 0) self.assertEqual(438, len(trigger_path)) self.assertEqual(438, len(trigger_decision))
for componentName in os.listdir(args.heppydir): if componentName in ('Chunks', 'failed'): continue componentPath = os.path.join(args.heppydir, componentName) if not os.path.isdir(componentPath): continue component = Component(componentPath) treeProducerSusyAlphaT = component.treeProducerSusyAlphaT inputPath = os.path.join(treeProducerSusyAlphaT.path, fileName) inputFile = ROOT.TFile.Open(inputPath) tree = inputFile.Get(treeName) boundaries = [10**(i * 0.1) for i in range(-10, 36)] counts = {} for event in Events(tree, args.nevents): met = event.met_pt met_bin = max([b for b in boundaries if b < met]) if not met_bin in counts: counts[met_bin] = 0 counts[met_bin] += 1 keys = counts.keys() keys.sort() for k in keys: print >> outFile, "{:>22s} {:12.6f} {:6d}".format( componentName, k, counts[k]) ##__________________________________________________________________||
def test_init(self): tree = MockTree() events = Events(tree) events = Events(tree, 100) self.assertIs(tree, events.tree)
def setUp(self): self.tree = MockTree() self.events = Events(self.tree)