finish untested integration of eta dependent scalings and rate plots

configs/V29/object_performance/electron_trigger.yaml

@@ -13,9 +13,9 @@ ElectronsTriggerBarrel:
       object:
         - "abs({eta}) < 2.8"
   test_objects:
-    EG:default: "Pt"
-    tkElectron:NoIso: "Pt"
-    tkElectron:Iso: "Pt"
+    EG:default:barrel: "Pt"
+    tkElectron:NoIso:barrel: "Pt"
+    tkElectron:Iso:barrel: "Pt"
   thresholds: [10, 20, 30, 40]
   scalings:
     method: "naive"
@@ -42,9 +42,9 @@ ElectronsTriggerEndcap:
       object:
         - "abs({eta}) < 2.8"
   test_objects:
-    EG:default: "Pt"
-    tkElectron:NoIso: "Pt"
-    tkElectron:Iso: "Pt"
+    EG:default:beyond_barrel: "Pt"
+    tkElectron:NoIso:beyond_barrel: "Pt"
+    tkElectron:Iso:beyond_barrel: "Pt"
   thresholds: [10, 20, 30, 40]
   scalings:
     method: "naive"

menu_tools/rate_plots/plotter.py

@@ -148,8 +148,6 @@ def __init__(
         self.sample = sample
         self.version = version
         self.apply_offline_conversion = apply_offline_conversion
-        if self.apply_offline_conversion:
-            self.scaling_params = scalings.load_scaling_params(obj)
         self.arrays = self._load_cached_arrays()
 
     def _transform_key(self, raw_key: str) -> str:
@@ -181,7 +179,8 @@ def _load_cached_arrays(self):
 
         # Apply scalings if so configured
         if self.apply_offline_conversion:
-            arr["pt"] = scalings.compute_offline_pt(arr, self.scaling_params, "pt")
+            arr = scalings.add_offline_pt(arr, self.object)
+        arr["pt"] = scalings.get_pt_branch[arr]
 
         return arr
 
@@ -194,7 +193,10 @@ def compute_rate(self, threshold: float) -> float:
             rate: rate computed after all object cuts are applied
         """
         mask = objects.compute_selection_mask_for_object_cuts(self.object, self.arrays)
-        mask = mask & (self.arrays.pt >= threshold)
+        if self.apply_offline_conversion:
+            mask = mask & (self.arrays.offline_pt >= threshold)
+        else:
+            mask = mask & (self.arrays.pt >= threshold)
         rate = ak.sum(ak.any(mask, axis=1)) / len(mask) * constants.RATE_NORM_FACTOR
         return rate
 

menu_tools/utils/objects.py

@@ -127,7 +127,7 @@ def plot_label(self) -> str:
         return self._object_params["label"]
 
     @property
-    def eta_ranges(self) -> dict[str, tuple]:
+    def eta_ranges(self) -> dict[str, tuple[float, float]]:
         _eta_ranges = self._object_params["eta_ranges"]
         if "inclusive" not in _eta_ranges:
             _eta_ranges["inclusive"] = [0, 7]

menu_tools/utils/scalings.py

@@ -1,22 +1,25 @@
 import os
 import warnings
-from typing import Optional
 
 import awkward as ak
 import yaml
 
 from menu_tools.utils.objects import Object
 
 
-def load_scaling_params(obj: Object) -> dict:
-    """Retrieves scalings for object (incl. id) from `outputs`
+def load_scaling_params(obj: Object, eta_range: str) -> tuple[float, float]:
+    """Retrieves scalings for object+id from `outputs`
+
+    obj: Object for which to retrive scaling parameters
+    eta_range: specifier of the range for which scalings are to be retrieved
 
     Returns:
         scaling_params: parameters computed in object_performance
         for the online-offline scaling
     """
+    fname = str(obj).replace("inclusive", eta_range)
     fpath = os.path.join(
-        "outputs", "object_performance", obj.version, "scalings", f"{str(obj)}.yaml"
+        "outputs", "object_performance", obj.version, "scalings", fname + ".yaml"
     )
     try:
         with open(fpath, "r") as f:
@@ -29,43 +32,42 @@ def load_scaling_params(obj: Object) -> dict:
             lineno=18,
         )
         raise UserWarning
-    return scaling_params
+    return scaling_params["slope"], scaling_params["offset"]
 
 
-def compute_offline_pt(
-    arr: ak.Array, obj_scaling_params: dict[str, float], pt_var: Optional[str] = None
-) -> ak.Array:
-    # initialise array of zeros identical to the original pt
-    if pt_var is not None:
-        pt_orig = arr[pt_var]
-    elif "et" in arr.fields:
+def get_pt_branch(arr: ak.Array) -> ak.Array:
+    if "et" in arr.fields:
         pt_orig = arr.et
     elif "pt" in arr.fields:
         pt_orig = arr.pt
     elif "" in arr.fields:
         pt_orig = arr[""][:, 0]
     else:
-        raise ValueError(
-            "No branch to which to apply the scalings."
-            " One of `et`, `pt` or `` must exist to compute offline pt/et."
-        )
-
-    # scale pt for non-masked elements of this eta region
-    offline_pt = pt_orig * obj_scaling_params["slope"] + obj_scaling_params["offset"]
+        raise RuntimeError("Unknown pt branch!")
+    return pt_orig
 
-    return offline_pt
 
-
-def add_offline_pt(
-    arr: ak.Array, obj_scaling_params: dict, pt_var: Optional[str] = None
-) -> ak.Array:
+def add_offline_pt(arr: ak.Array, obj: Object) -> ak.Array:
     """
-    Use the scalings to convert online pT to offline pT.
-    The `pt_var` argument can be used to specify which observables
-    should be used as "pT" for a given object.
-    If `pt_var` is not specified, `pt` or `et` are used.
-    For each object, a dedicated scaling in the barrel/endcap regions
-    is applied to the online pT.
+    Add offline pt to filed called `offline_pt` and return array
     """
-    new_pt = compute_offline_pt(arr, obj_scaling_params, pt_var)
+    pt_orig = get_pt_branch(arr)
+    new_pt = ak.zeros_like(pt_orig)
+
+    if len(obj.eta_ranges) == 1:
+        # if only a single eta range is configured, the scalings are applied
+        # inclusively on that region
+        slope, offset = load_scaling_params(obj, "inclusive")
+        new_pt = new_pt + (pt_orig * slope + offset)
+    else:
+        # if multiple eta ranges are found, the "inclusive" range is skipped
+        # and all other ranges are applied
+        for eta_range, eta_min_max in obj.eta_ranges.items():
+            if eta_range == "inclusive":
+                continue
+            slope, offset = load_scaling_params(obj, eta_range)
+            eta_mask = (abs(arr.eta) >= eta_min_max[0]) & (
+                abs(arr.eta) < eta_min_max[1]
+            )
+            new_pt = new_pt + eta_mask * (pt_orig * slope + offset)
     return ak.with_field(arr, new_pt, "offline_pt")