From ecc81e1f3d2c948745a1532f944aa2774d8a60ee Mon Sep 17 00:00:00 2001 From: Zichun Hao Date: Fri, 23 Oct 2026 22:57:57 -0700 Subject: [PATCH 1/6] Skimmer: 2022-2023 JECs and JER from the official correctionlib files corrections.JECs now builds the jet factories of every Run 3 year from the CMS analysis-corrections JSONs (jec_clib), pinned to the latest dated JME folders: 2022/2022EE Summer22(EE)_22Sep2023_V4 with JER JRV2, 2023/2023BPix Summer23(BPix)Prompt23_V4 with JER RunCv1234/RunD_JRV3 (run-dependent data residuals), 2024/2025 unchanged. The compiled text-file JECs are read for Run 2 only. Checked on 2022 and 2023 MC and data: correctionlib reproduces the text-file JECs of the same version to a few 1e-6, and the update to V4 leaves the AK8 jets unchanged; AK4 jets move only at 4.25 < |eta| < 4.5 in 2022 (below 0.2%) and at raw pT 15-30 GeV, 1.3 < |eta| < 2.5 in 2023 data (about -1%). --- src/HH4b/corrections/build_jec.py | 4 +- src/HH4b/processors/corrections.py | 95 ++++++++++++++++++------------ src/HH4b/processors/jec_clib.py | 5 +- 3 files changed, 61 insertions(+), 43 deletions(-) diff --git a/src/HH4b/corrections/build_jec.py b/src/HH4b/corrections/build_jec.py index 3e8b7a8c..8010759c 100644 --- a/src/HH4b/corrections/build_jec.py +++ b/src/HH4b/corrections/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=[ diff --git a/src/HH4b/processors/corrections.py b/src/HH4b/processors/corrections.py index 0922ea08..9d6975e1 100644 --- a/src/HH4b/processors/corrections.py +++ b/src/HH4b/processors/corrections.py @@ -21,6 +21,7 @@ import numpy as np import uproot from coffea.analysis_tools import Weights +from coffea.jetmet_tools import CorrectedMETFactory from coffea.nanoevents.methods import vector from coffea.nanoevents.methods.base import NanoEventsArray from coffea.nanoevents.methods.nanoaod import FatJetArray, JetArray @@ -78,6 +79,46 @@ }, } +""" +JECs and JER of every Run 3 year from the correctionlib files of the CAT area (processors/jec_clib.py), +no text files: pinned dated folders with the latest versions (identical to "latest" on 2026-10-20). +The 2022-2023 data residuals of these versions are run dependent, as the 2024 and 2025 ones. +""" +cat_jerc = { + "2022": { + "jme": f"{cat_correction_path}/JME/Run3-22CDSep23-Summer22-NanoAODv12/2026-06-05", + "jec": "Summer22_22Sep2023_V4", + "jer": "Summer22_22Sep2023_JRV2", + }, + "2022EE": { + "jme": f"{cat_correction_path}/JME/Run3-22EFGSep23-Summer22EE-NanoAODv12/2026-06-05", + "jec": "Summer22EE_22Sep2023_V4", + "jer": "Summer22EE_22Sep2023_JRV2", + }, + "2023": { + "jme": f"{cat_correction_path}/JME/Run3-23CSep23-Summer23-NanoAODv12/2026-07-15", + "jec": "Summer23Prompt23_V4", + "jer": "Summer23Prompt23_RunCv1234_JRV3", + }, + "2023BPix": { + "jme": f"{cat_correction_path}/JME/Run3-23DSep23-Summer23BPix-NanoAODv12/2026-07-15", + "jec": "Summer23BPixPrompt23_V4", + "jer": "Summer23BPixPrompt23_RunD_JRV3", + }, + **{y: {k: cat[k] for k in ("jme", "jec", "jer")} for y, cat in cat_corrections.items()}, +} + +# input names of CorrectedMETFactory -> MET and jet fields (as corrections/build_jec.py) +met_name_map = { + "JetPt": "pt", + "JetPhi": "phi", + "ptRaw": "pt_raw", + "METpt": "pt", + "METphi": "phi", + "UnClusteredEnergyDeltaX": "MetUnclustEnUpDeltaX", + "UnClusteredEnergyDeltaY": "MetUnclustEnUpDeltaY", +} + def get_UL_year(year: str) -> str: return f"{year}_UL" @@ -284,39 +325,30 @@ def get_scale_weights(events): class JECs: def __init__(self, year): - # TODO: add more years when available - if year in ["2022", "2022EE", "2023", "2023BPix", "2024", "2025"]: - jec_compiled = package_path + "/corrections/jec_compiled.pkl.gz" - elif year in ["2016", "2016APV", "2017", "2018"]: - jec_compiled = package_path + "/corrections/jec_compiled_run2.pkl.gz" - else: - jec_compiled = None - self.jet_factory = {} self.met_factory = None - if jec_compiled is not None: - with gzip.open(jec_compiled, "rb") as filehandler: - jmestuff = pickle.load(filehandler) - - self.jet_factory["ak4"] = jmestuff["jet_factory"] - self.jet_factory["ak8"] = jmestuff["fatjet_factory"] - self.met_factory = jmestuff["met_factory"] - - if year in cat_corrections: - # official correctionlib JECs and JER (run-dependent data residuals, AK8 MC truth) - # replace the text-file ones of the pickle, which predate them - cat = cat_corrections[year] + if year in cat_jerc: + # official correctionlib JECs and JER (run-dependent data residuals) + cat = cat_jerc[year] for jet_type, (jerc_file, algo) in { "ak4": ("jet_jerc", "AK4PFPuppi"), "ak8": ("fatJet_jerc", "AK8PFPuppi"), }.items(): path = f"{cat['jme']}/{jerc_file}.json.gz" self.jet_factory[jet_type] = { - **self.jet_factory[jet_type], f"{year}mc": clib_jet_factory(path, cat["jec"], cat["jer"], algo, isData=False), year: clib_jet_factory(path, cat["jec"], cat["jer"], algo, isData=True), } + self.met_factory = CorrectedMETFactory(met_name_map) + elif year in ["2016", "2016APV", "2017", "2018"]: + jec_compiled = package_path + "/corrections/jec_compiled_run2.pkl.gz" + with gzip.open(jec_compiled, "rb") as filehandler: + jmestuff = pickle.load(filehandler) + + self.jet_factory["ak4"] = jmestuff["jet_factory"] + self.jet_factory["ak8"] = jmestuff["fatjet_factory"] + self.met_factory = jmestuff["met_factory"] def _add_jec_variables( self, jets: JetArray, event_rho: ak.Array, isData: bool, run: ak.Array | None = None @@ -354,7 +386,7 @@ def get_jec_jets( if "Rho" in events.fields else events.fixedGridRhoFastjetAll ) - run = events.run if (isData and year in cat_corrections) else None + run = events.run if (isData and year in cat_jerc) else None jets = self._add_jec_variables(jets, rho, isData, run) apply_jecs = ak.any(jets.pt) if (applyData or not isData) else False @@ -379,22 +411,9 @@ def get_jec_jets( jec_cache = cachetools.Cache(np.inf) if isData: - if year == "2022": - corr_key = f"{year}_runCD" - elif year == "2022EE" and "Run2022E" in dataset: - corr_key = f"{year}_runE" - elif year == "2022EE" and "Run2022F" in dataset: - corr_key = f"{year}_runF" - elif year == "2022EE" and "Run2022G" in dataset: - corr_key = f"{year}_runG" - elif year == "2023": - corr_key = "2023_runCv4" if "Run2023Cv4" in dataset else "2023_runCv123" - elif year == "2023BPix": - corr_key = "2023BPix_runD" - elif year == "2024": - corr_key = "2024" - elif year == "2025": - corr_key = "2025" + if year in cat_jerc: + # run-dependent residuals: one correction for all eras of the year + corr_key = year else: print(dataset, year) print("warning, no valid dataset, JECs won't be applied to data") diff --git a/src/HH4b/processors/jec_clib.py b/src/HH4b/processors/jec_clib.py index a8077633..8d40b0c5 100644 --- a/src/HH4b/processors/jec_clib.py +++ b/src/HH4b/processors/jec_clib.py @@ -1,7 +1,6 @@ """ -JECs and JER from the official correctionlib files, for eras whose corrections are published only in -that format (2024 and later): correctionlib-backed stand-ins for the four parts of a coffea 0.7 -``JECStack``. The unchanged ``CorrectedJetsFactory`` then applies the nominal JEC, the hybrid JER +JECs and JER of every Run 3 year from the official correctionlib files: correctionlib-backed +stand-ins for the four parts of a coffea 0.7 ``JECStack``. The unchanged ``CorrectedJetsFactory`` then applies the nominal JEC, the hybrid JER smearing and the JES/JER variations, and the corrected jets keep the fields of the text-file path (``pt``, ``mass``, ``JES_jes``, ``JER``). From 0b6696373c3a7955e2d9f8ba9a668fd7280758bf Mon Sep 17 00:00:00 2001 From: Zichun Hao Date: Fri, 23 Oct 2026 23:11:57 -0700 Subject: [PATCH 2/6] Archive the Run 3 text-file JECs Since the correctionlib JECs of every Run 3 year (#363 for 2024/2025, #366 for 2022/2023), jec_compiled.pkl.gz, jec_compiled_upto2023BPix.pkl.gz and their builder build_jec.py are no longer read: moved to corrections/archive/ with a README on their contents and the skims that used them. The corrections README describes the correctionlib JECs; comments point to the new path. Run 2 keeps jec_compiled_run2.pkl.gz and build_jec_run2.py. --- src/HH4b/corrections/README.md | 11 ++++++----- src/HH4b/corrections/archive/README.md | 17 +++++++++++++++++ .../corrections/{ => archive}/build_jec.py | 0 .../{ => archive}/jec_compiled.pkl.gz | Bin .../jec_compiled_upto2023BPix.pkl.gz | Bin src/HH4b/postprocessing/CreateDatacard.py | 6 +++--- src/HH4b/processors/corrections.py | 2 +- src/HH4b/processors/jec_clib.py | 4 ++-- 8 files changed, 29 insertions(+), 11 deletions(-) create mode 100644 src/HH4b/corrections/archive/README.md rename src/HH4b/corrections/{ => archive}/build_jec.py (100%) rename src/HH4b/corrections/{ => archive}/jec_compiled.pkl.gz (100%) rename src/HH4b/corrections/{ => archive}/jec_compiled_upto2023BPix.pkl.gz (100%) diff --git a/src/HH4b/corrections/README.md b/src/HH4b/corrections/README.md index 3cee1bcb..169f31a3 100644 --- a/src/HH4b/corrections/README.md +++ b/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 diff --git a/src/HH4b/corrections/archive/README.md b/src/HH4b/corrections/archive/README.md new file mode 100644 index 00000000..db029012 --- /dev/null +++ b/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`. diff --git a/src/HH4b/corrections/build_jec.py b/src/HH4b/corrections/archive/build_jec.py similarity index 100% rename from src/HH4b/corrections/build_jec.py rename to src/HH4b/corrections/archive/build_jec.py diff --git a/src/HH4b/corrections/jec_compiled.pkl.gz b/src/HH4b/corrections/archive/jec_compiled.pkl.gz similarity index 100% rename from src/HH4b/corrections/jec_compiled.pkl.gz rename to src/HH4b/corrections/archive/jec_compiled.pkl.gz diff --git a/src/HH4b/corrections/jec_compiled_upto2023BPix.pkl.gz b/src/HH4b/corrections/archive/jec_compiled_upto2023BPix.pkl.gz similarity index 100% rename from src/HH4b/corrections/jec_compiled_upto2023BPix.pkl.gz rename to src/HH4b/corrections/archive/jec_compiled_upto2023BPix.pkl.gz diff --git a/src/HH4b/postprocessing/CreateDatacard.py b/src/HH4b/postprocessing/CreateDatacard.py index 196b9044..0c14fec7 100644 --- a/src/HH4b/postprocessing/CreateDatacard.py +++ b/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)") diff --git a/src/HH4b/processors/corrections.py b/src/HH4b/processors/corrections.py index 9d6975e1..baad5860 100644 --- a/src/HH4b/processors/corrections.py +++ b/src/HH4b/processors/corrections.py @@ -108,7 +108,7 @@ **{y: {k: cat[k] for k in ("jme", "jec", "jer")} for y, cat in cat_corrections.items()}, } -# input names of CorrectedMETFactory -> MET and jet fields (as corrections/build_jec.py) +# input names of CorrectedMETFactory -> MET and jet fields (as corrections/archive/build_jec.py) met_name_map = { "JetPt": "pt", "JetPhi": "phi", diff --git a/src/HH4b/processors/jec_clib.py b/src/HH4b/processors/jec_clib.py index 8d40b0c5..8037ad60 100644 --- a/src/HH4b/processors/jec_clib.py +++ b/src/HH4b/processors/jec_clib.py @@ -16,8 +16,8 @@ import numpy as np from coffea.jetmet_tools import CorrectedJetsFactory -# correctionlib input names -> jet fields; as build_jec.jec_name_map, plus the run number for the -# run-dependent data residuals +# correctionlib input names -> jet fields; as jec_name_map of corrections/archive/build_jec.py, +# plus the run number for the run-dependent data residuals jec_name_map = { "JetPt": "pt", "JetMass": "mass", From 25c84c37e733152901006de71c4f315f71c63697 Mon Sep 17 00:00:00 2001 From: Zichun Hao Date: Sat, 24 Oct 2026 14:57:48 -0700 Subject: [PATCH 3/6] Skimmer: Type-1 PUPPI MET recomputed with the re-applied JECs The data MET went through coffea's CorrectedMETFactory, which in coffea <= 0.7.26 (we pin 0.7.21) adds the change of the jet pT relative to the raw pT with the wrong sign: applied on top of the Type-1 PuppiMET, it undid the Type-1 correction (data MET close to RawPuppiMET), while MC kept NanoAOD PuppiMET. NanoAOD PuppiMET itself carries the Type-1 correction of the JECs of its production, not the ones the skimmer re-applies to the jets. For the correctionlib (Run 3) years the MET is now the Type-1 PUPPI MET recomputed from RawPuppiMET with the re-applied AK4 PUPPI JECs, as CMSSW does for PUPPI (Jet and CorrT1METJet, muon-subtracted corrected pT above 15 GeV, EM fraction below 0.9, unsmeared jets), for data and MC alike. MET_pt and the new MET_phi hold it; PuppiMET_pt/phi and RawPuppiMET_pt/phi are saved as well. Other years keep the NanoAOD MET, also for data. ttSkimmer gets the same MET. With the JECs of the production it reproduces NanoAOD PuppiMET (CMSSW) to 0.3 GeV RMS (2022 and 2024 ttbar), and its output is identical to the Type-1 MET of coffea's CorrectedMETFactory (scikit-hep/coffea#1537), as used in DAZSLE/hbb-run3#74. --- src/HH4b/processors/bbbbSkimmer.py | 32 +++++++++----------- src/HH4b/processors/corrections.py | 12 +++++++- src/HH4b/processors/jec_clib.py | 47 ++++++++++++++++++++++++++++++ src/HH4b/processors/ttSkimmer.py | 20 +++++++++---- 4 files changed, 86 insertions(+), 25 deletions(-) diff --git a/src/HH4b/processors/bbbbSkimmer.py b/src/HH4b/processors/bbbbSkimmer.py index 9742728d..71051388 100644 --- a/src/HH4b/processors/bbbbSkimmer.py +++ b/src/HH4b/processors/bbbbSkimmer.py @@ -708,23 +708,15 @@ 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. + met_vars = {} + met_type1 = JEC_loader.type1_puppi_met(events, year, isData) + if met_type1 is not None: + met_pt, met_phi = met_type1 + 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 @@ -732,6 +724,7 @@ def process(self, events: ak.Array): 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) @@ -915,13 +908,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() diff --git a/src/HH4b/processors/corrections.py b/src/HH4b/processors/corrections.py index baad5860..b9826ba6 100644 --- a/src/HH4b/processors/corrections.py +++ b/src/HH4b/processors/corrections.py @@ -26,7 +26,7 @@ from coffea.nanoevents.methods.base import NanoEventsArray from coffea.nanoevents.methods.nanoaod import FatJetArray, JetArray -from .jec_clib import clib_jet_factory +from .jec_clib import clib_jet_factory, type1_puppi_met from .utils import pad_val ak.behavior.update(vector.behavior) @@ -350,6 +350,16 @@ def __init__(self, year): self.jet_factory["ak8"] = jmestuff["fatjet_factory"] self.met_factory = jmestuff["met_factory"] + def type1_puppi_met( + self, events: NanoEventsArray, year: str, isData: bool + ) -> tuple[np.ndarray, np.ndarray] | None: + """``jec_clib.type1_puppi_met`` with the year's correctionlib JECs; None without them or + without RawPuppiMET.""" + if year not in cat_jerc or "RawPuppiMET" not in events.fields: + return None + cat = cat_jerc[year] + return type1_puppi_met(events, f"{cat['jme']}/jet_jerc.json.gz", cat["jec"], isData) + def _add_jec_variables( self, jets: JetArray, event_rho: ak.Array, isData: bool, run: ak.Array | None = None ) -> JetArray: diff --git a/src/HH4b/processors/jec_clib.py b/src/HH4b/processors/jec_clib.py index 8037ad60..1638eedc 100644 --- a/src/HH4b/processors/jec_clib.py +++ b/src/HH4b/processors/jec_clib.py @@ -132,3 +132,50 @@ def clib_jet_factory( ), ) return CorrectedJetsFactory(dict(jec_name_map), stack) + + +# TODO: use coffea's CorrectedMETFactory (Type-1 mode) once on coffea >= 2026.9.0 or a 0.7.x +# release after 0.7.31 +def type1_puppi_met(events, path: str, jec_tag: str, isData: bool) -> tuple[np.ndarray, np.ndarray]: + """(pt, phi) of the Type-1 PUPPI MET from RawPuppiMET with the AK4 PUPPI JEC ``jec_tag`` of + ``path``, as in CMSSW: every Jet and CorrT1METJet with muon-subtracted corrected pT > 15 GeV and + EM fraction < 0.9 shifts it by -(f_L1L2L3Res - f_L1) x muon-subtracted raw pT; no smearing. + Adapted from coffea's Type-1 CorrectedMETFactory (scikit-hep/coffea#1537); same output as + DAZSLE/hbb-run3#74 (correct_met).""" + cset = _correction_set(path) + level = "DATA" if isData else "MC" + corr_full = cset.compound[f"{jec_tag}_{level}_L1L2L3Res_AK4PFPuppi"] + corr_l1 = cset[f"{jec_tag}_{level}_L1FastJet_AK4PFPuppi"] + nev = len(events) + rho = _to_numpy(events.Rho.fixedGridRhoFastjetAll) + run = _to_numpy(events.run) + raw_met_pt = _to_numpy(events.RawPuppiMET.pt) + raw_met_phi = _to_numpy(events.RawPuppiMET.phi) + met_x = raw_met_pt * np.cos(raw_met_phi) + met_y = raw_met_pt * np.sin(raw_met_phi) + jets, low_pt_jets = events.Jet, events.CorrT1METJet + for coll, raw_pt, em_fraction in [ + (jets, (1 - jets.rawFactor) * jets.pt, jets.chEmEF + jets.neEmEF), + (low_pt_jets, low_pt_jets.rawPt, low_pt_jets.EmEF), + ]: + idx = np.repeat(np.arange(nev), ak.to_numpy(ak.num(raw_pt))) + raw = _to_numpy(ak.flatten(raw_pt)) + inputs = { + "JetA": _to_numpy(ak.flatten(coll.area)), + "JetEta": _to_numpy(ak.flatten(coll.eta)), + "JetPhi": _to_numpy(ak.flatten(coll.phi)), + "JetPt": raw, + "Rho": rho[idx], + "run": run[idx], + } + f = corr_full.evaluate(*[inputs[inp.name] for inp in corr_full.inputs]) + f_l1 = corr_l1.evaluate(*[inputs[inp.name] for inp in corr_l1.inputs]) + raw_nomu = raw * (1 - _to_numpy(ak.flatten(coll.muonSubtrFactor))) + phi = inputs["JetPhi"] + if "muonSubtrDeltaPhi" in coll.fields: + phi = phi + _to_numpy(ak.flatten(coll.muonSubtrDeltaPhi)) + use = (f * raw_nomu > 15) & (_to_numpy(ak.flatten(em_fraction)) < 0.9) + shift = np.where(use, (f - f_l1) * raw_nomu, 0.0) + met_x -= np.bincount(idx, weights=shift * np.cos(phi), minlength=nev) + met_y -= np.bincount(idx, weights=shift * np.sin(phi), minlength=nev) + return np.hypot(met_x, met_y), np.arctan2(met_y, met_x) diff --git a/src/HH4b/processors/ttSkimmer.py b/src/HH4b/processors/ttSkimmer.py index 3dea0a1d..099146cf 100644 --- a/src/HH4b/processors/ttSkimmer.py +++ b/src/HH4b/processors/ttSkimmer.py @@ -267,7 +267,17 @@ def process(self, events: ak.Array): nano_version=self._nano_version, ) - met = JEC_loader.met_factory.build(events.MET, ak4_jets, {}) if isData else events.MET + # MET as in bbbbSkimmer + met_vars = {} + met_type1 = JEC_loader.type1_puppi_met(events, year, isData) + if met_type1 is not None: + met_pt, met_phi = met_type1 + 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: + met = events.MET if hasattr(events, "MET") else events.PuppiMET + met_pt, met_phi = met.pt.to_numpy(), met.phi.to_numpy() num_fatjets = 2 # number to save fatjets = get_ak8jets(events.FatJet) @@ -346,14 +356,14 @@ def process(self, events: ak.Array): print("Lepton vars", f"{time.time() - start:.2f}") # event variable - 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) skimmed_events = { **genVars, @@ -404,7 +414,7 @@ def process(self, events: ak.Array): ) # MET - met_selection = met.pt >= self.met_selection["pt"] + met_selection = met_pt >= self.met_selection["pt"] cut_metfilters = np.ones(len(events), dtype="bool") for mf in self.met_filters: From 99c2b5e0fd956a56645e097fc49dc1508511e164 Mon Sep 17 00:00:00 2001 From: Zichun Hao Date: Sat, 24 Oct 2026 14:59:33 -0700 Subject: [PATCH 4/6] Skimmer: JES, JER and unclustered-energy variations of the Type-1 MET In MC the skims now also save MET_pt/MET_phi for JES_up/down, JER_up/down and UES_up/down, and the JES/JER-shifted BDT scores use the shifted MET instead of the nominal. The variations are those of the unsmeared Type-1 MET in CMSJMECalculators (Type1METVariationsCalculator, isT1SmearedMET false), the same as the MET_T1 variations of nanoAOD-tools' jetmetUncertainties, for the jets of the Type-1 sum at their corrected pT (muon-subtracted part corrected, muon added back): JES up/down puts each of them at that pT x (1 +- the total JES uncertainty there); JER up and down are both the Type-1 MET with the Jets smeared by their nominal smearing factors (the CorrT1METJets unsmeared, as coffea's Type-1 MET does without their factors); UES up/down adds/subtracts the unclustered-energy shift of NanoAOD PuppiMET. coffea's Type-1 CorrectedMETFactory instead scales the jets by the JER-varied over the nominal smeared pT, which blows up for jets that the nominal smearing pushes to ~0 GeV. --- src/HH4b/processors/bbbbSkimmer.py | 24 ++++++++-- src/HH4b/processors/corrections.py | 11 +++-- src/HH4b/processors/jec_clib.py | 77 ++++++++++++++++++++++++++---- src/HH4b/processors/ttSkimmer.py | 2 +- 4 files changed, 96 insertions(+), 18 deletions(-) diff --git a/src/HH4b/processors/bbbbSkimmer.py b/src/HH4b/processors/bbbbSkimmer.py index 71051388..3bc29738 100644 --- a/src/HH4b/processors/bbbbSkimmer.py +++ b/src/HH4b/processors/bbbbSkimmer.py @@ -709,11 +709,21 @@ def process(self, events: ak.Array): jec_shifted_jetvars = {} # 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. + # 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_type1 = JEC_loader.type1_puppi_met(events, year, isData) + 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 + 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() @@ -1032,7 +1042,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), diff --git a/src/HH4b/processors/corrections.py b/src/HH4b/processors/corrections.py index b9826ba6..befe2470 100644 --- a/src/HH4b/processors/corrections.py +++ b/src/HH4b/processors/corrections.py @@ -351,14 +351,19 @@ def __init__(self, year): self.met_factory = jmestuff["met_factory"] def type1_puppi_met( - self, events: NanoEventsArray, year: str, isData: bool - ) -> tuple[np.ndarray, np.ndarray] | None: + self, + events: NanoEventsArray, + year: str, + isData: bool, + jer_smear_factor: ak.Array | None = None, + ) -> dict[str, tuple[np.ndarray, np.ndarray]] | None: """``jec_clib.type1_puppi_met`` with the year's correctionlib JECs; None without them or without RawPuppiMET.""" if year not in cat_jerc or "RawPuppiMET" not in events.fields: return None cat = cat_jerc[year] - return type1_puppi_met(events, f"{cat['jme']}/jet_jerc.json.gz", cat["jec"], isData) + path = f"{cat['jme']}/jet_jerc.json.gz" + return type1_puppi_met(events, path, cat["jec"], isData, jer_smear_factor) def _add_jec_variables( self, jets: JetArray, event_rho: ak.Array, isData: bool, run: ak.Array | None = None diff --git a/src/HH4b/processors/jec_clib.py b/src/HH4b/processors/jec_clib.py index 1638eedc..3999fe91 100644 --- a/src/HH4b/processors/jec_clib.py +++ b/src/HH4b/processors/jec_clib.py @@ -135,13 +135,31 @@ def clib_jet_factory( # TODO: use coffea's CorrectedMETFactory (Type-1 mode) once on coffea >= 2026.9.0 or a 0.7.x -# release after 0.7.31 -def type1_puppi_met(events, path: str, jec_tag: str, isData: bool) -> tuple[np.ndarray, np.ndarray]: - """(pt, phi) of the Type-1 PUPPI MET from RawPuppiMET with the AK4 PUPPI JEC ``jec_tag`` of - ``path``, as in CMSSW: every Jet and CorrT1METJet with muon-subtracted corrected pT > 15 GeV and - EM fraction < 0.9 shifts it by -(f_L1L2L3Res - f_L1) x muon-subtracted raw pT; no smearing. +# release after 0.7.31. It gives the same nominal and, with the total JES uncertainty passed as +# jes_uncertainties, the same JES variations; its JER variation of the unsmeared MET differs (it +# scales each jet by the JER-varied over the nominal smeared pT, which blows up for jets smeared +# to ~0 GeV). +def type1_puppi_met( + events, path: str, jec_tag: str, isData: bool, jer_smear_factor: ak.Array | None = None +) -> dict[str, tuple[np.ndarray, np.ndarray]]: + """{"": (pt, phi)} of the Type-1 PUPPI MET from RawPuppiMET with the AK4 PUPPI JEC ``jec_tag`` + of ``path``, as in CMSSW: every Jet and CorrT1METJet with muon-subtracted corrected pT > 15 GeV + and EM fraction < 0.9 shifts it by -(f_L1L2L3Res - f_L1) x muon-subtracted raw pT; no smearing. Adapted from coffea's Type-1 CorrectedMETFactory (scikit-hep/coffea#1537); same output as - DAZSLE/hbb-run3#74 (correct_met).""" + DAZSLE/hbb-run3#74 (correct_met). + + With ``jer_smear_factor`` (MC: the nominal JER smearing factor of each events.Jet), also the + variations of the unsmeared Type-1 MET as in CMSJMECalculators + (gitlab.cern.ch/cms-analysis/general/CMSJMECalculators, Type1METVariationsCalculator with + isT1SmearedMET false; the MET_T1 variations of nanoAOD-tools' jetmetUncertainties are the + same). They move the same jets, at pT_L1L2L3, the corrected pT of their muon-subtracted part + plus the muon: + - JES_up/down: each jet at pT_L1L2L3 x (1 +- delta), delta the total JES uncertainty at + (eta, pT_L1L2L3); + - JER_up = JER_down: the Jets at pT_L1L2L3 x their smearing factor, i.e. the smeared Type-1 MET + (the CorrT1METJets unsmeared, as in coffea's Type-1 MET without their smearing factors); + - UES_up/down: plus/minus the unclustered-energy shift of NanoAOD PuppiMET. + """ cset = _correction_set(path) level = "DATA" if isData else "MC" corr_full = cset.compound[f"{jec_tag}_{level}_L1L2L3Res_AK4PFPuppi"] @@ -154,9 +172,10 @@ def type1_puppi_met(events, path: str, jec_tag: str, isData: bool) -> tuple[np.n met_x = raw_met_pt * np.cos(raw_met_phi) met_y = raw_met_pt * np.sin(raw_met_phi) jets, low_pt_jets = events.Jet, events.CorrT1METJet - for coll, raw_pt, em_fraction in [ - (jets, (1 - jets.rawFactor) * jets.pt, jets.chEmEF + jets.neEmEF), - (low_pt_jets, low_pt_jets.rawPt, low_pt_jets.EmEF), + type1_jets = [] # per collection: what the variations need + for coll, raw_pt, em_fraction, smear_factor in [ + (jets, (1 - jets.rawFactor) * jets.pt, jets.chEmEF + jets.neEmEF, jer_smear_factor), + (low_pt_jets, low_pt_jets.rawPt, low_pt_jets.EmEF, None), ]: idx = np.repeat(np.arange(nev), ak.to_numpy(ak.num(raw_pt))) raw = _to_numpy(ak.flatten(raw_pt)) @@ -178,4 +197,42 @@ def type1_puppi_met(events, path: str, jec_tag: str, isData: bool) -> tuple[np.n shift = np.where(use, (f - f_l1) * raw_nomu, 0.0) met_x -= np.bincount(idx, weights=shift * np.cos(phi), minlength=nev) met_y -= np.bincount(idx, weights=shift * np.sin(phi), minlength=nev) - return np.hypot(met_x, met_y), np.arctan2(met_y, met_x) + pt_l1l2l3 = f * raw_nomu + (raw - raw_nomu) + type1_jets.append((idx, use, pt_l1l2l3, inputs["JetEta"], phi, smear_factor)) + + def polar(x, y): + return np.hypot(x, y), np.arctan2(y, x) + + out = {"": polar(met_x, met_y)} + if jer_smear_factor is None: + return out + + def vector_sum(idx, values, phi): + return ( + np.bincount(idx, weights=values * np.cos(phi), minlength=nev), + np.bincount(idx, weights=values * np.sin(phi), minlength=nev), + ) + + # sums over the jets of their pT changes: pT_L1L2L3 x delta (JES), pT_L1L2L3 x (s - 1) (JER) + jes_unc = cset[f"{jec_tag}_MC_Total_AK4PFPuppi"] + jes_x = jes_y = jer_x = jer_y = 0.0 + for idx, use, pt_l1l2l3, eta, phi, smear_factor in type1_jets: + jes_inputs = {"JetEta": eta, "JetPt": pt_l1l2l3} + delta = jes_unc.evaluate(*[jes_inputs[inp.name] for inp in jes_unc.inputs]) + dx, dy = vector_sum(idx, np.where(use, pt_l1l2l3 * delta, 0.0), phi) + jes_x, jes_y = jes_x + dx, jes_y + dy + if smear_factor is not None: + smear = _to_numpy(ak.flatten(smear_factor)) + dx, dy = vector_sum(idx, np.where(use, pt_l1l2l3 * (smear - 1), 0.0), phi) + jer_x, jer_y = jer_x + dx, jer_y + dy + out["JES_up"] = polar(met_x - jes_x, met_y - jes_y) + out["JES_down"] = polar(met_x + jes_x, met_y + jes_y) + out["JER_up"] = out["JER_down"] = polar(met_x - jer_x, met_y - jer_y) + puppi = events.PuppiMET + ues_x = _to_numpy(puppi.ptUnclusteredUp) * np.cos(_to_numpy(puppi.phiUnclusteredUp)) + ues_y = _to_numpy(puppi.ptUnclusteredUp) * np.sin(_to_numpy(puppi.phiUnclusteredUp)) + ues_x -= _to_numpy(puppi.pt) * np.cos(_to_numpy(puppi.phi)) + ues_y -= _to_numpy(puppi.pt) * np.sin(_to_numpy(puppi.phi)) + out["UES_up"] = polar(met_x + ues_x, met_y + ues_y) + out["UES_down"] = polar(met_x - ues_x, met_y - ues_y) + return out diff --git a/src/HH4b/processors/ttSkimmer.py b/src/HH4b/processors/ttSkimmer.py index 099146cf..6ed9f710 100644 --- a/src/HH4b/processors/ttSkimmer.py +++ b/src/HH4b/processors/ttSkimmer.py @@ -271,7 +271,7 @@ def process(self, events: ak.Array): met_vars = {} met_type1 = JEC_loader.type1_puppi_met(events, year, isData) if met_type1 is not None: - met_pt, met_phi = met_type1 + met_pt, met_phi = met_type1[""] 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() From d83274d168ce75fff56b3508c0edf57566502854 Mon Sep 17 00:00:00 2001 From: Zichun Hao Date: Sat, 24 Oct 2026 21:08:34 -0700 Subject: [PATCH 5/6] Skimmer: JES/JER of the AK4 jets away from the bb jets and of the VBF jet mass The skims had no JES/JER-shifted columns for the two AK4 jets away from the bb jets and none for the VBF jet mass, so in every JES/JER variation the BDT inputs H1AK4JetAway1mass and H2AK4JetAway2mass stayed nominal and the VBF dijet mass used the nominal jet masses. The MC signal-region skims now save the shifted pT and mass of the AK4-away jets (AK4JetAway{Pt,Mass}_{JES,JER}_{up,down}) and the shifted VBF jet mass (VBFJetMass_*), from the AK4 PUPPI JERC of these jets, and the skimmer's shifted BDT scores use them. In a 2024 TTto4Q chunk nothing else changes; the JES-up shift of bdt_score_vbf widens from 4.7e-3 to 5.5e-3 RMS. --- src/HH4b/processors/bbbbSkimmer.py | 16 +++++++++++++++- src/HH4b/processors/utils.py | 3 +++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/src/HH4b/processors/bbbbSkimmer.py b/src/HH4b/processors/bbbbSkimmer.py index 3bc29738..975d2eda 100644 --- a/src/HH4b/processors/bbbbSkimmer.py +++ b/src/HH4b/processors/bbbbSkimmer.py @@ -901,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] @@ -1006,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): diff --git a/src/HH4b/processors/utils.py b/src/HH4b/processors/utils.py index a18615fd..8f2b866c 100644 --- a/src/HH4b/processors/utils.py +++ b/src/HH4b/processors/utils.py @@ -70,6 +70,9 @@ jec_vars = [ "bbFatJetPt", "VBFJetPt", + "VBFJetMass", + "AK4JetAwayPt", + "AK4JetAwayMass", "HHPt", "HHeta", "HHmass", From 4e424c04afb867aa4a7a2632bef7847ac1bebdb0 Mon Sep 17 00:00:00 2001 From: Zichun Hao Date: Sat, 24 Oct 2026 22:49:16 -0700 Subject: [PATCH 6/6] Postprocessing: JES/JER of the AK4-away jets, the VBF jet mass and the MET in the BDT inputs The JES/JER-shifted BDT inputs moved only the bb jets and the VBF jet pT; the AK4-away jets, the VBF jet mass and the MET stayed nominal. jec_vars now has VBFJetMass, AK4JetAwayPt, AK4JetAwayMass and MET_pt, the v15 BDT config reads the MET through key_map, and their shifted columns are loaded with the systematics where the skims have them (load_samples optional_columns). For skims made before they were saved, fill_ak4_jec_shifts derives the shifted VBF jet mass from the shifted VBF pT (the jet factory scales pT and mass by the same factor) and keeps the AK4-away jets and the MET nominal. On the current 2024 VBF HH skim only VBFjjMass of the shifted inputs changes (0.15% RMS, in the 57% of events with two VBF jets); the nominal inputs do not. --- .../bdt_trainings_run3/v13_glopartv3.py | 4 +- src/HH4b/hh_vars.py | 4 + src/HH4b/postprocessing/postprocessing.py | 79 +++++++++++++++++-- src/HH4b/utils.py | 14 ++++ .../test_postprocessing_load_run3_samples.py | 2 + 5 files changed, 94 insertions(+), 9 deletions(-) diff --git a/src/HH4b/boosted/bdt_trainings_run3/v13_glopartv3.py b/src/HH4b/boosted/bdt_trainings_run3/v13_glopartv3.py index 191a8909..ab4624eb 100644 --- a/src/HH4b/boosted/bdt_trainings_run3/v13_glopartv3.py +++ b/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 diff --git a/src/HH4b/hh_vars.py b/src/HH4b/hh_vars.py index 81272154..ff32cfee 100644 --- a/src/HH4b/hh_vars.py +++ b/src/HH4b/hh_vars.py @@ -488,6 +488,10 @@ jec_vars = [ "bbFatJetPt", "VBFJetPt", + "VBFJetMass", + "AK4JetAwayPt", + "AK4JetAwayMass", + "MET_pt", "bdt_score", "bdt_score_vbf", "HHPt", diff --git a/src/HH4b/postprocessing/postprocessing.py b/src/HH4b/postprocessing/postprocessing.py index 4a960294..26986f70 100644 --- a/src/HH4b/postprocessing/postprocessing.py +++ b/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: diff --git a/src/HH4b/utils.py b/src/HH4b/utils.py index 75e7340f..d49d22a4 100644 --- a/src/HH4b/utils.py +++ b/src/HH4b/utils.py @@ -1004,6 +1004,7 @@ def load_samples( year: str, filters: list = None, columns: list = None, + optional_columns: list | None = None, variations: bool = True, weight_shifts: dict[str, Syst] = None, reorder_txbb: bool = False, # temporary fix for sorting by given Txbb @@ -1024,6 +1025,9 @@ def load_samples( year (str): year. filters (List): Optional filters when loading data. columns (List): Optional columns to load. + optional_columns (List): further columns (as ``format_columns`` names them) loaded only from + the sample directories whose parquet files have them; a frame that combines directories + with and without them has NaN for the rows of the latter. variations (bool): Normalize variations as well (saves time to not do so). Defaults to True. weight_shifts (Dict[str, Syst]): dictionary of weight shifts to consider. ptqq_stitch (PtqqStitch): opt-in stitching of the overlapping open-ended 2024 V+jets @@ -1160,6 +1164,16 @@ def load_samples( ) if range_mode and sample in ranges: _check_range_sample(parquet_path, pickles_path, sample, *ranges[sample]) + # optional columns: only if this sample's skim has them (e.g. newer JES/JER shifts) + if optional_columns and sample_columns is not None and parquet_path.exists(): + first_file = next((f for f in sorted(parquet_path.glob("*.parquet")) if _parquet_has_rows(f)), None) # fmt: skip + if first_file is not None: + names = set(pq.read_schema(first_file).names) + sample_columns = sample_columns + [ + col + for col in optional_columns + if col in names and col not in sample_columns + ] # no parquet directory? if not parquet_path.exists(): diff --git a/tests/test_postprocessing_load_run3_samples.py b/tests/test_postprocessing_load_run3_samples.py index 9026a2b7..4b91d738 100644 --- a/tests/test_postprocessing_load_run3_samples.py +++ b/tests/test_postprocessing_load_run3_samples.py @@ -59,6 +59,8 @@ def fake_load_samples(_input_dir, samples, _year, **kwargs): return {label: {"weight": [1.0]} for label in samples} monkeypatch.setattr(pp.utils, "load_samples", fake_load_samples) + # the fake frames have no jet columns for the JES/JER fill of old skims (fill_ak4_jec_shifts) + monkeypatch.setattr(pp, "fill_ak4_jec_shifts", lambda events, _key: events) samples_run3 = _base_samples_run3_for_year("2024") pp.load_run3_samples(