implement bjet regression calibrator (WIP)

hbw/calibration/jet.py

@@ -4,8 +4,43 @@
 Custom jet energy calibration methods that disable data uncertainties (for searches).
 """
 
+from columnflow.util import maybe_import
+from columnflow.columnar_util import set_ak_column
 from columnflow.calibration.cms.jets import jec
-
+from columnflow.calibration import calibrator, Calibrator
 
 # custom jec calibrator that only runs nominal correction
 jec_nominal = jec.derive("jec_nominal", cls_dict={"uncertainty_sources": []})
+
+ak = maybe_import("awkward")
+
+
+@calibrator(
+    uses={
+        "Jet.pt", "Jet.phi", "Jet.eta", "Jet.mass",
+        "Jet.bRegCorr", "Jet.bRegRes",
+    },
+    produces={"Jet.pt"},
+)
+def bjet_regression(self: Calibrator, events: ak.Array, **kwargs):
+    """
+    Calibrator to apply the bjet regression.
+
+    NOTE: still looking for official documentation on how to implement correctly.
+    Question: should we apply this before or after jec/jer corrections?
+
+    Documentation: https://twiki.cern.ch/twiki/bin/view/Main/BJetRegression
+    """
+
+    # apply regression only for jet pt > 20 (docu: https://twiki.cern.ch/twiki/bin/view/Main/BJetRegression)
+    jet_mask = (events.Jet.pt > 20)
+
+    jet_pt = ak.where(jet_mask, events.Jet.pt * events.Jet.bRegCorr, events.Jet.pt)
+    jet_pt_bReg_up = ak.where(jet_mask, jet_pt * (1 + events.Jet.bRegRes), jet_pt)
+    jet_pt_bReg_down = ak.where(jet_mask, jet_pt * (1 - events.Jet.bRegRes), jet_pt)
+
+    events = set_ak_column(events, "Jet.pt", jet_pt)
+    events = set_ak_column(events, "Jet.pt_bReg_up", jet_pt_bReg_up)
+    events = set_ak_column(events, "Jet.pt_bReg_down", jet_pt_bReg_down)
+
+    return events