Add n_bjets and n_taus columns to preprocessing output

haddadanas · Jan 23, 2025 · d649bba · d649bba
1 parent ea8ec0e
commit d649bba
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/hbt/production/preprocessing/preprocess_funcs.py b/hbt/production/preprocessing/preprocess_funcs.py
@@ -133,7 +133,7 @@ def pp_jets(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
         pp_jets, pp_leptons, pp_channel_id, hh_mass, normalization_weights, process_ids
     },
     produces={
-        pp_jets, pp_leptons, pp_channel_id, hh_mass, normalization_weights, process_ids
+        pp_jets, pp_leptons, pp_channel_id, hh_mass, normalization_weights, process_ids, "n_bjets", "n_taus"
     },
 )
 def preprocess(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
@@ -151,4 +151,10 @@ def preprocess(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
 
     events = self[pp_channel_id](events)
 
+    n_bjets = ak.num(events.HHBJet, axis=1)
+    n_taus = ak.num(events.Tau, axis=1)
+
+    events = set_ak_column_f32(events, "n_bjets", n_bjets)
+    events = set_ak_column_f32(events, "n_taus", n_taus)
+
     return events