File indexing completed on 2026-11-09 10:46:03
0001
0002 """Compare lAger HepMC3 ASCII samples (e p -> e' J/psi p', J/psi -> e+ e- (gamma)).
0003
0004 usage: compare.py OUTDIR LABEL=file.hepmc [LABEL=file.hepmc ...]
0005 The first sample is the reference.
0006 """
0007 import sys
0008 import numpy as np
0009 from scipy import stats
0010 import matplotlib
0011 matplotlib.use("Agg")
0012 import matplotlib.pyplot as plt
0013 from matplotlib.backends.backend_pdf import PdfPages
0014
0015
0016 def mdot(a, b):
0017 return a[..., 3] * b[..., 3] - np.sum(a[..., :3] * b[..., :3], axis=-1)
0018
0019
0020 def read_hepmc(path):
0021 """Return dict of per-event 4-vectors (px,py,pz,E) and the raw particle rows."""
0022 out = {k: [] for k in ("k", "kp", "P", "Pp", "jpsi", "lm", "lp", "ngam", "egam")}
0023 raw = []
0024
0025 def flush(ev):
0026 if ev is None:
0027 return
0028 parts = ev
0029 beam_e = next(p for p in parts.values() if p["st"] == 4 and p["pid"] == 11)
0030 beam_p = next(p for p in parts.values() if p["st"] == 4 and p["pid"] == 2212)
0031 jpsi = next(p for p in parts.values() if p["pid"] == 443)
0032 scat = next(p for p in parts.values() if p["pid"] == 11 and p["st"] == 1 and p["par"] == beam_e["id"])
0033 rec = next(p for p in parts.values() if p["pid"] == 2212 and p["st"] == 1)
0034 lm = next(p for p in parts.values() if p["pid"] == 11 and p["par"] == jpsi["id"])
0035 lp = next(p for p in parts.values() if p["pid"] == -11 and p["par"] == jpsi["id"])
0036 gams = [p for p in parts.values() if p["pid"] == 22 and p["par"] == jpsi["id"]]
0037 for key, p in (("k", beam_e), ("kp", scat), ("P", beam_p), ("Pp", rec), ("jpsi", jpsi), ("lm", lm), ("lp", lp)):
0038 out[key].append(p["v"])
0039 out["ngam"].append(len(gams))
0040 out["egam"].append(sum(g["v"][3] for g in gams))
0041 raw.append(np.concatenate([p["v"] for p in sorted(parts.values(), key=lambda p: p["id"])]))
0042
0043 ev = None
0044 vtx_parent = {}
0045 with open(path) as f:
0046 for line in f:
0047 c = line[:2]
0048 if c == "E ":
0049 flush(ev)
0050 ev, vtx_parent = {}, {}
0051 elif c == "V ":
0052 tok = line.split()
0053 ins = tok[3].strip("[]").split(",")
0054 vtx_parent[int(tok[1])] = int(ins[0])
0055 elif c == "P ":
0056 tok = line.split()
0057 pid_, par = int(tok[1]), int(tok[2])
0058 if par < 0:
0059 par = vtx_parent.get(par, par)
0060 ev[pid_] = dict(id=pid_, par=par, pid=int(tok[3]), st=int(tok[9]),
0061 v=np.array([float(x) for x in tok[4:8]]))
0062 flush(ev)
0063 d = {k: np.array(v) for k, v in out.items()}
0064 return d, raw
0065
0066
0067 def observables(d):
0068 k, kp, P, Pp, V = d["k"], d["kp"], d["P"], d["Pp"], d["jpsi"]
0069 q = k - kp
0070 o = {}
0071 o["Q2 [GeV2]"] = -mdot(q, q)
0072 o["y"] = mdot(P, q) / mdot(P, k)
0073 o["W [GeV]"] = np.sqrt(mdot(P + q, P + q))
0074 dP = P - Pp
0075 o["-t [GeV2]"] = -mdot(dP, dP)
0076 o["x_B"] = o["Q2 [GeV2]"] / (2 * mdot(P, q))
0077
0078 def pt(v): return np.hypot(v[:, 0], v[:, 1])
0079 def eta(v):
0080 p = np.linalg.norm(v[:, :3], axis=1)
0081 return 0.5 * np.log((p + v[:, 2]) / (p - v[:, 2]))
0082 def mom(v): return np.linalg.norm(v[:, :3], axis=1)
0083
0084 o["J/psi pT [GeV]"] = pt(V)
0085 o["J/psi eta"] = eta(V)
0086 o["J/psi p [GeV]"] = mom(V)
0087 o["scat. e- eta"] = eta(kp)
0088 o["scat. e- p [GeV]"] = mom(kp)
0089 o["recoil p pT [GeV]"] = pt(Pp)
0090 o["recoil p eta"] = eta(Pp)
0091 o["decay e- eta"] = eta(d["lm"])
0092 o["decay e+ eta"] = eta(d["lp"])
0093 o["decay e- p [GeV]"] = mom(d["lm"])
0094 ll = d["lm"] + d["lp"]
0095 o["M(e+e-) [GeV]"] = np.sqrt(mdot(ll, ll))
0096 o["FSR photon E [GeV]"] = d["egam"]
0097 o["N FSR photons"] = d["ngam"].astype(float)
0098
0099 fin = kp + Pp + d["lm"] + d["lp"]
0100 o["|sum p_in - p_out| (no FSR) [GeV]"] = np.linalg.norm((k + P - fin)[:, :3], axis=1)
0101 return o
0102
0103
0104 LOGX = {"Q2 [GeV2]", "y", "-t [GeV2]", "x_B", "FSR photon E [GeV]"}
0105
0106
0107 def main():
0108 outdir = sys.argv[1]
0109 samples = [a.split("=", 1) for a in sys.argv[2:]]
0110 data = {}
0111 for label, path in samples:
0112 d, raw = read_hepmc(path)
0113 data[label] = (observables(d), raw)
0114 print(f"read {label}: {len(raw)} events from {path}")
0115
0116 labels = [s[0] for s in samples]
0117 ref = labels[0]
0118 ref_obs = data[ref][0]
0119
0120
0121 print("\n== event-by-event comparison vs", ref)
0122 for lab in labels[1:]:
0123 a, b = data[ref][1], data[lab][1]
0124 n = min(len(a), len(b))
0125 same_shape = sum(1 for i in range(n) if a[i].shape == b[i].shape)
0126 ident = sum(1 for i in range(n) if a[i].shape == b[i].shape and np.allclose(a[i], b[i], rtol=1e-9, atol=1e-12))
0127 first_diff = next((i for i in range(n) if not (a[i].shape == b[i].shape and np.allclose(a[i], b[i], rtol=1e-9, atol=1e-12))), None)
0128 print(f" {lab}: {ident}/{n} events identical (rel 1e-9); same topology {same_shape}/{n}; first differing event: {first_diff}")
0129
0130 print("\n== distribution comparison (KS two-sample p-value; means)")
0131 hdr = f" {'observable':38s}" + "".join(f"{l:>22s}" for l in labels) + "".join(f"{'KS p '+l:>18s}" for l in labels[1:])
0132 print(hdr)
0133 rows = []
0134 for name in ref_obs:
0135 line = f" {name:38s}"
0136 for lab in labels:
0137 x = data[lab][0][name]
0138 line += f"{np.mean(x):>12.5g} ±{np.std(x)/np.sqrt(len(x)):<9.2g}"
0139 ps = []
0140 for lab in labels[1:]:
0141 p = stats.ks_2samp(ref_obs[name], data[lab][0][name]).pvalue
0142 ps.append(p)
0143 line += f"{p:>18.3f}"
0144 rows.append((name, ps))
0145 print(line)
0146
0147 with PdfPages(f"{outdir}/comparison.pdf") as pdf:
0148 names = list(ref_obs)
0149 for i0 in range(0, len(names), 6):
0150 fig, axes = plt.subplots(2, 3, figsize=(15, 8))
0151 for ax, name in zip(axes.flat, names[i0:i0 + 6]):
0152 allx = np.concatenate([data[l][0][name] for l in labels])
0153 lo, hi = np.percentile(allx, [0.1, 99.9])
0154 if name in LOGX and lo > 0:
0155 bins = np.logspace(np.log10(lo), np.log10(hi), 50); ax.set_xscale("log")
0156 elif name == "N FSR photons":
0157 bins = np.arange(-0.5, allx.max() + 1.5)
0158 else:
0159 bins = np.linspace(lo, hi, 50)
0160 for lab in labels:
0161 x = data[lab][0][name]
0162 h, e = np.histogram(x, bins=bins)
0163 w = 1.0 / len(x)
0164 c = 0.5 * (e[1:] + e[:-1])
0165 if lab == ref:
0166 ax.hist(x, bins=bins, weights=np.full(len(x), w), histtype="stepfilled", alpha=0.3, label=lab)
0167 else:
0168 ax.errorbar(c, h * w, yerr=np.sqrt(h) * w, fmt=".", ms=3, label=lab)
0169 ax.set_xlabel(name); ax.set_ylabel("fraction / bin")
0170 if name in ("-t [GeV2]", "FSR photon E [GeV]", "M(e+e-) [GeV]", "|sum p_in - p_out| (no FSR) [GeV]"):
0171 ax.set_yscale("log")
0172 ax.legend(fontsize=7)
0173 fig.tight_layout(); pdf.savefig(fig); fig.savefig(f"{outdir}/comparison_p{i0 // 6 + 1}.png", dpi=90); plt.close(fig)
0174 print(f"\nplots: {outdir}/comparison.pdf")
0175
0176
0177 if __name__ == "__main__":
0178 main()
