up

Author: farakiko

boostedhiggs/hwwprocessor.py

@@ -147,13 +147,13 @@ def __init__(
     def accumulator(self):
         return self._accumulator
 
-    def save_dfs_parquet(self, fname, dfs_dict, ch):
+    def _save_dfs_parquet(self, fname, dfs_dict, ch):
         if self._output_location is not None:
             table = pa.Table.from_pandas(dfs_dict)
             if len(table) != 0:  # skip dataframes with empty entries
                 pq.write_table(table, self._output_location + ch + "/parquet/" + fname + ".parquet")
 
-    def ak_to_pandas(self, output_collection: ak.Array) -> pd.DataFrame:
+    def _ak_to_pandas(self, output_collection: ak.Array) -> pd.DataFrame:
         output = pd.DataFrame()
         for field in ak.fields(output_collection):
             output[field] = ak.to_numpy(output_collection[field])
@@ -287,8 +287,6 @@ def _build_objects(self, events):
         jets, jec_shifted_jetvars = get_jec_jets(events, events.Jet, self._year, not self.isMC, self.jecs, fatjets=False)
         met = met_factory.build(events.MET, jets, {}) if self.isMC else events.MET
 
-        ht = ak.sum(jets.pt, axis=1)
-
         jet_selector = (
             (jets.pt > 15)
             & (abs(jets.eta) < 5.0)
@@ -372,33 +370,25 @@ def _build_objects(self, events):
             "ak4_outside_ak8_selector": ak4_outside_ak8_selector,
             "ak4_outside_ak8": ak4_outside_ak8,
             "bjet_selector": bjet_selector,
-            "ht": ht,
         }
 
         return objects
 
     def _derive_variables(self, events, objects):
 
-        # unpack
-        good_fatjets = objects["good_fatjets"]
-        fj_idx_lep = objects["fj_idx_lep"]
-
         ak4_outside_ak8 = objects["ak4_outside_ak8"]
-
-        FirstFatjet = objects["FirstFatjet"]
-        SecondFatjet = objects["SecondFatjet"]
-
         bjet_selector = objects["bjet_selector"]
-        ht = objects["ht"]
+
+        ht = ak.sum(objects["jets"].pt, axis=1)
 
         # VH jet
-        minDeltaR = ak.argmin(objects["candidatelep_p4"].delta_r(good_fatjets), axis=1)
-        fatJetIndices = ak.local_index(good_fatjets, axis=1)
+        minDeltaR = ak.argmin(objects["candidatelep_p4"].delta_r(objects["good_fatjets"]), axis=1)
+        fatJetIndices = ak.local_index(objects["good_fatjets"], axis=1)
         mask_candidatefj = fatJetIndices != minDeltaR
 
-        allScores = VScore(good_fatjets)
+        allScores = VScore(objects["good_fatjets"])
         masked = allScores[mask_candidatefj]
-        VH_fj = ak.firsts(good_fatjets[allScores == ak.max(masked, axis=1)])
+        VH_fj = ak.firsts(objects["good_fatjets"][allScores == ak.max(masked, axis=1)])
 
         # nleptons
         n_loose_taus_mu = ak.sum(objects["loose_taus_mu"], axis=1)
@@ -429,7 +419,7 @@ def _derive_variables(self, events, objects):
         mjj = (ak.firsts(objects["jet1"]) + ak.firsts(objects["jet2"])).mass
 
         # njets
-        NumFatjets = ak.num(good_fatjets)
+        NumFatjets = ak.num(objects["good_fatjets"])
         NumOtherJets = ak.num(ak4_outside_ak8)
 
         # n-bjets
@@ -486,6 +476,7 @@ def _derive_variables(self, events, objects):
             nCjets = (ak.sum(goodgenjets.hadronFlavour == 4, axis=1)).to_numpy()
 
         derived_vars = {
+            "fj_idx_lep": objects["fj_idx_lep"],
             "ht": ht,
             # candidatefj
             "fj_pt": objects["candidatefj"].pt,
@@ -528,17 +519,17 @@ def _derive_variables(self, events, objects):
             "deta": deta,
             "mjj": mjj,
             # leading fatjet
-            "FirstFatjet_pt": FirstFatjet.pt,
-            "FirstFatjet_eta": FirstFatjet.eta,
-            "FirstFatjet_phi": FirstFatjet.phi,
-            "FirstFatjet_msd": FirstFatjet.msdcorr,
-            "FirstFatjet_Vscore": VScore(SecondFatjet),
+            "FirstFatjet_pt": objects["FirstFatjet"].pt,
+            "FirstFatjet_eta": objects["FirstFatjet"].eta,
+            "FirstFatjet_phi": objects["FirstFatjet"].phi,
+            "FirstFatjet_msd": objects["FirstFatjet"].msdcorr,
+            "FirstFatjet_Vscore": VScore(objects["SecondFatjet"]),
             # second leading fatjet
-            "SecondFatjet_pt": SecondFatjet.pt,
-            "SecondFatjet_eta": SecondFatjet.eta,
-            "SecondFatjet_phi": SecondFatjet.phi,
-            "SecondFatjet_msd": SecondFatjet.msdcorr,
-            "SecondFatjet_Vscore": VScore(FirstFatjet),
+            "SecondFatjet_pt": objects["SecondFatjet"].pt,
+            "SecondFatjet_eta": objects["SecondFatjet"].eta,
+            "SecondFatjet_phi": objects["SecondFatjet"].phi,
+            "SecondFatjet_msd": objects["SecondFatjet"].msdcorr,
+            "SecondFatjet_Vscore": VScore(objects["FirstFatjet"]),
             # second fatjet after candidate jet
             "VH_fj_pt": VH_fj.pt,
             "VH_fj_eta": VH_fj.eta,
@@ -549,7 +540,7 @@ def _derive_variables(self, events, objects):
             ).miniPFRelIso_all,
             "loose_lep1_pt": ak.firsts(objects["loose_muons1"][ak.argsort(objects["loose_muons1"].pt, ascending=False)]).pt,
             "msk_leptonic_taus": msk_leptonic_taus,
-            "fj_mass_raw": good_fatjets[fj_idx_lep].msdcorr,
+            "ht": ht,
         }
 
         if self.isMC:
