fix: read njet

cms-btv-pog · Jan 29, 2025 · 202b066 · 202b066
1 parent aff7239
commit 202b066
Show file tree

Hide file tree

Showing 11 changed files with 52 additions and 18 deletions.
diff --git a/src/BTVNanoCommissioning/utils/array_writer.py b/src/BTVNanoCommissioning/utils/array_writer.py
@@ -50,7 +50,7 @@ def array_writer(
     othersMC=["Pileup_nTrueInt", "Pileup_nPU"],  # other fields, for MC only
     empty=False,
 ):
-    if not isRealData:
+    if not isRealData and systname != ["nominal"]:
         pruned_event["weight"] = weights.weight()
         for ind_wei in weights.weightStatistics.keys():
             pruned_event[f"{ind_wei}_weight"] = weights.partial_weight(

diff --git a/src/BTVNanoCommissioning/utils/histogrammer.py b/src/BTVNanoCommissioning/utils/histogrammer.py
@@ -705,18 +705,38 @@ def histo_writter(pruned_ev, output, weights, systematics, isSyst, SF_map):
         "DeepJetC",
     ]  # exclude b-tag SFs for btag inputs
     # define Jet flavor
-    nj = 1 if type(pruned_ev.SelJet.pt[0]) == float else len(pruned_ev.SelJet.pt[0])
+
+    # find out minimum length of jets
+    nj = (
+        min(ak.min(ak.num(pruned_ev.SelJet.pt, axis=-1)), 4)
+        if len(flatten(pruned_ev.SelJet.pt)) != len(pruned_ev)
+        else 1
+    )  # set maximum 4 jets
+    pruned_ev.SelJet = (
+        pruned_ev.SelJet
+        if nj == 1 and len(flatten(pruned_ev.SelJet)) == len(pruned_ev)
+        else pruned_ev.SelJet[:, :nj]
+    )
+    if "var" in str(ak.type(pruned_ev.SelJet.pt)) and nj == 1:
+        pruned_ev.SelJet = pruned_ev.SelJet[:, 0]
     if "hadronFlavour" in pruned_ev.SelJet.fields:
         isRealData = False
         genflavor = ak.values_astype(
-            pruned_ev.SelJet.hadronFlavour + 1 * (pruned_ev.SelJet.partonFlavour == 0)
-            & (pruned_ev.SelJet.hadronFlavour == 0),
+            pruned_ev.SelJet.hadronFlavour
+            + 1
+            * (
+                (pruned_ev.SelJet.partonFlavour == 0)
+                & (pruned_ev.SelJet.hadronFlavour == 0)
+            ),
             int,
         )
         if "MuonJet" in pruned_ev.fields:
             smflav = ak.values_astype(
-                1 * (pruned_ev.MuonJet.partonFlavour == 0)
-                & (pruned_ev.MuonJet.hadronFlavour == 0)
+                1
+                * (
+                    (pruned_ev.MuonJet.partonFlavour == 0)
+                    & (pruned_ev.MuonJet.hadronFlavour == 0)
+                )
                 + pruned_ev.MuonJet.hadronFlavour,
                 int,
             )
@@ -888,6 +908,7 @@ def histo_writter(pruned_ev, output, weights, systematics, isSyst, SF_map):
                                 weight=weight,
                             )
                         else:
+                            print(weight, flav, genflavor)
                             h.fill(
                                 syst,
                                 flatten(flav),
@@ -945,10 +966,16 @@ def histo_writter(pruned_ev, output, weights, systematics, isSyst, SF_map):
                 syst, pruned_ev.posl.delta_r(pruned_ev.negl), weight=weight
             )
             output["dr_posljet"].fill(
-                syst, genflavor, pruned_ev.posl.delta_r(pruned_ev.SelJet), weight=weight
+                syst,
+                genflavor,
+                pruned_ev.posl.delta_r(pruned_ev.SelJet),
+                weight=weight,
             )
             output["dr_negljet"].fill(
-                syst, genflavor, pruned_ev.negl.delta_r(pruned_ev.SelJet), weight=weight
+                syst,
+                genflavor,
+                pruned_ev.negl.delta_r(pruned_ev.SelJet),
+                weight=weight,
             )
         # Muon enriched jet histograms
         if "MuonJet" in pruned_ev.fields:

