Update to cf changes, better type hints.

hbt/inference/test.py

@@ -76,7 +76,10 @@ def test(self):
     self.add_parameter(
         "CMS_pileup2",
         type=ParameterType.rate_uniform,
-        transformations=[ParameterTransformation.effect_from_shape, ParameterTransformation.symmetrize],
+        transformations=[
+            ParameterTransformation.effect_from_shape,
+            ParameterTransformation.symmetrize,
+        ],
         config_shift_source="minbias_xs",
     )
     self.add_parameter_to_group("CMS_pileup2", "experiment")

hbt/production/tau.py

@@ -98,7 +98,7 @@ def tau_weights(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
     if self.id_vs_jet_corrector.version == 0:
         # pt, dm, genmatch, jet wp, syst, sf type
         tau_args = lambda mask, syst: (pt[mask], dm[mask], match[mask], "VVLoose", syst, "dm")
-    elif self.id_vs_jet_corrector.version == 1:
+    elif self.id_vs_jet_corrector.version in (1, 2):
         # pt, dm, genmatch, jet wp, e wp, syst, sf type
         tau_args = lambda mask, syst: (pt[mask], dm[mask], match[mask], "Loose", "VVLoose", syst, "dm")
     else:
@@ -200,9 +200,9 @@ def tau_weights_setup(self: Producer, reqs: dict, inputs: dict, reader_targets:
     self.id_vs_mu_corrector = correction_set[f"{tagger_name}VSmu"]
 
     # check versions
-    assert self.id_vs_jet_corrector.version in [0, 1]
-    assert self.id_vs_e_corrector.version in [0, 1]
-    assert self.id_vs_mu_corrector.version in [0, 1]
+    assert self.id_vs_jet_corrector.version in (0, 1, 2)
+    assert self.id_vs_e_corrector.version in (0,)
+    assert self.id_vs_mu_corrector.version in (0,)
 
 
 @producer(

hbt/selection/lepton.py

@@ -388,7 +388,8 @@ def lepton_selection(
     sel_muon_indices = empty_indices
     sel_tau_indices = empty_indices
 
-    # perform each lepton election step separately per trigger
+    # perform each lepton election step separately per trigger, avoid caching
+    sel_kwargs = {**kwargs, "call_force": True}
     for trigger, trigger_fired, leg_masks in trigger_results.x.trigger_data:
         is_single = trigger.has_tag("single_trigger")
         is_cross = trigger.has_tag("cross_trigger")
@@ -398,17 +399,15 @@ def lepton_selection(
             events,
             trigger,
             leg_masks,
-            call_force=True,
-            **kwargs,
+            **sel_kwargs,
         )
 
         # muon selection
         muon_indices, muon_veto_indices = self[muon_selection](
             events,
             trigger,
             leg_masks,
-            call_force=True,
-            **kwargs,
+            **sel_kwargs,
         )
 
         # tau selection
@@ -418,8 +417,7 @@ def lepton_selection(
             leg_masks,
             electron_indices,
             muon_indices,
-            call_force=True,
-            **kwargs,
+            **sel_kwargs,
         )
 
         # lepton pair selecton per trigger via lepton counting

hbt/util.py

@@ -8,16 +8,15 @@
 
 __all__ = ["IF_NANO_V9", "IF_NANO_V11"]
 
-from typing import Any
-
+from columnflow.types import Any
 from columnflow.columnar_util import ArrayFunction, deferred_column
 
 
 @deferred_column
-def IF_NANO_V9(self, func: ArrayFunction) -> Any | set[Any]:
+def IF_NANO_V9(self: ArrayFunction.DeferredColumn, func: ArrayFunction) -> Any | set[Any]:
     return self.get() if func.config_inst.campaign.x.version == 9 else None
 
 
 @deferred_column
-def IF_NANO_V11(self, func: ArrayFunction) -> Any | set[Any]:
+def IF_NANO_V11(self: ArrayFunction.DeferredColumn, func: ArrayFunction) -> Any | set[Any]:
     return self.get() if func.config_inst.campaign.x.version == 11 else None