update cf (refactor taf_init) and first set of changes

Author: mafrahm

hbw/analysis/create_analysis.py

@@ -191,7 +191,7 @@ def analysis_factory(configs: od.UniqueObjectIndex):
         "analysis", "task_family", "config", "configs", "dataset", "shift", "version",
         "calibrator", "calibrators", "selector", "producer", "producers",
         "ml_model", "ml_data", "ml_models",
-        "weightprod", "inf_model",
+        "weight_producer", "inf_model",
         "plot", "shift_sources", "shifts", "datasets",
         # MLTraining
         "calib", "sel", "prod",
@@ -249,7 +249,7 @@ def reorganize_parts(task, store_parts):
             "config", "configs",
             "producers", "prod",
             "ml_data", "ml_model", "ml_models",
-            "weightprod", "inf_model",
+            "weight_producer", "inf_model",
             "task_family",
             "calibrator", "producer",
             "shift", "dataset",

hbw/calibration/default.py

@@ -71,12 +71,12 @@ def ele_init(self: Calibrator) -> None:
     # TODO: deterministic FatJet seeds
     produces={"FatJet.pt"},
 )
-def fatjet(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
+def fatjet(self: Calibrator, events: ak.Array, task, **kwargs) -> ak.Array:
     """
     FatJet calibrator, combining JEC and JER.
     Uses as JER uncertainty either only "Total" for MC or no uncertainty for data.
     """
-    if self.task.local_shift != "nominal":
+    if task.local_shift != "nominal":
         raise Exception("FatJet Calibrator should not be run for shifts other than nominal")
 
     # apply the fatjet JEC and JER
@@ -89,10 +89,6 @@ def fatjet(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
 
 @fatjet.init
 def fatjet_init(self: Calibrator) -> None:
-    if not self.task or self.task.task_family != "cf.CalibrateEvents":
-        # init only required for task itself
-        return
-
     # derive calibrators to add settings once
     flag = f"custom_fatjet_calibs_registered_{self.cls_name}"
     if not self.config_inst.x(flag, False):
@@ -185,9 +181,6 @@ def jet_base(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array:
 
 @jet_base.init
 def jet_base_init(self: Calibrator) -> None:
-    if not self.task or self.task.task_family != "cf.CalibrateEvents":
-        # init only required for task itself
-        return
 
     # derive calibrators to add settings once
     flag = f"custom_jet_calibs_registered_{self.cls_name}"

hbw/calibration/jet.py

@@ -88,23 +88,22 @@ def bjet_regression(
 @bjet_regression.init
 def bjet_regression_init(self: Calibrator):
     # setup only required by CalibrateEvents task itself
-    if self.task and self.task.task_family == "cf.CalibrateEvents":
-        self.b_tagger = {
-            2: "deepjet",
-            3: "particlenet",
-        }[self.config_inst.x.run]
-
-        self.b_score_column = {
-            "particlenet": "btagPNetB",
-            "deepjet": "btagDeepFlavB",
-        }[self.b_tagger]
-
-        self.b_reg_column = {
-            "particlenet": "PNetRegPtRawCorr",
-            "deepjet": "bRegCorr",
-        }[self.b_tagger]
-
-        self.btag_wp = self.config_inst.x("btag_wp", "medium")
-
-        self.uses.add(f"Jet.{self.b_score_column}")
-        self.uses.add(f"Jet.{self.b_reg_column}")
+    self.b_tagger = {
+        2: "deepjet",
+        3: "particlenet",
+    }[self.config_inst.x.run]
+
+    self.b_score_column = {
+        "particlenet": "btagPNetB",
+        "deepjet": "btagDeepFlavB",
+    }[self.b_tagger]
+
+    self.b_reg_column = {
+        "particlenet": "PNetRegPtRawCorr",
+        "deepjet": "bRegCorr",
+    }[self.b_tagger]
+
+    self.btag_wp = self.config_inst.x("btag_wp", "medium")
+
+    self.uses.add(f"Jet.{self.b_score_column}")
+    self.uses.add(f"Jet.{self.b_reg_column}")

hbw/categorization/categories.py

@@ -112,8 +112,8 @@ def catid_lep(
         muon = events.Muon
 
     mask = (
-        (ak.sum(electron.pt > 0, axis=-1) == self.n_electron) &
-        (ak.sum(muon.pt > 0, axis=-1) == self.n_muon)
+        (ak.sum(electron["pt"] > 0, axis=-1) == self.n_electron) &
+        (ak.sum(muon["pt"] > 0, axis=-1) == self.n_muon)
     )
     return events, mask
 
@@ -138,7 +138,7 @@ def catid_ge3lep(
         electron = events.Electron
         muon = events.Muon
 
-    mask = ak.sum(electron.pt > 0, axis=-1) + ak.sum(muon.pt > 0, axis=-1) >= 3
+    mask = ak.sum(electron["pt"] > 0, axis=-1) + ak.sum(muon["pt"] > 0, axis=-1) >= 3
     return events, mask
 
 

hbw/columnflow_patches.py

@@ -84,13 +84,22 @@ def patch_csp_versioning():
     Patches the TaskArrayFunction to add the version to the string representation of the task.
     """
 
+    from columnflow.tasks.framework.mixins import ArrayFunctionClassMixin
+
     def TaskArrayFunction_str(self):
         version = self.version() if callable(getattr(self, "version", None)) else getattr(self, "version", None)
         if version and not isinstance(version, (int, str)):
             raise Exception(f"version must be an integer or string, but is {version} ({type(version)})")
         version_str = f"V{version}" if version is not None else ""
         return f"{self.cls_name}{version_str}"
 
+    def array_function_cls_repr(self, array_function):
+        # NOTE: this might be a problem when we have identical names between different types of
+        # TaskArrayFunctions...
+        array_function_cls = TaskArrayFunction.get_cls(array_function)
+        return TaskArrayFunction_str(array_function_cls)
+
+    ArrayFunctionClassMixin.array_function_cls_repr = array_function_cls_repr
     TaskArrayFunction.__str__ = TaskArrayFunction_str
     logger.info(
         "patched TaskArrayFunction.__str__ to include the CSP version attribute",

hbw/config/defaults_and_groups.py

@@ -51,9 +51,9 @@ def default_ml_model(cls, container, task_params):
         if inference_model in (None, law.NO_STR, RESOLVE_DEFAULT):
             inference_model = container.x.default_inference_model
 
-        # get the default_ml_model from the inference_model_inst
-        inference_model_inst = InferenceModel.get_cls(inference_model)
-        default_ml_model = getattr(inference_model_inst, "ml_model_name", default_ml_model)
+        # get the default_ml_model from the inference_model_cls
+        inference_model_cls = InferenceModel.get_cls(inference_model)
+        default_ml_model = getattr(inference_model_cls, "ml_model_name", default_ml_model)
 
     return default_ml_model
 

hbw/config/styling.py

@@ -205,7 +205,7 @@ def stylize_processes(config: od.Config) -> None:
         if "hh_" in proc.name.lower():
             proc.add_tag("is_signal")
             proc.unstack = True
-            proc.scale = "stack"
+            # proc.scale = "stack"
 
         # labels used for ML categories
         proc.x.ml_label = ml_labels.get(proc.name, proc.name)

hbw/config/variables.py

@@ -113,21 +113,21 @@ def add_variables(config: od.Config) -> None:
         expression="event",
         binning=(1, 0.0, 1.0e9),
         x_title="Event number",
-        discrete_x=True,
+        # discrete_x=True,
     )
     config.add_variable(
         name="run",
         expression="run",
         binning=(1, 100000.0, 500000.0),
         x_title="Run number",
-        discrete_x=True,
+        # discrete_x=True,
     )
     config.add_variable(
         name="lumi",
         expression="luminosityBlock",
         binning=(1, 0.0, 5000.0),
         x_title="Luminosity block",
-        discrete_x=True,
+        # discrete_x=True,
     )
 
     #

hbw/ml/stats.py

@@ -50,6 +50,7 @@ def del_sub_proc_stats(
 def ml_preparation(
     self: Producer,
     events: ak.Array,
+    task,
     stats: dict = {},
     fold_indices: ak.Array | None = None,
     ml_model_inst: MLModel | None = None,
@@ -58,7 +59,7 @@ def ml_preparation(
     """
     Producer that is run as part of PrepareMLEvents to collect relevant stats
     """
-    if self.task.task_family == "cf.PrepareMLEvents":
+    if task.task_family == "cf.PrepareMLEvents":
         # pass category mask to only use events that belong to the main "signal region"
         # NOTE: we could also just require the pre_ml_cats Producer here
         sr_categorizer = catid_sr if self.config_inst.has_tag("is_sl") else catid_mll_low
@@ -70,7 +71,7 @@ def ml_preparation(
         "num_events": Ellipsis,  # all events
     }
 
-    if self.task.dataset_inst.is_mc:
+    if task.dataset_inst.is_mc:
         # full event weight
         events, weight = self[default_weight_producer](events, **kwargs)
         events = set_ak_column_f32(events, "event_weight", weight)
@@ -121,8 +122,6 @@ def ml_preparation(
 
 @ml_preparation.init
 def ml_preparation_init(self):
-    # TODO: we access self.task.dataset_inst instead of self.dataset_inst due to an issue
-    # with the preparation producer initialization
     if not getattr(self, "dataset_inst", None) or self.dataset_inst.is_data:
         return
 

hbw/production/categories.py

@@ -7,7 +7,7 @@
 import law
 
 from columnflow.production import Producer, producer
-from columnflow.util import maybe_import, InsertableDict
+from columnflow.util import maybe_import
 from columnflow.production.categories import category_ids
 
 from hbw.config.categories import add_categories_production, add_categories_ml
@@ -58,16 +58,18 @@ def cats_ml(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
 
 
 @cats_ml.requires
-def cats_ml_reqs(self: Producer, reqs: dict) -> None:
+def cats_ml_reqs(self: Producer, task: law.Task, reqs: dict) -> None:
     if "ml" in reqs:
         return
 
     from columnflow.tasks.ml import MLEvaluation
-    reqs["ml"] = MLEvaluation.req(self.task, ml_model=self.ml_model_name)
+    reqs["ml"] = MLEvaluation.req(task, ml_model=self.ml_model_name)
 
 
 @cats_ml.setup
-def cats_ml_setup(self: Producer, reqs: dict, inputs: dict, reader_targets: InsertableDict) -> None:
+def cats_ml_setup(
+    self: Producer, task: law.Task, reqs: dict, inputs: dict, reader_targets: law.util.InsertableDict,
+) -> None:
     reader_targets["mlcolumns"] = inputs["ml"]["mlcolumns"]
 
 

hbw/production/dataset_normalization.py

@@ -3,9 +3,10 @@
 """
 Column production methods related to sample normalization event weights.
 """
+import law
 
 from columnflow.production import Producer, producer
-from columnflow.util import maybe_import, InsertableDict
+from columnflow.util import maybe_import
 from columnflow.columnar_util import set_ak_column
 
 np = maybe_import("numpy")
@@ -39,7 +40,7 @@ def dataset_normalization_weight(self: Producer, events: ak.Array, **kwargs) ->
 
 
 @dataset_normalization_weight.requires
-def dataset_normalization_weight_requires(self: Producer, reqs: dict) -> None:
+def dataset_normalization_weight_requires(self: Producer, task: law.Task, reqs: dict) -> None:
     """
     Adds the requirements needed by the underlying py:attr:`task` to access selection stats into
     *reqs*.
@@ -49,17 +50,18 @@ def dataset_normalization_weight_requires(self: Producer, reqs: dict) -> None:
     #       (i.e. all datasets that might contain any of the sub processes found in a dataset)
     from columnflow.tasks.selection import MergeSelectionStats
     reqs["selection_stats"] = MergeSelectionStats.req(
-        self.task,
+        task,
         branch=-1,
     )
 
 
 @dataset_normalization_weight.setup
 def dataset_normalization_weight_setup(
     self: Producer,
+    task: law.Task,
     reqs: dict,
     inputs: dict,
-    reader_targets: InsertableDict,
+    reader_targets: law.util.InsertableDict,
 ) -> None:
     """
     Load inclusive selection stats and cross sections for the normalization weight calculation.

hbw/production/gen_v.py

@@ -6,8 +6,10 @@
 
 from __future__ import annotations
 
+import law
+
 from columnflow.production import Producer, producer
-from columnflow.util import maybe_import, InsertableDict, DotDict
+from columnflow.util import maybe_import, DotDict
 from columnflow.columnar_util import set_ak_column
 
 
@@ -206,24 +208,19 @@ def vjets_weight_skip(self: Producer) -> bool:
     )
 
 
-@vjets_weight.init
-def vjets_weight_init(self: Producer) -> None:
-    shift_inst = getattr(self, "local_shift_inst", None)
-    if not shift_inst:
-        return
-
-
 @vjets_weight.requires
-def vjets_weight_requires(self: Producer, reqs: dict) -> None:
+def vjets_weight_requires(self: Producer, task: law.Task, reqs: dict) -> None:
     if "external_files" in reqs:
         return
 
     from columnflow.tasks.external import BundleExternalFiles
-    reqs["external_files"] = BundleExternalFiles.req(self.task)
+    reqs["external_files"] = BundleExternalFiles.req(task)
 
 
 @vjets_weight.setup
-def vjets_weight_setup(self: Producer, reqs: dict, inputs: dict, reader_targets: InsertableDict) -> None:
+def vjets_weight_setup(
+    self: Producer, task: law.Task, reqs: dict, inputs: dict, reader_targets: law.util.InsertableDict,
+) -> None:
     bundle = reqs["external_files"]
 
     # create the L1 prefiring weight evaluator

hbw/production/jets.py

@@ -29,7 +29,7 @@ def jetId(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
     https://twiki.cern.ch/twiki/bin/view/CMS/JetID13p6TeV?rev=21
     """
     abseta = abs(events.Jet.eta)
-    print("start")
+
     # baseline mask (abseta < 2.7)
     passJetId_Tight = (events.Jet.jetId & 2 == 2)