rename cutbased if definition

ikrommyd · Jun 4, 2024 · faf3836 · faf3836
1 parent 4f2cec9
commit faf3836
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 25 deletions.
diff --git a/src/egamma_tnp/nanoaod_efficiency.py b/src/egamma_tnp/nanoaod_efficiency.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import ClassVar
+
 import dask_awkward as dak
 from coffea.lumi_tools import LumiMask
 from coffea.nanoevents import NanoAODSchema
@@ -9,6 +11,13 @@
 
 
 class ElectronTagNProbeFromNanoAOD(BaseTagNProbe):
+    cutbased_dict: ClassVar[dict] = {
+        "passingCutBasedVeto": 1,
+        "passingCutBasedLoose": 2,
+        "passingCutBasedMedium": 3,
+        "passingCutBasedTight": 4,
+    }
+
     def __init__(
         self,
         fileset,
@@ -61,8 +70,8 @@ def __init__(
             If it fails to do so, it will set it to 0.
         tags_abseta_cut: int or float, optional
             The absolute Eta cut to apply to the tag electrons. The default is 2.5.
-        cutbased_id: int, optional
-            The number of the cutbased ID to apply to the electrons.
+        cutbased_id: str, optional
+            The name of the cutbased ID to apply to the probes.
             If None, no cutbased ID is applied. The default is None.
         goldenjson: str, optional
             The golden json to use for luminosity masking. The default is None.
@@ -83,6 +92,8 @@ def __init__(
             The HLT filter to also require an event to have passed to consider a probe belonging to that event as passing.
             If None, no such requirement is applied. The default is None.
         """
+        if cutbased_id is not None and cutbased_id not in self.cutbased_dict:
+            raise ValueError(f"Incorrect cutbased ID: {cutbased_id}. Must be one of {list(self.cutbased_dict.keys())}.")
         if for_trigger is False:
             raise NotImplementedError("Only trigger efficiencies are supported at the moment.")
         if use_sc_phi and not egm_nano:
@@ -145,7 +156,7 @@ def find_probes(self, events, cut_and_count, vars):
             mask = lumimask(events.run, events.luminosityBlock)
             events = events[mask]
 
-        good_events = ElectronTagNProbeFromNanoAOD._filter_events(events, self.cutbased_id)
+        good_events = ElectronTagNProbeFromNanoAOD._filter_events(events, self.cutbased_dict[self.cutbased_id])
 
         ij = dak.argcartesian([good_events.Electron, good_events.Electron])
         is_not_diag = ij["0"] != ij["1"]
@@ -205,11 +216,11 @@ def _filter_events(events, cutbased_id):
         two_electrons = dak.num(events.Electron) >= 2
         abs_eta = abs(events.Electron.eta_to_use)
         if cutbased_id is not None:
-            pass_tight_id = events.Electron.cutBased >= cutbased_id
+            pass_cutbased_id = events.Electron.cutBased >= cutbased_id
         else:
-            pass_tight_id = True
+            pass_cutbased_id = True
         pass_eta = abs_eta <= 2.5
-        pass_selection = pass_hlt & two_electrons & pass_eta & pass_tight_id
+        pass_selection = pass_hlt & two_electrons & pass_eta & pass_cutbased_id
         n_of_good_electrons = dak.sum(pass_selection, axis=1)
         events["Electron"] = events.Electron[pass_selection]
         good_events = events[n_of_good_electrons >= 2]
@@ -292,6 +303,12 @@ def _process_zcands(
 
 
 class PhotonTagNProbeFromNanoAOD(BaseTagNProbe):
+    cutbased_dict: ClassVar[dict] = {
+        "passingCutBasedLoose": 1,
+        "passingCutBasedMedium": 2,
+        "passingCutBasedTight": 3,
+    }
+
     def __init__(
         self,
         fileset,
@@ -350,8 +367,8 @@ def __init__(
             If it fails to do so, it will set it to 0.
         tags_abseta_cut: int or float, optional
             The absolute Eta cut to apply to the tag photons. The default is 2.5.
-        cutbased_id: int, optional
-            The number of the cutbased ID to apply to the photons.
+        cutbased_id: str, optional
+            The name of the cutbased ID to apply to the probes.
             If None, no cutbased ID is applied. The default is None.
         goldenjson: str, optional
             The golden json to use for luminosity masking. The default is None.
@@ -372,6 +389,8 @@ def __init__(
             The HLT filter to also require an event to have passed to consider a probe belonging to that event as passing.
             If None, no such requirement is applied. The default is None.
         """
+        if cutbased_id is not None and cutbased_id not in self.cutbased_dict:
+            raise ValueError(f"Incorrect cutbased ID: {cutbased_id}. Must be one of {list(self.cutbased_dict.keys())}.")
         if for_trigger is False:
             raise NotImplementedError("Only trigger efficiencies are supported at the moment.")
         if use_sc_phi and not egm_nano:
@@ -445,9 +464,9 @@ def find_probes(self, events, cut_and_count, vars):
         events["Photon", "charge"] = 0.0 * events.Photon.pt
 
         if self.start_from_diphotons:
-            good_events = PhotonTagNProbeFromNanoAOD._filter_events_diphotons(events, self.cutbased_id)
+            good_events = PhotonTagNProbeFromNanoAOD._filter_events_diphotons(events, self.cutbased_dict[self.cutbased_id])
         else:
-            good_events = PhotonTagNProbeFromNanoAOD._filter_events_electron_photon(events, self.cutbased_id)
+            good_events = PhotonTagNProbeFromNanoAOD._filter_events_electron_photon(events, self.cutbased_dict[self.cutbased_id])
 
         if self.start_from_diphotons:
             ij = dak.argcartesian([good_events.Photon, good_events.Photon])
@@ -514,11 +533,11 @@ def _filter_events_diphotons(events, cutbased_id):
         two_photons = dak.num(events.Photon) >= 2
         abs_eta = abs(events.Photon.eta_to_use)
         if cutbased_id is not None:
-            pass_tight_id = events.Photon.cutBased >= cutbased_id
+            pass_cutbased_id = events.Photon.cutBased >= cutbased_id
         else:
-            pass_tight_id = True
+            pass_cutbased_id = True
         pass_eta = abs_eta <= 2.5
-        pass_selection = pass_hlt & two_photons & pass_eta & pass_tight_id
+        pass_selection = pass_hlt & two_photons & pass_eta & pass_cutbased_id
         n_of_good_photons = dak.sum(pass_selection, axis=1)
         events["Photon"] = events.Photon[pass_selection]
         good_events = events[n_of_good_photons >= 2]
@@ -533,15 +552,15 @@ def _filter_events_electron_photon(events, cutbased_id):
         abs_eta_electron = abs(events.Electron.eta_to_use)
         abs_eta_photon = abs(events.Photon.eta_to_use)
         if cutbased_id is not None:
-            pass_tight_id_electron = events.Electron.cutBased >= cutbased_id + 1
-            pass_tight_id_photon = events.Photon.cutBased >= cutbased_id
+            pass_cutbased_id_electron = events.Electron.cutBased >= cutbased_id + 1
+            pass_cutbased_id_photon = events.Photon.cutBased >= cutbased_id
         else:
-            pass_tight_id_electron = True
-            pass_tight_id_photon = True
+            pass_cutbased_id_electron = True
+            pass_cutbased_id_photon = True
         pass_eta_electron = abs_eta_electron <= 2.5
         pass_eta_photon = abs_eta_photon <= 2.5
-        pass_selection_electrons = pass_hlt & one_electron & pass_eta_electron & pass_tight_id_electron
-        pass_selection_photons = pass_hlt & one_photon & pass_eta_photon & pass_tight_id_photon
+        pass_selection_electrons = pass_hlt & one_electron & pass_eta_electron & pass_cutbased_id_electron
+        pass_selection_photons = pass_hlt & one_photon & pass_eta_photon & pass_cutbased_id_photon
         n_of_good_electrons = dak.sum(pass_selection_electrons, axis=1)
         n_of_good_photons = dak.sum(pass_selection_photons, axis=1)
         events["Electron"] = events.Electron[pass_selection_electrons]

diff --git a/src/egamma_tnp/triggers.py b/src/egamma_tnp/triggers.py
@@ -80,7 +80,7 @@ def __new__(
                 probes_pt_cut=probes_pt_cut,
                 tags_abseta_cut=2.5,
                 filterbit=1,
-                cutbased_id=4,
+                cutbased_id="passingCutBasedTight",
                 goldenjson=goldenjson,
                 extra_filter=extra_filter,
                 extra_filter_args=extra_filter_args,
@@ -178,7 +178,7 @@ def __new__(
                 probes_pt_cut=probes_pt_cut,
                 tags_abseta_cut=2.5,
                 filterbit=filterbit,
-                cutbased_id=4,
+                cutbased_id="passingCutBasedTight",
                 goldenjson=goldenjson,
                 extra_filter=extra_filter,
                 extra_filter_args=extra_filter_args,
@@ -272,7 +272,7 @@ def __new__(
                 probes_pt_cut=probes_pt_cut,
                 tags_abseta_cut=2.5,
                 filterbit=4,
-                cutbased_id=4,
+                cutbased_id="passingCutBasedTight",
                 goldenjson=goldenjson,
                 extra_filter=extra_filter,
                 extra_filter_args=extra_filter_args,
@@ -373,7 +373,7 @@ def __new__(
                 probes_pt_cut=probes_pt_cut,
                 tags_abseta_cut=2.5,
                 filterbit=5,
-                cutbased_id=4,
+                cutbased_id="passingCutBasedTight",
                 goldenjson=goldenjson,
                 extra_filter=extra_filter,
                 extra_filter_args=extra_filter_args,
@@ -469,7 +469,7 @@ def __new__(
                 probes_pt_cut=probes_pt_cut,
                 tags_abseta_cut=2.5,
                 filterbit=15,
-                cutbased_id=4,
+                cutbased_id="passingCutBasedTight",
                 goldenjson=goldenjson,
                 extra_filter=extra_filter,
                 extra_filter_args=extra_filter_args,
@@ -565,7 +565,7 @@ def __new__(
                 probes_pt_cut=probes_pt_cut,
                 tags_abseta_cut=2.5,
                 filterbit=16,
-                cutbased_id=4,
+                cutbased_id="passingCutBasedTight",
                 goldenjson=goldenjson,
                 extra_filter=extra_filter,
                 extra_filter_args=extra_filter_args,