python/tools/vid_id_tools.py

0001 import FWCore.ParameterSet.Config as cms
0002
0003 from PhysicsTools.SelectorUtils.centralIDRegistry import central_id_registry
0004
0005 import sys
0006
0007 #keep python2.6 compatibility for computing
0008 #import importlib
0009
0010 #general simple tools for various object types
0011 def setupVIDSelection(vidproducer,cutflow):
0012     if not hasattr(cutflow,'idName'):
0013         raise Exception('InvalidVIDCutFlow', 'The cutflow configuation provided is malformed and does not have a specified name!')
0014     if not hasattr(cutflow,'cutFlow'):
0015         raise Exception('InvalidVIDCutFlow', 'The cutflow configuration provided is malformed and does not have a specific cutflow!')
0016     cutflow_md5 = central_id_registry.getMD5FromName(cutflow.idName)
0017     isPOGApproved = False
0018     if hasattr(cutflow,'isPOGApproved'):
0019         isPOGApproved = cutflow.isPOGApproved.value()
0020     vidproducer.physicsObjectIDs.append(
0021         cms.PSet( idDefinition = cutflow,
0022                   isPOGApproved = cms.untracked.bool(isPOGApproved),
0023                   idMD5 = cms.string(cutflow_md5) )
0024     )
0025 #    sys.stderr.write('Added ID \'%s\' to %s\n'%(cutflow.idName.value(),vidproducer.label()))
0026
0027 def addVIDSelectionToPATProducer(patProducer,idProducer,idName,addUserData=True):
0028     patProducerIDs = None
0029     userDatas = None
0030     for key in patProducer.__dict__.keys():
0031         if 'IDSources' in key:
0032             patProducerIDs = getattr(patProducer,key)
0033         if 'userData' in key:
0034             userDatas = getattr(patProducer,key)
0035     if patProducerIDs is None:
0036         raise Exception('StrangePatModule','%s does not have ID sources!'%patProducer.label())
0037     if userDatas is None:
0038         raise Exception('StrangePatModule','%s does not have UserData sources!'%patProducer.label())
0039     setattr(patProducerIDs,idName,cms.InputTag('%s:%s'%(idProducer,idName)))
0040     if( addUserData ):
0041         if( len(userDatas.userClasses.src) == 1 and
0042             isinstance(userDatas.userClasses.src[0], str) and
0043             userDatas.userClasses.src[0] == ''            ):
0044             userDatas.userClasses.src = cms.VInputTag(cms.InputTag('%s:%s'%(idProducer,idName)))
0045         else:
0046             userDatas.userClasses.src.append(cms.InputTag('%s:%s'%(idProducer,idName)))
0047         sys.stderr.write('\t--- %s:%s added to %s\n'%(idProducer,idName,patProducer.label()))
0048
0049 def setupAllVIDIdsInModule(process,id_module_name,setupFunction,patProducer=None,addUserData=True,task=None):
0050 #    idmod = importlib.import_module(id_module_name)
0051     idmod= __import__(id_module_name, globals(), locals(), ['idName','cutFlow'])
0052     for name in dir(idmod):
0053         item = getattr(idmod,name)
0054         if hasattr(item,'idName') and hasattr(item,'cutFlow'):
0055             setupFunction(process,item,patProducer,addUserData,task)
0056
0057 # Supported data formats defined via "enum"
0058 from PhysicsTools.SelectorUtils.tools.DataFormat import DataFormat
0059
0060 ####
0061 # Electrons
0062 ####
0063
0064 #turns on the VID electron ID producer, possibly with extra options
0065 # for PAT and/or MINIAOD
0066 def switchOnVIDElectronIdProducer(process, dataFormat, task=None):
0067     process.load('RecoEgamma.ElectronIdentification.egmGsfElectronIDs_cff')
0068     if task is not None:
0069         task.add(process.egmGsfElectronIDTask)
0070     #*always* reset to an empty configuration
0071     if( len(process.egmGsfElectronIDs.physicsObjectIDs) > 0 ):
0072         process.egmGsfElectronIDs.physicsObjectIDs = cms.VPSet()
0073     dataFormatString = "Undefined"
0074     if dataFormat == DataFormat.AOD:
0075         # No reconfiguration is required, default settings are for AOD
0076         dataFormatString = "AOD"
0077     elif dataFormat == DataFormat.MiniAOD:
0078         # If we are dealing with MiniAOD, we overwrite the electron collection
0079         # name appropriately, for the fragment we just loaded above.
0080         process.egmGsfElectronIDs.physicsObjectSrc = cms.InputTag('slimmedElectrons')
0081         dataFormatString = "MiniAOD"
0082     else:
0083         raise Exception('InvalidVIDDataFormat', 'The requested data format is different from AOD or MiniAOD')
0084     #
0085 #    sys.stderr.write('Added \'egmGsfElectronIDs\' to process definition (%s format)!\n' % dataFormatString)
0086
0087 def setupVIDElectronSelection(process,cutflow,patProducer=None,addUserData=True,task=None):
0088     if not hasattr(process,'egmGsfElectronIDs'):
0089         raise Exception('VIDProducerNotAvailable','egmGsfElectronIDs producer not available in process!')
0090     setupVIDSelection(process.egmGsfElectronIDs,cutflow)
0091     #add to PAT electron producer if available or specified
0092     if hasattr(process,'patElectrons') or patProducer is not None:
0093         if patProducer is None:
0094             patProducer = process.patElectrons
0095         idName = cutflow.idName.value()
0096         addVIDSelectionToPATProducer(patProducer,'egmGsfElectronIDs',idName,addUserData)
0097
0098     #we know cutflow has a name otherwise an exception would have been thrown in setupVIDSelection
0099     #run this for all heep IDs except V60 standard for which it is not needed
0100     #it is needed for V61 and V70
0101     if (cutflow.idName.value().find("HEEP")!=-1 and
0102         cutflow.idName.value()!="heepElectronID-HEEPV60"):
0103
0104         #not the ideal way but currently the easiest
0105         useMiniAOD = process.egmGsfElectronIDs.physicsObjectSrc == cms.InputTag('slimmedElectrons')
0106
0107         from RecoEgamma.ElectronIdentification.Identification.heepElectronID_tools import addHEEPProducersToSeq
0108         addHEEPProducersToSeq(process=process,seq=process.egmGsfElectronIDSequence,
0109                               useMiniAOD=useMiniAOD,task=task)
0110
0111 ####
0112 # Muons
0113 ####
0114
0115 #turns on the VID electron ID producer, possibly with extra options
0116 # for PAT and/or MINIAOD
0117 def switchOnVIDMuonIdProducer(process, dataFormat, task=None):
0118     process.load('RecoMuon.MuonIdentification.muoMuonIDs_cff')
0119     if task is not None:
0120         task.add(process.muoMuonIDTask)
0121      #*always* reset to an empty configuration
0122     if( len(process.muoMuonIDs.physicsObjectIDs) > 0 ):
0123         process.muoMuonIDs.physicsObjectIDs = cms.VPSet()
0124     dataFormatString = "Undefined"
0125     if dataFormat == DataFormat.AOD:
0126         # No reconfiguration is required, default settings are for AOD
0127         dataFormatString = "AOD"
0128     elif dataFormat == DataFormat.MiniAOD:
0129         # If we are dealing with MiniAOD, we overwrite the muon collection
0130         # name appropriately, for the fragment we just loaded above.
0131         process.muoMuonIDs.physicsObjectSrc = cms.InputTag('slimmedMuons')
0132         dataFormatString = "MiniAOD"
0133     else:
0134         raise Exception('InvalidVIDDataFormat', 'The requested data format is different from AOD or MiniAOD')
0135     #
0136 #    sys.stderr.write('Added \'muoMuonIDs\' to process definition (%s format)!\n' % dataFormatString)
0137
0138 def setupVIDMuonSelection(process,cutflow,patProducer=None):
0139     moduleName = "muoMuonIDs"
0140     if not hasattr(process, moduleName):
0141         raise Exception("VIDProducerNotAvailable", "%s producer not available in process!" % moduleName)
0142     module = getattr(process, moduleName)
0143
0144     setupVIDSelection(module, cutflow)
0145     #add to PAT electron producer if available or specified
0146     #if hasattr(process,'patMuons') or patProducer is not None:
0147     #    if patProducer is None:
0148     #        patProducer = process.patMuons
0149     #    idName = cutflow.idName.value()
0150     #    addVIDSelectionToPATProducer(patProducer, moduleName, idName)
0151
0152 ####
0153 # Photons
0154 ####
0155
0156 #turns on the VID photon ID producer, possibly with extra options
0157 # for PAT and/or MINIAOD
0158 def switchOnVIDPhotonIdProducer(process, dataFormat, task=None):
0159     from RecoEgamma.PhotonIdentification.egmPhotonIDs_cff import  loadEgmIdSequence
0160     # Set up the ID task and sequence appropriate for this data format
0161     loadEgmIdSequence(process,dataFormat)
0162     if task is not None:
0163         task.add(process.egmPhotonIDTask)
0164     #*always* reset to an empty configuration
0165     if( len(process.egmPhotonIDs.physicsObjectIDs) > 0 ):
0166         process.egmPhotonIDs.physicsObjectIDs = cms.VPSet()
0167     dataFormatString = "Undefined"
0168     if dataFormat == DataFormat.AOD:
0169         # No reconfiguration is required, default settings are for AOD
0170         dataFormatString = "AOD"
0171     elif dataFormat == DataFormat.MiniAOD:
0172         # If we are dealing with MiniAOD, we overwrite the electron collection
0173         # name appropriately, for the fragment we just loaded above.
0174         process.egmPhotonIDs.physicsObjectSrc = cms.InputTag('slimmedPhotons')
0175         dataFormatString = "MiniAOD"
0176     else:
0177         raise Exception('InvalidVIDDataFormat', 'The requested data format is different from AOD or MiniAOD')
0178     #
0179 #    sys.stderr.write('Added \'egmPhotonIDs\' to process definition (%s format)!\n' % dataFormatString)
0180
0181 def setupVIDPhotonSelection(process,cutflow,patProducer=None,addUserData=True,task=None):
0182     if not hasattr(process,'egmPhotonIDs'):
0183         raise Exception('VIDProducerNotAvailable','egmPhotonIDs producer not available in process!\n')
0184     setupVIDSelection(process.egmPhotonIDs,cutflow)
0185     #add to PAT photon producer if available or specified
0186     if hasattr(process,'patPhotons') or patProducer is not None:
0187         if patProducer is None:
0188             patProducer = process.patPhotons
0189         idName = cutflow.idName.value()
0190         addVIDSelectionToPATProducer(patProducer,'egmPhotonIDs',idName,addUserData)
0191
0192