def eventLoop(modules,
inputFile,
outputFile,
inputTree,
wrappedOutputTree,
outputFileSmear=None,
outputTreeSmear=None,
typeofprocess=None,
maxEvents=-1,
eventRange=None,
progress=(10000, sys.stdout),
filterOutput=True):
for m in modules:
if typeofprocess == "smear":
m.beginFile(inputFile, outputFile, inputTree, wrappedOutputTree,
outputFileSmear, outputTreeSmear)
else:
m.beginFile(inputFile, outputFile, inputTree, wrappedOutputTree)
t0 = time.clock()
tlast = t0
doneEvents = 0
acceptedEvents = 0
entries = inputTree.entries
for i in xrange(entries) if eventRange == None else eventRange:
if maxEvents > 0 and i >= maxEvents - 1: break
e = Event(inputTree, i)
clearExtraBranches(inputTree)
doneEvents += 1
ret = True
for m in modules:
ret = m.analyze(e)
if not ret: break
if ret:
acceptedEvents += 1
if (
ret or not filterOutput
) and wrappedOutputTree != None and typeofprocess != "resp" and typeofprocess != "tau" and typeofprocess != "smear":
wrappedOutputTree.fill()
if progress:
if i > 0 and i % progress[0] == 0:
t1 = time.clock()
progress[1].write(
"Processed %8d/%8d entries (elapsed time %7.1fs, curr speed %8.3f kHz, avg speed %8.3f kHz), accepted %8d/%8d events (%5.2f%%)\n"
% (i, entries, t1 - t0, (progress[0] / 1000.) /
(max(t1 - tlast, 1e-9)), i / 1000. /
(max(t1 - t0, 1e-9)), acceptedEvents, doneEvents,
acceptedEvents / (0.01 * doneEvents)))
tlast = t1
for m in modules:
m.endFile(inputFile, outputFile, inputTree, wrappedOutputTree)
return (doneEvents, acceptedEvents, time.clock() - t0)
def eventLoop(modules,
inputFile,
outputFile,
inputTree,
wrappedOutputTree,
maxEvents=-1,
eventRange=None,
progress=(10000, sys.stdout),
filterOutput=True):
for m in modules:
m.beginFile(inputFile, outputFile, inputTree, wrappedOutputTree)
t0 = time.time()
tlast = t0
doneEvents = 0
acceptedEvents = 0
entries = inputTree.entries
if eventRange:
entries = len(eventRange)
if maxEvents > 0:
entries = min(entries, maxEvents)
for ie, i in enumerate(
range(entries) if eventRange == None else eventRange):
if maxEvents > 0 and ie >= maxEvents:
break
e = Event(inputTree, i)
clearExtraBranches(inputTree)
doneEvents += 1
ret = True
for m in modules:
ret = m.analyze(e)
if not ret:
break
if ret:
acceptedEvents += 1
if (ret or not filterOutput) and wrappedOutputTree != None:
wrappedOutputTree.fill()
if progress:
if ie > 0 and ie % progress[0] == 0:
t1 = time.time()
progress[1].write(
"Processed %8d/%8d entries, %5.2f%% (elapsed time %7.1fs, curr speed %8.3f kHz, avg speed %8.3f kHz), accepted %8d/%8d events (%5.2f%%)\n"
% (ie, entries, ie / float(0.01 * entries), t1 - t0,
(progress[0] / 1000.) /
(max(t1 - tlast, 1e-9)), ie / 1000. /
(max(t1 - t0, 1e-9)), acceptedEvents, doneEvents,
acceptedEvents / (0.01 * doneEvents)))
tlast = t1
for m in modules:
m.endFile(inputFile, outputFile, inputTree, wrappedOutputTree)
return (doneEvents, acceptedEvents, time.time() - t0)
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
if (x.pt > 500. and x.pt < 1000. and abs(x.eta) < 2.4)
]
for idx, gp in enumerate(particles):
if not hasattr(gp, 'dauIdx'):
gp.dauIdx = []
if gp.genPartIdxMother >= 0:
mom = particles[gp.genPartIdxMother]
if not hasattr(mom, 'dauIdx'):
def run(self):
outpostfix = self.postfix if self.postfix != None else (
"_Friend" if self.friend else "_Skim")
fullClone = False
outFileNames = []
totEntriesRead = 0
t0 = time.time()
for fileName in self.inputFiles:
# open file
print("Opening file %s" % fileName)
inFile = ROOT.TFile.Open(fileName)
if (not inFile): #check for null pointer
print("Unable to open file %s, exting \n" % fileName)
return 1
# get input tree
inTree = inFile.Get("Events")
nEntries = min(inTree.GetEntries() - self.firstEntry,
self.maxEntries)
totEntriesRead += nEntries
# pre-skimming
elist, jsonFilter = preSkim(inTree,
self.json,
self.cut,
maxEntries=self.maxEntries,
firstEntry=self.firstEntry)
# number of events to be processed
nTotal = elist.GetN() if elist else nEntries
print('Pre-select %d entries out of %s ' % (nTotal, nEntries))
inTree = InputTree(inTree, elist)
# output
outFileName = os.path.join(
self.outputDir,
os.path.basename(fileName).replace(".root",
outpostfix + ".root"))
#compressionAlgo = ROOT.ROOT.kLZMA
#compressionLevel = int(9)
compressionAlgo = ROOT.ROOT.kLZ4
compressionLevel = int(4)
outFile = ROOT.TFile.Open(outFileName, "RECREATE", "",
compressionLevel)
outFileNames.append(outFileName)
outFile.SetCompressionAlgorithm(compressionAlgo)
maxEntries = self.maxEntries
if self.perJet: #save two first jets
maxEntries = self.maxEntries * 2
outTree = FullOutput(inFile,
inTree,
outFile,
branchSelection=self.branchsel,
outputbranchSelection=self.outputbranchsel,
fullClone=fullClone,
maxEntries=maxEntries,
firstEntry=self.firstEntry,
jsonFilter=jsonFilter,
provenance=self.provenance)
t0 = time.time()
tlast = t0
doneEvents = 0
acceptedEvents = 0
if elist:
eventRange = [(elist.GetEntry(0) if i == 0 else elist.Next())
for i in range(elist.GetN())]
else:
eventRange = range(self.firstEntry, self.firstEntry +
nEntries) if nEntries > 0 else None
entries = inTree.entries
if eventRange:
entries = len(eventRange)
maxEvents = self.maxEntries
if maxEvents > 0:
entries = min(entries, self.maxEntries)
entriesRange = range(entries) if eventRange == None else eventRange
for m in self.modules:
m.beginFile(inFile, outFile, inTree, outTree, entriesRange)
for ie, i in enumerate(entriesRange):
if maxEvents > 0 and ie >= maxEvents: break
e = Event(inTree, ie)
ret = True
if self.perJet:
#print('ie ',ie)
for m in self.modules:
ret = m.analyze(e, ie)
if not ret: break
else:
clearExtraBranches(inTree)
m.fill(e, ie)
else:
clearExtraBranches(inTree)
for m in self.modules:
ret = m.analyze(e, ie)
if not ret: break
if ret and outTree is not None:
outTree.fill()
if ret:
acceptedEvents += 1
for m in self.modules:
m.endFile(inFile, outFile, inTree, outTree)
outTree.write()
outFile.Close()
print("Done %s" % outFileName)
for m in self.modules:
m.endJob()
print("Total time %.1f sec. to process %i events. Rate = %.1f Hz." %
((time.time() - t0), totEntriesRead, totEntriesRead /
(time.time() - t0)))
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()
count = 0
cut_1 = 0
cut_2 = 0
cut_3 = 0
cut_4 = 0
cut_5 = 0
cut_6 = 0
for entry in range(0, treeEntries):
count = count + 1
sys.stdout.write(
"%i / %i ... %.2f \r" %
(count, treeEntries, 100 * float(count) / float(treeEntries)))
sys.stdout.flush()
event = Event(inTree, entry)
fatJetsColl = Collection(event, "FatJet")
eColl = Collection(event, "Electron")
muColl = Collection(event, "Muon")
subJetsColl = Collection(event, 'SubJet')
fillHist(Jet1ptpass, fatJetsColl, 'pt', 0, [400, 2000])
fillHist(Jet2ptpass, fatJetsColl, 'pt', 1, [400, 2000])
fillHist(Jet1etapass,
fatJetsColl,
'eta',
0, [0.0, 2.4],
absval=True)
fillHist(Jet2etapass,
fatJetsColl,
def HarvestNanoAOD(inFilePath, outFilePath):
#
# Create the output file
#
print "Create Output File: %s" % (outFilePath)
f = ROOT.TFile(outFilePath, "RECREATE")
f.cd()
#
# Initialize the tree jet
#
treeName = "TreeFatJet"
print "Create Output Tree: %s" % (treeName)
TreeFatJet = ROOT.TTree(treeName, treeName)
#
# FatJet branch
#
nFatJetSizeMax = 10
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)
FatJetMSoftDrop = bookFloatArrayBranch(TreeFatJet, 'FatJet_msoftdrop',
nFatJetString, nFatJetSizeMax)
#
# GenPart branch
#
nGenPartSizeMax = 200
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 = 10
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)
#
# PV branch
#
PVnpvs = bookIntBranch(TreeFatJet, 'nPVnpvs')
PVnpvsGood = bookIntBranch(TreeFatJet, 'nPVnpvsGood')
PileUpNTrueInt = bookFloatBranch(TreeFatJet, 'nPileUpNTrueInt')
PileUpNPU = bookIntBranch(TreeFatJet, 'nPileUpNPU')
#
# SetupTChain
#
tree = ROOT.TChain("Events")
tree.Add(inFilePath)
#
# Use TChain and Setup TTreeReader.
#
inTree = InputTree(tree)
numEvents = inTree.entries
#
# Set max number of events to process
# Set to -1 if you want to run over all events
#
maxevents = -1
# maxevents = 5
#
# Loop over events
#
for iev in xrange(0, numEvents):
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()
FatJetTau21[i] = fj.tau2 / fj.tau1
FatJetTau31[i] = fj.tau3 / fj.tau1
FatJetTau32[i] = fj.tau3 / fj.tau2
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
nFatJet[0] += 1
#
# Loop over genparts
#
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
#
# 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 HarvestNanoAOD(inFilePath, outFilePath):
#
# Create the output file
#
print "Create Output File: %s" % (outFilePath)
f = ROOT.TFile(outFilePath, "RECREATE")
f.cd()
#
# Initialize the tree jet
#
treeName = "TreeFatJet"
print "Create Output Tree: %s" % (treeName)
TreeFatJet = ROOT.TTree(treeName, treeName)
#
# FatJet branch
#
nFatJetSizeMax = 10
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_m', nFatJetString,
nFatJetSizeMax)
#
# SetupTChain
#
tree = ROOT.TChain("Events")
tree.Add(inFilePath)
#
# Use TChain and Setup TTreeReader.
#
inTree = InputTree(tree)
numEvents = inTree.entries
#
# Set max number of events to process
# Set to -1 if you want to run over all events
#
maxevents = -1
# maxevents = 5
#
# Loop over events
#
for iev in xrange(0, numEvents):
if maxevents > 0 and iev > maxevents:
break
if (iev) % 250 == 0:
print "Processing event %d out of %d" % (iev, numEvents)
#
# Load Event
#
evt = Event(inTree, iev)
#
# GenParticles
#
particles = Collection(evt, "GenPart")
#
# Loop over jets
#
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()
nFatJet[0] += 1
#
# 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 eventLoop(modules,
inputFile,
outputFile,
inputTree,
wrappedOutputTree,
maxEvents=-1,
eventRange=None,
progress=(10000, sys.stdout),
filterOutput=True):
for m in modules:
m.beginFile(inputFile, outputFile, inputTree, wrappedOutputTree)
h_nevents = ROOT.TH1F('nEvents', 'nEvents', 4, 0, 4)
h_nevents.Sumw2()
h_nevents.GetXaxis().SetBinLabel(1, "total")
h_nevents.GetXaxis().SetBinLabel(2, "pos")
h_nevents.GetXaxis().SetBinLabel(3, "neg")
h_nevents.GetXaxis().SetBinLabel(4, "pass")
t0 = time.clock()
tlast = t0
doneEvents = 0
acceptedEvents = 0
entries = inputTree.entries
if eventRange: entries = len(eventRange)
if maxEvents > 0: entries = min(entries, maxEvents)
for ie, i in enumerate(
xrange(entries) if eventRange == None else eventRange):
if maxEvents > 0 and ie >= maxEvents: break
e = Event(inputTree, i)
weight = 1.
if hasattr(e, "genWeight"):
genWeight = getattr(e, "genWeight", None)
if genWeight:
weight = genWeight
h_nevents.Fill(0.5, weight)
if weight >= 0: h_nevents.Fill(1.5, weight)
else: h_nevents.Fill(2.5, weight)
clearExtraBranches(inputTree)
doneEvents += 1
ret = True
for m in modules:
ret = m.analyze(e)
if not ret: break
if ret:
acceptedEvents += 1
if (ret or not filterOutput) and wrappedOutputTree != None:
h_nevents.Fill(3.5, weight)
wrappedOutputTree.fill()
if progress:
if ie > 0 and ie % progress[0] == 0:
t1 = time.clock()
progress[1].write(
"Processed %8d/%8d entries, %5.2f%% (elapsed time %7.1fs, curr speed %8.3f kHz, avg speed %8.3f kHz), accepted %8d/%8d events (%5.2f%%)\n"
% (ie, entries, ie / float(0.01 * entries), t1 - t0,
(progress[0] / 1000.) /
(max(t1 - tlast, 1e-9)), ie / 1000. /
(max(t1 - t0, 1e-9)), acceptedEvents, doneEvents,
acceptedEvents / (0.01 * doneEvents)))
tlast = t1
for m in modules:
m.endFile(inputFile, outputFile, inputTree, wrappedOutputTree)
prevdir = ROOT.gDirectory
outputFile.cd()
h_nevents.Write()
prevdir.cd()
return (doneEvents, acceptedEvents, time.clock() - t0)
dominiisoscan = False
trigger = True
if not (dominiisoscan):
totalMCmu = 1.
totalMCele = 1.
totalnoMCmu = 1.
totalnoMCele = 1.
total_mu_ptrel_drmin = 1.
good_mu_ptrel_drmin = 1.
total_ele_ptrel_drmin = 1.
good_ele_ptrel_drmin = 1.
for i in xrange(0, tree.GetEntries()):
#for i in xrange(0,10000):
event = Event(tree, i)
electrons = Collection(event, "Electron")
muons = Collection(event, "Muon")
jets = Collection(event, "Jet")
fatjets = Collection(event, "FatJet")
PV = Object(event, "PV")
met = Object(event, "MET")
HLT = Object(event, "HLT")
Flag = Object(event, 'Flag')
LHEPdfWeight = Collection(event, 'LHEPdfWeight')
#print "no. of pdf weights " + str(len(LHEPdfWeight))
for pdfw in LHEPdfWeight:
print str(pdfw.__getattr__(""))
if trigger:
