analyzers/gen/GeneratorAnalyzer.py

0001 from __future__ import print_function
0002 from builtins import range
0003 from PhysicsTools.Heppy.analyzers.core.Analyzer import Analyzer
0004 from PhysicsTools.Heppy.analyzers.core.AutoHandle import AutoHandle
0005 from PhysicsTools.Heppy.physicsutils.genutils import isNotFromHadronicShower, realGenMothers, realGenDaughters
0006
0007 def interestingPdgId(id,includeLeptons=False):
0008     id = abs(id)
0009     return id in [6,7,8,17,18] or (includeLeptons and 11 <= id and id < 16) or (22 <= id and id < 40) or id > 1000000
0010
0011 class GeneratorAnalyzer( Analyzer ):
0012     """Save the hard-scattering final state of the event: top quarks, gauge & higgs bosons and BSM
0013        particles, plus their immediate decay products, and their siblings (in order to get the jets
0014        from matched X+jets generation.
0015        Incoming partons are not included, by design.
0016
0017        In the default configuration, leptons, light quarks and gluons are saved before FSR (a la status 3).
0018        Everything else is saved after any radiation, i.e. immediately before the decay.
0019
0020        Particles are saved in a list event.generatorSummary, with the index of their mothers ('motherIndex')
0021        if the mother is also in the list, and with the pdgId of the mother ('motherId') and grand-mother
0022        ('grandmotherId'). Particles also carry their index in the miniAOD genparticles collection ('rawIndex')
0023        In addition, a 'sourceId' is set to the pdgId of the heaviest ancestor (or of the particle itself)
0024        i.e.  in  top -> W -> lepton: the lepton sourceId will be 6
0025              in  tt+W with W -> lepton, the sourceId of the lepton will be 24.
0026        sourceId will be 99 for paricles from hard scattering whose mother is light
0027
0028        If requested, the full list of genParticles is also produced in event.genParticles (with indices
0029        aligned to the miniAOD one). For particles that are in the generatorSummary, the same object is used.
0030        An extra index 'genSummaryIndex' will be added to all particles, with the index in the generatorSummary
0031        or -1 if the particle is not in the generatorSummary.
0032
0033        Also, if requested it creates the splitted collections:
0034             event.genHiggsBosons = []
0035             event.genVBosons = []
0036             event.gennus     = []  # prompt neutrinos
0037             event.gennusFromTop = []  # Neutrinos from t->W decay
0038             event.genleps    = []  # leptons from direct decays
0039             event.gentauleps = []  # leptons from prompt taus
0040             event.gentaus    = []  # hadronically-decaying taus (if allGenTaus is False) or all taus (if allGenTaus is True)
0041             event.gentopquarks  = []
0042             event.genbquarks    = [] # b quarks from hard event (e.g. from top decays)
0043             event.genwzquarks   = [] # quarks from W,Z decays
0044             event.genbquarksFromTop = []
0045             event.genbquarksFromH   = []
0046             event.genlepsFromTop = [] #mu/ele that have a t->W chain as ancestor, also contained in event.genleps
0047        event.genwzquarks and event.genbquarks, might have overlaps
0048        event.genbquarksFromTop and event.genbquarksFromH are all contained in event.genbquarks
0049
0050        """
0051
0052     def __init__(self, cfg_ana, cfg_comp, looperName ):
0053         super(GeneratorAnalyzer,self).__init__(cfg_ana,cfg_comp,looperName)
0054         self.stableBSMParticleIds  = set(cfg_ana.stableBSMParticleIds) # neutralinos and such
0055         self.savePreFSRParticleIds = set(cfg_ana.savePreFSRParticleIds)
0056         self.makeAllGenParticles   = cfg_ana.makeAllGenParticles
0057         self.makeSplittedGenLists  = cfg_ana.makeSplittedGenLists
0058         self.allGenTaus            = cfg_ana.allGenTaus if self.makeSplittedGenLists else False
0059
0060     def declareHandles(self):
0061         super(GeneratorAnalyzer, self).declareHandles()
0062         self.mchandles['genParticles'] = AutoHandle( 'prunedGenParticles', 'std::vector<reco::GenParticle>' )
0063
0064     def beginLoop(self,setup):
0065         super(GeneratorAnalyzer,self).beginLoop(setup)
0066
0067     def makeMCInfo(self, event):
0068         verbose = getattr(self.cfg_ana, 'verbose', False)
0069         rawGenParticles = self.mchandles['genParticles'].product()
0070         good = []; keymap = {};
0071         allGenParticles = []
0072         for rawIndex,p in enumerate(rawGenParticles):
0073             if self.makeAllGenParticles: allGenParticles.append(p)
0074             id     = abs(p.pdgId())
0075             status = p.status()
0076             # particles must be status > 2, except for prompt leptons, photons, neutralinos
0077             if status <= 2:
0078                 if ((id not in self.stableBSMParticleIds) and
0079                     (id not in [11,12,13,14,15,16,22] or not isNotFromHadronicShower(p))):
0080                         continue
0081             # a particle must not be decaying into itself
0082             #print "  test %6d  : %+8d  %3d :  %8.2f   %+5.2f   %+5.2f : %d %d : %+8d {%6d}: %s" % ( rawIndex,
0083             #        p.pdgId(), p.status(), p.pt(), p.eta(), p.phi(), p.numberOfMothers(), p.numberOfDaughters(),
0084             #        p.motherRef(0).pdgId() if p.numberOfMothers() > 0 else -999, p.motherRef(0).key()   if p.numberOfMothers() > 0 else -999,
0085             #        "  ".join("%d[%d]" % (p.daughter(i).pdgId(), p.daughter(i).status()) for i in xrange(p.numberOfDaughters())))
0086             if id in self.savePreFSRParticleIds:
0087                 # for light objects, we want them pre-radiation
0088                 if any((p.mother(j).pdgId() == p.pdgId()) for j in range(p.numberOfMothers())):
0089                     #print "    fail auto-decay"
0090                     continue
0091             else:
0092                 # everything else, we want it after radiation, i.e. just before decay
0093                 if any((p.daughter(j).pdgId() == p.pdgId() and p.daughter(j).status() > 2) for j in range(p.numberOfDaughters())):
0094                     #print "    fail auto-decay"
0095                     continue
0096             # FIXME find a better criterion to discard there
0097             if status == 71:
0098                 #drop QCD radiation with unclear parentage
0099                 continue
0100             # is it an interesting particle?
0101             ok = False
0102             if interestingPdgId(id):
0103                 #print "    pass pdgId"
0104                 ok = True
0105             ### no: we don't select by decay, so that we keep the particle summary free of incoming partons and such
0106             # if not ok and any(interestingPdgId(d.pdgId()) for d in realGenDaughters(p)):
0107             #    #print "    pass dau"
0108             #     ok = True
0109             if not ok:
0110               for mom in realGenMothers(p):
0111                 if interestingPdgId(mom.pdgId()) or (getattr(mom,'rawIndex',-1) in keymap):
0112                     #print "    interesting mom"
0113                     # exclude extra x from p -> p + x
0114                     if not any(mom.daughter(j2).pdgId() == mom.pdgId() for j2 in range(mom.numberOfDaughters())):
0115                         #print "         pass no-self-decay"
0116                         ok = True
0117                     # Account for generator feature with Higgs decaying to itself with same four-vector but no daughters
0118                     elif mom.pdgId() == 25 and any(mom.daughter(j2).pdgId() == 25 and mom.daughter(j2).numberOfDaughters()==0 for j2 in range(mom.numberOfDaughters())):
0119                         ok = True
0120                 if abs(mom.pdgId()) == 15:
0121                     # if we're a tau daughter we're status 2
0122                     # if we passed all the previous steps, then we're a prompt lepton
0123                     # so we'd like to be included
0124                     ok = True
0125                 if not ok and p.pt() > 10 and id in [1,2,3,4,5,21,22] and any(interestingPdgId(d.pdgId()) for d in realGenDaughters(mom)):
0126                     # interesting for being a parton brother of an interesting particle (to get the extra jets in ME+PS)
0127                     ok = True
0128                 if not ok and id in [11, 13, 15] and abs(mom.pdgId()) in [1,2,3,4,5,21]:
0129                     # Lepton e.g. in off-shell DY with the mother being one of the incoming partons
0130                     ok = True
0131             if ok:
0132                 gp = p
0133                 gp.rawIndex = rawIndex # remember its index, so that we can set the mother index later
0134                 keymap[rawIndex] = len(good)
0135                 good.append(gp)
0136         # connect mother links
0137         for igp,gp in enumerate(good):
0138             gp.motherIndex = -1
0139             gp.sourceId    = 99
0140             gp.genSummaryIndex = igp
0141             ancestor = None if gp.numberOfMothers() == 0 else gp.motherRef(0)
0142             while ancestor != None and ancestor.isNonnull():
0143                 if ancestor.key() in keymap:
0144                     gp.motherIndex = keymap[ancestor.key()]
0145                     if ancestor.pdgId() != good[gp.motherIndex].pdgId():
0146                         print("Error keying %d: motherIndex %d, ancestor.pdgId %d, good[gp.motherIndex].pdgId() %d " % (igp, gp.motherIndex, ancestor.pdgId(),  good[gp.motherIndex].pdgId()))
0147                     break
0148                 ancestor = None if ancestor.numberOfMothers() == 0 else ancestor.motherRef(0)
0149             if abs(gp.pdgId()) not in [1,2,3,4,5,11,12,13,14,15,16,21]:
0150                 gp.sourceId = gp.pdgId()
0151             if gp.motherIndex != -1:
0152                 ancestor = good[gp.motherIndex]
0153                 if ancestor.sourceId != 99 and (ancestor.mass() > gp.mass() or gp.sourceId == 99):
0154                     gp.sourceId = ancestor.sourceId
0155         event.generatorSummary = good
0156         # add the ID of the mother to be able to recreate last decay chains
0157         for ip,p in enumerate(good):
0158             moms = realGenMothers(p)
0159             if len(moms)==0:
0160                 p.motherId = 0
0161                 p.grandmotherId = 0
0162             elif len(moms)==1:
0163                 p.motherId = moms[0].pdgId()
0164                 gmoms = realGenMothers(moms[0])
0165                 p.grandmotherId = (gmoms[0].pdgId() if len(gmoms)==1 else (0 if len(gmoms)==0 else -9999))
0166             else:
0167                 #print "    unclear what mothers to give to this particle, among ","  ".join("%d[%d]" % (m.pdgId(),m.status()) for m in moms)
0168                 p.motherId = -9999
0169                 p.grandmotherId = -9999
0170             if verbose:
0171                 print("%3d  {%6d}: %+8d  %3d :  %8.2f   %+5.2f   %+5.2f : %d %2d : %+8d {%3d}: %s" % ( ip,p.rawIndex,
0172                         p.pdgId(), p.status(), p.pt(), p.eta(), p.phi(), len(moms), p.numberOfDaughters(),
0173                         p.motherId, p.motherIndex,
0174                         "  ".join("%d[%d]" % (p.daughter(i).pdgId(), p.daughter(i).status()) for i in range(p.numberOfDaughters()))))
0175         if verbose:
0176             print("\n\n")
0177
0178         if self.makeAllGenParticles:
0179             event.genParticles = allGenParticles
0180
0181         if self.makeSplittedGenLists:
0182             event.genHiggsBosons = []
0183             event.genVBosons     = []
0184             event.gennus         = []
0185             event.gennusFromTop  = []
0186             event.genleps        = []
0187             event.gentauleps     = []
0188             event.gentaus        = []
0189             event.gentopquarks   = []
0190             event.genbquarks     = []
0191             event.genwzquarks    = []
0192             event.genbquarksFromTop = []
0193             event.genbquarksFromH   = []
0194             event.genlepsFromTop = []
0195             for p in event.generatorSummary:
0196                 id = abs(p.pdgId())
0197                 if id == 25:
0198                     event.genHiggsBosons.append(p)
0199                 elif id in {23,24}:
0200                     event.genVBosons.append(p)
0201                 elif id in {12,14,16}:
0202                     event.gennus.append(p)
0203
0204                     momids = [(m, abs(m.pdgId())) for m in realGenMothers(p)]
0205
0206                     #have a look at the lepton mothers
0207                     for mom, momid in momids:
0208                         #lepton from W
0209                         if momid == 24:
0210                             wmomids = [abs(m.pdgId()) for m in realGenMothers(mom)]
0211                             #W from t
0212                             if 6 in wmomids:
0213                                 #save mu,e from t->W->mu/e
0214                                 event.gennusFromTop.append(p)
0215
0216                 elif id in {11,13}:
0217                     #taus to separate vector
0218                     if abs(p.motherId) == 15:
0219                         event.gentauleps.append(p)
0220                     #all muons and electrons
0221                     else:
0222                         event.genleps.append(p)
0223                         momids = [(m, abs(m.pdgId())) for m in realGenMothers(p)]
0224
0225                         #have a look at the lepton mothers
0226                         for mom, momid in momids:
0227                             #lepton from W
0228                             if momid == 24:
0229                                 wmomids = [abs(m.pdgId()) for m in realGenMothers(mom)]
0230                                 #W from t
0231                                 if 6 in wmomids:
0232                                     #save mu,e from t->W->mu/e
0233                                     event.genlepsFromTop.append(p)
0234                 elif id == 15:
0235                     if self.allGenTaus or not any([abs(d.pdgId()) in {11,13} for d in realGenDaughters(p)]):
0236                         event.gentaus.append(p)
0237                 elif id == 6:
0238                     event.gentopquarks.append(p)
0239                 elif id == 5:
0240                     event.genbquarks.append(p)
0241                     momids = [abs(m.pdgId()) for m in realGenMothers(p)]
0242                     if  6 in momids: event.genbquarksFromTop.append(p)
0243                     if 25 in momids: event.genbquarksFromH.append(p)
0244                 if id <= 5 and any([abs(m.pdgId()) in {23,24} for m in realGenMothers(p)]):
0245                     event.genwzquarks.append(p)
0246
0247     def process(self, event):
0248         self.readCollections( event.input )
0249
0250         # if not MC, nothing to do
0251         if not self.cfg_comp.isMC:
0252             return True
0253         # do MC level analysis
0254         self.makeMCInfo(event)
0255         return True
0256
0257 import PhysicsTools.HeppyCore.framework.config as cfg
0258 setattr(GeneratorAnalyzer,"defaultConfig",
0259     cfg.Analyzer(GeneratorAnalyzer,
0260         # BSM particles that can appear with status <= 2 and should be kept
0261         stableBSMParticleIds = [ 1000022 ],
0262         # Particles of which we want to save the pre-FSR momentum (a la status 3).
0263         # Note that for quarks and gluons the post-FSR doesn't make sense,
0264         # so those should always be in the list
0265         savePreFSRParticleIds = [ 1,2,3,4,5, 11,12,13,14,15,16, 21 ],
0266         # Make also the list of all genParticles, for other analyzers to handle
0267         makeAllGenParticles = True,
0268         # Make also the splitted lists
0269         makeSplittedGenLists = True,
0270         allGenTaus = False,
0271         # Print out debug information
0272         verbose = False,
0273     )
0274 )