ElectronIdentification/python/FWLite.py

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, mvaVariablesFileRun3, mvaVariablesFileRun3NonIso, 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_RunIIIWinter22_iso_V1_cff \
0096         import mvaWeightFiles as RunIIIWinter22_iso_V1_weightFiles
0097 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_RunIIIWinter22_noIso_V1_cff \
0098         import mvaWeightFiles as RunIIIWinter22_noIso_V1_weightFiles
0099 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Spring16_GeneralPurpose_V1_cff \
0100         import mvaSpring16WeightFiles_V1 as mvaSpring16GPWeightFiles_V1
0101 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Spring16_HZZ_V1_cff \
0102         import mvaSpring16WeightFiles_V1 as mvaSpring16HZZWeightFiles_V1
0103 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Winter22_HZZ_V1_cff \
0104         import mvaWeightFiles as mvaWinter22HZZWeightFiles_V1
0105
0106 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Spring16_GeneralPurpose_V1_cff \
0107         import workingPoints as mvaSpring16GP_V1_workingPoints
0108 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Spring16_HZZ_V1_cff \
0109         import workingPoints as mvaSpring16HZZ_V1_workingPoints
0110 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Fall17_iso_V2_cff \
0111         import workingPoints as Fall17_iso_V2_workingPoints
0112 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Fall17_noIso_V2_cff \
0113         import workingPoints as Fall17_noIso_V2_workingPoints
0114 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_RunIIIWinter22_iso_V1_cff \
0115         import workingPoints as RunIIIWinter22_iso_V1_workingPoints
0116 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_RunIIIWinter22_noIso_V1_cff \
0117         import workingPoints as RunIIIWinter22_noIso_V1_workingPoints
0118 from RecoEgamma.ElectronIdentification.Identification.mvaElectronID_Winter22_HZZ_V1_cff \
0119         import workingPoints as Winter22_HZZ_V1_workingPoints
0120
0121 # Dictionary with the relecant e/gmma MVAs
0122
0123 electron_mvas = {
0124     "RunIIIWinter22IsoV1"   : ElectronMVAID("ElectronMVAEstimatorRun2","RunIIIWinter22IsoV1",
0125                                             EleMVA_6CategoriesCuts, RunIIIWinter22_iso_V1_weightFiles, mvaVariablesFileRun3),
0126     "RunIIIWinter22NoIsoV1" : ElectronMVAID("ElectronMVAEstimatorRun2","RunIIIWinter22NoIsoV1",
0127                                             EleMVA_6CategoriesCuts, RunIIIWinter22_noIso_V1_weightFiles, mvaVariablesFileRun3NonIso),
0128     "Fall17IsoV2"   : ElectronMVAID("ElectronMVAEstimatorRun2","Fall17IsoV2",
0129                                     EleMVA_6CategoriesCuts, Fall17_iso_V2_weightFiles, mvaVariablesFile),
0130     "Fall17NoIsoV2" : ElectronMVAID("ElectronMVAEstimatorRun2","Fall17NoIsoV2",
0131                                     EleMVA_6CategoriesCuts, Fall17_noIso_V2_weightFiles, mvaVariablesFile),
0132     "Spring16HZZV1" : ElectronMVAID("ElectronMVAEstimatorRun2","Spring16HZZV1",
0133                                     EleMVA_6CategoriesCuts, mvaSpring16HZZWeightFiles_V1, mvaVariablesFile),
0134     "Spring16GPV1"    : ElectronMVAID("ElectronMVAEstimatorRun2","Spring16GeneralPurposeV1",
0135                                     EleMVA_3CategoriesCuts, mvaSpring16GPWeightFiles_V1, mvaVariablesFile),
0136     "Winter22HZZV1"    : ElectronMVAID("ElectronMVAEstimatorRun2","Winter22HZZV1",
0137                                     EleMVA_6CategoriesCuts, mvaWinter22HZZWeightFiles_V1, mvaVariablesFileRun3),
0138     }
0139
0140 working_points = {
0141     "RunIIIWinter22IsoV1"   : WorkingPoints("ElectronMVAEstimatorRun2","RunIIIWinter22IsoV1",
0142                                     RunIIIWinter22_iso_V1_workingPoints),
0143     "RunIIIWinter22NoIsoV1" : WorkingPoints("ElectronMVAEstimatorRun2","RunIIIWinter22NoIsoV1",
0144                                     RunIIIWinter22_noIso_V1_workingPoints),
0145     "Fall17IsoV2"   : WorkingPoints("ElectronMVAEstimatorRun2","Fall17IsoV2",
0146                                     Fall17_iso_V2_workingPoints),
0147     "Fall17NoIsoV2" : WorkingPoints("ElectronMVAEstimatorRun2","Fall17NoIsoV2",
0148                                     Fall17_noIso_V2_workingPoints),
0149     "Spring16HZZV1" : WorkingPoints("ElectronMVAEstimatorRun2","Spring16HZZV1",
0150                                     mvaSpring16HZZ_V1_workingPoints, logistic_transform=True),
0151     "Spring16GPV1"    : WorkingPoints("ElectronMVAEstimatorRun2","Spring16GeneralPurposeV1",
0152                                     mvaSpring16GP_V1_workingPoints, logistic_transform=True),
0153     "Winter22HZZV1" : WorkingPoints("ElectronMVAEstimatorRun2","Winter22HZZV1",
0154                                     Winter22_HZZ_V1_workingPoints),
0155
0156     }