analyzers/gen/GeneratorAnalyzer.py

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