implement bjet regression calibrator (WIP)

hbw/calibration/default.py

@@ -4,16 +4,61 @@
 Calibration methods.
 """
 
+import law
+
 from columnflow.calibration import Calibrator, calibrator
 from columnflow.calibration.cms.jets import jec, jer
 from columnflow.production.cms.seeds import deterministic_seeds
 from columnflow.util import maybe_import
 
-from hbw.calibration.jet import jec_nominal
+from hbw.calibration.jet import jec_nominal, bjet_regression
 
 ak = maybe_import("awkward")
 
 
+logger = law.logger.get_logger(__name__)
+
+
+@calibrator(
+    uses={deterministic_seeds},
+    produces={deterministic_seeds},
+    skip_jecunc=True,
+    bjet_regression=True,
+)
+def base(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
+    events = self[deterministic_seeds](events, **kwargs)
+
+    logger.info(f"Running calibrators '{[calib.name for calib in self.calibrators]}' (in that order)")
+    for calibrator in self.calibrators:
+        events = self[calibrator](events, **kwargs)
+
+    return events
+
+
+@base.init
+def default_init(self: Calibrator) -> None:
+    if not getattr(self, "dataset_inst", None):
+        return
+
+    # list of calibrators to apply (in that order)
+    self.calibrators = []
+
+    if self.dataset_inst.is_data or self.skip_jecunc:
+        self.calibrators.append(jec_nominal)
+    else:
+        self.calibrators.append(jec)
+
+    if self.bjet_regression:
+        self.calibrators.append(bjet_regression)
+
+    if self.dataset_inst.is_mc:
+        # TODO: we might need to modify jer when using bjet calibration
+        self.calibrators.append(jer)
+
+    self.uses |= set(self.calibrators)
+    self.produces |= set(self.calibrators)
+
+
 @calibrator(
     uses={deterministic_seeds},
     produces={deterministic_seeds},

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