ElectronIdentification/test/test_eleid_fwlite.py

0001 from RecoEgamma.ElectronIdentification.FWLite import electron_mvas, working_points
0002 from DataFormats.FWLite import Events, Handle
0003
0004 # Small script to validate Electron MVA implementation in FWlite
0005
0006 import numpy as np
0007 import pandas as pd
0008
0009 print('open input file...')
0010
0011 events = Events('root://cms-xrd-global.cern.ch//store/relval/CMSSW_10_6_0/RelValZEE_13/'+ \
0012         'MINIAODSIM/PU25ns_106X_mcRun2_asymptotic_v3_FastSim-v1/10000/7BD68C01-3BA4-A74D-8B37-3EA162D34590.root')
0013
0014 # Get Handles on the electrons and other products needed to calculate the MVAs
0015 ele_handle  = Handle('std::vector<pat::Electron>')
0016 rho_handle  = Handle('double')
0017 conv_handle = Handle('reco::ConversionCollection')
0018 bs_handle   = Handle('reco::BeamSpot')
0019
0020 n = 100000
0021
0022 data = {"Fall17IsoV2"         : np.zeros(n),
0023         "Fall17IsoV2-wp80"    : np.zeros(n, dtype=bool),
0024         "Fall17IsoV2-wp90"    : np.zeros(n, dtype=bool),
0025         "Fall17IsoV2-wpLoose" : np.zeros(n, dtype=bool),
0026         "Fall17IsoV2-wpHZZ"   : np.zeros(n, dtype=bool),
0027
0028         "Fall17NoIsoV2"         : np.zeros(n),
0029         "Fall17NoIsoV2-wp80"    : np.zeros(n, dtype=bool),
0030         "Fall17NoIsoV2-wp90"    : np.zeros(n, dtype=bool),
0031         "Fall17NoIsoV2-wpLoose" : np.zeros(n, dtype=bool),
0032
0033         "Spring16HZZV1"         : np.zeros(n),
0034         "Spring16HZZV1-wpLoose" : np.zeros(n, dtype=bool),
0035
0036         "Spring16GPV1"         : np.zeros(n),
0037         "Spring16GPV1-wp80"    : np.zeros(n, dtype=bool),
0038         "Spring16GPV1-wp90"    : np.zeros(n, dtype=bool),
0039
0040         "nEvent"        : -np.ones(n, dtype=int),
0041         "pt"            : np.zeros(n)}
0042
0043 print('start processing')
0044
0045 accepted = 0
0046 for i,event in enumerate(events):
0047
0048     nEvent = event._event.id().event()
0049
0050     print("processing event {0}: {1}...".format(i, nEvent))
0051
0052     # Save information on the first electron in an event,
0053     # if there is any the first electron of the
0054
0055     event.getByLabel(('slimmedElectrons'), ele_handle)
0056     electrons = ele_handle.product()
0057
0058     if not len(electrons):
0059         continue
0060
0061     event.getByLabel(('fixedGridRhoFastjetAll'), rho_handle)
0062
0063     rho = rho_handle.product()
0064
0065     ele = electrons[0]
0066     i = accepted
0067
0068     if ele.pt() in data["pt"][i-10:i]:
0069         continue
0070
0071     data["nEvent"][i]           = nEvent
0072     data["pt"][i]               = ele.pt()
0073
0074     mva, category = electron_mvas["Fall17IsoV2"](ele, rho)
0075     data["Fall17IsoV2"][i] = mva
0076     data["Fall17IsoV2-wp80"][i] = working_points["Fall17IsoV2"].passed(ele, mva, category, 'wp80')
0077     data["Fall17IsoV2-wp90"][i] = working_points["Fall17IsoV2"].passed(ele, mva, category, 'wp90')
0078     data["Fall17IsoV2-wpLoose"][i] = working_points["Fall17IsoV2"].passed(ele, mva, category, 'wpLoose')
0079     data["Fall17IsoV2-wpHZZ"][i] = working_points["Fall17IsoV2"].passed(ele, mva, category, 'wpHZZ')
0080
0081     mva, category = electron_mvas["Fall17NoIsoV2"](ele, rho)
0082     data["Fall17NoIsoV2"][i] = mva
0083     data["Fall17NoIsoV2-wp80"][i] = working_points["Fall17NoIsoV2"].passed(ele, mva, category, 'wp80')
0084     data["Fall17NoIsoV2-wp90"][i] = working_points["Fall17NoIsoV2"].passed(ele, mva, category, 'wp90')
0085     data["Fall17NoIsoV2-wpLoose"][i] = working_points["Fall17NoIsoV2"].passed(ele, mva, category, 'wpLoose')
0086
0087     mva, category = electron_mvas["Spring16HZZV1"](ele, rho)
0088     data["Spring16HZZV1"][i] = mva
0089     data["Spring16HZZV1-wpLoose"][i] = working_points["Spring16HZZV1"].passed(ele, mva, category, 'wpLoose')
0090
0091     mva, category = electron_mvas["Spring16GPV1"](ele, rho)
0092     data["Spring16GPV1"][i] = mva
0093     data["Spring16GPV1-wp80"][i] = working_points["Spring16GPV1"].passed(ele, mva, category, 'wp80')
0094     data["Spring16GPV1-wp90"][i] = working_points["Spring16GPV1"].passed(ele, mva, category, 'wp90')
0095
0096     accepted += 1
0097
0098     if accepted==n:
0099         break
0100
0101 ele_df = pd.DataFrame(data)
0102 ele_df = ele_df[ele_df["nEvent"] > 0]
0103 ele_df.to_hdf("test_eleid_fwlite.h5", key="electron_data")