def test_chain_draw_hist_init_first(): if sys.version_info[0] >= 3: raise SkipTest("Python 3 support not implemented") hist = Hist(100, 0, 1) chain = TreeChain('tree', FILE_PATHS) chain.draw('a_x', hist=hist) assert_equal(hist.Integral() > 0, True)
def test_chain_draw(self): chain = TreeChain('tree', self.file_paths) hist = Hist(100, 0, 1) chain.draw('a_x', hist=hist) assert_equals(hist.Integral() > 0, True) # check that Draw can be repeated hist2 = Hist(100, 0, 1) chain.draw('a_x', hist=hist2) assert_equals(hist.Integral(), hist2.Integral())
def test_chain_draw(): chain = TreeChain("tree", FILE_PATHS) hist = Hist(100, 0, 1) chain.draw("a_x", hist=hist) assert_equal(hist.Integral() > 0, True) # check that Draw can be repeated hist2 = Hist(100, 0, 1) chain.draw("a_x", hist=hist2) assert_equal(hist.Integral(), hist2.Integral())
def get_histograms_from_trees( trees = [], branch = 'var', weightBranch = 'EventWeight', selection = '1', files = {}, verbose = False, nBins = 40, xMin = 0, xMax = 100, ignoreUnderflow = True, ): histograms = {} nHistograms = 0 # Setup selection and weight string for ttree draw weightAndSelection = '( %s ) * ( %s )' % ( weightBranch, selection ) for sample, input_file in files.iteritems(): histograms[sample] = {} for tree in trees: tempTree = tree if 'data' in sample and ( 'Up' in tempTree or 'Down' in tempTree ) : tempTree = tempTree.replace('_'+tempTree.split('_')[-1],'') chain = None; if isinstance( input_file, list ): for f in input_file: chain.Add(f) else: chain = TreeChain(tempTree, [input_file]); weightAndSelection = '( %s ) * ( %s )' % ( weightBranch, selection ) root_histogram = Hist( nBins, xMin, xMax, type='D') chain.Draw(branch, weightAndSelection, hist = root_histogram) if not is_valid_histogram( root_histogram, tree, input_file): return # When a tree is filled with a dummy variable, it will end up in the underflow, so ignore it if ignoreUnderflow: root_histogram.SetBinContent(0, 0) root_histogram.SetBinError(0,0) gcd() nHistograms += 1 histograms[sample][tree] = root_histogram.Clone() return histograms
def test_chain_draw(): if sys.version_info[0] >= 3: raise SkipTest("Python 3 support not implemented") chain = TreeChain('tree', FILE_PATHS) hist = Hist(100, 0, 1) chain.draw('a_x', hist=hist) assert_equal(hist.Integral() > 0, True) # check that Draw can be repeated hist2 = Hist(100, 0, 1) chain.draw('a_x', hist=hist2) assert_equal(hist.Integral(), hist2.Integral())
def load_tree(infiles): from rootpy.tree import TreeChain from rootpy.ROOT import gROOT gROOT.SetBatch(True) print('[INFO] Opening files: {0}'.format(infiles)) tree = TreeChain('ntupler/tree', infiles) tree.define_collection(name='hits', prefix='vh_', size='vh_size') tree.define_collection(name='simhits', prefix='vc_', size='vc_size') tree.define_collection(name='tracks', prefix='vt_', size='vt_size') tree.define_collection(name='particles', prefix='vp_', size='vp_size') return tree
def test_chain_iter(): if sys.version_info[0] >= 3: raise SkipTest("Python 3 support not implemented") chain = TreeChain('tree', FILE_PATHS) assert_equal(len(chain), 3) # 3 files entries = 0 for entry in chain: entries += 1 assert_equal(entries, 300) entries = 0 for entry in chain: entries += 1 assert_equal(entries, 300) assert_equal(chain.GetEntries(), 300) assert_equal(chain.GetEntriesFast(), 300)
def loop(self): tree = TreeChain(self.treeName, self.inFile) self.Nev = 0 for f in self.inFile: tmpf = root_open(f) for path, dirs, objects in tmpf.walk(): if "Events" in objects: evHist = tmpf.Get(path + "/Events") self.Nev += evHist.GetBinContent(1) break tmpf.close() counter = 0 for e in tree: self.Fill(0, "zp_pt", e.mu1_pt) #just profile stuff self.Profile(0, "test_profile", e.mu1_pt, e.mu1_pt / 14) if (e.mu1_pt > 30): #make cuts self.Fill(1, "zp_pt", e.mu1_pt / 15) self.writeFile()
def __init__(self, files, events=-1, load_emu_trees=False, load_reco_trees=True): from cmsl1t.utils.root_glob import glob input_files = [] for f in files: if '*' in f: input_files.extend(glob(f)) else: input_files.append(f) # this is not efficient self._trees = [] self._names = [] load_ROOT_library('L1TAnalysisDataformats.so') allTrees = get_trees(load_emu_trees, load_reco_trees) for name, path in allTrees.iteritems(): try: chain = TreeChain(path, input_files, cache=True, events=events) except RuntimeError: logger.warn("Cannot find tree: {0} in input file".format(path)) continue self._names.append(name) self._trees.append(chain)
def purge_bad_files(infiles): good_files = [] for infile in infiles: try: _ = TreeChain('ntupler/tree', infile) good_files.append(infile) except: pass return good_files
def __init__(self, files=None, type='MC', maxevents=-1): if type.upper() not in ['MC', 'DATA']: raise ValueError("Argument `type` need to be MC/DATA") self.Type = type self.MaxEvents = maxevents if not files: raise ValueError("Argument `files` need to be non-empty") if isinstance(files, str): files = [ files, ] self.Chain = TreeChain('ffNtuplizer/ffNtuple', files) ## register collections ### self.Chain.define_collection('pvs', prefix='pv_', size='pv_n') self.Chain.define_collection('muons', prefix='muon_', size='muon_n') self.Chain.define_collection('dsamuons', prefix='dsamuon_', size='dsamuon_n') self.Chain.define_collection('ak4jets', prefix='akjet_ak4PFJetsCHS_', size='akjet_ak4PFJetsCHS_n') self.Chain.define_collection('leptonjets', prefix='pfjet_', size='pfjet_n', mix=LeptonJetMix) self.Chain.define_collection('trigobjs', prefix='trigobj_', size='trigobj_n') self.Chain.define_object('hlt', prefix='HLT_') self.Chain.define_object('metfilters', prefix='metfilters_') self.Chain.define_object('cosmicveto', prefix='cosmicveto_') self.Histos = {} self.LookupWeight = root_open( os.path.join( os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/puWeights_10x_56ifb.root')).Get( 'puWeights') self.Scale = 1.
def _load_trees(self): for treeName in self._treeNames: try: self._trees[treeName] = TreeChain( treeName, self.input_files, cache=True, events=self.nevents, ) except RuntimeError: logger.warn( "Cannot find tree: {0} in input file".format(treeName)) continue
def copy_in_trigger_signal(in_files_name, out_name, tree_name, prefix, cdc_events, cth_events, rand_t=None): # Convert input lists to sets first set_cdc_events = set(cdc_events) set_cth_events = set(cth_events) # Define the chain of input trees in_chain = TreeChain(name=tree_name, files=in_files_name) # First create a new file to save the new tree in: out_file = root_open(out_name, "r+") # Add the time shift if we want it in the tree ExtraBranches = Tagged if rand_t is not None: ExtraBranches += Smeared # Get the new tree with its extra branches out_tree = Tree(tree_name, model=ExtraBranches.prefix(prefix)) # This creates all the same branches in the new tree but # their addresses point to the same memory used by the original tree. out_tree.create_branches(in_chain._buffer) out_tree.update_buffer(in_chain._buffer) # Now loop over the original tree(s) and fill the new tree for entry in in_chain: # Add in the new values this_event_number = entry[prefix + "EventNumber"].value out_tree.__setattr__(prefix + "GoodTrack", this_event_number in set_cdc_events) out_tree.__setattr__(prefix + "GoodTrig", this_event_number in set_cth_events) if rand_t is not None: try: out_tree.__setattr__(prefix + "SmearTime", rand_t[this_event_number]) except: for key, item in entry.iteritems(): print key, item # Fill, noting that most of the buffer is shared between the chain # and the output tree out_tree.Fill() # Close it up out_tree.Write() out_file.Close()
def test_chain_draw(): chain = TreeChain('tree', FILE_PATHS) hist = Hist(100, 0, 1) chain.draw('a_x', hist=hist) assert_equal(hist.Integral() > 0, True) # check that Draw can be repeated hist2 = Hist(100, 0, 1) chain.draw('a_x', hist=hist2) assert_equal(hist.Integral(), hist2.Integral())
def test_chain_draw(): if sys.version_info[0] >= 3: raise SkipTest("Python 3 support not implemented") chain = TreeChain('tree', FILE_PATHS) hist = Hist(100, 0, 1) chain.draw('a_x', hist=hist) assert_equal(hist.Integral() > 0, True) assert_equal(hist.GetEntries(), 300) # check that Draw can be repeated hist2 = Hist(100, 0, 1) chain.draw('a_x', hist=hist2) assert_equal(hist.Integral(), hist2.Integral()) # draw into a graph graph = chain.draw("a_x:a_y") assert_true(isinstance(graph, _GraphBase)) assert_equal(len(graph), chain.GetEntries()) assert_equal(len(graph), 300)
def __init__(self, files, events=-1): from cmsl1t.utils.root_glob import glob input_files = [] for f in files: if '*' in f: input_files.extend(glob(f)) else: input_files.append(f) # this is not efficient self._trees = [] self._names = [] for name, path in ALL_TREE.iteritems(): try: chain = TreeChain(path, input_files, cache=True, events=events) except RuntimeError: logger.warn("Cannot find tree: {0} in input file".format(path)) continue self._names.append(name) self._trees.append(chain)
def mergeChannel(channel, fileList): ''' Merge one channel for all files, return merged tree. Should be run with a file open. ''' if isinstance(fileList, str): fileList = [fileList] chain = TreeChain('{}/ntuple'.format(channel), fileList) out = Tree('ntuple'.format(channel)) out.set_buffer(chain._buffer, create_branches=True) found = set() for ev in chain: evID = (ev.run, ev.lumi, ev.evt) if evID in found: continue out.fill() found.add(evID) return out
def work(self): D4PD_model = RecoTauBlock + RecoJetBlock + EventVariables + TrueTauBlock + PartonBlock # initialize the TreeChain of all input files (each containing one tree named self.metadata.treename) tree = TreeChain(self.metadata.treename, files=self.files, events=self.events, cache=True, cache_size=10000000, learn_entries=30) # create output tree self.output.cd() D4PD = Tree(name=self.metadata.name, model=D4PD_model) copied_variables = ['actualIntPerXing', 'averageIntPerXing'] copied_variables += mc_triggers D4PD.set_buffer(tree.buffer, variables=copied_variables, create_branches=True, visible=False) tree.always_read(copied_variables) # set the event filters # passthrough for MC for trigger acceptance studies self.event_filters = EventFilterList([ GRLFilter(self.grl, passthrough = self.metadata.datatype != datasets.DATA), PriVertex(), LArError(), JetCleaningLoose(passthrough = self.metadata.datatype != datasets.DATA), #JetCleaningMedium(passthrough = self.metadata.datatype != datasets.DATA), LArHole(), #JetCrackVeto(), ElectronVeto(), MuonVeto(), TauElectronVeto(), TauMuonVeto(), TauAuthorTrack(), ]) self.event_filters.insert(1, MCTriggers()) tree.filters += self.event_filters cutflow = Cutflow() # define tree collections tree.define_collection(name="taus", prefix="tau_", size="tau_n", mix=TauFourMomentum) # jet_eta etc is AntiKt4LCTopo in tau-perf D3PDs tree.define_collection(name="jets", prefix="jet_", size="jet_n", mix=FourMomentum) tree.define_collection(name="truetaus", prefix="trueTau_", size="trueTau_n", mix=MCTauFourMomentum) tree.define_collection(name="mc", prefix="mc_", size="mc_n", mix=MCParticle) tree.define_collection(name="muons", prefix="mu_staco_", size="mu_staco_n") tree.define_collection(name="electrons", prefix="el_", size="el_n") tree.define_collection(name="vertices", prefix="vxp_", size="vxp_n") # define tree objects D4PD.define_object(name='tau1', prefix='tau1_') D4PD.define_object(name='tau2', prefix='tau2_') D4PD.define_object(name='jet1', prefix='jet1_') D4PD.define_object(name='jet2', prefix='jet2_') """ tree.define_association(origin='taus', target='truetaus', prefix='trueTauAssoc_', link='index') tree.define_association(origin='truetaus', target='taus', prefix='tauAssoc_', link='index') """ # entering the main event loop... for event in tree: D4PD.reset() cutflow.reset() # tau selection event.taus.select(lambda tau: tau.pt > 15*GeV) # Jet selection event.jets.select(lambda jet: jet.fourvect.P() > 25*GeV and abs(jet.emscale_eta) < 4.5) """ Get VBF jets """ # get partons (already sorted by eta in hepmc) parton1, parton2 = hepmc.get_VBF_partons(event) PartonBlock.set(D4PD, parton1, parton2) D4PD.dR_quarks = parton1.fourvect.DeltaR(parton2.fourvect) """ Get true taus """ event.truetaus.select(lambda tau: tau.vis_Et > 10 * GeV and abs(tau.vis_eta) < 2.5) if len(event.truetaus) > 2: print "ERROR: too many true taus: %i" % len(event.truetaus) D4PD.error = 1 D4PD.Fill() continue elif len(event.truetaus) < 2: print "ERROR: too few true taus: %i" % len(event.truetaus) D4PD.error = 2 D4PD.Fill() continue """ fourvects = [] colors = [] radii = [] for thing in event.jets: fourvects.append(thing.fourvect) colors.append('blue') radii.append(.4) for parton in (parton1, parton2): fourvects.append(parton.fourvect) colors.append('green') radii.append(.1) for thing in event.taus: fourvects.append(thing.fourvect) colors.append('red') radii.append(.2) for tau in event.truetaus: fourvects.append(tau.fourvect) colors.append('purple') radii.append(.1) eventview.draw(event, fourvects, colors=colors, radii=radii) """ TrueTauBlock.set(D4PD, 1, event.truetaus[0]) TrueTauBlock.set(D4PD, 2, event.truetaus[1]) D4PD.dR_truetaus = event.truetaus[0].fourvect.DeltaR(event.truetaus[1].fourvect) D4PD.dR_quark_tau = min([ parton1.fourvect.DeltaR(event.truetaus[0].fourvect), parton2.fourvect.DeltaR(event.truetaus[0].fourvect), parton1.fourvect.DeltaR(event.truetaus[1].fourvect), parton2.fourvect.DeltaR(event.truetaus[1].fourvect), ]) taus = [] if event.taus: for truetau in event.truetaus: closest_tau = min(event.taus, key=lambda tau: utils.dR(tau.seedCalo_eta, tau.seedCalo_phi, truetau.vis_eta, truetau.vis_phi)) if utils.dR(closest_tau.seedCalo_eta, closest_tau.seedCalo_phi, truetau.vis_eta, truetau.vis_phi) < 0.2: if closest_tau in taus: # collision D4PD.error = 3 break taus.append(closest_tau) if len(taus) < 2: # collision D4PD.Fill() continue """ # Overlap removal between taus and jets event.jets.select(lambda jet: not any([tau for tau in taus if utils.dR(jet.emscale_eta, jet.emscale_phi, tau.seedCalo_eta, tau.seedCalo_phi) < .2])) """ jets = [] if event.jets: for quark in (parton1, parton2): closest_jet = min(event.jets, key=lambda jet: utils.dR(jet.eta, jet.phi, quark.eta, quark.phi)) if utils.dR(closest_jet.eta, closest_jet.phi, quark.eta, quark.phi) < 0.4: if closest_jet in jets: # collision D4PD.error = 4 break jets.append(closest_jet) if len(jets) < 2: # collision D4PD.Fill() continue """ Jet variables """ RecoJetBlock.set(D4PD, jets[0], jets[1]) """ Reco tau variables This must come after the RecoJetBlock is filled since that sets the jet_beta for boosting the taus """ RecoTauBlock.set(event, D4PD, taus[0], taus[1]) D4PD.true_Mvis_tau1_tau2 = (D4PD.trueTau1_fourvect_vis + D4PD.trueTau2_fourvect_vis).M() """ MET """ METx = event.MET_RefFinal_etx METy = event.MET_RefFinal_ety MET_vect = Vector2(METx, METy) MET_3vect = Vector3(METx, METy, 0.) D4PD.MET = event.MET_RefFinal_et D4PD.MET_phi = event.MET_RefFinal_phi # HT TODO: Fix sumET = event.MET_RefFinal_sumet D4PD.HT = sumET D4PD.numVertices = len([vtx for vtx in event.vertices if (vtx.type == 1 and vtx.nTracks >= 4) or (vtx.type == 3 and vtx.nTracks >= 2)]) # fill output ntuple # use reset=True to reset all variables to their defaults after the fill # to avoid any values from this event carrying over into the next D4PD.cutflow = cutflow.int() D4PD.Fill() self.output.cd() D4PD.FlushBaskets() D4PD.Write()
def __init__(self, fns, outn, outd, model): print "Initializing Container!" #self.tin = r.TChain('EcalVeto') #for fn in fns: # self.tin.Add(fn) #self.tfile = root_open(fn,'r+') #with root_open(fn,'r+') as f: #self.tin = self.tfile.EcalVeto #self.tin.Print() self.tin = TreeChain('EcalVeto', fns) self.tin.create_branches({'discValue_gabrielle': 'F'}) self.model = model self.outdir = outd self.outname = outn self.outfile = root_open(outn, 'RECREATE') self.tout = Tree('EcalVeto') self.tout.set_buffer(self.tin._buffer, create_branches=True) self.events = [] #print self.tin.GetEntries() for event in self.tin: #if len(self.events)>10: # continue evt = [] ################################### Features ####################################### evt.append(event.nReadoutHits) evt.append(event.summedDet) evt.append(event.summedTightIso) evt.append(event.maxCellDep) evt.append(event.showerRMS) evt.append(event.xStd) evt.append(event.yStd) evt.append(event.avgLayerHit) evt.append(event.deepestLayerHit) evt.append(event.stdLayerHit) #new features evt.append(event.ele68ContEnergy) evt.append(event.ele68x2ContEnergy) evt.append(event.ele68x3ContEnergy) evt.append(event.ele68x4ContEnergy) evt.append(event.ele68x5ContEnergy) evt.append(event.photon68ContEnergy) evt.append(event.photon68x2ContEnergy) evt.append(event.photon68x3ContEnergy) evt.append(event.photon68x4ContEnergy) evt.append(event.photon68x5ContEnergy) evt.append(event.outside68ContEnergy) evt.append(event.outside68x2ContEnergy) evt.append(event.outside68x3ContEnergy) evt.append(event.outside68x4ContEnergy) evt.append(event.outside68x5ContEnergy) evt.append(event.outside68ContNHits) evt.append(event.outside68x2ContNHits) evt.append(event.outside68x3ContNHits) evt.append(event.outside68x4ContNHits) evt.append(event.outside68x5ContNHits) evt.append(event.outside68ContXstd) evt.append(event.outside68x2ContXstd) evt.append(event.outside68x3ContXstd) evt.append(event.outside68x4ContXstd) evt.append(event.outside68x5ContXstd) evt.append(event.outside68ContYstd) evt.append(event.outside68x2ContYstd) evt.append(event.outside68x3ContYstd) evt.append(event.outside68x4ContYstd) evt.append(event.outside68x5ContYstd) evt.append(event.ecalBackEnergy) evtarray = np.array([evt]) pred = float(model.predict(xgb.DMatrix(evtarray))[0]) #print pred event.discValue_gabrielle = pred self.tout.Fill() ###################################################################################### self.events.append(evt) if (len(self.events) % 10000 == 0 and len(self.events) > 0): print 'The shape of events = ', np.shape(self.events) self.outfile.cd() self.tout.Write() self.outfile.Close() #self.tfile.Close() print 'cp %s %s' % (self.outname, self.outdir) os.system('cp %s %s' % (self.outname, self.outdir))
tree = Tree("test") tree.create_branches(branches) for i in range(10000): tree.x = gauss(.5, 1.) tree.y = gauss(.3, 2.) tree.z = gauss(13., 42.) tree.i = i tree.fill() tree.write() # Make a histogram of the second tree hist2 = Hist(100, -10, 10, name='hist2') tree.Draw('x', 'y > 1', hist=hist2) hist2.SetDirectory(0) # memory resident print("The second tree has {0:f} entries where y > 1".format(hist2.Integral())) f.close() combined_hist = hist1 + hist2 print("Building TreeChain") chain = TreeChain('test', ['chaintest2.root', 'chaintest1.root']) # Make the equivalent of the combined_hist combined_hist_chain = Hist(100, -10, 10, name='combined') chain.Draw('x', 'y > 1', hist=combined_hist_chain) residual = combined_hist_chain - combined_hist print("The combined histogram (separately) minus " "the combined from the chain has {0:f} entries".format( residual.Integral()))
def __init__(self, files, events=-1): # this is not efficient self._trees = [ TreeChain(name, files, cache=True, events=events) for name in TREE_NAMES ]
def work(self): # trigger config tool to read trigger info in the ntuples trigger_config = get_trigger_config() OutputModel = (RecoTauBlock + EventVariables + SkimExtraModel + TrueTauBlock) onfilechange = [] # update the trigger config maps on every file change onfilechange.append((update_trigger_config, (trigger_config,))) cutflow = Hist(2, 0, 2, name='cutflow', type='D') # initialize the TreeChain of all input files (each containing one tree named self.metadata.treename) chain = TreeChain(self.metadata.treename, files=self.files, events=self.events, cache=True, cache_size=10000000, learn_entries=30, onfilechange=onfilechange) # create output tree self.output.cd() tree = Tree(name='higgstautauhh', model=OutputModel) copied_variables = ['actualIntPerXing', 'averageIntPerXing', 'RunNumber', 'EventNumber', 'lbn'] tree.set_buffer( chain.buffer, branches=copied_variables, create_branches=True, visible=False) chain.always_read(copied_variables) # set the event filters event_filters = EventFilterList([ #Triggers( # datatype=self.metadata.datatype, # year=YEAR, # skim=False), PriVertex(), LArError(), LArHole(datatype=self.metadata.datatype), JetCleaning( datatype=self.metadata.datatype, year=YEAR), TauAuthor(1), TauHasTrack(1), TauPT(1, thresh=25 * GeV), TauEta(1), TauCrack(1), TauLArHole(1), #TauTriggerMatch( # config=trigger_config, # year=YEAR, # datatype=self.metadata.datatype, # skim=False, # tree=tree, # min_taus=1), ]) self.filters['event'] = event_filters chain.filters += event_filters # define tree collections chain.define_collection(name="taus", prefix="tau_", size="tau_n", mix=TauFourMomentum) chain.define_collection(name="taus_EF", prefix="trig_EF_tau_", size="trig_EF_tau_n", mix=TauFourMomentum) # jet_* etc. is AntiKt4LCTopo_* in tau-perf D3PDs chain.define_collection(name="jets", prefix="jet_", size="jet_n", mix=FourMomentum) chain.define_collection(name="truetaus", prefix="trueTau_", size="trueTau_n", mix=MCTauFourMomentum) chain.define_collection(name="mc", prefix="mc_", size="mc_n", mix=MCParticle) chain.define_collection(name="muons", prefix="mu_staco_", size="mu_staco_n") chain.define_collection(name="electrons", prefix="el_", size="el_n") chain.define_collection(name="vertices", prefix="vxp_", size="vxp_n") from externaltools import PileupReweighting from ROOT import Root # Initialize the pileup reweighting tool pileup_tool = Root.TPileupReweighting() if YEAR == 2011: pileup_tool.AddConfigFile(PileupReweighting.get_resource('mc11b_defaults.prw.root')) pileup_tool.AddLumiCalcFile('lumi/2011/hadhad/ilumicalc_histograms_None_178044-191933.root') elif YEAR == 2012: pileup_tool.AddConfigFile(PileupReweighting.get_resource('mc12a_defaults.prw.root')) pileup_tool.SetDataScaleFactors(1./1.11) pileup_tool.AddLumiCalcFile('lumi/2012/hadhad/ilumicalc_histograms_None_200841-205113.root') else: raise ValueError('No pileup reweighting defined for year %d' % YEAR) # discard unrepresented data (with mu not simulated in MC) pileup_tool.SetUnrepresentedDataAction(2) pileup_tool.Initialize() # entering the main event loop... for event in chain: tree.reset() event.vertices.select(vertex_selection) tree.number_of_good_vertices = len(event.vertices) # match only with visible true taus event.truetaus.select(lambda tau: tau.vis_Et > 10 * GeV and abs(tau.vis_eta) < 2.5) if len(event.truetaus) == 1: true_tau = event.truetaus[0] TrueTauBlock.set(tree, 1, true_tau) else: continue # Truth-matching matched_reco = None reco_index = true_tau.tauAssoc_index tau = event.taus.getitem(reco_index) if tau in event.taus: matched_reco = tau else: continue tree.MET = event.MET_RefFinal_BDTMedium_et # fill tau block RecoTauBlock.set(event, tree, matched_reco, None) # set the event weight tree.pileup_weight = pileup_tool.GetCombinedWeight(event.RunNumber, event.mc_channel_number, event.averageIntPerXing) tree.mc_weight = event.mc_event_weight tree.Fill(reset=True) self.output.cd() tree.FlushBaskets() tree.Write() total_events = event_filters[0].total cutflow[0] = total_events cutflow[1] = total_events cutflow.Write()
def test_chain_draw_hist_init_first(self): hist = Hist(100, 0, 1) chain = TreeChain('tree', self.file_paths) chain.draw('a_x', hist=hist) assert_equals(hist.Integral() > 0, True)
class MuonTypeLJEvents(object): def __init__(self, files=None, type='MC', maxevents=-1): if type.upper() not in ['MC', 'DATA']: raise ValueError("Argument `type` need to be MC/DATA") self.Type = type self.MaxEvents = maxevents if not files: raise ValueError("Argument `files` need to be non-empty") if isinstance(files, str): files = [ files, ] self.Chain = TreeChain('ffNtuplizer/ffNtuple', files) ## register collections ### self.Chain.define_collection('pvs', prefix='pv_', size='pv_n') self.Chain.define_collection('muons', prefix='muon_', size='muon_n') self.Chain.define_collection('dsamuons', prefix='dsamuon_', size='dsamuon_n') self.Chain.define_collection('ak4jets', prefix='akjet_ak4PFJetsCHS_', size='akjet_ak4PFJetsCHS_n') self.Chain.define_collection('leptonjets', prefix='pfjet_', size='pfjet_n', mix=LeptonJetMix) self.Chain.define_collection('trigobjs', prefix='trigobj_', size='trigobj_n') self.Chain.define_object('hlt', prefix='HLT_') self.Chain.define_object('metfilters', prefix='metfilters_') self.Chain.define_object('cosmicveto', prefix='cosmicveto_') self.Histos = {} self.LookupWeight = root_open( os.path.join( os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/puWeights_10x_56ifb.root')).Get( 'puWeights') self.Scale = 1. def bookHisto(self, name, hist): self.Histos[name] = hist def setScale(self, scale): self.Scale = scale def process(self): for i, event in enumerate(self.Chain): if self.MaxEvents > 0 and i > self.MaxEvents: break ## event-level mask ## if not event.cosmicveto.result: continue if not event.metfilters.PrimaryVertexFilter: continue aux = {} ## event weight ## aux['wgt'] = self.Scale if self.Type == 'MC': aux['wgt'] *= event.weight # gen weight aux['wgt'] *= self.LookupWeight.GetBinContent( self.LookupWeight.GetXaxis().FindBin( event.trueInteractionNum)) ## pileup correction leptonjets = [ lj for lj in event.leptonjets if lj.passSelection(event) ] leptonjets.sort(key=lambda lj: lj.p4.pt(), reverse=True) leptonjets = leptonjets[:2] muontypeljs = [lj for lj in leptonjets if lj.isMuonType()] egmtypeljs = [lj for lj in leptonjets if lj.isEgmType()] # if not muontypeljs: continue aux['muontype'] = muontypeljs aux['egmtype'] = egmtypeljs aux['leptonjets'] = leptonjets looseMuonIdx, mediumMuonIdx = [], [] for i, mu in enumerate(event.muons): if mu.p4.pt() < 5: continue if abs(mu.p4.eta()) > 2.4: continue if (mu.selectors & (1 << 0)) != (1 << 0): continue # ID-loose if (mu.selectors & (1 << 7)) != (1 << 7): continue # Iso-loose looseMuonIdx.append(i) if (mu.selectors & (1 << 1)) != (1 << 1): continue # ID-medium mediumMuonIdx.append(i) if not mediumMuonIdx: continue muonpairs = [] for m in mediumMuonIdx: for l in looseMuonIdx: if m == l: continue if (l, m) in muonpairs: continue muonpairs.append((m, l)) if not muonpairs: continue aux['muonpairs'] = muonpairs self.processEvent(event, aux) def processEvent(self, event, aux): """To be override by daughter class""" pass @property def histos(self): return self.Histos
def work(self): """ This is the one function that all "ATLASStudent"s must implement. """ datatype = self.metadata.datatype year = self.metadata.year verbose = self.args.verbose OutputModel = C3POEvent if datatype == datasets.MC: # only create truth branches for MC OutputModel += ( FourVectModel.prefix('resonance_') + TrueTau.prefix('truetau1_') + TrueTau.prefix('truetau2_')) onfilechange = [] count_funcs = {} if datatype in (datasets.MC, datasets.EMBED): def mc_weight_count(event): return event.mc_event_weight count_funcs = { 'mc_weight': mc_weight_count, } trigger_config = None if datatype != datasets.EMBED: # trigger config tool to read trigger info in the ntuples trigger_config = get_trigger_config() # update the trigger config maps on every file change onfilechange.append((update_trigger_config, (trigger_config,))) if datatype == datasets.DATA: merged_grl = GRL() def update_grl(student, grl, name, file, tree): grl |= str(file.Get('Lumi/%s' % student.metadata.treename).GetString()) onfilechange.append((update_grl, (self, merged_grl,))) if datatype == datasets.DATA: merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D') else: merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D') def update_cutflow(student, cutflow, name, file, tree): year = student.metadata.year datatype = student.metadata.datatype if datatype == datasets.MC: cutflow[0] += file.cutflow_event[0] cutflow[1] += file.cutflow_event_mc_weight[0] else: cutflow[0] += file.cutflow_event[0] onfilechange.append((update_cutflow, (self, merged_cutflow,))) # initialize the TreeChain of all input files # (each containing one tree named self.metadata.treename) chain = TreeChain( self.metadata.treename, files=self.files, events=self.events, read_branches_on_demand=True, cache=True, onfilechange=onfilechange) # create output tree self.output.cd() tree = Tree(name='higgstautauhh', model=OutputModel) copied_variables = [ 'actualIntPerXing', 'averageIntPerXing', 'RunNumber', 'EventNumber', 'lbn'] tree.set_buffer( chain._buffer, branches=copied_variables, create_branches=True, visible=False) chain.always_read(copied_variables) # set the event filters event_filters = EventFilterList([ CoreFlags( count_funcs=count_funcs), TauSelected(2, count_funcs=count_funcs), TruthMatching( passthrough=datatype != datasets.MC, count_funcs=count_funcs), MCWeight( datatype=datatype, tree=tree, passthrough=datatype != datasets.MC, count_funcs=count_funcs) ]) self.filters['event'] = event_filters chain._filters += event_filters define_objects(chain, year, skim=False) # define tree objects taus = [ tree.define_object(name='tau1', prefix='tau1_'), tree.define_object(name='tau2', prefix='tau2_')] if datatype == datasets.MC: truetaus = [ tree.define_object(name='truetau1', prefix='truetau1_'), tree.define_object(name='truetau2', prefix='truetau2_')] tree.define_object(name='resonance', prefix='resonance_') # entering the main event loop... for event in chain: # sort taus and jets in decreasing order by pT event.taus.sort(key=lambda tau: tau.pt, reverse=True) tau1, tau2 = event.taus # MET METx = event.MET.etx METy = event.MET.ety MET_vect = Vector2(METx, METy) MET = event.MET.et MET_phi = event.MET.phi tree.MET = MET tree.MET_x = METx tree.MET_y = METy tree.MET_phi = MET_phi sumET = event.MET.sumet tree.sumET = sumET if sumET != 0: tree.MET_sig = ((2. * MET / GeV) / (utils.sign(sumET) * sqrt(abs(sumET / GeV)))) else: tree.MET_sig = -1. # use MMC values from skim mmc_mass = event.tau_MMC_mass mmc_resonance = event.tau_MMC_resonance mmc_met = Vector2(event.tau_MMC_MET_x, event.tau_MMC_MET_y) tree.mass_mmc_tau1_tau2 = mmc_mass tree.mmc_resonance.copy_from(mmc_resonance) if mmc_mass > 0: tree.mmc_resonance_pt = mmc_resonance.Pt() tree.MET_mmc = mmc_met.Mod() tree.MET_mmc_x = mmc_met.X() tree.MET_mmc_y = mmc_met.Y() tree.MET_mmc_phi = math.pi - mmc_met.Phi() # truth matching if datatype == datasets.MC: resonance, tau_decays = get_taus(event) if resonance is not None: FourVectModel.set(tree.resonance, resonance) matched_taus = [] decays = tau_decays[:] for itau, tau in enumerate(event.taus): for idecay, tau_decay in enumerate(decays): if tau.matches_vect(tau_decay.fourvect_visible): tau_decay.matched = True tau_decay.matched_object = tau tau.matched = True tau.matched_object = tau_decay TrueTau.set(truetaus[itau], tau_decay, verbose=verbose) decays.pop(idecay) matched_taus.append(itau) break if len(decays) > 0: for idecay, decay in enumerate(decays): reco_idx = -1 remaining_idx = range(2) for imatched in remaining_idx: if imatched not in matched_taus: reco_idx = imatched remaining_idx.remove(imatched) break TrueTau.set(truetaus[reco_idx], tau_decay, verbose=verbose) if len(tau_decays) == 2: # write truth met fourvect_missing = (tau_decays[0].fourvect_missing + tau_decays[1].fourvect_missing) tree.MET_true = fourvect_missing.Pt() tree.MET_phi_true = fourvect_missing.Phi() tree.MET_x_true = tree.MET_true * math.cos(tree.MET_phi_true) tree.MET_y_true = tree.MET_true * math.sin(tree.MET_phi_true) tree.MET_phi_diff = Vector2.Phi_mpi_pi(tree.MET_phi_true - MET_phi) # tau - vertex association tree.tau_same_vertex = ( tau1.privtx_x == tau2.privtx_x and tau1.privtx_y == tau2.privtx_y and tau1.privtx_z == tau2.privtx_z) # fill tau block for outtau, intau in zip(taus, event.taus): RecoTau.set(outtau, intau, verbose=verbose) # fill output ntuple tree.Fill(reset=True) self.output.cd() tree.FlushBaskets() tree.Write() if datatype == datasets.DATA: xml_string = ROOT.TObjString(merged_grl.str()) xml_string.Write('lumi') merged_cutflow.Write()
def test_chain_draw_hist_init_first(): hist = Hist(100, 0, 1) chain = TreeChain('tree', FILE_PATHS) chain.draw('a_x', hist=hist) assert_equal(hist.Integral() > 0, True)
class Events(object): def __init__(self, files=None, outname=None, type='MC', dtag='', maxevents=-1, channel=['4mu', '2mu2e'], ctau=None, chargedlj=False): if type.upper() not in ['MC', 'DATA']: raise ValueError("Argument `type` need to be MC/DATA") self.OutName = outname self.Type = type.upper() self.ChargedLJ = chargedlj self.MaxEvents = maxevents self.Channel = channel self.Dtag = dtag self.Ctau = ctau __signal_sample_param = dict( [substr.split('-') for substr in self.Dtag.split('_')]) self.SignalParam = { k.upper(): float(v.replace('p', '.')) for k, v in __signal_sample_param.items() } if not files: raise ValueError("Argument `files` need to be non-empty") if isinstance(files, str): files = [ files, ] self.Chain = TreeChain('ffNtuplizer/ffNtuple', files) ### register collections ### # self.Chain.define_collection('pvs', prefix='pv_', size='pv_n') self.Chain.define_collection('electrons', prefix='electron_', size='electron_n') self.Chain.define_collection('muons', prefix='muon_', size='muon_n') self.Chain.define_collection('dsamuons', prefix='dsamuon_', size='dsamuon_n') self.Chain.define_collection('photons', prefix='photon_', size='photon_n') self.Chain.define_collection('ak4jets', prefix='akjet_ak4PFJetsCHS_', size='akjet_ak4PFJetsCHS_n') self.Chain.define_collection('hftagscores', prefix='hftagscore_', size='hftagscore_n') self.Chain.define_collection('leptonjets', prefix='pfjet_', size='pfjet_n', mix=LeptonJetMix) self.Chain.define_collection('ljsources', prefix='ljsource_', size='ljsource_n') self.Chain.define_collection('cosmicmuons', prefix='cosmicmuon_', size='cosmicmuon_n') self.Chain.define_collection('trigobjs', prefix='trigobj_', size='trigobj_n') # self.Chain.define_collection('',) self.Chain.define_object('hlt', prefix='HLT_') self.Chain.define_object('metfilters', prefix='metfilters_') self.Chain.define_object('cosmicveto', prefix='cosmicveto_') self.Triggers = [ "DoubleL2Mu23NoVtx_2Cha", "DoubleL2Mu23NoVtx_2Cha_NoL2Matched", "DoubleL2Mu23NoVtx_2Cha_CosmicSeed", "DoubleL2Mu23NoVtx_2Cha_CosmicSeed_NoL2Matched", "DoubleL2Mu25NoVtx_2Cha_Eta2p4", "DoubleL2Mu25NoVtx_2Cha_CosmicSeed_Eta2p4", ] #self.addTRG = [ # "DoubleL2Mu23NoVtx_2Cha",] self.Histos = {} for chan in channel: self.Histos['{}/cutflow'.format(chan)] = ROOT.Hist( 20, 0, 20, title='cutflow', drawstyle='hist') self.KeepCutFlow = False self.RawCutFlow = False self.LookupWeight = root_open( os.path.join( os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/PUWeights_2018.root')).Get( 'puWeights') self.LookupMuonSFLowpT = root_open( os.path.join(os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/mu_Loose_pt7.root')).Get( 'ratio_syst') self.LookupMuonSF = root_open( os.path.join(os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/RunABCD_SF_ID.root')).Get( 'NUM_LooseID_DEN_TrackerMuons_pt_abseta_syst') self.LookupElectronSF = root_open( os.path.join( os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/2018_ElectronLoose.root')).Get( 'EGamma_SF2D') self.LookupPhotonSF = root_open( os.path.join( os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/2018_PhotonsLoose.root')).Get( 'EGamma_SF2D') self.Scale = 1. def setTriggers(self, triggers): self.Triggers = triggers def addTrigger(self, trigger): self.Triggers.append(trigger) def bookHisto(self, name, hist): self.Histos[name] = hist def setScale(self, scale): self.Scale = scale def process(self): for i, event in enumerate(self.Chain): if self.MaxEvents > 0 and i > self.MaxEvents: break ## event weight ## aux = {} aux['wgt'] = self.Scale if self.Type == 'MC': aux['wgt'] *= event.weight # gen weight aux['wgt'] *= self.LookupWeight.GetBinContent( self.LookupWeight.GetXaxis().FindBin( event.trueInteractionNum)) ## pileup correction for ch in self.Channel: if self.RawCutFlow: self.Histos['{}/cutflow'.format(ch)].Fill(0) else: self.Histos['{}/cutflow'.format(ch)].Fill(0, aux['wgt']) ## trigger ## #if not any([getattr(event.hlt, t) for t in self.Triggers]): continue #if not any([getattr(event.hlt, t) for t in self.addTRG]): continue for ch in self.Channel: if self.RawCutFlow: self.Histos['{}/cutflow'.format(ch)].Fill(1) else: self.Histos['{}/cutflow'.format(ch)].Fill(1, aux['wgt']) ## 2 leptonjets in channel definition ## leptonjets = [lj for lj in event.leptonjets] if len(leptonjets) < 2: continue #print len(leptonjets)#js #LJ2 = none leptonjets.sort(key=lambda lj: lj.p4.pt(), reverse=True) LJ0 = leptonjets[0] LJ1 = leptonjets[1] #LJ2 = leptonjets[2] #if len(event.leptonjets)>2:# #LJ2 = leptonjets[2]# #aux['lj2'] = LJ2 if self.ChargedLJ: if not LJ0.passChargedSelection(event): continue if not LJ1.passChargedSelection(event): continue #if not LJ2.passChargedSelection(event): continue else: if not LJ0.passSelection(event): continue if not LJ1.passSelection(event): continue #if not LJ2.passSelection(event): continue if LJ0.isMuonType() and LJ1.isMuonType(): aux['channel'] = '4mu' elif LJ0.isMuonType() and LJ1.isEgmType(): aux['channel'] = '2mu2e' elif LJ0.isEgmType() and LJ1.isMuonType(): aux['channel'] = '2mu2e' else: continue aux['lj0'] = LJ0 aux['lj1'] = LJ1 #aux['lj2'] = LJ2 if self.Type == 'MC': aux['sf_electron'] = 1. aux['sf_electron_up'] = 1. aux['sf_electron_low'] = 1. aux['sf_photon'] = 1. aux['sf_photon_up'] = 1. aux['sf_photon_low'] = 1. aux['sf_pfmuon'] = 1. aux['sf_pfmuon_up'] = 1. aux['sf_pfmuon_low'] = 1. for lj in [LJ0, LJ1]: ## muon scale factor for i in lj.pfcand_pfmuonIdx: _pfmu_p4 = event.muons[i].p4 pt, eta = _pfmu_p4.pt(), _pfmu_p4.eta() if pt >= 20: xbin = self.LookupMuonSF.xaxis.FindBin(pt) xbin = min(max(xbin, 1), self.LookupMuonSF.nbins(0)) ybin = self.LookupMuonSF.yaxis.FindBin(abs(eta)) sf = self.LookupMuonSF.GetBinContent(xbin, ybin) err_up = self.LookupMuonSF.GetBinErrorUp( xbin, ybin) err_lo = self.LookupMuonSF.GetBinErrorLow( xbin, ybin) else: for i in range(self.LookupMuonSFLowpT.num_points): x = self.LookupMuonSFLowpT.x(i) xh = x + self.LookupMuonSFLowpT.xerrh(i) xl = x - self.LookupMuonSFLowpT.xerrl(i) if xl <= eta and eta <= xh: sf = self.LookupMuonSFLowpT.y(i) err_up = self.LookupMuonSFLowpT.yerrh(i) err_lo = self.LookupMuonSFLowpT.yerrl(i) break aux['wgt'] *= sf aux['sf_pfmuon'] *= sf aux['sf_pfmuon_up'] *= sf + err_up aux['sf_pfmuon_low'] *= sf - err_lo ## NOTE DSA scale factor, nothing for now for i in lj.pfcand_dsamuonIdx: #print dsamu.pt.P4() sf = 1. aux['wgt'] *= sf ## electron scale factor for i in lj.pfcand_electronIdx: _electron = event.electrons[i] xbin = self.LookupElectronSF.xaxis.FindBin( _electron.scEta) ybin = self.LookupElectronSF.xaxis.FindBin( _electron.p4.pt()) ybin = min(max(ybin, 1), self.LookupElectronSF.nbins(1)) sf = self.LookupElectronSF.GetBinContent(xbin, ybin) aux['wgt'] *= sf aux['sf_electron'] *= sf aux['sf_electron_up'] *= sf + self.LookupElectronSF.GetBinErrorUp( xbin, ybin) aux['sf_electron_low'] *= sf - self.LookupElectronSF.GetBinErrorLow( xbin, ybin) ## photon scale factor for i in lj.pfcand_photonIdx: _photon = event.photons[i] xbin = self.LookupPhotonSF.xaxis.FindBin(_photon.scEta) ybin = self.LookupPhotonSF.xaxis.FindBin( _photon.p4.pt()) ybin = min(max(ybin, 1), self.LookupPhotonSF.nbins(1)) sf = self.LookupPhotonSF.GetBinContent(xbin, ybin) aux['wgt'] *= sf aux['sf_photon'] *= sf aux['sf_photon_up'] *= sf + self.LookupPhotonSF.GetBinErrorUp( xbin, ybin) aux['sf_photon_low'] *= sf - self.LookupPhotonSF.GetBinErrorLow( xbin, ybin) aux['wgt_electron_up'] = aux['wgt'] / aux['sf_electron'] * aux[ 'sf_electron_up'] aux['wgt_electron_low'] = aux['wgt'] / aux[ 'sf_electron'] * aux['sf_electron_low'] aux['wgt_photon_up'] = aux['wgt'] / aux['sf_photon'] * aux[ 'sf_photon_up'] aux['wgt_photon_low'] = aux['wgt'] / aux['sf_photon'] * aux[ 'sf_photon_low'] aux['wgt_pfmuon_up'] = aux['wgt'] / aux['sf_pfmuon'] * aux[ 'sf_pfmuon_up'] aux['wgt_pfmuon_low'] = aux['wgt'] / aux['sf_pfmuon'] * aux[ 'sf_pfmuon_low'] # for t, pt, eta in zip(list(lj.pfcand_type), list(lj.pfcand_pt), list(lj.pfcand_eta)): # ## muon scale factor, DSA same as muon for now # if t==3 or t==8: # xbin = self.LookupMuonSF.xaxis.FindBin(pt) # xbin = min(max(xbin, 1), self.LookupMuonSF.nbins(0)) # ybin = self.LookupMuonSF.yaxis.FindBin(abs(eta)) # sf = self.LookupMuonSF.GetBinContent(xbin, ybin) # aux['wgt'] *= sf # ## electron scale factor, using eta instead of SC eta for now # if t==2: # xbin = self.LookupElectronSF.xaxis.FindBin(eta) # ybin = self.LookupElectronSF.xaxis.FindBin(pt) # ybin = min(max(ybin, 1), self.LookupElectronSF.nbins(1)) # sf = self.LookupElectronSF.GetBinContent(xbin, ybin) # aux['wgt'] *= sf # ## photon scale factor, using eta instead of SC eta for now # if t==4: # xbin = self.LookupPhotonSF.xaxis.FindBin(eta) # ybin = self.LookupPhotonSF.xaxis.FindBin(pt) # ybin = min(max(ybin, 1), self.LookupPhotonSF.nbins(1)) # sf = self.LookupPhotonSF.GetBinContent(xbin, ybin) # aux['wgt'] *= sf for ch in self.Channel: if self.RawCutFlow: self.Histos['{}/cutflow'.format(ch)].Fill(2) else: self.Histos['{}/cutflow'.format(ch)].Fill(2, aux['wgt']) ## event-level mask ## if not event.metfilters.PrimaryVertexFilter: continue for ch in self.Channel: if self.RawCutFlow: self.Histos['{}/cutflow'.format(ch)].Fill(3) else: self.Histos['{}/cutflow'.format(ch)].Fill(3, aux['wgt']) nppCOSMIC, cosmicShowerTagged = globalCosmicShower( event.cosmicmuons, aux['channel']) if cosmicShowerTagged: continue for ch in self.Channel: if self.RawCutFlow: self.Histos['{}/cutflow'.format(ch)].Fill(4) else: self.Histos['{}/cutflow'.format(ch)].Fill(4, aux['wgt']) self.processEvent(event, aux) def processEvent(self, event, aux): """To be override by daughter class""" pass def postProcess(self): if self.KeepCutFlow: labels = [ 'total', 'trigger_pass', 'leptonjet_ge2', 'pv_good', 'cosmicveto_pass' ] for ch in self.Channel: xaxis = self.Histos['{}/cutflow'.format(ch)].axis(0) for i, s in enumerate(labels, start=1): xaxis.SetBinLabel(i, s) # binNum., labAngel, labSize, labAlign, labColor, labFont, labText xaxis.ChangeLabel(i, 315, -1, 11, -1, -1, s) else: for ch in self.Channel: self.Histos.pop('{}/cutflow'.format(ch)) @property def histos(self): return self.Histos @property def channel(self): return self.Channel
pass # ______________________________________________________________________________ # Analyzer mystate = 0 verbose = False batch_mode = True # Open file if mystate == 0: tree_name_i = '/home/jlow/L1MuonTrigger/CRAB3/P2_9_2_3_patch1/crab_projects/crab_ntuple_SingleMuon_PositiveEndCap/results/ntuple_SingleMuon_PositiveEndCap_%i.root' tree = TreeChain('ntupler/tree', [(tree_name_i % (i + 1)) for i in range(8)]) #tree_name_i = '/home/jlow/L1MuonTrigger/CRAB3/P2_9_2_3_patch1/crab_projects/crab_ntuple_SingleMuon_PositiveEndCap_PU200/results/ntuple_SingleMuon_PositiveEndCap_PU200_%i.root' #tree = TreeChain('ntupler/tree', [(tree_name_i % (i+1)) for i in range(50)]) elif mystate == 1: tree_name_i = '/home/jlow/L1MuonTrigger/CRAB3/P2_9_2_3_patch1/crab_projects/crab_ntuple_SingleNeutrino_PU200/results/ntuple_SingleNeutrino_PU200_%i.root' tree = TreeChain('ntupler/tree', [(tree_name_i % (i + 1)) for i in range(50)]) else: raise Exception("Unexpected state: %i" % mystate) maxEvents = -1 #maxEvents = 2000 # ROOT globals gROOT.SetBatch(batch_mode)
def __init__(self, files=None, outname=None, type='MC', dtag='', maxevents=-1, channel=['4mu', '2mu2e'], ctau=None, chargedlj=False): if type.upper() not in ['MC', 'DATA']: raise ValueError("Argument `type` need to be MC/DATA") self.OutName = outname self.Type = type.upper() self.ChargedLJ = chargedlj self.MaxEvents = maxevents self.Channel = channel self.Dtag = dtag self.Ctau = ctau __signal_sample_param = dict( [substr.split('-') for substr in self.Dtag.split('_')]) self.SignalParam = { k.upper(): float(v.replace('p', '.')) for k, v in __signal_sample_param.items() } if not files: raise ValueError("Argument `files` need to be non-empty") if isinstance(files, str): files = [ files, ] self.Chain = TreeChain('ffNtuplizer/ffNtuple', files) ### register collections ### # self.Chain.define_collection('pvs', prefix='pv_', size='pv_n') self.Chain.define_collection('electrons', prefix='electron_', size='electron_n') self.Chain.define_collection('muons', prefix='muon_', size='muon_n') self.Chain.define_collection('dsamuons', prefix='dsamuon_', size='dsamuon_n') self.Chain.define_collection('photons', prefix='photon_', size='photon_n') self.Chain.define_collection('ak4jets', prefix='akjet_ak4PFJetsCHS_', size='akjet_ak4PFJetsCHS_n') self.Chain.define_collection('hftagscores', prefix='hftagscore_', size='hftagscore_n') self.Chain.define_collection('leptonjets', prefix='pfjet_', size='pfjet_n', mix=LeptonJetMix) self.Chain.define_collection('ljsources', prefix='ljsource_', size='ljsource_n') self.Chain.define_collection('cosmicmuons', prefix='cosmicmuon_', size='cosmicmuon_n') self.Chain.define_collection('trigobjs', prefix='trigobj_', size='trigobj_n') # self.Chain.define_collection('',) self.Chain.define_object('hlt', prefix='HLT_') self.Chain.define_object('metfilters', prefix='metfilters_') self.Chain.define_object('cosmicveto', prefix='cosmicveto_') self.Triggers = [ "DoubleL2Mu23NoVtx_2Cha", "DoubleL2Mu23NoVtx_2Cha_NoL2Matched", "DoubleL2Mu23NoVtx_2Cha_CosmicSeed", "DoubleL2Mu23NoVtx_2Cha_CosmicSeed_NoL2Matched", "DoubleL2Mu25NoVtx_2Cha_Eta2p4", "DoubleL2Mu25NoVtx_2Cha_CosmicSeed_Eta2p4", ] #self.addTRG = [ # "DoubleL2Mu23NoVtx_2Cha",] self.Histos = {} for chan in channel: self.Histos['{}/cutflow'.format(chan)] = ROOT.Hist( 20, 0, 20, title='cutflow', drawstyle='hist') self.KeepCutFlow = False self.RawCutFlow = False self.LookupWeight = root_open( os.path.join( os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/PUWeights_2018.root')).Get( 'puWeights') self.LookupMuonSFLowpT = root_open( os.path.join(os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/mu_Loose_pt7.root')).Get( 'ratio_syst') self.LookupMuonSF = root_open( os.path.join(os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/RunABCD_SF_ID.root')).Get( 'NUM_LooseID_DEN_TrackerMuons_pt_abseta_syst') self.LookupElectronSF = root_open( os.path.join( os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/2018_ElectronLoose.root')).Get( 'EGamma_SF2D') self.LookupPhotonSF = root_open( os.path.join( os.getenv('CMSSW_BASE'), 'src/FireROOT/Analysis/data/2018_PhotonsLoose.root')).Get( 'EGamma_SF2D') self.Scale = 1.
def load_tree_multiple(infiles): print('[INFO] Opening file: %s' % ' '.join(infiles)) tree = TreeChain('ntupler/tree', infiles) define_collections(tree) return tree
def work(self): # get argument values local = self.args.local syst_terms = self.args.syst_terms datatype = self.metadata.datatype year = self.metadata.year verbose = self.args.student_verbose very_verbose = self.args.student_very_verbose redo_selection = self.args.redo_selection nominal_values = self.args.nominal_values # get the dataset name dsname = os.getenv('INPUT_DATASET_NAME', None) if dsname is None: # attempt to guess dsname from dirname if self.files: dsname = os.path.basename(os.path.dirname(self.files[0])) # is this a signal sample? # if so we will also keep some truth information in the output below is_signal = datatype == datasets.MC and ( '_VBFH' in dsname or '_ggH' in dsname or '_ZH' in dsname or '_WH' in dsname or '_ttH' in dsname) log.info("DATASET: {0}".format(dsname)) log.info("IS SIGNAL: {0}".format(is_signal)) # is this an inclusive signal sample for overlap studies? is_inclusive_signal = is_signal and '_inclusive' in dsname # is this a BCH-fixed sample? (temporary) is_bch_sample = 'r5470_r4540_p1344' in dsname if is_bch_sample: log.warning("this is a BCH-fixed r5470 sample") # onfilechange will contain a list of functions to be called as the # chain rolls over to each new file onfilechange = [] count_funcs = {} if datatype != datasets.DATA: # count the weighted number of events if local: def mc_weight_count(event): return event.hh_mc_weight else: def mc_weight_count(event): return event.mc_event_weight count_funcs = { 'mc_weight': mc_weight_count, } # three instances of the pileup reweighting tool are created to write # out the nominal, high and low pileup weights pileup_tool = None pileup_tool_high = None pileup_tool_low = None if local: # local means running on the skims, the output of this script # running on the grid if datatype == datasets.DATA: # merge the GRL fragments merged_grl = goodruns.GRL() def update_grl(student, grl, name, file, tree): grl |= str(file.Get('Lumi/%s' % student.metadata.treename).GetString()) onfilechange.append((update_grl, (self, merged_grl,))) if datatype == datasets.DATA: merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D') else: merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D') def update_cutflow(student, cutflow, name, file, tree): # record a cut-flow year = student.metadata.year datatype = student.metadata.datatype cutflow[1].value += file.cutflow_event[1].value if datatype != datasets.DATA: cutflow[2].value += file.cutflow_event_mc_weight[1].value onfilechange.append((update_cutflow, (self, merged_cutflow,))) else: # get pileup reweighting tool pileup_tool = get_pileup_reweighting_tool( year=year, use_defaults=True) pileup_tool_high = get_pileup_reweighting_tool( year=year, use_defaults=True, systematic='high') pileup_tool_low = get_pileup_reweighting_tool( year=year, use_defaults=True, systematic='low') if datatype not in (datasets.EMBED, datasets.MCEMBED): # merge TrigConfTrees metadirname = '%sMeta' % self.metadata.treename trigconfchain = ROOT.TChain('%s/TrigConfTree' % metadirname) map(trigconfchain.Add, self.files) metadir = self.output.mkdir(metadirname) metadir.cd() trigconfchain.Merge(self.output, -1, 'fast keep') self.output.cd() if datatype == datasets.DATA: # merge GRL XML strings merged_grl = goodruns.GRL() for fname in self.files: with root_open(fname) as f: for key in f.Lumi.keys(): merged_grl |= goodruns.GRL( str(key.ReadObj().GetString()), from_string=True) lumi_dir = self.output.mkdir('Lumi') lumi_dir.cd() xml_string= ROOT.TObjString(merged_grl.str()) xml_string.Write(self.metadata.treename) self.output.cd() self.output.cd() # create the output tree model = get_model(datatype, dsname, prefix=None if local else 'hh_', is_inclusive_signal=is_inclusive_signal) log.info("Output Model:\n\n{0}\n\n".format(model)) outtree = Tree(name=self.metadata.treename, model=model) if local: tree = outtree else: tree = outtree.define_object(name='tree', prefix='hh_') tree.define_object(name='tau', prefix='tau_') tree.define_object(name='tau1', prefix='tau1_') tree.define_object(name='tau2', prefix='tau2_') tree.define_object(name='truetau1', prefix='truetau1_') tree.define_object(name='truetau2', prefix='truetau2_') tree.define_object(name='jet1', prefix='jet1_') tree.define_object(name='jet2', prefix='jet2_') tree.define_object(name='jet3', prefix='jet3_') mmc_objects = [ tree.define_object(name='mmc0', prefix='mmc0_'), tree.define_object(name='mmc1', prefix='mmc1_'), tree.define_object(name='mmc2', prefix='mmc2_'), ] for mmc_obj in mmc_objects: mmc_obj.define_object(name='resonance', prefix='resonance_') trigger_emulation = TauTriggerEmulation( year=year, passthrough=local or datatype != datasets.MC or year > 2011, count_funcs=count_funcs) if not trigger_emulation.passthrough: onfilechange.append( (update_trigger_trees, (self, trigger_emulation,))) trigger_config = None if datatype not in (datasets.EMBED, datasets.MCEMBED): # trigger config tool to read trigger info in the ntuples trigger_config = get_trigger_config() # update the trigger config maps on every file change onfilechange.append((update_trigger_config, (trigger_config,))) # define the list of event filters if local and syst_terms is None and not redo_selection: event_filters = None else: tau_ntrack_recounted_use_ntup = False if year > 2011: # peek at first tree to determine if the extended number of # tracks is already stored with root_open(self.files[0]) as test_file: test_tree = test_file.Get(self.metadata.treename) tau_ntrack_recounted_use_ntup = ( 'tau_out_track_n_extended' in test_tree) event_filters = EventFilterList([ GRLFilter( self.grl, passthrough=( local or ( datatype not in (datasets.DATA, datasets.EMBED))), count_funcs=count_funcs), CoreFlags( passthrough=local, count_funcs=count_funcs), EmbeddingPileupPatch( passthrough=( local or year > 2011 or datatype != datasets.EMBED), count_funcs=count_funcs), averageIntPerXingPatch( passthrough=( local or year < 2012 or datatype != datasets.MC), count_funcs=count_funcs), PileupTemplates( year=year, passthrough=( local or is_bch_sample or datatype not in ( datasets.MC, datasets.MCEMBED)), count_funcs=count_funcs), RandomSeed( datatype=datatype, count_funcs=count_funcs), BCHSampleRunNumber( passthrough=not is_bch_sample, count_funcs=count_funcs), RandomRunNumber( tree=tree, datatype=datatype, pileup_tool=pileup_tool, passthrough=local, count_funcs=count_funcs), trigger_emulation, Triggers( year=year, tree=tree, datatype=datatype, passthrough=datatype in (datasets.EMBED, datasets.MCEMBED), count_funcs=count_funcs), PileupReweight( year=year, tool=pileup_tool, tool_high=pileup_tool_high, tool_low=pileup_tool_low, tree=tree, passthrough=( local or ( datatype not in (datasets.MC, datasets.MCEMBED))), count_funcs=count_funcs), PriVertex( passthrough=local, count_funcs=count_funcs), LArError( passthrough=local, count_funcs=count_funcs), TileError( passthrough=local, count_funcs=count_funcs), TileTrips( passthrough=( local or datatype in (datasets.MC, datasets.MCEMBED)), count_funcs=count_funcs), JetCopy( tree=tree, passthrough=local, count_funcs=count_funcs), # IMPORTANT! # JetCalibration MUST COME BEFORE ANYTHING THAT REFERS TO # jet.fourvect since jet.fourvect IS CACHED! JetCalibration( datatype=datatype, year=year, verbose=very_verbose, passthrough=local or nominal_values, count_funcs=count_funcs), # in situ TES shift for 2012 data TauEnergyShift( passthrough=( local or datatype != datasets.DATA or year < 2012 or nominal_values), count_funcs=count_funcs), # truth matching must come before systematics due to # TES_TRUE/FAKE TruthMatching( passthrough=datatype == datasets.DATA, count_funcs=count_funcs), NvtxJets( tree=tree, count_funcs=count_funcs), # PUT THE SYSTEMATICS "FILTER" BEFORE # ANY FILTERS THAT REFER TO OBJECTS # BUT AFTER CALIBRATIONS # Systematics must also come before anything that refers to # thing.fourvect since fourvect is cached! Systematics( terms=syst_terms, year=year, datatype=datatype, tree=tree, verbose=verbose, passthrough=not syst_terms, count_funcs=count_funcs), JetIsPileup( passthrough=( local or year < 2012 or datatype not in (datasets.MC, datasets.MCEMBED)), count_funcs=count_funcs), LArHole( tree=tree, passthrough=year > 2011, count_funcs=count_funcs), JetCleaning( datatype=datatype, year=year, count_funcs=count_funcs), ElectronVeto( count_funcs=count_funcs), MuonVeto( year=year, count_funcs=count_funcs), TauPT(2, thresh=20 * GeV, count_funcs=count_funcs), TauHasTrack(2, count_funcs=count_funcs), TauEta(2, count_funcs=count_funcs), TauElectronVeto(2, count_funcs=count_funcs), TauMuonVeto(2, count_funcs=count_funcs), TauAuthor(2, count_funcs=count_funcs), TauCrack(2, count_funcs=count_funcs), TauLArHole(2, tree=tree, passthrough=year > 2011, count_funcs=count_funcs), # before selecting the leading and subleading taus # be sure to only consider good candidates TauIDMedium(2, count_funcs=count_funcs), #TauTriggerMatchIndex( # config=trigger_config, # year=year, # datatype=datatype, # passthrough=datatype == datasets.EMBED, # count_funcs=count_funcs), # Select two leading taus at this point # 25 and 35 for data # 20 and 30 for MC to leave room for TES uncertainty TauLeadSublead( lead=( 35 * GeV if datatype == datasets.DATA or local else 30 * GeV), sublead=( 25 * GeV if datatype == datasets.DATA or local else 20 * GeV), count_funcs=count_funcs), # taus are sorted (in decreasing order) by pT from here on TauIDSelection( tree=tree, count_funcs=count_funcs), TaudR(3.2, count_funcs=count_funcs), #TauTriggerMatchThreshold( # datatype=datatype, # tree=tree, # count_funcs=count_funcs), TauTriggerEfficiency( year=year, datatype=datatype, tree=tree, tes_systematic=self.args.syst_terms and ( Systematics.TES_TERMS & self.args.syst_terms), passthrough=datatype == datasets.DATA, count_funcs=count_funcs), PileupScale( tree=tree, year=year, datatype=datatype, passthrough=local, count_funcs=count_funcs), TauIDScaleFactors( year=year, passthrough=datatype == datasets.DATA, count_funcs=count_funcs), TauFakeRateScaleFactors( year=year, datatype=datatype, tree=tree, tes_up=(self.args.syst_terms is not None and (Systematics.TES_FAKE_TOTAL_UP in self.args.syst_terms or Systematics.TES_FAKE_FINAL_UP in self.args.syst_terms)), tes_down=(self.args.syst_terms is not None and (Systematics.TES_FAKE_TOTAL_DOWN in self.args.syst_terms or Systematics.TES_FAKE_FINAL_DOWN in self.args.syst_terms)), passthrough=datatype in (datasets.DATA, datasets.EMBED), count_funcs=count_funcs), HiggsPT( year=year, tree=tree, passthrough=not is_signal or local, count_funcs=count_funcs), TauTrackRecounting( year=year, use_ntup_value=tau_ntrack_recounted_use_ntup, passthrough=local, count_funcs=count_funcs), MCWeight( datatype=datatype, tree=tree, passthrough=local or datatype == datasets.DATA, count_funcs=count_funcs), EmbeddingIsolation( tree=tree, passthrough=( local or year < 2012 or datatype not in (datasets.EMBED, datasets.MCEMBED)), count_funcs=count_funcs), EmbeddingCorrections( tree=tree, year=year, passthrough=( local or datatype not in (datasets.EMBED, datasets.MCEMBED)), count_funcs=count_funcs), EmbeddingTauSpinner( year=year, tree=tree, passthrough=( local or datatype not in ( datasets.EMBED, datasets.MCEMBED)), count_funcs=count_funcs), # put MET recalculation after tau selection but before tau-jet # overlap removal and jet selection because of the RefAntiTau # MET correction METRecalculation( terms=syst_terms, year=year, tree=tree, refantitau=not nominal_values, verbose=verbose, very_verbose=very_verbose, count_funcs=count_funcs), TauJetOverlapRemoval( count_funcs=count_funcs), JetPreselection( count_funcs=count_funcs), NonIsolatedJet( tree=tree, count_funcs=count_funcs), JetSelection( year=year, count_funcs=count_funcs), RecoJetTrueTauMatching( passthrough=datatype == datasets.DATA or local, count_funcs=count_funcs), BCHCleaning( tree=tree, passthrough=year == 2011 or local, datatype=datatype, count_funcs=count_funcs), ClassifyInclusiveHiggsSample( tree=tree, passthrough=not is_inclusive_signal, count_funcs=count_funcs), ]) # set the event filters self.filters['event'] = event_filters # peek at first tree to determine which branches to exclude with root_open(self.files[0]) as test_file: test_tree = test_file.Get(self.metadata.treename) ignore_branches = test_tree.glob( hhbranches.REMOVE, exclude=hhbranches.KEEP) ignore_branches_output = test_tree.glob( hhbranches.REMOVE_OUTPUT, exclude=hhbranches.KEEP_OUTPUT) # initialize the TreeChain of all input files chain = TreeChain( self.metadata.treename, files=self.files, ignore_branches=ignore_branches, events=self.events, onfilechange=onfilechange, filters=event_filters, cache=True, cache_size=50000000, learn_entries=100) if local: copied = [ 'EventNumber', ] hh_buffer = TreeBuffer() buffer = TreeBuffer() for name, value in chain._buffer.items(): if name.startswith('hh_'): hh_buffer[name[3:]] = value elif name in copied: buffer[name] = value outtree.set_buffer( hh_buffer, create_branches=False, visible=True) outtree.set_buffer( buffer, create_branches=True, visible=False) else: # additional decorations on existing objects if year > 2011 and datatype in (datasets.MC, datasets.MCEMBED): class Decorations(TreeModel): jet_ispileup = stl.vector('bool') chain.set_buffer(Decorations(), create_branches=True) # include the branches in the input chain in the output tree # set branches to be removed in ignore_branches outtree.set_buffer( chain._buffer, ignore_branches=ignore_branches + ignore_branches_output, create_branches=True, ignore_duplicates=True, transfer_objects=True, visible=False) # define tree objects define_objects(chain, year) # create the MMC mmc = mass.MMC(year=year) # report which packages have been loaded externaltools.report() self.output.cd() # The main event loop # the event filters above are automatically run for each event and only # the surviving events are looped on for event in chain: if local and syst_terms is None and not redo_selection: outtree.Fill() continue # sort taus and jets in decreasing order by pT event.taus.sort(key=lambda tau: tau.pt, reverse=True) event.jets.sort(key=lambda jet: jet.pt, reverse=True) # tau1 is the leading tau # tau2 is the subleading tau tau1, tau2 = event.taus jets = list(event.jets) jet1, jet2, jet3 = None, None, None beta = None if len(jets) >= 2: jet1, jet2 = jets[:2] # determine boost of system # determine jet CoM frame beta = (jet1.fourvect + jet2.fourvect).BoostVector() tree.jet_beta.copy_from(beta) jet1.fourvect_boosted.copy_from(jet1.fourvect) jet2.fourvect_boosted.copy_from(jet2.fourvect) jet1.fourvect_boosted.Boost(beta * -1) jet2.fourvect_boosted.Boost(beta * -1) tau1.fourvect_boosted.copy_from(tau1.fourvect) tau2.fourvect_boosted.copy_from(tau2.fourvect) tau1.fourvect_boosted.Boost(beta * -1) tau2.fourvect_boosted.Boost(beta * -1) tau1.min_dr_jet = min( tau1.fourvect.DeltaR(jet1.fourvect), tau1.fourvect.DeltaR(jet2.fourvect)) tau2.min_dr_jet = min( tau2.fourvect.DeltaR(jet1.fourvect), tau2.fourvect.DeltaR(jet2.fourvect)) #sphericity, aplanarity = eventshapes.sphericity_aplanarity( # [tau1.fourvect, # tau2.fourvect, # jet1.fourvect, # jet2.fourvect]) # sphericity #tree.sphericity = sphericity # aplanarity #tree.aplanarity = aplanarity #sphericity_boosted, aplanarity_boosted = eventshapes.sphericity_aplanarity( # [tau1.fourvect_boosted, # tau2.fourvect_boosted, # jet1.fourvect_boosted, # jet2.fourvect_boosted]) # sphericity #tree.sphericity_boosted = sphericity_boosted # aplanarity #tree.aplanarity_boosted = aplanarity_boosted # tau centrality (degree to which they are between the two jets) tau1.centrality = eventshapes.eta_centrality( tau1.fourvect.Eta(), jet1.fourvect.Eta(), jet2.fourvect.Eta()) tau2.centrality = eventshapes.eta_centrality( tau2.fourvect.Eta(), jet1.fourvect.Eta(), jet2.fourvect.Eta()) # boosted tau centrality tau1.centrality_boosted = eventshapes.eta_centrality( tau1.fourvect_boosted.Eta(), jet1.fourvect_boosted.Eta(), jet2.fourvect_boosted.Eta()) tau2.centrality_boosted = eventshapes.eta_centrality( tau2.fourvect_boosted.Eta(), jet1.fourvect_boosted.Eta(), jet2.fourvect_boosted.Eta()) # 3rd leading jet if len(jets) >= 3: jet3 = jets[2] jet3.fourvect_boosted.copy_from(jet3.fourvect) jet3.fourvect_boosted.Boost(beta * -1) elif len(jets) == 1: jet1 = jets[0] tau1.min_dr_jet = tau1.fourvect.DeltaR(jet1.fourvect) tau2.min_dr_jet = tau2.fourvect.DeltaR(jet1.fourvect) #sphericity, aplanarity = eventshapes.sphericity_aplanarity( # [tau1.fourvect, # tau2.fourvect, # jet1.fourvect]) # sphericity #tree.sphericity = sphericity # aplanarity #tree.aplanarity = aplanarity RecoJetBlock.set(tree, jet1, jet2, jet3, local=local) # mass of ditau + leading jet system if jet1 is not None: tree.mass_tau1_tau2_jet1 = ( tau1.fourvect + tau2.fourvect + jet1.fourvect).M() # full sphericity and aplanarity #sphericity_full, aplanarity_full = eventshapes.sphericity_aplanarity( # [tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets]) #tree.sphericity_full = sphericity_full #tree.aplanarity_full = aplanarity_full ##################################### # number of tracks from PV minus taus ##################################### ntrack_pv = 0 ntrack_nontau_pv = 0 for vxp in event.vertices: # primary vertex if vxp.type == 1: ntrack_pv = vxp.nTracks ntrack_nontau_pv = ntrack_pv - tau1.numTrack - tau2.numTrack break tree.ntrack_pv = ntrack_pv tree.ntrack_nontau_pv = ntrack_nontau_pv ######################### # MET variables ######################### METx = event.MET.etx METy = event.MET.ety MET = event.MET.et MET_vect = Vector2(METx, METy) MET_4vect = LorentzVector() MET_4vect.SetPxPyPzE(METx, METy, 0., MET) MET_4vect_boosted = LorentzVector() MET_4vect_boosted.copy_from(MET_4vect) if beta is not None: MET_4vect_boosted.Boost(beta * -1) tree.MET_et = MET tree.MET_etx = METx tree.MET_ety = METy tree.MET_phi = event.MET.phi dPhi_tau1_tau2 = abs(tau1.fourvect.DeltaPhi(tau2.fourvect)) dPhi_tau1_MET = abs(tau1.fourvect.DeltaPhi(MET_4vect)) dPhi_tau2_MET = abs(tau2.fourvect.DeltaPhi(MET_4vect)) tree.dPhi_tau1_tau2 = dPhi_tau1_tau2 tree.dPhi_tau1_MET = dPhi_tau1_MET tree.dPhi_tau2_MET = dPhi_tau2_MET tree.dPhi_min_tau_MET = min(dPhi_tau1_MET, dPhi_tau2_MET) tree.MET_bisecting = is_MET_bisecting( dPhi_tau1_tau2, dPhi_tau1_MET, dPhi_tau2_MET) sumET = event.MET.sumet tree.MET_sumet = sumET if sumET != 0: tree.MET_sig = ((2. * MET / GeV) / (utils.sign(sumET) * sqrt(abs(sumET / GeV)))) else: tree.MET_sig = -1. tree.MET_centrality = eventshapes.phi_centrality( tau1.fourvect, tau2.fourvect, MET_vect) tree.MET_centrality_boosted = eventshapes.phi_centrality( tau1.fourvect_boosted, tau2.fourvect_boosted, MET_4vect_boosted) tree.number_of_good_vertices = len(event.vertices) ########################## # Jet and sum pt variables ########################## tree.numJets = len(event.jets) # sum pT with only the two leading jets tree.sum_pt = sum( [tau1.pt, tau2.pt] + [jet.pt for jet in jets[:2]]) # sum pT with all selected jets tree.sum_pt_full = sum( [tau1.pt, tau2.pt] + [jet.pt for jet in jets]) # vector sum pT with two leading jets and MET tree.vector_sum_pt = sum( [tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets[:2]] + [MET_4vect]).Pt() # vector sum pT with all selected jets and MET tree.vector_sum_pt_full = sum( [tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets] + [MET_4vect]).Pt() # resonance pT tree.resonance_pt = sum( [tau1.fourvect, tau2.fourvect, MET_4vect]).Pt() ############################# # tau <-> vertex association ############################# tree.tau_same_vertex = ( tau1.privtx_x == tau2.privtx_x and tau1.privtx_y == tau2.privtx_y and tau1.privtx_z == tau2.privtx_z) tau1.vertex_prob = ROOT.TMath.Prob( tau1.privtx_chiSquared, int(tau1.privtx_numberDoF)) tau2.vertex_prob = ROOT.TMath.Prob( tau2.privtx_chiSquared, int(tau2.privtx_numberDoF)) ########################## # MMC Mass ########################## mmc_result = mmc.mass( tau1, tau2, METx, METy, sumET, njets=len(event.jets)) for mmc_method, mmc_object in enumerate(mmc_objects): mmc_mass, mmc_resonance, mmc_met = mmc_result[mmc_method] if verbose: log.info("MMC (method %d): %f" % (mmc_method, mmc_mass)) mmc_object.mass = mmc_mass mmc_object.MET_et = mmc_met.Mod() mmc_object.MET_etx = mmc_met.X() mmc_object.MET_ety = mmc_met.Y() mmc_object.MET_phi = math.pi - mmc_met.Phi() if mmc_mass > 0: FourMomentum.set(mmc_object.resonance, mmc_resonance) ############################ # collinear and visible mass ############################ vis_mass, collin_mass, tau1_x, tau2_x = mass.collinearmass( tau1, tau2, METx, METy) tree.mass_vis_tau1_tau2 = vis_mass tree.mass_collinear_tau1_tau2 = collin_mass tau1.collinear_momentum_fraction = tau1_x tau2.collinear_momentum_fraction = tau2_x ########################### # Match jets to VBF partons ########################### #if datatype == datasets.MC and 'VBF' in dsname and year == 2011: # # get partons (already sorted by eta in hepmc) FIXME!!! # parton1, parton2 = hepmc.get_VBF_partons(event) # tree.mass_true_quark1_quark2 = (parton1.fourvect + parton2.fourvect).M() # # order here needs to be revised since jets are no longer # # sorted by eta but instead by pT # PartonBlock.set(tree, parton1, parton2) # if len(jets) >= 2: # jet1, jet2 = jets[:2] # for i, jet in zip((1, 2), (jet1, jet2)): # for parton in (parton1, parton2): # if utils.dR(jet.eta, jet.phi, parton.eta, parton.phi) < .8: # setattr(tree, 'jet%i_matched' % i, True) # Fill the tau block # This must come after the RecoJetBlock is filled since # that sets the jet_beta for boosting the taus RecoTauBlock.set(event, tree, datatype, tau1, tau2, local=local) if datatype != datasets.DATA: TrueTauBlock.set(tree, tau1, tau2) # fill the output tree outtree.Fill(reset=True) externaltools.report() # flush any baskets remaining in memory to disk self.output.cd() outtree.FlushBaskets() outtree.Write() if local: if datatype == datasets.DATA: xml_string = ROOT.TObjString(merged_grl.str()) xml_string.Write('lumi') merged_cutflow.Write()
def test_chain_draw_hist_init_first(): hist = Hist(100, 0, 1) chain = TreeChain("tree", FILE_PATHS) chain.draw("a_x", hist=hist) assert_equal(hist.Integral() > 0, True)
# fill the tree for i in range(100): tree.x = gauss(.5, 1.) tree.y = gauss(.3, 2.) tree.z = gauss(13., 42.) tree.i = i tree.fill() tree.write() """ This section below takes the example trees and copies it while overwriting a branch with new values. """ # first define the chain of trees chain = TreeChain(name="test", files=fnames) # Now we want to copy the tree above into a new file while overwriting a branch # First create a new file to save the new tree in: f_copy = root_open("test_copy.root", "recreate") # You may not know the entire model of the original tree but only the branches # you intend to overwrite, so I am not specifying the model=Event below as an # example of how to deal with this in general: tree_copy = Tree("test_copy") # If the original tree was not handed to you through rootpy don't forget to: # >>> from rootpy import asrootpy # >>> tree = asrootpy(tree) # Here we specify the buffer for the new tree to use. We use the same buffer as
def work(self): # get argument values local = self.args.local syst_terms = self.args.syst_terms datatype = self.metadata.datatype year = self.metadata.year verbose = self.args.student_verbose very_verbose = self.args.student_very_verbose redo_selection = self.args.redo_selection nominal_values = self.args.nominal_values # get the dataset name dsname = os.getenv('INPUT_DATASET_NAME', None) if dsname is None: # attempt to guess dsname from dirname if self.files: dsname = os.path.basename(os.path.dirname(self.files[0])) # is this a signal sample? # if so we will also keep some truth information in the output below is_signal = datatype == datasets.MC and ( '_VBFH' in dsname or '_ggH' in dsname or '_ZH' in dsname or '_WH' in dsname or '_ttH' in dsname) log.info("DATASET: {0}".format(dsname)) log.info("IS SIGNAL: {0}".format(is_signal)) # is this an inclusive signal sample for overlap studies? is_inclusive_signal = is_signal and '_inclusive' in dsname # is this a BCH-fixed sample? (temporary) is_bch_sample = 'r5470_r4540_p1344' in dsname if is_bch_sample: log.warning("this is a BCH-fixed r5470 sample") # onfilechange will contain a list of functions to be called as the # chain rolls over to each new file onfilechange = [] count_funcs = {} if datatype != datasets.DATA: # count the weighted number of events if local: def mc_weight_count(event): return event.hh_mc_weight else: def mc_weight_count(event): return event.mc_event_weight count_funcs = { 'mc_weight': mc_weight_count, } # three instances of the pileup reweighting tool are created to write # out the nominal, high and low pileup weights pileup_tool = None pileup_tool_high = None pileup_tool_low = None if local: # local means running on the skims, the output of this script # running on the grid if datatype == datasets.DATA: # merge the GRL fragments merged_grl = goodruns.GRL() def update_grl(student, grl, name, file, tree): grl |= str( file.Get('Lumi/%s' % student.metadata.treename).GetString()) onfilechange.append((update_grl, ( self, merged_grl, ))) if datatype == datasets.DATA: merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D') else: merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D') def update_cutflow(student, cutflow, name, file, tree): # record a cut-flow year = student.metadata.year datatype = student.metadata.datatype cutflow[1].value += file.cutflow_event[1].value if datatype != datasets.DATA: cutflow[2].value += file.cutflow_event_mc_weight[1].value onfilechange.append((update_cutflow, ( self, merged_cutflow, ))) else: # get pileup reweighting tool pileup_tool = get_pileup_reweighting_tool(year=year, use_defaults=True) pileup_tool_high = get_pileup_reweighting_tool(year=year, use_defaults=True, systematic='high') pileup_tool_low = get_pileup_reweighting_tool(year=year, use_defaults=True, systematic='low') if datatype not in (datasets.EMBED, datasets.MCEMBED): # merge TrigConfTrees metadirname = '%sMeta' % self.metadata.treename trigconfchain = ROOT.TChain('%s/TrigConfTree' % metadirname) map(trigconfchain.Add, self.files) metadir = self.output.mkdir(metadirname) metadir.cd() trigconfchain.Merge(self.output, -1, 'fast keep') self.output.cd() if datatype == datasets.DATA: # merge GRL XML strings merged_grl = goodruns.GRL() for fname in self.files: with root_open(fname) as f: for key in f.Lumi.keys(): merged_grl |= goodruns.GRL(str( key.ReadObj().GetString()), from_string=True) lumi_dir = self.output.mkdir('Lumi') lumi_dir.cd() xml_string = ROOT.TObjString(merged_grl.str()) xml_string.Write(self.metadata.treename) self.output.cd() self.output.cd() # create the output tree model = get_model(datatype, dsname, prefix=None if local else 'hh_', is_inclusive_signal=is_inclusive_signal) log.info("Output Model:\n\n{0}\n\n".format(model)) outtree = Tree(name=self.metadata.treename, model=model) if local: tree = outtree else: tree = outtree.define_object(name='tree', prefix='hh_') tree.define_object(name='tau', prefix='tau_') tree.define_object(name='tau1', prefix='tau1_') tree.define_object(name='tau2', prefix='tau2_') tree.define_object(name='truetau1', prefix='truetau1_') tree.define_object(name='truetau2', prefix='truetau2_') tree.define_object(name='jet1', prefix='jet1_') tree.define_object(name='jet2', prefix='jet2_') tree.define_object(name='jet3', prefix='jet3_') mmc_objects = [ tree.define_object(name='mmc0', prefix='mmc0_'), tree.define_object(name='mmc1', prefix='mmc1_'), tree.define_object(name='mmc2', prefix='mmc2_'), ] for mmc_obj in mmc_objects: mmc_obj.define_object(name='resonance', prefix='resonance_') trigger_emulation = TauTriggerEmulation(year=year, passthrough=local or datatype != datasets.MC or year > 2011, count_funcs=count_funcs) if not trigger_emulation.passthrough: onfilechange.append((update_trigger_trees, ( self, trigger_emulation, ))) trigger_config = None if datatype not in (datasets.EMBED, datasets.MCEMBED): # trigger config tool to read trigger info in the ntuples trigger_config = get_trigger_config() # update the trigger config maps on every file change onfilechange.append((update_trigger_config, (trigger_config, ))) # define the list of event filters if local and syst_terms is None and not redo_selection: event_filters = None else: tau_ntrack_recounted_use_ntup = False if year > 2011: # peek at first tree to determine if the extended number of # tracks is already stored with root_open(self.files[0]) as test_file: test_tree = test_file.Get(self.metadata.treename) tau_ntrack_recounted_use_ntup = ('tau_out_track_n_extended' in test_tree) event_filters = EventFilterList([ averageIntPerXingPatch( passthrough=(local or year < 2012 or datatype != datasets.MC), count_funcs=count_funcs), PileupTemplates( year=year, passthrough=(local or is_bch_sample or datatype not in (datasets.MC, datasets.MCEMBED)), count_funcs=count_funcs), RandomSeed(datatype=datatype, count_funcs=count_funcs), RandomRunNumber(tree=tree, datatype=datatype, pileup_tool=pileup_tool, passthrough=local, count_funcs=count_funcs), PileupReweight( year=year, tool=pileup_tool, tool_high=pileup_tool_high, tool_low=pileup_tool_low, tree=tree, passthrough=(local or (datatype not in (datasets.MC, datasets.MCEMBED))), count_funcs=count_funcs), TruthMatching(passthrough=datatype == datasets.DATA, count_funcs=count_funcs), JetIsPileup( passthrough=(local or year < 2012 or datatype not in (datasets.MC, datasets.MCEMBED)), count_funcs=count_funcs), HiggsPT(year=year, tree=tree, passthrough=not is_signal or local, count_funcs=count_funcs), MCWeight(datatype=datatype, tree=tree, passthrough=local or datatype == datasets.DATA, count_funcs=count_funcs), ClassifyInclusiveHiggsSample( tree=tree, passthrough=not is_inclusive_signal, count_funcs=count_funcs), ]) # set the event filters self.filters['event'] = event_filters # peek at first tree to determine which branches to exclude with root_open(self.files[0]) as test_file: test_tree = test_file.Get(self.metadata.treename) ignore_branches = test_tree.glob(hhbranches.REMOVE, exclude=hhbranches.KEEP) ignore_branches_output = test_tree.glob( hhbranches.REMOVE_OUTPUT, exclude=hhbranches.KEEP_OUTPUT) # initialize the TreeChain of all input files chain = TreeChain(self.metadata.treename, files=self.files, ignore_branches=ignore_branches, events=self.events, onfilechange=onfilechange, filters=event_filters, cache=True, cache_size=50000000, learn_entries=100) if local: copied = [ 'EventNumber', ] hh_buffer = TreeBuffer() buffer = TreeBuffer() for name, value in chain._buffer.items(): if name.startswith('hh_'): hh_buffer[name[3:]] = value elif name in copied: buffer[name] = value outtree.set_buffer(hh_buffer, create_branches=False, visible=True) outtree.set_buffer(buffer, create_branches=True, visible=False) else: # additional decorations on existing objects if year > 2011 and datatype in (datasets.MC, datasets.MCEMBED): class Decorations(TreeModel): jet_ispileup = stl.vector('bool') chain.set_buffer(Decorations(), create_branches=True) # include the branches in the input chain in the output tree # set branches to be removed in ignore_branches outtree.set_buffer(chain._buffer, ignore_branches=ignore_branches + ignore_branches_output, create_branches=True, ignore_duplicates=True, transfer_objects=True, visible=False) # define tree objects define_objects(chain, year) # create the MMC mmc = mass.MMC(year=year) # report which packages have been loaded externaltools.report() self.output.cd() # The main event loop # the event filters above are automatically run for each event and only # the surviving events are looped on for event in chain: if local and syst_terms is None and not redo_selection: outtree.Fill() continue # sort taus and jets in decreasing order by pT event.taus.sort(key=lambda tau: tau.pt, reverse=True) event.jets.sort(key=lambda jet: jet.pt, reverse=True) # tau1 is the leading tau # tau2 is the subleading tau taus = list(event.taus) if len(taus) >= 2: tau1, tau2 = taus[0], taus[1] jets = list(event.jets) jet1, jet2, jet3 = None, None, None beta = None if len(jets) >= 2: jet1, jet2 = jets[:2] # determine boost of system # determine jet CoM frame beta = (jet1.fourvect + jet2.fourvect).BoostVector() tree.jet_beta.copy_from(beta) jet1.fourvect_boosted.copy_from(jet1.fourvect) jet2.fourvect_boosted.copy_from(jet2.fourvect) jet1.fourvect_boosted.Boost(beta * -1) jet2.fourvect_boosted.Boost(beta * -1) tau1.fourvect_boosted.copy_from(tau1.fourvect) tau2.fourvect_boosted.copy_from(tau2.fourvect) tau1.fourvect_boosted.Boost(beta * -1) tau2.fourvect_boosted.Boost(beta * -1) tau1.min_dr_jet = min(tau1.fourvect.DeltaR(jet1.fourvect), tau1.fourvect.DeltaR(jet2.fourvect)) tau2.min_dr_jet = min(tau2.fourvect.DeltaR(jet1.fourvect), tau2.fourvect.DeltaR(jet2.fourvect)) # sphericity, aplanarity = eventshapes.sphericity_aplanarity( # [tau1.fourvect, # tau2.fourvect, # jet1.fourvect, # jet2.fourvect]) # sphericity # tree.sphericity = sphericity # aplanarity # tree.aplanarity = aplanarity # sphericity_boosted, aplanarity_boosted = eventshapes.sphericity_aplanarity( # [tau1.fourvect_boosted, # tau2.fourvect_boosted, # jet1.fourvect_boosted, # jet2.fourvect_boosted]) # sphericity # tree.sphericity_boosted = sphericity_boosted # aplanarity # tree.aplanarity_boosted = aplanarity_boosted # tau centrality (degree to which they are between the two jets) tau1.centrality = eventshapes.eta_centrality( tau1.fourvect.Eta(), jet1.fourvect.Eta(), jet2.fourvect.Eta()) tau2.centrality = eventshapes.eta_centrality( tau2.fourvect.Eta(), jet1.fourvect.Eta(), jet2.fourvect.Eta()) # boosted tau centrality tau1.centrality_boosted = eventshapes.eta_centrality( tau1.fourvect_boosted.Eta(), jet1.fourvect_boosted.Eta(), jet2.fourvect_boosted.Eta()) tau2.centrality_boosted = eventshapes.eta_centrality( tau2.fourvect_boosted.Eta(), jet1.fourvect_boosted.Eta(), jet2.fourvect_boosted.Eta()) # 3rd leading jet if len(jets) >= 3: jet3 = jets[2] jet3.fourvect_boosted.copy_from(jet3.fourvect) jet3.fourvect_boosted.Boost(beta * -1) elif len(jets) == 1: jet1 = jets[0] tau1.min_dr_jet = tau1.fourvect.DeltaR(jet1.fourvect) tau2.min_dr_jet = tau2.fourvect.DeltaR(jet1.fourvect) # sphericity, aplanarity = eventshapes.sphericity_aplanarity( # [tau1.fourvect, # tau2.fourvect, # jet1.fourvect]) # sphericity # tree.sphericity = sphericity # aplanarity #tree.aplanarity = aplanarity RecoJetBlock.set(tree, jet1, jet2, jet3, local=local) # mass of ditau + leading jet system if jet1 is not None: tree.mass_tau1_tau2_jet1 = (tau1.fourvect + tau2.fourvect + jet1.fourvect).M() # full sphericity and aplanarity # sphericity_full, aplanarity_full = eventshapes.sphericity_aplanarity( # [tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets]) # tree.sphericity_full = sphericity_full # tree.aplanarity_full = aplanarity_full # #################################### # number of tracks from PV minus taus # #################################### ntrack_pv = 0 ntrack_nontau_pv = 0 for vxp in event.vertices: # primary vertex if vxp.type == 1: ntrack_pv = vxp.nTracks ntrack_nontau_pv = ntrack_pv - tau1.numTrack - tau2.numTrack break tree.ntrack_pv = ntrack_pv tree.ntrack_nontau_pv = ntrack_nontau_pv # ######################## # MET variables # ######################## METx = event.MET.etx METy = event.MET.ety MET = event.MET.et MET_vect = Vector2(METx, METy) MET_4vect = LorentzVector() MET_4vect.SetPxPyPzE(METx, METy, 0., MET) MET_4vect_boosted = LorentzVector() MET_4vect_boosted.copy_from(MET_4vect) if beta is not None: MET_4vect_boosted.Boost(beta * -1) tree.MET_et = MET tree.MET_etx = METx tree.MET_ety = METy tree.MET_phi = event.MET.phi dPhi_tau1_tau2 = abs(tau1.fourvect.DeltaPhi(tau2.fourvect)) dPhi_tau1_MET = abs(tau1.fourvect.DeltaPhi(MET_4vect)) dPhi_tau2_MET = abs(tau2.fourvect.DeltaPhi(MET_4vect)) tree.dPhi_tau1_tau2 = dPhi_tau1_tau2 tree.dPhi_tau1_MET = dPhi_tau1_MET tree.dPhi_tau2_MET = dPhi_tau2_MET tree.dPhi_min_tau_MET = min(dPhi_tau1_MET, dPhi_tau2_MET) tree.MET_bisecting = is_MET_bisecting(dPhi_tau1_tau2, dPhi_tau1_MET, dPhi_tau2_MET) sumET = event.MET.sumet tree.MET_sumet = sumET if sumET != 0: tree.MET_sig = ( (2. * MET / GeV) / (utils.sign(sumET) * sqrt(abs(sumET / GeV)))) else: tree.MET_sig = -1. tree.MET_centrality = eventshapes.phi_centrality( tau1.fourvect, tau2.fourvect, MET_vect) tree.MET_centrality_boosted = eventshapes.phi_centrality( tau1.fourvect_boosted, tau2.fourvect_boosted, MET_4vect_boosted) tree.number_of_good_vertices = len(event.vertices) # ######################### # Jet and sum pt variables # ######################### tree.numJets = len(event.jets) # sum pT with only the two leading jets tree.sum_pt = sum([tau1.pt, tau2.pt] + [jet.pt for jet in jets[:2]]) # sum pT with all selected jets tree.sum_pt_full = sum([tau1.pt, tau2.pt] + [jet.pt for jet in jets]) # vector sum pT with two leading jets and MET tree.vector_sum_pt = sum([tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets[:2]] + [MET_4vect]).Pt() # vector sum pT with all selected jets and MET tree.vector_sum_pt_full = sum([tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets] + [MET_4vect]).Pt() # resonance pT tree.resonance_pt = sum( [tau1.fourvect, tau2.fourvect, MET_4vect]).Pt() # ############################ # tau <-> vertex association # ############################ tree.tau_same_vertex = (tau1.privtx_x == tau2.privtx_x and tau1.privtx_y == tau2.privtx_y and tau1.privtx_z == tau2.privtx_z) tau1.vertex_prob = ROOT.TMath.Prob(tau1.privtx_chiSquared, int(tau1.privtx_numberDoF)) tau2.vertex_prob = ROOT.TMath.Prob(tau2.privtx_chiSquared, int(tau2.privtx_numberDoF)) # ######################### # MMC Mass # ######################### mmc_result = mmc.mass(tau1, tau2, METx, METy, sumET, njets=len(event.jets)) for mmc_method, mmc_object in enumerate(mmc_objects): mmc_mass, mmc_resonance, mmc_met = mmc_result[mmc_method] if verbose: log.info("MMC (method %d): %f" % (mmc_method, mmc_mass)) mmc_object.mass = mmc_mass mmc_object.MET_et = mmc_met.Mod() mmc_object.MET_etx = mmc_met.X() mmc_object.MET_ety = mmc_met.Y() mmc_object.MET_phi = math.pi - mmc_met.Phi() if mmc_mass > 0: FourMomentum.set(mmc_object.resonance, mmc_resonance) # ########################### # collinear and visible mass # ########################### vis_mass, collin_mass, tau1_x, tau2_x = mass.collinearmass( tau1, tau2, METx, METy) tree.mass_vis_tau1_tau2 = vis_mass tree.mass_collinear_tau1_tau2 = collin_mass tau1.collinear_momentum_fraction = tau1_x tau2.collinear_momentum_fraction = tau2_x ########################### # Match jets to VBF partons ########################### #if datatype == datasets.MC and 'VBF' in dsname and year == 2011: # # get partons (already sorted by eta in hepmc) FIXME!!! # parton1, parton2 = hepmc.get_VBF_partons(event) # tree.mass_true_quark1_quark2 = (parton1.fourvect + parton2.fourvect).M() # # order here needs to be revised since jets are no longer # # sorted by eta but instead by pT # PartonBlock.set(tree, parton1, parton2) # if len(jets) >= 2: # jet1, jet2 = jets[:2] # for i, jet in zip((1, 2), (jet1, jet2)): # for parton in (parton1, parton2): # if utils.dR(jet.eta, jet.phi, parton.eta, parton.phi) < .8: # setattr(tree, 'jet%i_matched' % i, True) # Fill the tau block # This must come after the RecoJetBlock is filled since # that sets the jet_beta for boosting the taus RecoTauBlock.set(event, tree, datatype, tau1, tau2, local=local) if datatype != datasets.DATA: TrueTauBlock.set(tree, tau1, tau2) # fill the output tree outtree.Fill(reset=True) externaltools.report() # flush any baskets remaining in memory to disk self.output.cd() outtree.FlushBaskets() outtree.Write() if local: if datatype == datasets.DATA: xml_string = ROOT.TObjString(merged_grl.str()) xml_string.Write('lumi') merged_cutflow.Write()
TrueTaus(count_funcs=count_funcs), ClassifyDecay(count_funcs=count_funcs, tree=outtree), TrueJets(count_funcs=count_funcs), ]) files = [args.input] # peek at first tree to determine which branches to exclude with root_open(files[0]) as test_file: test_tree = test_file.Get(args.tree_name) ignore_branches = test_tree.glob(branches.REMOVE) chain = TreeChain(args.tree_name, files=files, filters=event_filters, cache=True, cache_size=50000000, learn_entries=100, ignore_branches=ignore_branches) define_objects(chain) for event in chain: outtree.runnumber = event.RunNumber outtree.evtnumber = event.EventNumber outtree.weight = event.mc_event_weight # sort taus and jets in decreasing order by pT event.taus.sort(key=lambda tau: tau.decay.fourvect_vis.Pt(), reverse=True) event.jets.sort(key=lambda jet: jet.pt, reverse=True)