up

Author: farakiko

boostedhiggs/hwwprocessor.py

@@ -149,19 +149,24 @@ def accumulator(self):
         return self._accumulator
 
     def _save_dfs_parquet(self, fname, dfs_dict, ch):
+        """Save the given pandas DataFrame dictionary as a Parquet file to the output directory."""
         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:
+        """Converts an Awkward Array collection to a flat pandas DataFrame for to write the output."""
         output = pd.DataFrame()
         for field in ak.fields(output_collection):
             output[field] = ak.to_numpy(output_collection[field])
         return output
 
     def add_selection(self, name: str, sel: np.ndarray, channel: str = "all"):
-        """Adds selection to PackedSelection object and the cutflow dictionary"""
+        """
+        Adds a boolean mask to the PackedSelection object for the given channel(s) and updates the
+        cutflow dictionary with the sum of selected events or weights.
+        """
         channels = self._channels if channel == "all" else [channel]
 
         for ch in channels:
@@ -180,6 +185,7 @@ def add_selection(self, name: str, sel: np.ndarray, channel: str = "all"):
                 self.cutflows[ch][name] = np.sum(selection_ch)
 
     def _compute_lhe_pdf_weights(self, events):
+        """Computes and returns summed LHE scale and PDF weights (used for QCD and PDF uncertainty calculations)."""
         sumlheweight = {}
         if "LHEScaleWeight" in events.fields:
             if len(events.LHEScaleWeight[0]) == 9:
@@ -193,7 +199,8 @@ def _compute_lhe_pdf_weights(self, events):
 
         return sumlheweight, sumpdfweight
 
-    def _apply_trigger_and_metfilters(self, events):
+    def _compute_trigger_msk(self, events):
+        """Computes boolean mask for trigger."""
         nevents = len(events)
         trigger = {}
         for ch in ["ele", "mu_lowpt", "mu_highpt"]:
@@ -206,16 +213,21 @@ def _apply_trigger_and_metfilters(self, events):
         trigger["mu_highpt"] = trigger["mu_highpt"] & (~trigger["ele"])
         trigger["mu_lowpt"] = trigger["mu_lowpt"] & (~trigger["ele"])
 
+        return trigger
+
+    def _compute_metfilters_msk(self, events):
+        """Computes boolean mask for METfilters."""
+        nevents = len(events)
         metfilters = np.ones(nevents, dtype="bool")
         metfilterkey = "mc" if self.isMC else "data"
         for mf in self._metfilters[metfilterkey]:
             if mf in events.Flag.fields:
                 metfilters = metfilters & events.Flag[mf]
 
-        return trigger, metfilters
+        return metfilters
 
     def _build_objects(self, events):
-
+        """Constructs physics objects (e.g. leptons, jets) and applies basic kinematic and quality cuts."""
         # MUONS
         muons = ak.with_field(events.Muon, 0, "flavor")
 
@@ -273,6 +285,13 @@ def _build_objects(self, events):
             msk_good_muons = msk_tight_muons
             msk_good_electrons = msk_tight_electrons
 
+        # CANDIDATE LEP
+        goodleptons = ak.concatenate([muons[msk_good_muons], electrons[msk_good_electrons]], axis=1)
+        goodleptons = goodleptons[ak.argsort(goodleptons.pt, ascending=False)]  # sort by pt
+
+        candidatelep = ak.firsts(goodleptons)  # pick highest pt
+        candidatelep_p4 = build_p4(candidatelep)  # build p4 for candidate lepton
+
         # AK8 JETS
         fatjets = events.FatJet
         fatjets["msdcorr"] = fatjets.msoftdrop
@@ -284,6 +303,15 @@ def _build_objects(self, events):
             events, good_fatjets, self._year, not self.isMC, self.jecs, fatjets=True
         )
 
+        FirstFatjet = ak.firsts(good_fatjets[:, 0:1])
+        SecondFatjet = ak.firsts(good_fatjets[:, 1:2])
+
+        # HIGGS CANDIDATE
+        fj_idx_lep = ak.argmin(good_fatjets.delta_r(candidatelep_p4), axis=1, keepdims=True)
+        candidatefj = ak.firsts(good_fatjets[fj_idx_lep])
+
+        jmsr_shifted_fatjetvars = get_jmsr(good_fatjets[fj_idx_lep], num_jets=1, year=self._year, isData=not self.isMC)
+
         # AK4 JETS
         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
