---

jetTools.py

Go to the documentation of this file.

import FWCore.ParameterSet.Config as cms

from PhysicsTools.PatAlgos.tools.helpers import *

def switchJECSet(process,newName,oldName):
    """Replace tags in the JetCorrFactorsProducer
       For end-users, taking the process as first argument"""
    switchJECSet_(process.jetCorrFactors,newName,oldName)
    
def switchJECSet_(jetCorrFactors,newName,oldName,
                  steps=['L1Offset', 'L2Relative', 'L3Absolute', 'L4EMF', 'L5Flavor', 'L6UE', 'L7Parton']):
    """Replace tags in the JetCorrFactorsProducer
       Inner implementation taking a jetCorrFactors object"""
    found = False
    for k in steps:
        vv = getattr(jetCorrFactors, k).value();
        if vv.find(oldName) != -1: found = True
        if (vv != "none"): 
            setattr(jetCorrFactors, k, vv.replace(oldName,newName))
            # the first replace is good for L2, L3; the last for L7 (which don't have type dependency, at least not in the name)
    if not found:
        raise RuntimeError, """None of the JEC steps uses old name %s, so I can't replace %s with %s.
                               The full configuration is %s""" % (newName,oldName,newName,jetCorrFactors.dumpPython())

def switchJECParameters(jetCorrFactors,newalgo,newtype="Calo",oldalgo="IC5",oldtype="Calo"):
    """Replace tags in the JetCorrFactorsProducer -- L5Flavor is taken out as it is said not to be dependend on the specific jet algorithm"""
    for k in ['L1Offset', 'L2Relative', 'L3Absolute', 'L4EMF', 'L6UE', 'L7Parton']:
        vv = getattr(jetCorrFactors, k).value();
        if (vv != "none"): 
            setattr(jetCorrFactors, k, vv.replace(oldalgo+oldtype,newalgo+newtype).replace(oldalgo,newalgo) )
            # the first replace is good for L2, L3; the last for L7 (which don't have type dependency, at least not in the name)

