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