@@ -299,6 +327,14 @@ def _build_objects(self, events):
         jet_veto_map, cut_jetveto = get_JetVetoMap(jets, self._year)
         jets = jets[(jets.pt > 30) & jet_veto_map]
 
+        # AK4 JETS (OUTSIDE AK8)
+        msk_ak4_outside_ak8 = jets.delta_r(candidatefj) > 0.8
+        ak4_outside_ak8 = jets[msk_ak4_outside_ak8]
+
+        # VBF VARIABLES
+        jet1 = ak4_outside_ak8[:, 0:1]
+        jet2 = ak4_outside_ak8[:, 1:2]
+
         # TAUS
         loose_taus_mu = (events.Tau.pt > 20) & (abs(events.Tau.eta) < 2.3) & (events.Tau.idAntiMu >= 1)  # loose antiMu ID
         loose_taus_ele = (
@@ -309,30 +345,7 @@ def _build_objects(self, events):
         loose_taus = (events.Tau.pt > 20) & (abs(events.Tau.eta) < 2.3)
         loose_taus = events.Tau[loose_taus]
 
-        # CANDIDATE LEP
-        goodleptons = ak.concatenate([muons[msk_good_muons], electrons[msk_good_electrons]], axis=1)
-        goodleptons = goodleptons[ak.argsort(goodleptons.pt, ascending=False)]  # sort by pt
-
-        candidatelep = ak.firsts(goodleptons)  # pick highest pt
-        candidatelep_p4 = build_p4(candidatelep)  # build p4 for candidate lepton
-
-        # HIGGS CANDIDATE
-        fj_idx_lep = ak.argmin(good_fatjets.delta_r(candidatelep_p4), axis=1, keepdims=True)
-        candidatefj = ak.firsts(good_fatjets[fj_idx_lep])
-
-        jmsr_shifted_fatjetvars = get_jmsr(good_fatjets[fj_idx_lep], num_jets=1, year=self._year, isData=not self.isMC)
-
-        # ak4 jet outside ak8
-        ak4_outside_ak8_selector = jets.delta_r(candidatefj) > 0.8
-        ak4_outside_ak8 = jets[ak4_outside_ak8_selector]
-
-        # VBF variables
-        jet1 = ak4_outside_ak8[:, 0:1]
-        jet2 = ak4_outside_ak8[:, 1:2]
-
-        FirstFatjet = ak.firsts(good_fatjets[:, 0:1])
-        SecondFatjet = ak.firsts(good_fatjets[:, 1:2])
-
+        # b-jets (only for jets with abs(eta)<2.5)
         bjet_selector = (jets.delta_r(candidatefj) > 0.8) & (abs(jets.eta) < 2.5)
 
         objects = {
@@ -366,20 +379,16 @@ def _build_objects(self, events):
             "jmsr_shifted_fatjetvars": jmsr_shifted_fatjetvars,
             "jec_shifted_fatjetvars": jec_shifted_fatjetvars,
             "jec_shifted_jetvars": jec_shifted_jetvars,
-            # others
+            # others (easier to just store them here)
             "fj_idx_lep": fj_idx_lep,
-            "ak4_outside_ak8_selector": ak4_outside_ak8_selector,
-            "ak4_outside_ak8": ak4_outside_ak8,
+            "msk_ak4_outside_ak8": msk_ak4_outside_ak8,
             "bjet_selector": bjet_selector,
         }
 
         return objects
 
     def _derive_variables(self, events, objects):
-
-        ak4_outside_ak8 = objects["ak4_outside_ak8"]
-        bjet_selector = objects["bjet_selector"]
-
+        """Derives physics variables from `objects` (e.g. deltaR(jet,lep))."""
         ht = ak.sum(objects["jets"].pt, axis=1)
 
         # VH jet
@@ -388,8 +397,7 @@ def _derive_variables(self, events, objects):
         mask_candidatefj = fatJetIndices != minDeltaR
 
         allScores = VScore(objects["good_fatjets"])
-        masked = allScores[mask_candidatefj]
-        VH_fj = ak.firsts(objects["good_fatjets"][allScores == ak.max(masked, axis=1)])
+        VH_fj = ak.firsts(objects["good_fatjets"][allScores == ak.max(allScores[mask_candidatefj], axis=1)])
 
         # nleptons
         n_loose_taus_mu = ak.sum(objects["loose_taus_mu"], axis=1)
@@ -413,15 +421,15 @@ def _derive_variables(self, events, objects):
         lep_miso = objects["candidatelep"].miniPFRelIso_all  # miniso for candidate lepton
 
         # b-jets (only for jets with abs(eta)<2.5)
-        ak4_bjet_candidate = objects["jets"][bjet_selector]
+        ak4_bjet_candidate = objects["jets"][objects["bjet_selector"]]
 
         # VBF variables
         deta = abs(ak.firsts(objects["jet1"]).eta - ak.firsts(objects["jet2"]).eta)
         mjj = (ak.firsts(objects["jet1"]) + ak.firsts(objects["jet2"])).mass
 
         # njets
         NumFatjets = ak.num(objects["good_fatjets"])
-        NumOtherJets = ak.num(ak4_outside_ak8)
+        NumOtherJets = ak.sum(objects["msk_ak4_outside_ak8"])
 
         # n-bjets
         n_bjets_L = ak.sum(ak4_bjet_candidate.btagDeepFlavB > btagWPs["deepJet"][self._year]["L"], axis=1)
@@ -477,7 +485,6 @@ 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,
@@ -556,14 +563,11 @@ def _derive_variables(self, events, objects):
         return derived_vars
 
     def _apply_JEC(self, objects, variables):
-
+        """Computes JES/JER and JMSR shift variations and appends them to the set of output variables."""
         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 = {
             "fj_pt": objects["candidatefj"].pt,
             "fj_eta": objects["candidatefj"].eta,
@@ -576,7 +580,7 @@ def _apply_JEC(self, objects, variables):
             # JEC vars
             for shift, vals in jec_shifted_fatjetvars["pt"].items():
                 if shift != "":
-                    fatjetvars_sys[f"fj_pt{shift}"] = ak.firsts(vals[fj_idx_lep])
+                    fatjetvars_sys[f"fj_pt{shift}"] = ak.firsts(vals[objects["fj_idx_lep"]])
 
             # JMSR vars
             for shift, vals in jmsr_shifted_fatjetvars["msoftdrop"].items():
@@ -593,11 +597,11 @@ def _apply_JEC(self, objects, variables):
                     pt_1 = objects["jet1"].pt
                     pt_2 = objects["jet2"].pt
                     try:
-                        pt_1 = vals[ak4_outside_ak8_selector][:, 0]
+                        pt_1 = vals[objects["msk_ak4_outside_ak8"]][:, 0]
                     except Exception:
                         pt_1 = objects["jet1"].pt
                     try:
-                        pt_2 = vals[ak4_outside_ak8_selector][:, 1]
+                        pt_2 = vals[objects["msk_ak4_outside_ak8"]][:, 1]
                     except Exception:
                         pt_2 = objects["jet2"].pt
 
@@ -661,7 +665,7 @@ def _apply_pileup_cutoff(self, events, year, yearmod, cutoff: float = 4):
         self.add_selection(name="PU_cutoff", sel=pw_pass)
 
     def _add_selections(self, events, trigger, metfilters, objects, variables):
-
+        """Adds event pre-selection to PackedSelection for cutflow and filtering."""
         if self.isMC:
             self._apply_pileup_cutoff(events, self._year, self._yearmod, cutoff=4)
 
@@ -738,8 +742,7 @@ def _add_selections(self, events, trigger, metfilters, objects, variables):
             self.add_selection(name="dummy", sel=(variables["NumFatjets"] > -1))
 
     def _store_genVars(self, dataset, events, objects, variables):
-
-        # store gen-level matching variables
+        """Matches generator-level particles to reconstructed jets or leptons."""
         signal_mask = None
         if self.isMC:
             if self.isSignal:
@@ -774,12 +777,8 @@ def _store_genVars(self, dataset, events, objects, variables):
 
         return variables
 
-    def _store_MCweights(self, dataset, events, objects, variables):
-
-        candidatelep = objects["candidatelep"]
-        jets = objects["jets"]
-        bjet_selector = objects["bjet_selector"]
-
+    def _add_MC_weights(self, dataset, events, objects, variables):
+        """Applies event-level weights and corrections for MC, including."""
         for ch in self._channels:
             if self._year in ("2016", "2017"):
                 self.weights[ch].add(
@@ -795,12 +794,12 @@ def _store_MCweights(self, dataset, events, objects, variables):
                 nPU=ak.to_numpy(events.Pileup.nPU),
             )
 
-            add_pileupid_weights(self.weights[ch], self._year, self._yearmod, jets, events.GenJet, wp="L")
+            add_pileupid_weights(self.weights[ch], self._year, self._yearmod, objects["jets"], events.GenJet, wp="L")
 
             if ch == "mu":
-                add_lepton_weight(self.weights[ch], candidatelep, self._year + self._yearmod, "muon")
+                add_lepton_weight(self.weights[ch], objects["candidatelep"], self._year + self._yearmod, "muon")
             elif ch == "ele":
-                add_lepton_weight(self.weights[ch], candidatelep, self._year + self._yearmod, "electron")
+                add_lepton_weight(self.weights[ch], objects["candidatelep"], self._year + self._yearmod, "electron")
 
             ewk_corr, qcd_corr, alt_qcd_corr = add_VJets_kFactors(self.weights[ch], events.GenPart, dataset, events)
             # add corrections for plotting
@@ -880,8 +879,8 @@ def _store_MCweights(self, dataset, events, objects, variables):
                     **variables,
                     **get_btag_weights(
                         self._year,
-                        jets,
-                        bjet_selector,
+                        objects["jets"],
+                        objects["bjet_selector"],
                         wp=wp_,
                         algo="deepJet",
                         systematics=self._systematics,
@@ -890,6 +889,7 @@ def _store_MCweights(self, dataset, events, objects, variables):
         return variables
 
     def _add_lundplane_weights(self, events, dataset, selection_ch, objects):
+        """Computes variables used for Lund plane reweighting technique."""
         from boostedhiggs.corrections import getLPweights
 
         pf_cands, gen_parts_eta_phi, gen_parts_pt_mass, ak8_jets, bgen_parts_eta_phi, genlep = getLPweights(
@@ -925,6 +925,7 @@ def _add_lundplane_weights(self, events, dataset, selection_ch, objects):
         return lpvars
 
     def _run_inference(self, events, selection_ch, objects):
+        """Runs ParT tagger inference using Triton."""
         model_name = "ak8_MD_vminclv2ParT_manual_fixwrap_all_nodes"
         pnet_vars = runInferenceTriton(
             self.tagger_resources_path, events[selection_ch], objects["fj_idx_lep"][selection_ch], model_name=model_name
@@ -937,6 +938,7 @@ def _run_inference(self, events, selection_ch, objects):
         return {**scores, **reg_mass, **hidNeurons}
 
     def _build_output(self, events, dataset, variables, objects):
+        """Stores the events passing the selections in a pandas dataframes."""
         output = {}
         for ch in self._channels:
             selection_ch = self.selections[ch].all(*self.selections[ch].names)
@@ -975,7 +977,7 @@ def _build_output(self, events, dataset, variables, objects):
         return output
 
     def process(self, events: ak.Array):
-        """Returns skimmed events which pass preselection cuts and with the branches listed in self._skimvars"""
+        """Returns skimmed events which pass preselection cuts as pandas dataframe."""
 
         dataset = events.metadata["dataset"]
         nevents = len(events)
@@ -994,7 +996,8 @@ def process(self, events: ak.Array):
             for ch in self._channels:
                 self.weights[ch].add("genweight", events.genWeight)
 
-        trigger, metfilters = self._apply_trigger_and_metfilters(events)
+        trigger = self._compute_trigger_msk(events)
+        metfilters = self._compute_metfilters_msk(events)
 
         objects = self._build_objects(events)
         variables = self._derive_variables(events, objects)
@@ -1010,7 +1013,7 @@ def process(self, events: ak.Array):
         variables = self._store_genVars(dataset, events, objects, variables)
 
         if self.isMC:
-            variables = self._store_MCweights(dataset, events, objects, variables)
+            variables = self._add_MC_weights(dataset, events, objects, variables)
 
         output = self._build_output(events, dataset, variables, objects)
 
@@ -1038,4 +1041,5 @@ def process(self, events: ak.Array):
         }
 
     def postprocess(self, accumulator):
+        """Placeholder for Coffea postprocessing. Currently does nothing."""
         return accumulator