Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2022-01-26 04:12:36

0001 import ROOT
0002 import ctypes
0003 import pprint
0004 from numpy import exp
0005 
0006 # Python wrappers around the Electron MVAs.
0007 # Usage example in RecoEgamma/ElectronIdentification/test
0008 
0009 class ElectronMVAID:
0010     """ Electron MVA wrapper class.
0011     """
0012 
0013     def __init__(self, name, tag, categoryCuts, xmls, variablesFile, debug=False):
0014         self.name = name
0015         self.tag = tag
0016         self.categoryCuts = categoryCuts
0017         self.variablesFile = variablesFile
0018         self.xmls = ROOT.vector(ROOT.string)()
0019         for x in xmls: self.xmls.push_back(x)
0020         self._init = False
0021         self._debug = debug
0022 
0023     def __call__(self, ele, rho, debug=False):
0024         '''returns a tuple mva_value, category 
0025         ele: a reco::GsfElectron
0026         convs: conversions
0027         beam_spot: beam spot
0028         rho: energy density in the event
0029         debug: enable debugging mode. 
0030 
0031         example: 
0032         
0033             event.getByLabel(('slimmedElectrons'),       ele_handle)
0034             event.getByLabel(('fixedGridRhoFastjetAll'), rho_handle)
0035             
0036             electrons = ele_handle.product()
0037             rho       = rho_handle.product()
0038 
0039             mva, category = electron_mva_id(electron[0], rho)
0040         '''
0041         if not self._init:
0042             print('Initializing ' + self.name + self.tag)
0043             ROOT.gInterpreter.Declare('#include "RecoEgamma/ElectronIdentification/interface/ElectronMVAEstimatorRun2.h"')
0044             ROOT.gSystem.Load("libRecoEgammaElectronIdentification")
0045             categoryCutStrings =  ROOT.vector(ROOT.string)()
0046             for x in self.categoryCuts : 
0047                 categoryCutStrings.push_back(x)
0048             self.estimator = ROOT.ElectronMVAEstimatorRun2(
0049                 self.tag, self.name, len(self.xmls), 
0050                 self.variablesFile, categoryCutStrings, self.xmls, self._debug)
0051             self._init = True
0052         category = ctypes.c_int(0)
0053         mva = self.estimator.mvaValue(ele, rho[0], category)
0054         return mva, category.value
0055 
0056 
0057 class WorkingPoints(object):
0058     '''Working Points. Keeps track of the cuts associated to a given flavour of the MVA ID 
0059     for each working point and allows to test the working points'''
0060 
0061     def __init__(self, name, tag, working_points, logistic_transform=False):
0062         self.name = name 
0063         self.tag = tag
0064         self.working_points = self._reformat_cut_definitions(working_points)
0065         self.logistic_transform = logistic_transform
0066 
0067     def _reformat_cut_definitions(self, working_points):
0068         new_definitions = dict()
0069         for wpname, definitions in working_points.items():
0070             new_definitions[wpname] = dict()
0071             for name, cut in definitions.cuts.items():
0072                 categ_id = int(name.lstrip('cutCategory'))
0073                 cut = cut.replace('pt','x')
0074                 formula = ROOT.TFormula('_'.join([self.name, wpname, name]), cut)
0075                 new_definitions[wpname][categ_id] = formula
0076         return new_definitions
0077 
0078     def passed(self, ele, mva, category, wp):
0079         '''return true if ele passes wp'''
0080         threshold = self.working_points[wp][category].Eval(ele.pt())
0081         if self.logistic_transform:
0082             mva = 2.0/(1.0+exp(-2.0*mva))-1
0083         return mva > threshold
0084 
0085 
0086 # Import information needed to construct the e/gamma MVAs
0087 
0088 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_tools \
0089         import EleMVA_6CategoriesCuts, mvaVariablesFile, EleMVA_3CategoriesCuts
0090 
0091 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Fall17_iso_V2_cff \
0092         import mvaWeightFiles as Fall17_iso_V2_weightFiles
0093 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Fall17_noIso_V2_cff \
0094         import mvaWeightFiles as Fall17_noIso_V2_weightFiles
0095 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Spring16_GeneralPurpose_V1_cff \
0096         import mvaSpring16WeightFiles_V1 as mvaSpring16GPWeightFiles_V1
0097 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Spring16_HZZ_V1_cff \
0098         import mvaSpring16WeightFiles_V1 as mvaSpring16HZZWeightFiles_V1
0099 
0100 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Spring16_GeneralPurpose_V1_cff \
0101         import workingPoints as mvaSpring16GP_V1_workingPoints
0102 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Spring16_HZZ_V1_cff \
0103         import workingPoints as mvaSpring16HZZ_V1_workingPoints
0104 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Fall17_iso_V2_cff \
0105         import workingPoints as Fall17_iso_V2_workingPoints
0106 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Fall17_noIso_V2_cff \
0107         import workingPoints as Fall17_noIso_V2_workingPoints
0108 
0109 # Dictionary with the relecant e/gmma MVAs
0110 
0111 electron_mvas = {
0112     "Fall17IsoV2"   : ElectronMVAID("ElectronMVAEstimatorRun2","Fall17IsoV2",
0113                                     EleMVA_6CategoriesCuts, Fall17_iso_V2_weightFiles, mvaVariablesFile),
0114     "Fall17NoIsoV2" : ElectronMVAID("ElectronMVAEstimatorRun2","Fall17NoIsoV2",
0115                                     EleMVA_6CategoriesCuts, Fall17_noIso_V2_weightFiles, mvaVariablesFile),
0116     "Spring16HZZV1" : ElectronMVAID("ElectronMVAEstimatorRun2","Spring16HZZV1",
0117                                     EleMVA_6CategoriesCuts, mvaSpring16HZZWeightFiles_V1, mvaVariablesFile),
0118     "Spring16GPV1"    : ElectronMVAID("ElectronMVAEstimatorRun2","Spring16GeneralPurposeV1",
0119                                     EleMVA_3CategoriesCuts, mvaSpring16GPWeightFiles_V1, mvaVariablesFile),
0120     }
0121 
0122 working_points = {
0123     "Fall17IsoV2"   : WorkingPoints("ElectronMVAEstimatorRun2","Fall17IsoV2",
0124                                     Fall17_iso_V2_workingPoints),
0125     "Fall17NoIsoV2" : WorkingPoints("ElectronMVAEstimatorRun2","Fall17NoIsoV2",
0126                                     Fall17_noIso_V2_workingPoints),
0127     "Spring16HZZV1" : WorkingPoints("ElectronMVAEstimatorRun2","Spring16HZZV1",
0128                                     mvaSpring16HZZ_V1_workingPoints, logistic_transform=True),
0129     "Spring16GPV1"    : WorkingPoints("ElectronMVAEstimatorRun2","Spring16GeneralPurposeV1",
0130                                     mvaSpring16GP_V1_workingPoints, logistic_transform=True),
0131 
0132     }