diff --git a/hbw/calibration/jet.py b/hbw/calibration/jet.py index 6a70f39a..8b401f09 100644 --- a/hbw/calibration/jet.py +++ b/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