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/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 98% rename from src/HH4b/corrections/build_jec.py rename to src/HH4b/corrections/archive/build_jec.py index 3e8b7a8c..8010759c 100644 --- a/src/HH4b/corrections/build_jec.py +++ b/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=[ 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/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/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/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/processors/bbbbSkimmer.py b/src/HH4b/processors/bbbbSkimmer.py index 9742728d..975d2eda 100644 --- a/src/HH4b/processors/bbbbSkimmer.py +++ b/src/HH4b/processors/bbbbSkimmer.py @@ -708,23 +708,25 @@ 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 @@ -732,6 +734,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) @@ -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), diff --git a/src/HH4b/processors/corrections.py b/src/HH4b/processors/corrections.py index 0922ea08..befe2470 100644 --- a/src/HH4b/processors/corrections.py +++ b/src/HH4b/processors/corrections.py @@ -21,11 +21,12 @@ 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 -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) @@ -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/archive/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,45 @@ 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 type1_puppi_met( + 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] + 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 @@ -354,7 +401,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 +426,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..3999fe91 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``). @@ -17,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", @@ -133,3 +132,107 @@ 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. 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). + + 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"] + 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 + 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)) + 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) + 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 3dea0a1d..6ed9f710 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: 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", 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(