python/Identification/mvaElectronID_tools.py

0001 import FWCore.ParameterSet.Config as cms
0002
0003 # =========================
0004 # Various commom parameters
0005 # =========================
0006
0007 # This MVA implementation class name
0008 mvaClassName = "ElectronMVAEstimatorRun2"
0009
0010 # The locatoins of value maps with the actual MVA values and categories
0011 # for all particles.
0012 # The names for the maps are "<module name>:<MVA class name>Values"
0013 # and "<module name>:<MVA class name>Categories"
0014 mvaProducerModuleLabel = "electronMVAValueMapProducer"
0015
0016 # The files with the variable definitions
0017 mvaVariablesFile            = "RecoEgamma/ElectronIdentification/data/ElectronMVAEstimatorRun2Variables.txt"
0018 mvaVariablesFileRun3     = "RecoEgamma/ElectronIdentification/data/ElectronIDVariablesRun3.txt"
0019 mvaVariablesFileRun3NonIso  = "RecoEgamma/ElectronIdentification/data/ElectronIDVariablesRun3NonIso.txt"
0020
0021 # =======================================
0022 # Define some commonly used category cuts
0023 # =======================================
0024
0025 EleMVA_3CategoriesCuts = [
0026     "abs(superCluster.eta) < 0.800",
0027     "abs(superCluster.eta) >= 0.800 && abs(superCluster.eta) < 1.479",
0028     "abs(superCluster.eta) >= 1.479"
0029     ]
0030
0031 EleMVA_6CategoriesCuts = [
0032     "pt < 10. && abs(superCluster.eta) < 0.800",
0033     "pt < 10. && abs(superCluster.eta) >= 0.800 && abs(superCluster.eta) < 1.479",
0034     "pt < 10. && abs(superCluster.eta) >= 1.479",
0035     "pt >= 10. && abs(superCluster.eta) < 0.800",
0036     "pt >= 10. && abs(superCluster.eta) >= 0.800 && abs(superCluster.eta) < 1.479",
0037     "pt >= 10. && abs(superCluster.eta) >= 1.479",
0038     ]
0039
0040 # =======================================================
0041 # Define simple containers for MVA cut values and related
0042 # =======================================================
0043
0044 class EleMVA_WP:
0045     """
0046     This is a container class to hold MVA cut values for a n-category MVA
0047     as well as the names of the value maps that contain the MVA values computed
0048     for all particles in a producer upstream.
0049
0050     IMPORTANT: the cuts need to be given in alphabetical order, which must
0051     be the order in which they are used by the cut class.
0052     """
0053     def __init__(self,
0054                  idName,
0055                  mvaTag,
0056                  **cuts
0057                  ):
0058         self.idName               = idName
0059         # map with MVA values for all particles
0060         self.mvaValueMapName      = mvaProducerModuleLabel + ":" + mvaClassName + mvaTag + "Values"
0061         # map with category index for all particles
0062         self.mvaCategoriesMapName = mvaProducerModuleLabel + ":" + mvaClassName + mvaTag + "Categories"
0063         self.cuts = cuts
0064
0065     def getCutStrings(self):
0066         keylist = sorted(self.cuts.keys())
0067         return [self.cuts[key] for key in keylist]
0068
0069 class EleMVARaw_WP:
0070     """
0071     This is a container class to hold MVA cut values for a n-category MVA
0072     as well as the names of the value maps that contain the MVA values computed
0073     for all particles in a producer upstream.
0074
0075     IMPORTANT: the cuts need to be given in alphabetical order, which must
0076     be the order in which they are used by the cut class.
0077     """
0078     def __init__(self,
0079                  idName,
0080                  mvaTag,
0081                  **cuts
0082                  ):
0083         self.idName               = idName
0084         # map with MVA values for all particles
0085         self.mvaValueMapName      = mvaProducerModuleLabel + ":" + mvaClassName + mvaTag + "RawValues"
0086         # map with category index for all particles
0087         self.mvaCategoriesMapName = mvaProducerModuleLabel + ":" + mvaClassName + mvaTag + "Categories"
0088         self.cuts = cuts
0089
0090     def getCutStrings(self):
0091         keylist = sorted(self.cuts.keys())
0092         return [self.cuts[key] for key in keylist]
0093
0094 # ================================
0095 # Define the complete MVA cut sets
0096 # ================================
0097
0098 def configureVIDMVAEleID(mvaWP, cutName="GsfEleMVACut"):
0099     """
0100     This function configures the full cms.PSet for a VID ID and returns it.
0101     The inputs: an object of the class EleMVA_WP or similar
0102     that contains all necessary parameters for this MVA.
0103     """
0104     pSet = cms.PSet(
0105         #
0106         idName = cms.string( mvaWP.idName ),
0107         cutFlow = cms.VPSet(
0108             cms.PSet( cutName = cms.string(cutName),
0109                       mvaCuts = cms.vstring( mvaWP.getCutStrings() ),
0110                       mvaValueMapName = cms.InputTag( mvaWP.mvaValueMapName ),
0111                       mvaCategoriesMapName = cms.InputTag( mvaWP.mvaCategoriesMapName ),
0112                       needsAdditionalProducts = cms.bool(True),
0113                       isIgnored = cms.bool(False)
0114                       )
0115             )
0116         )
0117     #
0118     return pSet