diff --git a/src/BTVNanoCommissioning/workflows/QCD_soft_mu_validation.py b/src/BTVNanoCommissioning/workflows/QCD_soft_mu_validation.py
@@ -131,7 +131,8 @@ def process_shift(self, events, shift_name):
                     self,
                     events[event_level],
                     events,
-                    "nominal",
+                    None,
+                    ["nominal"],
                     dataset,
                     isRealData,
                     empty=True,

diff --git a/src/BTVNanoCommissioning/workflows/QCD_validation.py b/src/BTVNanoCommissioning/workflows/QCD_validation.py
@@ -96,7 +96,6 @@ def process_shift(self, events, shift_name):
             valid_events = (ak.count(pruned_ev.Jet.pt, axis=1) > 0) & (
                 ak.count(pruned_ev.JetSVs.pt, axis=1) > 0
             )
-            print("valid_events", valid_events, len(valid_events), len(events))
             filtered_events = pruned_ev[valid_events]
 
             # Pad pruned_ev.JetSVs.jetIdx to match the length of pruned_ev.Jet

diff --git a/src/BTVNanoCommissioning/workflows/ctag_DY_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_DY_valid_sf.py
@@ -198,7 +198,8 @@ def process_shift(self, events, shift_name):
                     self,
                     events[event_level],
                     events,
-                    "nominal",
+                    None,
+                    ["nominal"],
                     dataset,
                     isRealData,
                     empty=True,
@@ -222,7 +223,7 @@ def process_shift(self, events, shift_name):
         )
         # Keep the structure of events and pruned the object size
         pruned_ev = events[event_level]
-        pruned_ev["SelJet"] = event_jet[event_level][:, 0]
+        pruned_ev["SelJet"] = event_jet[event_level]
         if isMu:
             pruned_ev["MuonPlus"] = sposmu
             pruned_ev["MuonMinus"] = snegmu

diff --git a/src/BTVNanoCommissioning/workflows/ctag_Wctt_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_Wctt_valid_sf.py
@@ -316,7 +316,8 @@ def process_shift(self, events, shift_name):
                     self,
                     events[event_level],
                     events,
-                    "nominal",
+                    None,
+                    ["nominal"],
                     dataset,
                     isRealData,
                     empty=True,

diff --git a/src/BTVNanoCommissioning/workflows/ctag_dileptt_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_dileptt_valid_sf.py
@@ -210,7 +210,8 @@ def process_shift(self, events, shift_name):
                     self,
                     events[event_level],
                     events,
-                    "nominal",
+                    None,
+                    ["nominal"],
                     dataset,
                     isRealData,
                     empty=True,

diff --git a/src/BTVNanoCommissioning/workflows/ctag_emdileptt_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_emdileptt_valid_sf.py
@@ -202,7 +202,8 @@ def process_shift(self, events, shift_name):
                     self,
                     events[event_level],
                     events,
-                    "nominal",
+                    None,
+                    ["nominal"],
                     dataset,
                     isRealData,
                     empty=True,

diff --git a/src/BTVNanoCommissioning/workflows/example.py b/src/BTVNanoCommissioning/workflows/example.py
@@ -140,7 +140,8 @@ def process_shift(self, events, shift_name):
                     self,
                     events[event_level],
                     events,
-                    "nominal",
+                    None,
+                    ["nominal"],
                     dataset,
                     isRealData,
                     empty=True,

diff --git a/src/BTVNanoCommissioning/workflows/ttdilep_valid_sf.py b/src/BTVNanoCommissioning/workflows/ttdilep_valid_sf.py
@@ -154,7 +154,8 @@ def process_shift(self, events, shift_name):
                     self,
                     events[event_level],
                     events,
-                    "nominal",
+                    None,
+                    ["nominal"],
                     dataset,
                     isRealData,
                     empty=True,

diff --git a/src/BTVNanoCommissioning/workflows/ttsemilep_valid_sf.py b/src/BTVNanoCommissioning/workflows/ttsemilep_valid_sf.py
@@ -215,7 +215,8 @@ def process_shift(self, events, shift_name):
                     self,
                     events[event_level],
                     events,
-                    "nominal",
+                    None,
+                    ["nominal"],
                     dataset,
                     isRealData,
                     empty=True,