def runBTagging(process,jetCollection,label) :
    """Define a sequence to run BTagging on AOD on top of jet collection 'jetCollection', appending 'label' to module labels.
       The sequence will be called "btaggingAOD" + 'label', and will already be added to the process (but not to any Path)
       The sequence will include a JetTracksAssociatorAtVertex with name "jetTracksAssociatorAtVertex" + 'label'
       The method will return a pair (sequence, labels) where 'sequence' is the cms.Sequence object, and 'labels' contains
         labels["jta"]      = the name of the JetTrackAssociator module
         labels["tagInfos"] = list of names of TagInfo modules
         labels["jetTags "] = list of names of JetTag modules
       these labels are meant to be used for PAT BTagging tools
       NOTE: 'label' MUST NOT BE EMPTY
     """
    if (label == ''):
        raise ValueError, "Label for re-running BTagging can't be empty, it will crash CRAB." 
    process.load("RecoJets.JetAssociationProducers.ic5JetTracksAssociatorAtVertex_cfi")
    process.load("RecoBTag.Configuration.RecoBTag_cff")
    from RecoJets.JetAssociationProducers.ic5JetTracksAssociatorAtVertex_cfi import ic5JetTracksAssociatorAtVertex
    import RecoBTag.Configuration.RecoBTag_cff as btag
    
    # quickly make VInputTag from strings
    def vit(*args) : return cms.VInputTag( *[ cms.InputTag(x) for x in args ] )
    
    # define labels
    jtaLabel =  'jetTracksAssociatorAtVertex' + label
    ipTILabel = 'impactParameterTagInfos'     + label
    svTILabel = 'secondaryVertexTagInfos'     + label
    seTILabel = 'softElectronTagInfos'        + label
    smTILabel = 'softMuonTagInfos'            + label
    
    # make JTA and TagInfos
    setattr( process, jtaLabel,  ic5JetTracksAssociatorAtVertex.clone(jets = jetCollection))
    setattr( process, ipTILabel, btag.impactParameterTagInfos.clone(jetTracks = cms.InputTag(jtaLabel)) )
    setattr( process, svTILabel, btag.secondaryVertexTagInfos.clone(trackIPTagInfos = cms.InputTag(ipTILabel)) )
    setattr( process, seTILabel, btag.softElectronTagInfos.clone(jets = jetCollection) )
    setattr( process, smTILabel, btag.softMuonTagInfos.clone(jets = jetCollection) )
    setattr( process, 'jetBProbabilityBJetTags'+label, btag.jetBProbabilityBJetTags.clone(tagInfos = vit(ipTILabel)) )
    setattr( process, 'jetProbabilityBJetTags' +label,  btag.jetProbabilityBJetTags.clone(tagInfos = vit(ipTILabel)) )
    setattr( process, 'trackCountingHighPurBJetTags'+label, btag.trackCountingHighPurBJetTags.clone(tagInfos = vit(ipTILabel)) )
    setattr( process, 'trackCountingHighEffBJetTags'+label, btag.trackCountingHighEffBJetTags.clone(tagInfos = vit(ipTILabel)) )
    setattr( process, 'impactParameterMVABJetTags'+label, btag.impactParameterMVABJetTags.clone(tagInfos = vit(ipTILabel)) )
    setattr( process, 'simpleSecondaryVertexBJetTags'+label, btag.simpleSecondaryVertexBJetTags.clone(tagInfos = vit(svTILabel)) )
    setattr( process, 'combinedSecondaryVertexBJetTags'+label, btag.combinedSecondaryVertexBJetTags.clone(tagInfos = vit(ipTILabel, svTILabel)) )
    setattr( process, 'combinedSecondaryVertexMVABJetTags'+label, btag.combinedSecondaryVertexMVABJetTags.clone(tagInfos = vit(ipTILabel, svTILabel)) )
    setattr( process, 'softElectronBJetTags'+label, btag.softElectronBJetTags.clone(tagInfos = vit(seTILabel)) )
    setattr( process, 'softMuonBJetTags'+label, btag.softMuonBJetTags.clone(tagInfos = vit(smTILabel)) )
    setattr( process, 'softMuonNoIPBJetTags'+label, btag.softMuonNoIPBJetTags.clone(tagInfos = vit(smTILabel)) )
    
    def mkseq(process, firstlabel, *otherlabels):
       seq = getattr(process, firstlabel)
       for x in otherlabels: seq += getattr(process, x)
       return cms.Sequence(seq)
    
    labels = { 'jta' : jtaLabel, 
               'tagInfos' : (ipTILabel,svTILabel,seTILabel,smTILabel), 
               'jetTags'  : [ (x + label) for x in ('jetBProbabilityBJetTags',
                                                'jetProbabilityBJetTags',
                                                'trackCountingHighPurBJetTags',
                                                'trackCountingHighEffBJetTags',
                                                'impactParameterMVABJetTags',
                                                'simpleSecondaryVertexBJetTags',
                                                'combinedSecondaryVertexBJetTags',
                                                'combinedSecondaryVertexMVABJetTags',
                                                'softElectronBJetTags',
                                                'softMuonBJetTags',
                                                'softMuonNoIPBJetTags') ]
    }
    
    setattr( process, 'btaggingTagInfos' + label, mkseq(process, *(labels['tagInfos']) ) )
    setattr( process, 'btaggingJetTags' + label,  mkseq(process, *(labels['jetTags'])  ) )
    seq = mkseq(process, jtaLabel, 'btaggingTagInfos' + label, 'btaggingJetTags' + label) 
    setattr( process, 'btagging' + label, seq )
    return (seq, labels)

