fix infinite values in MET during calibration

hbw/calibration/default.py

@@ -10,11 +10,13 @@
 from columnflow.calibration.cms.met import met_phi
 from columnflow.calibration.cms.jets import jec, jer
 from columnflow.production.cms.seeds import deterministic_seeds
-from columnflow.util import maybe_import
+from columnflow.util import maybe_import, try_float
+from columnflow.columnar_util import set_ak_column, EMPTY_FLOAT
 
 from hbw.calibration.jet import bjet_regression
 
 ak = maybe_import("awkward")
+np = maybe_import("numpy")
 
 
 logger = law.logger.get_logger(__name__)
@@ -24,6 +26,9 @@
     # jec uncertainty_sources: set to None to use config default
     jec_sources=["Total"],
     version=1,
+    # add dummy produces such that this calibrator will always be run when requested
+    # (temporary workaround until init's are only run as often as necessary)
+    produces={"FatJet.pt"},
 )
 def fatjet(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
     """
@@ -87,10 +92,29 @@ def fatjet_init(self: Calibrator) -> None:
 def jet_base(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
     events = self[deterministic_seeds](events, **kwargs)
 
+    # keep a copy of non-propagated MET to replace infinite values
+    pre_calib_met = events[self.config_inst.x.met_name]
+
     logger.info(f"Running calibrators '{[calib.cls_name for calib in self.calibrators]}' (in that order)")
     for calibrator_inst in self.calibrators:
         events = self[calibrator_inst](events, **kwargs)
 
+    # workaround for infinite values in MET pt/phi
+    for route in self.produced_columns:
+        col = route.string_column
+        m = ~np.isfinite(route.apply(events))
+        if ak.any(m):
+            # replace infinite values
+            replace_value = EMPTY_FLOAT
+            if self.config_inst.x.met_name in col:
+                # use pre-calibrated MET to replace infinite values of MET pt/phi
+                replace_value = pre_calib_met[col.split(".")[-1].split("_")[0]]
+            logger.info(
+                f"Found infinite values in {col}; Values will be replaced with "
+                f"{replace_value if try_float(replace_value) else replace_value[m]}"
+            )
+            events = set_ak_column(events, col, ak.where(m, replace_value, route.apply(events)))
+
     return events
 
 

hbw/selection/stats.py

@@ -11,7 +11,7 @@
 from columnflow.columnar_util import optional_column as optional
 
 from columnflow.util import maybe_import
-from hbw.util import has_tag
+from hbw.util import has_tag, RAW_MET_COLUMN
 
 np = maybe_import("numpy")
 ak = maybe_import("awkward")
@@ -50,7 +50,7 @@ def hbw_selection_step_stats(
 
 
 @selector(
-    uses={increment_stats, event_weights_to_normalize},
+    uses={increment_stats, event_weights_to_normalize, RAW_MET_COLUMN("pt")},
 )
 def hbw_increment_stats(
     self: Selector,
@@ -79,6 +79,10 @@ def hbw_increment_stats(
         weight_map["num_negative_weights"] = (events.mc_weight < 0)
         weight_map["num_pu_0"] = (events.pu_weight == 0)
         weight_map["num_pu_100"] = (events.pu_weight >= 100)
+
+        raw_puppi_met = events[self.config_inst.x.raw_met_name]
+        weight_map["num_raw_met_isinf"] = (~np.isfinite(raw_puppi_met.pt))
+        weight_map["num_raw_met_isinf_selected"] = (~np.isfinite(raw_puppi_met.pt) & event_mask)
         # "sum" operations
         weight_map["sum_mc_weight"] = events.mc_weight  # weights of all events
         weight_map["sum_mc_weight_selected"] = (events.mc_weight, event_mask)  # weights of selected events

hbw/util.py

@@ -635,6 +635,17 @@ def my_producer(self, events):
     return f"{met_name}.{self.get()}"
 
 
+@deferred_column
+def RAW_MET_COLUMN(self: ArrayFunction.DeferredColumn, func: ArrayFunction) -> Any | set[Any]:
+    """
+    Similar to MET_COLUMN, see MET_COLUMN for more information.
+    """
+    raw_met_name = func.config_inst.x("raw_met_name", None)
+    if not raw_met_name:
+        raise Exception("the raw_met_name has not been configured")
+    return f"{raw_met_name}.{self.get()}"
+
+
 @deferred_column
 def IF_DATASET_HAS_LHE_WEIGHTS(
     self: ArrayFunction.DeferredColumn,