diff --git a/scripts/rabbit/scetlib_ad/backend_check.py b/scripts/rabbit/scetlib_ad/backend_check.py new file mode 100755 index 000000000..3e5883172 --- /dev/null +++ b/scripts/rabbit/scetlib_ad/backend_check.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""Standalone sanity check of the SCETlib autodiff backend. + +Runs without rabbit and without a datacard: load a cache, verify the anchor +round trip, finite-difference a couple of Jacobian columns, check the Hessian, +and exercise the bin permutation. This is the first thing to run after building +a new cache. + + source /setup.sh + python scripts/rabbit/scetlib_ad/backend_check.py \ + --conf /examples/matched_ad/matched.conf \ + --cache /examples/matched_ad/cache_debug.npz +""" + +import argparse +import os +import sys +import time + +import numpy as np + +sys.path.insert( + 0, + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ), +) + +from wremnants.postprocessing.scetlib_ad import params as adp # noqa: E402 +from wremnants.postprocessing.scetlib_ad.xsec_backend import ( # noqa: E402 + ScetlibADXsec, +) + + +def gen_axes_from_bins(bins): + """Recover ((qT, edges), (|Y|, edges)) and the Q window from a cache bin list. + + Only valid for a cache that is a full product grid, which is what + prepare_cache builds; it lets the check run without a datacard. + """ + Q = np.unique(bins[:, 0:2], axis=0) + if Q.shape[0] != 1: + raise SystemExit(f"expected a single Q bin, got {Q}") + y = np.unique(bins[:, 2:4], axis=0) + t = np.unique(bins[:, 4:6], axis=0) + y = y[np.argsort(y[:, 0])] + t = t[np.argsort(t[:, 0])] + y_edges = np.concatenate([y[:, 0], y[-1:, 1]]) + t_edges = np.concatenate([t[:, 0], t[-1:, 1]]) + return [("qT", t_edges), ("absY", y_edges)], float(Q[0, 0]), float(Q[0, 1]) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument( + "--conf", required=True, help="SCETlib runcard the cache was built from" + ) + ap.add_argument("--cache", required=True, help="cache .npz") + ap.add_argument("--threads", type=int, default=0) + ap.add_argument( + "--fd-params", + default=None, + help="comma-separated SCETlib parameter names to " + "finite-difference (default: alphas + the first NP one)", + ) + args = ap.parse_args() + + t0 = time.time() + core = ScetlibADXsec(args.conf, args.cache, threads=args.threads) + print(f"loaded in {time.time() - t0:.1f} s: {core}") + print(f" parameters ({core.n_params}):") + for n, v in zip(core.param_names, core.anchor): + print(f" {n:<24} {v:12.6g} -> rabbit {adp.rabbit_name(n)}") + + # ---- values and Jacobian at the anchor. The first call pays the worker-pool + # build and the LHAPDF/grid touch, so time a displaced repeat too -- that + # second number is what a minimizer iteration actually costs. + t0 = time.time() + val, jac = core.values_and_jacobian(core.anchor) + dt = time.time() - t0 + p_off = core.anchor * (1.0 + 1e-3) + 1e-6 + t0 = time.time() + core.values_and_jacobian(p_off) + dt_warm = time.time() - t0 + print( + f"\nvalue+jacobian at the anchor: {dt * 1e3:.0f} ms " + f"({dt / core.n_bins * 1e3:.2f} ms/bin) first call, " + f"{dt_warm * 1e3:.0f} ms ({dt_warm / core.n_bins * 1e3:.2f} ms/bin) warm; " + f"sum(sigma) = {val.sum():.8g}" + ) + if not np.all(np.isfinite(val)): + raise SystemExit("non-finite values at the anchor") + if np.any(val <= 0): + n = int((val <= 0).sum()) + print( + f" NOTE: {n} bin(s) are non-positive at the anchor " + f"(lowest-qT bins can be, with NP off / no large-bT damping)" + ) + + # ---- ratio to itself must be exactly 1 + val2, _ = core.values_and_jacobian(core.anchor) + print(f" anchor re-evaluation is bit-identical: {np.array_equal(val, val2)}") + + # ---- finite-difference a couple of Jacobian columns + names = core.param_names + if args.fd_params: + which = [names.index(n) for n in args.fd_params.split(",")] + else: + which = [0] + for i, n in enumerate(names): + if n.startswith("np_gnu_lambda2") or n.startswith("np_eff_lambda2"): + which.append(i) + break + print("\nfinite-difference check of d(sum sigma)/dp:") + worst = 0.0 + for i in which: + h = 1e-4 * max(abs(core.anchor[i]), 1e-3) + pp, pm = core.anchor.copy(), core.anchor.copy() + pp[i] += h + pm[i] -= h + fd = ( + core.values_and_jacobian(pp)[0].sum() + - core.values_and_jacobian(pm)[0].sum() + ) / (2 * h) + an = jac[:, i].sum() + rel = abs(an - fd) / max(abs(fd), 1e-300) + worst = max(worst, rel) + print(f" {names[i]:<24} analytic {an:14.8g} FD {fd:14.8g} rel {rel:.2e}") + # The rule reproduces the direct calculation to ~1e-15, so the only error + # here is the FD truncation; anything above ~1e-5 means the gradient is wrong. + status = "OK" if worst < 1e-5 else "FAIL" + print(f" worst relative disagreement {worst:.2e} -> {status}") + + # ---- Hessian + t0 = time.time() + H = core.hessian(core.anchor) + dt = time.time() - t0 + asym = np.max(np.abs(H - np.transpose(H, (0, 2, 1)))) + scale = max(np.max(np.abs(H)), 1e-300) + print( + f"\nhessian: {dt * 1e3:.0f} ms ({dt / core.n_bins * 1e3:.2f} ms/bin), " + f"shape {H.shape}, max|H - H^T|/max|H| = {asym / scale:.2e}" + ) + print( + f" -> hessian costs {dt / max(dt_warm, 1e-12):.0f}x a value+jacobian " + f"call; that ratio is why curvature is off by default during the fit " + f"and only turned on for the covariance pass" + ) + + # ---- fold onto the gen grid, and the sum rule that proves it is exact + gen_axes, q_lo, q_hi = gen_axes_from_bins(core.bins) + fold = core.fold_for(gen_axes, q_lo, q_hi) + folded = fold(val) + print( + f"\nfold onto the gen grid " + f"({gen_axes[0][0]}: {len(gen_axes[0][1]) - 1}, " + f"{gen_axes[1][0]}: {len(gen_axes[1][1]) - 1}, Q [{q_lo:g}, {q_hi:g}]): " + f"{fold.describe()}" + ) + # Summing bin-integrated cross sections is exact, so the folded total must + # equal the total over the cache bins the fold used (fp round-off only). + used = val.sum() if fold.n_dropped == 0 else None + if used is not None: + rel = abs(folded.sum() / used - 1.0) + print( + f" sum rule: folded total vs cache total, rel {rel:.2e} " + f"-> {'OK' if rel < 1e-12 else 'FAIL'}" + ) + if rel >= 1e-12: + raise SystemExit("the fold does not conserve the total cross section") + + if status != "OK": + raise SystemExit("gradient check failed") + print("\nall checks passed") + + +if __name__ == "__main__": + main() diff --git a/scripts/rabbit/scetlib_ad/compare_cards.py b/scripts/rabbit/scetlib_ad/compare_cards.py new file mode 100644 index 000000000..a3ed2d0b9 --- /dev/null +++ b/scripts/rabbit/scetlib_ad/compare_cards.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +"""Compare a CORRECTED and an UNCORRECTED datacard, to justify moving the +MiNNLO->SCETlib theory correction out of the histmaker and into the param model. + +The correction is applied by the histmaker as a per-event weight looked up from +the CorrZ histogram -- a BIN LOOKUP, not an interpolation +(``correctionsTensor_helper.py``: "takes a histogram and returns what is in the +bin"). So the event-weight route is itself piecewise-constant and + + Sum_events w(g_event) == Sum_g w(g) * R_raw(b, g) + +is an IDENTITY at matched binning: the two routes are the same object written +differently. What differs in practice is grid coarseness -- the card's response +gen grid (ptVGen 21 x absYVGen 10) against the correction's own grid +(qT 70 x absY 17) -- so the residual is the WITHIN-GEN-BIN variation of the +correction, expected worst at low qT where it moves fastest. + +Two tests, neither needing a cache, SCETlib, or the corr file: + +T1 row-sum audit, per card + R_rowsum(b) / norm_signal(b), R_rowsum = Sum_g R_raw(b,g) == R @ N_gen + The theory cancels identically, so this tests ONLY the plumbing: reco + marginalization, cropping, axis order, channel slicing, and how much gen + truth leaks outside the response grid. Must give the same profile on both + cards, since the leakage is a property of R's coverage, not of the weights. + +T2 granularity -- THE DECISIVE TEST + cbar(g) = N_gen_corr(g) / N_gen_unc(g) the coarse-grained correction + M(b) = [Sum_g R_raw(b,g) cbar(g)] / [Sum_g R_raw(b,g)] matrix route + E(b) = norm_corr(b) / norm_unc(b) event route + M/E - 1 IS the event-vs-matrix residual, with the theory, the pb/fb + factor, the Y convention and sigma_SCETlib all cancelled out. + Decision rule: adopt the matrix route if the yield-weighted |M/E - 1| is + below ~0.1% and no sensitivity-carrying bin exceeds ~1%. Otherwise the + response's gen binning needs refining toward the correction's grid. + +NB the two cards must differ ONLY by the correction -- same reco axes, bin +count, processes and zero pattern -- or E(b) is not the event route and T2 is +void. That is asserted, not assumed. +""" + +import argparse +import os +import sys + +import numpy as np + +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..") +) + +from wremnants.postprocessing.scetlib_ad.response import ( # noqa: E402 + R_info_from_auxiliary, + crop_R_to_fit, + marginalize_R_reco, +) + +AXIS_LABELS = {"ptll": r"$p_{T}^{\ell\ell}$ [GeV]", "yll": r"$y^{\ell\ell}$"} + + +def load_card(path, signal_proc): + """R (cropped to the fit channel), N_gen, the signal norm column, and axes.""" + from rabbit.inputdata import FitInputData + + ind = FitInputData(path) + info = R_info_from_auxiliary(ind) + + non_masked = [ + (n, i) for n, i in ind.channel_info.items() if not i.get("masked", False) + ] + if len(non_masked) != 1: + raise SystemExit( + f"{path}: expected one non-masked channel, got {len(non_masked)}" + ) + chan, cinfo = non_masked[0] + fit_axes = [(a.name, np.asarray(a.edges, float)) for a in cinfo["axes"]] + + R_full, R_axes = marginalize_R_reco( + info["R"], info["reco_axes"], [n for n, _ in fit_axes] + ) + R = crop_R_to_fit(R_full, R_axes, fit_axes) + n_reco = int(np.prod(R.shape[: len(fit_axes)])) + R = R.reshape(n_reco, -1) + + procs = [p.decode() if isinstance(p, bytes) else str(p) for p in ind.procs] + if signal_proc not in procs: + raise SystemExit(f"{path}: signal {signal_proc!r} not in {procs}") + norm = np.asarray( + ind.norm.numpy() if hasattr(ind.norm, "numpy") else ind.norm, dtype=float + ) + # Slice the channel's own rows rather than assuming it starts at 0. + lo = int(cinfo.get("start", 0) or 0) + hi = int(cinfo.get("stop", lo + n_reco) or lo + n_reco) + norm_sig = norm[lo:hi, procs.index(signal_proc)] + + return dict( + path=path, + R=R, + N_gen=np.asarray(info["N_gen"], float).reshape(-1), + norm_sig=norm_sig, + fit_axes=fit_axes, + gen_axes=info["gen_axes"], + procs=procs, + channel=chan, + lumi=cinfo.get("lumi"), + ) + + +def summarize(name, ratio, weights, axes, shape, top_n=6): + """Yield-weighted stats plus the worst bins, with their coordinates.""" + good = np.isfinite(ratio) & (weights > 0) + r, w = ratio[good], weights[good] + dev = np.abs(r - 1.0) + ywm = float(np.average(dev, weights=w)) + print( + f" {name}: bins {r.size} mean {r.mean():.6f} " + f"min {r.min():.6f} max {r.max():.6f}\n" + f" YIELD-WEIGHTED mean|ratio-1| = {ywm:.6f} " + f"({100*ywm:.4f}%) worst |dev| = {dev.max():.3e}" + ) + idx = np.argsort(-dev)[:top_n] + flat = np.where(good)[0] + print(f" worst {top_n} bins:") + for j in idx: + b = flat[j] + coord = np.unravel_index(b, shape) + loc = ", ".join(f"{nm}[{e[c]:g},{e[c+1]:g}]" for (nm, e), c in zip(axes, coord)) + print(f" {loc:<44} ratio {r[j]:.6f} yield {w[j]:.4g}") + return ywm + + +def profile(name, num, den, axis_idx, axes, shape): + """1-D profile: sum numerator and denominator over the OTHER axes, then divide. + + Not a mean of per-bin ratios -- that would weight the low-yield forward bins + like the peak. + """ + nm, edges = axes[axis_idx] + other = tuple(i for i in range(len(shape)) if i != axis_idx) + n = num.reshape(shape).sum(axis=other) + d = den.reshape(shape).sum(axis=other) + r = np.where(d > 0, n / np.where(d == 0, np.nan, d), np.nan) + print(f"\n {name} profiled in {nm}:") + print(f" {'bin':>16} {'ratio':>12} {'dev':>12}") + for k in range(r.size): + print( + f" [{edges[k]:6g},{edges[k+1]:6g}] {r[k]:12.6f} " f"{r[k]-1.0:+12.3e}" + ) + return r + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--card-corrected", required=True) + ap.add_argument("--card-uncorrected", required=True) + ap.add_argument("--signal-proc", default="Zmumu") + ap.add_argument("--profile-axis", default="ptll") + args = ap.parse_args() + + print("loading cards ...", flush=True) + C = load_card(args.card_corrected, args.signal_proc) + U = load_card(args.card_uncorrected, args.signal_proc) + + # ---- preconditions: the cards must differ ONLY by the correction -------- + print("\n=== preconditions ===") + errs = [] + if [n for n, _ in C["fit_axes"]] != [n for n, _ in U["fit_axes"]]: + errs.append("fit axis NAMES differ") + for (nc, ec), (nu, eu) in zip(C["fit_axes"], U["fit_axes"]): + if ec.shape != eu.shape or not np.allclose(ec, eu): + errs.append(f"fit axis {nc} EDGES differ") + if C["procs"] != U["procs"]: + errs.append(f"procs differ: {C['procs']} vs {U['procs']}") + if C["R"].shape != U["R"].shape: + errs.append(f"R shape {C['R'].shape} vs {U['R'].shape}") + if C["norm_sig"].shape != U["norm_sig"].shape: + errs.append("norm shape differs") + zc, zu = C["norm_sig"] == 0, U["norm_sig"] == 0 + if not np.array_equal(zc, zu): + errs.append(f"zero pattern differs in {int((zc ^ zu).sum())} bins") + for e in errs: + print(f" FAIL {e}") + if errs: + raise SystemExit("cards are not comparable; T2 would be meaningless") + print( + f" OK axes {[f'{n}({len(e)-1})' for n, e in C['fit_axes']]}, " + f"R {C['R'].shape}, procs {C['procs']}" + ) + print(f" lumi: corrected {C['lumi']}, uncorrected {U['lumi']}") + shape = tuple(len(e) - 1 for _, e in C["fit_axes"]) + + # ---- T1: row-sum audit, per card --------------------------------------- + print("\n=== T1 row-sum audit (theory cancels; plumbing only) ===") + for tag, D in (("corrected ", C), ("uncorrected", U)): + rowsum = D["R"].sum(axis=1) + with np.errstate(divide="ignore", invalid="ignore"): + ratio = rowsum / np.where(D["norm_sig"] == 0, np.nan, D["norm_sig"]) + print( + f"\n [{tag}] sum R_rowsum={rowsum.sum():.8g} " + f"sum norm_sig={D['norm_sig'].sum():.8g} " + f"global={rowsum.sum()/D['norm_sig'].sum():.8f}" + ) + summarize(tag, ratio, D["norm_sig"], D["fit_axes"], shape) + + # ---- T2: granularity ---------------------------------------------------- + print("\n=== T2 granularity: matrix route vs event route (DECISIVE) ===") + ng_c, ng_u = C["N_gen"], U["N_gen"] + pos = ng_u > 0 + if not pos.all(): + print(f" note: {int((~pos).sum())} gen bins have N_gen_unc <= 0; excluded") + cbar = np.ones_like(ng_u) + cbar[pos] = ng_c[pos] / ng_u[pos] + print( + f" cbar = N_gen_corr/N_gen_unc over {int(pos.sum())} gen bins: " + f"min {cbar[pos].min():.6f} max {cbar[pos].max():.6f} " + f"mean {cbar[pos].mean():.6f}" + ) + print(" (that IS the correction, coarse-grained onto the response gen grid)") + + Ru = U["R"] + num = Ru @ cbar + den = Ru.sum(axis=1) + with np.errstate(divide="ignore", invalid="ignore"): + M = num / np.where(den == 0, np.nan, den) + E = C["norm_sig"] / np.where(U["norm_sig"] == 0, np.nan, U["norm_sig"]) + MoverE = M / np.where(E == 0, np.nan, E) + print(f"\n M (matrix) min/max: {np.nanmin(M):.6f} / {np.nanmax(M):.6f}") + print(f" E (event) min/max: {np.nanmin(E):.6f} / {np.nanmax(E):.6f}") + ywm = summarize("M/E", MoverE, U["norm_sig"], C["fit_axes"], shape) + + # R invariance: R = R_raw/N_gen must be identical between the two cards, + # since a gen-level reweighting multiplies numerator and denominator alike. + # The whole construction rests on this, so check it rather than assume it. + print("\n=== R invariance (R = R_raw/N_gen must be card-independent) ===") + both = (C["N_gen"] > 0) & (U["N_gen"] > 0) + Pc = C["R"][:, both] / C["N_gen"][both][None, :] + Pu = U["R"][:, both] / U["N_gen"][both][None, :] + with np.errstate(divide="ignore", invalid="ignore"): + rel = np.abs(Pc / np.where(Pu == 0, np.nan, Pu) - 1.0) + ok = np.isfinite(rel) + print( + f" |P_corr/P_unc - 1| over {int(ok.sum())} entries: " + f"max {rel[ok].max():.3e} median {np.median(rel[ok]):.3e}" + ) + if Pu[ok].sum() > 0: + print( + f" response-weighted mean = " f"{np.average(rel[ok], weights=Pu[ok]):.3e}" + ) + print( + " (big values in near-empty response entries are MC noise, not a " + "violation; the weighted mean is the meaningful number)" + ) + + names = [n for n, _ in C["fit_axes"]] + # Which axis drives the residual? The axis whose M/E profile carries the + # most structure is the one whose gen binning is too coarse. + for ai, nm in enumerate(names): + pm = profile(f"M (matrix) [{nm}]", num, den, ai, C["fit_axes"], shape) + pe = profile( + f"E (event) [{nm}]", + C["norm_sig"], + U["norm_sig"], + ai, + C["fit_axes"], + shape, + ) + d = pm / pe - 1.0 + print( + f"\n ==> M/E - 1 profiled in {nm}: " + f"max|.| = {np.nanmax(np.abs(d)):.3e}, " + f"rms = {np.sqrt(np.nanmean(d**2)):.3e}" + ) + + print("\n=== verdict ===") + print(f" yield-weighted |M/E - 1| = {100*ywm:.4f}%") + if ywm < 1e-3: + print(" -> below the 0.1% decision threshold: the matrix route reproduces") + print(" the event route on this gen grid. Refining it is unnecessary.") + else: + print(" -> ABOVE 0.1%: the response gen grid is too coarse for the") + print(" correction it has to carry. Consider refining it toward the") + print(" correction's own grid (qT 70 x absY 17), weighing the extra") + print(" statistical noise per gen bin in R.") + + +if __name__ == "__main__": + main() diff --git a/scripts/rabbit/scetlib_ad/compare_to_np_model.py b/scripts/rabbit/scetlib_ad/compare_to_np_model.py new file mode 100755 index 000000000..0a0689f76 --- /dev/null +++ b/scripts/rabbit/scetlib_ad/compare_to_np_model.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python3 +"""Cross-check the SCETlib autodiff model's lambda response against scetlib_np. + +The two models compute sigma_gen by completely different routes -- the AD one +replays SCETlib's own compressed bin rules, the NP one reconstructs the Hankel +integral from a cached bT grid in TensorFlow -- so agreement of the RESPONSE + + R(lambda) = sigma_gen(lambda) / sigma_gen(lambda_central) + +is a sharp test of the new model. The NP model's response is itself validated +against the histmaker templates at the 0.02-0.05% level, so a disagreement here +well above that points at the new model (a wrong name map, a transposed fold, a +sign). + +The two do NOT share a nonsingular: the AD path uses SCETlib's own analytic +NLO V+jet at O(alphas^2), the NP path uses DYTurbo minus the SCETlib singular +expansion. That difference largely cancels in the ratio, but not exactly, and it +grows where the fixed order dominates -- hence the low-qT / high-qT split in the +report. Read the low-qT column as the real test. + + source /setup.sh + python scripts/rabbit/scetlib_ad/compare_to_np_model.py \ + --conf <...>/matched.conf --cache <...>/cache_debug.npz +""" + +import argparse +import os +import sys + +import numpy as np + +sys.path.insert( + 0, + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ), +) + +from wremnants.postprocessing.scetlib_ad import params as adp # noqa: E402 +from wremnants.postprocessing.scetlib_ad.xsec_backend import ( # noqa: E402 + ScetlibADXsec, +) + +# rabbit-facing name -> which scetlib_np dict it belongs in +EFF = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") +GNU = ("lambda2_nu", "lambda4_nu", "lambda6_nu", "lambda_inf_nu") + + +def gen_axes_from_bins(bins): + Q = np.unique(bins[:, 0:2], axis=0) + if Q.shape[0] != 1: + raise SystemExit(f"expected a single Q bin, got {Q}") + y = np.unique(np.abs(bins[:, 2:4]), axis=0) + t = np.unique(bins[:, 4:6], axis=0) + y = y[np.argsort(y[:, 0])] + t = t[np.argsort(t[:, 0])] + return ( + [ + ("ptVGen", np.concatenate([t[:, 0], t[-1:, 1]])), + ("absYVGen", np.concatenate([y[:, 0], y[-1:, 1]])), + ], + float(Q[0, 0]), + float(Q[0, 1]), + ) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--conf", required=True) + ap.add_argument("--cache", required=True) + ap.add_argument( + "--btgrid", + default=None, + help="scetlib_np bT-grid dir (default: the CT18Z one the " "package resolves)", + ) + ap.add_argument("--threads", type=int, default=0) + ap.add_argument( + "--rel-shifts", + default="0.05,0.2,1.0", + help="comma-separated RELATIVE displacements applied to each " + "lambda in turn. Relative, because the compressed bin " + "rules are trained in a neighbourhood of the anchor " + "(scale=0.15 by default) and a fixed absolute shift " + "would move a small lambda far outside it while barely " + "moving a large one", + ) + ap.add_argument( + "--abs-floor", + type=float, + default=0.02, + help="displacement used for a lambda whose anchor is ~0 " + "(delta_lambda2), where a relative shift is meaningless", + ) + ap.add_argument( + "--qt-split", + type=float, + default=10.0, + help="qT (GeV) separating the resummation-dominated region, " + "where the two nonsingulars agree best, from the rest", + ) + args = ap.parse_args() + + import tensorflow as tf # noqa: F401 (imported by the NP core anyway) + + from wremnants.postprocessing.scetlib_np.sigma_gen import SigmaGenModel + + core = ScetlibADXsec(args.conf, args.cache, threads=args.threads) + gen_axes, q_lo, q_hi = gen_axes_from_bins(core.bins) + fold = core.fold_for(gen_axes, q_lo, q_hi) + shape = fold.gen_shape + print( + f"gen grid {shape[0]} x {shape[1]}, Q [{q_lo:g}, {q_hi:g}]; " + f"{fold.describe()}" + ) + + # The NP model's central must be the AD cache's anchor, or the two are not + # anchored at the same point and every ratio below is meaningless. + npsec = core.conf["Nonperturbative"] + eff_c = {"np_model": npsec.get("np_model", "tanh_2")} + gnu_c = {"np_model_nu": npsec.get("np_model_nu", "tanh_2")} + for sname, value in zip(core.param_names, core.anchor): + try: + rname = adp.rabbit_name(sname) + except KeyError: + continue + if rname in EFF: + eff_c[rname] = float(value) + elif rname in GNU: + gnu_c[rname] = float(value) + print(f"anchored at eff {eff_c}\n gnu {gnu_c}") + + np_core = SigmaGenModel( + btgrid_dir=args.btgrid, + lambda_central={"eff_params": eff_c, "gnu_params": gnu_c}, + gen_axes=gen_axes, + Q_lo=q_lo, + Q_hi=q_hi, + ) + + def sigma_ad(**over): + p = core.anchor.copy() + for rname, val in over.items(): + p[core.param_names.index(adp.scetlib_name(rname))] = val + vals, _ = core.values_and_jacobian(p) + return fold(np.asarray(vals, dtype=np.float64)).reshape(shape) + + def sigma_np(**over): + eff, gnu = dict(eff_c), dict(gnu_c) + for rname, val in over.items(): + (eff if rname in EFF else gnu)[rname] = val + return np.asarray(np_core.sigma_gen(eff, gnu)) + + ad0, np0 = sigma_ad(), sigma_np() + qt_lo = gen_axes[0][1][:-1] + low = qt_lo < args.qt_split + + # ABSOLUTE normalisation. The fit never sees this, because compute() divides + # by the model's own central -- which is exactly why it is worth printing: + # it is what would have to be right if the anchor ratio were ever dropped in + # favour of an absolute prediction. + print( + f"\nABSOLUTE totals (hidden by the anchor ratio): " + f"AD {ad0.sum():.6g}, NP {np0.sum():.6g}, AD/NP = {ad0.sum() / np0.sum():.6f}" + ) + print( + " NB the AD cache integrates the POSITIVE Y side only, so a factor 2 per " + "|Y| bin\n is expected here and cancels in the ratio; anything beyond " + "that is a real\n normalisation convention to nail before going absolute." + ) + + # Informational: the two predictions do not have to agree in absolute shape + # (different nonsingular), and compute() never uses the absolute value. + shape_ratio = (ad0 / ad0.sum()) / (np0 / np0.sum()) + print( + "\ncentral shape, AD vs NP (normalised; NOT required to agree -- the " + "models use different nonsingulars, and the fit only ever uses the " + "ratio to each model's own central):" + ) + print( + f" qT < {args.qt_split:g}: max |ratio - 1| = " + f"{np.max(np.abs(shape_ratio[low] - 1)):.2%} " + f"qT > {args.qt_split:g}: {np.max(np.abs(shape_ratio[~low] - 1)):.2%}" + ) + + common = [ + n for n in EFF + GNU if n in [adp.rabbit_name(s) for s in core.param_names] + ] + rel_shifts = [float(x) for x in args.rel_shifts.split(",") if x.strip()] + print( + f"\nlambda response agreement, |R_AD - R_NP| as a fraction of the " + f"response |R_NP - 1| (max over qT < {args.qt_split:g})." + ) + print( + " Scanning the displacement matters: the rules are compressed around " + "the anchor,\n so the agreement should IMPROVE towards small shifts. " + "A flat or growing\n trend towards zero displacement is a bug, not " + "rule locality." + ) + header = ( + " " + + f"{'parameter':<18}" + + "".join(f"{'x' + format(1 + r, 'g'):>12}" for r in rel_shifts) + + f"{'response(x' + format(1 + rel_shifts[-1], 'g') + ')':>16}" + ) + print("\n" + header) + worst_small = 0.0 + for name in common: + base = float(core.anchor[core.param_names.index(adp.scetlib_name(name))]) + cells, size = [], 0.0 + for r in rel_shifts: + step = base * r if abs(base) > 1e-9 else args.abs_floor * r / rel_shifts[-1] + if abs(step) < 1e-12: + cells.append(" n/a") + continue + r_ad = sigma_ad(**{name: base + step}) / ad0 + r_np = sigma_np(**{name: base + step}) / np0 + size = float(np.max(np.abs(r_np[low] - 1.0))) + if size < 1e-4: + cells.append(" inert") + continue + frac = float(np.max(np.abs(r_ad[low] - r_np[low])) / size) + cells.append(f"{frac:11.1%}") + if r == min(rel_shifts): + worst_small = max(worst_small, frac) + print(f" {name:<18}" + "".join(f"{c:>12}" for c in cells) + f"{size:15.2%}") + print(f"\n worst disagreement at the smallest displacement: {worst_small:.1%}") + print( + " A wrong name map, a transposed fold or a sign error would show up " + "as O(1) at every\n displacement. What remains at small displacement " + "is the genuine difference between\n the two predictions: different " + "nonsingulars, and the NP model rebinning a POINT\n grid by " + "quadrature where SCETlib integrates each bin exactly." + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/rabbit/scetlib_ad/compare_to_scetlib_run.py b/scripts/rabbit/scetlib_ad/compare_to_scetlib_run.py new file mode 100644 index 000000000..f1b5efc68 --- /dev/null +++ b/scripts/rabbit/scetlib_ad/compare_to_scetlib_run.py @@ -0,0 +1,395 @@ +#!/usr/bin/env python3 +"""Validate an autodiff cache against a native SCETlib prediction. + +Two references are meaningful, and they test different things: + +* a production run with ``calculation_piece = sing`` -- the RESUMMED cross + section. Replaying only our cache's compressed bin rules gives the same object, + with no matching, no fixed-order generator and no MC in between, so any + disagreement is ours: runcard, quadrature, Q integration, rule compression. +* a ``*Corr.pkl.lz4`` theory correction -- the MATCHED prediction the + analysis actually uses, resummed plus (fixed-order generator minus the singular + expansion). Our cache computes its own matched total with SCETlib's in-house + analytic V+jet instead, so this comparison measures the deliberate change of + nonsingular as well as everything the resummed test covers. + +Both sides are bin-integrated, so the two are summed onto their COMMON bin edges +-- exact, no interpolation -- and the script refuses to run unless each side +tiles that common grid exactly. A signed-Y cache is folded onto |Y| when the +reference is binned in |Y|. + + source /setup.sh + python scripts/rabbit/scetlib_ad/compare_to_scetlib_run.py \ + --conf .conf --cache .npz \ + --reference <...>_combined.pkl --piece resummed + python scripts/rabbit/scetlib_ad/compare_to_scetlib_run.py \ + --conf .conf --cache .npz \ + --reference <...>_CorrZ.pkl.lz4 --piece matched +""" + +import argparse +import os +import pickle +import sys + +import numpy as np + +sys.path.insert( + 0, + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ), +) + +from wremnants.postprocessing.scetlib_ad.xsec_backend import ( # noqa: E402 + ScetlibADXsec, +) + +TOL = 1e-9 + + +def load_reference(path, var): + """(values (Y, qT), Q edges, Y edges, qT edges, is_absY, config) from either + a SCETlib production pkl or a theory-correction pkl.lz4.""" + if path.endswith(".lz4"): + import lz4.frame + + with lz4.frame.open(path, "rb") as fh: + d = pickle.load(fh) + proc = next(iter(d)) # "Z" / "W" + block = d[proc] + key = next( + k for k in block if k.endswith("_hist") and not k.startswith("minnlo") + ) + h, cfg = block[key], {} + for v in (d.get("file_meta_data") or {}).values(): + if isinstance(v, dict) and "config" in v: + cfg = v["config"] + break + kind = f"theory correction ({key})" + else: + with open(path, "rb") as fh: + d = pickle.load(fh) + h, cfg = d["hist"], d.get("config", {}) + piece = cfg.get("Calculation_settings", {}).get("calculation_piece") + kind = f"production run (calculation_piece={piece})" + + names = [a.name for a in h.axes] + yname = "absY" if "absY" in names else "Y" + sel = {} + if "vars" in names: + sel["vars"] = var + if "charge" in names: + sel["charge"] = sum # a Z corr has one charge bin; sum is a no-op there + hs = h[sel] if sel else h + order = [a.name for a in hs.axes] + vals = np.asarray(hs.values()) + # -> (Q, Y, qT) + vals = np.transpose(vals, [order.index(n) for n in ("Q", yname, "qT")]) + return ( + vals, + np.asarray(h.axes["Q"].edges), + np.asarray(h.axes[yname].edges), + np.asarray(h.axes["qT"].edges), + yname == "absY", + cfg, + kind, + ) + + +def common_edges(a, b): + """Edges present in both, over their overlapping range.""" + lo, hi = max(a[0], b[0]), min(a[-1], b[-1]) + out = [e for e in a if np.any(np.abs(b - e) < TOL) and lo - TOL <= e <= hi + TOL] + if len(out) < 2: + raise SystemExit( + f"the two binnings share fewer than two edges in their overlap " + f"[{lo:g}, {hi:g}]; they cannot be compared without interpolating." + ) + return np.asarray(out) + + +def rebin(vals, axis, fine, coarse, what): + """Sum `vals` along `axis` from `fine` edges onto `coarse` edges (exact).""" + idx = np.full(fine.size - 1, -1, dtype=np.int64) + for i in range(fine.size - 1): + lo, hi = fine[i], fine[i + 1] + for j in range(coarse.size - 1): + if lo >= coarse[j] - TOL and hi <= coarse[j + 1] + TOL: + idx[i] = j + break + else: + if hi > coarse[0] + TOL and lo < coarse[-1] - TOL: + raise SystemExit( + f"{what} bin [{lo:g}, {hi:g}] straddles a common edge; the sum " + f"onto the common grid would not be exact." + ) + covered = np.zeros(coarse.size - 1) + for i in np.nonzero(idx >= 0)[0]: + covered[idx[i]] += fine[i + 1] - fine[i] + want = np.diff(coarse) + bad = np.abs(covered - want) > TOL * np.maximum(want, 1.0) + if bad.any(): + j = int(np.argmax(bad)) + raise SystemExit( + f"{what} common bin [{coarse[j]:g}, {coarse[j + 1]:g}] is not exactly " + f"tiled ({covered[j]:g} of {want[j]:g})." + ) + out = np.zeros( + vals.shape[:axis] + (coarse.size - 1,) + vals.shape[axis + 1 :], dtype=float + ) + v = np.moveaxis(vals, axis, 0) + o = np.moveaxis(out, axis, 0) + for i in np.nonzero(idx >= 0)[0]: + o[idx[i]] += v[i] + return out + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--conf", required=True) + ap.add_argument("--cache", required=True) + ap.add_argument("--reference", required=True) + ap.add_argument( + "--piece", + choices=["matched", "resummed"], + default="matched", + help="which part of OUR cache to compare: the matched total (against a " + "theory correction) or the resummed piece alone (against a " + "calculation_piece=sing production run)", + ) + ap.add_argument("--var", default="central", help="reference vars entry") + ap.add_argument("--threads", type=int, default=0) + ap.add_argument("--plot-dir", default=None) + ap.add_argument( + "--y-range", + default=None, + help="restrict the comparison to this |Y| (or Y) range, 'LO,HI'. Both " + "endpoints must be edges of the common grid.", + ) + ap.add_argument( + "--qt-range", + default=None, + help="restrict the comparison to this qT range, 'LO,HI'. Both endpoints " + "must be edges of the common grid.", + ) + ap.add_argument( + "--tag", + default=None, + help="suffix for the output filenames, so a restricted slice does not " + "overwrite the full-range plots", + ) + args = ap.parse_args() + + rvals, rQ, rY, rT, ref_absY, cfg, kind = load_reference(args.reference, args.var) + print(f"reference: {kind}") + + core = ScetlibADXsec(args.conf, args.cache, threads=args.threads) + ours = ( + core.resummed_only(core.anchor) + if args.piece == "resummed" + else core.values_and_jacobian(core.anchor)[0] + ) + print(f"ours: {args.piece} total from the cache") + + _report_config(cfg, core) + + # --- our grid, from the cache's own bin list + b = core.bins + ourQ = np.unique(b[:, 0:2], axis=0) + if ourQ.shape[0] != 1: + raise SystemExit("this script expects a single Q bin in the cache") + oY = np.unique(b[:, 2:4], axis=0) + oT = np.unique(b[:, 4:6], axis=0) + oY, oT = oY[np.argsort(oY[:, 0])], oT[np.argsort(oT[:, 0])] + Ye = np.concatenate([oY[:, 0], oY[-1:, 1]]) + Te = np.concatenate([oT[:, 0], oT[-1:, 1]]) + ours = np.asarray(ours).reshape(Te.size - 1, Ye.size - 1).T # (Y, qT) + + # --- fold ours onto |Y| if the reference is + if ref_absY: + if not np.allclose(Ye, -Ye[::-1]): + raise SystemExit( + "the reference is binned in |Y| but the cache's Y grid is not " + "symmetric about 0, so it cannot be folded." + ) + n = (Ye.size - 1) // 2 + ours = ours[n:] + ours[:n][::-1] + Ye = Ye[n:] + print(f" folded our signed Y onto |Y|: {n} bins") + + # --- reference: pick the Q bin matching ours, never sum over Q + q = next( + ( + j + for j in range(rQ.size - 1) + if abs(rQ[j] - ourQ[0, 0]) < TOL and abs(rQ[j + 1] - ourQ[0, 1]) < TOL + ), + None, + ) + if q is None: + raise SystemExit( + f"no reference Q bin matches ours [{ourQ[0, 0]:g}, {ourQ[0, 1]:g}]; " + f"reference Q edges {list(rQ)}" + ) + rvals = rvals[q] + print(f" reference Q bin {q} = [{rQ[q]:g}, {rQ[q + 1]:g}]") + + # --- sum both onto the common edges + cY, cT = common_edges(Ye, rY), common_edges(Te, rT) + cY = _restrict(cY, args.y_range, "Y") + cT = _restrict(cT, args.qt_range, "qT") + print( + f" common grid: {cY.size - 1} x {cT.size - 1} bins " + f"(|Y| {cY[0]:g}..{cY[-1]:g}, qT {cT[0]:g}..{cT[-1]:g})" + ) + ours_c = rebin(rebin(ours, 0, Ye, cY, "our Y"), 1, Te, cT, "our qT") + ref_c = rebin(rebin(rvals, 0, rY, cY, "reference Y"), 1, rT, cT, "reference qT") + + rel = ours_c / np.where(ref_c != 0, ref_c, np.nan) - 1.0 + print( + f"\ntotals: ours {ours_c.sum():.8g} reference {ref_c.sum():.8g} " + f"ours/ref = {ours_c.sum() / ref_c.sum():.6f}" + ) + print( + f"per-bin ours/ref - 1: max |.| = {np.nanmax(np.abs(rel)):.3e} " + f"median |.| = {np.nanmedian(np.abs(rel)):.3e}" + ) + print("\nby qT (max over |Y|):") + for j in range(cT.size - 1): + print( + f" qT [{cT[j]:6.1f},{cT[j + 1]:6.1f}] " + f"max |rel| = {np.nanmax(np.abs(rel[:, j])):+.3e}" + ) + + if args.plot_dir: + _plot(ours_c, ref_c, cY, cT, args, kind) + + +def _restrict(edges, spec, what): + """Clip a common-edge array to 'LO,HI', requiring both to be real edges.""" + if not spec: + return edges + lo, hi = (float(x) for x in spec.split(",")) + for v in (lo, hi): + if not np.any(np.abs(edges - v) < TOL): + raise SystemExit( + f"{v:g} is not an edge of the common {what} grid " + f"{[float(e) for e in edges]}; pick one of those." + ) + out = edges[(edges >= lo - TOL) & (edges <= hi + TOL)] + if out.size < 2: + raise SystemExit(f"the requested {what} range leaves no bins") + return out + + +def _report_config(cfg, core): + """Cross-check the settings we transcribed against the reference's own.""" + if not cfg: + print(" (the reference records no config; nothing to cross-check)") + return + # configparser lowercases keys, so compare case-insensitively or every + # camelCase setting reads as missing. + rc = {k.lower(): v for k, v in cfg.get("Calculation_settings", {}).items()} + oc = {k.lower(): v for k, v in dict(core.conf["Calculation_settings"]).items()} + rn = {k.lower(): v for k, v in cfg.get("Nonperturbative", {}).items()} + on = {k.lower(): v for k, v in dict(core.conf["Nonperturbative"]).items()} + diff = [ + (k, rc[k], oc.get(k)) + for k in ( + "lambda", + "transition_points", + "mu0_min", + "mub_min", + "mus_min", + "muf_min", + "compensate_fo", + "form_np_prescription", + "muf_follows_mub", + "disable_asymmetry", + "run_order", + "fixed_order", + ) + if k in rc and not _same(rc[k], oc.get(k)) + ] + npdiff = [ + (k, v, on.get(k)) for k, v in rn.items() if k in on and not _same(v, on[k]) + ] + print( + f" settings cross-check: " + f"{'OK' if not diff else f'{len(diff)} DIFFER: {diff}'}" + ) + print( + f" NP anchor cross-check: {len(rn)} entries, " + f"{'all agree' if not npdiff else f'DIFFER: {npdiff}'}" + ) + + +def _plot(ours_c, ref_c, Ye, Te, args, kind): + """qT spectrum integrated over Y, with a ratio-to-reference panel.""" + import hist + + from wums import output_tools, plot_tools + + # save_pdf_and_png does not create the directory; write_index_and_log adds the + # index.php gallery, which globs *.png and links each to its .log and .pdf. + os.makedirs(args.plot_dir, exist_ok=True) + tag = args.piece + (f"_{args.tag}" if args.tag else "") + + def h1(v, edges, name): + h = hist.Hist( + hist.axis.Variable(edges, name=name, overflow=False, underflow=False) + ) + h.values()[...] = v + return h + + meta = { + "reference": args.reference, + "reference kind": kind, + "cache": args.cache, + "runcard": args.conf, + "piece": args.piece, + "vars entry": args.var, + "ours/reference (total)": f"{ours_c.sum() / ref_c.sum():.6f}", + } + + # binwnorm because the qT bins differ hugely in width, so raw contents would + # show the binning rather than the spectrum. NOT logx: the first bin starts at + # qT = 0, and a log x-axis collapses the whole plot onto that edge. + fig = plot_tools.makePlotWithRatioToRef( + [h1(ref_c.sum(axis=0), Te, "qT"), h1(ours_c.sum(axis=0), Te, "qT")], + labels=["SCETlib reference", "autodiff cache"], + colors=["#5790fc", "#e42536"], + linestyles=["solid", "dashed"], + xlabel=r"boson $q_\mathrm{T}$ (GeV)", + ylabel=r"$d\sigma/dq_\mathrm{T}$ (a.u.)", + rlabel=["cache / reference"], + rrange=[[0.95, 1.05]], + binwnorm=1, + logy=True, + yerr=False, + nlegcols=1, + cms_label="Work in progress", + grid=True, + ) + plot_tools.save_pdf_and_png(args.plot_dir, f"{tag}_qT", fig=fig) + output_tools.write_index_and_log( + args.plot_dir, f"{tag}_qT", analysis_meta_info=meta, args=args + ) + + print(f"\n plots -> {args.plot_dir}") + + +def _same(a, b): + if a is None or b is None: + return a == b + try: + return abs(float(str(a).strip()) - float(str(b).strip())) < 1e-9 + except ValueError: + return str(a).strip().lower() == str(b).strip().lower() + + +if __name__ == "__main__": + main() diff --git a/scripts/rabbit/scetlib_ad/conf/Z_CT18Z_N3p0LL_FranksVals.conf b/scripts/rabbit/scetlib_ad/conf/Z_CT18Z_N3p0LL_FranksVals.conf new file mode 100644 index 000000000..ddbbeefa7 --- /dev/null +++ b/scripts/rabbit/scetlib_ad/conf/Z_CT18Z_N3p0LL_FranksVals.conf @@ -0,0 +1,134 @@ +#------------------------------------------------------------------------------- +# SCETlib autodiff runcard reproducing the CURRENT analysis central prediction: +# Z at 13 TeV, CT18Z, N3+0LL + NNLO, LatticeNP "FranksVals" tune. +# +# Every setting below is copied from the config embedded in the production run +# /work/submit/lavezzo/alphaS/TheoryCorrections/SCETlib/ +# com13_ct18z_newnps_n3+0ll_lattice_lambda4bugfix_franksvalsvars_fine/ +# inclusive_Z_..._franksvalsvars_fine_combined.pkl +# which is the resummed (calculation_piece = sing) prediction the analysis's +# theory correction is built from. Keeping this runcard compatible with that +# config is what makes the autodiff cache validatable against it bin by bin +# (scripts/rabbit/scetlib_ad/compare_to_scetlib_run.py). +# +# TWO deliberate departures, both required: +# * calculation_piece = matched -- the reference runs 'sing' only and takes the +# nonsingular from DYTurbo; the autodiff cache carries SCETlib's own matched +# total. The resummed sub-piece stays directly comparable to the reference. +# * fo_order2_analytic = yes -- needed to run the nonsingular at nnlo; it +# touches only the fixed-order pieces, not the resummed one. +# +# The [TNPs] block is what makes this N^{3+0}LL rather than N3LL, and it is also +# what registers the 10 TNP gradient parameters, so the vector is 19 long +# (alphas + 4 eff NP + 4 gamma_nu NP + 10 TNPs). +# +# Grid_* sections are supplied by prepare_cache_for_card.py, so none is given. +#------------------------------------------------------------------------------- + +[Process] +boson = Z + +[Calculation_settings] +calculation_piece = matched +run_order = n3ll +fixed_order = nnlo +# Do NOT set pick_structure_function: it silently swaps the exact O(alphas^2) U+L +# for the NLO/LO graft (see examples/matched_ad/README.md). +fo_order2_analytic = yes + +lambda = 0. +b0_over_bmax = 0. +b0_over_bmax_global = 0. +mu0_min = 1. +muB_min = 1. +muS_min = 1. +nuS_min = 0. +muf_min = 1.40 +transition_points = [0.2, 0.6, 1.0] +transition_type = slope +scale_setting = spectrum +weight_qT = none +muFO_fixed = 0. +kappaFO = 1. +kappaf = 1. +phase_muH = 1. +muf_follows_muB = no +compensate_fo = yes +form_np_prescription = collins_soper4 +scheme_RadISH = no +recoil_scheme = collins_soper +disable_asymmetry = yes +alphas_solution = analytic +alphas_rel_precision = 1.e-6 +rge_solution = analytic +profile_functional_form = slope +muf_max = 13000. + +[Singlet_scheme] +nonsinglet_use_exact_top_mass = no +vector_singlet_enabled = yes +vector_singlet_use_exact_top_mass = no +axial_singlet_enabled = yes +axial_singlet_use_exact_top_mass = no + +[Electroweak] +alphaem = 0.0077624494296896114 +sin2_thw = 0.23153999447822571 +mZ = 91.153509740726733 # GeV +GammaZ = 2.4932018986110700 # GeV +mW = 79.906853549493746 # GeV +GammaW = 2.0904310808144846 # GeV +ckm = [ [0.97446, 0.222, 0.00365], [0.22438, 0.97359, 0.04214], [0.00896, 0.04133, 0.999105] ] + +[QCD] +nf = 5 +alphas_mu0 = 0.118 +mu0 = 91.1876 +alphas_order = n3ll +pdf_set = CT18ZNNLO +pdf_member = 0 +Ecm = 13000. + +[Integration] +report_error_estimate = yes +target_precision_rel = 1.e-4 +target_precision_abs = 1.e-8 +precision_buffer_bT = 1.e-1 +precision_buffer_decay = 1.e-1 +tolerance_qT = 1. +precision_buffer_full = 1. +abs_precision_buffer_inner = 1. +change_var_Q = arctan_Q2 +change_var_Q_q0 = 91.15348061918276 +change_var_qT = none +change_var_qT_q0 = 1. + +[Nonperturbative] +# The rule ANCHOR = the FranksVals tune, i.e. the reference run's central. +np_model_nu = tanh_2 +lambda2_nu = 0.15 +lambda4_nu = 0. +lambda6_nu = 0. +lambda_inf_nu = 2. +b0_over_bmax_nu = 1. +np_model = tanh_2 +lambda2 = 0.4 +lambda4 = 0.4 +lambda6 = 0. +lambda_inf = 1. +delta_lambda2 = 0.0 +np_model_tmd = off +lambda4_i = 0. + +[TNPs] +# theta = 0 with 'level0' IS the N^{3+0}LL prescription; not a no-op. +gamma_cusp = (0., 'level0') +gamma_mu_q = (0., 'level0') +gamma_nu = (0., 'level0') +h_qqV = (0., 'level0') +s = (0., 'level0') +b_qqV = (0., 'relative') +b_qqbarV = (0., 'relative') +b_qqS = (0., 'relative') +b_qqDS = (0., 'relative') +b_qg = (0., 'relative') diff --git a/scripts/rabbit/scetlib_ad/conf/Z_CT18Z_N3p0LL_analysis.conf b/scripts/rabbit/scetlib_ad/conf/Z_CT18Z_N3p0LL_analysis.conf new file mode 100644 index 000000000..46e8ca7b5 --- /dev/null +++ b/scripts/rabbit/scetlib_ad/conf/Z_CT18Z_N3p0LL_analysis.conf @@ -0,0 +1,136 @@ +#------------------------------------------------------------------------------- +# SCETlib autodiff runcard matching the ANALYSIS setup for Z at 13 TeV, CT18Z, +# N3+0LL + NNLO. Use it as --base-conf for +# scripts/rabbit/scetlib_ad/prepare_cache_for_card.py, which replaces the Grid_* +# sections with the card's own gen binning. +# +# Provenance: every [Calculation_settings], [Integration] and [TNPs] entry below +# is copied from the bT-grid production cards for this tune, +# /scratch/submit/cms/wmass/scetlib_np/Z_COM13_CT18Z_N3p0LL_btgrid_fineall/ +# base.conf +# inclusive_Z_COM13_CT18Z_N3p0LL_btgrid_fineall_b0nu1.ini +# so that the autodiff prediction is the SAME calculation the analysis's theory +# correction was built from. The examples/matched_ad/matched.conf runcard is NOT +# a substitute: it leaves lambda, the transition points, the scale floors, +# muf_follows_muB, compensate_fo and form_np_prescription at the SCETlib +# defaults, and it has no [TNPs] block -- i.e. plain N3LL with a different +# profile. Measured on the 30-bin debug grid, that mismatch moves the lambda +# RESPONSE by 7% (lambda2, delta_lambda2), 10% (lambda2_nu) and 27-35% +# (lambda4, lambda4_nu) relative to the response itself. +# +# Two deliberate departures from the bT-grid cards, both required: +# * calculation_piece = matched (they run 'sing' only, and get the nonsingular +# from DYTurbo; the autodiff cache carries SCETlib's own matched total); +# * a real [Nonperturbative] block -- the bT-grid production calls +# force_np_off() and applies the NP factors afterwards in TensorFlow, while +# here the NP lambdas ARE the differentiated parameters and their values are +# the rule anchor. Set them to the card's lambda_central. +# +# NOTE the [TNPs] block is what makes this N^{3+0}LL rather than N3LL, and it is +# also what registers the 10 TNP gradient parameters. Matching the analysis order +# and getting differentiable TNPs are the same action; the parameter vector is 19 +# long (alphas + 4 eff NP + 4 gamma_nu NP + 10 TNPs), not 9. +#------------------------------------------------------------------------------- + +[Process] +boson = Z + +[Calculation_settings] +calculation_piece = matched +run_order = n3ll +fixed_order = nnlo +# Required for the nnlo nonsingular: the native V+jet matrix elements are +# O(alphas) only, and this supplies the O(alphas^2) term. Do NOT set +# pick_structure_function -- it silently swaps the exact O(alphas^2) U+L for the +# NLO/LO graft (see examples/matched_ad/README.md). +fo_order2_analytic = yes + +# --- profile / scale block, verbatim from the bT-grid base.conf +recoil_scheme = collins_soper +alphas_solution = analytic +rge_solution = analytic +disable_asymmetry = yes +profile_functional_form = slope +form_np_prescription = collins_soper4 +muf_max = 13000. +b0_over_bmax = 0. +b0_over_bmax_global = 0. +lambda = 0. +mu0_min = 1. +muB_min = 1. +muS_min = 1. +nuS_min = 0. +muf_min = 1.40 +compensate_fo = yes +transition_points: [0.2, 0.6, 1.0] +muf_follows_muB = no + +[Singlet_scheme] +vector_singlet_enabled = yes +axial_singlet_enabled = yes +axial_singlet_use_exact_top_mass = no + +[Electroweak] +alphaem = 0.0077624494296896114 +sin2_thw = 0.23153999447822571 +mZ = 91.153509740726733 # GeV +GammaZ = 2.4932018986110700 # GeV +mW = 79.906853549493746 # GeV +GammaW = 2.0904310808144846 # GeV +ckm = [ [0.97446, 0.222, 0.00365], [0.22438, 0.97359, 0.04214], [0.00896, 0.04133, 0.999105] ] + +[QCD] +nf = 5 +alphas_mu0 = 0.118 +mu0 = 91.1876 +alphas_order = n3ll +pdf_set = CT18ZNNLO +pdf_member = 0 +Ecm = 13000. + +[Integration] +# The bT-grid production ran at 1e-5, but that is a POINT grid; here every bin is +# an adaptive 3-D cubature. The subdivision is driven by the VALUE component, and +# the derivative components are less well resolved than the value at the same +# target (doc/autodiff-design.md: ~4e-6 on the value against ~1e-2 median on the +# gradients at a 1e-3 target), so this is the knob to tighten if the lambda +# response looks noisy -- at roughly a factor 2 in cost per decade. +target_precision_rel = 1.e-4 +# MUST be non-zero: pieces whose integrand diverges as qT -> 0 can never satisfy +# a purely relative tolerance. +target_precision_abs = 1.e-9 +precision_buffer_bT = 1.e-2 +precision_buffer_decay = 1.e-1 +change_var_Q = arctan_Q2 +change_var_Q_q0 = 91.15348061918276 + +[Nonperturbative] +# The rule ANCHOR. Must equal the nonperturbative anchor the card records -- +# the model refuses to run otherwise. Values below are the LatticeNP tune. +np_model_nu = tanh_2 +lambda2_nu = 0.0870 +lambda4_nu = 0.0074 +lambda_inf_nu = 1.6853 +b0_over_bmax_nu = 1. +np_model = tanh_2 +lambda2 = 0.25 +delta_lambda2 = 0.0 +lambda4 = 0.06 +lambda_inf = 1. + +[TNPs] +# theta = 0 with 'level0' IS the N^{3+0}LL prescription the analysis uses (the +# "N3p0LL" in the correction names); it is not a no-op. It also registers each +# tag as a differentiable gradient parameter. +gamma_cusp = (0., 'level0') +gamma_mu_q = (0., 'level0') +gamma_nu = (0., 'level0') +h_qqV = (0., 'level0') +s = (0., 'level0') +# The beam TNPs have no native estimate; 'relative' multiplies the highest-order +# coefficient by (1 + theta), which at theta = 0 is the N^{3+0}LL baseline. +b_qqV = (0., 'relative') +b_qqbarV = (0., 'relative') +b_qqS = (0., 'relative') +b_qqDS = (0., 'relative') +b_qg = (0., 'relative') diff --git a/scripts/rabbit/scetlib_ad/make_debug_card.py b/scripts/rabbit/scetlib_ad/make_debug_card.py new file mode 100755 index 000000000..7b9612397 --- /dev/null +++ b/scripts/rabbit/scetlib_ad/make_debug_card.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +"""Build a gen-level rabbit card straight from a SCETlib autodiff cache. + +The card's single channel IS the cache's (qT, |Y|) grid, its signal template is +sigma_gen at the cache anchor, and its data is sigma_gen at whatever point you +ask for. That makes it a self-contained closure test of +``SCETlibADParamModel``: everything the card knows comes from the same backend +the model fits with, so a failure to recover the injected point is a bug in the +model, not a mismatch between two predictions. + +It deliberately does NOT replace the real gen-level sigmaUL card +(``feedRabbitSigmaUL.py``); it exists so the plumbing can be exercised before a +production cache is built. + + source /setup.sh + python scripts/rabbit/scetlib_ad/make_debug_card.py \ + --conf <...>/matched.conf --cache <...>/cache_debug.npz \ + --inject np_gnu_lambda2=0.12,alphas=0.1195 \ + -o /tmp/scetlib_ad_debug +""" + +import argparse +import os +import sys + +import hist +import numpy as np + +sys.path.insert( + 0, + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ), +) + +from rabbit import tensorwriter # noqa: E402 +from wremnants.postprocessing.scetlib_ad import params as adp # noqa: E402 +from wremnants.postprocessing.scetlib_ad.response import ( # noqa: E402 + NP_ANCHOR_META_KEY, +) +from wremnants.postprocessing.scetlib_ad.xsec_backend import ( # noqa: E402 + ScetlibADXsec, +) + +QT_AXIS = "ptVGen" +ABSY_AXIS = "absYVGen" +PROC = "Zmumu" +CHANNEL = "chSigmaUL" + + +def gen_axes_from_bins(bins): + """Recover the (qT, |Y|) product grid and the Q window from the cache bins.""" + Q = np.unique(bins[:, 0:2], axis=0) + if Q.shape[0] != 1: + raise SystemExit(f"expected a single Q bin in the cache, got {Q}") + y = np.unique(bins[:, 2:4], axis=0) + t = np.unique(bins[:, 4:6], axis=0) + y = y[np.argsort(y[:, 0])] + t = t[np.argsort(t[:, 0])] + if y.shape[0] * t.shape[0] != bins.shape[0]: + raise SystemExit( + "the cache bins are not a full (qT x |Y|) product grid; this debug " + "card maker only handles product grids" + ) + return ( + [ + (QT_AXIS, np.concatenate([t[:, 0], t[-1:, 1]])), + (ABSY_AXIS, np.concatenate([y[:, 0], y[-1:, 1]])), + ], + float(Q[0, 0]), + float(Q[0, 1]), + ) + + +def parse_kv(spec): + if not spec: + return {} + out = {} + for item in spec.split(","): + item = item.strip() + if not item: + continue + k, v = item.split("=", 1) + out[k.strip()] = float(v) + return out + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--conf", required=True) + ap.add_argument("--cache", required=True) + ap.add_argument("-o", "--outdir", required=True) + ap.add_argument("--outname", default="debug_card") + ap.add_argument("--threads", type=int, default=0) + ap.add_argument( + "--inject", + default=None, + help="SCETlib parameter values the DATA is generated at, " + "'np_gnu_lambda2=0.12,alphas=0.1195'. Default: the anchor (Asimov).", + ) + ap.add_argument( + "--yield-total", + type=float, + default=1.0e6, + help="scale sigma_gen to this total number of events, which is what sets " + "the Poisson stat power (and hence the recovered uncertainties)", + ) + args = ap.parse_args() + + core = ScetlibADXsec(args.conf, args.cache, threads=args.threads) + gen_axes, q_lo, q_hi = gen_axes_from_bins(core.bins) + fold = core.fold_for(gen_axes, q_lo, q_hi) + shape = fold.gen_shape + + def sigma(p): + vals, _ = core.values_and_jacobian(p) + return fold(np.asarray(vals, dtype=np.float64)).reshape(shape) + + nominal = sigma(core.anchor) + if np.any(nominal <= 0): + bad = int((nominal <= 0).sum()) + raise SystemExit( + f"{bad} bin(s) of sigma_gen at the anchor are non-positive; a Poisson " + f"card cannot be built from them. Restrict the cache's qT range." + ) + scale = args.yield_total / nominal.sum() + + p_data = core.anchor.copy() + injected = parse_kv(args.inject) + for name, val in injected.items(): + if name not in core.param_names: + raise SystemExit( + f"--inject: {name!r} is not a cache parameter {core.param_names}" + ) + p_data[core.param_names.index(name)] = val + data = sigma(p_data) + + axes = [ + hist.axis.Variable(edges, name=name, overflow=False, underflow=False) + for name, edges in gen_axes + ] + + def as_hist(values, weighted): + h = hist.Hist( + *axes, + storage=hist.storage.Weight() if weighted else hist.storage.Double(), + ) + if weighted: + h.values()[...] = values * scale + # Poisson-equivalent variance for an Asimov-style template. + h.variances()[...] = values * scale + else: + h.values()[...] = values * scale + return h + + writer = tensorwriter.TensorWriter() + writer.add_channel(axes, CHANNEL) + writer.add_data(as_hist(data, weighted=False), CHANNEL) + writer.add_process(as_hist(nominal, weighted=True), PROC, CHANNEL, signal=True) + + # The anchor, in the shape the model cross-checks against + # Writing it means the debug card exercises the + # anchor guard rather than skipping it. + eff, gnu = {}, {} + for sname, value in zip(core.param_names, core.anchor): + try: + rname = adp.rabbit_name(sname) + except KeyError: + continue + if sname.startswith("np_eff_"): + eff[rname] = float(value) + elif sname.startswith("np_gnu_"): + gnu[rname] = float(value) + conf_np = _np_models_from_conf(core.conf) + eff["np_model"] = conf_np[0] + gnu["np_model_nu"] = conf_np[1] + meta = {NP_ANCHOR_META_KEY: {"Z": {"eff_params": eff, "gnu_params": gnu}}} + + os.makedirs(args.outdir, exist_ok=True) + writer.write( + outfolder=args.outdir, outfilename=args.outname + ".hdf5", meta_data_dict=meta + ) + out = os.path.join(args.outdir, args.outname + ".hdf5") + print(f"\nwrote {out}") + print( + f" channel {CHANNEL}: {shape[0]} x {shape[1]} = {int(np.prod(shape))} bins " + f"({gen_axes[0][0]}, {gen_axes[1][0]}), Q [{q_lo:g}, {q_hi:g}]" + ) + print(f" template = sigma_gen(anchor), scaled to {args.yield_total:g} events") + print(f" data = sigma_gen({'anchor' if not injected else injected})") + if injected: + rel = np.max(np.abs(data / nominal - 1.0)) + print(f" injected shape change: max |data/nominal - 1| = {rel:.3%}") + + +def _np_models_from_conf(conf): + """(np_model, np_model_nu) from the runcard, for the lambda_central block.""" + npsec = conf["Nonperturbative"] if conf.has_section("Nonperturbative") else {} + return ( + npsec.get("np_model", "tanh_2"), + npsec.get("np_model_nu", "tanh_2"), + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/rabbit/scetlib_ad/prepare_cache_for_card.py b/scripts/rabbit/scetlib_ad/prepare_cache_for_card.py new file mode 100755 index 000000000..012ed1b8b --- /dev/null +++ b/scripts/rabbit/scetlib_ad/prepare_cache_for_card.py @@ -0,0 +1,418 @@ +#!/usr/bin/env python3 +"""Build a SCETlib autodiff cache for the gen binning of a specific rabbit card. + +The cache is only valid for the bins it was built for, so the binning should +come from the card rather than being kept in sync by hand. This reads the gen +axes out of a datacard -- the fit channel's own axes for a gen-level sigmaUL +card, or the response auxiliary's gen axes for a reco card -- writes the +matching SCETlib runcard next to the output, and runs the expensive build: + + compressed bin rules (resummed) + frozen fixed-order grid (nonsingular) + +Cost scales with the number of gen bins. Measured on this SCETlib build: +~0.34 s/bin of rule building and ~2 s/bin of fixed-order warming, and ~0.84 MB +of cache per bin. A 200-bin gen-level card is therefore minutes and ~170 MB; the +5740-bin correction grid is hours and several GB. + + source /setup.sh + python scripts/rabbit/scetlib_ad/prepare_cache_for_card.py \ + --card .hdf5 --base-conf /examples/matched_ad/matched.conf \ + -o /path/to/cachedir +""" + +import argparse +import configparser +import os +import sys +import time + +import h5py +import numpy as np + +sys.path.insert( + 0, + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ), +) + +from wremnants.postprocessing.scetlib_ad.response import ( # noqa: E402 + DEFAULT_RESPONSE_GROUP, +) +from wremnants.postprocessing.scetlib_ad.xsec_backend import ( # noqa: E402 + _scetlib_src, + bins_from_gen_axes, + configure, +) + + +def gen_axes_from_card(path, gen_level): + """[(qT name, edges), (|Y| name, edges)] for the card's gen grid.""" + from wums.ioutils import pickle_load_h5py + + with h5py.File(path, "r") as f: + if gen_level: + meta = pickle_load_h5py(f["meta"]) + channels = { + n: i + for n, i in meta["channel_info"].items() + if not i.get("masked", False) + } + if len(channels) != 1: + raise SystemExit( + f"expected a single non-masked channel, got {list(channels)}" + ) + info = next(iter(channels.values())) + axes = [ + (ax.name, np.asarray(ax.edges, dtype=np.float64)) for ax in info["axes"] + ] + if len(axes) != 2: + raise SystemExit( + f"expected 2 gen axes (qT, |Y|), got {[n for n, _ in axes]}" + ) + return axes + from rabbit.auxiliary import read_auxiliary_from_h5 + + aux = read_auxiliary_from_h5(f.get("auxiliary")) or {} + if DEFAULT_RESPONSE_GROUP not in aux: + raise SystemExit( + f"the card has no {DEFAULT_RESPONSE_GROUP!r} auxiliary, so its " + "gen binning is " + "unknown. Rebuild it with setupRabbit --storeResponseMatrix, or " + "pass --gen-level for a gen-level sigmaUL card." + ) + b = aux[DEFAULT_RESPONSE_GROUP] + names = [n.decode() if isinstance(n, bytes) else str(n) for n in b["gen_axes"]] + return [(n, np.asarray(b[f"edges__{n}"], dtype=np.float64)) for n in names] + + +def _upstream_prepare_cache(): + """The upstream example module, for its PDF-variation helpers. + + ``alphas_of`` / ``find_alphas_pair`` / ``pdf_set_size`` / + ``ensure_beamfunc_grids`` are non-trivial (the last one fans out one + single-threaded ~3.5 min process per PDF member and works around a shared + .info race), and they live in the SCETlib checkout we already depend on. + Loading them by path keeps one implementation, so an upstream fix arrives + for free -- importing is safe because its ``main()`` is guarded. + """ + import importlib.util + + path = os.path.join(_scetlib_src(), "examples", "matched_ad", "prepare_cache.py") + if not os.path.exists(path): + raise SystemExit( + f"scetlib_ad: cannot find {path}, needed for the PDF/alphaS/muF " + f"variation helpers. Pass --no-pdf to build a physics-only cache." + ) + spec = importlib.util.spec_from_file_location("_scetlib_prepare_cache", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def build_variations(sing, nons, bins, p0, names, conf, args): + """PDF eigenvector + alphaS-pair + muF variation members. + + Returns ``(n_eig, has_as, has_muf)`` for ``ScetlibCachedXsecTF``. These are + what turn the model from "alpha_s at fixed PDF plus the NP lambdas" into a + continuous parametrisation of the whole SCETlib theory uncertainty, so the + corresponding card templates can be dropped instead of double-counted. + """ + up = _upstream_prepare_cache() + repo = _scetlib_src() + pdf_set = conf["QCD"]["pdf_set"] + n_eig = args.pdf_eig if args.pdf_eig >= 0 else (up.pdf_set_size(pdf_set) - 1) // 2 + nf = conf["QCD"].getint("nf", fallback=5) + as_cen = float(p0[names.index("alphas")]) if "alphas" in names else 0.0 + as_step = 0.0 + pair = up.find_alphas_pair(pdf_set, args.as_pair, as_cen) + # muF is not live in the kernel (the beam convolutions are frozen at their + # own muF), so it rides on two extra members at kappa_F = 0.5 / 2.0, + # interpolated in t = ln(kappa_F)/ln(muf_hi): exact at 0.5, 1, 2. + muf_lo, muf_hi = (0.0, 0.0) if args.no_muf else (0.5, 2.0) + if not (n_eig or pair or muf_hi): + return 0, False, False + + members = list(range(1, 2 * n_eig + 1)) + sets = [pdf_set] * len(members) + if n_eig: + up.ensure_beamfunc_grids(repo, pdf_set, members, args.grid_jobs) + if pair: + # LAST and in (down, up) order -- both builders index the pair from the + # end. Its effect is added to the EXISTING alphas slot, so one parameter + # moves the calculation and the PDF together; without it, alphas is a + # partial derivative at fixed PDF. + down, up_set, as_step = pair + sets += [down, up_set] + members += [0, 0] + up.ensure_beamfunc_grids(repo, down, [0], args.grid_jobs) + up.ensure_beamfunc_grids(repo, up_set, [0], args.grid_jobs) + print( + f" alphaS pair: {down} / {up_set}, central {as_cen:.4f} " + f"+- {as_step:.4f}", + flush=True, + ) + if muf_hi: + # After the alphaS pair, in (lo, hi) order. The set/member entries are + # unused for these two -- the scale moves, not the PDF -- but one entry + # per member is still required. + sets += [pdf_set, pdf_set] + members += [0, 0] + print(f" muF pair: kappa_F = {muf_lo} / {muf_hi}", flush=True) + + mem = np.array(members, dtype=np.int32) + t0 = time.time() + sing.build_pdf_variations( + sets, + mem, + nf, + p0, + n_train_var=3, + n_eig=n_eig, + as_cen=as_cen, + as_step=as_step, + muf_lo=muf_lo, + muf_hi=muf_hi, + ) + print( + f" {n_eig} PDF eigenvector pairs for the resummed piece in " + f"{(time.time() - t0) / 60:.1f} min", + flush=True, + ) + t0 = time.time() + nons.build_fo_pdf_variations( + sets, + mem, + nf, + bins, + np.asarray(nons.gradient_central()), + n_eig=n_eig, + as_cen=as_cen, + as_step=as_step, + muf_lo=muf_lo, + muf_hi=muf_hi, + ) + print( + f" ... and for the fixed-order piece in " f"{(time.time() - t0) / 60:.1f} min", + flush=True, + ) + return n_eig, as_step > 0.0, muf_hi > 0.0 + + +def write_runcard(base_conf, out_path, gen_axes, Q_lo, Q_hi): + """Base runcard + the card's grids, written where the cache can find it.""" + conf = configparser.ConfigParser(inline_comment_prefixes="#") + conf.optionxform = str # SCETlib option names are case-sensitive + if not conf.read(base_conf): + raise SystemExit(f"cannot read base runcard {base_conf!r}") + for key, values in ( + ("Q", [Q_lo, Q_hi]), + ("Y", list(gen_axes[1][1])), + ("qT", list(gen_axes[0][1])), + ): + sec = f"Grid_{key}" + if not conf.has_section(sec): + conf.add_section(sec) + conf[sec]["custom_grid"] = "true" + conf[sec]["bins"] = "true" + conf[sec]["values"] = "[" + ", ".join(f"{v:g}" for v in values) + "]" + header = ( + "# Generated by scripts/rabbit/scetlib_ad/prepare_cache_for_card.py.\n" + f"# Base runcard: {os.path.abspath(base_conf)}\n" + "# The Grid_* sections are the card's gen binning; everything else is\n" + "# inherited. Keep this file next to the cache -- the fit needs it to\n" + "# rebuild the identical calculation the rules attach to.\n" + ) + with open(out_path, "w") as f: + f.write(header) + conf.write(f) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--card", default=None, help="rabbit datacard hdf5") + ap.add_argument( + "--grid-json", + default=None, + help="explicit gen grid instead of a card, as JSON: " + '\'{"Q": [60, 120], "Y": [...], "qT": [...]}\'. The Y edges may be ' + "signed (no |Y| folding assumed) -- useful for validating against a " + "SCETlib reference run on its own signed grid, and required for W.", + ) + ap.add_argument( + "--base-conf", + required=True, + help="SCETlib runcard supplying the physics settings " + "(orders, PDF, nonperturbative model); its Grid_* sections are replaced", + ) + ap.add_argument("-o", "--outdir", required=True) + ap.add_argument("--outname", default="cache") + ap.add_argument( + "-g", + "--gen-level", + action="store_true", + help="take the gen binning from the fit channel's own axes " + "(a gen-level sigmaUL card) instead of the response " + "auxiliary", + ) + ap.add_argument("--Q-lo", type=float, default=60.0) + ap.add_argument("--Q-hi", type=float, default=120.0) + ap.add_argument("--threads", type=int, default=0) + ap.add_argument( + "--n-train", + type=int, + default=9, + help="training points for the rule compression. Accuracy " + "tracks n_train/n_params, but the non-negative " + "least-squares solve grows roughly like n_train^2, so " + "raising it with the parameter count is the expensive " + "knob (see doc/autodiff-design.md)", + ) + ap.add_argument( + "--pdf-eig", + type=int, + default=-1, + help="number of PDF eigenvector pairs (default: all of the set). 0 " + "still keeps the alphaS pair, which is a separate direction.", + ) + ap.add_argument( + "--as-pair", + default="auto", + help="'auto' finds _as_0116/_as_0120, 'off' disables (alphas then " + "moves the calculation but NOT the PDF), or 'down,up' explicitly.", + ) + ap.add_argument( + "--no-muf", + action="store_true", + help="skip the muF member pair (then resumScaleMuF does nothing and the " + "card's resumFOScale* must be kept).", + ) + ap.add_argument( + "--no-pdf", + action="store_true", + help="physics-only cache: no PDF eigenvectors, no alphaS pair, no muF. " + "alphaS is then a fixed-PDF derivative -- do not quote it.", + ) + ap.add_argument( + "--grid-jobs", + type=int, + default=0, + help="parallel beamfunc-grid generation jobs (0 = one per core).", + ) + ap.add_argument( + "--dry-run", + action="store_true", + help="write the runcard and report the bin count and the " + "projected cost, without building", + ) + args = ap.parse_args() + + if (args.card is None) == (args.grid_json is None): + raise SystemExit("give exactly one of --card and --grid-json") + if args.grid_json: + import json + + g = json.loads(args.grid_json) + args.Q_lo, args.Q_hi = float(g["Q"][0]), float(g["Q"][-1]) + gen_axes = [ + ("qT", np.asarray(g["qT"], dtype=np.float64)), + ("Y", np.asarray(g["Y"], dtype=np.float64)), + ] + else: + gen_axes = gen_axes_from_card(args.card, args.gen_level) + n_bins = int(np.prod([len(e) - 1 for _, e in gen_axes])) + os.makedirs(args.outdir, exist_ok=True) + runcard = os.path.join(args.outdir, args.outname + ".conf") + write_runcard(args.base_conf, runcard, gen_axes, args.Q_lo, args.Q_hi) + print(f"gen binning from {args.card or 'explicit --grid-json'}:") + for name, edges in gen_axes: + print(f" {name:<12} {len(edges) - 1:4d} bins [{edges[0]:g}, {edges[-1]:g}]") + print(f" Q 1 bin [{args.Q_lo:g}, {args.Q_hi:g}]") + print(f" -> {n_bins} SCETlib bins; runcard written to {runcard}") + # Per-bin costs measured on a 5740-bin build with all cores busy. A small + # cache does not reach that: with 30 bins the same build ran ~10x slower per + # bin, because there is not enough work to fill the pool. + print( + f" projected (at full parallelism): ~{n_bins * 0.34 / 60:.0f} min of " + f"rules, ~{n_bins * 2.0 / 60:.0f} min of fixed-order warming, " + f"~{n_bins * 0.84:.0f} MB of cache" + ) + if args.dry_run: + return + + conf, sigma = configure(runcard, args.threads) + bins = bins_from_gen_axes(gen_axes, args.Q_lo, args.Q_hi) + p0 = np.asarray(sigma.gradient_central(), dtype=np.float64) + names = list(sigma.gradient_param_names()) + print(f"\n{len(p0)} differentiable parameters:") + for n, v in zip(names, p0): + print(f" {n:<24} {v:.6g}") + + sing, nons = sigma.sub_pieces() + t0 = time.time() + info = sing.build_bin_rules( + bins, + p0, + n_train=args.n_train, + n_hvp=1, + seed=4242, + n_jobs=args.threads or 0, + ) + nodes = [d["nodes"] for d in info] + resid = max(d["resid"] for d in info) + # The upstream docstring quotes ~1e-15 for a fully converged solve, but the + # active-set iteration cap deliberately trades residual for build time + # (see the build_bin_rules comment in py/qT/qT.cpp): ~1e-7..1e-9 is normal. + print( + f"\nrules built in {(time.time() - t0) / 60:.1f} min " + f"(median {int(np.median(nodes))} nodes/bin, worst training residual " + f"{resid:.1e})", + flush=True, + ) + if resid > 1e-6: + print( + f" WARNING: residual {resid:.1e} is large; the rules may not " + f"reproduce the direct calculation. Cross-check before using.", + flush=True, + ) + + # Populate the fixed-order grid over EVERY bin before saving. It fills lazily + # per evaluated region, so warming one bin would leave the rest to be computed + # on first use -- i.e. minutes of fixed-order work inside what is supposed to + # be the cheap stage. + t0 = time.time() + sigma.sigma_binned_batch(bins, p0) + print(f"fixed-order grid warmed in {(time.time() - t0) / 60:.1f} min", flush=True) + + if args.no_pdf: + n_eig, has_as, has_muf = 0, False, False + print( + "\n--no-pdf: physics-only cache. alphaS will be a derivative at " + "FIXED PDF, and the card's pdf*/pdfAlphaS/resumFOScale* templates " + "must be kept.", + flush=True, + ) + else: + n_eig, has_as, has_muf = build_variations( + sing, nons, bins, p0, names, conf, args + ) + + from scetlib_tf import ScetlibCachedXsecTF + + out = os.path.join(args.outdir, args.outname) + ScetlibCachedXsecTF( + sing, nons, bins=bins, n_eig=n_eig, has_as=has_as, has_muf=has_muf + ).save(out) + path = out + ".npz" + print(f"\nwrote {path} ({os.path.getsize(path) / 1e6:.1f} MB)") + print( + "\nNow check it:\n" + f" python scripts/rabbit/scetlib_ad/backend_check.py " + f"--conf {runcard} --cache {path}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/rabbit/scetlib_ad/validate_reco.py b/scripts/rabbit/scetlib_ad/validate_reco.py new file mode 100644 index 000000000..adaf2d381 --- /dev/null +++ b/scripts/rabbit/scetlib_ad/validate_reco.py @@ -0,0 +1,425 @@ +#!/usr/bin/env python3 +"""Validate the model's sigma_reco against the histmaker's corrected reco hist. + +This is the reco counterpart of ``compare_to_scetlib_run.py`` (which validates +sigma_gen against a native SCETlib run) and the same test the scetlib_np model +was validated with, so the numbers are directly comparable to its 0.14%. + +What is being compared, and why it is a real test: + + model : sigma_reco = R @ sigma_gen(p_anchor), with R = R_raw / N_gen read from + the datacard's response auxiliary. Pure theory + the response. + ref : the histmaker 'nominal' for the signal sample, whose central weight + carries the scetlib_dyturbo theory correction (i.e. NOT --theoryCorrAltOnly). + +Both are then the same physical object, so the per-bin ratio should be flat at 1. +Only the SHAPE is tested: the model is a cross section in pb, 'nominal' is an +xsec-weighted event yield, so summarize() applies ONE global scale and the plots +density-normalize. Never hard-code a pb->fb factor -- the residual total-sigma +difference would tilt the ratio by ~0.2%. + +Four traps carried over from the scetlib_np validation, all still live here: + + * R SUMS helicitySig while N_gen takes UL. Taking UL of R discards the angular + partition that the cosThetaStar*/phiStar* reco bins encode -- doing so blew + the closure from 0.14% to 2.0%. Handled inside response.py; do not "tidy". + * The gen ptVGen axis must carry the trailing overflow bin (e.g. [44, 100]): + ~6% of events reconstructed in the last ptll bin have true qT above the last + gen edge, and with no gen column there sigma_reco is low by exactly that. + Present in the card's gen axes; the cache must span it. + * NEVER hist.project() an aligned/cropped hist -- project() folds the summed + axes' FLOW bins back in, re-adding the content align_nominal just cropped + (measured as a spurious ~1.6% U-shape in yll). Use project_inrange(). + * pb vs fb: see above. + +Run inside the wmass singularity with the SCETlib setup sourced. +""" + +import argparse +import os + +import numpy as np + +NOMINAL_HIST = "nominal" +SIGNAL_SAMPLE = "Zmumu_2016PostVFP" +AXIS_LABELS = { + "ptll": r"$p_{T}^{\ell\ell}$ [GeV]", + "yll": r"$y^{\ell\ell}$", + "cosThetaStarll_quantile": r"$\cos\theta^{*}$ quantile", + "phiStarll_quantile": r"$\phi^{*}$ quantile", +} + + +def load_nominal(histmaker_path, sample_key, hist_name=NOMINAL_HIST): + """Load the histmaker 'nominal' Hist for the signal sample.""" + import h5py + + from wums import ioutils as wums_io + + with h5py.File(histmaker_path, "r") as f: + if sample_key not in f: + raise KeyError( + f"{histmaker_path}: no {sample_key!r} group. " + f"top-level: {[k for k in f.keys()][:12]}" + ) + sample = wums_io.pickle_load_h5py(f[sample_key]) + out = sample["output"] + if hist_name not in out: + raise KeyError( + f"{sample_key}: no {hist_name!r} hist. available: " + f"{list(out.keys())[:15]}" + ) + proxy = out[hist_name] + return proxy.get() if hasattr(proxy, "get") else proxy + + +def align_nominal(h, reco_axes_meta, tol=1e-6): + """Reorder + crop the nominal Hist onto the model's reco binning. + + project(*names) reorders (and sums out unlisted axes); then an integer-bin + slice per axis handles the histmaker axis being a SUPERSET of the fit axis + (ptll has a trailing [44, 100] bin while the fit stops at 44). The cropped + content lands in that axis' overflow -- which is why project_inrange, not + project, must be used downstream. + """ + names = [n for n, _ in reco_axes_meta] + have = [a.name for a in h.axes] + missing = [n for n in names if n not in have] + if missing: + raise ValueError(f"nominal hist missing axes {missing}; has {have}") + h = h.project(*names) + + crop = {} + for name, medges in reco_axes_meta: + medges = np.asarray(medges, dtype=np.float64) + hedges = np.asarray(h.axes[name].edges, dtype=np.float64) + nb = medges.size - 1 + hits = np.where(np.isclose(hedges, medges[0], atol=tol))[0] + if hits.size == 0: + raise ValueError( + f"axis {name}: model low edge {medges[0]} not found in hist " + f"edges [{hedges[0]} .. {hedges[-1]}]" + ) + i0 = int(hits[0]) + if i0 + nb + 1 > hedges.size or not np.allclose( + hedges[i0 : i0 + nb + 1], medges, atol=tol + ): + raise ValueError( + f"axis {name}: hist edges from index {i0} don't match model " + f"edges. hist={hedges[i0 : i0 + nb + 1]} model={medges}" + ) + if i0 != 0 or nb != h.axes[name].size: + crop[name] = slice(i0, i0 + nb) + return h[crop] if crop else h + + +def project_inrange(h, axis): + """1D Hist on ``axis``, summing the OTHER axes over in-range bins only.""" + import hist as _hist + + names = [a.name for a in h.axes] + ai = names.index(axis) + vals = h.values(flow=False) + other = tuple(i for i in range(vals.ndim) if i != ai) + out = _hist.Hist( + _hist.axis.Variable( + h.axes[axis].edges, name=axis, underflow=False, overflow=False + ), + storage=_hist.storage.Double(), + ) + out.view(flow=False)[...] = vals.sum(axis=other) + return out + + +def card_signal_column(indata, signal_proc_idx, reco_shape): + """indata.norm's signal column for the single unmasked fit channel. + + This is the reference the NEW (uncorrected-histmaker) construction actually + divides by, so validating against it is not the same test as validating + against the histmaker's own 'nominal' -- it carries the card's units and its + process decomposition. Slice by start:stop, NOT [:nbins]: a card with masked + channels puts this channel at an offset. + """ + info = next(i for _, i in indata.channel_info.items() if not i.get("masked", False)) + norm = indata.norm + if hasattr(norm, "todense"): + norm = norm.todense() + if hasattr(norm, "numpy"): + norm = norm.numpy() + norm = np.asarray(norm, dtype=np.float64) + col = norm[int(info["start"]) : int(info["stop"]), signal_proc_idx] + return col.reshape(reco_shape), float(info["lumi"]) + + +def summarize(model_sigma, nominal, names, match_norm=True): + """Diagnostics between the two reco tensors. + + match_norm=True applies ONE global scale first, so the result is a SHAPE + comparison and the absolute normalisation is divided out by construction. + That is the right test for the ratio construction (where the normalisation + cancels in sigma(p)/sigma(anchor)) and the WRONG one for the k*sigma_SC + construction, where the normalisation enters the prediction at first order. + """ + assert model_sigma.shape == nominal.shape, (model_sigma.shape, nominal.shape) + m = model_sigma.astype(np.float64) + n = nominal.astype(np.float64) + msum, nsum = m.sum(), n.sum() + print(f"\n sum model sigma_reco : {msum:.6g}") + print(f" sum nominal : {nsum:.6g}") + print(f" total model/nominal : {msum / nsum:.6g}") + if match_norm: + m_scaled = m * (nsum / msum) + print(" -> SHAPE mode: one global scale applied, total divided out") + else: + m_scaled = m + print(" -> ABSOLUTE mode: no scale matching, the total counts") + + good = n > 0 + if int((~good).sum()): + print(f" ({int((~good).sum())} of {n.size} bins have nominal<=0; excluded)") + r = (m_scaled[good] / n[good]).astype(np.float64) + w = n[good] + wmad = float(np.average(np.abs(r - 1.0), weights=w)) + lab = "scale*model / nominal" if match_norm else "model / nominal (ABSOLUTE)" + print(f"\n per-bin ratio ({lab}), should be ~1:") + print(f" bins : {r.size}") + print(f" mean / median : {r.mean():.5f} / {np.median(r):.5f}") + print(f" min / max : {r.min():.5f} / {r.max():.5f}") + for q in (1, 5, 50, 95, 99): + print(f" p{q:<2d} : {np.percentile(r, q):.5f}") + print(f" YIELD-WEIGHTED mean|ratio-1| : {wmad:.5f} <-- the headline") + for ax, name in enumerate(names): + other = tuple(i for i in range(m.ndim) if i != ax) + mp, npj = m_scaled.sum(axis=other), n.sum(axis=other) + with np.errstate(divide="ignore", invalid="ignore"): + rp = np.where(npj > 0, mp / npj, np.nan) + fin = rp[np.isfinite(rp)] + print( + f" projection {name:<26} max|ratio-1| = " + f"{np.max(np.abs(fin - 1)):.5f}" + ) + return wmad + + +def plot_axis( + model_h, + ref_h, + axis, + outdir, + tag, + meta, + density=True, + ref_label="histmaker nominal (theory-corrected)", +): + """Overlay + ratio panel on one reco axis. + + density=False keeps the absolute normalisation in the picture, which is + the whole point when the construction under test is k*sigma_SC. + """ + import hist + + from wums import output_tools, plot_tools + + os.makedirs(outdir, exist_ok=True) + m1, r1 = project_inrange(model_h, axis), project_inrange(ref_h, axis) + + def dens(h): + v = h.values(flow=False).astype(np.float64) + out = hist.Hist( + hist.axis.Variable( + h.axes[axis].edges, name=axis, underflow=False, overflow=False + ), + storage=hist.storage.Double(), + ) + out.view(flow=False)[...] = v / v.sum() if density else v + return out + + fig = plot_tools.makePlotWithRatioToRef( + [dens(r1), dens(m1)], + labels=[ref_label, "model $\\sigma_{reco}$"], + colors=["#5790fc", "#e42536"], + linestyles=["solid", "dashed"], + xlabel=AXIS_LABELS.get(axis, axis), + ylabel="normalized" if density else "yield", + rlabel=["model / nominal"], + rrange=[[0.97, 1.03]], + binwnorm=1, + logy=False, + yerr=False, + nlegcols=1, + cms_label="Work in progress", + grid=True, + ) + name = f"reco_{tag}_{axis}" if density else f"reco_{tag}_{axis}_abs" + plot_tools.save_pdf_and_png(outdir, name, fig=fig) + output_tools.write_index_and_log(outdir, name, analysis_meta_info=meta, args=None) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument( + "--datacard", required=True, help="reco card WITH the response auxiliary" + ) + ap.add_argument( + "--histmaker", + default=None, + help="histmaker hdf5 carrying reco 'nominal' (required for " + "--reference histmaker)", + ) + ap.add_argument("--cache", required=True) + ap.add_argument("--conf", required=True) + ap.add_argument("--sample", default=SIGNAL_SAMPLE) + ap.add_argument("--hist", default=NOMINAL_HIST) + ap.add_argument("--threads", type=int, default=32) + ap.add_argument("--plot-dir", default=None) + ap.add_argument("--tag", default="nominal") + ap.add_argument( + "--fit-params", + default="lambda2", + help="which parameters to register. Affects ONLY the " + "double-counting guards -- sigma_reco at the anchor is " + "independent of it -- so the default is one lambda, letting the " + "check run against a card that still carries pdfAlphaS / pdf* / " + "resum* templates.", + ) + ap.add_argument( + "--reference", + choices=("histmaker", "card"), + default="histmaker", + help="what to compare against. 'histmaker' = the theory-corrected reco " + "'nominal' (needs --histmaker). 'card' = indata.norm's signal column, " + "which is what the k*sigma_SC construction actually divides by; the " + "model is then scaled by the PHYSICAL k = lumi*1000.", + ) + ap.add_argument( + "--y-fold", + default="auto", + help="factor putting sigma on the card's |Y| convention. 'auto' reads " + "the cache's own declaration: a positive-side-only cache holds HALF the " + "|Y|-binned cross section, so it needs 2.0. Only matters when the " + "normalisation is not divided out, which is why the ratio construction " + "never had to care.", + ) + ap.add_argument( + "--no-match-norm", + dest="match_norm", + action="store_false", + help="do NOT apply a global scale before comparing, and do not " + "density-normalize the plots: compare absolutely.", + ) + ap.add_argument( + "--plot-axes", + nargs="*", + default=["ptll", "yll"], + help="reco axes to plot 1D ratios for", + ) + args = ap.parse_args() + + from rabbit.inputdata import FitInputData + from wremnants.postprocessing.scetlib_ad.param_model import SCETlibADParamModel + + if args.reference == "histmaker" and not args.histmaker: + raise SystemExit("--reference histmaker needs --histmaker") + print(f"datacard : {args.datacard}") + print(f"reference : {args.reference}") + print(f"histmaker : {args.histmaker}") + print(f"cache : {args.cache}") + indata = FitInputData(args.datacard) + # jitCompile: the model refuses to construct without it, because rabbit + # XLA-compiles compute() by default and XLA cannot lower tf.py_function. + # Constructing it outside rabbit_fit still has to answer that question. + model = SCETlibADParamModel( + indata, + cache=args.cache, + conf=args.conf, + gen_level=0, + threads=args.threads, + fit_params=args.fit_params, + # keep the POI inside fit_params (rabbit's layout contract); which + # parameter it is does not matter for the anchor prediction. + poi_params=args.fit_params.split(",")[0], + jitCompile="off", + ) + reco_axes = model._fit_axes(indata) + names = [n for n, _ in reco_axes] + print("reco axes : " + ", ".join(f"{n}({len(e)-1})" for n, e in reco_axes)) + print(f"reco shape: {model.reco_shape} gen shape: {model.gen_shape}") + + sig = model.sigma_reco_central + if sig is None: + raise SystemExit("model.sigma_reco_central is None -- did gen_level stay set?") + m = np.asarray(sig.numpy() if hasattr(sig, "numpy") else sig, dtype=np.float64) + m = m.reshape(model.reco_shape) + + if args.reference == "card": + import hist as _hist + + n, lumi = card_signal_column(indata, model.signal_proc_idx, model.reco_shape) + k = lumi * 1000.0 + conv = getattr(getattr(model, "_fold", None), "y_convention", "unknown") + if args.y_fold == "auto": + y_fold = 2.0 if conv == "positive-side-only" else 1.0 + else: + y_fold = float(args.y_fold) + print(f"reference : card indata.norm, proc idx {model.signal_proc_idx}") + print(f"k : lumi*1000 = {lumi} * 1000 = {k:.6g} (physical, pb->yield)") + print(f"y fold : x{y_fold:g} (cache Y convention: {conv})") + m = m * k * y_fold + ref = _hist.Hist( + *[ + _hist.axis.Variable(e, name=nm, underflow=False, overflow=False) + for nm, e in reco_axes + ], + storage=_hist.storage.Double(), + ) + ref.view(flow=False)[...] = n + else: + ref = align_nominal( + load_nominal(args.histmaker, args.sample, args.hist), reco_axes + ) + n = np.asarray(ref.values(flow=False), dtype=np.float64) + wmad = summarize(m, n, names, match_norm=args.match_norm) + + if args.plot_dir: + import hist + + axes = [ + hist.axis.Variable(e, name=nm, underflow=False, overflow=False) + for nm, e in reco_axes + ] + mh = hist.Hist(*axes, storage=hist.storage.Double()) + mh.view(flow=False)[...] = m + meta = { + "datacard": args.datacard, + "histmaker": args.histmaker, + "cache": args.cache, + "runcard": args.conf, + "hist": args.hist, + "reference": args.reference, + "norm matching": "on (SHAPE only)" if args.match_norm else "off (ABSOLUTE)", + "y fold": str(args.y_fold), + "yield-weighted mean|ratio-1|": f"{wmad:.5f}", + } + for ax in args.plot_axes: + if ax in names: + plot_axis( + mh, + ref, + ax, + args.plot_dir, + args.tag, + meta, + density=args.match_norm, + ref_label=( + "card nominal ($\\mathrm{indata.norm}$)" + if args.reference == "card" + else "histmaker nominal (theory-corrected)" + ), + ) + print(f"\n plots -> {args.plot_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/rabbit/scetlib_ad/validate_variations.py b/scripts/rabbit/scetlib_ad/validate_variations.py new file mode 100644 index 000000000..a0e8680ed --- /dev/null +++ b/scripts/rabbit/scetlib_ad/validate_variations.py @@ -0,0 +1,413 @@ +#!/usr/bin/env python3 +"""Validate EVERY theory variation of the model against the old template. + +The central-value validations (``compare_to_scetlib_run.py``, ``validate_reco.py``) +test one point. This tests the RESPONSE, which is what a fit actually uses: for +each variation the theory-correction file carries, compare + + model : sigma_gen(p_var) / sigma_gen(p_anchor) -- the AD prediction + ref : Corr[var] / Corr[central] -- the template it replaces + +Both sides are variation/central ratios, so no normalisation enters and the +comparison is non-circular. A flat 1.0 means the continuous direction reproduces +the discrete template it is meant to retire. + +Needs no datacard: a gen-level ratio is cache + runcard only. The reference's +fine (absY, qT) bins are summed onto the cache's gen grid, which is exact +because both sides are bin-integrated -- summing numerator and denominator +separately, then dividing, which is what the fit's per-bin rnorm does. + +Note the two JOINT variations (muf and kappaFO together): those test the model's +cross-term between muR and muF, which a template outer product cannot represent +and which is a large part of why we want the continuous treatment. +""" + +import argparse +import os +import re +import sys + +import numpy as np + +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..") +) + +from wremnants.postprocessing.scetlib_ad.xsec_backend import ScetlibADXsec # noqa: E402 + +# Variation label -> {SCETlib parameter: PHYSICAL value}. Absolute values, not +# offsets. Derived from prod/scetlib_run's variations_resummed.conf; the beam +# TNPs are 'relative' mode at +-0.5, the others 'level0' at +-1. +VARIATIONS = { + "lambda2_nu0.05": {"np_gnu_lambda2": 0.05}, + "lambda2_nu0.25": {"np_gnu_lambda2": 0.25}, + "lambda20.0": {"np_eff_lambda2": 0.0}, + "lambda21.0": {"np_eff_lambda2": 1.0}, + "delta_lambda2-0.02": {"np_eff_delta_lambda2": -0.02}, + "delta_lambda20.02": {"np_eff_delta_lambda2": 0.02}, + "lambda40.0": {"np_eff_lambda4": 0.0}, + "lambda41.0": {"np_eff_lambda4": 1.0}, + # kappaFO and kappaf move together so that muF is HELD: that combination is + # what our single kappa_R direction represents (see the upstream commit that + # made the scales differentiable -- a bare kappaFO is the wrong reference). + "kappaFO0.5-kappaf2.": {"scale_kappa_R": 0.5}, + "kappaFO2.-kappaf0.5": {"scale_kappa_R": 2.0}, + # muf up/down is a full factor of 2 (Scale_provider.cpp: pow(2., _vary.muf), + # enum up=+1 / down=-1). NB it also rescales muf_min, while our direction + # comes from members built at kappa_F = 0.5/2.0 -- whether those agree + # exactly is part of what this test measures. + "mufdown": {"scale_kappa_F": 0.5}, + "mufup": {"scale_kappa_F": 2.0}, + "mufdown-kappaFO0.5-kappaf2.": {"scale_kappa_F": 0.5, "scale_kappa_R": 0.5}, + "mufup-kappaFO2.-kappaf0.5": {"scale_kappa_F": 2.0, "scale_kappa_R": 2.0}, + "transition_points0.2_0.35_1.0": {"scale_x2": 0.35}, + "transition_points0.2_0.75_1.0": {"scale_x2": 0.75}, + # old central values, a cross-check variation rather than an uncertainty -- + # and the only one that needs x1/x3, which are frozen in the fit by default. + "transition_points0.3_0.6_0.9": {"scale_x1": 0.3, "scale_x3": 0.9}, +} +for _t in ("gamma_cusp", "gamma_mu_q", "gamma_nu", "h_qqV", "s"): + for _v in (-1.0, 1.0): + VARIATIONS[f"{_t}{_v:g}."] = {f"tnp_{_t}": _v} +for _t in ("b_qqV", "b_qqbarV", "b_qqS", "b_qqDS", "b_qg"): + for _v in (-0.5, 0.5): + VARIATIONS[f"{_t}{_v:g}"] = {f"tnp_{_t}": _v} + +CORR_HIST_SUFFIX = "_hist" +CENTRAL = "central" + +# The alphaS variations live in a SEPARATE corr file (`*_pdfas_CorrZ`), whose +# labels carry the PDF set name -- pdfCT18ZNNLO_as_0116, or ALPHAS_116 for +# HERAPDF -- so they are resolved by pattern rather than listed. as_0118 is that +# file's CENTRAL, not a variation. +_AS_RE = re.compile(r"(?:_as_0|ALPHAS_)(\d{3})$", re.I) +# PDF eigenvectors live in `*_pdfvars_CorrZ`: pdf0 is the central and +# pdf(2i+1)/pdf(2i+2) are eigenvector i up/down, i.e. c_e = +-1. Needs a cache +# built with n_eig > 0; otherwise these are reported as skipped, not silently +# passed. +_PDF_RE = re.compile(r"^pdf(\d+)$") + + +def variation_for(label): + """Label -> {SCETlib parameter: physical value}, or None if unmapped.""" + if label in VARIATIONS: + return VARIATIONS[label] + m = _AS_RE.search(label) + if m: + return {"alphas": float("0." + m.group(1))} + m = _PDF_RE.match(label) + if m: + n = int(m.group(1)) + if n == 0: + return None # the central of that file + i, side = (n - 1) // 2, (n - 1) % 2 # up first, then down + return {f"pdf_eig{i}": 1.0 if side == 0 else -1.0} + return None + + +def central_label(labels): + """The label a file's variations are ratios to.""" + if CENTRAL in labels: + return CENTRAL + for cand in labels: + m = _AS_RE.search(cand) + if m and m.group(1) == "118": + return cand + if "pdf0" in labels: + return "pdf0" + raise SystemExit(f"cannot identify the central among {labels[:6]}") + + +def load_corr(path): + """The theory-correction sigma hist (Q, absY, qT, charge, vars).""" + import pickle + + import lz4.frame + + with lz4.frame.open(path, "rb") as f: + d = pickle.load(f) + boson = next(k for k in d if k in ("Z", "W", "Wplus", "Wminus")) + inner = d[boson] + key = next(k for k in inner if k.endswith(CORR_HIST_SUFFIX) and "minnlo" not in k) + print(f"reference: {boson} / {key}") + return inner[key] + + +def merge_matrix(fine, coarse, name, tol=1e-9): + """(n_coarse, n_fine) 0/1 matrix summing fine bins into coarse ones.""" + fine = np.asarray(fine, float) + coarse = np.asarray(coarse, float) + M = np.zeros((coarse.size - 1, fine.size - 1)) + for k in range(coarse.size - 1): + lo, hi = coarse[k], coarse[k + 1] + idx = [ + i + for i in range(fine.size - 1) + if fine[i] >= lo - tol and fine[i + 1] <= hi + tol + ] + if not idx: + raise SystemExit(f"{name}: no reference bins inside [{lo}, {hi}]") + if abs(fine[idx[0]] - lo) > tol or abs(fine[idx[-1] + 1] - hi) > tol: + raise SystemExit( + f"{name}: coarse edges [{lo}, {hi}] are not reference edges; " + f"the model grid must be a sub-binning of the correction's." + ) + M[k, idx] = 1.0 + return M + + +def plot_response(label, Te, r_model, r_ref, outdir, meta): + """Model vs template RESPONSE on one variation, |Y| integrated. + + Both curves are variation/central ratios. Y is integrated by summing sigma + over |Y| for numerator and denominator SEPARATELY and then dividing -- not by + averaging per-bin ratios, which would weight the low-yield forward bins the + same as the peak. + """ + import hist + + from wums import output_tools, plot_tools + + os.makedirs(outdir, exist_ok=True) + + def h1(v): + h = hist.Hist( + hist.axis.Variable(Te, name="qT", overflow=False, underflow=False), + storage=hist.storage.Double(), + ) + h.view(flow=False)[...] = v + return h + + # Both curves are ratios, so centre the top panel on 1 rather than letting + # matplotlib pick an offset range -- an off-centre axis makes a symmetric + # response look like a trend. The floor keeps a near-null variation (e.g. + # b_qqDS, which is identically 1) from getting a degenerate range. + dev = max( + float(np.max(np.abs(np.asarray(r_ref, float) - 1.0))), + float(np.max(np.abs(np.asarray(r_model, float) - 1.0))), + ) + pad = max(1.2 * dev, 2.0e-3) + fig = plot_tools.makePlotWithRatioToRef( + [h1(r_ref), h1(r_model)], + labels=[f"template {label}", f"model {label}"], + ylim=[1.0 - pad, 1.0 + pad], + # mplhep loc: 0 = above the axes, 2 = top-left INSIDE the box (default). + logoPos=0, + colors=["#5790fc", "#e42536"], + linestyles=["solid", "dashed"], + xlabel=r"boson $q_\mathrm{T}$ (GeV)", + ylabel=r"$\sigma_\mathrm{var}/\sigma_\mathrm{central}$", + rlabel=["model / template"], + rrange=[[0.995, 1.005]], + binwnorm=None, + logy=False, + yerr=False, + nlegcols=1, + cms_label="Work in progress", + grid=True, + ) + safe = re.sub(r"[^A-Za-z0-9]+", "_", label).strip("_") + plot_tools.save_pdf_and_png(outdir, f"var_{safe}", fig=fig) + output_tools.write_index_and_log( + outdir, f"var_{safe}", analysis_meta_info=meta, args=None + ) + + +def _one_file( + path, + todo, + cen_lab, + ref_on_grid, + r_cen, + s_cen, + model_on_grid, + Te, + args, + rows, + skipped, +): + """Compare every mapped variation in one correction file.""" + for L in todo: + ov = variation_for(L) + if ov is None: + skipped.append((L, "no mapping")) + continue + s_var = model_on_grid(ov) + if s_var is None: + skipped.append((L, f"cache lacks {list(ov)}")) + continue + rm = s_var / s_cen + rr = ref_on_grid(L) / r_cen + good = np.isfinite(rm) & np.isfinite(rr) & (rr != 0) + dev = np.abs(rm[good] / rr[good] - 1.0) + # WHERE in qT the disagreement sits, which is the question that decides + # whether a residual is the known low-qT cutoff mismatch or something + # else. Arrays are (|Y|, qT) and C-ordered, so the qT index of the worst + # bin is the flat argmax modulo the number of qT bins. + with np.errstate(divide="ignore", invalid="ignore"): + dev2 = np.where(good, np.abs(rm / rr - 1.0), np.nan) + nqt = dev2.shape[1] + iq = int(np.nanargmax(dev2) % nqt) + rows.append((L, float(dev.max()), float(dev.mean()))) + print( + f"{L:<32} {dev.max():10.2e} {dev.mean():10.2e} " + f"[{rm[good].min():.4f},{rm[good].max():.4f}] " + f"[{rr[good].min():.4f},{rr[good].max():.4f}] " + f"{'[' + format(Te[iq], 'g') + ',' + format(Te[iq + 1], 'g') + ']':>12}" + ) + if args.profile: + per_qt = np.nanmax(dev2, axis=0) + print(" qT profile of max|dev| over |Y|:") + for k in range(nqt): + bar = "#" * int(min(40, round(40 * per_qt[k] / np.nanmax(per_qt)))) + print(f" [{Te[k]:6g},{Te[k + 1]:6g}] {per_qt[k]:10.2e} {bar}") + if args.plot_dir: + # Y-integrated response: sum sigma over |Y| first, then divide. + rm1 = s_var.sum(axis=0) / s_cen.sum(axis=0) + rr1 = ref_on_grid(L).sum(axis=0) / r_cen.sum(axis=0) + plot_response( + L, + Te, + rm1, + rr1, + args.plot_dir, + { + "variation": L, + "model setting": str(ov), + "reference": os.path.basename(path), + "cache": os.path.basename(args.cache), + "both curves": "variation / central (a RESPONSE, not a xsec)", + "max|model/template - 1| (all bins)": f"{dev.max():.3e}", + "mean|.| (all bins)": f"{dev.mean():.3e}", + }, + ) + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument( + "--corr", + required=True, + nargs="+", + help="one or more scetlib_dyturbo_*_Corr.pkl.lz4. Pass the matching " + "*_pdfas_CorrZ alongside the main file to validate alphaS: the alphaS " + "variations are NOT in the main correction.", + ) + ap.add_argument("--cache", required=True) + ap.add_argument("--conf", required=True) + ap.add_argument("--threads", type=int, default=32) + ap.add_argument( + "--only", nargs="*", default=None, help="restrict to these variation labels" + ) + ap.add_argument("--plot-dir", default=None) + ap.add_argument( + "--profile", + action="store_true", + help="print the qT profile of max|dev| over |Y| for each variation, " + "which separates a low-qT cutoff artefact from a genuine " + "disagreement in the response", + ) + args = ap.parse_args() + + core = ScetlibADXsec(args.conf, args.cache, threads=args.threads) + names = list(core.param_names) + print(f"cache: {core.n_bins} bins, {core.n_params} params") + + # the cache's gen grid, from the bins themselves + b = core.bins + yl = np.unique(np.round(b[:, 2:4], 12), axis=0) + tl = np.unique(np.round(b[:, 4:6], 12), axis=0) + yl, tl = yl[np.argsort(yl[:, 0])], tl[np.argsort(tl[:, 0])] + Ye = np.concatenate([yl[:, 0], yl[-1:, 1]]) + Te = np.concatenate([tl[:, 0], tl[-1:, 1]]) + print( + f"model grid: |Y| {Ye.size-1} bins [{Ye[0]:g}, {Ye[-1]:g}], " + f"qT {Te.size-1} bins [{Te[0]:g}, {Te[-1]:g}]" + ) + + def make_reference(path): + """(labels, central, ref_on_grid) for one correction file.""" + h = load_corr(path) + ax = {a.name: a for a in h.axes} + labels = [str(x) for x in ax["vars"]] + vals = np.asarray(h.values(flow=False)) + dims = [a.name for a in h.axes] + iQ, ich = dims.index("Q"), dims.index("charge") + if vals.shape[iQ] != 1 or vals.shape[ich] != 1: + raise SystemExit(f"{path}: expected a single Q and charge bin") + vals = np.squeeze(vals, axis=(iQ, ich)) + order = [d for d in dims if d not in ("Q", "charge")] + vals = np.moveaxis( + vals, + [order.index("absY"), order.index("qT"), order.index("vars")], + [0, 1, 2], + ) + MY = merge_matrix(ax["absY"].edges, Ye, "absY") + MT = merge_matrix(ax["qT"].edges, Te, "qT") + + def ref_on_grid(label): + return MY @ vals[:, :, labels.index(label)] @ MT.T # (nY, nT) + + return labels, central_label(labels), ref_on_grid + + # GenFold indexes in the order the gen axes are GIVEN, and the cache was + # built from the card as (ptVGen, absYVGen) -- gen shape (21, 10). Passing + # them Y-first makes it read the Y edges as qT and reject the cache. + fold = core.fold_for([("ptVGen", Te), ("absYVGen", Ye)], b[0, 0], b[0, 1]) + anchor = core.anchor.copy() + + def model_on_grid(overrides): + p = anchor.copy() + for k, val in overrides.items(): + if k not in names: + return None + p[names.index(k)] = val + vals_, _ = core.values_and_jacobian(p) + # fold -> (ptVGen, absYVGen); transpose to the (|Y|, qT) convention the + # reference side uses. + return fold(np.asarray(vals_, float)).reshape(Te.size - 1, Ye.size - 1).T + + s_cen = model_on_grid({}) + + print( + f"\n{'variation':<32} {'max|dev|':>10} {'mean|dev|':>10} " + f"{'model rng':>18} {'ref rng':>18} {'worst qT':>12}" + ) + rows, skipped = [], [] + for path in args.corr: + labels, cen_lab, ref_on_grid = make_reference(path) + r_cen = ref_on_grid(cen_lab) + todo = [ + L for L in labels if L != cen_lab and (args.only is None or L in args.only) + ] + if len(args.corr) > 1: + print(f" -- {os.path.basename(path)} (central: {cen_lab})") + _one_file( + path, + todo, + cen_lab, + ref_on_grid, + r_cen, + s_cen, + model_on_grid, + Te, + args, + rows, + skipped, + ) + if skipped: + print("\nskipped:") + for L, why in skipped: + print(f" {L:<32} {why}") + if rows: + worst = max(rows, key=lambda r: r[1]) + print( + f"\n{len(rows)} variations compared; worst max|dev| = " + f"{worst[1]:.2e} ({worst[0]})" + ) + + +if __name__ == "__main__": + main() diff --git a/wremnants/postprocessing/scetlib_ad/README.md b/wremnants/postprocessing/scetlib_ad/README.md new file mode 100644 index 000000000..3f48b62dd --- /dev/null +++ b/wremnants/postprocessing/scetlib_ad/README.md @@ -0,0 +1,136 @@ +# `scetlib_ad` — a fully differentiable SCETlib prediction for rabbit + +A rabbit `ParamModel` in which **every theory parameter SCETlib exposes is a +continuous fit parameter with exact derivatives**, rather than a discrete template +morph whose joint response with the others is an outer product: + +| parameter group | count (Z) | notes | +|---|---|---| +| `alphaS` | 1 | physical units; with a PDF α_s member pair it is the PDF-consistent coupling, not α_s at fixed PDF | +| nonperturbative λ | 8 | Collins–Soper and TMD form factors, `tanh_2`/`tanh_6` | +| theory nuisance parameters | 10 | `gamma_cusp`, `gamma_mu_q`, `gamma_nu`, `s`, `h_qqV` + 5 beam-function TNPs | +| PDF eigenvector coefficients | `n_eig` | extra differentiable columns, exact at `c_e = 0, ±1` | + +Which of these exist is a property of the **cache**, not of this code: the model +reads `gradient_param_names()` and registers what it finds. Only the +profile-scale parameters — `kappaFO`, `kappaf`, `muf`, the transition points — +are outside SCETlib's autodiff (they need d/dμ of the PDF convolution grids) and +still need template nuisances. + +The physics comes from the SCETlib `autodiff-sigmaul` branch: `ScetlibCachedXsecTF` +(`scetlib-cms/py/scetlib_tf.py`) replays a prepared cache — compressed bin rules +for the resummed piece plus a frozen fixed-order grid for the nonsingular one — +and returns exact first and second derivatives from clad. + +| file | role | +|---|---| +| `params.py` | SCETlib ↔ rabbit name map, prior σ, POI/POU defaults, impact groups | +| `xsec_backend.py` | `ScetlibADXsec` (configure + cache load + value/J/K) and `GenFold` (cache bins → the card's gen grid) | +| `response.py` | the reco fold: `P(b|g) = R_raw/N_gen`, the datacard response auxiliary, the positivity floor | +| `param_model.py` | `SCETlibADParamModel` — the rabbit adapter | + +Scripts live in `scripts/rabbit/scetlib_ad/`: + +| script | role | +|---|---| +| `backend_check.py` | standalone cache sanity: anchor round trip, FD-checked Jacobian, Hessian symmetry, fold sum rule | +| `prepare_cache_for_card.py` | build a cache for a card's gen binning, or an explicit `--grid-json` | +| `make_debug_card.py` | a self-contained gen-level card built from a cache, for closure tests | +| `compare_to_scetlib_run.py` | validate the resummed piece against a native SCETlib production run | +| `conf/Z_CT18Z_N3p0LL_FranksVals.conf` | runcard reproducing the current analysis central (see below) | + +## Running + +Inside the WRemnants singularity, `source setup.sh` then the SCETlib one: + +```bash +source $WREM_BASE/scetlib-cms/setup.sh # PYTHONPATH, LD_LIBRARY_PATH, ulimit -s + +# ONE job: minimize, postfit Hessian, impacts. +rabbit_fit.py .hdf5 -v 3 \ + --paramModel wremnants.postprocessing.scetlib_ad.SCETlibADParamModel \ + cache=.npz conf=.conf gen_level=1 threads=32 \ + --jitCompile off -t 0 --doImpacts -o +``` + +`--jitCompile off` is **mandatory** — the model reaches SCETlib through +`tf.py_function`, which XLA cannot compile. The model refuses to construct +otherwise, with that message. + +**One pass, not two.** The fit and the postfit covariance are a single job. The +model always includes the exact second-derivative term, so the composite Hessian +is exact and there is nothing to configure — no `--noHessian` fit followed by an +`--externalPostfit --noFit` pass. Measured on a 30-bin debug card, one pass +reproduces the two-pass numbers exactly (α_s = 0.1195 ± 0.00045, identical λ +uncertainties) with a much harder-converged EDM (6.1e-22 vs 3.5e-17). + +A Gauss-Newton variant dropping that term was measured ~5× faster and, on Asimov, +numerically identical — but it is an approximation on real data, so it is +deliberately **not** offered. The exact Hessian costs ~1 s/bin of serial work, +scaling as `1 + P(P+1)/2` in the parameter count: 3–80 s per minimiser iteration +on 64 threads at the few-hundred-to-1200-bin gen binnings we fit on. Reintroduce a +switch only if a fit on the full correction grid ever needs it. + +Two things keep that affordable: the `(value, J, K)` triple is cached on the +parameter vector, so one C++ Hessian is computed per *distinct* point rather than +per HVP (the minimiser takes many HVPs at fixed `x`); and the exact Hessian makes +the minimiser converge harder, so it needs fewer iterations. + +## How the derivatives get into TensorFlow + +`ScetlibCachedXsecTF` is an ordinary TF-differentiable function — its backward pass +is itself a `custom_gradient` whose own gradient contracts Hessian-vector +products — so nested `GradientTape`s work and TF drives every C++ call. The model +just calls it inside the graph, exactly as `examples/matched_ad/tf_gradients.py` +does. **There is no surrogate anywhere**: autodiff differentiates the real +prediction, and rabbit's postfit Hessian is the real Hessian. + +One requirement that imposes, and it is easy to break by accident: + +> Map rabbit's fit vector into SCETlib's layout with a **constant 0/1 matrix +> multiply**, never `tensor_scatter_nd_update`. + +Some mapping is unavoidable — rabbit's vector holds only the fitted parameters, +POIs first, while SCETlib's holds every registered parameter in registry order. But +a scatter's backward pass contains a gather, whose gradient TF represents as +`tf.IndexedSlices`, and the bridge's second-order py_function payloads call +`.numpy()` on the incoming cotangent and fail on it — so **everything past first +order breaks**, while first order keeps working, which makes it a nasty way to +fail. Isolated: a nested-tape HVP works on a bare `Variable`, and fails with a +scatter in front even when the scatter covers the whole vector. A constant matmul +is bit-identical (the entries are exactly 0 and 1) and free at these sizes. + +Fixing it upstream — densifying the incoming cotangent before `.numpy()` in +`_uhvp_py` — would remove the trap rather than leave us relying on a comment. + +## Validating a cache + +Two independent checks, both cheap relative to a fit: + +- `compare_to_scetlib_run.py` — against a native SCETlib production run with + `calculation_piece = sing`. That pkl **is** the resummed cross section, + bin-integrated, so replaying only our cache's rules gives the same object: no + matching, no fixed-order generator, no MC, no correction file in between. Any + disagreement is therefore ours — runcard, quadrature, Q integration, rule + compression. The reference's finer bins are summed onto ours, which is exact + because both sides are bin-integrated, and the script refuses to run unless our + edges really are a subset of the reference's. +- `backend_check.py` — the cache's own consistency, plus the fold sum rule. + +**Use an analysis runcard, not `examples/matched_ad/matched.conf`.** The example +runcard is plain N3LL with SCETlib's default profile scales; the analysis is +N³⁺⁰LL — every TNP at `theta = 0` with `'level0'` — with `lambda = 0`, transition +points `[0.2, 0.6, 1.0]`, scale floors, `compensate_fo` and `collins_soper4`. +Measured on a 30-bin grid, that difference moves the λ response by 7–35% of the +response itself. `conf/Z_CT18Z_N3p0LL_FranksVals.conf` transcribes the config of +the production run the current analysis correction is built from, with the +provenance in comments. + +Note the consequence: the analysis order is *defined* by the `[TNPs]` block, and a +non-`off` TNP scheme is exactly what registers a TNP as a gradient parameter. So +an analysis-faithful cache has 19 parameters, not 9. They are not fitted by +default — pass them in `fit_params`, and then `priors=1` is required (the model +refuses to float a TNP free, and refuses any parameter whose Jacobian column is +identically zero, which `resumTNP_b_qqDS` is for the Z). + +Study logbook: `WRemnantsHelpers/studies/scetlib-ad-param-model/LOGBOOK.md`. diff --git a/wremnants/postprocessing/scetlib_ad/__init__.py b/wremnants/postprocessing/scetlib_ad/__init__.py new file mode 100644 index 000000000..7c1be1fcb --- /dev/null +++ b/wremnants/postprocessing/scetlib_ad/__init__.py @@ -0,0 +1,29 @@ +"""SCETlib autodiff postprocessing package. + +``SCETlibADParamModel`` is a rabbit ParamModel whose prediction is differentiable +in every parameter SCETlib exposes -- alpha_s, the nonperturbative lambdas, the +theory nuisance parameters, and PDF eigenvector coefficients -- via the cached +matched cross section of the SCETlib ``autodiff-sigmaul`` branch. + +Imported lazily: the heavy dependencies (TensorFlow and the compiled +``scetlib_qT`` extension) should not be pulled in by a tool that merely wants the +parameter-name map or the response helpers. The package-level re-export +``wremnants.postprocessing.scetlib_ad.SCETlibADParamModel`` -- the name rabbit's +``--paramModel`` loader resolves -- still works, via PEP 562. +""" + +__all__ = ["SCETlibADParamModel", "ScetlibADXsec"] + + +def __getattr__(name): + if name == "SCETlibADParamModel": + from wremnants.postprocessing.scetlib_ad.param_model import ( + SCETlibADParamModel, + ) + + return SCETlibADParamModel + if name == "ScetlibADXsec": + from wremnants.postprocessing.scetlib_ad.xsec_backend import ScetlibADXsec + + return ScetlibADXsec + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/wremnants/postprocessing/scetlib_ad/param_model.py b/wremnants/postprocessing/scetlib_ad/param_model.py new file mode 100644 index 000000000..359dc7e57 --- /dev/null +++ b/wremnants/postprocessing/scetlib_ad/param_model.py @@ -0,0 +1,889 @@ +"""rabbit ParamModel for a fully differentiable SCETlib prediction. + +The prediction is assembled in three steps:: + + 1 SCETlib cached rule replay -> sigma(p; g) boson level, gen grid + 2 fold through the response R -> sigma_reco(p; b) gen -> reco + 3 ratio to the reference -> rnorm(b, proc) handed to rabbit + +``p`` is SCETlib's own differentiation vector, so every theory parameter the +calculation exposes is a continuous fit parameter with exact derivatives: + +* ``alphaS`` -- the strong coupling; +* the nonperturbative lambdas -- the Collins-Soper and TMD form factors; +* the theory nuisance parameters -- ``gamma_cusp``, ``gamma_mu_q``, + ``gamma_nu``, ``s``, ``h_qqV`` and the five beam-function TNPs, present + whenever the runcard declares a ``[TNPs]`` block (which is also what makes the + prediction N^{3+0}LL rather than N3LL); +* PDF eigenvector coefficients -- carried as additional differentiable columns + by a cache built with PDF variations, exact at ``c_e = 0, +-1``. Supplying the + PDF set's alpha_s member pair alongside them folds the + ``dsigma/dPDF . dPDF/dalpha_s`` piece into the ``alphaS`` slot, so ``alphaS`` + becomes the PDF-consistent coupling rather than alpha_s at fixed PDF. + +* the profile scales -- ``set_diff_scales(1)`` registers the + resummation ``kappa_R`` and the three matching transition points + ``x1..x3`` as differentiable, so these no longer need template nuisances. + ``kappa_F`` gets a slot too but is INERT in the kernel: it does nothing + unless the cache was built with the muF member pair, and a fit that tries to + float it is refused (see :meth:`_check_no_inert_params`). + +Which of these are present is a property of the cache, not of this file: the +model reads ``gradient_param_names()`` and registers what it finds. + +VALIDATION STATUS of the scale directions, which is NOT uniform -- see +``scripts/rabbit/scetlib_ad/validate_variations.py``: + +* the TNPs reproduce their templates to 1e-4..1e-16, the NP lambdas to ~1e-3; +* ``kappa_R`` reproduces ``kappaFO2.-kappaf0.5`` to 4.5e-03 but + ``kappaFO0.5-kappaf2.`` only to 4.0e-02 -- the down direction is 10x worse + than the up direction, and this is the direction that dominates + sigma(alpha_s) (rho(alphaS, resumScaleMuR) = +0.93); +* the TRANSITION POINTS DISAGREE IN SIGN with their templates, and the cause is + upstream, not here. All three ``transition_points*`` variations move the + prediction the opposite way from the reference (e.g. model [1.0000, 1.1593] + against reference [0.9602, 1.0000]). Making the identical physical change + through the RUNCARD with ``set_diff_scales`` off reproduces the template to + 2e-6; through the registered parameter with it on, the response is sign-flipped + and roughly -7x in slope. Moving the transition points moves ``muF`` by ~20% + (``muF`` has its own profile over the same points) while the per-node beam + convolutions stay frozen at the config's ``muF`` -- they shift 7-16% over that + range. ``kappa_R`` escapes this because ``set_muR_factor`` holds ``muF`` fixed + by construction. Not fixable from Python: the ``muF`` machinery interpolates a + GLOBAL member while the induced shift is PER NODE. All three + ``resumTransition*`` are therefore in :data:`params.DEFAULT_FROZEN`; that + removes the transition-point uncertainty from the fit, which is a known gap + rather than a fix. Re-run ``validate_variations.py`` before unfreezing. + +How the derivatives get into TensorFlow +-------------------------------------- +``ScetlibCachedXsecTF`` is an ordinary TF-differentiable function: its backward +pass is itself a ``custom_gradient`` whose own gradient contracts Hessian-vector +products, so nested ``GradientTape``s work and TF drives every C++ call. The model +simply calls it inside the graph, exactly as +``examples/matched_ad/tf_gradients.py`` does. There is no surrogate anywhere -- +autodiff differentiates the real prediction. + +One requirement that imposes, and that is easy to break by accident: map rabbit's +fit vector into SCETlib's layout with a CONSTANT 0/1 MATRIX MULTIPLY, never +``tensor_scatter_nd_update``. rabbit's vector holds only the fitted parameters, +POIs first, while SCETlib's holds every registered parameter in registry order, so +some mapping is unavoidable. A scatter's backward pass contains a gather, whose +gradient TF represents as ``tf.IndexedSlices``, and the bridge's second-order +py_function payloads call ``.numpy()`` on the incoming cotangent and fail on it -- +so anything past first order breaks. The matmul is bit-identical (entries are +exactly 0 and 1) and free at these sizes (at most ~25 x 25). + +XLA cannot compile a ``PyFunc``, so the fit MUST run with ``--jitCompile off``. +The model checks this at construction. +""" + +import copy +import re + +import numpy as np +import tensorflow as tf + +from rabbit.param_models.param_model import ParamModel +from wremnants.postprocessing.scetlib_ad import params as adp +from wremnants.postprocessing.scetlib_ad.response import ( + DEFAULT_RESPONSE_GROUP, + RATIO_FLOOR_MIN, + RATIO_FLOOR_SCALE, + R_info_from_auxiliary, + crop_R_to_fit, + marginalize_R_reco, + np_anchor_from_meta, +) +from wremnants.postprocessing.scetlib_ad.xsec_backend import ScetlibADXsec + +DTYPE = tf.float64 + +# Substring -> the registered parameter that makes a card syst a double count. +# Checked case-insensitively against indata.systs; only the entries whose model +# parameter is actually fitted are enforced, so a lambda-only run still tolerates +# a card carrying pdfAlphaS. +# Regex (matched case-insensitively against indata.systs) -> the registered +# parameter that makes a card syst a double count. Only the entries whose model +# parameter is actually FITTED are enforced, so a lambda-only run still tolerates +# a card carrying pdfAlphaS. +# +# Regex, not substring: the PDF eigenvector templates are pdfSym{Avg,Diff} +# and must be told apart from pdfAlphaS, which is a different physics direction +# living in a different model parameter. +_CONFLICTS = ( + ( + r"scetlibnp", + "any NP lambda", + lambda names: any(n.startswith("lambda") for n in names), + ), + (r"^pdfalphas", "alphaS", lambda names: "alphaS" in names), + ( + r"^resumtnp", + "a resummation TNP", + lambda names: any(n.startswith(adp.TNP_PREFIX_OUT) for n in names), + ), + ( + r"^pdf\d+", + "a PDF eigenvector coefficient", + lambda names: any(n.startswith(adp.PDF_PREFIX_OUT) for n in names), + ), + ( + r"^resumfoscale", + "the muR / muF profile scales", + lambda names: any(n in ("resumScaleMuR", "resumScaleMuF") for n in names), + ), + ( + r"^resumtransition", + "a matching transition point", + lambda names: any(n.startswith("resumTransition") for n in names), + ), +) + + +def _as_name_tuple(value): + """Spec tokens arrive as strings; accept ``a,b`` as well as a real tuple.""" + if value is None: + return () + if isinstance(value, str): + return tuple(s.strip() for s in value.split(",") if s.strip()) + return tuple(value) + + +class SCETlibADParamModel(ParamModel): + """Fit SCETlib's own differentiable parameters directly. + + Usage:: + + --paramModel wremnants.postprocessing.scetlib_ad.SCETlibADParamModel \\ + cache=.npz conf=.conf gen_level=1 [key=value ...] + """ + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + """``key=value`` spec tokens, typed off the ``__init__`` default. + + A later duplicate key wins, which is what lets a driver append overrides + to a spec it inherited from a previous step's recorded arguments. + """ + import inspect + + sig = inspect.signature(cls.__init__) + valid = { + n: p + for n, p in sig.parameters.items() + if n not in ("self", "indata") + and p.kind is not inspect.Parameter.VAR_KEYWORD + } + positional = [] + for tok in args: + key = tok.split("=", 1)[0] if isinstance(tok, str) and "=" in tok else None + if key is not None and key not in valid: + # Do NOT fall through to positional: every argument of this model + # is keyword-with-default, so an unknown key=value is a typo (or a + # token that has been removed), and silently treating it as + # positional produces "got multiple values for argument 'cache'". + raise TypeError( + f"{cls.__name__}: unknown spec token {key!r}. Valid tokens: " + + ", ".join(sorted(valid)) + ) + if key is not None: + val = tok.split("=", 1)[1] + default = valid[key].default + if isinstance(default, bool): + val = str(val).strip().lower() in ("1", "true", "yes", "on") + elif isinstance(default, float): + val = float(val) + elif isinstance(default, int): + val = int(val) + kwargs[key] = val + else: + positional.append(tok) + return cls(indata, *positional, **kwargs) + + def __init__( + self, + indata, + cache=None, + conf=None, + gen_level=False, + signal_proc="Zmumu", + Q_lo=60.0, + Q_hi=120.0, + fit_params=None, + poi_params="alphaS", + threads=0, + priors=False, + prior_sigmas=None, + xparam_default=None, + check_anchor=True, + anchor_tol=1e-6, + response_group=DEFAULT_RESPONSE_GROUP, + **kwargs, + ): + """ + Parameters + ---------- + cache, conf + The ``.npz`` written by ``ScetlibCachedXsecTF.save`` and the SCETlib + runcard it was built from. Both are required: the cache holds the + compressed rules and the frozen fixed-order grid, and the runcard is + what rebuilds the identical calculation they attach to. + gen_level + Gen-level sigmaUL mode. The fit channel IS the gen (qT, |Y|) binning, + so there is no response matrix and no fold: ``compute()`` returns the + per-gen-bin ratio ``sigma_gen(p) / sigma_gen(p_anchor)``. Otherwise + the reco fold reads R from the datacard's response auxiliary + (see ``response_group``). + signal_proc + Process whose column carries the ratio; the others stay at 1. + Q_lo, Q_hi + The mass window the cache's single Q bin must span. + fit_params + Comma-separated rabbit-facing names to expose to the fit. Default: + every parameter the cache carries except ``params.DEFAULT_FROZEN`` + (the tanh saturation scales and the b* convention, which are shape + constants). Parameters not listed are held at their cache anchor and + never reach rabbit, so they cannot contribute a zero-derivative + (singular) Hessian row. + poi_params + Subset of ``fit_params`` reported as POIs (they must come first in + the fitter's layout). Default ``alphaS``. + threads + SCETlib worker threads for the batch replay (0 = one per hardware + thread). + priors + Declare Gaussian priors. rabbit applies priors whenever a model + declares ``prior_sigmas``, so this token IS the decision. Off by + default -- everything floats free. Note the TNP defaults are sigma=1 + (they are genuine nuisances), unlike the lambdas. + prior_sigmas + Per-name override, ``name=value,...`` or a Mapping. ``nan`` frees a + parameter. + xparam_default + ``name=value,...`` shifting the fit START (and the prior mean) off + the cache anchor, for injection / closure tests. The ratio + DENOMINATOR is not moved -- it always stays the anchor. + check_anchor + Cross-check the cache anchor against the nonperturbative values the + card records for its own prediction. An anchor that disagrees + is the silent-wrong-answer trap documented in + ``knowledge/20_frameworks/gen_level_sigmaul_fit.md``: the ratio is + still 1 at the start, so nothing looks broken, but the response is + evaluated at the wrong point. + """ + self.indata = indata + if cache is None or conf is None: + raise ValueError( + "SCETlibADParamModel needs both cache=.npz and " + "conf=.conf spec tokens." + ) + self._require_no_xla(kwargs) + + self._response_group = str(response_group) + self.gen_level = bool(gen_level) + + # ---- Backend: rebuild the calculation and load the cache. + self.core = ScetlibADXsec(conf, cache, threads=threads) + self.scetlib_names = list(self.core.param_names) + self.rabbit_names = [adp.rabbit_name(n) for n in self.scetlib_names] + self._anchor = np.asarray(self.core.anchor, dtype=np.float64) + + # ---- Gen binning, and (reco path) the response matrix. + self._setup_binning(indata, Q_lo, Q_hi) + + # ---- Parameter registration. Everything the fit does NOT expose stays + # pinned at the anchor, so the SCETlib vector is always complete. + self._register_params(fit_params, poi_params, xparam_default) + + # ---- Central: the ratio denominator, evaluated by the model itself at + # the anchor so the ratio is exactly 1 at the start whatever the card's + # own template looks like. + sigma_gen_anchor = self._sigma_gen_np(self._p_base_anchor) + self.sigma_gen_central_flat = tf.constant(sigma_gen_anchor, dtype=DTYPE) + if self.gen_level: + self.sigma_reco_central = None + else: + self.sigma_reco_central = tf.linalg.matvec( + self.R, self.sigma_gen_central_flat + ) + n_bad = int(tf.reduce_sum(tf.cast(self.sigma_reco_central <= 0, tf.int32))) + if n_bad: + raise ValueError( + f"SCETlibADParamModel: {n_bad} reco bins have non-positive " + f"sigma_reco at the anchor. Likely a binning mismatch between " + f"R and the fit-tensor reco axes." + ) + + if check_anchor: + self._check_anchor_against_card(anchor_tol) + self._check_double_counting() + self._check_no_inert_params() + + # ---- Process column. + procs = [p.decode() if isinstance(p, bytes) else str(p) for p in indata.procs] + if signal_proc not in procs: + raise ValueError( + f"SCETlibADParamModel: signal_proc={signal_proc!r} not in " + f"indata.procs={procs[:10]}..." + ) + self.signal_proc_idx = procs.index(signal_proc) + self.nproc = indata.nproc + self._signal_col_mask = tf.reshape( + tf.one_hot(self.signal_proc_idx, self.nproc, dtype=indata.dtype), + [1, self.nproc], + ) + + self._setup_priors(priors, prior_sigmas) + + print( + f"[SCETlibADParamModel] {self.core} | {self._fold.describe()} | " + f"{'gen-level' if self.gen_level else 'reco'} | " + f"fitting {self.nparams} of {self.core.n_params} " + f"({self.npoi} POI: {[n for n in self._param_order[: self.npoi]]})", + flush=True, + ) + + # ========================================================================= + # construction helpers + # ========================================================================= + + def _require_no_xla(self, kwargs): + """Refuse to build under XLA -- a PyFunc has no XLA lowering. + + rabbit resolves ``--jitCompile auto`` to True in dense mode + (Fitter.__init__), and the failure is an opaque compile error deep in the + first loss evaluation, so trip here with the fix in the message. + """ + opt = str(kwargs.get("jitCompile", "auto")).lower() + sparse = bool(getattr(self.indata, "sparse", False)) + # --eager turns every tf.function into eager execution, so jit_compile + # never applies and a PyFunc is fine. + if kwargs.get("eager") or opt == "off" or (opt == "auto" and sparse): + return + raise ValueError( + "SCETlibADParamModel calls into SCETlib through tf.py_function, " + "which XLA cannot compile. Re-run with --jitCompile off " + f"(got --jitCompile {opt}" + (", dense input)." if not sparse else ").") + ) + + def _setup_binning(self, indata, Q_lo, Q_hi): + """Resolve the gen grid (and R, in the reco path) and map it onto the cache.""" + if self.gen_level: + gen_axes = self._fit_axes(indata) + if len(gen_axes) != 2: + raise NotImplementedError( + "gen_level SCETlibADParamModel expects a single fit channel " + "with 2 gen axes (qT, |Y|); got " + f"{[n for n, _ in gen_axes]}" + ) + self.R = None + self.reco_shape = None + else: + R_info = R_info_from_auxiliary(indata, self._response_group) + fit_reco_axes = self._fit_axes(indata) + R_full, R_reco_axes = marginalize_R_reco( + R_info["R"], R_info["reco_axes"], [n for n, _ in fit_reco_axes] + ) + R_arr = crop_R_to_fit(R_full, R_reco_axes, fit_reco_axes) + self.reco_shape = R_arr.shape[: len(fit_reco_axes)] + gen_axes = R_info["gen_axes"] + if R_info.get("N_gen") is None: + raise ValueError( + f"SCETlibADParamModel: the {self._response_group!r} " + "auxiliary has no N_gen " + "(gen-total). Rebuild the datacard from a histmaker output " + "that carries the 'prefsr' xnorm hist." + ) + n_reco = int(np.prod(self.reco_shape)) + n_gen = int(np.prod([len(e) - 1 for _, e in gen_axes])) + R_raw = tf.constant(R_arr.reshape(n_reco, n_gen), dtype=DTYPE) + n_gen_flat = tf.constant( + np.asarray(R_info["N_gen"]).reshape(-1), dtype=DTYPE + ) + # R must encode only the gen->reco mapping, not the MC's absolute gen + # spectrum: normalise each gen column by the gen-total (empty gen bins + # keep a zero column). + safe = tf.where(n_gen_flat > 0, n_gen_flat, tf.ones_like(n_gen_flat)) + self.R = R_raw / safe[tf.newaxis, :] + + self.gen_axes = [ + (name, np.asarray(edges, dtype=np.float64)) for name, edges in gen_axes + ] + self.gen_shape = tuple(len(e) - 1 for _, e in self.gen_axes) + self.Q_lo, self.Q_hi = float(Q_lo), float(Q_hi) + + # compute() returns one row per fit bin, so the shape it builds must be + # the card's. A mismatch here would surface as an opaque broadcast error + # inside the first loss evaluation. + n_rows = int(np.prod(self.gen_shape if self.gen_level else self.reco_shape)) + if n_rows != int(indata.nbins): + raise ValueError( + f"SCETlibADParamModel: the model produces {n_rows} bins " + f"({'gen' if self.gen_level else 'reco'} shape " + f"{self.gen_shape if self.gen_level else self.reco_shape}) but the " + f"card has {int(indata.nbins)}." + ) + # Exact sum of cache bins onto the gen grid: handles a different nesting + # order, a signed-Y cache folded onto |Y|, and a cache finer than the fit's + # gen binning. Coverage is verified, so a cache that does not tile this + # card's gen bins raises here rather than integrating over less phase space. + self._fold = self.core.fold_for(self.gen_axes, self.Q_lo, self.Q_hi) + + def _fit_axes(self, indata): + """(name, edges) of each axis of the single non-masked channel.""" + non_masked = [ + (name, info) + for name, info in indata.channel_info.items() + if not info.get("masked", False) + ] + if len(non_masked) != 1: + raise NotImplementedError( + f"SCETlibADParamModel supports a single non-masked channel; got " + f"{len(non_masked)}: {[n for n, _ in non_masked]}" + ) + _, info = non_masked[0] + return [ + (ax.name, np.asarray(ax.edges, dtype=np.float64)) for ax in info["axes"] + ] + + def _register_params(self, fit_params, poi_params, xparam_default): + """Decide which SCETlib parameters rabbit sees, and their start values.""" + available = list(self.rabbit_names) + # Default: alpha_s and the NP lambdas, minus the shape constants. TNPs are + # NOT included by default even when the cache carries them -- an + # analysis-faithful runcard registers all ten (theta=0 'level0' IS the + # N^{3+0}LL prescription), and silently floating ten unconstrained theory + # nuisances is not what "fit the NP model" should mean. Ask for them + # explicitly via fit_params, with priors. + # 'all' = every registered direction except the frozen shape constants, + # i.e. hand the whole SCETlib theory uncertainty to the model. That + # includes the TNPs, so priors are required (see the check below). + if _as_name_tuple(fit_params) == ("all",): + requested = tuple(n for n in available if n not in adp.DEFAULT_FROZEN) + else: + requested = _as_name_tuple(fit_params) or tuple( + n + for n in available + if n not in adp.DEFAULT_FROZEN and not n.startswith(adp.TNP_PREFIX_OUT) + ) + unknown = [n for n in requested if n not in available] + if unknown: + raise ValueError( + f"SCETlibADParamModel: fit_params {unknown} are not in this " + f"cache's parameter set {available}. A parameter can only be " + f"fitted if the runcard declared it before the rules were built." + ) + # Floating a direction whose response we have measured to be wrong is + # allowed -- it is how the fix gets tested -- but never silently. + for n in requested: + if n in adp.KNOWN_BAD_RESPONSE: + print( + f"[SCETlibADParamModel] WARNING: floating {n!r}, whose " + f"response is {adp.KNOWN_BAD_RESPONSE[n]}. It is in " + f"DEFAULT_FROZEN for that reason; this fit is a study, not a " + f"physics result." + ) + pois = _as_name_tuple(poi_params) + bad_pois = [n for n in pois if n not in requested] + if bad_pois: + raise ValueError( + f"SCETlibADParamModel: poi_params {bad_pois} are not in " + f"fit_params {list(requested)}." + ) + # rabbit's layout contract: all POIs first, then the POUs. + nou = tuple(n for n in requested if n not in pois) + self._param_order = tuple(pois) + nou + self.npoi = len(pois) + self.npou = len(nou) + self.params = np.array([p.encode() for p in self._param_order]) + + # Position of each fitted parameter inside SCETlib's own vector. rabbit's + # vector is NOT SCETlib's: it holds only what we fit, POIs first, while + # SCETlib's has every registered parameter in its registry order. + self._fit_idx = np.array( + [available.index(n) for n in self._param_order], dtype=np.int64 + ) + # The map from rabbit's vector to SCETlib's, as a constant 0/1 matrix. + # Deliberately NOT tensor_scatter_nd_update: a scatter's backward pass + # contains a gather, whose gradient TF represents as tf.IndexedSlices, and + # the SCETlib bridge's second-order py_function payloads call .numpy() on + # the incoming cotangent and die on it. A matmul against a constant gives + # a dense gradient, so differentiation survives at second + # order. Bit-identical to the scatter (the entries are exactly 0 and 1) + # and negligible in cost: (n_scetlib, n_fit) is at most ~25 x 25. + self._select = np.zeros((len(available), len(self._param_order))) + self._select[self._fit_idx, np.arange(len(self._param_order))] = 1.0 + + # Reparametrisation (see params.REPARAM): for the profile scales the + # FITTED parameter is a unit nuisance theta and the PHYSICAL value handed + # to SCETlib is a function of it. Stored as coefficient vectors so one + # vectorised expression covers every parameter, identity included, and + # the TF path stays a handful of elementwise ops with exact derivatives. + n_fit = len(self._param_order) + self._rp_log = np.zeros(n_fit, dtype=bool) + self._rp_quad = np.zeros(n_fit, dtype=bool) + self._rp_L = np.zeros(n_fit) + self._rp_c = np.zeros((3, n_fit)) + for i, name in enumerate(self._param_order): + spec = adp.reparam(name) + if spec is None: + continue + kind, coeffs = spec + if kind == "log": + self._rp_log[i] = True + (self._rp_L[i],) = coeffs + elif kind == "quad": + self._rp_quad[i] = True + self._rp_c[:, i] = coeffs + else: + raise ValueError(f"params.REPARAM: unknown kind {kind!r}") + self._rp_id = ~(self._rp_log | self._rp_quad) + self._reparametrised = tuple( + n for n, f in zip(self._param_order, ~self._rp_id) if f + ) + + # Start values: the cache anchor, optionally shifted for injection tests. + # _p_base_anchor is the UNSHIFTED full vector and stays the ratio + # denominator; _p_base carries the shift for the non-fitted slots only + # (fitted slots are overwritten from the fit vector on every call). + self._p_base_anchor = self._anchor.copy() + self._p_base = self._anchor.copy() + defaults = self._anchor[self._fit_idx].copy() + # A reparametrised parameter starts at theta = 0, NOT at its physical + # anchor. The check below is what guarantees theta = 0 maps back onto the + # anchor, so the ratio-to-central is exactly 1 at the start; a mistyped + # coefficient would otherwise shift the whole prediction silently. + defaults[~self._rp_id] = 0.0 + round_trip = self._physical(defaults) + if not np.allclose(round_trip, self._anchor[self._fit_idx], rtol=0, atol=1e-12): + bad = [ + (n, float(a), float(b)) + for n, a, b in zip( + self._param_order, round_trip, self._anchor[self._fit_idx] + ) + if abs(a - b) > 1e-12 + ] + raise ValueError( + "scetlib_ad: the REPARAM maps do not reproduce the cache anchor " + f"at theta = 0, so sigma_gen/sigma_central would not be 1: {bad}" + ) + for name, val in _parse_kv(xparam_default).items(): + if name not in available: + raise KeyError(f"xparam_default: unknown parameter {name!r}") + if name in self._param_order: + defaults[self._param_order.index(name)] = val + else: + # not fitted: pin the held value at the shifted point + self._p_base[available.index(name)] = val + print( + f"[SCETlibADParamModel] xparam_default {name}={val:g} applies " + f"to a NON-fitted parameter; it is pinned there, not floated.", + flush=True, + ) + if xparam_default and self._reparametrised: + print( + "[SCETlibADParamModel] NB xparam_default for " + f"{list(self._reparametrised)} is in THETA units (unit nuisance), " + "not physical units.", + flush=True, + ) + if xparam_default: + print( + "[SCETlibADParamModel] start shifted: " + f"{dict(zip(self._param_order, defaults))}", + flush=True, + ) + + # lambdas can be legitimately zero or negative (delta_lambda2), so store + # POIs directly rather than as sqrt(value). + self.allowNegativeParam = True + self.is_linear = False + self.xparamdefault = tf.constant(defaults, dtype=self.indata.dtype) + + active = set(self._param_order) + groups = { + label: tuple(p for p in members if p in active) + for label, members in adp.IMPACT_GROUP_MEMBERS.items() + } + tnps = adp.tnp_group(self._param_order) + if tnps: + groups["resumTNP"] = tnps + self.param_impact_groups = {k: v for k, v in groups.items() if v} + + def _setup_priors(self, priors, prior_sigmas): + """Declare ``prior_sigmas`` only when asked; rabbit's Fitter keys off it.""" + tnps = adp.tnp_group(self._param_order) + if tnps and not priors: + # theta is normalised upstream so |theta| = 1 IS the recommended + # variation; floating a TNP free discards that, and the fit will + # happily absorb a physical effect into an unconstrained nuisance. + raise ValueError( + f"SCETlibADParamModel: {len(tnps)} theory nuisance parameter(s) " + f"are being fitted ({', '.join(tnps[:4])}" + f"{', ...' if len(tnps) > 4 else ''}) but priors are off. TNPs " + f"carry an N(0,1) constraint by construction. Pass priors=1 (the " + f"registry gives every TNP sigma=1 and leaves the lambdas free), " + f"or drop them from fit_params." + ) + if not priors: + return + overrides = ( + _parse_kv(prior_sigmas) + if not isinstance(prior_sigmas, dict) + else dict(prior_sigmas) + ) + unknown = [k for k in overrides if k not in self._param_order] + if unknown: + raise KeyError( + f"prior_sigmas: {unknown} are not fitted parameters " + f"({list(self._param_order)})" + ) + sigmas = np.empty(self.nparams, dtype=np.float64) + for i, name in enumerate(self._param_order): + s = overrides.get(name, adp.prior_sigma(name)) + sigmas[i] = np.nan if s is None else float(s) + self.prior_sigmas = sigmas + constrained = { + n: s for n, s in zip(self._param_order, sigmas) if np.isfinite(s) and s > 0 + } + print( + f"[SCETlibADParamModel] Gaussian priors on {len(constrained)} " + f"parameter(s): {constrained}", + flush=True, + ) + + def _check_anchor_against_card(self, tol): + """Compare the cache anchor with the card's propagated lambda_central.""" + meta = getattr(self.indata, "metadata", None) + if not meta: + print( + "[SCETlibADParamModel] WARNING: the card carries no metadata, so " + "the cache anchor could not be cross-checked against the " + "nonperturbative values the card was produced with. Pass " + "check_anchor=0 to silence, but verify by hand.", + flush=True, + ) + return + card = np_anchor_from_meta(meta) + if not card: + print( + "[SCETlibADParamModel] WARNING: the card records no " + "nonperturbative anchor, so the cache anchor was NOT " + "cross-checked.", + flush=True, + ) + return + mismatched = [] + for card_key, rabbit in adp.LAMBDA_CENTRAL_KEYS.items(): + if card_key not in card or rabbit not in self.rabbit_names: + continue + want = float(card[card_key]) + have = float(self._anchor[self.rabbit_names.index(rabbit)]) + if abs(have - want) > tol * max(1.0, abs(want)): + mismatched.append((rabbit, want, have)) + if mismatched: + raise ValueError( + "SCETlibADParamModel: the cache anchor does not match the card's " + "NP central. The ratio would still be 1 at the start, so this " + "fails silently -- rebuild the cache at the card's runcard " + "values, or the card at the cache's.\n" + + "\n".join( + f" {n}: card {w:.6g} vs cache {h:.6g}" for n, w, h in mismatched + ) + ) + + def _check_no_inert_params(self): + """Refuse a fitted parameter the prediction does not depend on. + + Not hypothetical: with the analysis runcard ``tnp_b_qqDS`` has an + identically zero gradient for the Z (the channel it scales does not + contribute), and ``lambda_inf`` / ``lambda_inf_nu`` are nearly inert at + the anchor. A zero Jacobian column is a zero row and column of the NLL + Hessian, i.e. a singular covariance and a meaningless "uncertainty". + """ + p_start = np.asarray(self.xparamdefault.numpy(), dtype=np.float64) + _, jac = self.core.values_and_jacobian(self._full_vector(p_start)) + J = self._fold(np.asarray(jac, dtype=np.float64))[:, self._fit_idx] + scale = np.max(np.abs(J)) or 1.0 + dead = [ + n + for i, n in enumerate(self._param_order) + if np.max(np.abs(J[:, i])) <= 1e-12 * scale + ] + if dead: + raise ValueError( + f"SCETlibADParamModel: {len(dead)} fitted parameter(s) have an " + f"identically zero derivative at the start point and would make " + f"the covariance singular: {', '.join(dead)}. Drop them from " + f"fit_params." + ) + + def _check_double_counting(self): + """Refuse a card that still carries the templates our parameters replace.""" + systs = getattr(self.indata, "systs", None) + if systs is None or len(systs) == 0: + return + names = [s.decode() if isinstance(s, bytes) else str(s) for s in systs] + lowered = [(s, s.lower()) for s in names] + for pattern, what, applies in _CONFLICTS: + if not applies(self._param_order): + continue + rx = re.compile(pattern) + clash = [s for s, low in lowered if rx.search(low)] + if clash: + raise ValueError( + f"[SCETlibADParamModel] {len(clash)} card syst(s) matching " + f"{pattern!r} describe the same physics as the fitted " + f"{what}; running both double-counts. Remake the datacard " + f"with setupRabbit --excludeNuisances '{pattern}' " + f"(case as in the card). Conflicting systs:\n" + + "\n".join(f" {s}" for s in clash[:20]) + ) + + # ========================================================================= + # evaluation + # ========================================================================= + + def _physical(self, theta): + """Fit values -> the PHYSICAL values SCETlib expects (numpy). + + Identity for everything except the reparametrised profile scales; see + params.REPARAM for why those are unit nuisances. + """ + t = np.asarray(theta, dtype=np.float64) + return ( + np.where(self._rp_id, t, 0.0) + + np.where(self._rp_log, np.exp(t * self._rp_L), 0.0) + + np.where( + self._rp_quad, + self._rp_c[0] + self._rp_c[1] * t + self._rp_c[2] * t * t, + 0.0, + ) + ) + + def _physical_tf(self, theta): + """:meth:`_physical` in TensorFlow, so the map is differentiated too. + + The chain rule does the rest: TF differentiates the map, SCETlib supplies + d(sigma)/d(physical), and the composite gradient and Hessian stay exact. + """ + t = tf.cast(theta, DTYPE) + c = tf.constant(self._rp_c, dtype=DTYPE) + return ( + tf.constant(self._rp_id.astype(np.float64), dtype=DTYPE) * t + + tf.constant(self._rp_log.astype(np.float64), dtype=DTYPE) + * tf.exp(t * tf.constant(self._rp_L, dtype=DTYPE)) + + tf.constant(self._rp_quad.astype(np.float64), dtype=DTYPE) + * (c[0] + c[1] * t + c[2] * t * t) + ) + + def _full_vector(self, fit_values): + """Fitted values -> the complete SCETlib parameter vector.""" + p = self._p_base.copy() + p[self._fit_idx] = self._physical(fit_values) + return p + + def _sigma_gen_np(self, p_full): + """sigma_gen on the gen grid (flattened) at a full SCETlib vector.""" + vals, _ = self.core.values_and_jacobian(p_full) + return self._fold(np.asarray(vals, dtype=np.float64)) + + def _sigma_gen(self, param): + """sigma_gen on the gen grid, differentiable via the SCETlib bridge. + + ``ScetlibCachedXsecTF.__call__`` is an ordinary TF-differentiable function + -- its backward pass is itself a ``custom_gradient`` whose own gradient + contracts Hessian-vector products -- so nested tapes work and TF drives + every C++ call. Nothing here is a surrogate; autodiff sees the real thing. + """ + p = self._physical_tf(param) + # held = the non-fitted slots at their anchor, zero where we fit, so + # held + S.p reconstructs the full vector. See _select on why this is a + # matmul and not a scatter. + held = self._p_base.copy() + held[self._fit_idx] = 0.0 + p_full = tf.constant(held, dtype=DTYPE) + tf.linalg.matvec( + tf.constant(self._select, dtype=DTYPE), p + ) + return self._fold.fold_tf(self.core.tf_fn(p_full)) + + def _ratio_from_param(self, param): + """Per-fit-bin ratio to the anchor prediction, softly floored positive.""" + sigma_gen = self._sigma_gen(param) + if self.gen_level: + ratio = sigma_gen / self.sigma_gen_central_flat + else: + ratio = tf.linalg.matvec(self.R, sigma_gen) / self.sigma_reco_central + # The rules put no wall in front of pathological parameters: a bad point + # can drive sigma negative, which is a NaN Poisson NLL and a dead + # gradient. Soft-floor so such a point is a large-but-finite penalty with + # a usable gradient. The scale is far below any physical response, so the + # anchor closure is untouched. + scale = tf.constant(RATIO_FLOOR_SCALE, dtype=ratio.dtype) + return tf.maximum( + scale * tf.math.softplus(ratio / scale), + tf.constant(RATIO_FLOOR_MIN, dtype=ratio.dtype), + ) + + def compute(self, param, full=False): + """(N_bins, N_proc) multiplicative scaling; only the signal column moves.""" + ratio = self._ratio_from_param(param) + ratio_col = tf.cast(tf.reshape(ratio, [-1, 1]), self.indata.dtype) + rnorm = 1.0 + (ratio_col - 1.0) * self._signal_col_mask + + n_masked = int(getattr(self.indata, "nbinsmasked", 0) or 0) + if full and n_masked: + # Masked channels sit after the fit bins in the full tensor and carry + # no model prediction, so they scale by 1. + rnorm = tf.concat( + [rnorm, tf.ones([n_masked, self.nproc], dtype=rnorm.dtype)], axis=0 + ) + return rnorm + + # ========================================================================= + # introspection (validation scripts) + # ========================================================================= + + def sigma_gen_at(self, **overrides): + """sigma_gen on the gen grid at the anchor with named overrides applied. + + ``model.sigma_gen_at(lambda2_nu=0.12)`` -- used by the validation + scripts to compare against a native SCETlib run. + """ + p = self._p_base_anchor.copy() + for name, val in overrides.items(): + if name not in self.rabbit_names: + raise KeyError(f"sigma_gen_at: unknown parameter {name!r}") + p[self.rabbit_names.index(name)] = float(val) + return self._sigma_gen_np(p).reshape(self.gen_shape) + + def __deepcopy__(self, memo): + """Copy the model but share the (immutable, expensive) SCETlib backend.""" + cls = self.__class__ + new = cls.__new__(cls) + memo[id(self)] = new + for k, v in self.__dict__.items(): + if k == "core": + new.core = v + else: + new.__dict__[k] = copy.deepcopy(v, memo) + return new + + +def _parse_kv(spec): + """``"a=1,b=2"`` (or a Mapping, or None) -> ``{"a": 1.0, "b": 2.0}``.""" + if not spec: + return {} + if not isinstance(spec, str): + return {str(k): float(v) for k, v in dict(spec).items()} + out = {} + for item in spec.split(","): + item = item.strip() + if not item: + continue + if "=" not in item: + raise ValueError(f"expected name=value, got {item!r}") + k, v = item.split("=", 1) + out[k.strip()] = float(v) + return out diff --git a/wremnants/postprocessing/scetlib_ad/params.py b/wremnants/postprocessing/scetlib_ad/params.py new file mode 100644 index 000000000..802c52830 --- /dev/null +++ b/wremnants/postprocessing/scetlib_ad/params.py @@ -0,0 +1,313 @@ +"""Parameter registry for the SCETlib autodiff param model. + +SCETlib owns the parameter vector: :meth:`DrellYan.gradient_param_names` returns +the differentiable parameters in a FIXED order that is baked into the cached bin +rules (the rule fingerprint hashes the names in order, so any addition or +reordering invalidates a cache). This module is the translation layer between +that vector and the names rabbit sees. + +Rabbit-facing names use the spelling the analysis tooling already reads +(``lambda2``, ``lambda2_nu``, …) so the postfit readers, the cross-run +fit-summary tools and the impact-group labels work unchanged. ``alphas`` becomes +``alphaS`` in PHYSICAL units (0.118-ish), NOT the ``pdfAlphaS`` template's +Delta(alpha_s) = 0.002-per-theta convention. +""" + +import math + +# --- SCETlib gradient name -> rabbit-facing name ----------------------------- +# +# Every entry here is an exact-match rename. Names not listed fall through +# :func:`rabbit_name`, which handles the two open-ended families (TNPs and, in +# a later phase, PDF eigenvector coefficients). +EXPLICIT_NAMES = { + "alphas": "alphaS", + "np_eff_lambda_inf": "lambda_inf", + "np_eff_lambda2": "lambda2", + "np_eff_lambda4": "lambda4", + "np_eff_lambda6": "lambda6", + "np_eff_delta_lambda2": "delta_lambda2", + "np_gnu_lambda_inf": "lambda_inf_nu", + "np_gnu_lambda2": "lambda2_nu", + "np_gnu_lambda4": "lambda4_nu", + "np_gnu_lambda6": "lambda6_nu", + "np_gnu_b0_bmax": "b0_over_bmax_nu", + # Profile scales and matching transition points, registered by + # set_diff_scales(1). scale_kappa_F is inert in the kernel -- the slot exists + # only for build_pdf_variations to tie the muF member pair to -- so it does + # nothing unless the cache was built with has_muf. + "scale_kappa_R": "resumScaleMuR", + "scale_kappa_F": "resumScaleMuF", + "scale_x1": "resumTransition1", + "scale_x2": "resumTransition2", + "scale_x3": "resumTransition3", +} + +# TNPs: ``tnp_gamma_cusp`` -> ``resumTNP_gamma_cusp``. The prefix matches the +# group setupRabbit gives the discrete TNP templates (``resumTNP``), so a +# grouped-impact bar stays comparable between the template and model treatments. +TNP_PREFIX_IN = "tnp_" +TNP_PREFIX_OUT = "resumTNP_" + +# PDF eigenvector coefficients, appended after the physics parameters when the +# cache carries PDF variations (phase 4). ``c_e`` are standard N(0,1) Hessian +# coefficients. +PDF_COEFF_FMT = "pdfEig{:d}" +PDF_PREFIX_IN = "pdf_eig" +PDF_PREFIX_OUT = "pdfEig" + + +def rabbit_name(scetlib_name): + """SCETlib gradient-parameter name -> the name rabbit reports.""" + if scetlib_name in EXPLICIT_NAMES: + return EXPLICIT_NAMES[scetlib_name] + if scetlib_name.startswith(TNP_PREFIX_IN): + return TNP_PREFIX_OUT + scetlib_name[len(TNP_PREFIX_IN) :] + if scetlib_name.startswith(PDF_PREFIX_IN): + return PDF_PREFIX_OUT + scetlib_name[len(PDF_PREFIX_IN) :] + raise KeyError( + f"scetlib_ad.params: no rabbit name for SCETlib parameter " + f"{scetlib_name!r}. Add it to EXPLICIT_NAMES (and give it a prior in " + f"PRIOR_SIGMAS / a group in IMPACT_GROUP_MEMBERS)." + ) + + +def scetlib_name(rabbit): + """Inverse of :func:`rabbit_name` (raises on an unknown name).""" + for k, v in EXPLICIT_NAMES.items(): + if v == rabbit: + return k + if rabbit.startswith(TNP_PREFIX_OUT): + return TNP_PREFIX_IN + rabbit[len(TNP_PREFIX_OUT) :] + if rabbit.startswith(PDF_PREFIX_OUT): + return PDF_PREFIX_IN + rabbit[len(PDF_PREFIX_OUT) :] + raise KeyError(f"scetlib_ad.params: unknown rabbit parameter {rabbit!r}") + + +# --- Priors ------------------------------------------------------------------ +# +# Only consulted when the model is constructed with ``priors=1``; otherwise every +# parameter floats free. ``None`` means "free even when priors are on". +# +# The lambda sigmas are the ones the analysis uses for the corresponding template +# nuisances, so the nonperturbative sector is constrained the same way whether it +# is fitted continuously or morphed. TNPs are genuine N(0,1) nuisances -- +# theta is normalised upstream so |theta|=1 IS the recommended variation +# (prod/scetlib_run/examples/theory_nuisance_parameters/*.conf) -- and get +# sigma = 1 by default, unlike the free lambdas. +PRIOR_SIGMAS = { + "alphaS": None, + "lambda2": 0.50, + "lambda4": 0.50, + "lambda6": 0.10, + "delta_lambda2": 0.50, + "lambda_inf": None, + "lambda2_nu": 0.10, + "lambda4_nu": 0.50, + "lambda6_nu": 0.10, + "lambda_inf_nu": None, + "b0_over_bmax_nu": None, + # --- profile scales and transition points ------------------------------- + # These REPLACE card nuisances, so their priors have to reproduce the + # variation the card encodes, and two of the three only do so approximately. + # + # resumScaleMuF: exact. The muF pair is built at kappa_F = 0.5 / 2.0 and + # interpolated in t = ln(kappa_F)/ln(muf_hi), so the direction is already + # unit-normalised and |theta| = 1 IS the card's variation. + # resumScaleMuR: APPROXIMATE. The parameter is kappa_R itself, central 1 and + # linear, while the card varies kappaFO by x2 and /2. sigma = 0.5 gives + # +-1 sigma = [0.5, 1.5], so the up-variation is understated ([0.5, 2.0]). + # A log-parametrised kappa_R would fix this properly; until then this is a + # deliberate approximation, not an equivalence. + # resumTransition2: APPROXIMATE. Central 0.6, the card's variations are + # 0.35 and 0.75 (variations_resummed.conf [35]/[36]), i.e. -0.25/+0.15; + # 0.2 is the symmetric stand-in. + # resumTransition1/3 are FROZEN by default -- the analysis varies only the + # CENTRAL transition point ("new recommendation from Frank for variation of + # central transition parameter only"), so floating the outer two would ADD + # uncertainty the card does not carry. + # resumScaleMuR / MuF / Transition2 are reparametrised unit nuisances -- + # see REPARAM below; prior_sigma() returns 1.0 for them. + "resumTransition1": None, + "resumTransition3": None, +} +TNP_PRIOR_SIGMA = 1.0 + + +def prior_sigma(rabbit): + """Default Gaussian prior sigma for a rabbit-facing name (None = free).""" + if rabbit in REPARAM: + # Unit nuisance by construction: the map carries the physical range, so + # theta = +-1 IS the variation the replaced template encoded. + return 1.0 + if rabbit.startswith(TNP_PREFIX_OUT): + return TNP_PRIOR_SIGMA + if rabbit.startswith(PDF_PREFIX_OUT): + # Hessian eigenvector coefficients: c = +-1 IS the member 1 sigma, by + # construction of build_pdf_variations (exact at c = 0, +-1). + return 1.0 + return PRIOR_SIGMAS.get(rabbit, None) + + +# --- Reparametrisation: unit nuisances for the profile scales ----------------- +# +# SCETlib registers the profile scales as the PHYSICAL quantities: +# ``scale_kappa_R`` and ``scale_kappa_F`` are kappa itself with central 1, and +# ``scale_x1..x3`` are the transition points themselves. That is the right +# interface for a calculation, but it is the wrong one for a nuisance, because +# the template variations these REPLACE are not symmetric in the physical +# variable: +# +# kappaFO x2 and /2 -> symmetric in ln(kappa), not in kappa +# x2 0.6 -> 0.35, 0.75 -> genuinely asymmetric, -0.25 / +0.15 +# +# and rabbit's ParamModel priors are a single symmetric Gaussian per parameter +# (fitter.py: cw = 1/sigma^2, one scalar, no up/down hook). Tuning sigma cannot +# reproduce either variation: sigma = 0.5 on a linear kappa_R gives [0.5, 1.5], +# understating the up side. +# +# So the fitted parameter is a UNIT nuisance theta and the model maps it to the +# physical value, exactly as SCETlib itself does for the PDF eigenvectors and +# the muF pair ("exact at 0, +-1, quadratic in between"). Every replaced-template +# direction is then sigma = 1, the same convention as the TNPs and pdfEig*. +# +# "log" : value = exp(theta * L) theta = +-1 -> exp(+-L) +# "quad" : value = c0 + c1*theta + c2*theta^2 +# +# The log form has a second benefit: exp() is positive by construction, so it +# cannot trip SCETlib's silent `p[_muf_index] > 0. ? ... : 1.` fallback, which +# would drop the muF variation with no error at all. +LN2 = math.log(2.0) +REPARAM = { + # kappa_R: theta = +-1 -> kappa_R = 2 / 0.5, matching kappaFO x2 and /2. + "resumScaleMuR": ("log", (LN2,)), + # kappa_F: same. SCETlib converts internally to t = ln(kappa_F)/ln(2), so + # theta = +-1 lands exactly on the two members that were built (0.5, 2.0). + "resumScaleMuF": ("log", (LN2,)), + # x2: the quadratic through the three points the analysis actually uses -- + # theta = -1 -> 0.35, theta = 0 -> 0.6, theta = +1 -> 0.75. Monotone for + # |theta| < 2 (the derivative 0.20 - 0.10*theta vanishes at theta = 2). + "resumTransition2": ("quad", (0.6, 0.20, -0.05)), + # resumTransition1/3 are deliberately NOT reparametrised: they are frozen by + # default and no reference variation exists for them, so a study that floats + # them should do so in the physical variable and choose its own range. +} + + +def reparam(rabbit): + """``(kind, coeffs)`` for a reparametrised name, else ``None``.""" + return REPARAM.get(rabbit) + + +# --- Defaults ---------------------------------------------------------------- +# +# Parameters frozen unless the user asks otherwise. These are shape constants of +# the SCETlib nonperturbative forms, not physics we fit: lambda_inf sets the +# saturation of the tanh forms and b0_over_bmax_nu the b* convention. +DEFAULT_FROZEN = ( + "lambda_inf", + "lambda_inf_nu", + "b0_over_bmax_nu", + # Only the CENTRAL matching transition point is varied in the analysis; see + # the PRIOR_SIGMAS comment. Float these two only as a deliberate study. + "resumTransition1", + "resumTransition3", + # resumTransition2 is frozen NOT as a physics choice but because the + # derivative SCETlib hands us for it is wrong. Making the identical physical + # change through the runcard with autodiff OFF reproduces the production + # template to 2e-6; making it through the registered scale_x2 parameter with + # set_diff_scales(1) comes out with the OPPOSITE SIGN (0.966985 where the + # template says 1.159163 at qT [33,44]), because moving the transition points + # moves muF ~20% while the per-node beam convolutions stay frozen at the + # config's muF. See studies/scetlib-ad-param-model/ for the elimination chain + # and the upstream issue. Freezing it removes the transition-point + # uncertainty from the fit -- that is a KNOWN GAP, not a fix; it is preferable + # only to profiling against a wrong-sign response. Unfreeze when upstream + # lands a per-node muF, and re-run validate_variations.py before trusting it. + "resumTransition2", +) + +# Directions whose response has been MEASURED to disagree with the template it +# replaces, keyed by rabbit-facing name -> why. These are not frozen for physics +# reasons, so anyone who floats one anyway deserves to be told once, loudly, +# rather than to find out from a pull. Keep the strings short; the detail lives in +# studies/scetlib-ad-param-model/. +KNOWN_BAD_RESPONSE = { + "resumTransition1": "sign-inverted vs the template (upstream: per-node muF)", + "resumTransition2": "sign-inverted vs the template (upstream: per-node muF)", + "resumTransition3": "sign-inverted vs the template (upstream: per-node muF)", +} + + +# Grouped impacts over the model's own parameters (rabbit resolves these labels +# to floating x-indices; see Fitter._resolved_param_impact_groups). Membership is +# intersected with the parameters actually registered. +IMPACT_GROUP_MEMBERS = { + "resumNonpert": ( + "lambda2", + "lambda4", + "lambda6", + "delta_lambda2", + "lambda_inf", + "lambda2_nu", + "lambda4_nu", + "lambda6_nu", + "lambda_inf_nu", + "b0_over_bmax_nu", + ), + "scetlibNPFeff": ( + "lambda2", + "lambda4", + "lambda6", + "delta_lambda2", + "lambda_inf", + ), + "scetlibNPgammaNu": ( + "lambda2_nu", + "lambda4_nu", + "lambda6_nu", + "lambda_inf_nu", + "b0_over_bmax_nu", + ), + # Named to line up with the card groups they replace, so a grouped-impact + # bar stays comparable between the template and model treatments. + "resumScale": ("resumScaleMuR", "resumScaleMuF"), + "resumTransition": ( + "resumTransition1", + "resumTransition2", + "resumTransition3", + ), +} + + +def tnp_group(names): + """The ``resumTNP`` impact group for whichever TNPs are registered.""" + return tuple(n for n in names if n.startswith(TNP_PREFIX_OUT)) + + +def pdf_group(names): + """The ``pdf`` impact group for whichever eigenvector coefficients exist.""" + return tuple(n for n in names if n.startswith(PDF_PREFIX_OUT)) + + +# --- lambda_central cross-check ---------------------------------------------- +# +# A histmaker output records the nonperturbative values its theory correction was +# generated at, under two sub-dicts using the histmaker's own spelling, and that +# is propagated into the datacard. Map those names onto rabbit-facing ones so the +# card's anchor can be compared against the cache's. +LAMBDA_CENTRAL_KEYS = { + # card metadata key -> rabbit-facing name (identical here, but + # spelled out so a future divergence is a one-line fix rather than a silent + # mismatch) + "lambda2": "lambda2", + "lambda4": "lambda4", + "lambda6": "lambda6", + "delta_lambda2": "delta_lambda2", + "lambda_inf": "lambda_inf", + "lambda2_nu": "lambda2_nu", + "lambda4_nu": "lambda4_nu", + "lambda6_nu": "lambda6_nu", + "lambda_inf_nu": "lambda_inf_nu", + "b0_over_bmax_nu": "b0_over_bmax_nu", +} diff --git a/wremnants/postprocessing/scetlib_ad/response.py b/wremnants/postprocessing/scetlib_ad/response.py new file mode 100644 index 000000000..d7b753821 --- /dev/null +++ b/wremnants/postprocessing/scetlib_ad/response.py @@ -0,0 +1,208 @@ +"""Response matrix: folding the boson-level prediction to reco bins. + +The differentiable prediction from SCETlib is a boson-level cross section on a +(Q, Y, qT) grid. Turning it into expected reco yields needs the detector, and +that comes from the MC through a response matrix. Two histograms out of the +unfolding histmaker: + +* ``R_raw(b, g)`` -- MC weight of events *generated* in gen bin ``g`` and + *reconstructed* in reco bin ``b``. Its reco axes carry everything the boson + calculation does not: the decay angles, smearing, efficiency, acceptance. +* ``N_gen(g)`` -- MC weight of ALL events generated in gen bin ``g``, with no + reco selection. The gen-side marginal of the same sample. + +The model uses the ratio + + P(b|g) = R_raw(b, g) / N_gen(g) + +which is a conditional probability -- "given generated in ``g``, land in ``b``" +-- i.e. pure efficiency times migration, carrying no information about how much +cross section there is. That last part is the point: ``R_raw`` alone has the MC's +own gen spectrum baked into it, so folding a theory prediction through ``R_raw`` +would count the gen spectrum twice. Dividing by ``N_gen`` removes it. + + sigma_reco(b) = sum_g P(b|g) . sigma(p; g) + +A second consequence of the same division: any gen-level reweighting of the MC +(for instance a theory correction applied at histmaker time) multiplies ``R_raw`` +and ``N_gen`` by the same factor and cancels, so ``P`` is unchanged by it. That +holds wherever the reweighting is constant within a gen bin, which is why the gen +binning should resolve whatever correction the MC carries. +""" + +import numpy as np + +# Name of the auxiliary group the datacard carries the response in. A DATACARD +# convention fixed by whichever setupRabbit run embedded it -- not a property of +# this package -- so it is overridable, and the historical name is the default so +# existing cards keep working. +DEFAULT_RESPONSE_GROUP = "scetlib_np" + +# A pathological parameter point can drive the folded cross section -- and so the +# predicted yield -- negative, which gives a NaN Poisson NLL and a dead gradient +# that strands the minimiser. Soft-floor the ratio (softplus, not a hard clamp) so +# such a point is a large-but-finite penalty with a usable gradient. +# RATIO_FLOOR_SCALE -- softplus transition width, far below any physical +# response, so healthy ratios pass through to machine precision: +# scale * softplus(r / scale) == r for r >> scale. +# RATIO_FLOOR_MIN -- hard positive ground, since softplus underflows to exactly +# zero for extreme negatives; keeps the yield strictly > 0. +RATIO_FLOOR_SCALE = 1.0e-4 +RATIO_FLOOR_MIN = 1.0e-9 + + +def crop_R_to_fit(R, R_reco_axes, fit_reco_axes, tol=1e-9): + """Crop R's reco bins so its reco shape matches the fit channel's. + + R is usually stored on a superset of the fit's reco binning. The fit's edges + must appear as a *contiguous sub-range* of R's, but need not start at R's + first edge: a low-side acceptance cut (``ptll > 5``, an ``--axlim`` dropping + leading bins) makes the fit an interior slice of an R stored from zero. For + each reco axis, find the offset where R's edges line up and crop to + ``[offset, offset + nbins)``. + """ + if len(R_reco_axes) != len(fit_reco_axes): + raise ValueError( + f"Reco axis count mismatch: R has {len(R_reco_axes)}, " + f"fit has {len(fit_reco_axes)}" + ) + axis_slices = [] + for (rname, redges), (fname, fedges) in zip(R_reco_axes, fit_reco_axes): + if rname != fname: + raise ValueError(f"Reco axis name mismatch: R={rname!r} vs fit={fname!r}") + redges = np.asarray(redges) + fedges = np.asarray(fedges) + fnb = len(fedges) + if len(redges) < fnb: + raise ValueError( + f"Reco axis {rname}: R has {len(redges) - 1} bins, fit needs " + f"{fnb - 1}. R is missing edges." + ) + offset = next( + ( + k + for k in range(len(redges) - fnb + 1) + if np.allclose(redges[k : k + fnb], fedges, atol=tol) + ), + None, + ) + if offset is None: + raise ValueError( + f"Reco axis {rname}: fit edges are not a contiguous sub-range of " + f"R's edges. R={list(redges)} vs fit={list(fedges)}" + ) + axis_slices.append(slice(offset, offset + fnb - 1)) + slices = tuple(axis_slices) + # keep every gen axis (the trailing axes of R) + slices += (slice(None),) * (R.ndim - len(fit_reco_axes)) + return R[slices] + + +def marginalize_R_reco(R, R_reco_axes, fit_axis_names, log_prefix="scetlib_ad"): + """Sum R over the reco axes the fit channel does not have. + + The datacard embeds R at the full reco binning it was produced with, + whatever the fit channel's dimensionality. R is a counts response, so the + response for a lower-dimensional channel (a 1D ptll or 2D ptll-yll fit) is + exactly the marginal over the dropped axes. Gen axes are untouched. The kept + axes must appear in the fit's order. + """ + R_names = [n for n, _ in R_reco_axes] + missing = [n for n in fit_axis_names if n not in R_names] + if missing: + raise ValueError(f"Fit reco axes {missing} not among R's reco axes {R_names}") + kept = [n for n in R_names if n in fit_axis_names] + if kept != list(fit_axis_names): + raise ValueError( + f"Fit reco-axis order {list(fit_axis_names)} doesn't match R's " + f"stored order {kept}" + ) + drop = tuple(i for i, n in enumerate(R_names) if n not in fit_axis_names) + if drop: + R = R.sum(axis=drop) + print( + f"[{log_prefix}] marginalized R over reco axes " + f"{[R_names[i] for i in drop]} (fit channel is " + f"{len(fit_axis_names)}D: {list(fit_axis_names)})", + flush=True, + ) + reco_axes = [(n, e) for n, e in R_reco_axes if n in fit_axis_names] + return R, reco_axes + + +def R_info_from_auxiliary(indata, group=DEFAULT_RESPONSE_GROUP): + """Read the response bundle out of the datacard's auxiliary group. + + setupRabbit extracts R -- plus the gen total ``N_gen`` and the reco/gen axis + names and edges -- once from the unfolding histmaker output and embeds it via + rabbit's ``add_auxiliary``; rabbit exposes it as ``FitInputData.auxiliary``. + Reading it only from there means R is always the one consistent with the run + that produced this datacard, with no fit-time file path to get wrong. + + Returns ``R``, ``N_gen``, and ``reco_axes`` / ``gen_axes`` as ordered + ``(name, edges)`` lists. + """ + aux = getattr(indata, "auxiliary", None) or {} + if group not in aux: + raise ValueError( + f"scetlib_ad: the datacard has no {group!r} auxiliary (the reco x gen " + f"response matrix R). Rebuild it with a setupRabbit that embeds the " + f"response from a mz_dilepton --unfolding input carrying " + f"'nominal_prefsr_yieldsUnfolding' and the 'prefsr' gen total, or pass " + f"response_group= if the card names it differently." + ) + bundle = aux[group] + n_gen = bundle.get("N_gen") + return dict( + R=np.asarray(bundle["R"], dtype=np.float64), + N_gen=None if n_gen is None else np.asarray(n_gen, dtype=np.float64), + reco_axes=[ + (name, np.asarray(bundle[f"edges__{name}"], dtype=np.float64)) + for name in bundle["reco_axes"] + ], + gen_axes=[ + (name, np.asarray(bundle[f"edges__{name}"], dtype=np.float64)) + for name in bundle["gen_axes"] + ], + ) + + +# --- the card's recorded NP anchor ------------------------------------------- +# +# Histmaker outputs record the nonperturbative values their theory correction was +# generated at, and that key is propagated into the datacard. It is worth +# cross-checking against the cache anchor: a mismatch leaves the ratio equal to 1 +# at the start, so nothing looks wrong, while the response is being evaluated at +# the wrong point. +# Likewise a datacard/histmaker convention, not ours. +NP_ANCHOR_META_KEY = "scetlib_np_lambda_central" + + +def np_anchor_from_meta(meta, proc="Z", max_depth=8): + """The card's recorded NP anchor as a flat ``{name: value}`` dict, or None. + + rabbit nests the histmaker's ``meta_info`` under ``meta_info_input``, and + again in a fitresult, so walk that chain. Returns None when the key is absent + rather than raising -- the check is a guard, not a requirement. + """ + cur = meta + for _ in range(max_depth): + if not isinstance(cur, dict): + return None + entry = cur.get(NP_ANCHOR_META_KEY) + if isinstance(entry, dict) and entry: + per_proc = entry.get(proc) + if per_proc is None and len(entry) == 1: + per_proc = next(iter(entry.values())) + if isinstance(per_proc, dict): + out = {} + for sub in ("eff_params", "gnu_params"): + for k, v in dict(per_proc.get(sub, {})).items(): + if isinstance(v, (int, float)): + out[k] = float(v) + return out or None + nxt = cur.get("meta_info_input") + if not isinstance(nxt, dict) or nxt is cur: + return None + cur = nxt + return None diff --git a/wremnants/postprocessing/scetlib_ad/xsec_backend.py b/wremnants/postprocessing/scetlib_ad/xsec_backend.py new file mode 100644 index 000000000..00d79755f --- /dev/null +++ b/wremnants/postprocessing/scetlib_ad/xsec_backend.py @@ -0,0 +1,445 @@ +"""SCETlib autodiff backend: cached matched sigma_UL with exact derivatives. + +Thin, numpy-facing wrapper around ``ScetlibCachedXsecTF`` +(``scetlib-cms/py/scetlib_tf.py``). It owns three things the param model should +not have to know about: + +* rebuilding the SCETlib calculation exactly as the cache was prepared with + (``configure_calculation`` is NOT the whole configuration -- the profile-scale + / b* block and the nonperturbative models go through + ``variations.set_vary``, so skipping it + silently keeps the C++ defaults); +* loading the one-file cache (compressed bin rules for the resummed piece + + frozen fixed-order grid for the nonsingular one) and checking it belongs to + this configuration; +* mapping the cache's bin list onto the fit's gen grid, by an explicit + permutation rather than by assuming both were generated in the same order. + +Derivatives come from C++: ``values_and_jacobian`` returns value and Jacobian +from one call, ``hessian`` the exact per-bin Hessian. The param model injects +those into the TF graph as a local quadratic, so nothing here needs to be +differentiable. + +Runtime requirements (inside the WRemnants singularity): +``source /setup.sh`` -- it puts ``scetlib_qT`` and ``scetlib_run`` +on ``PYTHONPATH``, the shared libraries on ``LD_LIBRARY_PATH``, and lifts the +stack limit (the bT integrators are deep). LHAPDF and the per-member beamfunc +grids under ``share/scetlib/beamfunc/`` must be reachable. +""" + +import configparser +import os + +import numpy as np + +# Bin-edge matching tolerance. Gen edges come from hist axes (fp64 literals) and +# cache edges from a .npz round trip of the same literals, so exact equality is +# the expectation and this only absorbs formatting round trips. +EDGE_TOL = 1e-9 + + +def _import_scetlib(): + """Import the SCETlib python modules, with an actionable error if absent.""" + try: + import scetlib_qT # noqa: F401 + from scetlib_run import config as sl_config + from scetlib_run import variations as sl_variations + from scetlib_tf import ScetlibCachedXsecTF + except ImportError as e: + raise ImportError( + "scetlib_ad needs the SCETlib autodiff build on PYTHONPATH. Run\n" + " source /scetlib-cms/setup.sh\n" + "inside the container before rabbit_fit.py (it also sets " + "LD_LIBRARY_PATH and lifts the stack limit).\n" + f"Original error: {e}" + ) from e + return sl_config, sl_variations, ScetlibCachedXsecTF + + +def _scetlib_src(): + """Repository root of the SCETlib checkout (for prod/scetlib_run/defaults.conf).""" + src = os.environ.get("SCETLIB_SRC") + if src: + return src + # setup.sh puts /prod/scetlib_run on PYTHONPATH, so the package sits at + # /prod/scetlib_run/scetlib_run/. + import scetlib_run + + return os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(scetlib_run.__file__))) + ) + ) + + +def configure(config_path, threads=0, diff_scales=True, fo_resolve_muR=True): + """Rebuild the calculation the cache was prepared with. + + Follows what ``prod/scetlib_run/scetlib-run-qT.py`` does, which is the + path every production correction was made with: layer the runcard on + ``prod/scetlib_run/defaults.conf``, configure the calculation, apply the + card's electroweak parameters and fiducial volumes, then apply variation 0 + (central) through ``set_vary`` -- which is what installs the profile-scale + / b* block and the nonperturbative models. Finally enable the frozen-node + cache on both sub-pieces, without which the rule replay has nothing to + replay against. + + ``configure_ew_parameters`` and ``configure_fiducial_volumes`` are NOT + optional and are NOT part of ``configure_calculation``. Skipping them + silently keeps SCETlib's default electroweak inputs while the card + specifies its own -- measured as a flat **1.61%** normalization error + against the production driver on the analysis card, which sets + ``mZ = 91.1535``, ``GammaZ = 2.4932``, a custom ``alphaem``, ``sin2_thw`` + and CKM. With them applied, ``operator()`` here reproduces + ``scetlib-run-qT.py`` to 1e-9. ``examples/matched_ad/prepare_cache.py`` + upstream still omits both; do not "simplify" this back to match it. + + ``diff_scales`` registers muR and the three matching transition points as + differentiable parameters (``scale_kappa_R``, ``scale_x1..x3``, plus an inert + ``scale_kappa_F`` slot that ``build_pdf_variations`` ties the muF pair to). + It REQUIRES ``muf_follows_muB = no``: with muf tied to muB a live kappa_R + moves muF while the beam convolutions stay frozen at their own muF. It also + changes the parameter registry, so a cache built without it cannot be loaded + with it -- :class:`ScetlibADXsec` therefore reads the flag off the cache + rather than assuming. + + ``fo_resolve_muR`` resolves the fixed-order muR dependence into the frozen + grid so the FO piece follows kappa_R in closed form. It must be set BEFORE + the grid is built -- enabling it afterwards invalidates the cache -- and it + samples several scales, so the warm costs roughly 3x. + + Returns ``(conf, sigma)``. + """ + sl_config, sl_variations, _ = _import_scetlib() + src = _scetlib_src() + conf = configparser.ConfigParser(inline_comment_prefixes="#") + conf.read(os.path.join(src, "prod", "scetlib_run", "defaults.conf")) + if not conf.read(config_path): + raise FileNotFoundError(f"scetlib_ad: cannot read runcard {config_path!r}") + + order, alphas, decay, scales, sigma = sl_config.configure_calculation(conf) + sl_config.configure_ew_parameters(conf, sigma) + sl_config.configure_fiducial_volumes(conf, decay) + if diff_scales: + follows = conf["Calculation_settings"].get("muf_follows_muB", "no") + if str(follows).strip().lower() in ("yes", "true", "1"): + raise ValueError( + "scetlib_ad: diff_scales needs muf_follows_muB = no, but the " + f"runcard sets muf_follows_muB = {follows}. With muf tied to muB " + "a live kappa_R moves muF while the beam convolutions stay " + "frozen at their own muF." + ) + sigma.set_diff_scales(1) + if fo_resolve_muR: + for piece in sigma.sub_pieces(): + piece.set_fo_resolve_muR(True) + varis = sl_variations.configure_variations( + conf, + os.path.join(os.path.dirname(os.path.abspath(config_path)), "variations.conf"), + ) + sl_variations.set_vary(varis[0], order, alphas, scales, sigma) + + nthreads = int(threads) if threads else (os.cpu_count() or 8) + for piece in sigma.sub_pieces(): + piece.set_gradient_threads(nthreads) + piece.set_gradient_node_cache(True) + return conf, sigma + + +def bins_from_gen_axes(gen_axes, Q_lo, Q_hi): + """(N, 6) SCETlib bins for a (qT, |Y|) gen grid, in the grid's flatten order. + + ``gen_axes`` is ``[(qT_name, edges), (Y_name, edges)]`` -- the order the param + model flattens in, qT-major. The rapidity edges are used as given: pass the + positive side only to build a folded-|Y| cache (correct for the Z, where the + factor 2 cancels in the ratio), or signed edges to build a signed one (needed + for W, and for a bin-by-bin comparison against a SCETlib reference run on its + own signed grid). :class:`GenFold` accepts either. + """ + qT_edges = np.asarray(gen_axes[0][1], dtype=np.float64) + absY_edges = np.asarray(gen_axes[1][1], dtype=np.float64) + return np.array( + [ + [Q_lo, Q_hi, absY_edges[j], absY_edges[j + 1], qT_edges[i], qT_edges[i + 1]] + for i in range(qT_edges.size - 1) + for j in range(absY_edges.size - 1) + ], + dtype=np.float64, + ) + + +def _containing_bin(lo, hi, edges, what): + """Index of the ``edges`` bin that [lo, hi] falls inside, or None.""" + i = int(np.searchsorted(edges, lo + EDGE_TOL) - 1) + if i < 0 or i >= edges.size - 1: + return None + if hi > edges[i + 1] + EDGE_TOL: + raise ValueError( + f"scetlib_ad: cache {what} bin [{lo:g}, {hi:g}] straddles the gen " + f"edge {edges[i + 1]:g}. The cache binning must nest inside the " + f"fit's gen binning (a gen bin may be the sum of several cache bins, " + f"but a cache bin may not span two gen bins)." + ) + return i + + +class GenFold: + """Exact sum of cache bins onto the fit's gen grid. + + A plain reindexing is not enough in general. Three things can differ between + the cache and the fit's gen binning, and all three are handled here by + SUMMING bin-integrated cross sections, which is exact: + + * **order** -- ``prepare_cache`` nests (Q, Y, qT) with qT innermost, while + the gen grid (and the response matrix's gen columns) flatten qT-major; + * **rapidity sign** -- a production cache is built on the SIGNED Y grid of + the theory correction, so a gen |Y| bin is the sum of the +Y and -Y cache + bins. A cache built on the positive side only is also accepted, and then + every gen bin is short by the same factor 2 -- which cancels in the ratio + to the central, and is reported rather than silently applied; + * **granularity** -- a fine cache (e.g. the 70 x 82 correction grid) folds + exactly onto any coarser gen grid whose edges are a subset of its own. + + Coverage is checked, not assumed: every gen bin must be tiled exactly by the + cache bins assigned to it, so a cache that only partially covers a gen bin + raises instead of quietly integrating over less phase space. + """ + + def __init__(self, bins, gen_axes, Q_lo, Q_hi): + bins = np.asarray(bins, dtype=np.float64) + qT_edges = np.asarray(gen_axes[0][1], dtype=np.float64) + absY_edges = np.asarray(gen_axes[1][1], dtype=np.float64) + n_qt, n_y = qT_edges.size - 1, absY_edges.size - 1 + self.n_gen = n_qt * n_y + self.gen_shape = (n_qt, n_y) + + rows = np.full(bins.shape[0], -1, dtype=np.int64) + # (qT, |Y|) area each gen bin actually receives, to verify exact tiling. + covered = np.zeros((n_qt, n_y), dtype=np.float64) + sides = np.zeros((n_qt, n_y, 2), dtype=np.int64) # [-Y, +Y] contributions + for k, (q_lo, q_hi, y_lo, y_hi, t_lo, t_hi) in enumerate(bins): + if q_lo < Q_lo - EDGE_TOL or q_hi > Q_hi + EDGE_TOL: + raise ValueError( + f"scetlib_ad: cache Q bin [{q_lo:g}, {q_hi:g}] lies outside " + f"the fit's mass window [{Q_lo:g}, {Q_hi:g}]." + ) + if y_lo < -EDGE_TOL and y_hi > EDGE_TOL: + raise ValueError( + f"scetlib_ad: cache Y bin [{y_lo:g}, {y_hi:g}] straddles " + f"Y = 0 and cannot be folded onto |Y|." + ) + negative = y_hi <= EDGE_TOL + a, b = (abs(y_hi), abs(y_lo)) if negative else (y_lo, y_hi) + j = _containing_bin(a, b, absY_edges, "|Y|") + i = _containing_bin(t_lo, t_hi, qT_edges, "qT") + if i is None or j is None: + continue # outside the fit's gen range: dropped, accounted below + rows[k] = i * n_y + j + covered[i, j] += (t_hi - t_lo) * (b - a) + sides[i, j, 0 if negative else 1] += 1 + + self.n_used = int((rows >= 0).sum()) + self.n_dropped = int(bins.shape[0] - self.n_used) + + empty = sides.sum(axis=-1) == 0 + if empty.any(): + i, j = np.argwhere(empty)[0] + raise ValueError( + f"scetlib_ad: {int(empty.sum())} gen bin(s) receive no cache bin " + f"at all, e.g. qT [{qT_edges[i]:g}, {qT_edges[i + 1]:g}] x |Y| " + f"[{absY_edges[j]:g}, {absY_edges[j + 1]:g}]. The cache does not " + f"cover this card's gen range; rebuild it for this card." + ) + both = (sides[..., 0] > 0) & (sides[..., 1] > 0) + pos_only = (sides[..., 0] == 0) & (sides[..., 1] > 0) + if both.all(): + self.y_convention = "signed" + y_factor = 2.0 + elif pos_only.all(): + self.y_convention = "positive-side-only" + y_factor = 1.0 + else: + raise ValueError( + "scetlib_ad: the cache covers |Y| on both signed sides for some " + "gen bins and only one side for others, so the fold would apply " + "an inconsistent factor 2 across the grid. Rebuild the cache on " + "a uniform Y convention." + ) + # An exactly tiled gen bin receives y_factor x its own (qT, |Y|) area. + want = np.diff(qT_edges)[:, None] * np.diff(absY_edges)[None, :] * y_factor + bad = np.abs(covered - want) > EDGE_TOL * np.maximum(want, 1.0) + if bad.any(): + i, j = np.argwhere(bad)[0] + raise ValueError( + f"scetlib_ad: {int(bad.sum())} gen bin(s) are not exactly tiled " + f"by the cache, e.g. qT [{qT_edges[i]:g}, {qT_edges[i + 1]:g}] x " + f"|Y| [{absY_edges[j]:g}, {absY_edges[j + 1]:g}] receives " + f"{covered[i, j]:g} of {want[i, j]:g} in (qT, |Y|) area. The cache " + f"does not cover this card's gen binning; rebuild it." + ) + + # Group the surviving cache rows by destination so the fold is a single + # reduceat rather than a scatter-add (which is slow, and is on the hot + # path once per loss/gradient evaluation). + keep = np.nonzero(rows >= 0)[0] + self._order = keep[np.argsort(rows[keep], kind="stable")] + self._starts = np.searchsorted(rows[self._order], np.arange(self.n_gen)) + # Destination gen bin per cache bin, -1 for the ones outside the gen + # range. tf.math.unsorted_segment_sum drops negative ids, so this is the + # same fold expressed for the in-graph path (see fold_tf). + self.segment_ids = rows + + def __call__(self, a): + """Fold a cache-indexed array (first axis = cache bin) onto the gen grid.""" + return np.add.reduceat(np.asarray(a)[self._order], self._starts, axis=0) + + def fold_tf(self, a): + """The same fold, as TensorFlow ops, for the differentiate-through path.""" + import tensorflow as tf + + return tf.math.unsorted_segment_sum( + a, tf.constant(self.segment_ids, dtype=tf.int32), self.n_gen + ) + + def describe(self): + return ( + f"{self.n_used} cache bin(s) -> {self.n_gen} gen bin(s), " + f"Y convention {self.y_convention}" + + ( + f", {self.n_dropped} cache bin(s) outside the gen range" + if self.n_dropped + else "" + ) + ) + + +class ScetlibADXsec: + """Cached matched sigma_UL(p) on a fixed bin set, with exact derivatives. + + Parameters + ---------- + conf_path + SCETlib runcard the cache was built from (layered on defaults.conf). + cache_path + The ``.npz`` written by ``ScetlibCachedXsecTF.save`` -- rules, frozen + fixed-order grid, bins, anchor and parameter names in one file. + threads + Worker threads for the batch replay (0 = one per hardware thread). + """ + + @staticmethod + def cache_param_names(cache_path): + """The parameter names a cache was built with, without loading it. + + ``.npz`` is a zip of individual arrays, so this reads only ``names``. + Used to configure the calculation the way the cache expects instead of + the way we would prefer -- a mismatch is otherwise a hard load failure + (the fingerprint hashes the names in order). + """ + with np.load(cache_path, allow_pickle=False) as z: + if "names" not in z.files: + return None + return [str(n) for n in z["names"]] + + def __init__(self, conf_path, cache_path, threads=0): + _, _, ScetlibCachedXsecTF = _import_scetlib() + self.conf_path = os.path.abspath(conf_path) + self.cache_path = os.path.abspath(cache_path) + # Match the cache's direction set: caches built before the scale + # directions existed have no scale_* entries, and configuring them in + # would change the registry and fail the fingerprint check. + cached_names = self.cache_param_names(self.cache_path) + want_scales = cached_names is None or any( + n.startswith("scale_") for n in cached_names + ) + # fo_resolve_muR likewise has to match: it changes the frozen FO grid, + # and the docs are explicit that enabling it after the grid was built + # invalidates the cache. Both travel together with the scale directions. + self.conf, self._sigma = configure( + self.conf_path, + threads, + diff_scales=want_scales, + fo_resolve_muR=want_scales, + ) + sing, nons = self._sigma.sub_pieces() + self._fn = ScetlibCachedXsecTF.load(self.cache_path, sing, nons) + + self.param_names = list(self._fn.param_names) + self.n_params = len(self.param_names) + self.bins = np.asarray(self._fn._points, dtype=np.float64) + self.n_bins = self.bins.shape[0] + # The anchor the rules were compressed around. Accuracy is best in its + # neighbourhood and degrades gracefully away from it (the upstream guard + # that used to refuse large excursions was removed), so the fit start + # should sit here and a postfit re-check is good practice. + self.anchor = np.asarray(self._fn.anchor, dtype=np.float64) + + # ------------------------------------------------------------------ # + # bin bookkeeping # + # ------------------------------------------------------------------ # + + def fold_for(self, gen_axes, Q_lo, Q_hi): + """:class:`GenFold` summing this cache's bins onto the fit's gen grid. + + Matching is by VALUE, not by assuming the cache was generated in the + grid's order: a transposed or differently-nested cache would otherwise be + a silent wrong answer rather than an error. + """ + return GenFold(self.bins, gen_axes, float(Q_lo), float(Q_hi)) + + # ------------------------------------------------------------------ # + # evaluation # + # ------------------------------------------------------------------ # + + @property + def tf_fn(self): + """The raw ``ScetlibCachedXsecTF``: a TF-differentiable function of the + FULL parameter vector, whose gradients come from nested + ``tf.custom_gradient`` wrappers that call back into C++. Used by the + param model's ``differentiate=through`` mode.""" + return self._fn + + def values_and_jacobian(self, p): + """(values (N,), jacobian (N, P)) at parameter vector ``p``.""" + return self._fn.values_and_jacobian(np.asarray(p, dtype=np.float64)) + + def hessian(self, p): + """Exact per-bin Hessians, shape (N, P, P).""" + return np.asarray(self._fn.hessian(np.asarray(p, dtype=np.float64))) + + def resummed_only(self, p): + """The RESUMMED (singular) piece alone, without the nonsingular. + + The cached matched total is sing + nons, and the two are stored + separately -- compressed bin rules for the resummed piece, a frozen grid + for the fixed-order one. Replaying only the rules gives the resummed + piece, which is what a ``calculation_piece = sing`` SCETlib production run + computes. That makes it the apples-to-apples reference for validating our + configuration and quadrature, with no matching, no DYTurbo and no MiNNLO + in between. + """ + res = self._fn._sing.sigma_binned_rule_batch( + self.bins, np.ascontiguousarray(p, dtype=np.float64) + ) + return np.asarray(res["value"], dtype=np.float64) + + # ------------------------------------------------------------------ # + # lifecycle # + # ------------------------------------------------------------------ # + + def __deepcopy__(self, memo): + """Share the backend across a deepcopy instead of copying it. + + rabbit's ``Fitter.__deepcopy__`` copies the param model (saturated model, + toys). The pybind handles here wrap immutable configuration -- the loaded + rules and the frozen fixed-order grid -- and a real copy would mean a + second ``configure`` (LHAPDF reload, beamfunc grid re-read) for no gain. + """ + memo[id(self)] = self + return self + + def __repr__(self): + return ( + f"ScetlibADXsec(bins={self.n_bins}, params={self.n_params}, " + f"cache={os.path.basename(self.cache_path)})" + )