add 3lepton category

Author: mafrahm

hbw/categorization/categories.py

@@ -13,7 +13,7 @@
 from columnflow.selection import SelectionResult
 from columnflow.columnar_util import has_ak_column, optional_column
 
-from hbw.util import MET_COLUMN
+from hbw.util import MET_COLUMN, BTAG_COLUMN
 
 np = maybe_import("numpy")
 ak = maybe_import("awkward")
@@ -125,6 +125,23 @@ def catid_lep(
 catid_emu = catid_lep.derive("catid_emu", cls_dict={"n_electron": 1, "n_muon": 1})
 
 
+@categorizer(
+    uses={"{Muon,Electron}.pt"},
+)
+def catid_ge3lep(
+    self: Categorizer, events: ak.Array, results: SelectionResult | None = None, **kwargs,
+) -> tuple[ak.Array, ak.Array]:
+    if results:
+        electron = events.Electron[results.objects.Electron.Electron]
+        muon = events.Muon[results.objects.Muon.Muon]
+    else:
+        electron = events.Electron
+        muon = events.Muon
+
+    mask = ak.sum(electron.pt > 0, axis=-1) + ak.sum(muon.pt > 0, axis=-1) >= 3
+    return events, mask
+
+
 #
 # Categorizer for ABCD (either during cf.SelectEvents or cf.ProduceColumns)
 #
@@ -246,9 +263,6 @@ def catid_njet3(
     return events, mask
 
 
-from hbw.util import BTAG_COLUMN
-
-
 @categorizer(uses={BTAG_COLUMN("Jet")})
 def catid_1b(self: Categorizer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]:
     btag_column = self.config_inst.x.btag_column

hbw/config/categories.py

@@ -1,3 +1,4 @@
+
 # coding: utf-8
 
 """
@@ -136,7 +137,7 @@ def add_mll_categories(config: od.Config) -> None:
 def add_lepton_categories(config: od.Config) -> None:
     config.x.lepton_channels = {
         "sl": ("1e", "1mu"),
-        "dl": ("2e", "2mu", "emu"),
+        "dl": ("2e", "2mu", "emu", "ge3lep"),
     }[config.x.lepton_tag]
 
     cat_1e = config.add_category(  # noqa: F841
@@ -174,6 +175,13 @@ def add_lepton_categories(config: od.Config) -> None:
         label="1 Electron 1 Muon",
     )
 
+    cat_emu = config.add_category(  # noqa: F841
+        name="ge3lep",
+        id=60,
+        selection="catid_ge3lep",
+        label=r"$N_{lep} \geq 3$",
+    )
+
 
 @call_once_on_config()
 def add_njet_categories(config: od.Config) -> None: