lepton id weights

src/HHbbVV/processors/bbVVSkimmer.py

@@ -35,6 +35,7 @@
     get_jec_jets,
     get_jmsr,
     get_lund_SFs,
+    add_lepton_id_weights,
 )
 from .common import LUMI, HLTs, btagWPs, jec_shifts, jmsr_shifts
 from . import common
@@ -668,6 +669,16 @@ def process(self, events: ak.Array):
                     # to check in postprocessing for xsec & lumi normalisation
                     skimmed_events["weight_noxsec"] = weight
 
+            # separate lepton ID weights for now TODO: incorporate above if lepton vetoes are useful
+            lepton_weights = Weights(len(events), storeIndividual=True)
+            add_lepton_id_weights(lepton_weights, year, goodelectronHbb, "electron", "Loose")
+            add_lepton_id_weights(lepton_weights, year, goodelectronHH, "electron", "wp90noiso")
+            add_lepton_id_weights(lepton_weights, year, goodmuon, "muon", "Loose")
+
+            for systematic in lepton_weights.variations:
+                var_name = f"single_weight_{systematic}"
+                skimmed_events[var_name] = lepton_weights.partial_weight([var_name])
+
         # reshape and apply selections
         sel_all = selection.all(*selection.names)
 

src/HHbbVV/processors/corrections.py

@@ -432,6 +432,7 @@ def _get_lepton_clipped(lep_pt, lep_eta, lepton_type, corr=None):
     return lepton_pt, lepton_eta
 
 
+# Used only for validation region right now
 def add_lepton_weights(weights: Weights, year: str, lepton: MuonArray, lepton_type: str = "muon"):
     ul_year = get_UL_year(year)
 
@@ -455,6 +456,44 @@ def add_lepton_weights(weights: Weights, year: str, lepton: MuonArray, lepton_ty
         weights.add(f"{lepton_type}_{corr}", values["nominal"], values["up"], values["down"])
 
 
+# For analysis region
+def add_lepton_id_weights(
+    weights: Weights, year: str, lepton: NanoEventsArray, lepton_type: str, wp: str
+):
+    ul_year = get_UL_year(year)
+
+    cset = correctionlib.CorrectionSet.from_file(get_pog_json(lepton_type, year))
+
+    lep_pt = np.array(ak.fill_none(lepton.pt, 0.0))
+    lep_eta = np.abs(np.array(ak.fill_none(lepton.eta, 0.0)))
+
+    # some voodoo from cristina
+    lepton_pt, lepton_eta = _get_lepton_clipped(lep_pt, lep_eta, lepton_type)
+    values = {}
+
+    if lepton_type == "electron":
+        # https://cms-nanoaod-integration.web.cern.ch/commonJSONSFs/summaries/EGM_2018_UL_electron.html
+        cset_map = cset["UL-Electron-ID-SF"]
+
+        values["nominal"] = cset_map.evaluate(ul_year, "sf", wp, lepton_eta, lepton_pt)
+        values["up"] = cset_map.evaluate(ul_year, "systup", wp, lepton_eta, lepton_pt)
+        values["down"] = cset_map.evaluate(ul_year, "systdown", wp, lepton_eta, lepton_pt)
+    else:
+        cset_map = cset[f"NUM_{wp}ID_DEN_TrackerMuons"]
+
+        values["nominal"] = cset_map.evaluate(ul_year, lepton_eta, lepton_pt, "sf")
+        values["up"] = cset_map.evaluate(ul_year, lepton_eta, lepton_pt, "systup")
+        values["down"] = cset_map.evaluate(ul_year, lepton_eta, lepton_pt, "systdown")
+
+    for key, value in values.items():
+        # efficiency for a single lepton passing is 1 - (1 - eff1) * (1 - eff2) * ...
+        val = 1 - np.prod(1 - value, axis=1)
+        values[key] = np.nan_to_num(val, nan=1)
+
+    # add weights (for now only the nominal weight)
+    weights.add(f"{lepton_type}_id_{wp}", values["nominal"], values["up"], values["down"])
+
+
 TOP_PDGID = 6
 GEN_FLAGS = ["fromHardProcess", "isLastCopy"]