"""AK8 re-clustering closure test on collide2v (read-only on EOS; writes results to OUTDIR). For each test file and each mode: puppi: anti-kT R=0.8 on FullReco_PUPPIPart 4-vectors scaled by PuppiW (PuppiW>0) vs FullReco_JetPuppiAK8 plain: anti-kT R=0.8 on unweighted FullReco_PFPart vs FullReco_JetAK8 Reports (a) closure vs the provided jets (which only exist above 300 GeV): matching efficiency, fake rate, pT/eta/phi/mass agreement, constituent-set agreement; and (b) response/efficiency vs Gen_JetAK8 (stored from 200 GeV) down to 200 GeV, for both the re-clustered and the provided jets. """ import json import os import sys import time import awkward as ak import fastjet import fsspec import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import pyarrow.parquet as pq import vector vector.register_awkward() REDIR = "root://eosproject-f.cern.ch/" BASE = "/eos/project/f/foundational-model-dataset/samples/production_final" OUTDIR = os.environ.get("OUTDIR", "/traindatavol/collide2v/ak8closure") TEST_FILES = [ ("tt_hadronic_1", "tt0123j_5f_ckm_LO_MLM_hadronic/tt0123j_5f_ckm_LO_MLM_hadronic-NEVENT10000-RS28000001.parquet"), ("tt_hadronic_2", "tt0123j_5f_ckm_LO_MLM_hadronic/tt0123j_5f_ckm_LO_MLM_hadronic-NEVENT10000-RS28000002.parquet"), ("qcd_ht250_1", "QCD_HT250toInf/QCD_HT250toInf-NEVENT10000-RS61000001.parquet"), ] N_EVENTS = int(os.environ.get("N_EVENTS", "2000")) # per file BATCH = 250 RECO_MIN_PT = 150.0 # re-clustering threshold (to probe below the provided 300 GeV) PROVIDED_MIN_PT = 300.0 MATCH_DR = 0.2 GEN_MATCH_DR = 0.4 MODES = { "puppi": dict(cands="FullReco_PUPPIPart", jets="FullReco_JetPuppiAK8", weighted=True), "plain": dict(cands="FullReco_PFPart", jets="FullReco_JetAK8", weighted=False), } JETDEF = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.8) def log(*a): print(*a, flush=True) def dphi(a, b): return (a - b + np.pi) % (2 * np.pi) - np.pi def iter_events(path, mode): m = MODES[mode] cols = [f"{m['cands']}_{f}" for f in ("PT", "Eta", "Phi", "Mass", "PuppiW")] cols += [f"{m['jets']}_{f}" for f in ("PT", "Eta", "Phi", "Mass", "ConstituentsIdx")] cols += [f"Gen_JetAK8_{f}" for f in ("PT", "Eta", "Phi", "Mass")] n = 0 fs, _ = fsspec.core.url_to_fs(REDIR + BASE) with fs.open(f"{BASE}/{path}", "rb") as fh: pf = pq.ParquetFile(fh) assert (pf.metadata.metadata or {}).get(b"dataset_version") == b"collide2v_v1.0" for batch in pf.iter_batches(batch_size=BATCH, columns=cols): flat = ak.from_arrow(batch) g = lambda c, f: ak.values_astype(flat[f"{c}_{f}"], np.float64) if f != "ConstituentsIdx" else flat[f"{c}_{f}"] yield { "c": {f: g(m["cands"], f) for f in ("PT", "Eta", "Phi", "Mass", "PuppiW")}, "j": {f: g(m["jets"], f) for f in ("PT", "Eta", "Phi", "Mass", "ConstituentsIdx")}, "gen": {f: g("Gen_JetAK8", f) for f in ("PT", "Eta", "Phi", "Mass")}, } n += len(flat) if n >= N_EVENTS: return def recluster(c, weighted): pt, eta, phi, mass = c["PT"], c["Eta"], c["Phi"], ak.where(c["Mass"] > 0, c["Mass"], 0.0) w = c["PuppiW"] if weighted else ak.ones_like(pt) keep = (w > 0) & (pt > 0) idx = ak.local_index(pt)[keep] v = ak.zip({"pt": pt[keep], "eta": eta[keep], "phi": phi[keep], "mass": mass[keep]}, with_name="Momentum4D") ww = w[keep] p4 = ak.zip({"px": v.px * ww, "py": v.py * ww, "pz": v.pz * ww, "E": v.E * ww}, with_name="Momentum4D") cs = fastjet.ClusterSequence(p4, JETDEF) jets = cs.inclusive_jets(min_pt=RECO_MIN_PT) cons = cs.constituent_index(min_pt=RECO_MIN_PT) order = ak.argsort(jets.pt, ascending=False) jets, cons = jets[order], cons[order] # map back to original candidate indices (per event; idx is per-event, cons is per-event-per-jet) cons_orig = ak.Array([[np.asarray(ie)[np.asarray(cj, dtype=np.int64)].tolist() for cj in ce] for ie, ce in zip(ak.to_list(idx), ak.to_list(cons))]) return ak.zip({"pt": jets.pt, "eta": jets.eta, "phi": jets.phi, "mass": jets.mass}), cons_orig def match(a_eta, a_phi, b_eta, b_phi, rmax): """for each object in a, index of closest b within rmax (else -1); per event lists (python)""" out = [] for ae, ap, be, bp in zip(a_eta, a_phi, b_eta, b_phi): ae, ap, be, bp = map(lambda x: np.asarray(x, dtype=np.float64), (ae, ap, be, bp)) if len(ae) == 0: out.append(np.zeros(0, dtype=int)); continue if len(be) == 0: out.append(-np.ones(len(ae), dtype=int)); continue d = np.hypot(ae[:, None] - be[None, :], dphi(ap[:, None], bp[None, :])) j = d.argmin(1) out.append(np.where(d[np.arange(len(ae)), j] < rmax, j, -1)) return out def run(mode): R = dict(n_events=0, n_prov=0, n_prov_matched=0, n_reco300=0, n_reco300_matched=0, n_reco320=0, n_reco320_matched=0, pt_ratio=[], mass_ratio=[], deta=[], dphi=[], jaccard=[], prov_pt=[], gen_pt=[], gen_reco_matched=[], gen_reco_resp=[], gen_prov_matched=[], gen_prov_resp=[], reco_mass_prov_matched=[], n_reco_150_300=0, t_cluster=0.0) for tag, path in TEST_FILES: log(f"[{mode}] {tag}: {path}") for ev in iter_events(path, mode): t0 = time.time() reco, cons = recluster(ev["c"], MODES[mode]["weighted"]) R["t_cluster"] += time.time() - t0 prov = ev["j"] nev = len(prov["PT"]); R["n_events"] += nev # provided -> reco m_pr = match(prov["Eta"], prov["Phi"], reco.eta, reco.phi, MATCH_DR) # reco (>=300 / >=320) -> provided m_rp = match(reco.eta, reco.phi, prov["Eta"], prov["Phi"], MATCH_DR) for e in range(nev): ppt = np.asarray(prov["PT"][e]); rpt = np.asarray(reco.pt[e]) R["n_prov"] += len(ppt); R["n_prov_matched"] += int((m_pr[e] >= 0).sum()) R["n_reco_150_300"] += int(((rpt >= 150) & (rpt < 300)).sum()) for thr, key in ((300, "n_reco300"), (320, "n_reco320")): s = rpt >= thr R[key] += int(s.sum()); R[key + "_matched"] += int((m_rp[e][s] >= 0).sum()) for i, j in enumerate(m_pr[e]): if j < 0: continue R["prov_pt"].append(float(ppt[i])) R["pt_ratio"].append(float(rpt[j] / ppt[i])) pm = float(prov["Mass"][e][i]); rm = float(reco.mass[e][j]) R["mass_ratio"].append(rm / pm if pm > 1 else np.nan) R["deta"].append(float(reco.eta[e][j] - prov["Eta"][e][i])) R["dphi"].append(float(dphi(reco.phi[e][j], prov["Phi"][e][i]))) a = set(np.asarray(prov["ConstituentsIdx"][e][i]).tolist()); b = set(np.asarray(cons[e][j]).tolist()) R["jaccard"].append(len(a & b) / max(len(a | b), 1)) # gen response (gen AK8 stored from 200 GeV) m_gr = match(ev["gen"]["Eta"], ev["gen"]["Phi"], reco.eta, reco.phi, GEN_MATCH_DR) m_gp = match(ev["gen"]["Eta"], ev["gen"]["Phi"], prov["Eta"], prov["Phi"], GEN_MATCH_DR) for e in range(nev): gpt = np.asarray(ev["gen"]["PT"][e]); geta = np.asarray(ev["gen"]["Eta"][e]) for i in range(len(gpt)): if abs(geta[i]) > 2.4: continue R["gen_pt"].append(float(gpt[i])) jr, jp = m_gr[e][i], m_gp[e][i] R["gen_reco_matched"].append(bool(jr >= 0)); R["gen_reco_resp"].append(float(reco.pt[e][jr] / gpt[i]) if jr >= 0 else np.nan) R["gen_prov_matched"].append(bool(jp >= 0)); R["gen_prov_resp"].append(float(prov["PT"][e][jp] / gpt[i]) if jp >= 0 else np.nan) log(f"[{mode}] after {tag}: events={R['n_events']} provided jets={R['n_prov']} matched={R['n_prov_matched']}") return R def q(x): x = np.asarray(x, dtype=float); x = x[np.isfinite(x)] if x.size == 0: return None return {k: round(float(v), 4) for k, v in zip(("q02", "q16", "q50", "q84", "q98"), np.quantile(x, [0.02, 0.16, 0.5, 0.84, 0.98]))} | {"n": int(x.size)} def summarize(mode, R): S = {"mode": mode, "events": R["n_events"], "cluster_seconds": round(R["t_cluster"], 1)} S["provided_jets"] = R["n_prov"] S["eff_provided_to_reclustered"] = round(R["n_prov_matched"] / max(R["n_prov"], 1), 4) S["reclustered_ge300"] = R["n_reco300"]; S["frac_reclustered_ge300_with_provided_match"] = round(R["n_reco300_matched"] / max(R["n_reco300"], 1), 4) S["reclustered_ge320"] = R["n_reco320"]; S["frac_reclustered_ge320_with_provided_match"] = round(R["n_reco320_matched"] / max(R["n_reco320"], 1), 4) S["reclustered_150_300_per_event"] = round(R["n_reco_150_300"] / max(R["n_events"], 1), 3) pr = np.asarray(R["pt_ratio"]); mr = np.asarray(R["mass_ratio"]); jc = np.asarray(R["jaccard"]) S["pt_ratio_reco_over_provided"] = q(pr) S["frac_pt_within_0p5pct"] = round(float(np.mean(np.abs(pr - 1) < 0.005)), 4) if pr.size else None S["frac_pt_within_2pct"] = round(float(np.mean(np.abs(pr - 1) < 0.02)), 4) if pr.size else None S["mass_ratio_reco_over_provided"] = q(mr) S["frac_mass_within_2pct"] = round(float(np.nanmean(np.abs(mr - 1) < 0.02)), 4) if mr.size else None S["deta"] = q(R["deta"]); S["dphi"] = q(R["dphi"]) S["constituent_jaccard"] = q(jc); S["frac_identical_constituents"] = round(float(np.mean(jc == 1.0)), 4) if jc.size else None gpt = np.asarray(R["gen_pt"]) bins = [200, 250, 300, 400, 600, 2000] S["gen_response"] = [] for lo, hi in zip(bins[:-1], bins[1:]): s = (gpt >= lo) & (gpt < hi) row = {"gen_pt": f"{lo}-{hi}", "n_gen": int(s.sum())} for k in ("reco", "prov"): mtc = np.asarray(R[f"gen_{k}_matched"])[s]; resp = np.asarray(R[f"gen_{k}_resp"])[s] row[f"{k}_eff"] = round(float(mtc.mean()), 3) if s.any() else None row[f"{k}_resp_median"] = round(float(np.nanmedian(resp)), 3) if np.isfinite(resp).any() else None S["gen_response"].append(row) return S def plots(mode, R): fig, ax = plt.subplots(2, 3, figsize=(16, 9)) pr = np.asarray(R["pt_ratio"]); mr = np.asarray(R["mass_ratio"]) ax[0, 0].hist(np.clip(pr, 0.9, 1.1), bins=100); ax[0, 0].set(xlabel="pT(reclustered)/pT(provided)", yscale="log") ax[0, 1].hist(np.clip(mr[np.isfinite(mr)], 0.8, 1.2), bins=100); ax[0, 1].set(xlabel="m(reclustered)/m(provided)", yscale="log") ax[0, 2].hist(R["jaccard"], bins=50); ax[0, 2].set(xlabel="constituent Jaccard", yscale="log") ax[1, 0].scatter(R["prov_pt"], pr, s=2); ax[1, 0].set(xlabel="provided pT", ylabel="pT ratio", ylim=(0.9, 1.1)) gpt = np.asarray(R["gen_pt"]); bins = np.array([200, 225, 250, 275, 300, 350, 400, 500, 600, 800, 1200]) c = 0.5 * (bins[1:] + bins[:-1]) for k, lab in (("reco", "reclustered (>150)"), ("prov", "provided (>300)")): mt = np.asarray(R[f"gen_{k}_matched"]); rs = np.asarray(R[f"gen_{k}_resp"]) eff = [mt[(gpt >= lo) & (gpt < hi)].mean() if ((gpt >= lo) & (gpt < hi)).any() else np.nan for lo, hi in zip(bins[:-1], bins[1:])] rsp = [np.nanmedian(rs[(gpt >= lo) & (gpt < hi)]) if np.isfinite(rs[(gpt >= lo) & (gpt < hi)]).any() else np.nan for lo, hi in zip(bins[:-1], bins[1:])] ax[1, 1].plot(c, eff, "o-", label=lab); ax[1, 2].plot(c, rsp, "o-", label=lab) ax[1, 1].set(xlabel="gen AK8 pT", ylabel="match efficiency (dR<0.4)", ylim=(0, 1.05)); ax[1, 1].legend() ax[1, 2].set(xlabel="gen AK8 pT", ylabel="median reco/gen pT"); ax[1, 2].legend() fig.suptitle(f"AK8 re-clustering closure: {mode}") fig.tight_layout(); fig.savefig(os.path.join(OUTDIR, f"closure_{mode}.png"), dpi=110); plt.close(fig) def main(): os.makedirs(OUTDIR, exist_ok=True) out = {"files": TEST_FILES, "n_events_per_file": N_EVENTS, "reco_min_pt": RECO_MIN_PT, "match_dR": MATCH_DR, "fastjet": fastjet.__version__} for mode in MODES: R = run(mode) S = summarize(mode, R) log(json.dumps(S, indent=1)) out[mode] = S plots(mode, R) with open(os.path.join(OUTDIR, "closure_summary.json"), "w") as f: json.dump(out, f, indent=1) log("wrote", OUTDIR) if __name__ == "__main__": main()