def switchJetCollection(process,jetCollection,doJTA=True,doBTagging=True,jetCorrLabel=None,doType1MET=True,
                                genJetCollection=cms.InputTag("iterativeCone5GenJets")):
    """Switch the collection of jets in PAT from the default value.
          doBTagging  : True to run the BTagging sequence on top of this jets, and import it into PAT.
          doJTA       : Run Jet Tracks Association and Jet Charge (will be forced to True if doBTagging is true)
          jetCorrLabel: Name of the algorithm and jet type JEC to pick corrections from, or None for no JEC 
                        Examples are ('IC5','Calo'), ('SC7','Calo'), ('KT4','PF')
                        It tries to find a 'L2L3JetCorrector' + algo + type , or otherwise to create if as a 
                        JetCorrectionServiceChain of 'L2RelativeJetCorrector' and 'L3AbsoluteJetCorrector'
          doType1MET  : If jetCorrLabel is not 'None', set this to 'True' to remake Type1 MET from these jets
                        NOTE: at the moment it must be False for non-CaloJets otherwise the JetMET POG module crashes.
          genJetCollection : GenJet collection to match to."""
    oldLabel = process.allLayer1Jets.jetSource;
    process.jetPartonMatch.src        = jetCollection
    process.jetGenJetMatch.src        = jetCollection
    process.jetGenJetMatch.matched    = genJetCollection
    process.jetPartonAssociation.jets = jetCollection
    massSearchReplaceParam(process.patTrigMatch, 'src', oldLabel, jetCollection)
    process.allLayer1Jets.jetSource = jetCollection
    # quickly make VInputTag from strings
    def vit(*args) : return cms.VInputTag( *[ cms.InputTag(x) for x in args ] )
    if doBTagging :
        (btagSeq, btagLabels) = runBTagging(process, jetCollection, 'AOD')
        process.patAODCoreReco += btagSeq # must add to Core, as it's needed by ExtraReco
        process.patJetCharge.src                     = btagLabels['jta']
        process.allLayer1Jets.trackAssociationSource = btagLabels['jta']
        process.allLayer1Jets.tagInfoSources       = cms.VInputTag( *[ cms.InputTag(x) for x in btagLabels['tagInfos'] ] )
        process.allLayer1Jets.discriminatorSources = cms.VInputTag( *[ cms.InputTag(x) for x in btagLabels['jetTags']  ] )
    else:
        process.patAODReco.remove(process.patBTagging)
        process.allLayer1Jets.addBTagInfo = False
    if doJTA or doBTagging:
        if not doBTagging:
            process.load("RecoJets.JetAssociationProducers.ic5JetTracksAssociatorAtVertex_cfi")
            from RecoJets.JetAssociationProducers.ic5JetTracksAssociatorAtVertex_cfi import ic5JetTracksAssociatorAtVertex
            process.jetTracksAssociatorAtVertex = ic5JetTracksAssociatorAtVertex.clone(jets = jetCollection)
            process.patAODReco.replace(process.patJetTracksCharge, process.jetTracksAssociatorAtVertex + process.patJetTracksCharge)
            process.patJetCharge.src                     = 'jetTracksAssociatorAtVertex'
            process.allLayer1Jets.trackAssociationSource = 'jetTracksAssociatorAtVertex'
    else: ## no JTA
        process.patAODReco.remove(process.patJetTracksCharge)
        process.allLayer1Jets.addAssociatedTracks = False
        process.allLayer1Jets.addJetCharge = False
    if jetCorrLabel != None:
        if jetCorrLabel == False : raise ValueError, "In switchJetCollection 'jetCorrLabel' must be set to None, not False"
        if jetCorrLabel == "None": raise ValueError, "In switchJetCollection 'jetCorrLabel' must be set to None (without quotes), not 'None'"
        if type(jetCorrLabel) != type(('IC5','Calo')): 
            raise ValueError, "In switchJetCollection 'jetCorrLabel' must be None, or a tuple ('Algo', 'Type')"
        if not hasattr( process, 'L2L3JetCorrector%s%s' % jetCorrLabel ):
            setattr( process, 
                        'L2L3JetCorrector%s%s' % jetCorrLabel, 
                        cms.ESSource("JetCorrectionServiceChain",
                            correctors = cms.vstring('L2RelativeJetCorrector%s%s' % jetCorrLabel,
                                                     'L3AbsoluteJetCorrector%s%s' % jetCorrLabel),
                            label      = cms.string('L2L3JetCorrector%s%s' % jetCorrLabel)
                        )
                    )
        switchJECParameters(process.jetCorrFactors, jetCorrLabel[0], jetCorrLabel[1], oldalgo='IC5',oldtype='Calo')
        process.jetCorrFactors.jetSource = jetCollection
        if doType1MET:
            process.corMetType1Icone5.inputUncorJetsLabel = jetCollection.value() # FIXME it's corMetType1Icone5 that's broken
            process.corMetType1Icone5.corrector           = 'L2L3JetCorrector%s%s' % jetCorrLabel
    else:
        process.patJetMETCorrections.remove(process.jetCorrFactors)
        process.allLayer1Jets.addJetCorrFactors = False
    ## Add this to the summary tables (not strictly needed, but useful)
    if oldLabel in process.aodSummary.candidates: 
        process.aodSummary.candidates[process.aodSummary.candidates.index(oldLabel)] = jetCollection
    else:
        process.aodSummary.candidates += [jetCollection]
        

