INDIR += y elif ip == 2: INDIR += "_%s" % (y) ### GET LIST OF FILES ### inFileList = [EOSUSER + f for f in glob.glob(INDIR + ".root")] # print inFileList ### SETUP TCHAIN ### tree = ROOT.TChain(treeName) for inFile in inFileList: tree.Add(inFile) ### USE TCHAIN AND SETUP TTREEREADER ### inTree = InputTree(tree) nEvents = inTree.GetEntries() # nEvents=100 ### LOOP OVER EVENTS ### print "Total events: ", nEvents for i in xrange(0, nEvents): if i % 10000 == 0: print "Running %d out of %d" % (i, nEvents) evt = Event(inTree, i) particles = Collection(evt, "GenPart") fatjet = Collection(evt, "FatJet") jetak8 = Collection(evt, "GenJetAK8") fjInak8 = [ x for x in fatjet
def HarvestNanoAOD(inFileList, outFilePath, sample): # # Create the output file # print "Create Output File: %s" % (outFilePath) f = ROOT.TFile(outFilePath, "RECREATE") f.cd() # # Initialize the tree jet # treeName = "TreeFatJet" isMC_QCD = "QCD" in sample print "Create Output Tree: %s" % (treeName) TreeFatJet = ROOT.TTree(treeName, treeName) # # FatJet branch # nFatJetSizeMax = 25 nFatJetString = 'nFatJet' nFatJet = bookIntBranch(TreeFatJet, nFatJetString) FatJetPt = bookFloatArrayBranch(TreeFatJet, 'FatJet_pt', nFatJetString, nFatJetSizeMax) FatJetEta = bookFloatArrayBranch(TreeFatJet, 'FatJet_eta', nFatJetString, nFatJetSizeMax) FatJetPhi = bookFloatArrayBranch(TreeFatJet, 'FatJet_phi', nFatJetString, nFatJetSizeMax) FatJetM = bookFloatArrayBranch(TreeFatJet, 'FatJet_mass', nFatJetString, nFatJetSizeMax) FatJetTau21 = bookFloatArrayBranch(TreeFatJet, 'FatJet_tau21', nFatJetString, nFatJetSizeMax) FatJetTau31 = bookFloatArrayBranch(TreeFatJet, 'FatJet_tau31', nFatJetString, nFatJetSizeMax) FatJetTau32 = bookFloatArrayBranch(TreeFatJet, 'FatJet_tau32', nFatJetString, nFatJetSizeMax) FatJetDeepTagTvsQCD = bookFloatArrayBranch(TreeFatJet, 'FatJet_deepTag_TvsQCD', nFatJetString, nFatJetSizeMax) FatJetDeepTagWvsQCD = bookFloatArrayBranch(TreeFatJet, 'FatJet_deepTag_WvsQCD', nFatJetString, nFatJetSizeMax) FatJetDeepTagZvsQCD = bookFloatArrayBranch(TreeFatJet, 'FatJet_deepTag_ZvsQCD', nFatJetString, nFatJetSizeMax) FatJetDeepTagQCD = bookFloatArrayBranch(TreeFatJet, 'FatJet_deepTag_QCD', nFatJetString, nFatJetSizeMax) FatJetDeepTagQCDOthers = bookFloatArrayBranch(TreeFatJet, 'FatJet_deepTag_QCDothers', nFatJetString, nFatJetSizeMax) FatJetDeepTagMDTvsQCD = bookFloatArrayBranch(TreeFatJet, 'FatJet_deepTagMD_TvsQCD', nFatJetString, nFatJetSizeMax) FatJetDeepTagMDWvsQCD = bookFloatArrayBranch(TreeFatJet, 'FatJet_deepTagMD_WvsQCD', nFatJetString, nFatJetSizeMax) FatJetDeepTagMDZvsQCD = bookFloatArrayBranch(TreeFatJet, 'FatJet_deepTagMD_ZvsQCD', nFatJetString, nFatJetSizeMax) FatJetMSoftDrop = bookFloatArrayBranch(TreeFatJet, 'FatJet_msoftdrop', nFatJetString, nFatJetSizeMax) FatJetRawFactor = bookFloatArrayBranch(TreeFatJet, 'FatJet_rawFactor', nFatJetString, nFatJetSizeMax) FatJetJetId = bookIntArrayBranch(TreeFatJet, 'FatJet_jetId', nFatJetString, nFatJetSizeMax) FatJetSubJetIdx1 = bookIntArrayBranch(TreeFatJet, 'FatJet_subJetIdx1', nFatJetString, nFatJetSizeMax) FatJetSubJetIdx2 = bookIntArrayBranch(TreeFatJet, 'FatJet_subJetIdx2', nFatJetString, nFatJetSizeMax) FatJetGenJetAK8Idx = bookIntArrayBranch(TreeFatJet, 'FatJet_genJetAK8Idx', nFatJetString, nFatJetSizeMax) # # GenPart branch # if isMC_QCD == False: nGenPartSizeMax = 1000 nGenPartString = 'nGenPart' nGenPart = bookIntBranch(TreeFatJet, nGenPartString) GenPartPt = bookFloatArrayBranch(TreeFatJet, 'GenPart_pt', nGenPartString, nGenPartSizeMax) GenPartEta = bookFloatArrayBranch(TreeFatJet, 'GenPart_eta', nGenPartString, nGenPartSizeMax) GenPartPhi = bookFloatArrayBranch(TreeFatJet, 'GenPart_phi', nGenPartString, nGenPartSizeMax) GenPartM = bookFloatArrayBranch(TreeFatJet, 'GenPart_mass', nGenPartString, nGenPartSizeMax) GenPartPdgId = bookIntArrayBranch(TreeFatJet, 'GenPart_pdgId', nGenPartString, nGenPartSizeMax) GenPartStatus = bookIntArrayBranch(TreeFatJet, 'GenPart_status', nGenPartString, nGenPartSizeMax) GenPartStatusFlags = bookIntArrayBranch(TreeFatJet, 'GenPart_statusFlags', nGenPartString, nGenPartSizeMax) GenPartGenPartIdxMother = bookIntArrayBranch( TreeFatJet, 'GenPart_genPartIdxMother', nGenPartString, nGenPartSizeMax) # # GenJetAK8 branch # nGenJetAK8SizeMax = 25 nGenJetAK8String = 'nGenJetAK8' nGenJetAK8 = bookIntBranch(TreeFatJet, nGenJetAK8String) GenJetAK8Pt = bookFloatArrayBranch(TreeFatJet, 'GenJetAK8_pt', nGenJetAK8String, nGenJetAK8SizeMax) GenJetAK8Eta = bookFloatArrayBranch(TreeFatJet, 'GenJetAK8_eta', nGenJetAK8String, nGenJetAK8SizeMax) GenJetAK8Phi = bookFloatArrayBranch(TreeFatJet, 'GenJetAK8_phi', nGenJetAK8String, nGenJetAK8SizeMax) GenJetAK8M = bookFloatArrayBranch(TreeFatJet, 'GenJetAK8_mass', nGenJetAK8String, nGenJetAK8SizeMax) GenJetAK8HadronFlavour = bookIntArrayBranch(TreeFatJet, 'GenJetAK8_hadronFlavour', nGenJetAK8String, nGenJetAK8SizeMax) GenJetAK8PartonFlavour = bookIntArrayBranch(TreeFatJet, 'GenJetAK8_partonFlavour', nGenJetAK8String, nGenJetAK8SizeMax) # # Subjet branch # nSubJetSizeMax = 50 nSubJetString = 'nSubJet' nSubJet = bookIntBranch(TreeFatJet, nSubJetString) SubJetPt = bookFloatArrayBranch(TreeFatJet, 'SubJet_pt', nSubJetString, nSubJetSizeMax) SubJetEta = bookFloatArrayBranch(TreeFatJet, 'SubJet_eta', nSubJetString, nSubJetSizeMax) SubJetPhi = bookFloatArrayBranch(TreeFatJet, 'SubJet_phi', nSubJetString, nSubJetSizeMax) SubJetM = bookFloatArrayBranch(TreeFatJet, 'SubJet_mass', nSubJetString, nSubJetSizeMax) SubJetRawFactor = bookFloatArrayBranch(TreeFatJet, 'SubJet_rawFactor', nSubJetString, nSubJetSizeMax) SubJetNBHadrons = bookIntArrayBranch(TreeFatJet, 'SubJet_nBHadrons', nSubJetString, nSubJetSizeMax) SubJetNCHadrons = bookIntArrayBranch(TreeFatJet, 'SubJet_nCHadrons', nSubJetString, nSubJetSizeMax) # # SubGenJetAK8 branch # nSubGenJetAK8SizeMax = 50 nSubGenJetAK8String = 'nSubGenJetAK8' nSubGenJetAK8 = bookIntBranch(TreeFatJet, nSubGenJetAK8String) SubGenJetAK8Pt = bookFloatArrayBranch(TreeFatJet, 'SubGenJetAK8_pt', nSubGenJetAK8String, nSubGenJetAK8SizeMax) SubGenJetAK8Eta = bookFloatArrayBranch(TreeFatJet, 'SubGenJetAK8_eta', nSubGenJetAK8String, nSubGenJetAK8SizeMax) SubGenJetAK8Phi = bookFloatArrayBranch(TreeFatJet, 'SubGenJetAK8_phi', nSubGenJetAK8String, nSubGenJetAK8SizeMax) SubGenJetAK8M = bookFloatArrayBranch(TreeFatJet, 'SubGenJetAK8_mass', nSubGenJetAK8String, nSubGenJetAK8SizeMax) # # PV branch # PVnpvs = bookIntBranch(TreeFatJet, 'nPVnpvs') PVnpvsGood = bookIntBranch(TreeFatJet, 'nPVnpvsGood') PileUpNTrueInt = bookFloatBranch(TreeFatJet, 'nPileUpNTrueInt') PileUpNPU = bookIntBranch(TreeFatJet, 'nPileUpNPU') # # SetupTChain # tree = ROOT.TChain("Events") for inFilePath in inFileList: print 'Adding files: %s' % (inFilePath) tree.Add(inFilePath) tree.ls() # # Use TChain and Setup TTreeReader. # inTree = InputTree(tree) # # # if isMC_QCD: branchSel = BranchSelection("branchSel_QCD.txt") else: branchSel = BranchSelection("branchSel.txt") branchSel.selectBranches(inTree) numEvents = inTree.GetEntries() # # Set max number of events to process # Set to -1 if you want to run over all events # maxevents = -1 # maxevents = 1000 # # Loop over events # print numEvents for iev in xrange(0, numEvents): # print iev if maxevents > 0 and iev > maxevents: break if (iev) % 1000 == 0: print "Processing event %d out of %d" % (iev, numEvents) # # Load Event # evt = Event(inTree, iev) # # Loop over fatjets # fatjets = Collection(evt, "FatJet") nFatJet[0] = 0 for i, fj in enumerate(fatjets): fj_p4 = fj.p4() FatJetPt[i] = fj_p4.Pt() FatJetEta[i] = fj_p4.Eta() FatJetPhi[i] = fj_p4.Phi() FatJetM[i] = fj_p4.M() if fj.tau1 > 0: FatJetTau21[i] = fj.tau2 / fj.tau1 else: FatJetTau21[i] = -1 if fj.tau1 > 0: FatJetTau31[i] = fj.tau3 / fj.tau1 else: FatJetTau31[i] = -1 if fj.tau2 > 0: FatJetTau32[i] = fj.tau3 / fj.tau2 else: FatJetTau32[i] = -1 FatJetDeepTagMDTvsQCD[i] = fj.deepTagMD_TvsQCD FatJetDeepTagMDWvsQCD[i] = fj.deepTagMD_WvsQCD FatJetDeepTagMDZvsQCD[i] = fj.deepTagMD_ZvsQCD FatJetDeepTagTvsQCD[i] = fj.deepTag_TvsQCD FatJetDeepTagWvsQCD[i] = fj.deepTag_WvsQCD FatJetDeepTagZvsQCD[i] = fj.deepTag_ZvsQCD FatJetDeepTagQCD[i] = fj.deepTag_QCD FatJetDeepTagQCDOthers[i] = fj.deepTag_QCDothers FatJetMSoftDrop[i] = fj.msoftdrop FatJetRawFactor[i] = fj.rawFactor FatJetJetId[i] = fj.jetId FatJetSubJetIdx1[i] = fj.subJetIdx1 FatJetSubJetIdx2[i] = fj.subJetIdx2 FatJetGenJetAK8Idx[i] = fj.genJetAK8Idx nFatJet[0] += 1 # # Loop over genparts # if isMC_QCD == False: particles = Collection(evt, "GenPart") nGenPart[0] = 0 for i, gp in enumerate(particles): GenPartPt[i] = gp.pt GenPartEta[i] = gp.eta GenPartPhi[i] = gp.phi GenPartM[i] = gp.mass GenPartPdgId[i] = gp.pdgId GenPartStatus[i] = gp.status GenPartStatusFlags[i] = gp.statusFlags GenPartGenPartIdxMother[i] = gp.genPartIdxMother nGenPart[0] += 1 # # Loop over GenJetAK8 # jets = Collection(evt, "GenJetAK8") nGenJetAK8[0] = 0 for i, gj in enumerate(jets): GenJetAK8Pt[i] = gj.pt GenJetAK8Eta[i] = gj.eta GenJetAK8Phi[i] = gj.phi GenJetAK8M[i] = gj.mass GenJetAK8HadronFlavour[i] = gj.hadronFlavour GenJetAK8PartonFlavour[i] = gj.partonFlavour nGenJetAK8[0] += 1 # # Subjet over GenJetAK8 # subjets = Collection(evt, "SubJet") nSubJet[0] = 0 for i, sj in enumerate(subjets): SubJetPt[i] = sj.pt SubJetEta[i] = sj.eta SubJetPhi[i] = sj.phi SubJetM[i] = sj.mass SubJetRawFactor[i] = sj.rawFactor SubJetNBHadrons[i] = sj.nBHadrons SubJetNCHadrons[i] = sj.nCHadrons nSubJet[0] += 1 # # Subjet over GenJetAK8 # subjets = Collection(evt, "SubGenJetAK8") nSubGenJetAK8[0] = 0 for i, sj in enumerate(subjets): SubGenJetAK8Pt[i] = sj.pt SubGenJetAK8Eta[i] = sj.eta SubGenJetAK8Phi[i] = sj.phi SubGenJetAK8M[i] = sj.mass nSubGenJetAK8[0] += 1 # # Loop over PV # PVnpvs[0] = evt.PV_npvs PVnpvsGood[0] = evt.PV_npvsGood PileUpNTrueInt[0] = evt.Pileup_nTrueInt PileUpNPU[0] = evt.Pileup_nPU # # Fill the tree for this event # TreeFatJet.Fill() # # Save the output ttree in the output file # print "Write tree to file" f.Write() # # Gracefully close the output file # print "Closing output" f.Close()
def run(self): if not self.noOut: outpostfix = self.postfix if self.postfix != None else ( "_Friend" if self.friend else "_Skim") if self.compression != "none": ROOT.gInterpreter.ProcessLine("#include <Compression.h>") (algo, level) = self.compression.split(":") compressionLevel = int(level) if algo == "LZMA": compressionAlgo = ROOT.ROOT.kLZMA elif algo == "ZLIB": compressionAlgo = ROOT.ROOT.kZLIB else: raise RuntimeError("Unsupported compression %s" % algo) else: compressionLevel = 0 print "Will write selected trees to " + self.outputDir if not self.justcount: if not os.path.exists(self.outputDir): os.system("mkdir -p " + self.outputDir) if self.noOut: if len(self.modules) == 0: raise RuntimeError( "Running with --noout and no modules does nothing!") for m in self.modules: m.beginJob() fullClone = (len(self.modules) == 0) outFileNames = [] t0 = time.clock() totEntriesRead = 0 for fname in self.inputFiles: # open input file inFile = ROOT.TFile.Open(fname) #get input tree inTree = inFile.Get("Events") totEntriesRead += inTree.GetEntries() # pre-skimming elist, jsonFilter = preSkim(inTree, self.json, self.cut) if self.justcount: print 'Would select %d entries from %s' % ( elist.GetN() if elist else inTree.GetEntries(), fname) continue else: print 'Pre-select %d entries out of %s ' % ( elist.GetN() if elist else inTree.GetEntries(), inTree.GetEntries()) if fullClone: # no need of a reader (no event loop), but set up the elist if available if elist: inTree.SetEntryList(elist) else: # initialize reader inTree = InputTree(inTree, elist) # prepare output file outFileName = os.path.join( self.outputDir, os.path.basename(fname).replace(".root", outpostfix + ".root")) outFile = ROOT.TFile.Open(outFileName, "RECREATE", "", compressionLevel) outFileNames.append(outFileName) if compressionLevel: outFile.SetCompressionAlgorithm(compressionAlgo) # prepare output tree if self.friend: outTree = FriendOutput(inFile, inTree, outFile) else: outTree = FullOutput(inFile, inTree, outFile, branchSelection=self.branchsel, fullClone=fullClone, jsonFilter=jsonFilter, provenance=self.provenance) # process events, if needed if not fullClone: (nall, npass, timeLoop) = eventLoop(self.modules, inFile, outFile, inTree, outTree) print 'Processed %d preselected entries from %s (%s entries). Finally selected %d entries' % ( nall, fname, inTree.GetEntries(), npass) else: print 'Selected %d entries from %s' % ( outTree.tree().GetEntries(), fname) # now write the output outTree.write() outFile.Close() print "Done %s" % outFileName if self.jobReport: self.jobReport.addInputFile(fname, nall) for m in self.modules: m.endJob() print totEntriesRead / (time.clock() - t0), "Hz" if self.haddFileName: os.system( "./haddnano.py %s %s" % (self.haddFileName, " ".join(outFileNames)) ) #FIXME: remove "./" once haddnano.py is distributed with cms releases if self.jobReport: self.jobReport.addOutputFile(self.haddFileName) self.jobReport.save()
def run(self): outpostfix = self.postfix if self.postfix != None else ( "_Friend" if self.friend else "_Skim") if not self.noOut: if self.compression != "none": ROOT.gInterpreter.ProcessLine("#include <Compression.h>") (algo, level) = self.compression.split(":") compressionLevel = int(level) if algo == "LZMA": compressionAlgo = ROOT.ROOT.kLZMA elif algo == "ZLIB": compressionAlgo = ROOT.ROOT.kZLIB else: raise RuntimeError("Unsupported compression %s" % algo) else: compressionLevel = 0 print "Will write selected trees to " + self.outputDir if not self.justcount: if not os.path.exists(self.outputDir): os.system("mkdir -p " + self.outputDir) else: compressionLevel = 0 if self.noOut: if len(self.modules) == 0: raise RuntimeError( "Running with --noout and no modules does nothing!") # Open histogram file, if desired if (self.histFileName != None and self.histDirName == None) or ( self.histFileName == None and self.histDirName != None): raise RuntimeError( "Must specify both histogram file and histogram directory!") elif self.histFileName != None and self.histDirName != None: self.histFile = ROOT.TFile.Open(self.histFileName, "RECREATE") else: self.histFile = None for m in self.modules: if hasattr(m, 'writeHistFile') and m.writeHistFile: m.beginJob(histFile=self.histFile, histDirName=self.histDirName) else: m.beginJob() fullClone = (len(self.modules) == 0) outFileNames = [] t0 = time.clock() totEntriesRead = 0 for fname in self.inputFiles: # open input file inFile = ROOT.TFile.Open(fname) #get input tree inTree = inFile.Get("Events") totEntriesRead += inTree.GetEntries() self.hcount.SetBinContent(1, inTree.GetEntries()) ROOT.gROOT.SetBatch(True) if self.SMSMasses != None: inTree.Draw( "MaxIf$(GenPart_mass, abs(GenPart_pdgId) == %i):MaxIf$(GenPart_mass, abs(GenPart_pdgId) == %i) >> hSMS(2000, -0.5, 1999.5, 2000, -0.5, 1999.5)" % (self.SMSMasses[0], self.SMSMasses[1])) self.hsmscount = ROOT.gDirectory.Get('hSMS') if self.doISR != None: #Dirty ISR recipe for EWKinos #Need to correct for each mass point #Can't correct per sample (wrong normalization), need to save whole unskimmed histogram per point an then postprocess pt1 = "MaxIf$(GenPart_pt, abs(GenPart_pdgId) == %i && GenPart_status == 22)" % self.doISR[ 0] pt2 = "MaxIf$(GenPart_pt, abs(GenPart_pdgId) == %i && GenPart_status == 22)" % self.doISR[ 1] phi1 = "MaxIf$(GenPart_pt, abs(GenPart_pdgId) == %i && GenPart_status == 22)" % self.doISR[ 0] phi2 = "MaxIf$(GenPart_pt, abs(GenPart_pdgId) == %i && GenPart_status == 22)" % self.doISR[ 1] pt_ISR = "hypot(%s + %s * cos(%s-%s), %s*sin(%s - %s))" % ( pt1, pt2, phi2, phi1, pt2, phi2, phi1) inTree.Draw( " %s : MaxIf$(GenPart_mass, abs(GenPart_pdgId) == %i) : MaxIf$(GenPart_mass, abs(GenPart_pdgId) == %i) >> hISR(1000, -0.5, 1999.5, 1000, -0.5, 1999.5, 20, 0, 1000)" % (pt_ISR, self.SMSMasses[0], self.SMSMasses[1])) self.hISR = ROOT.gDirectory.Get("hISR") if inTree.GetBranchStatus("genWeight"): inTree.Project("SumWeightsTemp", "1.0", "genWeight") sow = ROOT.gROOT.FindObject("SumWeightsTemp").Integral() self.hsumofweights.SetBinContent(1, sow) # pre-skimming elist, jsonFilter = preSkim(inTree, self.json, self.cut) if self.justcount: print 'Would select %d entries from %s' % ( elist.GetN() if elist else inTree.GetEntries(), fname) continue else: print 'Pre-select %d entries out of %s ' % ( elist.GetN() if elist else inTree.GetEntries(), inTree.GetEntries()) if fullClone: # no need of a reader (no event loop), but set up the elist if available if elist: inTree.SetEntryList(elist) else: # initialize reader inTree = InputTree(inTree, elist) # prepare output file if not self.noOut: outFileName = os.path.join( self.outputDir, os.path.basename(fname).replace(".root", outpostfix + ".root")) outFile = ROOT.TFile.Open(outFileName, "RECREATE", "", compressionLevel) outFileNames.append(outFileName) if compressionLevel: outFile.SetCompressionAlgorithm(compressionAlgo) # prepare output tree if self.friend: outTree = FriendOutput(inFile, inTree, outFile) else: outTree = FullOutput( inFile, inTree, outFile, branchSelection=self.branchsel, outputbranchSelection=self.outputbranchsel, fullClone=fullClone, jsonFilter=jsonFilter, provenance=self.provenance) else: outFile = None outTree = None # process events, if needed if not fullClone: (nall, npass, timeLoop) = eventLoop(self.modules, inFile, outFile, inTree, outTree) print 'Processed %d preselected entries from %s (%s entries). Finally selected %d entries' % ( nall, fname, inTree.GetEntries(), npass) else: nall = inTree.GetEntries() print 'Selected %d entries from %s' % ( outTree.tree().GetEntries(), fname) # now write the output if not self.noOut: print "Start writing" self.hcount.Write() print "Start writing" if self.SMSMasses != None: self.hsmscount.Write() print "Start writing" if self.doISR != None: self.hISR.Write() print "Start writing" self.hsumofweights.Write() outTree.write() outFile.Close() print "Done %s" % outFileName if self.jobReport: self.jobReport.addInputFile(fname, nall) for m in self.modules: m.endJob() print totEntriesRead / (time.clock() - t0), "Hz" if self.haddFileName: os.system( "./haddnano.py %s %s" % (self.haddFileName, " ".join(outFileNames)) ) #FIXME: remove "./" once haddnano.py is distributed with cms releases if self.jobReport: self.jobReport.addOutputFile(self.haddFileName) self.jobReport.save()