@@ -565,11 +556,11 @@ def _derive_variables(self, events, objects):
 
     def _apply_JEC(self, objects, variables):
 
-        fj_idx_lep = objects["fj_idx_lep"]
         jec_shifted_fatjetvars = objects["jec_shifted_fatjetvars"]
         jmsr_shifted_fatjetvars = objects["jmsr_shifted_fatjetvars"]
         jec_shifted_jetvars = objects["jec_shifted_jetvars"]
 
+        fj_idx_lep = objects["fj_idx_lep"]
         ak4_outside_ak8_selector = objects["ak4_outside_ak8_selector"]
 
         fatjetvars = {
@@ -668,7 +659,7 @@ def _apply_pileup_cutoff(self, events, year, yearmod, cutoff: float = 4):
 
         self.add_selection(name="PU_cutoff", sel=pw_pass)
 
-    def _apply_selections(self, events, trigger, metfilters, objects, variables):
+    def _add_selections(self, events, trigger, metfilters, objects, variables):
 
         if self.isMC:
             self._apply_pileup_cutoff(events, self._year, self._yearmod, cutoff=4)
@@ -937,7 +928,7 @@ def _run_inference(self, events, selection_ch, objects):
         pnet_vars = runInferenceTriton(
             self.tagger_resources_path, events[selection_ch], objects["fj_idx_lep"][selection_ch], model_name=model_name
         )
-        pnet_df = self.ak_to_pandas(pnet_vars)
+        pnet_df = self._ak_to_pandas(pnet_vars)
         scores = {"fj_ParT_score": pnet_df[sigs].sum(axis=1).values}
         hidNeurons = {k: v for k, v in pnet_vars.items() if "hidNeuron" in k}
         reg_mass = {"fj_ParT_mass": pnet_vars["fj_ParT_mass"]}
@@ -967,7 +958,7 @@ def _build_output(self, events, dataset, variables, objects):
                 output[ch] = {}
 
             if not isinstance(output[ch], pd.DataFrame):
-                output[ch] = self.ak_to_pandas(output[ch])
+                output[ch] = self._ak_to_pandas(output[ch])
 
             for var_ in [
                 "rec_higgs_m",
@@ -1013,7 +1004,7 @@ def process(self, events: ak.Array):
 
         variables = self._apply_JEC(objects, variables)
 
-        self._apply_selections(events, trigger, metfilters, objects, variables)
+        self._add_selections(events, trigger, metfilters, objects, variables)
 
         variables = self._store_genVars(dataset, events, objects, variables)
 
@@ -1030,7 +1021,7 @@ def process(self, events: ak.Array):
                 os.makedirs(self._output_location + ch)
             if not os.path.exists(self._output_location + ch + "/parquet"):
                 os.makedirs(self._output_location + ch + "/parquet")
-            self.save_dfs_parquet(fname, output[ch], ch)
+            self._save_dfs_parquet(fname, output[ch], ch)
 
         return {
             dataset: {