def addJetCollection(process,jetCollection,postfixLabel,
                        doJTA=True,doBTagging=True,jetCorrLabel=None,doType1MET=True,doL1Counters=False,
                        genJetCollection=cms.InputTag("iterativeCone5GenJets"),doTrigMatch=False):
    """Add a new collection of jets in PAT from the default value.
          postfixLabel: Postpone this label to the name of all modules that work with these jet collection.
                        it can't be an empty string
          doBTagging  : True to run the BTagging sequence on top of this jets, and import it into PAT.
          doJTA       : Run Jet Tracks Association and Jet Charge (will be forced to True if doBTagging is true)
          jetCorrLabel: Name of the algorithm and jet type JEC to pick corrections from, or None for no JEC 
                        Examples are ('IC5','Calo'), ('SC7','Calo'), ('KT4','PF')
                        It tries to find a 'L2L3JetCorrector' + algo + type , or otherwise to create if as a 
                        JetCorrectionServiceChain of 'L2RelativeJetCorrector' and 'L3AbsoluteJetCorrector'
          doType1MET  : Make also a new MET (NOT IMPLEMENTED)
          doL1Counters: copy also the filter modules that accept/reject the event looking at the number of jets
          genJetCollection : GenJet collection to match to.

       Note: This takes the configuration from the already-configured jets, so if you do 
             replaces before calling addJetCollection then they will affect also the new jets
    """
    def addAlso (label,value):
        existing = getattr(process, label)
        setattr( process, label + postfixLabel, value)
        process.patDefaultSequence.replace( existing, existing * value )
    def addClone(label,**replaceStatements):
        new      = getattr(process, label).clone(**replaceStatements)
        addAlso(label, new)
    addClone('allLayer1Jets', jetSource = jetCollection, addTrigMatch = doTrigMatch)
    l1Jets = getattr(process, 'allLayer1Jets'+postfixLabel)
    addClone('selectedLayer1Jets', src=cms.InputTag('allLayer1Jets'+postfixLabel))
    addClone('cleanLayer1Jets', src=cms.InputTag('selectedLayer1Jets'+postfixLabel))
    if doL1Counters:
        addClone('countLayer1Jets', src=cms.InputTag('cleanLayer1Jets'+postfixLabel))
    addClone('jetPartonMatch',       src = jetCollection)
    addClone('jetGenJetMatch',       src = jetCollection, matched = genJetCollection)
    addClone('jetPartonAssociation', jets = jetCollection)
    addClone('jetFlavourAssociation',srcByReference = cms.InputTag('jetPartonAssociation' + postfixLabel))
    def fixInputTag(x): x.setModuleLabel(x.moduleLabel+postfixLabel)
    def fixVInputTag(x): x[0].setModuleLabel(x[0].moduleLabel+postfixLabel)
    fixInputTag(l1Jets.JetPartonMapSource)
    fixInputTag(l1Jets.genJetMatch)
    fixInputTag(l1Jets.genPartonMatch)
    triggers = MassSearchParamVisitor('src', process.allLayer1Jets.jetSource)
    process.patTrigMatch.visit(triggers)
    for mod in triggers.modules():
        if doTrigMatch:
            newmod = mod.clone(src = jetCollection)
            setattr( process, mod.label() + postfixLabel, newmod )
            process.patTrigMatch.replace( mod, mod * newmod )
    for it in l1Jets.trigPrimMatch.value(): fixInputTag(it)
    def vit(*args) : return cms.VInputTag( *[ cms.InputTag(x) for x in args ] )
    if doBTagging :
        (btagSeq, btagLabels) = runBTagging(process, jetCollection, postfixLabel)
        process.patAODCoreReco += btagSeq  # must add to Core, as it's needed by Extra
        addClone('patJetCharge', src=cms.InputTag(btagLabels['jta']))
        l1Jets.trackAssociationSource = cms.InputTag(btagLabels['jta'])
        l1Jets.tagInfoSources         = cms.VInputTag( *[ cms.InputTag(x) for x in btagLabels['tagInfos'] ] )
        l1Jets.discriminatorSources   = cms.VInputTag( *[ cms.InputTag(x) for x in btagLabels['jetTags']  ] )
        fixInputTag(l1Jets.jetChargeSource)
    else:
       l1Jets.addBTagInfo = False 
    if doJTA or doBTagging:
        if not doBTagging:
            process.load("RecoJets.JetAssociationProducers.ic5JetTracksAssociatorAtVertex_cfi")
            from RecoJets.JetAssociationProducers.ic5JetTracksAssociatorAtVertex_cfi import ic5JetTracksAssociatorAtVertex
            jtaLabel = 'jetTracksAssociatorAtVertex' + postfixLabel
            setattr( process, jtaLabel, ic5JetTracksAssociatorAtVertex.clone(jets = jetCollection) )
            process.patAODReco.replace(process.patJetTracksCharge, getattr(process,jtaLabel) + process.patJetTracksCharge)
            l1Jets.trackAssociationSource = cms.InputTag(jtaLabel)
            addClone('patJetCharge', src=cms.InputTag(jtaLabel)),
            fixInputTag(l1Jets.jetChargeSource)
    else: ## no JTA
        l1Jets.addAssociatedTracks = False
        l1Jets.addJetCharge = False
    if jetCorrLabel != None:
        if jetCorrLabel == False : raise ValueError, "In addJetCollection 'jetCorrLabel' must be set to None, not False"
        if jetCorrLabel == "None": raise ValueError, "In addJetCollection 'jetCorrLabel' must be set to None (without quotes), not 'None'"
        if type(jetCorrLabel) != type(('IC5','Calo')): 
            raise ValueError, "In switchJetCollection 'jetCorrLabel' must be None, or a tuple ('Algo', 'Type')"
        if not hasattr( process, 'L2L3JetCorrector%s%s' % jetCorrLabel ):
            setattr( process, 
                        'L2L3JetCorrector%s%s' % jetCorrLabel, 
                        cms.ESSource("JetCorrectionServiceChain",
                            correctors = cms.vstring('L2RelativeJetCorrector%s%s' % jetCorrLabel,
                                                     'L3AbsoluteJetCorrector%s%s' % jetCorrLabel),
                            label      = cms.string('L2L3JetCorrector%s%s' % jetCorrLabel)
                        )
                    )
        addClone('jetCorrFactors',       jetSource           = jetCollection, 
                                         defaultJetCorrector = cms.string('L2L3JetCorrector%s%s' % jetCorrLabel))
        switchJECParameters( getattr(process,'jetCorrFactors'+postfixLabel), jetCorrLabel[0], jetCorrLabel[1], oldalgo='IC5',oldtype='Calo' )
        fixVInputTag(l1Jets.jetCorrFactorsSource)
        if doType1MET:
            addClone('corMetType1Icone5', inputUncorJetsLabel = jetCollection.value(),
                                          corrector = cms.string('L2L3JetCorrector%s%s' % jetCorrLabel))
            addClone('corMetType1Icone5Muons', uncorMETInputTag = cms.InputTag("corMetType1Icone5"+postfixLabel))
            addClone('layer1METs',              metSource = cms.InputTag("corMetType1Icone5Muons"+postfixLabel), addTrigMatch = doTrigMatch)
            l1MET = getattr(process, 'layer1METs'+postfixLabel)
            mettriggers = MassSearchParamVisitor('src', process.layer1METs.metSource)
            process.patTrigMatch.visit(mettriggers)
            for mod in mettriggers.modules():
                if doTrigMatch:
                    newmod = mod.clone(src = l1MET.metSource)
                    setattr( process, mod.label() + postfixLabel, newmod )
                    process.patTrigMatch.replace( mod, mod * newmod )
            for it in l1MET.trigPrimMatch.value(): fixInputTag(it)
            process.allLayer1Summary.candidates += [ cms.InputTag('layer1METs'+postfixLabel) ]
    else:
        l1Jets.addJetCorrFactors = False
    ## Add this to the summary tables (not strictly needed, but useful)
    if jetCollection not in process.aodSummary.candidates: 
        process.aodSummary.candidates += [ jetCollection ]
    process.allLayer1Summary.candidates      += [ cms.InputTag('allLayer1Jets'+postfixLabel) ]
    process.selectedLayer1Summary.candidates += [ cms.InputTag('selectedLayer1Jets'+postfixLabel) ]

---