Merge pull request #65 from uhh-cms/feature/nbjets

Adding nbjets variables

hbt/config/configs_hbt.py

@@ -422,7 +422,7 @@ def if_not_era(*, values: list[str | None] | None = None, **kwargs) -> list[str]
     cfg.x.default_ml_model = None
     cfg.x.default_inference_model = "default_no_shifts"
     cfg.x.default_categories = ("all",)
-    cfg.x.default_variables = ("n_jet", "n_btag", "res_pdnn_hh", "res_dnn_hh")
+    cfg.x.default_variables = ("njet", "nbtag", "res_pdnn_hh", "res_dnn_hh")
     cfg.x.default_weight_producer = "default"
 
     # process groups for conveniently looping over certain processs

hbt/config/variables.py

@@ -346,6 +346,41 @@ def build_hh(events, which=None):
         x_title=r"$\Delta R_{bb}$",
     )
 
+    def build_nbjets(events, which=None):
+        wp = "medium"
+        if which == "btagPNetB":
+            wp_value = config.x.btag_working_points["particleNet"][wp]
+        elif which == "btagDeepFlavB":
+            wp_value = config.x.btag_working_points["deepjet"][wp]
+        else:
+            raise ValueError(f"Unknown which: {which}")
+        bjet_mask = events.Jet[which] >= wp_value
+        objects = events.Jet[bjet_mask]
+        objects_num = ak.num(objects, axis=1)
+        return objects_num
+
+    build_nbjets.inputs = ["Jet.{btagPNetB,btagDeepFlavB}"]
+
+    add_variable(
+        config,
+        name="nbjets_pnet",
+        expression=partial(build_nbjets, which="btagPNetB"),
+        aux={"inputs": build_nbjets.inputs},
+        binning=(11, -0.5, 10.5),
+        x_title=r"Number of ParticleNet b-jets",
+        discrete_x=True,
+    )
+
+    add_variable(
+        config,
+        name="nbjets_djet",
+        expression=partial(build_nbjets, which="btagDeepFlavB"),
+        aux={"inputs": build_nbjets.inputs},
+        binning=(11, -0.5, 10.5),
+        x_title=r"Number of DeepJet b-jets",
+        discrete_x=True,
+    )
+
     # dilepton variables
     add_variable(
         config,