python/Tools/pfIsolation.py

0001 # Colin, March 2012
0002
0003 import FWCore.ParameterSet.Config as cms
0004 from PhysicsTools.PatAlgos.tools.helpers import cloneProcessingSnippet
0005
0006 def _getattrGenerator( process, postfix ):
0007     '''A function generator to simplify the getattr syntax'''
0008     def fun(name):
0009         return getattr(process, name+postfix)
0010     return fun
0011
0012
0013 _PFBRECO_loaded = False
0014
0015 def _loadPFBRECO(process):
0016     '''The particle-flow based reconstruction sequence should be loaded once in the process.
0017     Not optimal, should load it only if it is not detected (hasattr)'''
0018     global _PFBRECO_loaded
0019     if _PFBRECO_loaded is False:
0020         _PFBRECO_loaded = True
0021         process.load("CommonTools.ParticleFlow.PFBRECO_cff")
0022
0023
0024 def setupPFIso(process, leptonCollection, particleName, newpostfix='PFIso', postfix='', runPF2PAT = False):
0025     '''Generic function to setup particle-based isolation for a given lepton collection.
0026     Returns the isolation sequence.
0027     You are responsible for adding it to your path.
0028
0029     leptonCollection could e.g. be "gsfElectrons" or "muons"
0030     particleName must be either "Electron" or "Muon".
0031     newpostfix can be specified to define several particle-flow isolation sequences
0032     '''
0033     lepshort = None
0034     if particleName=='Electron':
0035         lepshort='el'
0036     elif particleName=='Muon':
0037         lepshort='mu'
0038     else:
0039         raise ValueError('particleName should be equal to "Electron" or "Muon"')
0040
0041     if runPF2PAT != True :
0042       _loadPFBRECO(process)
0043
0044     # ADD VETOES IN ENDCAPS!
0045     fullpostfix = postfix+newpostfix
0046     ga = _getattrGenerator( process, postfix )
0047     ganew = _getattrGenerator( process, fullpostfix )
0048
0049     leptonSeq = cms.Sequence(
0050         ga('pf{lepton}IsolationSequence'.format(lepton=particleName))
0051         )
0052     setattr( process, 'std{lepton}Sequence{postfix}'.format(lepton=particleName,
0053                                                    postfix=postfix), leptonSeq)
0054
0055     leptonSource = leptonCollection
0056     cloneProcessingSnippet(process,
0057                            ga('std{lepton}Sequence'.format(lepton=particleName)),
0058                            newpostfix)
0059
0060     ganew("{lepshort}PFIsoDepositCharged".format(lepshort=lepshort) ).src = leptonSource
0061     ganew("{lepshort}PFIsoDepositChargedAll".format(lepshort=lepshort)).src = leptonSource
0062     ganew("{lepshort}PFIsoDepositNeutral".format(lepshort=lepshort)).src = leptonSource
0063     ganew("{lepshort}PFIsoDepositGamma".format(lepshort=lepshort)).src = leptonSource
0064     ganew("{lepshort}PFIsoDepositPU".format(lepshort=lepshort)).src = leptonSource
0065
0066     return ganew('std{lepton}Sequence'.format(lepton=particleName))
0067
0068 def setupPFIsoPhoton(process, photonCollection, particleName, newpostfix='PFIso'):
0069     '''Generic function to setup particle-based isolation for a given lepton collection.
0070     Returns the isolation sequence.
0071     You are responsible for adding it to your path.
0072
0073     leptonCollection could e.g. be "gsfElectrons" or "muons"
0074     particleName must be either "Electron" or "Muon".
0075     newpostfix can be specified to define several particle-flow isolation sequences
0076     '''
0077     phoshort = None
0078     if particleName=='Photon':
0079         phoshort='ph'
0080     else:
0081         raise ValueError('particleName should be equal to "Photon"')
0082
0083     _loadPFBRECO(process)
0084
0085     postfix = ''
0086     # ADD VETOES IN ENDCAPS!
0087     fullpostfix = postfix+newpostfix
0088     #fullpostfix = ''
0089     ga = _getattrGenerator( process, postfix )
0090     ganew = _getattrGenerator( process, fullpostfix )
0091
0092     photonSeq = cms.Sequence(
0093         ga('pf{photon}IsolationSequence'.format(photon=particleName))
0094         )
0095     setattr( process, 'std{photon}Sequence{postfix}'.format(photon=particleName,
0096                                                    postfix=postfix), photonSeq)
0097
0098     photonSource = photonCollection
0099     cloneProcessingSnippet(process,
0100                            ga('std{photon}Sequence'.format(photon=particleName)),
0101                            newpostfix)
0102
0103     ganew("{phoshort}PFIsoDepositCharged".format(phoshort=phoshort) ).src = photonSource
0104     ganew("{phoshort}PFIsoDepositChargedAll".format(phoshort=phoshort)).src = photonSource
0105     ganew("{phoshort}PFIsoDepositNeutral".format(phoshort=phoshort)).src = photonSource
0106     ganew("{phoshort}PFIsoDepositGamma".format(phoshort=phoshort)).src = photonSource
0107     ganew("{phoshort}PFIsoDepositPU".format(phoshort=phoshort)).src = photonSource
0108
0109     return ganew('std{photon}Sequence'.format(photon=particleName))
0110
0111
0112 def setupPFMuonIso(process, muonCollection, postfix='PFIso' ):
0113     '''Set up particle-based isolation for the muons in muonCollection.
0114
0115     Calls setupPFIso.
0116     '''
0117     return setupPFIso( process, muonCollection, 'Muon', postfix)
0118
0119
0120
0121 def setupPFElectronIso(process, electronCollection, newpostfix='PFIso', postfix='', runPF2PAT = False ):
0122     '''Set up particle-based isolation for the electrons in electronCollection.
0123
0124     Calls setupPFIso.
0125     '''
0126     #    print 'WARNING!!! the vetoes are the ones defined for the PF e-s (no veto...).'
0127     #    print 'Vetoes should be applied in the endcaps when doing particle-based isolation on gsfElectrons.'
0128     #    print 'Need a volunteer to implement that.'
0129     return setupPFIso( process, electronCollection, 'Electron', newpostfix, postfix, runPF2PAT)
0130
0131
0132 def setupPFPhotonIso(process, photonCollection, postfix='PFIso' ):
0133     '''Set up particle-based isolation for the electrons in electronCollection.
0134
0135     Calls setupPFIsoPhoton.
0136     '''
0137     #    print 'WARNING!!! the vetoes are the ones defined for the PF e-s (no veto...).'
0138     #    print 'Please make sure that your file with vetoes is up to date'
0139     return setupPFIsoPhoton( process, photonCollection, 'Photon', postfix)
0140