Skimmer: correctionlib JECs for 2022-2023 and Type-1 PUPPI MET with the re-applied JECs

src/HH4b/boosted/bdt_trainings_run3/v13_glopartv3.py

@@ -139,8 +139,8 @@ def mask(arr, pad_mask):
             key_map("HHPt"): mask(hh.pt, hh_pad),
             key_map("HHeta"): mask(hh.eta, hh_pad),
             key_map("HHmass"): mask(hh.mass, hh_pad),
-            # met in the event
-            key_map("MET"): events.MET_pt[0],
+            # MET (Type-1 PUPPI, recomputed in the skimmer with the re-applied JECs)
+            key_map("MET"): events[(key_map("MET_pt"), 0)],
             # fatjet tau32
             key_map("H1T32"): mask(
                 events[key_map("bbFatJetTau3OverTau2")].to_numpy()[:, 0], h1_pad

src/HH4b/corrections/README.md

@@ -1,10 +1,11 @@
-## To build JECs
+## JECs
 
-```
-python -m build_jec jec_compiled.pkl.gz
-```
+Run 3 (2022-2025): the official JME correctionlib files on CVMFS, pinned in
+`processors/corrections.py` (`cat_jerc`) and read at run time through `processors/jec_clib.py`
+(https://github.com/LPC-HH/HH4b/pull/366). A version update is a new dated folder and tag there;
+nothing to build. The earlier text-file pickles and their builder are in `archive/`.
 
-or:
+Run 2:
 
 ```
 python -m build_jec_run2 jec_compiled_run2.pkl.gz

src/HH4b/corrections/archive/README.md

@@ -0,0 +1,17 @@
+# Archived Run 3 text-file JECs
+
+Superseded by the official correctionlib JECs and JER (2024 and 2025 since #363, 2022 and 2023 since
+https://github.com/LPC-HH/HH4b/pull/366): `processors/corrections.py` (`cat_jerc`) reads every Run 3
+year from the JME correctionlib files on CVMFS through `processors/jec_clib.py`. Nothing in this
+folder is read by the code; it is kept for reference and for reproducing older skims.
+
+- `jec_compiled.pkl.gz`: coffea JEC stacks built from text files by `build_jec.py`: 2022/2022EE
+  Summer22(EE)_22Sep2023 V2 with JER JRV1, 2023/2023BPix Summer23(BPix)Prompt23 V1 with JER JRV1,
+  and provisional 2024/2025 stacks (Summer24Prompt24 RunCnib1 V1 residual for all data, AK4 files
+  for AK8, 2023BPix JER). The skims made before the switch (e.g. `20260820_glopartv3_v15_signal`)
+  used it.
+- `jec_compiled_upto2023BPix.pkl.gz`: an older build with 2022-2023 only.
+- `build_jec.py`: the builder (`python -m build_jec --output jec_compiled.pkl.gz` in this folder,
+  with the text files in `data/jecs/`, which are not in the repository).
+
+Run 2 still uses `../jec_compiled_run2.pkl.gz`, built by `../build_jec_run2.py`.

src/HH4b/corrections/build_jec.py → src/HH4b/corrections/archive/build_jec.py

@@ -34,8 +34,8 @@ def jet_factory_factory(files):
     return CorrectedJetsFactory(jec_name_map, jec_stack)
 
 
-# 2024 and 2025: the processors take the official correctionlib JECs/JER instead
-# (processors/corrections.py cat_corrections, processors/jec_clib.py); these entries are unused
+# Run 3 (2022-2025): the processors take the official correctionlib JECs/JER instead
+# (processors/corrections.py cat_jerc, processors/jec_clib.py); these entries are unused
 jet_factory = {
     "2022mc": jet_factory_factory(
         files=[

src/HH4b/hh_vars.py

@@ -488,6 +488,10 @@
 jec_vars = [
     "bbFatJetPt",
     "VBFJetPt",
+    "VBFJetMass",
+    "AK4JetAwayPt",
+    "AK4JetAwayMass",
+    "MET_pt",
     "bdt_score",
     "bdt_score_vbf",
     "HHPt",

src/HH4b/postprocessing/CreateDatacard.py

@@ -491,9 +491,9 @@ def templates_have_split_shared_mc(templates_dir: str) -> bool:
 # 2025 JER: PLACEHOLDER borrowed from 2024. 2025 has no MC of its own (PostProcess
 # --split-shared-mc builds it from half of the 2024 MC) and no JER measurement, so its JER
 # templates carry the 2024 MC JER variation (the 2024 MC JER SFs, which are the 2023BPix
-# Summer23BPixPrompt23 RunD JRV1 SFs, see corrections/build_jec.py). It gets its own nuisance,
-# CMS_res_j_2025, uncorrelated with the other eras like every JER nuisance; replace the values
-# once 2025 MC and 2025 JER SFs exist. Cards without 2025 are unchanged.
+# Summer23BPixPrompt23 RunD JRV1 SFs, see corrections/archive/build_jec.py). It gets its own
+# nuisance, CMS_res_j_2025, uncorrelated with the other eras like every JER nuisance; replace the
+# values once 2025 MC and 2025 JER SFs exist. Cards without 2025 are unchanged.
 if "2025" in years:
     uncorr_year_shape_systs["JER"].uncorr_years["2025"] = ["2025"]
     logging.info("2025: separate JER nuisance (placeholder: 2024 JER variation of the 2024 MC)")

src/HH4b/postprocessing/postprocessing.py

@@ -201,6 +201,19 @@ class Region:
     ]
 load_columns_syst += [("bbFatJetParTmassVis_raw", 2)]
 
+# JES/JER shifts of the VBF jet mass, the AK4-away jets and the MET: read where the skims have them,
+# filled by fill_ak4_jec_shifts where they do not
+load_columns_syst_ak4 = [
+    col
+    for jshift in jec_shifts
+    for col in [
+        (f"VBFJetMass_{jshift}", 2),
+        (f"AK4JetAwayPt_{jshift}", 2),
+        (f"AK4JetAwayMass_{jshift}", 2),
+        (f"MET_pt_{jshift}", 1),
+    ]
+]
+
 # load scale weights for ttbar
 load_columns_ttbar = [
     ("scale_weights", 6),
@@ -572,6 +585,10 @@ def load_run3_samples(
         if sample not in syst_keys
     }
 
+    optional_syst_columns = (
+        utils.format_columns(load_columns_syst_ak4) if load_systematics else None
+    )
+
     # load ggf samples
     events_dict_syst_ggf = {
         **utils.load_samples(
@@ -587,6 +604,7 @@ def load_run3_samples(
                 )
                 + load_columns_jmsr
             ),
+            optional_columns=optional_syst_columns,
             reorder_txbb=reorder_txbb,
             txbb_str=txbb_str,
             ptqq_stitch=ptqq_stitch,
@@ -613,6 +631,7 @@ def load_run3_samples(
                 )
                 + load_columns_jmsr
             ),
+            optional_columns=optional_syst_columns,
             reorder_txbb=reorder_txbb,
             txbb_str=txbb_str,
             ptqq_stitch=ptqq_stitch,
@@ -643,6 +662,7 @@ def load_run3_samples(
                 )
                 + load_columns_jmsr
             ),
+            optional_columns=optional_syst_columns,
             reorder_txbb=reorder_txbb,
             txbb_str=txbb_str,
             ptqq_stitch=ptqq_stitch,
@@ -669,6 +689,7 @@ def load_run3_samples(
                 )
                 + load_columns_jmsr
             ),
+            optional_columns=optional_syst_columns,
             reorder_txbb=reorder_txbb,
             txbb_str=txbb_str,
             ptqq_stitch=ptqq_stitch,
@@ -698,16 +719,25 @@ def load_run3_samples(
         ),
     }
 
+    # every MC except QCD (syst_keys), loaded with the systematics; data and QCD MC are in
+    # events_dict_nosyst
+    events_dicts_mc_except_qcd = [
+        events_dict_ttbar,
+        events_dict_syst_bg,
+        events_dict_syst_ggf,
+        events_dict_syst_vbf,
+    ]
+
+    if load_systematics:
+        for events_dict_mc in events_dicts_mc_except_qcd:
+            for key in events_dict_mc:
+                events_dict_mc[key] = fill_ak4_jec_shifts(events_dict_mc[key], key)
+
     if jmsr_correct:
         if scale_and_smear:
             raise ValueError("jmsr_correct and scale_and_smear both re-derive the mass; pick one")
-        # all MC except QCD (samples_nosyst = data + QCD MC are never corrected)
-        for events_dict_mc in [
-            events_dict_ttbar,
-            events_dict_syst_bg,
-            events_dict_syst_ggf,
-            events_dict_syst_vbf,
-        ]:
+        # data and QCD MC are never corrected
+        for events_dict_mc in events_dicts_mc_except_qcd:
             for key, events in events_dict_mc.items():
                 jmsr_correct_mass(
                     events, jmsr_year if jmsr_year is not None else year, mass_str, key
@@ -738,6 +768,41 @@ def load_run3_samples(
     return events_dict
 
 
+def fill_ak4_jec_shifts(events: pd.DataFrame, key: str) -> pd.DataFrame:
+    """The JES/JER shifts of ``load_columns_syst_ak4`` for the events of skims made before they were
+    saved (the columns are missing, or NaN for the rows of such skims). The VBF jet mass follows
+    from its saved shifted pT, since the jets' CorrectedJetsFactory scales pT and mass by the same
+    factor; the AK4-away jets and the MET keep their nominal values, as before."""
+    fallback = {}
+    for jshift in jec_shifts:
+        for i in range(2):
+            pt = events[("VBFJetPt", i)].to_numpy()
+            ratio = np.divide(
+                events[(f"VBFJetPt_{jshift}", i)].to_numpy(), pt, out=np.ones(len(pt)), where=pt > 0
+            )
+            fallback[(f"VBFJetMass_{jshift}", i)] = events[("VBFJetMass", i)].to_numpy() * ratio
+            for var in ["AK4JetAwayPt", "AK4JetAwayMass"]:
+                fallback[(f"{var}_{jshift}", i)] = events[(var, i)].to_numpy()
+        fallback[(f"MET_pt_{jshift}", 0)] = events[("MET_pt", 0)].to_numpy()
+
+    filled, missing_cols = set(), {}
+    for col, vals in fallback.items():
+        if col not in events.columns:
+            missing_cols[col] = vals
+            filled.add(col[0].rsplit("_", 2)[0])
+        elif events[col].isna().any():
+            events[col] = events[col].fillna(pd.Series(vals, index=events.index))
+            filled.add(col[0].rsplit("_", 2)[0])
+    if missing_cols:
+        events = pd.concat([events, pd.DataFrame(missing_cols, index=events.index)], axis=1)
+    if filled:
+        print(
+            f"{key}: JES/JER shifts of {', '.join(sorted(filled))} not in (all of) the skims: VBF "
+            "jet mass from its shifted pT, AK4-away jets and MET nominal"
+        )
+    return events
+
+
 def jmsr_correct_mass(
     events: pd.DataFrame, year: str, mass_str: str, key: str, seed: int = 42
 ) -> None:

src/HH4b/processors/bbbbSkimmer.py

@@ -708,30 +708,33 @@ def process(self, events: ak.Array):
             jets = events.Jet
             jec_shifted_jetvars = {}
 
-        if JEC_loader.met_factory is not None:
-            # check if "MET" attribute exists
-            if hasattr(events, "MET"):
-                events_met = events.MET
-            elif hasattr(events, "PuppiMET"):
-                events_met = events.PuppiMET
-                # No deltaX and deltaY in PuppiMET, so we have to calculate them
-                # by definition: up - nominal
-                deltaX_up = events_met.ptUnclusteredUp * np.cos(events_met.phiUnclusteredUp)
-                deltaY_up = events_met.ptUnclusteredUp * np.sin(events_met.phiUnclusteredUp)
-                deltaX_nom = events_met.pt * np.cos(events_met.phi)
-                deltaY_nom = events_met.pt * np.sin(events_met.phi)
-                events_met["MetUnclustEnUpDeltaX"] = deltaX_up - deltaX_nom
-                events_met["MetUnclustEnUpDeltaY"] = deltaY_up - deltaY_nom
-            else:
-                raise AttributeError("Neither 'MET' nor 'PuppiMET' attribute found in events.")
-            met = JEC_loader.met_factory.build(events_met, jets, {}) if isData else events_met
+        # MET: Type-1 PUPPI MET from RawPuppiMET with the JECs re-applied above (NanoAOD PuppiMET
+        # carries the JECs of its production); PuppiMET and RawPuppiMET are saved too. In MC, also
+        # its JES/JER/unclustered-energy variations, which the shifted BDTs use.
+        met_vars = {}
+        met_pt_shifted = {}
+        jer_smear_factor = None
+        if not isData and isJECs and jec_shifted_jetvars:
+            jer_smear_factor = jets.jet_energy_resolution_correction
+        met_type1 = JEC_loader.type1_puppi_met(events, year, isData, jer_smear_factor)
+        if met_type1 is not None:
+            met_pt, met_phi = met_type1[""]
+            for shift, (pt, phi) in met_type1.items():
+                if shift:
+                    met_vars[f"MET_pt_{shift}"] = pt
+                    met_vars[f"MET_phi_{shift}"] = phi
+                    met_pt_shifted[shift] = pt
+            for name in ["PuppiMET", "RawPuppiMET"]:
+                met_vars[f"{name}_pt"] = events[name].pt.to_numpy()
+                met_vars[f"{name}_phi"] = events[name].phi.to_numpy()
         else:
             if hasattr(events, "MET"):
                 met = events.MET
             elif hasattr(events, "PuppiMET"):
                 met = events.PuppiMET
             else:
                 raise AttributeError("Neither 'MET' nor 'PuppiMET' attribute found in events.")
+            met_pt, met_phi = met.pt.to_numpy(), met.phi.to_numpy()
 
         jets = good_ak4jets(jets, year, self._nano_version)
         ht = ak.sum(jets.pt, axis=1)
@@ -898,6 +901,20 @@ def process(self, events: ak.Array):
                     if shift != "":
                         ak4JetVars[f"ak4Jet{key}_{shift}"] = pad_val(vals, num_jets, axis=1)
 
+            # AK4JetAway JEC variables
+            away_jets = (
+                ak4_jets_awayfromak8 if len(ak4_jets_awayfromak8) == 2 else [ak4_jets_awayfromak8]
+            )
+            for var in ["pt", "mass"]:
+                key = self.skim_vars["Jet"][var]
+                for label, shift in self.jecs.items():
+                    if shift in ak.fields(away_jets[0]):
+                        for vari in ["up", "down"]:
+                            vals = [jets_away[shift][vari][var] for jets_away in away_jets]
+                            ak4JetAwayVars[f"AK4JetAway{key}_{label}_{vari}"] = pad_val(
+                                ak.concatenate(vals, axis=1), 2, axis=1
+                            )
+
             # FatJet JEC variables
             for var in ["pt"]:
                 key = self.skim_vars["FatJet"][var]
@@ -915,13 +932,14 @@ def process(self, events: ak.Array):
                     bbFatJetVars[f"bbFatJet{key}{label}"] = vals
 
         # Event variables
-        met_pt = met.pt
         eventVars = {
             key: events[val].to_numpy()
             for key, val in self.skim_vars["Event"].items()
             if key in events.fields
         }
-        eventVars["MET_pt"] = met_pt.to_numpy()
+        eventVars["MET_pt"] = met_pt
+        eventVars["MET_phi"] = met_phi
+        eventVars.update(met_vars)
         eventVars["ht"] = ht.to_numpy()
         eventVars["nJets"] = ak.sum(jets_sel, axis=1).to_numpy()
         eventVars["nFatJets"] = ak.num(fatjets).to_numpy()
@@ -1002,7 +1020,7 @@ def process(self, events: ak.Array):
 
         # JEC variations for VBF Jets
         if self._region == "signal" and isJECs:
-            for var in ["pt"]:
+            for var in ["pt", "mass"]:
                 key = self.skim_vars["Jet"][var]
                 for label, shift in self.jecs.items():
                     if shift in ak.fields(vbf_jets):
@@ -1038,7 +1056,13 @@ def process(self, events: ak.Array):
                 ]
             for jshift in jshifts:
                 try:
-                    bdtVars = self.getBDT(bbFatJetVars, vbfJetVars, ak4JetAwayVars, met_pt, jshift)
+                    bdtVars = self.getBDT(
+                        bbFatJetVars,
+                        vbfJetVars,
+                        ak4JetAwayVars,
+                        met_pt_shifted.get(jshift, met_pt),
+                        jshift,
+                    )
                 except KeyError as e:
                     # A chunk with no fat jets never has its JEC/JMSR-shifted fat-jet
                     # branches built (those live inside the `if shift != ""` blocks above),