JPsiDataSet/validation/compare_invariants.py

0001 #!/usr/bin/env python3
0002 """Event-by-event comparison of Lorentz-invariant quantities between lAger samples,
0003 so that afterburned (boosted/rotated) events can be compared with generator output.
0004
0005 usage: compare_invariants.py LABEL=file.hepmc LABEL=file.hepmc [...]
0006 The first sample is the one every other sample is compared to.
0007 """
0008 import importlib.util
0009 import os
0010 import sys
0011 import numpy as np
0012 from scipy import stats
0013
0014 spec = importlib.util.spec_from_file_location(
0015     "compare", os.path.join(os.path.dirname(os.path.abspath(__file__)), "compare.py"))
0016 compare = importlib.util.module_from_spec(spec)
0017 spec.loader.exec_module(compare)
0018 mdot, read_hepmc = compare.mdot, compare.read_hepmc
0019
0020
0021 def gamma_vectors(path):
0022     """Summed FSR photon 4-vector per event (zero if no FSR)."""
0023     out, cur, jpsi_id, gams = [], None, None, None
0024     with open(path) as f:
0025         for line in f:
0026             if line.startswith("E "):
0027                 if cur is not None:
0028                     out.append(cur)
0029                 cur, jpsi_id = np.zeros(4), None
0030             elif line.startswith("P "):
0031                 t = line.split()
0032                 if int(t[3]) == 443:
0033                     jpsi_id = int(t[1])
0034                 elif int(t[3]) == 22 and int(t[2]) == jpsi_id:
0035                     cur += np.array([float(x) for x in t[4:8]])
0036         out.append(cur)
0037     return np.array(out)
0038
0039
0040 def invariants(path):
0041     d, _ = read_hepmc(path)
0042     k, kp, P, Pp, V, lm, lp = (d[x] for x in ("k", "kp", "P", "Pp", "jpsi", "lm", "lp"))
0043     q = k - kp
0044     MV = np.sqrt(mdot(V, V))
0045     g = gamma_vectors(path)
0046     ll = lm + lp
0047     return {
0048         "Q2": -mdot(q, q),
0049         "y": mdot(P, q) / mdot(P, k),
0050         "W": np.sqrt(mdot(P + q, P + q)),
0051         "-t": -mdot(P - Pp, P - Pp),
0052         "M(J/psi)": MV,
0053         "N FSR": d["ngam"].astype(float),
0054         "E_gamma (J/psi frame)": mdot(V, g) / MV,
0055         "M(e+e-)": np.sqrt(mdot(ll, ll)),
0056         "decay angle (q.l-/q.V)": mdot(q, lm) / mdot(q, V),
0057     }
0058
0059
0060 PRODUCTION = ("Q2", "y", "W", "-t", "M(J/psi)")
0061 DECAY = ("N FSR", "E_gamma (J/psi frame)", "M(e+e-)", "decay angle (q.l-/q.V)")
0062
0063
0064 def same(a, b, rtol=1e-6, atol=1e-9):
0065     return np.isclose(a, b, rtol=rtol, atol=atol)
0066
0067
0068 def main():
0069     samples = [a.split("=", 1) for a in sys.argv[1:]]
0070     inv = {}
0071     for label, path in samples:
0072         inv[label] = invariants(path)
0073         print(f"read {label}: {len(inv[label]['Q2'])} events")
0074     ref = samples[0][0]
0075     for label, _ in samples[1:]:
0076         a, b = inv[ref], inv[label]
0077         n = min(len(a["Q2"]), len(b["Q2"]))
0078         print(f"\n== {ref} vs {label}: first {n} events, event by event (rel. tol. 1e-6)")
0079         prod = np.all([same(a[x][:n], b[x][:n]) for x in PRODUCTION], axis=0)
0080         dec = np.all([same(a[x][:n], b[x][:n]) for x in DECAY], axis=0)
0081         for x in PRODUCTION + DECAY:
0082             m = same(a[x][:n], b[x][:n])
0083             print(f"  {x:26s} identical in {100 * m.mean():6.2f}% of events")
0084         print(f"  {'all production invariants':26s} identical in {100 * prod.mean():6.2f}% of events")
0085         print(f"  {'all decay invariants':26s} identical in {100 * dec.mean():6.2f}% of events")
0086         print("  distributions (KS p-value, mean A / mean B):")
0087         for x in PRODUCTION + DECAY:
0088             p = stats.ks_2samp(a[x][:n], b[x][:n]).pvalue
0089             print(f"    {x:26s} p={p:6.3f}   {np.mean(a[x][:n]):.5g} / {np.mean(b[x][:n]):.5g}")
0090
0091
0092 if __name__ == "__main__":
0093     main()