NanoAODTools/test/exampleGenDump.py

0001 #! /usr/bin/env python3
0002 # Sources:
0003 #   https://cms-nanoaod-integration.web.cern.ch/integration/master-106X/mc106Xul18_doc.html#GenPart
0004 #   https://github.com/cms-sw/cmssw/blob/master/PhysicsTools/NanoAOD/python/genparticles_cff.py
0005 #   https://github.com/cms-sw/cmssw/blob/master/PhysicsTools/NanoAOD/plugins/LHETablesProducer.cc
0006 from PhysicsTools.NanoAODTools.postprocessing.framework.postprocessor import PostProcessor
0007 from PhysicsTools.NanoAODTools.postprocessing.framework.eventloop import Module
0008 from PhysicsTools.NanoAODTools.postprocessing.framework.datamodel import Collection
0009 import PhysicsTools.NanoAODTools.postprocessing.framework.datamodel as datamodel
0010 datamodel.statusflags['isHardPrompt'] = datamodel.statusflags['isPrompt'] + datamodel.statusflags['fromHardProcess']
0011
0012 def hasbit(value,bit):
0013   """Check if i'th bit is set to 1, i.e. binary of 2^i,
0014   from the right to the left, starting from position i=0.
0015   Example: hasbit(GenPart_statusFlags,0) -> isPrompt"""
0016   return (value & (1 << bit))>0
0017
0018
0019 def getprodchain(part,genparts=None,event=None,decay=-1):
0020   """Print production chain recursively."""
0021   chain = "%3s"%(part.pdgId)
0022   imoth = part.genPartIdxMother
0023   while imoth>=0:
0024     if genparts is not None:
0025       moth = genparts[imoth]
0026       chain = "%3s -> "%(moth.pdgId)+chain
0027       imoth = moth.genPartIdxMother
0028     elif event is not None:
0029       chain = "%3s -> "%(event.GenPart_pdgId[imoth])+chain
0030       imoth = event.GenPart_genPartIdxMother[imoth]
0031   if genparts is not None and decay>0:
0032     chain = chain[:-3] # remove last particle
0033     chain += getdecaychain(part,genparts,indent=len(chain),depth=decay-1)
0034   return chain
0035
0036
0037 def getdecaychain(part,genparts,indent=0,depth=999):
0038   """Print decay chain recursively."""
0039   chain   = "%3s"%(part.pdgId)
0040   imoth   = part._index
0041   ndaus   = 0
0042   indent_ = len(chain)+indent
0043   for idau in range(imoth+1,genparts._len):
0044     dau = genparts[idau]
0045     if dau.genPartIdxMother==imoth: # found daughter
0046       if ndaus>=1:
0047         chain += '\n'+' '*indent_
0048       if depth>=2:
0049         chain += " -> "+getdecaychain(dau,genparts,indent=indent_+4,depth=depth-1)
0050       else: # stop recursion
0051         chain += " -> %3s"%(dau.pdgId)
0052       ndaus += 1
0053   return chain
0054
0055
0056 # DUMPER MODULE
0057 class LHEDumper(Module):
0058
0059   def __init__(self):
0060     self.nleptonic = 0
0061     self.ntauonic  = 0
0062     self.nevents   = 0
0063
0064   def analyze(self,event):
0065     """Dump gen information for each gen particle in given event."""
0066     print("\n%s Event %s %s"%('-'*10,event.event,'-'*70))
0067     self.nevents += 1
0068     leptonic = False
0069     tauonic = False
0070     bosons = [ ]
0071     taus = [ ]
0072     particles = Collection(event,'GenPart')
0073     #particles = Collection(event,'LHEPart')
0074     print(" \033[4m%7s %7s %7s %7s %7s %7s %7s %7s %8s %9s %10s  \033[0m"%(
0075       "index","pdgId","moth","mothId","dR","pt","eta","status","prompt","taudecay","last copy"))
0076     for i, particle in enumerate(particles):
0077       mothidx  = particle.genPartIdxMother
0078       eta      = max(-999,min(999,particle.eta))
0079       prompt   = particle.statusflag('isPrompt') #hasbit(particle.statusFlags,0)
0080       taudecay = particle.statusflag('isTauDecayProduct') #hasbit(particle.statusFlags,2)
0081       lastcopy = particle.statusflag('isLastCopy') #hasbit(particle.statusFlags,13)
0082       #ishardprompt = particle.statusflag('isHardPrompt')
0083       if 0<=mothidx<particles._len:
0084         moth    = particles[mothidx]
0085         mothpid = moth.pdgId
0086         mothdR  = max(-999,min(999,particle.DeltaR(moth))) #particle.p4().DeltaR(moth.p4())
0087         print(" %7d %7d %7d %7d %7.2f %7.2f %7.2f %7d %8s %9s %10s"%(
0088           i,particle.pdgId,mothidx,mothpid,mothdR,particle.pt,eta,particle.status,prompt,taudecay,lastcopy))
0089       else:
0090         print(" %7d %7d %7s %7s %7s %7.2f %7.2f %7d %8s %9s %10s"%(
0091           i,particle.pdgId,"","","",particle.pt,eta,particle.status,prompt,taudecay,lastcopy))
0092       if lastcopy:
0093         if abs(particle.pdgId) in [11,13,15]:
0094           leptonic = True
0095           if abs(particle.pdgId)==15:
0096             tauonic = True
0097             taus.append(particle)
0098         elif abs(particle.pdgId) in [23,24]:
0099           bosons.append(particle)
0100     for boson in bosons: # print production chain
0101       print("Boson production:")
0102       print(getprodchain(boson,particles,decay=2))
0103     for tau in taus: # print decay chain
0104       print("Tau decay:")
0105       print(getdecaychain(tau,particles))
0106     if leptonic:
0107       self.nleptonic += 1
0108     if tauonic:
0109       self.ntauonic += 1
0110
0111   def endJob(self):
0112     print('\n'+'-'*96)
0113     if self.nevents>0:
0114       print("  %-10s %4d / %-4d (%4.1f%%)"%('Tauonic: ',self.ntauonic, self.nevents,100.0*self.ntauonic/self.nevents))
0115       print("  %-10s %4d / %-4d (%4.1f%%)"%('Leptonic:',self.nleptonic,self.nevents,100.0*self.nleptonic/self.nevents))
0116     print("%s Done %s\n"%('-'*10,'-'*80))
0117
0118
0119 # PROCESS NANOAOD
0120 url = "root://cms-xrd-global.cern.ch/"
0121 from argparse import ArgumentParser
0122 parser = ArgumentParser()
0123 parser.add_argument('-i', '--infiles', nargs='+')
0124 parser.add_argument('-o', '--outdir', default='.')
0125 parser.add_argument('-n', '--maxevts', type=int, default=20)
0126 args = parser.parse_args()
0127 infiles = args.infiles or [
0128   url+'/store/mc/RunIISummer20UL18NanoAODv9/DYJetsToLL_M-50_TuneCP5_13TeV-madgraphMLM-pythia8/NANOAODSIM/106X_upgrade2018_realistic_v16_L1v1-v1/280000/525CD279-3344-6043-98B9-2EA8A96623E4.root',
0129   #url+'/store/mc/RunIISummer20UL18NanoAODv9/TTTo2L2Nu_TuneCP5_13TeV-powheg-pythia8/NANOAODSIM/106X_upgrade2018_realistic_v16_L1v1-v1/130000/44187D37-0301-3942-A6F7-C723E9F4813D.root',
0130 ]
0131 processor = PostProcessor(args.outdir,infiles,noOut=True,modules=[LHEDumper()],maxEntries=args.maxevts)
0132 processor.run()