From e0b47e54dbc05e3ded1a043cae3fb45cc8531d7d Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 29 May 2026 16:39:01 -0400 Subject: [PATCH 01/31] add module to compute scetlib NP model as a rabbit ParamModel --- .../postprocessing/scetlib_np/__init__.py | 13 + .../postprocessing/scetlib_np/btgrid_cache.py | 85 ++ .../scetlib_np/btgrid_integrate.py | 198 ++++ .../postprocessing/scetlib_np/btgrid_numpy.py | 713 ++++++++++++++ .../postprocessing/scetlib_np/btgrid_tf.py | 318 +++++++ .../scetlib_np/lambda_central.py | 197 ++++ .../postprocessing/scetlib_np/param_model.py | 872 ++++++++++++++++++ .../scetlib_np/response_matrix.py | 128 +++ 8 files changed, 2524 insertions(+) create mode 100644 wremnants/postprocessing/scetlib_np/__init__.py create mode 100644 wremnants/postprocessing/scetlib_np/btgrid_cache.py create mode 100644 wremnants/postprocessing/scetlib_np/btgrid_integrate.py create mode 100644 wremnants/postprocessing/scetlib_np/btgrid_numpy.py create mode 100644 wremnants/postprocessing/scetlib_np/btgrid_tf.py create mode 100644 wremnants/postprocessing/scetlib_np/lambda_central.py create mode 100644 wremnants/postprocessing/scetlib_np/param_model.py create mode 100644 wremnants/postprocessing/scetlib_np/response_matrix.py diff --git a/wremnants/postprocessing/scetlib_np/__init__.py b/wremnants/postprocessing/scetlib_np/__init__.py new file mode 100644 index 000000000..439644204 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/__init__.py @@ -0,0 +1,13 @@ +"""SCETlib NP continuous-λ param model and its bT-grid factorisation port. + +The fit-time rabbit ParamModel lives in :mod:`.param_model`; the supporting +bT-grid numpy/TF code, Q-integration, caching, response-matrix loader and +λ_central reader are the sibling modules. ``SCETlibNPParamModel`` is +re-exported here so it can be referenced by the short dotted path: + + --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel +""" + +from wremnants.postprocessing.scetlib_np.param_model import SCETlibNPParamModel + +__all__ = ["SCETlibNPParamModel"] diff --git a/wremnants/postprocessing/scetlib_np/btgrid_cache.py b/wremnants/postprocessing/scetlib_np/btgrid_cache.py new file mode 100644 index 000000000..db836e42f --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/btgrid_cache.py @@ -0,0 +1,85 @@ +"""One-shot pickle cache for the combined SCETlib bT-grid. + +Loading the fineall btgrid as 1519 individual shards takes ~110s. After the +first call, this module writes a single ``combined_btgrid.pkl`` in the btgrid +directory; subsequent calls load that in a few seconds. + +Usage: + from wremnants.postprocessing.scetlib_np import btgrid_cache + grid = btgrid_cache.load(BTGRID_DIR) +""" + +import glob +import os +import pickle +import time + +from wremnants.postprocessing.scetlib_np import btgrid_numpy as fz + +_COMBINED_BASENAME = "combined_btgrid.pkl" + + +def _shard_glob(submitdir): + for pat in ( + os.path.join(submitdir, "scetlib_outputs", "*_btgrid.pkl"), + os.path.join(submitdir, "*_btgrid.pkl"), + ): + files = glob.glob(pat) + if files: + return files + return [] + + +def _combined_path(submitdir): + return os.path.join(submitdir, _COMBINED_BASENAME) + + +def _cache_is_fresh(combined, shards): + if not os.path.exists(combined): + return False + if not shards: + return True # nothing to compare against; trust the cache + mtime = os.path.getmtime(combined) + return mtime >= max(os.path.getmtime(s) for s in shards) + + +def load(submitdir, rebuild=False, verbose=True): + """Load the combined bT-grid for ``submitdir``. + + On first call (or when ``rebuild=True``, or when any shard is newer than + the cached combined file), assembles the shards via + :func:`wremnants.postprocessing.scetlib_np.btgrid_numpy.load_btgrid_shards`, + writes ``combined_btgrid.pkl``, and returns the dict. + + On subsequent calls, loads the pickle directly. + """ + if not os.path.isdir(submitdir): + raise ValueError(f"{submitdir!r} is not a directory") + + combined = _combined_path(submitdir) + shards = _shard_glob(submitdir) + + if not rebuild and _cache_is_fresh(combined, shards): + t0 = time.time() + with open(combined, "rb") as f: + grid = pickle.load(f) + if verbose: + print(f"[btgrid_cache] loaded combined pickle in {time.time()-t0:.1f}s") + return grid + + if not shards: + raise FileNotFoundError(f"No btgrid shards found under {submitdir!r}") + + t0 = time.time() + grid = fz.load_btgrid_shards(submitdir) + if verbose: + print( + f"[btgrid_cache] assembled {grid['n_shards']} shards in " + f"{time.time()-t0:.1f}s; writing {combined}" + ) + + tmp = combined + ".tmp" + with open(tmp, "wb") as f: + pickle.dump(grid, f, protocol=pickle.HIGHEST_PROTOCOL) + os.replace(tmp, combined) + return grid diff --git a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py new file mode 100644 index 000000000..57acb50db --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py @@ -0,0 +1,198 @@ +"""Q integration and Y/qT rebin helpers for the SCETlib bT-grid ParamModel. + +All weight-construction is numpy (runs once at construction time); runtime +contractions are simple ``tf.tensordot`` / ``tf.einsum`` calls. + +Three pieces: + +1. :func:`dense_index_map` — build a ``(NQ, NY, NqT)`` int array mapping each + rectangular grid cell to a flat bin index in the sparse btgrid; ``-1`` + marks missing combos. Use ``tf.gather`` with a sentinel to pad a sparse + ``(Nbins,)`` σ tensor into a dense ``(NQ, NY, NqT)``. + +2. :func:`q_integrate_weights` — produces a ``(NQ,)`` weight vector for + ``arctan_Q²``-method Simpson integration over the Z mass window. Apply + via ``tf.einsum('q, qyz -> yz', w, sigma)``. + +3. :func:`rebin_weights` — given a fine source grid and a coarser target-bin + edge list, builds a ``(N_target, N_source)`` Simpson weight matrix. + Apply via ``tf.tensordot``. + +Parity tests against the numpy reference in :mod:`scetlib_btgrid_numpy` are in +:mod:`scetlib_btgrid_tf_parity` (added in Phase 3). +""" + +import numpy as np +import tensorflow as tf + +from wremnants.postprocessing.scetlib_np.btgrid_tf import ( + _as_dtype, + simpson_weights, +) + +# Z resonance defaults (matches integrate_over_Q in btgrid_numpy). +MZ_PDG = 91.1876 +GAMMAZ_PDG = 2.4952 + + +# ============================================================================= +# Sparse → dense index map +# ============================================================================= + + +def dense_index_map(bins, Q_unique=None, Y_unique=None, qT_unique=None): + """Build a dense rectangular index map for a sparse btgrid. + + Parameters + ---------- + bins : list of tuples + From ``load_btgrid_shards``: each element is ``(Q, Y, qT, lep)``, + sorted lexicographically. + + Returns + ------- + dict with keys: + Q_unique, Y_unique, qT_unique : sorted unique axis values + flat_idx : ndarray of shape (NQ, NY, NqT), int64. -1 marks missing. + missing_count : int + """ + if Q_unique is None: + Q_unique = sorted({b[0] for b in bins}) + if Y_unique is None: + Y_unique = sorted({b[1] for b in bins}) + if qT_unique is None: + qT_unique = sorted({b[2] for b in bins}) + + Q_unique = np.asarray(Q_unique, dtype=np.float64) + Y_unique = np.asarray(Y_unique, dtype=np.float64) + qT_unique = np.asarray(qT_unique, dtype=np.float64) + + Q_pos = {Q: i for i, Q in enumerate(Q_unique)} + Y_pos = {Y: i for i, Y in enumerate(Y_unique)} + qT_pos = {qT: i for i, qT in enumerate(qT_unique)} + + flat_idx = np.full( + (Q_unique.size, Y_unique.size, qT_unique.size), -1, dtype=np.int64 + ) + for k, (Q, Y, qT, _lep) in enumerate(bins): + flat_idx[Q_pos[Q], Y_pos[Y], qT_pos[qT]] = k + + missing = int(np.sum(flat_idx == -1)) + return dict( + Q_unique=Q_unique, + Y_unique=Y_unique, + qT_unique=qT_unique, + flat_idx=flat_idx, + missing_count=missing, + ) + + +def sparse_to_dense_tf(sigma_flat, flat_idx): + """Reshape a sparse ``(Nbins,)`` σ tensor to dense ``(NQ, NY, NqT)``. + + Missing cells (``flat_idx == -1``) are padded with 0. Implemented via + ``tf.gather`` with a 0-padded sentinel row. + """ + sigma_flat = _as_dtype(sigma_flat) + # Append one extra "zero" entry that the -1 indices will gather. + extended = tf.concat([sigma_flat, tf.zeros([1], dtype=sigma_flat.dtype)], axis=0) + sentinel = tf.cast(tf.shape(sigma_flat)[0], tf.int64) # index of the appended zero + idx_safe = tf.where(tf.equal(flat_idx, -1), sentinel, flat_idx) + return tf.gather(extended, idx_safe) + + +# ============================================================================= +# Q-integration weights (arctan_Q² method, matches numpy integrate_over_Q) +# ============================================================================= + + +def q_integrate_weights(Q_grid, Q_lo, Q_hi, q0=MZ_PDG, Gamma=GAMMAZ_PDG): + """Simpson weights for integrating over Q ∈ [Q_lo, Q_hi] in arctan-Q² space. + + Implements the same change of variable as + :func:`scetlib_btgrid_numpy.integrate_over_Q` with ``method="arctan_Q2"``: + x = arctan((Q² - q0²) / (q0 Γ)) flattens the Breit-Wigner peak, then + Simpson on x with the Jacobian dQ/dx. + + Returns a ``(NQ,)`` weight vector with zeros outside ``[Q_lo, Q_hi]``. + """ + Q_grid = np.asarray(Q_grid, dtype=np.float64) + mask = (Q_grid >= Q_lo) & (Q_grid <= Q_hi) + if mask.sum() < 2: + raise ValueError( + f"q_integrate_weights: need ≥ 2 Q samples in [{Q_lo}, {Q_hi}]; " + f"got {mask.sum()} from Q_grid={Q_grid}" + ) + Q_sub = Q_grid[mask] + x = np.arctan((Q_sub**2 - q0**2) / (q0 * Gamma)) + jac = (q0 * Gamma + (Q_sub**2 - q0**2) ** 2 / (q0 * Gamma)) / (2.0 * Q_sub) + w_simpson = simpson_weights(x) # weights in x-space + w_full = w_simpson * jac # at each Q sample + w_padded = np.zeros_like(Q_grid) + w_padded[mask] = w_full + return w_padded + + +def integrate_over_Q_tf(sigma_QYqT, Q_weights): + """Apply precomputed Q weights. ``sigma_QYqT`` shape ``(NQ, NY, NqT)``.""" + Q_weights = _as_dtype(Q_weights) + return tf.einsum("q, qyz -> yz", Q_weights, sigma_QYqT) + + +# ============================================================================= +# Y / qT rebin weights +# ============================================================================= + + +def rebin_weights(source_grid, target_edges, name="axis", tol=1e-9): + """Build a ``(N_target, N_source)`` Simpson rebin matrix. + + For each target bin ``[target_edges[i], target_edges[i+1]]``, find the + source samples that fall in (with ``tol`` slack) the bin's interior + + edges, then compute Simpson weights for those samples. Returns a dense + matrix; entries are 0 for source samples not contributing to a given + target bin. + + Mirrors the per-bin call pattern of + :func:`scetlib_btgrid_numpy.integrate_over_axis_bin`. + """ + source_grid = np.asarray(source_grid, dtype=np.float64) + target_edges = np.asarray(target_edges, dtype=np.float64) + if target_edges.ndim != 1 or target_edges.size < 2: + raise ValueError(f"rebin_weights[{name}]: need ≥ 2 target edges") + + N_target = target_edges.size - 1 + N_source = source_grid.size + W = np.zeros((N_target, N_source), dtype=np.float64) + for i in range(N_target): + lo, hi = target_edges[i], target_edges[i + 1] + mask = (source_grid >= lo - tol) & (source_grid <= hi + tol) + if mask.sum() < 2: + raise ValueError( + f"rebin_weights[{name}]: bin [{lo}, {hi}] has only " + f"{mask.sum()} source samples; need ≥ 2" + ) + sub_grid = source_grid[mask] + w_sub = simpson_weights(sub_grid) + W[i, mask] = w_sub + return W + + +def rebin_axis_tf(values, axis, weights): + """Apply rebin weights along ``axis`` of ``values``. + + ``values`` shape ``(..., N_source, ...)`` (source axis at position ``axis``). + ``weights`` shape ``(N_target, N_source)``. Returns shape + ``(..., N_target, ...)`` with the source axis replaced by the target axis. + """ + values = _as_dtype(values) + weights = _as_dtype(weights) + rank = len(values.shape) + if axis < 0: + axis += rank + # tensordot contracts values[axis] with weights[1]; target axis ends up + # at the END of the result. Permute it back to position ``axis``. + out = tf.tensordot(values, weights, axes=[[axis], [1]]) + perm = list(range(rank - 1)) + perm.insert(axis, rank - 1) + return tf.transpose(out, perm) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_numpy.py b/wremnants/postprocessing/scetlib_np/btgrid_numpy.py new file mode 100644 index 000000000..e4d271114 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/btgrid_numpy.py @@ -0,0 +1,713 @@ +# ------------------------------------------------------------------------------- +# NP factorization library for the bT-grid workflow. +# +# Vendored copy of: +# /work/submit/lavezzo/alphaS/scetlib-cms-newnp-lambda4fix/prod/scetlib_run/ +# scetlib_run/factorize.py +# Numpy-only; serves as the reference implementation against which the TF port +# (scetlib_btgrid_tf.py, Phase 2) is parity-tested. Kept in sync manually with +# the upstream scetlib repo — when the upstream changes, recopy this file and +# rerun the parity tests. +# +# Provides: +# - Pure-numpy transcriptions of NP_model_effective (F_eff) and +# NP_model_gammanu (gamma_nu^NP) that match the C++ code byte-for-byte. +# - A vectorised Hankel reconstruction of sigma(qT) from a cached bT-grid: +# sigma(qT) ~ int dbT bT J0(qT bT) * I_pert(b_bar) +# * exp(C_nu(bT) * gamma_nu^NP(b_bar)) +# * F_eff(Y, b_bar) +# - Loaders for the bT-grid pickle shards produced by --bt-grid and for the +# prior-art spectrum-mode "combined" pickles. +# +# Self-contained: depends on numpy only, no scipy / SCETlib runtime at import. +# ------------------------------------------------------------------------------- + +import glob +import os +import pickle +import sys + +import numpy as np + + +# Compat shim: pickles produced with numpy >= 2.0 reference `numpy._core`, +# which does not exist in numpy < 2.0. Alias the old `numpy.core` under the +# new name (and a few of its submodules) so unpickling succeeds. +def _ensure_numpy_core_alias(): + if "numpy._core" in sys.modules: + return + try: + import numpy._core # noqa: F401 + except ImportError: + try: + from numpy import core as _np_core + except ImportError: + return + sys.modules["numpy._core"] = _np_core + for name in ( + "multiarray", + "numeric", + "fromnumeric", + "umath", + "shape_base", + "_methods", + ): + sub = getattr(_np_core, name, None) + if sub is not None: + sys.modules[f"numpy._core.{name}"] = sub + + +_ensure_numpy_core_alias() + + +# ============================================================================= +# Numerics: bare-bones bessel J0 and Simpson, both numpy-only (the container +# in which the fit runs has no scipy). +# ============================================================================= + + +def bessel_j0(x): + """J_0(x) via Abramowitz & Stegun 9.4.1 / 9.4.3. Accurate to ~1.6e-8.""" + x = np.asarray(x, dtype=float) + out = np.empty_like(x) + small = np.abs(x) < 3.0 + xs = x[small] / 3.0 + y = xs * xs + out[small] = ( + 1.0 + - 2.2499997 * y + + 1.2656208 * y**2 + - 0.3163866 * y**3 + + 0.0444479 * y**4 + - 0.0039444 * y**5 + + 0.0002100 * y**6 + ) + xl = np.abs(x[~small]) + z = 3.0 / xl + f0 = ( + 0.79788456 + - 0.00000077 * z + - 0.00552740 * z**2 + - 0.00009512 * z**3 + + 0.00137237 * z**4 + - 0.00072805 * z**5 + + 0.00014476 * z**6 + ) + theta0 = ( + xl + - 0.78539816 + - 0.04166397 * z + - 0.00003954 * z**2 + + 0.00262573 * z**3 + - 0.00054125 * z**4 + - 0.00029333 * z**5 + + 0.00013558 * z**6 + ) + out[~small] = f0 * np.cos(theta0) / np.sqrt(xl) + return out + + +def simpson(y, x, axis=-1): + """Composite Simpson on a (possibly non-uniform) 1-D grid, vectorised + along the given axis. Falls back to trapezoid for the last segment when + the number of intervals is odd.""" + y = np.asarray(y, dtype=float) + x = np.asarray(x, dtype=float) + if x.ndim != 1: + raise ValueError("simpson expects 1-D x") + n = x.size - 1 + if n < 1: + return np.zeros( + y.shape[:-1] if axis == -1 else y.shape[:axis] + y.shape[axis + 1 :] + ) + # move integration axis to the end for simpler slicing + y_moved = np.moveaxis(y, axis, -1) + if n % 2 == 1: + lead = simpson(np.moveaxis(y_moved[..., :-1], -1, axis), x[:-1], axis=axis) + tail = 0.5 * (y_moved[..., -1] + y_moved[..., -2]) * (x[-1] - x[-2]) + return lead + tail + h = np.diff(x) + h0 = h[0::2] + h1 = h[1::2] + s = ( + (h0 + h1) + / 6.0 + * ( + y_moved[..., 0:-1:2] * (2.0 - h1 / h0) + + y_moved[..., 1::2] * (h0 + h1) ** 2 / (h0 * h1) + + y_moved[..., 2::2] * (2.0 - h0 / h1) + ) + ) + return np.sum(s, axis=-1) + + +# ============================================================================= +# NP model transcriptions. Mirror the C++ implementations one-to-one; expect +# bT to be the b̄T (= b_star) at which the NP factor enters. +# ============================================================================= + +# NP_model_effective.np_model enum values supported here: +EFF_MODELS = { + "identity", + "tanh_2", + "tanh_6", + "tanh_4", + "frac_2", + "frac_4", + "exp_2", + "exp_4", + "signed_lambda", + "hyp_tangent", + "square_root", +} + +# NP_model_gammanu.np_model_nu enum values supported here: +GNU_MODELS = { + "tanh_1", + "tanh_2", + "tanh_6", + "frac_1", + "frac_2", + "exp_1", + "exp_2", + "hyp_tangent", + "linear", +} + + +def F_eff(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_model): + """F_eff(Y, b̄T) — NP_model_effective::operator() from NP_models.hpp.""" + bT = np.asarray(bT, dtype=float) + + if np_model == "signed_lambda": + lambda2_Y = lambda2 + delta_lambda2 * Y * Y + if lambda4 <= 0.0 and (lambda2 != 0.0 or delta_lambda2 != 0.0): + raise ValueError( + "signed_lambda requires lambda4 > 0 when lambda2 or delta_lambda2 != 0" + ) + return (1.0 + lambda2_Y * bT**2) ** 2 * np.exp(-2.0 * lambda4 * bT**4) + + lambda2_Y = lambda2 + delta_lambda2 * Y * Y + arg = (lambda2_Y + lambda4 * bT**2) * bT + + if np_model == "identity": + return np.exp(-2.0 * bT * arg) + + if lambda_inf == 0.0: + return np.ones_like(bT) + + arg = arg / lambda_inf + # alias support + model = {"hyp_tangent": "tanh_2", "square_root": "frac_2"}.get(np_model, np_model) + + if model == "tanh_2": + arg = arg + (1.0 / 3.0) * (lambda2_Y * bT / lambda_inf) ** 3 + func = np.tanh(arg) + elif model == "tanh_6": + arg = arg + lambda6 * bT**5 / lambda_inf + arg = arg + (1.0 / 3.0) * (lambda2_Y * bT / lambda_inf) ** 3 + func = np.tanh(arg) + elif model == "tanh_4": + func = np.sqrt(np.tanh(arg**2)) + elif model == "frac_2": + arg = arg + 0.5 * (lambda2_Y * bT / lambda_inf) ** 3 + func = arg / np.sqrt(1.0 + arg**2) + elif model == "frac_4": + func = arg / np.sqrt(np.sqrt(1.0 + arg**4)) + elif model == "exp_2": + arg = arg + 0.25 * (lambda2_Y * bT / lambda_inf) ** 3 + func = np.sqrt(-np.expm1(-(arg**2))) + elif model == "exp_4": + func = np.sqrt(np.sqrt(-np.expm1(-(arg**4)))) + else: + raise ValueError(f"F_eff: unsupported np_model {np_model!r}") + + return np.exp(-2.0 * lambda_inf * bT * func) + + +def gamma_nu_NP(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): + """gamma_nu^NP(b̄T) — NP_model_gammanu::model_gammanu() from Gamma_nu.hpp. + + Note: NP_model_gammanu has its own b_star() (b0_bmax_nu), but the C++ + code calls model_gammanu(bT) with the raw bT passed into Gamma_nu — which + is b_star_global by the time it reaches us. So pass b̄T = b_star_global(bT) + here; b0_bmax_nu does NOT enter model_gammanu directly. + """ + bT = np.asarray(bT, dtype=float) + if lambda_inf_nu == 0.0: + return np.zeros_like(bT) + + bT2 = bT * bT + arg = (lambda2_nu + lambda4_nu * bT2) * bT2 / lambda_inf_nu + + model = {"hyp_tangent": "tanh_2", "linear": "frac_1"}.get(np_model_nu, np_model_nu) + + if model == "tanh_1": + arg = arg + (2.0 / 3.0) * (lambda2_nu * bT2 / lambda_inf_nu) ** 2 + func = np.tanh(np.sqrt(arg)) ** 2 + elif model == "tanh_2": + func = np.tanh(arg) + elif model == "tanh_6": + # NP_model_gammanu hardcodes lambda6_nu = 0.0007 (Gamma_nu.hpp:102) + arg = arg + 0.0007 * bT2**3 / lambda_inf_nu + func = np.tanh(arg) + elif model == "frac_1": + arg = arg + (lambda2_nu * bT2 / lambda_inf_nu) ** 2 + func = arg / (1.0 + arg) + elif model == "frac_2": + func = arg / np.sqrt(1.0 + arg**2) + elif model == "exp_1": + arg = arg + 0.5 * (lambda2_nu * bT2 / lambda_inf_nu) ** 2 + func = -np.expm1(-arg) + elif model == "exp_2": + func = np.sqrt(-np.expm1(-(arg**2))) + else: + raise ValueError(f"gamma_nu_NP: unsupported np_model_nu {np_model_nu!r}") + + return -lambda_inf_nu * func + + +# Convenience defaults: "all NP knobs off". These reproduce the SCETlib +# configuration that --bt-grid mode caches. +NP_ZERO_EFF = dict( + lambda_inf=0.0, + lambda2=0.0, + lambda4=0.0, + lambda6=0.0, + delta_lambda2=0.0, + np_model="identity", +) +NP_ZERO_GNU = dict( + lambda_inf_nu=0.0, lambda2_nu=0.0, lambda4_nu=0.0, np_model_nu="tanh_2" +) + + +def eff_params_from_conf(conf): + """Read NP_model_effective parameters from a configparser [Nonperturbative] + section (defaults to NP_ZERO_EFF if a field is absent).""" + if "Nonperturbative" not in conf: + return dict(NP_ZERO_EFF) + sec = conf["Nonperturbative"] + return dict( + lambda_inf=sec.getfloat("lambda_inf", fallback=NP_ZERO_EFF["lambda_inf"]), + lambda2=sec.getfloat("lambda2", fallback=NP_ZERO_EFF["lambda2"]), + lambda4=sec.getfloat("lambda4", fallback=NP_ZERO_EFF["lambda4"]), + lambda6=sec.getfloat("lambda6", fallback=NP_ZERO_EFF["lambda6"]), + delta_lambda2=sec.getfloat( + "delta_lambda2", fallback=NP_ZERO_EFF["delta_lambda2"] + ), + np_model=sec.get("np_model", fallback=NP_ZERO_EFF["np_model"]), + ) + + +def gnu_params_from_conf(conf): + """Read NP_model_gammanu parameters from a configparser [Nonperturbative] + section (defaults to NP_ZERO_GNU if a field is absent).""" + if "Nonperturbative" not in conf: + return dict(NP_ZERO_GNU) + sec = conf["Nonperturbative"] + return dict( + lambda_inf_nu=sec.getfloat( + "lambda_inf_nu", fallback=NP_ZERO_GNU["lambda_inf_nu"] + ), + lambda2_nu=sec.getfloat("lambda2_nu", fallback=NP_ZERO_GNU["lambda2_nu"]), + lambda4_nu=sec.getfloat("lambda4_nu", fallback=NP_ZERO_GNU["lambda4_nu"]), + np_model_nu=sec.get("np_model_nu", fallback=NP_ZERO_GNU["np_model_nu"]), + ) + + +# ============================================================================= +# Hankel reconstruction +# ============================================================================= + + +def reconstruct_one(qT, bT, I_pert, C_nu, b_bar, Y, eff_params, gnu_params): + """Hankel-reconstruct one sigma(qT) from cached arrays at a single (Q,Y,qT). + + Returns the differential structure-function value at the point (Q,Y,qT), + matching SCETlib's spectrum-mode ang.c convention (i.e. including the + qT factor that arises from the integration-variable choice in SCETlib's + integrator_de_oscillatory). + + qT : scalar + bT : (Nb,) integration variable (raw bT) + I_pert : (Nb,) cached SCETlib integrand at NP off + C_nu : (Nb,) coefficient of gamma_nu^NP in log evolution + b_bar : (Nb,) b_star_global(bT) where NP factors evaluate + Y : scalar rapidity + eff_params : dict for NP_model_effective parameters + gnu_params : dict for NP_model_gammanu parameters + """ + g_NP = gamma_nu_NP(b_bar, **gnu_params) + Feff = F_eff(Y, b_bar, **eff_params) + integrand = bT * bessel_j0(qT * bT) * I_pert * np.exp(C_nu * g_NP) * Feff + return qT * simpson(integrand, bT) + + +def reconstruct_grid_QYqT( + Q_grid, + Y_grid, + qT_grid, + bT, + I_pert, + C_nu, + b_bar, + eff_params, + gnu_params, + verbose=True, +): + """Reconstruct sigma at every (Q, Y, qT) sample point of a regularly-indexed + grid. Computes J0(qT*bT) once on (NqT, Nbt) (rather than the redundant + (NQ*NY*NqT, Nbt) the naive batch would build) and processes one Q-slice at + a time to bound peak memory and emit progress. + + Q_grid : (NQ,) point values in Q + Y_grid : (NY,) point values in Y + qT_grid : (NqT,) point values in qT + bT : (Nbt,) bT integration variable shared by all points + I_pert : (Npts, Nbt) cached perturbative integrand, Npts = NQ*NY*NqT, + indexed as flat list in the same order load_btgrid_shards + returns its `bins` list (sorted lexicographically by Q, Y, qT). + C_nu : (Npts, Nbt) rapidity-evolution-log coefficient + b_bar : (Nbt,) b_star_global(bT) + + Returns: (NQ, NY, NqT) ndarray of sigma_factorised values. + """ + Q_grid = np.asarray(Q_grid, dtype=float) + Y_grid = np.asarray(Y_grid, dtype=float) + qT_grid = np.asarray(qT_grid, dtype=float) + bT = np.asarray(bT, dtype=float) + NQ, NY, NqT = Q_grid.size, Y_grid.size, qT_grid.size + Nbt = bT.size + Npts = NQ * NY * NqT + if I_pert.shape != (Npts, Nbt): + raise ValueError( + f"I_pert shape {I_pert.shape} doesn't match expected " + f"({Npts}, {Nbt}) from {NQ} Q x {NY} Y x {NqT} qT" + ) + + # --- shared (Y-independent) factors over bT --- + g_NP = gamma_nu_NP(b_bar, **gnu_params) # (Nbt,) + delta_l2 = eff_params.get("delta_lambda2", 0.0) + if delta_l2 == 0.0: + Feff_bT = F_eff(0.0, b_bar, **eff_params) # (Nbt,) + else: + Feff_bT = None # per-Y below + + # --- bT*J0(qT*bT) cached once over (NqT, Nbt) --- + bT_J0 = bT[np.newaxis, :] * bessel_j0(qT_grid[:, np.newaxis] * bT[np.newaxis, :]) + # shape (NqT, Nbt); used in every Q-slice below. + + # I_pert and C_nu are stored as (NQ*NY*NqT, Nbt); reshape view to + # (NQ, NY, NqT, Nbt) and process Q-by-Q. + I_pert_r = I_pert.reshape(NQ, NY, NqT, Nbt) + C_nu_r = C_nu.reshape(NQ, NY, NqT, Nbt) + + out = np.empty((NQ, NY, NqT), dtype=float) + + import time + + t0 = time.time() + for iQ in range(NQ): + # exp(C_nu * g_NP) on the (NY, NqT, Nbt) Q-slice -- the only piece + # that doesn't factor across the slice + exp_g_factor = np.exp( + C_nu_r[iQ] * g_NP[np.newaxis, np.newaxis, :] + ) # (NY, NqT, Nbt) + if delta_l2 == 0.0: + # Feff_bT shape (Nbt,); broadcast over (NY, NqT) + integrand = ( + bT_J0[np.newaxis, :, :] + * I_pert_r[iQ] + * exp_g_factor + * Feff_bT[np.newaxis, np.newaxis, :] + ) + else: + # Feff depends on Y (per-Y row), shape (NY, Nbt) broadcast over qT + Feff = np.stack( + [F_eff(Y_i, b_bar, **eff_params) for Y_i in Y_grid] + ) # (NY, Nbt) + integrand = ( + bT_J0[np.newaxis, :, :] + * I_pert_r[iQ] + * exp_g_factor + * Feff[:, np.newaxis, :] + ) + # Simpson over bT (last axis) + sigma_Q = simpson(integrand, bT, axis=-1) # (NY, NqT) + # qT factor (SCETlib's x = qT*bT integration convention) + out[iQ] = qT_grid[np.newaxis, :] * sigma_Q + + if verbose: + elapsed = time.time() - t0 + print( + f" [hankel] Q-slice {iQ+1}/{NQ} done " + f"({elapsed:.1f}s elapsed, ETA {elapsed*(NQ-iQ-1)/(iQ+1):.1f}s)", + flush=True, + ) + return out + + +def integrate_over_Q( + sigma_QYqT, Q_grid, Q_lo, Q_hi, method="arctan_Q2", q0=91.1876, Gamma=2.4952 +): + """Integrate sigma(Q, Y, qT) over Q in [Q_lo, Q_hi]. Q-samples outside + [Q_lo, Q_hi] are dropped. + + sigma_QYqT : (NQ, NY, NqT) point values + Q_grid : (NQ,) Q sample positions (need not be uniform) + Q_lo, Q_hi : integration limits + method : "arctan_Q2" (default; uses x = arctan((Q²-q0²)/(q0*Gamma)) so + the Breit-Wigner Z resonance becomes smooth) or "simpson" + (Simpson directly in Q) or "trapz" (trapezoid in Q). + q0, Gamma : resonance mass and width for the arctan_Q2 transform. + Defaults match Z-boson values (mZ = 91.1876, ΓZ = 2.4952). + + Returns: (NY, NqT) integrated values. + """ + Q_grid = np.asarray(Q_grid, dtype=float) + mask = (Q_grid >= Q_lo) & (Q_grid <= Q_hi) + if mask.sum() < 2: + raise ValueError(f"Need >= 2 Q samples in [{Q_lo}, {Q_hi}]; got {mask.sum()}") + Q_sub = Q_grid[mask] + s_sub = sigma_QYqT[mask] # (NQ_sub, NY, NqT) + s_moved = np.moveaxis(s_sub, 0, -1) # (..., NQ_sub) + + if method == "simpson": + return simpson(s_moved, Q_sub, axis=-1) + if method == "trapz": + return np.trapz(s_moved, Q_sub, axis=-1) + if method == "arctan_Q2": + # x = arctan((Q² - q0²) / (q0 * Gamma)) → flattens the Breit-Wigner peak + x = np.arctan((Q_sub**2 - q0**2) / (q0 * Gamma)) + # dQ/dx = ( q0*Gamma + (Q² - q0²)² / (q0*Gamma) ) / (2 Q) + jac = (q0 * Gamma + (Q_sub**2 - q0**2) ** 2 / (q0 * Gamma)) / (2.0 * Q_sub) + return simpson(s_moved * jac, x, axis=-1) + raise ValueError(f"integrate_over_Q: unknown method {method!r}") + + +def integrate_over_axis_bin(values, axis_grid, axis_lo, axis_hi, name="axis"): + """Integrate a 1-D array of sample values over a single bin [axis_lo, axis_hi] + using Simpson's rule on whatever sample points fall in (and on) the bin edges. + + Designed for use with a btgrid sampled at bin-edges + bin-centres (3 samples + per bin minimum). When the grid carries the bin's 2 edges + central point, + this is a 3-point Simpson (4th-order accurate per bin); when more samples + fall inside the bin, simpson naturally extends to higher order via composite + rule. + + values : (..., N_axis) array values at axis_grid sample points + axis_grid : (N_axis,) sample positions (sorted, not necessarily uniform) + axis_lo, axis_hi : integration limits (bin edges) + + Returns: integrated value with the axis collapsed. + """ + axis_grid = np.asarray(axis_grid, dtype=float) + values = np.asarray(values, dtype=float) + # tolerance for "on the edge" lookups (FP-imprecise bin centres show up as + # e.g. -2.3499999999999996; bin widths are at least ~0.025 so 1e-9 is safe) + tol = 1e-9 + mask = (axis_grid >= axis_lo - tol) & (axis_grid <= axis_hi + tol) + if mask.sum() < 2: + raise ValueError( + f"integrate_over_axis_bin({name}): need >= 2 samples in " + f"[{axis_lo}, {axis_hi}]; got {mask.sum()}" + ) + sub = values[..., mask] + g = axis_grid[mask] + return simpson(sub, g, axis=-1) + + +def integrate_over_Y_bin(sigma_YqT, Y_grid, Y_lo, Y_hi): + """Integrate sigma(Y, qT) over Y in [Y_lo, Y_hi] using sample points of + Y_grid that fall in or on the bin edges. Returns (NqT,) array. + + Expects Y_grid to include Y_lo and Y_hi as samples (the edges) and ideally + the bin centre too — Simpson uses all available samples in [Y_lo, Y_hi].""" + # move Y axis (axis 0) to the end for integrate_over_axis_bin which uses last axis + return integrate_over_axis_bin( + np.moveaxis(sigma_YqT, 0, -1), Y_grid, Y_lo, Y_hi, name="Y" + ) + + +def integrate_over_qT_bin(sigma_YqT, qT_grid, qT_lo, qT_hi): + """Integrate sigma(Y, qT) over qT in [qT_lo, qT_hi]. Returns (NY,) array. + + Expects qT_grid to include qT_lo and qT_hi as samples (the bin edges).""" + return integrate_over_axis_bin(sigma_YqT, qT_grid, qT_lo, qT_hi, name="qT") + + +def reconstruct_batch( + qT_per_bin, bT, I_pert, C_nu, b_bar, Y_per_bin, eff_params, gnu_params +): + """Vectorised reconstruction for many bins at once. + + qT_per_bin : (Nbins,) qT for each bin + Y_per_bin : (Nbins,) Y for each bin (used by F_eff's δλ_2·Y² term) + bT : (Nbt,) integration variable (raw bT), shared by all bins + I_pert : (Nbins, Nbt) cached perturbative bT integrand + C_nu : (Nbins, Nbt) rapidity log coefficient + b_bar : (Nbt,) b_star_global(bT), shared by all bins + + Returns: (Nbins,) sigma_factorized for each bin + """ + qT_per_bin = np.asarray(qT_per_bin, dtype=float) + Y_per_bin = np.asarray(Y_per_bin, dtype=float) + bT = np.asarray(bT, dtype=float) + I_pert = np.asarray(I_pert, dtype=float) + C_nu = np.asarray(C_nu, dtype=float) + b_bar = np.asarray(b_bar, dtype=float) + + # gamma_nu^NP and the gamma_nu exponential factor are bin-shared + # (bT-dependent only); F_eff depends on Y through δλ_2·Y², so we evaluate + # it per bin if δλ_2 != 0. + g_NP = gamma_nu_NP(b_bar, **gnu_params) # (Nbt,) + exp_g_factor = np.exp(C_nu * g_NP[np.newaxis, :]) # (Nbins, Nbt) + + delta_l2 = eff_params.get("delta_lambda2", 0.0) + if delta_l2 == 0.0: + # F_eff has no Y dependence beyond a constant Y² factor, so identical + # across bins -> compute once + Feff = F_eff(0.0, b_bar, **eff_params) # (Nbt,) + bT_J0 = bT * bessel_j0( + qT_per_bin[:, np.newaxis] * bT[np.newaxis, :] + ) # (Nbins, Nbt) + integrand = bT_J0 * I_pert * exp_g_factor * Feff[np.newaxis, :] + else: + # need per-bin F_eff because of Y dependence in lambda2_Y + Feff = np.empty_like(I_pert) + for i, Y_i in enumerate(Y_per_bin): + Feff[i] = F_eff(Y_i, b_bar, **eff_params) + bT_J0 = bT * bessel_j0(qT_per_bin[:, np.newaxis] * bT[np.newaxis, :]) + integrand = bT_J0 * I_pert * exp_g_factor * Feff + + # Multiply by qT to match SCETlib's spectrum-mode integration convention + # (SCETlib's _int_bT integrates in x = qT*bT, picking up an explicit qT + # factor via the Jacobian; we integrate in bT directly so we must add it + # back). + return qT_per_bin * simpson(integrand, bT, axis=-1) + + +# ============================================================================= +# Loaders for the artefacts produced by the bT-grid condor run and by the +# spectrum-mode "combined" pickles. +# ============================================================================= + + +def load_btgrid_shards(submitdir_or_glob, runcard_basename=None): + """Combine bT-grid shards produced by --bt-grid mode. + + `submitdir_or_glob` may be: + - a directory: we look for ``*_btgrid.pkl`` inside (recursively only + one level via scetlib_outputs/). + - a glob pattern: used directly. + + Returns dict with: + bT : (Nbt,) + b_bar : (Nbt,) + bins : list of (Q, Y, qT, lep) bin centres, length Nbins + vars : dict of variation index -> setting dict (copied from the + first shard; all shards are expected to carry the same set) + I_pert : (Nvars, Nbins, Nbt) + C_nu : (Nvars, Nbins, Nbt) + config : dict from the first shard (perturbative configuration the + grid was generated against) + n_shards: int + """ + if os.path.isdir(submitdir_or_glob): + candidates = [ + os.path.join(submitdir_or_glob, "scetlib_outputs", "*_btgrid.pkl"), + os.path.join(submitdir_or_glob, "*_btgrid.pkl"), + ] + else: + candidates = [submitdir_or_glob] + + files = [] + for pat in candidates: + files = sorted(glob.glob(pat)) + if files: + break + if not files: + raise FileNotFoundError(f"No btgrid shards found under {submitdir_or_glob!r}") + + # First shard sets the schema; later shards must match. + with open(files[0], "rb") as f: + first = pickle.load(f) + if first.get("schema_version") != "bt_grid_v1": + raise ValueError( + f"Unexpected schema {first.get('schema_version')!r} in {files[0]}" + ) + bT = np.asarray(first["bT"], dtype=float) + b_bar = np.asarray(first["b_bar"], dtype=float) + varis = first["vars"] + config = first["config"] + n_vars = len(varis) + n_bt = bT.size + + # We don't know Nbins ahead of time without scanning all shards. Walk them + # once: build a dict of bin -> (var_idx -> (I_pert_row, C_nu_row)). + bin_to_data = {} + for path in files: + with open(path, "rb") as f: + d = pickle.load(f) + if d.get("schema_version") != "bt_grid_v1": + raise ValueError( + f"Mixed schema versions: {path} has {d.get('schema_version')}" + ) + if d["bT"].shape != bT.shape or not np.allclose(d["bT"], bT): + raise ValueError(f"bT grid mismatch in {path}") + bins_local = d["bins"] + I_local = np.asarray( + d["I_pert"], dtype=float + ) # (Nvars_local, Nbins_local, Nbt) + C_local = np.asarray(d["C_nu"], dtype=float) + # Map local variation indices to the union variation order. We assume + # all shards share the same vars dict (true when produced by the same + # condor submission). + var_order_local = list(d["vars"].keys()) + for b_idx, b_tup in enumerate(bins_local): + slot = bin_to_data.setdefault(tuple(b_tup), {}) + for v_pos, v_idx in enumerate(var_order_local): + slot[v_idx] = (I_local[v_pos, b_idx], C_local[v_pos, b_idx]) + + var_order = list(varis.keys()) + bins_sorted = sorted(bin_to_data.keys(), key=lambda t: (t[0], t[1], t[2])) + n_bins = len(bins_sorted) + I_pert = np.full((n_vars, n_bins, n_bt), np.nan, dtype=float) + C_nu = np.full((n_vars, n_bins, n_bt), np.nan, dtype=float) + for b_pos, b_tup in enumerate(bins_sorted): + per_var = bin_to_data[b_tup] + for v_pos, v_idx in enumerate(var_order): + if v_idx in per_var: + I_pert[v_pos, b_pos] = per_var[v_idx][0] + C_nu[v_pos, b_pos] = per_var[v_idx][1] + + return { + "bT": bT, + "b_bar": b_bar, + "bins": bins_sorted, + "vars": varis, + "var_order": var_order, + "I_pert": I_pert, + "C_nu": C_nu, + "config": config, + "n_shards": len(files), + } + + +def load_spectrum_reference(combined_pkl): + """Load a 'combined' spectrum-mode pickle written by scetlib-run-qT.py. + + Returns a dict + hist : the hist.Hist object stored in the pickle + config : the [section]->dict mapping the run used + meta_data: the meta-data dict + Use the returned hist directly (axes are typically Q, Y, qT, lep, vars). + """ + with open(combined_pkl, "rb") as f: + d = pickle.load(f) + return { + "hist": d.get("hist"), + "config": d.get("config", {}), + "meta_data": d.get("meta_data", {}), + } diff --git a/wremnants/postprocessing/scetlib_np/btgrid_tf.py b/wremnants/postprocessing/scetlib_np/btgrid_tf.py new file mode 100644 index 000000000..f73e80b92 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/btgrid_tf.py @@ -0,0 +1,318 @@ +"""TensorFlow port of the bT-grid factorization library. + +Mirrors :mod:`scetlib_btgrid_numpy` function-by-function. The numpy module is +the byte-for-byte transcription of SCETlib C++ and the parity test in +:mod:`scetlib_btgrid_tf_parity` keeps the two in sync. + +Design choices: + * ``np_model`` / ``np_model_nu`` strings are fixed at trace time (the SCETlib + runcard sets them once per fit). The TF functions dispatch on the string at + Python level — no ``tf.cond``. + * λ parameters are TF tensors (typically scalars, but broadcasting follows + the same rules as numpy). + * All ops are differentiable in λ. Branches on λ values use ``tf.where`` + with a safe denominator to avoid NaN gradients. + * ``b_star_global`` is not ported — the cached ``b_bar`` array in the bT-grid + shards is precomputed and travels as a ``tf.constant``. + * Simpson weights are precomputed at trace time from the (static) bT, Y, qT + grids; the runtime cost is just ``tf.reduce_sum(w * y)``. +""" + +from typing import Mapping + +import numpy as np +import tensorflow as tf + +# Set the dtype used for all ops in this module. Match the numpy reference +# (which uses ``float`` ≡ float64) to keep parity tight. +DTYPE = tf.float64 + + +def _as_dtype(x, dtype=DTYPE): + """Coerce ``x`` to ``dtype`` without losing precision on Python scalars. + + ``tf.cast(0.4, tf.float64)`` round-trips through float32 (returning + ``0.4000000059604645``); ``tf.constant(0.4, dtype=tf.float64)`` does not. + Use this helper everywhere a possibly-Python-float input enters the graph. + """ + if isinstance(x, (int, float)): + return tf.constant(x, dtype=dtype) + return tf.cast(x, dtype) + + +# ============================================================================= +# Simpson on a static 1-D non-uniform grid. +# ============================================================================= + + +def simpson_weights(x): + """Return weights ``w`` such that Simpson(y, x) == sum(w * y, axis=-1). + + ``x`` is a numpy array with size ``N``. Implementation mirrors the numpy + ``simpson`` in :mod:`scetlib_btgrid_numpy` (composite Simpson with + trapezoid fallback on the last segment when N-1 is odd). + """ + x = np.asarray(x, dtype=np.float64) + n_intervals = x.size - 1 + if n_intervals < 1: + return np.zeros_like(x) + + if n_intervals % 2 == 1: + # leading n-1 intervals get Simpson, last segment gets trapezoid + w_lead = simpson_weights(x[:-1]) + w = np.concatenate([w_lead, [0.0]]) + h_last = x[-1] - x[-2] + w[-2] += 0.5 * h_last + w[-1] += 0.5 * h_last + return w + + h = np.diff(x) + h0 = h[0::2] + h1 = h[1::2] + coef = (h0 + h1) / 6.0 + w_left = coef * (2.0 - h1 / h0) + w_mid = coef * (h0 + h1) ** 2 / (h0 * h1) + w_right = coef * (2.0 - h0 / h1) + + w = np.zeros_like(x) + w[0:-1:2] += w_left + w[1::2] += w_mid + w[2::2] += w_right + return w + + +def simpson_tf(y, weights): + """Simpson reduction along the last axis using precomputed ``weights``.""" + weights = tf.cast(weights, y.dtype) + return tf.reduce_sum(y * weights, axis=-1) + + +# ============================================================================= +# F_eff and gamma_nu^NP — TF transcriptions +# ============================================================================= + +EFF_MODELS = { + "identity", + "tanh_2", + "tanh_6", + "tanh_4", + "frac_2", + "frac_4", + "exp_2", + "exp_4", + "signed_lambda", + "hyp_tangent", + "square_root", +} +GNU_MODELS = { + "tanh_1", + "tanh_2", + "tanh_6", + "frac_1", + "frac_2", + "exp_1", + "exp_2", + "hyp_tangent", + "linear", +} + + +def _safe_div(num, den): + """``num / den`` with the denominator clamped away from zero, masked by + ``tf.where`` at the call site. Keeps gradients finite.""" + den_safe = tf.where(tf.equal(den, 0), tf.ones_like(den), den) + return num / den_safe + + +def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_model): + """TF port of :func:`scetlib_btgrid_numpy.F_eff` for a fixed ``np_model``.""" + if np_model not in EFF_MODELS: + raise ValueError(f"F_eff_tf: unsupported np_model {np_model!r}") + + bT = _as_dtype(bT) + Y = _as_dtype(Y) + lambda_inf = _as_dtype(lambda_inf) + lambda2 = _as_dtype(lambda2) + lambda4 = _as_dtype(lambda4) + lambda6 = _as_dtype(lambda6) + delta_lambda2 = _as_dtype(delta_lambda2) + + if np_model == "signed_lambda": + lambda2_Y = lambda2 + delta_lambda2 * Y * Y + return (1.0 + lambda2_Y * bT**2) ** 2 * tf.exp(-2.0 * lambda4 * bT**4) + + lambda2_Y = lambda2 + delta_lambda2 * Y * Y + arg = (lambda2_Y + lambda4 * bT**2) * bT + + if np_model == "identity": + return tf.exp(-2.0 * bT * arg) + + # lambda_inf == 0 returns ones (matches numpy short-circuit). We compute + # the full formula with a safe denominator and mask at the end. + arg_inf = _safe_div(arg, lambda_inf) + model = {"hyp_tangent": "tanh_2", "square_root": "frac_2"}.get(np_model, np_model) + + if model == "tanh_2": + a = arg_inf + (1.0 / 3.0) * _safe_div(lambda2_Y * bT, lambda_inf) ** 3 + func = tf.tanh(a) + elif model == "tanh_6": + a = arg_inf + _safe_div(lambda6 * bT**5, lambda_inf) + a = a + (1.0 / 3.0) * _safe_div(lambda2_Y * bT, lambda_inf) ** 3 + func = tf.tanh(a) + elif model == "tanh_4": + func = tf.sqrt(tf.tanh(arg_inf**2)) + elif model == "frac_2": + a = arg_inf + 0.5 * _safe_div(lambda2_Y * bT, lambda_inf) ** 3 + func = a / tf.sqrt(1.0 + a**2) + elif model == "frac_4": + func = arg_inf / tf.sqrt(tf.sqrt(1.0 + arg_inf**4)) + elif model == "exp_2": + a = arg_inf + 0.25 * _safe_div(lambda2_Y * bT, lambda_inf) ** 3 + func = tf.sqrt(-tf.math.expm1(-(a**2))) + elif model == "exp_4": + func = tf.sqrt(tf.sqrt(-tf.math.expm1(-(arg_inf**4)))) + else: # pragma: no cover — guarded above + raise ValueError(f"F_eff_tf: unsupported np_model {np_model!r}") + + full = tf.exp(-2.0 * lambda_inf * bT * func) + return tf.where(tf.equal(lambda_inf, 0), tf.ones_like(full), full) + + +def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): + """TF port of :func:`scetlib_btgrid_numpy.gamma_nu_NP` for fixed ``np_model_nu``.""" + if np_model_nu not in GNU_MODELS: + raise ValueError(f"gamma_nu_NP_tf: unsupported np_model_nu {np_model_nu!r}") + + bT = _as_dtype(bT) + lambda_inf_nu = _as_dtype(lambda_inf_nu) + lambda2_nu = _as_dtype(lambda2_nu) + lambda4_nu = _as_dtype(lambda4_nu) + + bT2 = bT * bT + arg = _safe_div((lambda2_nu + lambda4_nu * bT2) * bT2, lambda_inf_nu) + + model = {"hyp_tangent": "tanh_2", "linear": "frac_1"}.get(np_model_nu, np_model_nu) + + if model == "tanh_1": + a = arg + (2.0 / 3.0) * _safe_div(lambda2_nu * bT2, lambda_inf_nu) ** 2 + func = tf.tanh(tf.sqrt(a)) ** 2 + elif model == "tanh_2": + func = tf.tanh(arg) + elif model == "tanh_6": + # NP_model_gammanu hardcodes lambda6_nu = 0.0007 (Gamma_nu.hpp:102) + a = arg + _safe_div(0.0007 * bT2**3, lambda_inf_nu) + func = tf.tanh(a) + elif model == "frac_1": + a = arg + _safe_div(lambda2_nu * bT2, lambda_inf_nu) ** 2 + func = a / (1.0 + a) + elif model == "frac_2": + func = arg / tf.sqrt(1.0 + arg**2) + elif model == "exp_1": + a = arg + 0.5 * _safe_div(lambda2_nu * bT2, lambda_inf_nu) ** 2 + func = -tf.math.expm1(-a) + elif model == "exp_2": + func = tf.sqrt(-tf.math.expm1(-(arg**2))) + else: # pragma: no cover + raise ValueError(f"gamma_nu_NP_tf: unsupported np_model_nu {np_model_nu!r}") + + full = -lambda_inf_nu * func + return tf.where(tf.equal(lambda_inf_nu, 0), tf.zeros_like(full), full) + + +# ============================================================================= +# Hankel reconstruction (batched, λ-differentiable) +# ============================================================================= + + +def reconstruct_batch_tf( + qT_per_bin, + bT, + I_pert, + C_nu, + b_bar, + Y_per_bin, + eff_params: Mapping, + gnu_params: Mapping, + *, + np_model: str, + np_model_nu: str, + bT_simpson_weights=None, + bT_J0_kernel=None, + Y_unique=None, + Y_inverse_idx=None, +): + """TF port of :func:`scetlib_btgrid_numpy.reconstruct_batch`. + + All array-shape arguments are TF tensors or numpy arrays (will be cast). + The λ values inside ``eff_params`` / ``gnu_params`` are the differentiable + parameters — pass them as TF scalars (Variables or constants). + + ``bT_simpson_weights`` and ``bT_J0_kernel`` are optional precomputed + constants. Pass them in the ParamModel to avoid recomputing per-step. + + ``Y_unique`` / ``Y_inverse_idx`` are an optional precomputed unique-Y map + (``Y_unique`` = sorted distinct Y values, shape ``(NY,)``; ``Y_inverse_idx`` + = per-bin index into ``Y_unique``, shape ``(Nbins,)``). ``F_eff`` depends on + the bin only through Y, so when this map is supplied the NP transcendentals + are evaluated on the ``NY`` unique rows and gathered back to ``(Nbins, Nbt)`` + — bit-for-bit identical to the per-bin evaluation, but the expensive ops and + their λ-gradients run on ``NY`` rows instead of ``Nbins`` (Q and qT don't + enter ``F_eff``). Without the map it falls back to the full per-bin path. + """ + qT_per_bin = _as_dtype(qT_per_bin) # (Nbins,) + Y_per_bin = _as_dtype(Y_per_bin) # (Nbins,) + bT = _as_dtype(bT) # (Nbt,) + b_bar = _as_dtype(b_bar) # (Nbt,) + I_pert = _as_dtype(I_pert) # (Nbins, Nbt) + C_nu = _as_dtype(C_nu) # (Nbins, Nbt) + + if bT_J0_kernel is None: + bT_J0_kernel = build_bT_J0_kernel(qT_per_bin, bT) + bT_J0_kernel = _as_dtype(bT_J0_kernel) # (Nbins, Nbt) + + if bT_simpson_weights is None: + # bT is a tf.Tensor here; convert to numpy for the Python-side weights + bT_simpson_weights = simpson_weights(np.asarray(bT)) + bT_simpson_weights = _as_dtype(bT_simpson_weights) # (Nbt,) + + g_NP = gamma_nu_NP_tf(b_bar, **gnu_params, np_model_nu=np_model_nu) # (Nbt,) + exp_g_factor = tf.exp(C_nu * g_NP[tf.newaxis, :]) # (Nbins, Nbt) + + delta_l2_in = eff_params.get("delta_lambda2", 0.0) + if isinstance(delta_l2_in, (int, float)) and float(delta_l2_in) == 0.0: + # static fast path: F_eff has no Y dependence + Feff = F_eff_tf(0.0, b_bar, **eff_params, np_model=np_model) # (Nbt,) + Feff_b = Feff[tf.newaxis, :] + elif Y_unique is not None and Y_inverse_idx is not None: + # per-bin F_eff, but evaluated on the unique Y rows then gathered. + # Exact: identical Y -> identical F_eff row for any λ; the gather only + # replicates rows (its backward scatter-adds the cotangents, so λ-grads + # are unchanged). Transcendentals run on (NY, Nbt), not (Nbins, Nbt). + Y_u = _as_dtype(Y_unique)[:, tf.newaxis] # (NY, 1) + b_b = b_bar[tf.newaxis, :] + Feff_u = F_eff_tf(Y_u, b_b, **eff_params, np_model=np_model) # (NY, Nbt) + Feff_b = tf.gather(Feff_u, Y_inverse_idx) # (Nbins, Nbt) + else: + # per-bin F_eff due to Y dependence (delta_lambda2 * Y^2) + # build (Nbins, Nbt) by broadcasting Y_per_bin over bT + Y_b = Y_per_bin[:, tf.newaxis] + b_b = b_bar[tf.newaxis, :] + Feff_b = F_eff_tf(Y_b, b_b, **eff_params, np_model=np_model) # (Nbins, Nbt) + + integrand = bT_J0_kernel * I_pert * exp_g_factor * Feff_b # (Nbins, Nbt) + sigma = simpson_tf(integrand, bT_simpson_weights) # (Nbins,) + + # qT factor from SCETlib's x = qT*bT integration convention. + return qT_per_bin * sigma + + +def build_bT_J0_kernel(qT_per_bin, bT): + """Precompute ``bT * J_0(qT*bT)`` on the (Nbins, Nbt) grid. + + λ-independent — call once at ParamModel construction and pass into + :func:`reconstruct_batch_tf` as ``bT_J0_kernel``. + """ + qT_per_bin = _as_dtype(qT_per_bin) + bT = _as_dtype(bT) + arg = qT_per_bin[:, tf.newaxis] * bT[tf.newaxis, :] + return bT[tf.newaxis, :] * tf.math.special.bessel_j0(arg) diff --git a/wremnants/postprocessing/scetlib_np/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py new file mode 100644 index 000000000..d094eeec4 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -0,0 +1,197 @@ +"""λ_central auto-detect from a fit-input hdf5. + +The SCETlib correction's NP runcard is preserved in the upstream +``*_Corr.pkl.lz4`` file (under ``file_meta_data..config.Nonperturbative``), +but it is **not** propagated through the histmaker into the fit-input hdf5. We +read the correction tag from the hdf5 (``meta_info_input.args.theoryCorr``), +resolve to the upstream pkl, and extract the Nonperturbative section. + +Single entry point: + + read_lambda_central(hdf5_path, proc="Z") -> dict +""" + +import os +import pickle +import sys + +import h5py +import lz4.frame + +from wremnants.utilities import common as wrem_common +from wums import ioutils as wums_io + +# Names of the parameters the ParamModel cares about, split by which scetlib +# C++ struct consumes them. Values in the Nonperturbative section are strings; +# numeric ones get parsed to float, model names stay as strings. +GNU_NUMERIC = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") +GNU_STRING = ("np_model_nu",) +EFF_NUMERIC = ("lambda_inf", "lambda2", "lambda4", "lambda6", "delta_lambda2") +EFF_STRING = ("np_model",) + + +def _correction_pkl_path(tag, proc): + """Resolve a theoryCorr tag to its upstream pkl.lz4 path.""" + return os.path.join( + wrem_common.data_dir, "TheoryCorrections", f"{tag}_Corr{proc}.pkl.lz4" + ) + + +def _load_correction_pkl(tag, proc): + path = _correction_pkl_path(tag, proc) + if not os.path.exists(path): + raise FileNotFoundError( + f"SCETlib correction pkl not found: {path!r}. " + f"Cannot extract λ_central for theoryCorr={tag!r}." + ) + with lz4.frame.open(path, "rb") as f: + return pickle.load(f) + + +def _find_nonperturbative(corr_dict): + """Search the upstream correction dict for ``Nonperturbative`` configs. + + Returns a list of (basename, Nonperturbative dict) tuples. There are + typically multiple basenames (resummed-singular, fixed-order, etc.) — all + are expected to share the same Nonperturbative section since SCETlib runs + them with one runcard. + """ + out = [] + meta = corr_dict.get("file_meta_data") + if not isinstance(meta, dict): + raise KeyError( + "Correction pkl has no 'file_meta_data' entry — schema mismatch." + ) + for basename, file_meta in meta.items(): + if not isinstance(file_meta, dict): + continue + cfg = file_meta.get("config") + if not isinstance(cfg, dict): + continue + npert = cfg.get("Nonperturbative") + if isinstance(npert, dict): + out.append((basename, npert)) + if not out: + raise KeyError( + "No Nonperturbative section found in any basename of " + "file_meta_data — λ_central undefined." + ) + return out + + +def _parse_section(npert): + """Parse one Nonperturbative dict into the two parameter groups. + + Numeric params absent from the runcard default to 0 — SCETlib runcards + only set the keys relevant to the chosen np_model. e.g. tanh_2 setups + typically omit ``lambda6`` (the bT⁵ coefficient only used by tanh_6). + """ + eff_params = {"np_model": npert[EFF_STRING[0]]} + gnu_params = {"np_model_nu": npert[GNU_STRING[0]]} + for k in EFF_NUMERIC: + eff_params[k] = float(npert.get(k, 0.0)) + for k in GNU_NUMERIC: + gnu_params[k] = float(npert.get(k, 0.0)) + return eff_params, gnu_params + + +def read_lambda_central_from_meta(meta, proc="Z", _source=""): + """Same as :func:`read_lambda_central`, but takes the already-loaded + metadata dict (e.g. ``indata.metadata``) instead of an hdf5 path. + + Avoids re-opening the input HDF5 when the caller already has the meta + in hand. ``_source`` is used only in error messages. + """ + try: + theory_corr = meta["meta_info_input"]["args"]["theoryCorr"] + except (KeyError, TypeError) as exc: + raise KeyError( + f"{_source}: meta_info_input.args.theoryCorr missing — " + "cannot identify the central SCETlib correction." + ) from exc + + if not theory_corr: + raise ValueError(f"{_source}: theoryCorr list is empty.") + tag = theory_corr[0] # first entry = central; rest are pdfvars/pdfas + + return _resolve_tag_to_lambda(tag, proc) + + +def read_lambda_central(hdf5_path, proc="Z"): + """Extract λ_central from the SCETlib correction referenced by the hdf5. + + Returns a dict with: + + tag : the central theoryCorr tag (first entry of meta args) + pkl_path : path to the upstream correction pkl + basename : the file_meta basename whose runcard was used + eff_params : dict for NP_model_effective (F_eff). Keys: + np_model, lambda_inf, lambda2, lambda4, lambda6, + delta_lambda2. + gnu_params : dict for NP_model_gammanu (γ_ν^NP). Keys: + np_model_nu, lambda_inf_nu, lambda2_nu, lambda4_nu. + + Raises with a clear message if any link in the chain is missing. + """ + with h5py.File(hdf5_path, "r") as f: + if "meta" not in f: + raise KeyError(f"{hdf5_path}: no 'meta' group — wrong file type?") + meta = wums_io.pickle_load_h5py(f["meta"]) + return read_lambda_central_from_meta(meta, proc=proc, _source=hdf5_path) + + +def _resolve_tag_to_lambda(tag, proc): + """Internal: given a theoryCorr tag, load the pkl and parse λ_central.""" + + corr_dict = _load_correction_pkl(tag, proc) + sections = _find_nonperturbative(corr_dict) + + # Some correction pkls bundle multiple NP variants in one file (e.g. a + # "lattice" central + a "FranksVals" variant). We need to pick the + # basename that matches the analysis's NP tag. Heuristic: look for a + # substring in the basename that also appears in the tag (case-insensitive). + tag_lower = tag.lower() + KEYWORDS = ("franksvals", "lattice", "newvars", "lambda6") + matched_kw = next((k for k in KEYWORDS if k in tag_lower), None) + + def _basename_score(name): + name_lower = name.lower() + score = 0 + if matched_kw and matched_kw in name_lower: + score += 10 # strong preference: matches the analysis variant + if ( + "nnlo_sing" in name_lower + or "_sing_" in name_lower + or name_lower.endswith("sing.pkl") + ): + score += 1 # weak preference: resummed-singular carries the full NP set + return score + + sections_sorted = sorted(sections, key=lambda item: -_basename_score(item[0])) + basename, npert = sections_sorted[0] + eff_params, gnu_params = _parse_section(npert) + + return dict( + tag=tag, + pkl_path=_correction_pkl_path(tag, proc), + basename=basename, + eff_params=eff_params, + gnu_params=gnu_params, + ) + + +if __name__ == "__main__": + import sys + + if len(sys.argv) < 2: + print( + f"usage: python -m {__name__.replace('.', '/')} [proc]", + file=sys.stderr, + ) + sys.exit(2) + out = read_lambda_central(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else "Z") + print(f"tag : {out['tag']}") + print(f"pkl_path : {out['pkl_path']}") + print(f"basename : {out['basename']}") + print(f"eff_params: {out['eff_params']}") + print(f"gnu_params: {out['gnu_params']}") diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py new file mode 100644 index 000000000..a2b98ed0e --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -0,0 +1,872 @@ +"""SCETlibNPParamModel — continuous-λ rabbit ParamModel for SCETlib NP. + +Architecture (template-style fit): + + σ_reco(λ; b) = Σ_g R(b, g) · σ_gen(λ; g) + ratio(b) = σ_reco(λ; b) / σ_reco(λ_central; b) + rnorm = ratio per reco bin, broadcast over signal proc; ones elsewhere. + +R is the (reco × gen) response matrix loaded from the upstream unfolding +histmaker output (a separate hdf5 from the fit-tensor input). + +λ_central is read from the fit-tensor's meta_info_input via the upstream +SCETlib correction pkl (see :mod:`scetlib_lambda_central`). + +σ_gen(λ; g) is evaluated on the btgrid then integrated over Q (arctan_Q² +Simpson) and rebinned (Simpson) onto the unfolding hist's gen edges +(ptVGen, absYVGen). The absYVGen-side rebin folds the signed btgrid Y axis +into |Y| bins (NP is Y-symmetric: F_eff depends on Y², γ_ν^NP doesn't depend +on Y at all). + +The 8 v1 parameters (all factorisable through the current btgrid): + + γ_ν^NP (CS-side): lambda2_nu, lambda4_nu, lambda_inf_nu + F_eff (TMD-effective): lambda2, lambda4, lambda6, delta_lambda2, lambda_inf + +The np_model and np_model_nu strings are fixed at construction (from +λ_central). All λ values are TF Variables — differentiable in the fit. +""" + +from typing import Mapping, Optional + +import numpy as np +import tensorflow as tf + +from rabbit.param_models.param_model import ParamModel +from wremnants.postprocessing.scetlib_np import ( + btgrid_cache, +) +from wremnants.postprocessing.scetlib_np import btgrid_integrate as fz_int +from wremnants.postprocessing.scetlib_np import btgrid_tf as fz_tf +from wremnants.postprocessing.scetlib_np import lambda_central as scetlib_lambda_central +from wremnants.postprocessing.scetlib_np import response_matrix as fz_R + +# Ordered list of the v1 continuous λ. CS-side first, then TMD-effective. +GNU_PARAMS = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") +EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") +ALL_PARAMS = GNU_PARAMS + EFF_PARAMS + +# Physical lower bounds used by the Taylor surrogate so the finite-difference +# step (λ_central − h) can't cross into a region where F_eff / γ_ν^NP blow up. +# - lambda6 enters as `lambda6 · bT⁵ / lambda_inf` inside a tanh; if lambda6 is +# negative the tanh saturates to −1 at large bT and the outer exp inverts +# sign → exp(+huge). So lambda6 ≥ 0 is enforced. +# - lambda_inf / lambda_inf_nu sit in denominators (and in `exp(−2·lambda_inf·bT)`); +# require strictly positive with a safety margin. +PARAM_MIN_VALUE = { + "lambda2_nu": None, + "lambda4_nu": None, + "lambda_inf_nu": 0.05, + "lambda2": None, + "lambda4": None, + "lambda6": 0.0, + "delta_lambda2": None, + "lambda_inf": 0.05, +} + + +# Theorist-recommended Gaussian prior widths for the SCETlib NP λ parameters. +# Source: NP-NP discussion slide, central values 2026-05, plus a wide +# in-house default for delta_lambda2 (the theorist hasn't quoted a width +# for it; 0 ± 0.2 is comfortably wider than its expected scale). +# +# λ₂^ν = 0.15 ± 0.10 +# Λ₂ = 0.40 ⁺⁰·⁶₋₀.₄ (asymmetric) +# Λ₄ = 0.40 ⁺⁰·⁶₋₀.₄ (asymmetric) +# δ Λ₂ = 0.00 ± 0.20 (in-house wide default; not from the slide) +# +# Asymmetric uncertainties (Λ₂, Λ₄) are approximated by a symmetric Gaussian +# with σ = (σ⁺ + σ⁻) / 2. Slightly conservative on the upper side, slightly +# loose on the lower side. A future patch can implement a split-Gaussian if +# needed. +# +# Any λ not listed here gets σ = NaN by default → no prior, floats free. +# Currently those are: lambda_inf, lambda_inf_nu, lambda4_nu, lambda6. +# They are expected to be FROZEN via rabbit's --freezeParameters until +# the theorist provides priors for them. +THEORIST_PRIOR_SIGMAS = { + "lambda2_nu": 0.10, + "lambda2": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average + "lambda4": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average + "delta_lambda2": 0.20, # 0 ± 0.20 wide default (no theorist value yet) +} + + +def _crop_R_to_fit(R, R_reco_axes, fit_reco_axes, tol=1e-9): + """Crop R's trailing reco bins so its reco shape matches the fit. + + R's reco binning is typically a superset of the fit's (e.g. R has one + extra overflow ptll bin past the fit's last edge). For each reco axis, + require that R's leading edges match the fit's edges; crop R along that + axis to keep only the matching bins. + """ + 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)}" + ) + 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}") + 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." + ) + if not np.allclose(redges[:fnb], fedges, atol=tol): + raise ValueError( + f"Reco axis {rname}: leading R edges don't match fit edges. " + f"R[:{fnb}]={list(redges[:fnb])} vs fit={list(fedges)}" + ) + slices = tuple(slice(0, len(fedges) - 1) for (_, fedges) in fit_reco_axes) + # Keep all gen axes (the remaining axes of R). + slices += (slice(None),) * (R.ndim - len(fit_reco_axes)) + return R[slices] + + +class SCETlibNPParamModel(ParamModel): + + def __init__( + self, + indata, + unfolding_hdf5_path: str, + btgrid_dir: str, + lambda_central: Optional[Mapping] = None, + signal_proc: str = "Zmumu", + Q_lo: float = 60.0, + Q_hi: float = 120.0, + poi_params: Optional[tuple] = (), + prior_sigmas: Optional[Mapping] = None, + use_taylor: bool = False, + taylor_order: int = 2, + taylor_h_rel: float = 0.10, + taylor_h_min: float = 0.005, + **kwargs, + ): + """Construct the ParamModel. + + Parameters + ---------- + indata + rabbit's input-data structure (passed by ``ph.load_models``). + unfolding_hdf5_path + Path to the upstream histmaker output containing + ``nominal_postfsr_yieldsUnfolding`` for R. + btgrid_dir + Directory of the SCETlib bT-grid shards (fineall). + lambda_central + Dict with two sub-dicts ``eff_params`` and ``gnu_params`` (same + shape as returned by :func:`scetlib_lambda_central.read_lambda_central`). + By default λ_central is auto-detected from ``indata.metadata``; + pass this only to override or to support hand-built indata that + lacks metadata. + signal_proc + Name of the signal process whose reco yields get the per-bin + ratio. Other processes get factor 1. + Q_lo, Q_hi + Z mass window for the Q-integration on the btgrid. + poi_params + Tuple of parameter names (subset of ``ALL_PARAMS``) to treat as + POIs (reported as POIs in the fit output). The rest are reported + as model nuisances (npou). The POI vs POU split is independent + of the prior assignment (see ``prior_sigmas``). + prior_sigmas + Per-name override dict for the Gaussian prior σ on each + parameter. Defaults come from ``THEORIST_PRIOR_SIGMAS``: + + lambda2_nu : 0.10 + lambda2 : 0.50 (symmetric approx of +0.6/-0.4) + lambda4 : 0.50 (symmetric approx of +0.6/-0.4) + + All other params default to ``NaN`` → no prior, float free; in + practice they are expected to be frozen with rabbit's + ``--freezeParameters`` until the theorist provides priors for + them. Pass ``np.nan`` here to free a constrained param, or a + finite value to add a prior on one that defaults to NaN. + Only consumed when the fitter is invoked with + ``--paramModelPriors``; otherwise everything floats free. + Prior mean for each param is ``self.xparamdefault`` (the + runcard's λ_central). + """ + self.indata = indata + + # ---- Double-counting guard + # If the histmaker baked discrete NP κ-template variations into the + # input HDF5, those systs describe the same physics as our continuous + # λ POUs. Running both → double-counting (the discrete syst absorbs + # whatever shape variation our ParamModel should describe). Warn + # loudly if any such systs are present and unfrozen. + self._check_discrete_np_double_counting(kwargs.get("freezeParameters")) + + # ---- λ_central + # Three sources of λ_central, in priority order: + # 1. ``lambda_central`` constructor arg (explicit dict). + # 2. ``SCETLIB_NP_LAMBDA_CENTRAL_JSON`` env var (JSON-encoded dict + # with ``eff_params`` and ``gnu_params``). Useful when the upstream + # SCETlib pkl isn't accessible (e.g. colleague's input). + # 3. Auto-detect from the fit hdf5's theoryCorr → upstream pkl. + import json + import os + + env_lc = os.environ.get("SCETLIB_NP_LAMBDA_CENTRAL_JSON", "").strip() + if lambda_central is None and env_lc: + try: + lambda_central = json.loads(env_lc) + except json.JSONDecodeError as exc: + raise ValueError( + f"SCETLIB_NP_LAMBDA_CENTRAL_JSON must be valid JSON; got {exc}" + ) + print(f"[SCETlibNPParamModel] λ_central from env var", flush=True) + if lambda_central is None: + # Auto-detect from indata.metadata (loaded by rabbit's + # FitInputData from the input HDF5's "meta" group). + indata_meta = getattr(indata, "metadata", None) or {} + if not indata_meta: + raise ValueError( + "SCETlibNPParamModel: indata has no metadata; pass " + "lambda_central explicitly or use an indata that " + "carries metadata." + ) + lambda_central = scetlib_lambda_central.read_lambda_central_from_meta( + indata_meta, _source="indata.metadata" + ) + self.eff_central = dict(lambda_central["eff_params"]) + self.gnu_central = dict(lambda_central["gnu_params"]) + self.np_model = self.eff_central["np_model"] + self.np_model_nu = self.gnu_central["np_model_nu"] + + # ---- btgrid + dense layout + grid = btgrid_cache.load(btgrid_dir) + self._btgrid_meta = dict( + shards=grid["n_shards"], + n_bins=len(grid["bins"]), + ) + idx_map = fz_int.dense_index_map(grid["bins"]) + self.Q_unique = idx_map["Q_unique"] + self.Y_unique = idx_map["Y_unique"] + self.qT_unique = idx_map["qT_unique"] + self.flat_idx = tf.constant(idx_map["flat_idx"], dtype=tf.int64) + + # Cache btgrid arrays as TF constants. + self.bT = tf.constant(grid["bT"], dtype=fz_tf.DTYPE) + self.b_bar = tf.constant(grid["b_bar"], dtype=fz_tf.DTYPE) + self.I_pert = tf.constant(grid["I_pert"][0], dtype=fz_tf.DTYPE) # (Nbins, Nbt) + self.C_nu = tf.constant(grid["C_nu"][0], dtype=fz_tf.DTYPE) + + # Per-bin qT and Y (from the bin tuple), for reconstruct_batch_tf. + bins = grid["bins"] + self.qT_per_bin = tf.constant( + np.array([b[2] for b in bins], dtype=np.float64), dtype=fz_tf.DTYPE + ) + Y_pb_np = np.array([b[1] for b in bins], dtype=np.float64) + self.Y_per_bin = tf.constant(Y_pb_np, dtype=fz_tf.DTYPE) + + # F_eff depends on the bin only through Y (not Q or qT), and Y takes few + # distinct values across the grid. Precompute the unique-Y map so + # reconstruct_batch_tf evaluates the NP transcendentals on NY rows and + # gathers, instead of recomputing identical rows for every (Q, qT). + Y_feff_unique_np, Y_feff_inv_np = np.unique(Y_pb_np, return_inverse=True) + self.Y_feff_unique = tf.constant(Y_feff_unique_np, dtype=fz_tf.DTYPE) + self.Y_feff_inverse_idx = tf.constant( + Y_feff_inv_np.reshape(-1).astype(np.int32), dtype=tf.int32 + ) + + # Precompute the bT·J0(qT·bT) kernel (λ-independent). + self.bT_J0_kernel = fz_tf.build_bT_J0_kernel(self.qT_per_bin, self.bT) + self.bT_simpson_w = tf.constant( + fz_tf.simpson_weights(np.asarray(self.bT)), dtype=fz_tf.DTYPE + ) + + # ---- Q-integration weights (arctan_Q² Simpson on Z mass window). + self.Q_weights = tf.constant( + fz_int.q_integrate_weights(self.Q_unique, Q_lo, Q_hi), + dtype=fz_tf.DTYPE, + ) + + # ---- R matrix + R_info = fz_R.load_R(unfolding_hdf5_path) + # The fit-tensor's reco binning may differ from R's by trailing + # overflow bins (e.g. R has ptll [0, …, 44, 100] while the fit ends + # at 44). Crop R's trailing bins so the reco shape matches. + fit_reco_axes = self._fit_reco_axes(indata) + R_arr = _crop_R_to_fit(R_info["R"], R_info["reco_axes"], fit_reco_axes) + # Tighten the metadata to match the cropped R. + self.reco_shape = R_arr.shape[: len(fit_reco_axes)] + self.gen_shape = R_arr.shape[len(fit_reco_axes) :] + N_reco = int(np.prod(self.reco_shape)) + N_gen = int(np.prod(self.gen_shape)) + self.R = tf.constant(R_arr.reshape(N_reco, N_gen), dtype=fz_tf.DTYPE) + self._reco_axes_meta = [ + (name, fit_axes[1]) + for (name, fit_axes) in zip( + [a[0] for a in R_info["reco_axes"]], + fit_reco_axes, + ) + ] + self._gen_axes_meta = R_info["gen_axes"] + + # ---- Rebin weights: btgrid (NY signed) → (NabsYVGen) via |Y| folding + # and (NqT) → (NptVGen). + absY_edges = self._gen_axes_meta[1][1] # absYVGen edges + ptVGen_edges = self._gen_axes_meta[0][1] # ptVGen edges + + # |Y| folding: σ(Y) is symmetric in Y so the absY-bin integral is + # 2·∫_{absY_lo}^{absY_hi} σ(Y) dY. Use Y >= 0 source samples and + # multiply by 2. + Y_pos_mask = self.Y_unique >= 0 + Y_pos = self.Y_unique[Y_pos_mask] + absY_rebin_pos = fz_int.rebin_weights(Y_pos, absY_edges, name="absY") + # Pad to full NY: zero on negative-Y columns. + W_absY = np.zeros((absY_edges.size - 1, self.Y_unique.size), dtype=np.float64) + W_absY[:, Y_pos_mask] = 2.0 * absY_rebin_pos + self.W_absY = tf.constant(W_absY, dtype=fz_tf.DTYPE) + + # qT rebin: btgrid qT (signed nonneg, NqT=141) → ptVGen edges (e.g. 20 bins, 0-44). + # Anything past ptVGen_max is out of fit range; we drop it silently for now + # (events with gen qT > ptVGen_max are routed through R's overflow which is + # not present in our materialised R — see plan doc, dyturbo-handoff item). + self.W_ptVGen = tf.constant( + fz_int.rebin_weights(self.qT_unique, ptVGen_edges, name="ptVGen"), + dtype=fz_tf.DTYPE, + ) + + # ---- Cache σ_reco(λ_central): the denominator of the ratio. + sigma_gen_central = self._sigma_gen_at(self.eff_central, self.gnu_central) + # σ_reco(λ_central) = R · σ_gen(λ_central). Flatten the gen axes. + gen_flat = tf.reshape(sigma_gen_central, [-1]) + self.sigma_reco_central = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) + # Sanity floor — if any reco bin has zero or negative central yield, + # the ratio would blow up. Use the central yield directly; any genuine + # zero is a binning issue to flag. + if tf.reduce_any(self.sigma_reco_central <= 0).numpy(): + n_bad = int(tf.reduce_sum(tf.cast(self.sigma_reco_central <= 0, tf.int32))) + raise ValueError( + f"SCETlibNPParamModel: {n_bad} reco bins have non-positive " + f"σ_reco(λ_central). Likely a binning mismatch between R and " + f"the fit-tensor reco axes." + ) + + # ---- Process index: signal column gets the ratio, others get 1. + procs = [p.decode() if isinstance(p, bytes) else str(p) for p in indata.procs] + if signal_proc not in procs: + raise ValueError( + f"SCETlibNPParamModel: signal_proc={signal_proc!r} not in " + f"indata.procs={procs[:10]}..." + ) + self.signal_proc_idx = procs.index(signal_proc) + self.nproc = indata.nproc + + # ---- ParamModel registration (POIs first, then NOUs). + poi_params = tuple(poi_params or ()) + nou_params = tuple(p for p in ALL_PARAMS if p not in poi_params) + self._param_order = poi_params + nou_params + self.npoi = len(poi_params) + self.npou = len(nou_params) + self.params = np.array([p.encode() for p in self._param_order]) + + # Defaults: λ_central values per parameter. Optionally overridden by + # the ``SCETLIB_NP_XPARAMDEFAULT`` env var — comma-separated + # ``name=value`` pairs (for closure tests where the data-generating + # / fit-start point should differ from the card's λ_central). + central_lookup = {**self.eff_central, **self.gnu_central} + defaults = np.array( + [central_lookup[p] for p in self._param_order], dtype=np.float64 + ) + import os + + env_override = os.environ.get("SCETLIB_NP_XPARAMDEFAULT", "").strip() + if env_override: + overrides = dict( + tuple(s.split("=")) for s in env_override.split(",") if s.strip() + ) + for name, val in overrides.items(): + name = name.strip() + if name not in self._param_order: + raise KeyError(f"SCETLIB_NP_XPARAMDEFAULT: unknown param {name!r}") + i = self._param_order.index(name) + defaults[i] = float(val) + print( + f"[SCETlibNPParamModel] xparamdefault overridden: {dict(zip(self._param_order, defaults))}" + ) + # rabbit's set_param_default expects an internal-storage convention + # where POIs (npoi entries) are SQRT(value) if not allowNegativeParam. + # For our λ which can in principle be tiny / zero (delta_lambda2), + # default to allowNegativeParam=True so the stored value == λ directly. + self.allowNegativeParam = True + self.is_linear = False + self.xparamdefault = tf.constant(defaults, dtype=indata.dtype) + + # Gaussian priors (consumed by rabbit's Fitter when --paramModelPriors + # is set; ignored otherwise). Default σ values come from + # ``THEORIST_PRIOR_SIGMAS`` (lambda2_nu, lambda2, lambda4 — the only + # params the theorist provides widths for as of 2026-05). + # All other params default to σ = NaN → no prior, float free; in + # practice those should be frozen via rabbit's --freezeParameters + # until the theorist gives priors for them. + # The ``prior_sigmas`` kwarg is a per-name override dict; pass NaN to + # force a parameter free, or a finite value to add / change a prior. + # The mean of each prior is λ_central (i.e. self.xparamdefault). + prior_sigmas = dict(prior_sigmas or {}) + sigmas_arr = np.empty(self.nparams, dtype=np.float64) + for i, p in enumerate(self._param_order): + if p in prior_sigmas: + sigmas_arr[i] = float(prior_sigmas[p]) # explicit override (may be NaN) + elif p in THEORIST_PRIOR_SIGMAS: + sigmas_arr[i] = THEORIST_PRIOR_SIGMAS[p] # theorist recommendation + else: + sigmas_arr[i] = np.nan # free (expected to be frozen) + self.prior_sigmas = sigmas_arr + # prior_means defaults to xparamdefault if not set, so don't store + # redundantly — Fitter will fall back to xparamdefault. + + # ---- Taylor surrogate (optional, default OFF) + # Precompute σ_gen(λ_central) plus first/second derivatives so that + # per-fit-step compute() is a polynomial in (λ − λ_central) instead of + # a full Hankel/Simpson integral. Drops per-step cost from O(10s) to + # O(ms). Accuracy: validated against the full integral; expected + # sub-percent within typical NP variation ranges for quadratic order. + # Off by default — opt in via ``use_taylor=True`` kwarg (or env var + # ``SCETLIB_NP_USE_TAYLOR=1``). Env var ``SCETLIB_NP_USE_TAYLOR=0`` + # also forces it off (useful when the kwarg is set by a CLI wrapper). + env_taylor = os.environ.get("SCETLIB_NP_USE_TAYLOR", "").strip().lower() + if env_taylor in ("0", "false", "no", "off"): + use_taylor = False + print( + "[SCETlibNPParamModel] Taylor surrogate disabled by env var", flush=True + ) + elif env_taylor in ("1", "true", "yes", "on"): + use_taylor = True + print( + "[SCETlibNPParamModel] Taylor surrogate enabled by env var", flush=True + ) + self.use_taylor = use_taylor + self.taylor_order = int(taylor_order) + if self.use_taylor: + print( + f"[SCETlibNPParamModel] Taylor surrogate ON " + f"(order={int(taylor_order)}, h_rel={taylor_h_rel}, " + f"h_min={taylor_h_min}) — per-step compute() is a polynomial " + f"in (λ − λ_central), not the full Hankel integral", + flush=True, + ) + self._build_taylor_surrogate( + h_rel=taylor_h_rel, h_min=taylor_h_min, order=self.taylor_order + ) + else: + print( + "[SCETlibNPParamModel] Taylor surrogate OFF — per-step " + "compute() runs the full Hankel/Simpson integral", + flush=True, + ) + + # ========================================================================= + # Helpers + # ========================================================================= + + # Substrings (case-insensitive) that mark indata.systs as discrete + # NP-template variations of one of our 8 continuous λ parameters. + # Matching is case-insensitive because the histmaker uses inconsistent + # casing (canonical names are uppercase Lambda, but some configurations + # serialize them lowercase). + # + # Canonical names (see theory_variation_labels.py): + # chargeVgenNP0scetlibNPZLambda2 → catches "scetlibnpzlambda" + # chargeVgenNP0scetlibNPZLambda4 → catches "scetlibnpzlambda" + # chargeVgenNP0scetlibNPZDelta_Lambda2 → catches "scetlibnpzdelta" + # chargeVgenNP0scetlibNPLambda2 (W-side) → catches "scetlibnplambda" + # chargeVgenNP0scetlibNPLambda4 (W-side) → catches "scetlibnplambda" + # chargeVgenNP0scetlibNPDelta_Lambda2 → catches "scetlibnpdelta" + # scetlibNPgamma → catches "scetlibnpgamma" + # scetlibNPgammaEigvar{1,2,3} → catches "scetlibnpgamma" + # scetlibNPgammaLambda{2,4,Inf} → catches "scetlibnpgamma" + _DISCRETE_NP_PATTERNS = ( + "scetlibnpzlambda", # Z-side Lambda2 / Lambda4 templates + "scetlibnpzdelta", # Z-side Delta_Lambda2 template + "scetlibnplambda", # W-side Lambda2 / Lambda4 templates + "scetlibnpdelta", # W-side Delta_Lambda2 template + "scetlibnpgamma", # all γ_ν templates (Lambda2/4/Inf, Eigvar1/2/3, "gamma") + ) + + def _check_discrete_np_double_counting(self, freeze_patterns): + """Detect indata systs that overlap with our continuous λ POUs. + + Three outcomes: + - The discrete NP syst is **absent** from ``indata.systs`` entirely + (histmaker didn't include it): nothing to do, silent return. + - The discrete NP syst is **present and matched** by + ``freeze_patterns``: it's frozen at central → no double-counting, + silent return. + - The discrete NP syst is **present and unfrozen**: it overlaps + with one of our continuous λ POUs and will absorb shape variation + the ParamModel should describe → print a loud banner with the + exact freeze args to add. + """ + import re + + systs = getattr(self.indata, "systs", None) + if systs is None or len(systs) == 0: + return + syst_names = [s.decode() if isinstance(s, bytes) else str(s) for s in systs] + + # Case-insensitive substring match: canonical names use uppercase + # Lambda but some histmaker outputs lowercase the names. + conflicting = [ + s + for s in syst_names + if any(pat in s.lower() for pat in self._DISCRETE_NP_PATTERNS) + ] + if not conflicting: + return # not in indata.systs at all → nothing to warn about + + # Which of those are NOT already covered by a user-supplied freeze + # pattern (exact match or anchored regex)? + patterns = list(freeze_patterns or []) + unfrozen = [] + for s in conflicting: + covered = False + for pat in patterns: + if pat == s: + covered = True + break + try: + if re.fullmatch(pat, s): + covered = True + break + except re.error: + continue + if not covered: + unfrozen.append(s) + + if not unfrozen: + return # all conflicting systs are already frozen by the user + + print( + "\n" + "===================================================================\n" + "[SCETlibNPParamModel] DOUBLE-COUNTING WARNING\n" + "===================================================================\n" + f"Detected {len(unfrozen)} discrete NP κ-template syst(s) in the\n" + "input HDF5 that describe the same physics as this ParamModel's\n" + "continuous λ parameters. Running both leads to double-counting:\n" + "the discrete syst absorbs shape variation that the ParamModel\n" + "should describe (the indata syst will show a spurious pull, and\n" + "the postfit λ values are not what the data actually prefers).\n\n" + "Unfrozen conflicting systs:\n" + + "\n".join(f" {s}" for s in unfrozen) + + "\n\n" + "Fix by adding to --freezeParameters, e.g.:\n" + " --freezeParameters '.*scetlibNPZ.*lambda.*' " + "'.*scetlibNPgammaLambda.*' ...\n" + "or list them explicitly:\n" + " --freezeParameters " + " ".join(repr(s) for s in unfrozen) + "\n" + "===================================================================", + flush=True, + ) + + def _fit_reco_axes(self, indata): + """Read the (name, edges) of each reco axis from the (single) channel. + + Multi-channel support is deferred to v2; this raises if there's >1 + 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"SCETlibNPParamModel v1 supports a single non-masked channel; " + f"got {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"] + ] + + # ========================================================================= + # σ_gen evaluation + # ========================================================================= + + def _sigma_gen_at(self, eff_params, gnu_params): + """Evaluate σ_gen(λ) on R's gen binning. Returns shape (NptVGen, NabsYVGen).""" + # 1. Reconstruct σ on the btgrid's flat (Nbins,) layout. + sigma_flat = fz_tf.reconstruct_batch_tf( + qT_per_bin=self.qT_per_bin, + bT=self.bT, + I_pert=self.I_pert, + C_nu=self.C_nu, + b_bar=self.b_bar, + Y_per_bin=self.Y_per_bin, + eff_params={k: v for k, v in eff_params.items() if k != "np_model"}, + gnu_params={k: v for k, v in gnu_params.items() if k != "np_model_nu"}, + np_model=self.np_model, + np_model_nu=self.np_model_nu, + bT_J0_kernel=self.bT_J0_kernel, + bT_simpson_weights=self.bT_simpson_w, + Y_unique=self.Y_feff_unique, + Y_inverse_idx=self.Y_feff_inverse_idx, + ) + # 2. Sparse → dense (NQ, NY, NqT). Missing cells get 0. + sigma_dense = fz_int.sparse_to_dense_tf(sigma_flat, self.flat_idx) + # 3. Integrate over Q (arctan_Q² Simpson) → (NY, NqT). + sigma_YqT = fz_int.integrate_over_Q_tf(sigma_dense, self.Q_weights) + # 4. Rebin Y (signed) → absYVGen (|Y|-folded): (NabsYVGen, NqT). + sigma_absY_qT = fz_int.rebin_axis_tf(sigma_YqT, axis=0, weights=self.W_absY) + # 5. Rebin qT → ptVGen: (NabsYVGen, NptVGen). + sigma_absY_ptV = fz_int.rebin_axis_tf( + sigma_absY_qT, axis=1, weights=self.W_ptVGen + ) + # 6. Reorder to (NptVGen, NabsYVGen) to match R's gen axis order. + return tf.transpose(sigma_absY_ptV, perm=[1, 0]) + + # ========================================================================= + # Taylor surrogate + # ========================================================================= + + def _eff_gnu_from_array(self, lambdas_np): + """Helper: numpy 8-vector → (eff_params dict, gnu_params dict).""" + eff = {"np_model": self.np_model} + gnu = {"np_model_nu": self.np_model_nu} + for i, name in enumerate(self._param_order): + v = float(lambdas_np[i]) + if name in EFF_PARAMS: + eff[name] = v + elif name in GNU_PARAMS: + gnu[name] = v + else: + raise KeyError(name) + return eff, gnu + + def _sigma_gen_at_lambdas(self, lambdas_np): + """Wrapper around :meth:`_sigma_gen_at` taking a flat numpy λ vector.""" + eff, gnu = self._eff_gnu_from_array(lambdas_np) + return self._sigma_gen_at(eff, gnu).numpy() + + def _build_taylor_surrogate(self, h_rel, h_min, order): + """Precompute σ_gen(λ_central), ∂σ_gen/∂λ_i, and (order≥2) ∂²σ_gen/∂λ_i∂λ_j. + + Uses central finite differences with step size h_i = max(|λ_i|·h_rel, h_min). + Total full-Hankel evaluations: 1 + 16 + (28 if order≥2) = 45. + + Stored as ``tf.constant`` so the per-step ``_sigma_gen_taylor`` is + a small polynomial evaluation. + """ + import time + + t_start = time.time() + lambda_central = self.xparamdefault.numpy().astype(np.float64) + n = len(lambda_central) + h_i = np.maximum(np.abs(lambda_central) * h_rel, h_min) + # Per-parameter strategy: + # "central" — symmetric ±h FD: D = (σ⁺ − σ⁻)/(2h), H_ii = (σ⁺ − 2σ₀ + σ⁻)/h² + # "forward" — one-sided FD using σ₀, σ⁺ at h, σ⁺⁺ at 2h. Used when + # λ_c is at the physical lower bound (e.g. lambda6 = 0 in + # FranksVals). σ⁻ would be unphysical and produce overflow. + # Off-diagonal Hessian uses the σ⁺⁺_ij stencil which doesn't need σ⁻, + # so it works for both modes. + fd_mode = ["central"] * n + for i, name in enumerate(self._param_order): + min_v = PARAM_MIN_VALUE.get(name) + if min_v is None: + continue + max_h_allowed = lambda_central[i] - min_v + if max_h_allowed <= 0: + # Boundary case: use forward FD. Keep h_i at its preferred value. + fd_mode[i] = "forward" + else: + h_i[i] = min(h_i[i], 0.95 * max_h_allowed) + print( + f"[SCETlibNPParamModel] building Taylor surrogate (order={order}); " + f"h_i={dict(zip(self._param_order, h_i.round(4).tolist()))}", + flush=True, + ) + + # σ₀ at λ_central + sigma_0 = self._sigma_gen_at_lambdas(lambda_central) # (NptVGen, NabsYVGen) + elapsed = time.time() - t_start + print(f" central done in {elapsed:.1f}s; shape {sigma_0.shape}", flush=True) + + # σ₊ᵢ at λ_central + h_i e_i; σ₋ᵢ at λ_central − h_i e_i + sigma_plus = np.empty((n,) + sigma_0.shape, dtype=np.float64) + sigma_minus = np.empty_like(sigma_plus) + # Storage: for "central" mode params we keep σ⁻; for "forward" mode + # params we keep σ⁺⁺ (at +2h) in the same slot. The use site selects + # the correct stencil per-param. + sigma_other = np.empty_like(sigma_plus) + for i in range(n): + lp = lambda_central.copy() + lp[i] += h_i[i] + sigma_plus[i] = self._sigma_gen_at_lambdas(lp) + if fd_mode[i] == "central": + lm = lambda_central.copy() + lm[i] -= h_i[i] + sigma_other[i] = self._sigma_gen_at_lambdas(lm) + tag = "±h" + else: # forward + lpp = lambda_central.copy() + lpp[i] += 2.0 * h_i[i] + sigma_other[i] = self._sigma_gen_at_lambdas(lpp) + tag = "+h,+2h (forward)" + print( + f" {tag} for {self._param_order[i]} ({i+1}/{n}) at " + f"t+{time.time()-t_start:.1f}s", + flush=True, + ) + sigma_minus = sigma_other # retain old name for the central-FD slots + + # First derivatives D_i and diagonal Hessian H_ii — stencil depends on mode. + D = np.empty_like(sigma_plus) + H_full = None + if order >= 2: + H_full = np.zeros((n, n) + sigma_0.shape, dtype=np.float64) + for i in range(n): + if fd_mode[i] == "central": + D[i] = (sigma_plus[i] - sigma_minus[i]) / (2.0 * h_i[i]) + if H_full is not None: + H_full[i, i] = (sigma_plus[i] - 2.0 * sigma_0 + sigma_minus[i]) / ( + h_i[i] ** 2 + ) + else: # forward; sigma_minus[i] actually holds σ⁺⁺ at +2h + D[i] = (4.0 * sigma_plus[i] - sigma_minus[i] - 3.0 * sigma_0) / ( + 2.0 * h_i[i] + ) + if H_full is not None: + H_full[i, i] = (sigma_minus[i] - 2.0 * sigma_plus[i] + sigma_0) / ( + h_i[i] ** 2 + ) + if H_full is not None: + # Off-diagonals: 1-sided stencil + # H_ij ≈ [σ(λ_c + h_i e_i + h_j e_j) − σ(λ_c + h_i e_i) + # − σ(λ_c + h_j e_j) + σ(λ_c)] / (h_i h_j) + # 8·7/2 = 28 additional evaluations. + for i in range(n): + for j in range(i + 1, n): + lpp = lambda_central.copy() + lpp[i] += h_i[i] + lpp[j] += h_i[j] + sigma_pp = self._sigma_gen_at_lambdas(lpp) + H_ij = (sigma_pp - sigma_plus[i] - sigma_plus[j] + sigma_0) / ( + h_i[i] * h_i[j] + ) + H_full[i, j] = H_ij + H_full[j, i] = H_ij # symmetric + print( + f" H[{self._param_order[i]},{self._param_order[j]}] " + f"at t+{time.time()-t_start:.1f}s", + flush=True, + ) + + # Stash as tf.constants + dtype = self.indata.dtype + self._taylor_lambda_central = tf.constant(lambda_central, dtype=dtype) + self._taylor_sigma_central_2d = tf.constant(sigma_0, dtype=dtype) + self._taylor_D = tf.constant(D, dtype=dtype) # (n, NptVGen, NabsYVGen) + self._taylor_H = ( + tf.constant(H_full, dtype=dtype) if H_full is not None else None + ) + self._taylor_h_i = h_i # kept for diagnostics + print( + f"[SCETlibNPParamModel] surrogate ready in " + f"{time.time()-t_start:.1f}s; storage " + f"~{(1 + n + (n*n if H_full is not None else 0)) * sigma_0.size * 8 / 1e6:.1f} MB", + flush=True, + ) + + def _sigma_gen_taylor(self, lambdas_tf): + """Evaluate σ_gen(λ) via the precomputed Taylor surrogate. + + Returns shape (NptVGen, NabsYVGen). All ops are linear in (λ−λ_central) + for order=1, plus a quadratic correction for order=2. + """ + delta = lambdas_tf - self._taylor_lambda_central # (n,) + sigma = self._taylor_sigma_central_2d + tf.tensordot( + delta, self._taylor_D, axes=1 + ) # (NptVGen, NabsYVGen) + if self._taylor_H is not None: + quad = 0.5 * tf.einsum("i,j,ijab->ab", delta, delta, self._taylor_H) + sigma = sigma + quad + return sigma + + # ========================================================================= + # compute + # ========================================================================= + + def _unpack_params(self, param): + """Map flat param tensor to eff/gnu dicts in canonical order.""" + eff_params = {"np_model": self.np_model} + gnu_params = {"np_model_nu": self.np_model_nu} + # param values are stored directly (allowNegativeParam=True). + for i, name in enumerate(self._param_order): + v = param[i] + if name in EFF_PARAMS: + eff_params[name] = v + elif name in GNU_PARAMS: + gnu_params[name] = v + else: + raise KeyError(name) + return eff_params, gnu_params + + def compute(self, param, full=False): + """Return per-(bin, proc) scaling tensor. + + Shape: (N_reco, N_proc). Signal-proc column carries the per-reco-bin + ratio σ_reco(λ; b) / σ_reco(λ_central; b); other columns are 1. + """ + import time + + t_start = time.perf_counter() + if self.use_taylor: + # Fast path: polynomial in (λ − λ_central) with precomputed coefficients. + sigma_gen = self._sigma_gen_taylor(param) # (NptVGen, NabsYVGen) + else: + eff_params, gnu_params = self._unpack_params(param) + sigma_gen = self._sigma_gen_at(eff_params, gnu_params) + gen_flat = tf.reshape(sigma_gen, [-1]) + sigma_reco = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) + ratio = sigma_reco / self.sigma_reco_central # (N_reco,) + + # Build (N_reco, N_proc) scaling: ones except signal column. + N_reco = int(self.sigma_reco_central.shape[0]) + col_ones = tf.ones([N_reco, self.nproc], dtype=self.indata.dtype) + # Place ratio into the signal column. + ratio_col = tf.cast(tf.reshape(ratio, [N_reco, 1]), self.indata.dtype) + # mask: one-hot vector for signal_proc_idx. + mask = tf.one_hot( + self.signal_proc_idx, self.nproc, dtype=self.indata.dtype + ) # (N_proc,) + # rnorm = ones + (ratio - 1) * one_hot_signal + # → (ratio at signal col, 1 elsewhere) + rnorm = col_ones + (ratio_col - 1.0) * mask[tf.newaxis, :] + + # ---- Diagnostic timing ---- + if not hasattr(self, "_n_compute_calls"): + self._n_compute_calls = 0 + self._t_compute_total = 0.0 + import atexit + + atexit.register(self._print_compute_summary) + self._n_compute_calls += 1 + self._t_compute_total += time.perf_counter() - t_start + if self._n_compute_calls % 100 == 0: + n, t = self._n_compute_calls, self._t_compute_total + print( + f"[SCETlibNPParamModel] compute() running: {n} calls, " + f"total {t:.1f}s, mean {t*1000/n:.1f} ms/call", + flush=True, + ) + return rnorm + + def _print_compute_summary(self): + """Print final tally of compute() calls and time. Atexit hook.""" + n = getattr(self, "_n_compute_calls", 0) + t = getattr(self, "_t_compute_total", 0.0) + if n == 0: + return + print( + f"[SCETlibNPParamModel] compute() FINAL: {n} calls, " + f"total {t:.2f}s wall in compute(), " + f"mean {t*1000/n:.2f} ms/call", + flush=True, + ) diff --git a/wremnants/postprocessing/scetlib_np/response_matrix.py b/wremnants/postprocessing/scetlib_np/response_matrix.py new file mode 100644 index 000000000..e258f3dd2 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/response_matrix.py @@ -0,0 +1,128 @@ +"""Load the (reco × gen) response matrix R for the ParamModel. + +R lives in the unfolding histmaker output (a separate hdf5 from the fit tensor) +as ``nominal_postfsr_yieldsUnfolding`` under the Z sample group. We select +``acceptance=True``, project to the reco × (ptVGen, absYVGen) axes (summing +over helicitySig per Luca's guidance), and return a numpy array plus axis +metadata. + +Single entry point: + + load_R(unfolding_hdf5_path, + sample_key="Zmumu_2016PostVFP", + hist_name="nominal_postfsr_yieldsUnfolding") -> dict +""" + +import h5py +import numpy as np + +from wums import ioutils as wums_io + +DEFAULT_HIST = "nominal_postfsr_yieldsUnfolding" +DEFAULT_SAMPLE = "Zmumu_2016PostVFP" + +# Axes we keep, in canonical order: reco first, then gen. +RECO_AXES = ("ptll", "yll", "cosThetaStarll_quantile", "phiStarll_quantile") +GEN_AXES = ("ptVGen", "absYVGen") +# Axes we collapse: acceptance via {True} slice, helicitySig via project-out. +SUM_AXES = ("helicitySig",) + + +def load_R( + unfolding_hdf5_path, + sample_key=DEFAULT_SAMPLE, + hist_name=DEFAULT_HIST, + reco_axes=RECO_AXES, + gen_axes=GEN_AXES, +): + """Load R from the unfolding histmaker output. + + Returns a dict with: + R : ndarray shape (*reco_sizes, *gen_sizes), float64 + reco_axes : list of (name, edges) tuples in the canonical order + gen_axes : list of (name, edges) tuples + reco_shape : tuple of axis sizes (reco) + gen_shape : tuple of axis sizes (gen) + source : (path, sample_key, hist_name) for traceability + """ + with h5py.File(unfolding_hdf5_path, "r") as f: + if sample_key not in f: + raise KeyError( + f"{unfolding_hdf5_path}: no '{sample_key}' group. " + f"Available top-level: {list(f.keys())[:10]}" + ) + sample = wums_io.pickle_load_h5py(f[sample_key]) + try: + output = sample["output"] + except (KeyError, TypeError) as exc: + raise KeyError( + f"{sample_key}: no 'output' dict — schema mismatch?" + ) from exc + if hist_name not in output: + joint_candidates = [ + k for k in output.keys() if "yieldsUnfolding" in k or "Unfolding" in k + ] + raise KeyError( + f"{sample_key}: '{hist_name}' missing. " + f"Joint-hist candidates: {joint_candidates[:5]}" + ) + proxy = output[hist_name] + # Force materialization while the file is open. + h = proxy.get() if hasattr(proxy, "get") else proxy + + # Sanity-check the axes. + ax_names = [a.name for a in h.axes] + required = set(reco_axes) | set(gen_axes) | {"acceptance"} | set(SUM_AXES) + missing = required - set(ax_names) + if missing: + raise ValueError( + f"{hist_name}: missing expected axes {missing}. " f"Got: {ax_names}" + ) + + # Select acceptance=True (gen events in fiducial), keep gen + reco axes, + # sum over helicitySig. + h_sel = h[{"acceptance": True}] + h_proj = h_sel.project(*reco_axes, *gen_axes) + # .project sums over un-listed axes (i.e. helicitySig here). + + # Out from the with-block: hist is materialized. + R = h_proj.values(flow=False).astype(np.float64) + + reco_meta = [(name, h_proj.axes[name].edges) for name in reco_axes] + gen_meta = [(name, h_proj.axes[name].edges) for name in gen_axes] + reco_shape = tuple(h_proj.axes[name].size for name in reco_axes) + gen_shape = tuple(h_proj.axes[name].size for name in gen_axes) + + return dict( + R=R, + reco_axes=reco_meta, + gen_axes=gen_meta, + reco_shape=reco_shape, + gen_shape=gen_shape, + source=(unfolding_hdf5_path, sample_key, hist_name), + ) + + +if __name__ == "__main__": + import sys + + if len(sys.argv) < 2: + print( + f"usage: python -m {__name__.replace('.', '/')} [sample_key]", + file=sys.stderr, + ) + sys.exit(2) + sample = sys.argv[2] if len(sys.argv) > 2 else DEFAULT_SAMPLE + info = load_R(sys.argv[1], sample_key=sample) + print(f"R shape : {info['R'].shape}") + print(f"R sum : {info['R'].sum():.6g}") + print(f"reco_shape : {info['reco_shape']}") + print(f"gen_shape : {info['gen_shape']}") + for name, edges in info["reco_axes"]: + print( + f" reco {name}: size={len(edges)-1} edges=[{edges[0]:.3g}, {edges[-1]:.3g}]" + ) + for name, edges in info["gen_axes"]: + print( + f" gen {name}: size={len(edges)-1} edges=[{edges[0]:.3g}, {edges[-1]:.3g}]" + ) From ce90cfad72a6587c59a125cb6d3356eaafa44acf Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 29 May 2026 16:41:12 -0400 Subject: [PATCH 02/31] bump rabbit --- rabbit | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rabbit b/rabbit index 73f346aa2..843fbbd0a 160000 --- a/rabbit +++ b/rabbit @@ -1 +1 @@ -Subproject commit 73f346aa2fbc05f3f5823bda2588b0539fd355ae +Subproject commit 843fbbd0a9bbf48ab7726e0cc9cb0a9e5629616a From c4c1a32344298f63fafb66db7f05a32fcef5d739 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 3 Jun 2026 10:19:15 -0400 Subject: [PATCH 03/31] some fixes to the NP model --- .../postprocessing/scetlib_np/param_model.py | 558 ++++++++++-------- 1 file changed, 305 insertions(+), 253 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index a2b98ed0e..5ca91c2ec 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -2,11 +2,26 @@ Architecture (template-style fit): - σ_reco(λ; b) = Σ_g R(b, g) · σ_gen(λ; g) + P(b, g) = R_raw(b, g) / N_gen(g) (normalized response) + σ_reco(λ; b) = Σ_g P(b, g) · σ_gen(λ; g) ratio(b) = σ_reco(λ; b) / σ_reco(λ_central; b) rnorm = ratio per reco bin, broadcast over signal proc; ones elsewhere. -R is the (reco × gen) response matrix loaded from the upstream unfolding +The raw response counts R_raw(b, g) carry the MC's absolute gen spectrum, which +is theory-dependent. A response matrix should encode only the gen→reco +*mapping*, so each gen column is normalized by the gen-total N_gen(g) → +P(b, g) = efficiency × migration (theory-independent). N_gen(g) is the xnorm +"postfsr" histogram from the unfolding output — the generated fiducial yield +per gen bin BEFORE reco selection — loaded alongside R by response_matrix. +Then σ_gen(λ) is folded through P. NB normalizing instead by the reco-passing +marginal Σ_b R_raw(b, g) is wrong: R is post-reco-selection so that marginal +already includes efficiency, and dividing by it cancels efficiency +(migration-only), which closes far worse — efficiency is not flat in gen bin. +(If the gen-total hist is absent, σ_gen(λ_c) is used as a proxy for N_gen, but +then σ_gen cancels in σ_reco(λ_c) = R_raw·1 and the λ_central closure can't +test the integral.) + +R_raw is the (reco × gen) response matrix loaded from the upstream unfolding histmaker output (a separate hdf5 from the fit-tensor input). λ_central is read from the fit-tensor's meta_info_input via the upstream @@ -27,6 +42,9 @@ λ_central). All λ values are TF Variables — differentiable in the fit. """ +import json +import os +import re from typing import Mapping, Optional import numpy as np @@ -40,30 +58,28 @@ from wremnants.postprocessing.scetlib_np import btgrid_tf as fz_tf from wremnants.postprocessing.scetlib_np import lambda_central as scetlib_lambda_central from wremnants.postprocessing.scetlib_np import response_matrix as fz_R +from wremnants.utilities import common as wrem_common + +# Default fixed-order inputs for the NP-independent nonsingular term +# σ_ns = DYTurbo − SCETlib_singular, resolved relative to this package via +# wrem_common.data_dir (= /wremnants-data/data). These are the CT18Z +# N3+0LL fixed-order pieces; both are NP-independent (same for any λ tune). +_NONSING_FO_SING_DEFAULT = os.path.join( + wrem_common.data_dir, + "TheoryCorrections", + "inclusive_Z_COM13_CT18Z_N3+0LL_lattice_lambda4bugfix_fine_nnlo_sing_combined.pkl", +) +_NONSING_DYTURBO_DEFAULT = os.path.join( + wrem_common.data_dir, + "TheoryCorrections", + "results_z-2d-nnlo-vj-CT18ZNNLO-{scale}-scetlibmatch.txt", +) # Ordered list of the v1 continuous λ. CS-side first, then TMD-effective. GNU_PARAMS = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") ALL_PARAMS = GNU_PARAMS + EFF_PARAMS -# Physical lower bounds used by the Taylor surrogate so the finite-difference -# step (λ_central − h) can't cross into a region where F_eff / γ_ν^NP blow up. -# - lambda6 enters as `lambda6 · bT⁵ / lambda_inf` inside a tanh; if lambda6 is -# negative the tanh saturates to −1 at large bT and the outer exp inverts -# sign → exp(+huge). So lambda6 ≥ 0 is enforced. -# - lambda_inf / lambda_inf_nu sit in denominators (and in `exp(−2·lambda_inf·bT)`); -# require strictly positive with a safety margin. -PARAM_MIN_VALUE = { - "lambda2_nu": None, - "lambda4_nu": None, - "lambda_inf_nu": 0.05, - "lambda2": None, - "lambda4": None, - "lambda6": 0.0, - "delta_lambda2": None, - "lambda_inf": 0.05, -} - # Theorist-recommended Gaussian prior widths for the SCETlib NP λ parameters. # Source: NP-NP discussion slide, central values 2026-05, plus a wide @@ -92,6 +108,45 @@ } +def _load_lambda_central_file(path): + """Load a λ_central override from a JSON or YAML file. + + The file must decode to a dict with ``eff_params`` and ``gnu_params`` + sub-dicts (same shape as :func:`scetlib_lambda_central.read_lambda_central`). + Format is chosen by extension (``.yaml``/``.yml`` → YAML, else JSON); + YAML's loader also accepts JSON, so this is forgiving either way. + """ + + if not os.path.exists(path): + raise FileNotFoundError( + f"SCETLIB_NP_LAMBDA_CENTRAL_FILE points to a missing file: {path!r}" + ) + with open(path) as f: + text = f.read() + if path.lower().endswith((".yaml", ".yml")): + import yaml + + data = yaml.safe_load(text) + else: + try: + data = json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError( + f"SCETLIB_NP_LAMBDA_CENTRAL_FILE {path!r} is not valid JSON; got {exc}" + ) from exc + if ( + not isinstance(data, dict) + or "eff_params" not in data + or "gnu_params" not in data + ): + raise ValueError( + f"SCETLIB_NP_LAMBDA_CENTRAL_FILE {path!r} must decode to a dict with " + f"'eff_params' and 'gnu_params' keys; got {type(data).__name__} " + f"with keys {list(data) if isinstance(data, dict) else ''}." + ) + return data + + def _crop_R_to_fit(R, R_reco_axes, fit_reco_axes, tol=1e-9): """Crop R's trailing reco bins so its reco shape matches the fit. @@ -125,6 +180,97 @@ def _crop_R_to_fit(R, R_reco_axes, fit_reco_axes, tol=1e-9): return R[slices] +def _bin_sum_matrix(src_centers, target_edges, tol=1e-6): + """(N_target, N_src) 0/1 matrix that SUMS bin-integrated source bins whose + centre falls in each target bin. Source bins outside all target bins are + dropped — a natural truncation to the target range (qT>ptVGen_max, |Y|>absY_max).""" + src = np.asarray(src_centers, dtype=np.float64) + edges = np.asarray(target_edges, dtype=np.float64) + W = np.zeros((edges.size - 1, src.size), dtype=np.float64) + for i in range(edges.size - 1): + m = (src >= edges[i] - tol) & (src <= edges[i + 1] + tol) + W[i, m] = 1.0 + return W + + +def compute_nonsingular_gen( + fo_sing_path, + dyturbo_path, + gen_axes_meta, + charge=0, + q_lo=60.0, + q_hi=120.0, + qt_cutoff=1.0, + dyturbo_axes=("Q", "Y", "qT"), +): + """Nonsingular FO term on the model gen grid (NptVGen, NabsYVGen). + + The fixed-order/DYTurbo matching adds a NP-INDEPENDENT piece to σ_gen: + σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns , + where the nonsingular is read straight from the original fixed-order inputs: + σ_ns = (DYTurbo fixed order) − (SCETlib singular fixed order) + — exactly the ``-hfo_sing + hfo`` that ``read_matched_scetlib_hist`` forms. + ``fo_sing_path`` is the SCETlib singular ``…_nnlo_sing…combined.pkl``; + ``dyturbo_path`` is the DYTurbo fixed-order ``results_…scetlibmatch.txt`` + (use ``{scale}`` → mur1-muf1 for the central). The nonsingular is zeroed below + ``qt_cutoff`` (as make_theory_corr does), Q-windowed to [q_lo, q_hi], |Y|-folded, + and projected onto the coarse (ptVGen, absYVGen) gen bins by SUMMING the + bin-integrated native bins. + """ + from wremnants.utilities.io_tools import input_tools + from wums import boostHistHelpers as hh + + def _central(h): + if "vars" in h.axes.name: + names = list(h.axes["vars"]) + idx = 0 + for c in ("central", "pdf0", "nominal"): + if c in names: + idx = names.index(c) + break + h = h[{"vars": idx}] + return h + + # SCETlib singular fixed order, and DYTurbo fixed order, from their own files. + dyturbo_path = ( + dyturbo_path.format(scale="mur1-muf1") + if "{scale}" in dyturbo_path + else dyturbo_path + ) + hfo_sing = _central(input_tools.read_scetlib_hist(fo_sing_path, charge=charge)) + hfo = input_tools.read_dyturbo_hist( + [dyturbo_path], axes=list(dyturbo_axes), charge=charge + ) + if "vars" in hfo.axes.name: + hfo = _central(hfo) + + # Align shared physics axes (DYTurbo is coarser), then σ_ns = DYTurbo − singular. + for ax in ("Y", "Q", "qT"): + if ax in set(hfo.axes.name) & set(hfo_sing.axes.name): + hfo, hfo_sing = hh.rebinHistsToCommon([hfo, hfo_sing], ax) + nonsing_h = hh.addHists(-1.0 * hfo_sing, hfo, flow=False, by_ax_name=False) + + if "charge" in nonsing_h.axes.name: + nonsing_h = nonsing_h[{"charge": sum}] + # Q-window: slice(...,sum) sums ONLY the in-range Q bins (no underflow leak). + Qe = np.asarray(nonsing_h.axes["Q"].edges, dtype=np.float64) + qi = int(np.argmin(np.abs(Qe - q_lo))) + qj = int(np.argmin(np.abs(Qe - q_hi))) + nonsing_h = nonsing_h[{"Q": slice(qi, qj, sum)}] + nonsing_h = hh.makeAbsHist(nonsing_h, "Y") # signed Y -> |Y| + + qT_c = np.asarray(nonsing_h.axes["qT"].centers, dtype=np.float64) + absY_c = np.asarray(nonsing_h.axes["absY"].centers, dtype=np.float64) + v = nonsing_h.project("qT", "absY").values(flow=False) # (qT, absY) + v[qT_c < qt_cutoff, :] = 0.0 # zero the nonsingular below the cutoff + + ptV_edges = np.asarray(gen_axes_meta[0][1], dtype=np.float64) + absY_edges = np.asarray(gen_axes_meta[1][1], dtype=np.float64) + Wp = _bin_sum_matrix(qT_c, ptV_edges) # (NptVGen, NqT) + Wa = _bin_sum_matrix(absY_c, absY_edges) # (NabsYVGen, NabsYsrc) + return Wp @ v @ Wa.T # (NptVGen, NabsYVGen) + + class SCETlibNPParamModel(ParamModel): def __init__( @@ -138,10 +284,10 @@ def __init__( Q_hi: float = 120.0, poi_params: Optional[tuple] = (), prior_sigmas: Optional[Mapping] = None, - use_taylor: bool = False, - taylor_order: int = 2, - taylor_h_rel: float = 0.10, - taylor_h_min: float = 0.005, + include_nonsingular: bool = True, + nonsingular_fo_sing: str = _NONSING_FO_SING_DEFAULT, + nonsingular_dyturbo: str = _NONSING_DYTURBO_DEFAULT, + nonsingular_qt_cutoff: float = 1.0, **kwargs, ): """Construct the ParamModel. @@ -202,22 +348,18 @@ def __init__( # ---- λ_central # Three sources of λ_central, in priority order: # 1. ``lambda_central`` constructor arg (explicit dict). - # 2. ``SCETLIB_NP_LAMBDA_CENTRAL_JSON`` env var (JSON-encoded dict - # with ``eff_params`` and ``gnu_params``). Useful when the upstream - # SCETlib pkl isn't accessible (e.g. colleague's input). + # 2. ``SCETLIB_NP_LAMBDA_CENTRAL_FILE`` env var — path to a JSON or + # YAML file with ``eff_params`` and ``gnu_params``. Overrides the + # metadata auto-detect; useful when the upstream SCETlib pkl isn't + # accessible (e.g. a colleague's input). # 3. Auto-detect from the fit hdf5's theoryCorr → upstream pkl. - import json - import os - - env_lc = os.environ.get("SCETLIB_NP_LAMBDA_CENTRAL_JSON", "").strip() - if lambda_central is None and env_lc: - try: - lambda_central = json.loads(env_lc) - except json.JSONDecodeError as exc: - raise ValueError( - f"SCETLIB_NP_LAMBDA_CENTRAL_JSON must be valid JSON; got {exc}" - ) - print(f"[SCETlibNPParamModel] λ_central from env var", flush=True) + env_lc_file = os.environ.get("SCETLIB_NP_LAMBDA_CENTRAL_FILE", "").strip() + if lambda_central is None and env_lc_file: + lambda_central = _load_lambda_central_file(env_lc_file) + print( + f"[SCETlibNPParamModel] λ_central from file {env_lc_file!r}", + flush=True, + ) if lambda_central is None: # Auto-detect from indata.metadata (loaded by rabbit's # FitInputData from the input HDF5's "meta" group). @@ -231,6 +373,14 @@ def __init__( lambda_central = scetlib_lambda_central.read_lambda_central_from_meta( indata_meta, _source="indata.metadata" ) + print( + f"[SCETlibNPParamModel] λ_central auto-detected from indata.metadata", + flush=True, + ) + print(f"[SCETlibNPParamModel] λ_central:", flush=True) + for key, value in lambda_central.items(): + print(f" {key} = {value!r}", flush=True) + self.eff_central = dict(lambda_central["eff_params"]) self.gnu_central = dict(lambda_central["gnu_params"]) self.np_model = self.eff_central["np_model"] @@ -238,10 +388,6 @@ def __init__( # ---- btgrid + dense layout grid = btgrid_cache.load(btgrid_dir) - self._btgrid_meta = dict( - shards=grid["n_shards"], - n_bins=len(grid["bins"]), - ) idx_map = fz_int.dense_index_map(grid["bins"]) self.Q_unique = idx_map["Q_unique"] self.Y_unique = idx_map["Y_unique"] @@ -296,7 +442,24 @@ def __init__( self.gen_shape = R_arr.shape[len(fit_reco_axes) :] N_reco = int(np.prod(self.reco_shape)) N_gen = int(np.prod(self.gen_shape)) - self.R = tf.constant(R_arr.reshape(N_reco, N_gen), dtype=fz_tf.DTYPE) + # Raw response counts; normalized to a response below. + self._R_raw = tf.constant(R_arr.reshape(N_reco, N_gen), dtype=fz_tf.DTYPE) + # Gen-total denominator N_gen(g) from the xnorm hist ("postfsr"): the + # generated fiducial yield per gen bin (pre-reco-selection). Dividing R + # by this gives the theory-independent efficiency×migration response. + # Falls back to the σ_gen(λ_c) proxy if the gen-total isn't in the file. + if R_info.get("N_gen") is not None: + self._N_gen_flat = tf.constant( + R_info["N_gen"].reshape(-1), dtype=fz_tf.DTYPE + ) + else: + self._N_gen_flat = None + print( + "[SCETlibNPParamModel] WARNING: no gen-total hist in unfolding " + "output; falling back to σ_gen(λ_c) as the response normalizer " + "(circular closure — see param_model docstring).", + flush=True, + ) self._reco_axes_meta = [ (name, fit_axes[1]) for (name, fit_axes) in zip( @@ -331,14 +494,88 @@ def __init__( dtype=fz_tf.DTYPE, ) - # ---- Cache σ_reco(λ_central): the denominator of the ratio. + # ---- Normalize the response, then cache σ_reco(λ_central). + # A response matrix must encode only the gen→reco *mapping*, not the + # MC's absolute gen spectrum. Normalize each gen column by the gen-total + # N_gen(g) (the xnorm "postfsr" hist — generated fiducial yield before + # reco selection) → P(b|g) = eff×migration (theory-independent): + # P(b|g) = R_raw(b,g) / N_gen(g) + # σ_reco(λ;b) = Σ_g P(b|g) · σ_gen(λ;g) + # σ_reco(λ_c;b) = Σ_g P(b|g) · σ_gen(λ_c;g) + # NB the reco-passing marginal Σ_b R_raw(b,g) is the WRONG normalizer: + # it already includes efficiency (R is post-reco-selection), so dividing + # by it cancels efficiency (migration-only) — closes far worse, ε is not + # flat in gen bin. We use the true gen-total N_gen instead. Because + # σ_reco_central then depends on σ_gen(λ_c) (it does NOT collapse to + # R_raw·1), the λ_central closure is a genuine test of the integral. + # Fallback: if the gen-total hist is absent, use σ_gen(λ_c) as a proxy + # for N_gen (∝ σ_gen^MC) — keeps efficiency but makes the closure + # circular (σ_gen cancels to R_raw·1). + # Native-binning Q-integrated reconstruction (NY, NqT) on the signed-Y / + # qT grid, BEFORE the |Y|-fold and qT-rebin — exposed so the native-binning + # validation can compare it to the SCETlib reference / numpy factorize + # without the projection layer. + self.sigma_YqT_central = self._sigma_YqT_native_at( + self.eff_central, self.gnu_central + ) + # ---- Optional fixed-order/DYTurbo nonsingular term (NP-independent). + # σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns. Added at GEN level, so it folds + # through the same response R as the resummed piece. Because rnorm is a + # ratio, this correctly DILUTES the NP variation where the FO dominates + # (high qT). σ_ns is a constant (no λ dependence). + self.include_nonsingular = bool(include_nonsingular) + if self.include_nonsingular: + _dy0 = ( + nonsingular_dyturbo.format(scale="mur1-muf1") + if (nonsingular_dyturbo and "{scale}" in nonsingular_dyturbo) + else nonsingular_dyturbo + ) + missing = [ + p for p in (nonsingular_fo_sing, _dy0) if not (p and os.path.exists(p)) + ] + if missing: + raise FileNotFoundError( + "include_nonsingular=True needs the fixed-order inputs for " + "σ_ns = DYTurbo − SCETlib_singular, but these are missing:\n " + + "\n ".join(missing) + + "\nThey live under wremnants-data/data/TheoryCorrections (the " + "SCETlib singular …_nnlo_sing…combined.pkl and the DYTurbo " + "results_…scetlibmatch.txt). Pass nonsingular_fo_sing / " + "nonsingular_dyturbo, or set include_nonsingular=False for resum-only." + ) + sigma_ns_np = compute_nonsingular_gen( + nonsingular_fo_sing, + nonsingular_dyturbo, + self._gen_axes_meta, + q_lo=Q_lo, + q_hi=Q_hi, + qt_cutoff=nonsingular_qt_cutoff, + ) + if sigma_ns_np.shape != tuple(self.gen_shape): + raise ValueError( + f"nonsingular gen shape {sigma_ns_np.shape} != model gen shape " + f"{tuple(self.gen_shape)}" + ) + self.sigma_ns = tf.constant(sigma_ns_np, dtype=fz_tf.DTYPE) + else: + self.sigma_ns = tf.zeros(self.gen_shape, dtype=fz_tf.DTYPE) sigma_gen_central = self._sigma_gen_at(self.eff_central, self.gnu_central) - # σ_reco(λ_central) = R · σ_gen(λ_central). Flatten the gen axes. + # The pure gen-level integral (NptVGen, NabsYVGen), BEFORE folding through + # the response — used by the gen-level validation to test the integral + # in isolation (no R). + self.sigma_gen_central = sigma_gen_central gen_flat = tf.reshape(sigma_gen_central, [-1]) - self.sigma_reco_central = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) - # Sanity floor — if any reco bin has zero or negative central yield, - # the ratio would blow up. Use the central yield directly; any genuine - # zero is a binning issue to flag. + if tf.reduce_any(gen_flat <= 0).numpy(): + n_bad = int(tf.reduce_sum(tf.cast(gen_flat <= 0, tf.int32))) + raise ValueError( + f"SCETlibNPParamModel: {n_bad} gen bins have non-positive " + f"σ_gen(λ_central); cannot normalize / fold the response." + ) + N_gen = self._N_gen_flat if self._N_gen_flat is not None else gen_flat + # Guard empty gen bins (no generated events): leave column at 0. + safe_N_gen = tf.where(N_gen > 0, N_gen, tf.ones_like(N_gen)) + self.R = self._R_raw / safe_N_gen[tf.newaxis, :] # P(b|g) = R_raw / N_gen + self.sigma_reco_central = tf.linalg.matvec(self.R, gen_flat) # Σ_g P·σ_gen(λ_c) if tf.reduce_any(self.sigma_reco_central <= 0).numpy(): n_bad = int(tf.reduce_sum(tf.cast(self.sigma_reco_central <= 0, tf.int32))) raise ValueError( @@ -373,7 +610,6 @@ def __init__( defaults = np.array( [central_lookup[p] for p in self._param_order], dtype=np.float64 ) - import os env_override = os.environ.get("SCETLIB_NP_XPARAMDEFAULT", "").strip() if env_override: @@ -420,46 +656,6 @@ def __init__( # prior_means defaults to xparamdefault if not set, so don't store # redundantly — Fitter will fall back to xparamdefault. - # ---- Taylor surrogate (optional, default OFF) - # Precompute σ_gen(λ_central) plus first/second derivatives so that - # per-fit-step compute() is a polynomial in (λ − λ_central) instead of - # a full Hankel/Simpson integral. Drops per-step cost from O(10s) to - # O(ms). Accuracy: validated against the full integral; expected - # sub-percent within typical NP variation ranges for quadratic order. - # Off by default — opt in via ``use_taylor=True`` kwarg (or env var - # ``SCETLIB_NP_USE_TAYLOR=1``). Env var ``SCETLIB_NP_USE_TAYLOR=0`` - # also forces it off (useful when the kwarg is set by a CLI wrapper). - env_taylor = os.environ.get("SCETLIB_NP_USE_TAYLOR", "").strip().lower() - if env_taylor in ("0", "false", "no", "off"): - use_taylor = False - print( - "[SCETlibNPParamModel] Taylor surrogate disabled by env var", flush=True - ) - elif env_taylor in ("1", "true", "yes", "on"): - use_taylor = True - print( - "[SCETlibNPParamModel] Taylor surrogate enabled by env var", flush=True - ) - self.use_taylor = use_taylor - self.taylor_order = int(taylor_order) - if self.use_taylor: - print( - f"[SCETlibNPParamModel] Taylor surrogate ON " - f"(order={int(taylor_order)}, h_rel={taylor_h_rel}, " - f"h_min={taylor_h_min}) — per-step compute() is a polynomial " - f"in (λ − λ_central), not the full Hankel integral", - flush=True, - ) - self._build_taylor_surrogate( - h_rel=taylor_h_rel, h_min=taylor_h_min, order=self.taylor_order - ) - else: - print( - "[SCETlibNPParamModel] Taylor surrogate OFF — per-step " - "compute() runs the full Hankel/Simpson integral", - flush=True, - ) - # ========================================================================= # Helpers # ========================================================================= @@ -502,7 +698,6 @@ def _check_discrete_np_double_counting(self, freeze_patterns): the ParamModel should describe → print a loud banner with the exact freeze args to add. """ - import re systs = getattr(self.indata, "systs", None) if systs is None or len(systs) == 0: @@ -589,8 +784,12 @@ def _fit_reco_axes(self, indata): # σ_gen evaluation # ========================================================================= - def _sigma_gen_at(self, eff_params, gnu_params): - """Evaluate σ_gen(λ) on R's gen binning. Returns shape (NptVGen, NabsYVGen).""" + def _sigma_YqT_native_at(self, eff_params, gnu_params): + """Reconstruct σ(λ) on the btgrid and Q-integrate, returning the result + in the btgrid's *native* binning: shape (NY, NqT) on the signed-Y / + qT grid (Y_unique, qT_unique), BEFORE the |Y|-fold and qT-rebin. This + is the object that the native-binning validation compares against the + SCETlib spectrum reference (curve 1) and the numpy `factorize` (curve 2).""" # 1. Reconstruct σ on the btgrid's flat (Nbins,) layout. sigma_flat = fz_tf.reconstruct_batch_tf( qT_per_bin=self.qT_per_bin, @@ -611,7 +810,11 @@ def _sigma_gen_at(self, eff_params, gnu_params): # 2. Sparse → dense (NQ, NY, NqT). Missing cells get 0. sigma_dense = fz_int.sparse_to_dense_tf(sigma_flat, self.flat_idx) # 3. Integrate over Q (arctan_Q² Simpson) → (NY, NqT). - sigma_YqT = fz_int.integrate_over_Q_tf(sigma_dense, self.Q_weights) + return fz_int.integrate_over_Q_tf(sigma_dense, self.Q_weights) + + def _sigma_gen_at(self, eff_params, gnu_params): + """Evaluate σ_gen(λ) on R's gen binning. Returns shape (NptVGen, NabsYVGen).""" + sigma_YqT = self._sigma_YqT_native_at(eff_params, gnu_params) # 4. Rebin Y (signed) → absYVGen (|Y|-folded): (NabsYVGen, NqT). sigma_absY_qT = fz_int.rebin_axis_tf(sigma_YqT, axis=0, weights=self.W_absY) # 5. Rebin qT → ptVGen: (NabsYVGen, NptVGen). @@ -619,10 +822,12 @@ def _sigma_gen_at(self, eff_params, gnu_params): sigma_absY_qT, axis=1, weights=self.W_ptVGen ) # 6. Reorder to (NptVGen, NabsYVGen) to match R's gen axis order. - return tf.transpose(sigma_absY_ptV, perm=[1, 0]) + sigma_resum = tf.transpose(sigma_absY_ptV, perm=[1, 0]) + # 7. Add the NP-independent fixed-order/DYTurbo nonsingular (zeros if off). + return sigma_resum + self.sigma_ns # ========================================================================= - # Taylor surrogate + # λ-vector helpers # ========================================================================= def _eff_gnu_from_array(self, lambdas_np): @@ -639,156 +844,6 @@ def _eff_gnu_from_array(self, lambdas_np): raise KeyError(name) return eff, gnu - def _sigma_gen_at_lambdas(self, lambdas_np): - """Wrapper around :meth:`_sigma_gen_at` taking a flat numpy λ vector.""" - eff, gnu = self._eff_gnu_from_array(lambdas_np) - return self._sigma_gen_at(eff, gnu).numpy() - - def _build_taylor_surrogate(self, h_rel, h_min, order): - """Precompute σ_gen(λ_central), ∂σ_gen/∂λ_i, and (order≥2) ∂²σ_gen/∂λ_i∂λ_j. - - Uses central finite differences with step size h_i = max(|λ_i|·h_rel, h_min). - Total full-Hankel evaluations: 1 + 16 + (28 if order≥2) = 45. - - Stored as ``tf.constant`` so the per-step ``_sigma_gen_taylor`` is - a small polynomial evaluation. - """ - import time - - t_start = time.time() - lambda_central = self.xparamdefault.numpy().astype(np.float64) - n = len(lambda_central) - h_i = np.maximum(np.abs(lambda_central) * h_rel, h_min) - # Per-parameter strategy: - # "central" — symmetric ±h FD: D = (σ⁺ − σ⁻)/(2h), H_ii = (σ⁺ − 2σ₀ + σ⁻)/h² - # "forward" — one-sided FD using σ₀, σ⁺ at h, σ⁺⁺ at 2h. Used when - # λ_c is at the physical lower bound (e.g. lambda6 = 0 in - # FranksVals). σ⁻ would be unphysical and produce overflow. - # Off-diagonal Hessian uses the σ⁺⁺_ij stencil which doesn't need σ⁻, - # so it works for both modes. - fd_mode = ["central"] * n - for i, name in enumerate(self._param_order): - min_v = PARAM_MIN_VALUE.get(name) - if min_v is None: - continue - max_h_allowed = lambda_central[i] - min_v - if max_h_allowed <= 0: - # Boundary case: use forward FD. Keep h_i at its preferred value. - fd_mode[i] = "forward" - else: - h_i[i] = min(h_i[i], 0.95 * max_h_allowed) - print( - f"[SCETlibNPParamModel] building Taylor surrogate (order={order}); " - f"h_i={dict(zip(self._param_order, h_i.round(4).tolist()))}", - flush=True, - ) - - # σ₀ at λ_central - sigma_0 = self._sigma_gen_at_lambdas(lambda_central) # (NptVGen, NabsYVGen) - elapsed = time.time() - t_start - print(f" central done in {elapsed:.1f}s; shape {sigma_0.shape}", flush=True) - - # σ₊ᵢ at λ_central + h_i e_i; σ₋ᵢ at λ_central − h_i e_i - sigma_plus = np.empty((n,) + sigma_0.shape, dtype=np.float64) - sigma_minus = np.empty_like(sigma_plus) - # Storage: for "central" mode params we keep σ⁻; for "forward" mode - # params we keep σ⁺⁺ (at +2h) in the same slot. The use site selects - # the correct stencil per-param. - sigma_other = np.empty_like(sigma_plus) - for i in range(n): - lp = lambda_central.copy() - lp[i] += h_i[i] - sigma_plus[i] = self._sigma_gen_at_lambdas(lp) - if fd_mode[i] == "central": - lm = lambda_central.copy() - lm[i] -= h_i[i] - sigma_other[i] = self._sigma_gen_at_lambdas(lm) - tag = "±h" - else: # forward - lpp = lambda_central.copy() - lpp[i] += 2.0 * h_i[i] - sigma_other[i] = self._sigma_gen_at_lambdas(lpp) - tag = "+h,+2h (forward)" - print( - f" {tag} for {self._param_order[i]} ({i+1}/{n}) at " - f"t+{time.time()-t_start:.1f}s", - flush=True, - ) - sigma_minus = sigma_other # retain old name for the central-FD slots - - # First derivatives D_i and diagonal Hessian H_ii — stencil depends on mode. - D = np.empty_like(sigma_plus) - H_full = None - if order >= 2: - H_full = np.zeros((n, n) + sigma_0.shape, dtype=np.float64) - for i in range(n): - if fd_mode[i] == "central": - D[i] = (sigma_plus[i] - sigma_minus[i]) / (2.0 * h_i[i]) - if H_full is not None: - H_full[i, i] = (sigma_plus[i] - 2.0 * sigma_0 + sigma_minus[i]) / ( - h_i[i] ** 2 - ) - else: # forward; sigma_minus[i] actually holds σ⁺⁺ at +2h - D[i] = (4.0 * sigma_plus[i] - sigma_minus[i] - 3.0 * sigma_0) / ( - 2.0 * h_i[i] - ) - if H_full is not None: - H_full[i, i] = (sigma_minus[i] - 2.0 * sigma_plus[i] + sigma_0) / ( - h_i[i] ** 2 - ) - if H_full is not None: - # Off-diagonals: 1-sided stencil - # H_ij ≈ [σ(λ_c + h_i e_i + h_j e_j) − σ(λ_c + h_i e_i) - # − σ(λ_c + h_j e_j) + σ(λ_c)] / (h_i h_j) - # 8·7/2 = 28 additional evaluations. - for i in range(n): - for j in range(i + 1, n): - lpp = lambda_central.copy() - lpp[i] += h_i[i] - lpp[j] += h_i[j] - sigma_pp = self._sigma_gen_at_lambdas(lpp) - H_ij = (sigma_pp - sigma_plus[i] - sigma_plus[j] + sigma_0) / ( - h_i[i] * h_i[j] - ) - H_full[i, j] = H_ij - H_full[j, i] = H_ij # symmetric - print( - f" H[{self._param_order[i]},{self._param_order[j]}] " - f"at t+{time.time()-t_start:.1f}s", - flush=True, - ) - - # Stash as tf.constants - dtype = self.indata.dtype - self._taylor_lambda_central = tf.constant(lambda_central, dtype=dtype) - self._taylor_sigma_central_2d = tf.constant(sigma_0, dtype=dtype) - self._taylor_D = tf.constant(D, dtype=dtype) # (n, NptVGen, NabsYVGen) - self._taylor_H = ( - tf.constant(H_full, dtype=dtype) if H_full is not None else None - ) - self._taylor_h_i = h_i # kept for diagnostics - print( - f"[SCETlibNPParamModel] surrogate ready in " - f"{time.time()-t_start:.1f}s; storage " - f"~{(1 + n + (n*n if H_full is not None else 0)) * sigma_0.size * 8 / 1e6:.1f} MB", - flush=True, - ) - - def _sigma_gen_taylor(self, lambdas_tf): - """Evaluate σ_gen(λ) via the precomputed Taylor surrogate. - - Returns shape (NptVGen, NabsYVGen). All ops are linear in (λ−λ_central) - for order=1, plus a quadratic correction for order=2. - """ - delta = lambdas_tf - self._taylor_lambda_central # (n,) - sigma = self._taylor_sigma_central_2d + tf.tensordot( - delta, self._taylor_D, axes=1 - ) # (NptVGen, NabsYVGen) - if self._taylor_H is not None: - quad = 0.5 * tf.einsum("i,j,ijab->ab", delta, delta, self._taylor_H) - sigma = sigma + quad - return sigma - # ========================================================================= # compute # ========================================================================= @@ -817,12 +872,9 @@ def compute(self, param, full=False): import time t_start = time.perf_counter() - if self.use_taylor: - # Fast path: polynomial in (λ − λ_central) with precomputed coefficients. - sigma_gen = self._sigma_gen_taylor(param) # (NptVGen, NabsYVGen) - else: - eff_params, gnu_params = self._unpack_params(param) - sigma_gen = self._sigma_gen_at(eff_params, gnu_params) + eff_params, gnu_params = self._unpack_params(param) + sigma_gen = self._sigma_gen_at(eff_params, gnu_params) + # Fold σ_gen(λ) through the normalized migration response P (self.R). gen_flat = tf.reshape(sigma_gen, [-1]) sigma_reco = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) ratio = sigma_reco / self.sigma_reco_central # (N_reco,) From a289902a3a14b75621fca0d796bcbbacd433a62c Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Thu, 4 Jun 2026 06:32:42 -0400 Subject: [PATCH 04/31] some work on the NP model --- .../postprocessing/scetlib_np/btgrid_numpy.py | 9 +- .../postprocessing/scetlib_np/btgrid_tf.py | 5 + .../postprocessing/scetlib_np/param_model.py | 328 +++++++++++++----- .../scetlib_np/response_matrix.py | 164 ++++++++- 4 files changed, 396 insertions(+), 110 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_numpy.py b/wremnants/postprocessing/scetlib_np/btgrid_numpy.py index e4d271114..6c0a18ae1 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_numpy.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_numpy.py @@ -12,10 +12,11 @@ # Provides: # - Pure-numpy transcriptions of NP_model_effective (F_eff) and # NP_model_gammanu (gamma_nu^NP) that match the C++ code byte-for-byte. -# - A vectorised Hankel reconstruction of sigma(qT) from a cached bT-grid: -# sigma(qT) ~ int dbT bT J0(qT bT) * I_pert(b_bar) -# * exp(C_nu(bT) * gamma_nu^NP(b_bar)) -# * F_eff(Y, b_bar) +# - A vectorised Hankel reconstruction of sigma(qT) from a cached bT-grid. +# The full integrand — every factor and its bare-bT / b*(bT) / (Q,Y,qT) / +# lambda dependence — is written out ONCE in the module docstring of +# wremnants/postprocessing/scetlib_np/param_model.py (single source of +# truth). The reconstruct_* functions below implement it. # - Loaders for the bT-grid pickle shards produced by --bt-grid and for the # prior-art spectrum-mode "combined" pickles. # diff --git a/wremnants/postprocessing/scetlib_np/btgrid_tf.py b/wremnants/postprocessing/scetlib_np/btgrid_tf.py index f73e80b92..00edf32f8 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_tf.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_tf.py @@ -243,6 +243,11 @@ def reconstruct_batch_tf( ): """TF port of :func:`scetlib_btgrid_numpy.reconstruct_batch`. + Implements the per-(Q, Y, qT) bT-space integrand. The full formula with + every factor and its bare-bT / b*(bT) / (Q,Y,qT) / λ dependence is written + out once in the :mod:`param_model` module docstring (single source of + truth) — consult it rather than re-deriving the factors from this code. + All array-shape arguments are TF tensors or numpy arrays (will be cast). The λ values inside ``eff_params`` / ``gnu_params`` are the differentiable parameters — pass them as TF scalars (Variables or constants). diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 5ca91c2ec..725b01cc8 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -1,45 +1,174 @@ """SCETlibNPParamModel — continuous-λ rabbit ParamModel for SCETlib NP. -Architecture (template-style fit): - - P(b, g) = R_raw(b, g) / N_gen(g) (normalized response) - σ_reco(λ; b) = Σ_g P(b, g) · σ_gen(λ; g) - ratio(b) = σ_reco(λ; b) / σ_reco(λ_central; b) - rnorm = ratio per reco bin, broadcast over signal proc; ones elsewhere. - -The raw response counts R_raw(b, g) carry the MC's absolute gen spectrum, which -is theory-dependent. A response matrix should encode only the gen→reco -*mapping*, so each gen column is normalized by the gen-total N_gen(g) → -P(b, g) = efficiency × migration (theory-independent). N_gen(g) is the xnorm -"postfsr" histogram from the unfolding output — the generated fiducial yield -per gen bin BEFORE reco selection — loaded alongside R by response_matrix. -Then σ_gen(λ) is folded through P. NB normalizing instead by the reco-passing -marginal Σ_b R_raw(b, g) is wrong: R is post-reco-selection so that marginal -already includes efficiency, and dividing by it cancels efficiency -(migration-only), which closes far worse — efficiency is not flat in gen bin. +This ParamModel scales the signal reco template by a per-bin ratio of the +SCETlib nonperturbative (NP) prediction at the fitted λ vs. at λ_central. The +prediction is built in THREE STEPS, written out in order below — read top to +bottom for the full maths. This module docstring is the SINGLE SOURCE OF TRUTH: +:mod:`response_matrix`, :func:`btgrid_tf.reconstruct_batch_tf`, the numpy +reference :mod:`btgrid_numpy`, ``sigma_reco_central.md``, and the validation +scripts all point here rather than restate it. + +Pipeline at a glance (everything is a function of the NP parameters λ): + + Step 1 btgrid Hankel + Q integral → σ_resum(λ; g) resummed, on the gen grid + Step 2 + fixed-order matching → σ_gen(λ; g) = σ_resum(λ; g) + σ_ns(g) + Step 3 fold through response R → σ_reco(λ; b) gen → reco + Step 4 ratio vs λ_central → rnorm(b, proc) the shape handed to rabbit + +Steps 1–3 build the absolute physical cross section; Step 4 alone produces the +per-bin variation the fit consumes — they are kept separate on purpose. + +Indices: Q, Y, qT are the SCETlib btgrid axes (boson mass / rapidity / qT); +g = flattened gen bin (ptVGen, absYVGen); b = flattened reco bin (ptll, yll, +cosThetaStarll_quantile, phiStarll_quantile). λ splits into λ_eff (for F_eff) +and λ_ν (for γ_ν^NP) — the 8 differentiable parameters listed at the end. + +═════════════════════════════════════════════════════════════════════════════ +Step 1 — σ_resum(λ; g): the resummed prediction from the bT grid +═════════════════════════════════════════════════════════════════════════════ + +(1a) Per-(Q, Y, qT) bT-space (Hankel) integral — the NP parameters enter HERE: + + σ(Q, Y, qT; λ) = qT · ∫ dbT bT · J₀(qT·bT) + · I_pert(Q, Y, qT; b*(bT)) + · exp[ C_ν(Q, Y, qT; bT) · γ_ν^NP(b*(bT); λ_ν) ] + · F_eff(Y; b*(bT); λ_eff) + +Where each factor lives — bare bT vs the b*-frozen b̄T: + + bT bare bT, the Hankel integration variable. Enters ONLY the + measure factor bT and the Fourier kernel J₀(qT·bT). The leading + qT· prefactor is SCETlib's x = qT·bT integration convention. + b*(bT) the b*-prescription b_star_global(bT), cached as ``b_bar``: bT + frozen below b_max so the NP factors never reach the Landau pole. + EVERY NP-carrying factor is evaluated at b*(bT) — I_pert, γ_ν^NP, + and F_eff — never at bare bT. + I_pert perturbative bT-space integrand with NP off. Cached per + (Q, Y, qT) along the bT axis; λ-independent. + C_ν coefficient of γ_ν^NP in the rapidity (CS) log evolution. Cached + per (Q, Y, qT) AND bT — a full (Nbins, Nbt) array; λ-independent. + Taken at the BARE bT: the lone NP-exponent factor NOT frozen to + b* (per the SCETlib convention; see btgrid_numpy NP-model notes). + Because it varies with bT it sits inside the bT integral (not a + constant out front); because it varies with Q the Q integral + cannot be collapsed ahead of the fit — the λ-dependence does not + factor through ∫dQ (exp is nonlinear in C_ν·γ_ν^NP). + γ_ν^NP CS-side NP rapidity anomalous dimension; depends on b*(bT) and + λ_ν only (no Q/Y/qT). λ-dependent. + F_eff TMD-effective NP factor; depends on Y² and b*(bT) and λ_eff only + (no Q/qT). λ-dependent. + +Only γ_ν^NP and F_eff carry λ; everything else (bT, J₀, I_pert, C_ν, b*) is +λ-independent and precomputed once — the bT·J₀ kernel, the bT Simpson weights, +and the arctan_Q² Q-integration weights are all built at construction. + +(1b) Integrate over Q, then rebin onto the gen grid: + + σ_resum(λ; g) = rebin_{qT→ptVGen, |Y|→absYVGen} [ ∫_{Q_lo}^{Q_hi} dQ σ(Q, Y, qT; λ) ] + +The Q integral uses an arctan_Q² Simpson rule (the x = arctan((Q²−q0²)/(q0·Γ)) +transform flattens the Breit-Wigner Z peak). The result (Y, qT) is then rebinned +(Simpson) onto the unfolding hist's gen edges: qT → ptVGen, and the signed btgrid +Y axis folded into |Y| → absYVGen. The |Y| fold is valid because NP is +Y-symmetric (F_eff depends on Y², γ_ν^NP doesn't depend on Y at all). + +═════════════════════════════════════════════════════════════════════════════ +Step 2 — σ_gen(λ; g): add the NP-independent fixed-order nonsingular +═════════════════════════════════════════════════════════════════════════════ + + σ_gen(λ; g) = σ_resum(λ; g) + σ_ns(g) + + σ_ns(g) = rebin_{qT→ptVGen, |Y|→absYVGen} [ σ_DYTurbo^FO − σ_SCETlib-sing^FO ] + +The fixed-order matching adds the nonsingular piece σ_ns = (DYTurbo fixed order) +− (SCETlib singular fixed order), read straight from the original FO inputs +(the ``…_nnlo_sing…combined.pkl`` and the DYTurbo ``results_…scetlibmatch.txt``), +Q-windowed to [Q_lo, Q_hi], |Y|-folded, zeroed below qt_cutoff, and summed onto +the (ptVGen, absYVGen) gen bins (see :func:`compute_nonsingular_gen`). + +σ_ns is NP-INDEPENDENT (the same for every λ) and is added at GEN level, so it +folds through the same response R as σ_resum in Step 3. Because the fit uses a +ratio (Step 3), σ_ns correctly DILUTES the NP variation where the FO dominates +(high qT). Set include_nonsingular=False for a resum-only model (σ_ns = 0). + +═════════════════════════════════════════════════════════════════════════════ +Step 3 — σ_reco(λ; b): fold gen → reco through the response matrix +═════════════════════════════════════════════════════════════════════════════ + + P(b | g) = R_raw(b, g) / N_gen(g) (efficiency × migration) + σ_reco(λ; b) = Σ_g P(b | g) · σ_gen(λ; g) + +where, in these expressions, + g = the flattened GEN bin (ptVGen, absYVGen) — the grid σ_gen from Steps + 1–2 lives on (boson qT and |Y|), summed over by Σ_g; + b = the flattened RECO bin (ptll, yll, cosThetaStarll_quantile, + phiStarll_quantile) — the measured dilepton observables σ_reco lives on. +P(b | g) is the gen→reco mapping (one reco column per gen bin), so σ_reco(λ; b) +is just σ_gen pushed through the detector. The Σ_g is the matvec +``tf.linalg.matvec(self.R, σ_gen_flat)``. This step is pure detector folding — +no λ_central and no ratio enter here; that is Step 4. + +Factors (all loaded by :mod:`response_matrix`; see it for the hist mechanics): + + R_raw(b, g) reco×gen yield from the unfolding histmaker output + (``nominal_prefsr_yieldsUnfolding``, sample Zmumu): slice + acceptance=True (gen-fiducial), then project to reco×gen — + which SUMS the helicitySig axis. R is filled with + ``nominal_weight_helicity``, a PARTITION of the event weight + into the 8 helicity pieces that ADD BACK UP, so the physical + yield is the helicitySig SUM. Units: reco-selected, + gen-fiducial weighted event counts (already × reco efficiency). + N_gen(g) gen-total normalizer: the xnorm ``prefsr`` hist, gen-fiducial + but BEFORE reco selection. Takes the UL component + (helicitySig = −1), NOT the sum — N_gen is filled with + ``csAngularMoments``, a moment expansion whose A_i bins (0..7) + do NOT sum to σ (some are negative); UL is the + angular-integrated total. + ⚠ Same axis as R, OPPOSITE reduction: R is a weight partition + (SUM helicitySig), N_gen is a moment expansion (take UL). + Taking UL of R instead would discard the angular partition and + inflate the closure ~15×. + +Why N_gen and not the reco-passing marginal Σ_b R_raw(b, g): R is +post-reco-selection so that marginal already carries efficiency, and dividing +by it cancels efficiency (migration-only) — closes far worse, since efficiency +is strongly gen-dependent (ε ≈ 0.07–0.54 across gen bins on the current file). +N_gen(g) is the true generated total, so P = R_raw/N_gen is the +theory-independent gen→reco map. Pre-FSR: σ_gen, R, and N_gen must all sit at +the same QCD/boson gen level (the postfsr variants in the file close ~1% worse +— FSR mismatch). + +═════════════════════════════════════════════════════════════════════════════ +Step 4 — rnorm(b, proc): the per-reco-bin variation handed to rabbit +═════════════════════════════════════════════════════════════════════════════ + + ratio(b) = σ_reco(λ; b) / σ_reco(λ_central; b) + rnorm(b, proc) = 1 + (ratio(b) − 1) · [proc is signal] (1 in every other proc) + +This is the only object that leaves the model: compute() returns rnorm(b, proc), +and rabbit multiplies the signal process's reco template (reco bin b) by it, +leaving every other process at 1. Dividing by σ_reco(λ_central) cancels the +event-count↔cross-section scale and any overall normalization, so rnorm carries +purely the SHAPE of the NP variation per reco bin — Steps 1–3 build the absolute +σ_reco(λ; b), and Step 4 reduces it to the bin-by-bin template scaling the fit +needs. σ_reco(λ_central) is precomputed once at construction as the denominator. (If the gen-total hist is absent, σ_gen(λ_c) is used as a proxy for N_gen, but -then σ_gen cancels in σ_reco(λ_c) = R_raw·1 and the λ_central closure can't -test the integral.) +then σ_gen cancels in σ_reco(λ_c) = R_raw·1 and the λ_central closure can't test +the integral.) -R_raw is the (reco × gen) response matrix loaded from the upstream unfolding -histmaker output (a separate hdf5 from the fit-tensor input). +───────────────────────────────────────────────────────────────────────────── +Parameters and inputs +───────────────────────────────────────────────────────────────────────────── -λ_central is read from the fit-tensor's meta_info_input via the upstream -SCETlib correction pkl (see :mod:`scetlib_lambda_central`). - -σ_gen(λ; g) is evaluated on the btgrid then integrated over Q (arctan_Q² -Simpson) and rebinned (Simpson) onto the unfolding hist's gen edges -(ptVGen, absYVGen). The absYVGen-side rebin folds the signed btgrid Y axis -into |Y| bins (NP is Y-symmetric: F_eff depends on Y², γ_ν^NP doesn't depend -on Y at all). - -The 8 v1 parameters (all factorisable through the current btgrid): +The 8 v1 parameters λ (all factorisable through the current btgrid): γ_ν^NP (CS-side): lambda2_nu, lambda4_nu, lambda_inf_nu F_eff (TMD-effective): lambda2, lambda4, lambda6, delta_lambda2, lambda_inf -The np_model and np_model_nu strings are fixed at construction (from -λ_central). All λ values are TF Variables — differentiable in the fit. +λ_central is read from the fit-tensor's meta_info_input via the upstream +SCETlib correction pkl (see :mod:`scetlib_lambda_central`). The np_model and +np_model_nu strings are fixed at construction (from λ_central). All λ values +are TF Variables — differentiable in the fit. """ import json @@ -80,6 +209,20 @@ EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") ALL_PARAMS = GNU_PARAMS + EFF_PARAMS +# Positivity floor for the per-bin reco ratio in compute() (see its docstring). +# A pathological λ (e.g. λ4 < 0 with the bounded-tanh model) can drive σ_reco — +# and hence the predicted signal yield — negative, which makes the Poisson NLL +# NaN and stalls the minimizer. We soft-floor the ratio to a small positive value +# so a bad point becomes a LARGE-BUT-FINITE penalty the fit can back off from, +# with a non-zero gradient through the transition (softplus, not a hard clamp). +# RATIO_FLOOR_SCALE — softplus transition width. Chosen FAR below any physical +# response so healthy ratios (~0.9–1.1, and every validated λ-variation) pass +# through to machine precision: scale·softplus(r/scale) == r for r ≫ scale. +# RATIO_FLOOR_MIN — hard positive ground: softplus underflows to exactly 0 for +# the extreme (r ~ -1e43) case, so this keeps the yield strictly > 0 (no NaN). +RATIO_FLOOR_SCALE = 1.0e-4 +RATIO_FLOOR_MIN = 1.0e-9 + # Theorist-recommended Gaussian prior widths for the SCETlib NP λ parameters. # Source: NP-NP discussion slide, central values 2026-05, plus a wide @@ -298,7 +441,8 @@ def __init__( rabbit's input-data structure (passed by ``ph.load_models``). unfolding_hdf5_path Path to the upstream histmaker output containing - ``nominal_postfsr_yieldsUnfolding`` for R. + ``nominal_prefsr_yieldsUnfolding`` for R (and the ``prefsr`` xnorm + hist for N_gen) — see :mod:`response_matrix` for the defaults. btgrid_dir Directory of the SCETlib bT-grid shards (fineall). lambda_central @@ -444,7 +588,7 @@ def __init__( N_gen = int(np.prod(self.gen_shape)) # Raw response counts; normalized to a response below. self._R_raw = tf.constant(R_arr.reshape(N_reco, N_gen), dtype=fz_tf.DTYPE) - # Gen-total denominator N_gen(g) from the xnorm hist ("postfsr"): the + # Gen-total denominator N_gen(g) from the xnorm hist ("prefsr"): the # generated fiducial yield per gen bin (pre-reco-selection). Dividing R # by this gives the theory-independent efficiency×migration response. # Falls back to the σ_gen(λ_c) proxy if the gen-total isn't in the file. @@ -485,10 +629,14 @@ def __init__( W_absY[:, Y_pos_mask] = 2.0 * absY_rebin_pos self.W_absY = tf.constant(W_absY, dtype=fz_tf.DTYPE) - # qT rebin: btgrid qT (signed nonneg, NqT=141) → ptVGen edges (e.g. 20 bins, 0-44). - # Anything past ptVGen_max is out of fit range; we drop it silently for now - # (events with gen qT > ptVGen_max are routed through R's overflow which is - # not present in our materialised R — see plan doc, dyturbo-handoff item). + # qT rebin: btgrid qT (signed nonneg, NqT=141) → ptVGen edges. When + # load_R was built with ptVGen_overflow=True (default), ptVGen_edges ends + # in the overflow bin [last_gen_edge, PTVGEN_OVERFLOW_EDGE] (e.g. [44, 100]), + # so rebin_weights' last row Simpson-integrates the btgrid tail qT∈(44,100] + # into that overflow gen bin — matching R's gen-overflow column (true qT>44 + # migrating into the high-ptll reco bins). btgrid qT past the last edge + # (>100, beyond the grid) is dropped; negligible. Without the overflow + # column ptVGen_edges ends at 44 and that tail is simply truncated. self.W_ptVGen = tf.constant( fz_int.rebin_weights(self.qT_unique, ptVGen_edges, name="ptVGen"), dtype=fz_tf.DTYPE, @@ -497,7 +645,7 @@ def __init__( # ---- Normalize the response, then cache σ_reco(λ_central). # A response matrix must encode only the gen→reco *mapping*, not the # MC's absolute gen spectrum. Normalize each gen column by the gen-total - # N_gen(g) (the xnorm "postfsr" hist — generated fiducial yield before + # N_gen(g) (the xnorm "prefsr" hist — generated fiducial yield before # reco selection) → P(b|g) = eff×migration (theory-independent): # P(b|g) = R_raw(b,g) / N_gen(g) # σ_reco(λ;b) = Σ_g P(b|g) · σ_gen(λ;g) @@ -559,7 +707,11 @@ def __init__( self.sigma_ns = tf.constant(sigma_ns_np, dtype=fz_tf.DTYPE) else: self.sigma_ns = tf.zeros(self.gen_shape, dtype=fz_tf.DTYPE) - sigma_gen_central = self._sigma_gen_at(self.eff_central, self.gnu_central) + # Reuse the native (NY, NqT) integral already computed above for + # sigma_YqT_central — no need to run the bT reconstruction at λ_central twice. + sigma_gen_central = self._sigma_gen_at( + self.eff_central, self.gnu_central, sigma_YqT=self.sigma_YqT_central + ) # The pure gen-level integral (NptVGen, NabsYVGen), BEFORE folding through # the response — used by the gen-level validation to test the integral # in isolation (no R). @@ -575,6 +727,9 @@ def __init__( # Guard empty gen bins (no generated events): leave column at 0. safe_N_gen = tf.where(N_gen > 0, N_gen, tf.ones_like(N_gen)) self.R = self._R_raw / safe_N_gen[tf.newaxis, :] # P(b|g) = R_raw / N_gen + # Free the raw counts: only the normalized response self.R is used from + # here on (compute() never touches _R_raw) — no need to hold both. + del self._R_raw self.sigma_reco_central = tf.linalg.matvec(self.R, gen_flat) # Σ_g P·σ_gen(λ_c) if tf.reduce_any(self.sigma_reco_central <= 0).numpy(): n_bad = int(tf.reduce_sum(tf.cast(self.sigma_reco_central <= 0, tf.int32))) @@ -593,6 +748,12 @@ def __init__( ) self.signal_proc_idx = procs.index(signal_proc) self.nproc = indata.nproc + # Precompute the (1, N_proc) one-hot selecting the signal column, reused + # in every compute() to place the per-bin ratio (others stay at 1). + self._signal_col_mask = tf.reshape( + tf.one_hot(self.signal_proc_idx, self.nproc, dtype=indata.dtype), + [1, self.nproc], + ) # ---- ParamModel registration (POIs first, then NOUs). poi_params = tuple(poi_params or ()) @@ -812,9 +973,15 @@ def _sigma_YqT_native_at(self, eff_params, gnu_params): # 3. Integrate over Q (arctan_Q² Simpson) → (NY, NqT). return fz_int.integrate_over_Q_tf(sigma_dense, self.Q_weights) - def _sigma_gen_at(self, eff_params, gnu_params): - """Evaluate σ_gen(λ) on R's gen binning. Returns shape (NptVGen, NabsYVGen).""" - sigma_YqT = self._sigma_YqT_native_at(eff_params, gnu_params) + def _sigma_gen_at(self, eff_params, gnu_params, sigma_YqT=None): + """Evaluate σ_gen(λ) on R's gen binning. Returns shape (NptVGen, NabsYVGen). + + ``sigma_YqT`` lets a caller pass an already-computed native (NY, NqT) + integral to skip the (expensive) bT reconstruction — used at construction + to reuse ``self.sigma_YqT_central`` instead of integrating λ_central twice. + """ + if sigma_YqT is None: + sigma_YqT = self._sigma_YqT_native_at(eff_params, gnu_params) # 4. Rebin Y (signed) → absYVGen (|Y|-folded): (NabsYVGen, NqT). sigma_absY_qT = fz_int.rebin_axis_tf(sigma_YqT, axis=0, weights=self.W_absY) # 5. Rebin qT → ptVGen: (NabsYVGen, NptVGen). @@ -868,57 +1035,40 @@ def compute(self, param, full=False): Shape: (N_reco, N_proc). Signal-proc column carries the per-reco-bin ratio σ_reco(λ; b) / σ_reco(λ_central; b); other columns are 1. - """ - import time - t_start = time.perf_counter() + Positivity floor: the bT integral has no hard wall against pathological λ + (e.g. λ4 < 0 with the bounded-tanh models), which — via the b*-saturated, + hugely enhanced large-b region — can make σ_reco, and thus the predicted + signal yield, NEGATIVE. That would give a NaN Poisson NLL and a flat + gradient (tanh saturates), trapping the minimizer. We soft-floor the ratio + to a small positive value (RATIO_FLOOR_SCALE / RATIO_FLOOR_MIN) so a bad + point is a large-but-finite penalty with a usable gradient, NOT a crash. + The scale is far below any physical response, so the validated central and + λ-variation closures are unchanged (ratio == 1 at λ_central, to fp). This + is a numerical safety net, not a physics constraint — it does not stop the + fit from exploring negative-λ; it just keeps that exploration finite. + """ eff_params, gnu_params = self._unpack_params(param) sigma_gen = self._sigma_gen_at(eff_params, gnu_params) # Fold σ_gen(λ) through the normalized migration response P (self.R). gen_flat = tf.reshape(sigma_gen, [-1]) sigma_reco = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) ratio = sigma_reco / self.sigma_reco_central # (N_reco,) + # Soft positivity floor (see docstring): scale·softplus(r/scale) ≈ r for + # healthy r and → 0⁺ smoothly for r ≤ 0; the max() with a tiny ground + # keeps it strictly positive even when softplus underflows (r ~ -1e43). + scale = tf.constant(RATIO_FLOOR_SCALE, dtype=ratio.dtype) + ratio = tf.maximum( + scale * tf.math.softplus(ratio / scale), + tf.constant(RATIO_FLOOR_MIN, dtype=ratio.dtype), + ) - # Build (N_reco, N_proc) scaling: ones except signal column. - N_reco = int(self.sigma_reco_central.shape[0]) - col_ones = tf.ones([N_reco, self.nproc], dtype=self.indata.dtype) - # Place ratio into the signal column. - ratio_col = tf.cast(tf.reshape(ratio, [N_reco, 1]), self.indata.dtype) - # mask: one-hot vector for signal_proc_idx. - mask = tf.one_hot( - self.signal_proc_idx, self.nproc, dtype=self.indata.dtype - ) # (N_proc,) - # rnorm = ones + (ratio - 1) * one_hot_signal - # → (ratio at signal col, 1 elsewhere) - rnorm = col_ones + (ratio_col - 1.0) * mask[tf.newaxis, :] - - # ---- Diagnostic timing ---- - if not hasattr(self, "_n_compute_calls"): - self._n_compute_calls = 0 - self._t_compute_total = 0.0 - import atexit - - atexit.register(self._print_compute_summary) - self._n_compute_calls += 1 - self._t_compute_total += time.perf_counter() - t_start - if self._n_compute_calls % 100 == 0: - n, t = self._n_compute_calls, self._t_compute_total - print( - f"[SCETlibNPParamModel] compute() running: {n} calls, " - f"total {t:.1f}s, mean {t*1000/n:.1f} ms/call", - flush=True, - ) - return rnorm + # Build (N_reco, N_proc) scaling: 1 everywhere except the signal column, + # which carries the per-bin ratio. Broadcasting the precomputed (1, N_proc) + # one-hot avoids materializing a separate ones tensor / rebuilding one_hot. + ratio_col = tf.cast( + tf.reshape(ratio, [-1, 1]), self.indata.dtype + ) # (N_reco, 1) + rnorm = 1.0 + (ratio_col - 1.0) * self._signal_col_mask - def _print_compute_summary(self): - """Print final tally of compute() calls and time. Atexit hook.""" - n = getattr(self, "_n_compute_calls", 0) - t = getattr(self, "_t_compute_total", 0.0) - if n == 0: - return - print( - f"[SCETlibNPParamModel] compute() FINAL: {n} calls, " - f"total {t:.2f}s wall in compute(), " - f"mean {t*1000/n:.2f} ms/call", - flush=True, - ) + return rnorm diff --git a/wremnants/postprocessing/scetlib_np/response_matrix.py b/wremnants/postprocessing/scetlib_np/response_matrix.py index e258f3dd2..c80a43226 100644 --- a/wremnants/postprocessing/scetlib_np/response_matrix.py +++ b/wremnants/postprocessing/scetlib_np/response_matrix.py @@ -1,16 +1,22 @@ """Load the (reco × gen) response matrix R for the ParamModel. R lives in the unfolding histmaker output (a separate hdf5 from the fit tensor) -as ``nominal_postfsr_yieldsUnfolding`` under the Z sample group. We select -``acceptance=True``, project to the reco × (ptVGen, absYVGen) axes (summing -over helicitySig per Luca's guidance), and return a numpy array plus axis -metadata. +as ``nominal_prefsr_yieldsUnfolding`` under the Z sample group. We slice +``acceptance=True`` (gen-fiducial) and project to the reco × (ptVGen, absYVGen) +axes — which SUMS the helicitySig axis: R is filled with the weight PARTITION +``nominal_weight_helicity`` whose 8 pieces add back up, so the physical yield +is the helicitySig SUM (NOT UL). The gen-total normalizer N_gen, by contrast, +is filled with ``csAngularMoments`` and takes the UL component +(``helicitySig=-1``); see ``_select_ul_helicity`` and the inline comment in +``load_R``. The full response-fold formula and this SUM-vs-UL subtlety are +documented once in the :mod:`param_model` module docstring (single source of +truth) — consult it rather than re-deriving here. Single entry point: load_R(unfolding_hdf5_path, sample_key="Zmumu_2016PostVFP", - hist_name="nominal_postfsr_yieldsUnfolding") -> dict + hist_name="nominal_prefsr_yieldsUnfolding") -> dict """ import h5py @@ -18,14 +24,75 @@ from wums import ioutils as wums_io -DEFAULT_HIST = "nominal_postfsr_yieldsUnfolding" +# Pre-FSR gen level: the SCETlib btgrid σ_gen is the resummed *boson* qT/Y +# (QCD, before QED FSR), so the response and gen-total must also be pre-FSR for +# σ_gen, R, and N_gen to live at the same gen level. (postfsr variants exist in +# the same file — nominal_postfsr_yieldsUnfolding / "postfsr" — for comparison.) +DEFAULT_HIST = "nominal_prefsr_yieldsUnfolding" +DEFAULT_GENTOTAL = "prefsr" # xnorm gen-total denominator (pre-reco-selection) DEFAULT_SAMPLE = "Zmumu_2016PostVFP" # Axes we keep, in canonical order: reco first, then gen. RECO_AXES = ("ptll", "yll", "cosThetaStarll_quantile", "phiStarll_quantile") GEN_AXES = ("ptVGen", "absYVGen") -# Axes we collapse: acceptance via {True} slice, helicitySig via project-out. -SUM_AXES = ("helicitySig",) +# helicitySig is an angular-moment axis; we take the UL component (value -1), +# the angular-integrated total — see _select_ul_helicity. acceptance is sliced +# to True (gen-fiducial) at use. +HELICITY_AXIS = "helicitySig" + +# Gen ptVGen overflow handling. The unfolding hist's ptVGen axis ends at 44 (the +# last reco-ptll edge), but ~3.6% of the Z yield has true gen qT > 44 and +# resolution-migrates into the high-ptll reco bins (6.2% of the last reco bin). +# Dropping that gen-overflow column makes σ_reco low there. We instead fold the +# ptVGen overflow into an extra gen bin so the model can supply a σ_gen for it +# (the btgrid integral over qT ∈ (44, PTVGEN_OVERFLOW_EDGE]). The edge must be +# ≤ the btgrid qT max (the fineall grid runs to 100) and should coincide with a +# gen-histmaker ptVgen edge so the gen-level cross-check's _merge_matrix is exact +# — 100 satisfies both. (absYVGen has zero overflow: |Y| ≤ 2.5 is fully contained, +# so only ptVGen needs this.) +PTVGEN_OVERFLOW_EDGE = 100.0 + + +def _select_ul_helicity(h): + """Select the UL angular component (helicitySig = -1), if that axis exists. + + Applied to the gen-total denominator N_gen ONLY — NOT to R. N_gen is filled + with ``csAngularMoments``, a moment expansion whose A_i bins (0..7) are + signed and do NOT sum to σ; only UL (value -1) is the angular-integrated + total, so N_gen takes UL. R is the opposite: filled with the weight + PARTITION ``nominal_weight_helicity``, it is recovered by SUMMING + helicitySig (``project``), not by taking UL — see the inline comment in + ``load_R``. Same axis, opposite reduction; taking UL of R would inflate the + closure ~15×. Returning h unchanged when the axis is absent keeps + non-helicity inputs working. + """ + if HELICITY_AXIS in [a.name for a in h.axes]: + ul_idx = h.axes[HELICITY_AXIS].index(-1) + h = h[{HELICITY_AXIS: ul_idx}] + return h + + +def _append_axis_overflow(h, axis_name): + """Values array for ``h`` with ``axis_name``'s OVERFLOW bin appended as one + extra in-range bin along that axis; every other axis is taken in-range (no + flow). I.e. ``flow=False`` everywhere except that we keep ``axis_name``'s + overflow as a genuine trailing bin.""" + full = h.values(flow=True) + inr = h.values(flow=False).astype(np.float64) + idx, pos = [], None + for p, ax in enumerate(h.axes): + uf = 1 if ax.traits.underflow else 0 + if ax.name == axis_name: + pos = p + idx.append(slice(uf + ax.size, uf + ax.size + 1)) # the overflow bin + else: + idx.append(slice(uf, uf + ax.size)) # in-range only + if pos is None: + raise ValueError( + f"_append_axis_overflow: no {axis_name!r} axis in {[a.name for a in h.axes]}" + ) + over = full[tuple(idx)].astype(np.float64) + return np.concatenate([inr, over], axis=pos) def load_R( @@ -34,6 +101,8 @@ def load_R( hist_name=DEFAULT_HIST, reco_axes=RECO_AXES, gen_axes=GEN_AXES, + gen_total_name=DEFAULT_GENTOTAL, + ptVGen_overflow=True, ): """Load R from the unfolding histmaker output. @@ -44,6 +113,12 @@ def load_R( reco_shape : tuple of axis sizes (reco) gen_shape : tuple of axis sizes (gen) source : (path, sample_key, hist_name) for traceability + + ``ptVGen_overflow`` (default True): append the gen ptVGen overflow (true + qT > last gen edge) as a trailing gen bin in R and N_gen, with edge + ``PTVGEN_OVERFLOW_EDGE``. This lets the model fold a σ_gen(qT>44) through the + migration into the high-ptll reco bins (see the PTVGEN_OVERFLOW_EDGE note); + set False for the legacy in-range-only response. """ with h5py.File(unfolding_hdf5_path, "r") as f: if sample_key not in f: @@ -72,29 +147,84 @@ def load_R( # Sanity-check the axes. ax_names = [a.name for a in h.axes] - required = set(reco_axes) | set(gen_axes) | {"acceptance"} | set(SUM_AXES) + required = set(reco_axes) | set(gen_axes) | {"acceptance", HELICITY_AXIS} missing = required - set(ax_names) if missing: raise ValueError( f"{hist_name}: missing expected axes {missing}. " f"Got: {ax_names}" ) - # Select acceptance=True (gen events in fiducial), keep gen + reco axes, - # sum over helicitySig. + # Select acceptance=True (gen events in fiducial), then keep the reco + + # gen axes. project() SUMS helicitySig out — which is correct *for R*: + # the joint yield is filled with `nominal_weight_helicity` + # (= nominal_weight × helWeight_tensor, see helicity_utils), a PARTITION + # of the event weight into the 8 helicity pieces g_i(cosθ,φ) that ADD + # BACK UP to the full angular weight. So Σ_helicitySig R is the physical + # angular-resolved reco×gen yield (the angular dependence lives in the + # cosThetaStar*/phiStar* reco bins). NB this is the OPPOSITE reduction + # from N_gen below: N_gen is filled with `csAngularMoments` (a moment + # expansion whose 0..7 bins do NOT sum to σ), so it takes UL (-1). Same + # axis, different fill tensor → different recovery. Taking UL of R would + # discard the angular partition and inflate the closure (~15×). h_sel = h[{"acceptance": True}] h_proj = h_sel.project(*reco_axes, *gen_axes) - # .project sums over un-listed axes (i.e. helicitySig here). - # Out from the with-block: hist is materialized. - R = h_proj.values(flow=False).astype(np.float64) + # Gen-total denominator N_gen(g): the xnorm histogram (e.g. "postfsr"), + # filled on fiducial gen events BEFORE reco selection. Its gen marginal + # is the generated total per gen bin (efficiency NOT yet applied), so + # N_reco(b,g)/N_gen(g) = efficiency × migration — the theory-independent + # gen→reco response. (The gen marginal of R itself is reco-passing, i.e. + # already × efficiency, which is the wrong normalizer.) + N_gen_hist = None + if gen_total_name is not None and gen_total_name in output: + gp = output[gen_total_name] + hg = gp.get() if hasattr(gp, "get") else gp + # postfsr/prefsr axes: (count, ptVGen, absYVGen, helicitySig). This + # one is filled with `csAngularMoments`: a moment expansion whose UL + # bin (-1) is the angular-integrated total σ, while bins 0..7 are the + # A_i moments (orthogonal, signed) that do NOT sum to σ — summing them + # overcounts by ~19% here. So take UL. (R above is the opposite: a + # weight partition, recovered by SUMMING helicitySig.) Then project + # to the gen axes (sums the trivial 'count' axis) to match R's binning. + hg = _select_ul_helicity(hg) + hg_gen = hg.project(*gen_axes) + # Fold the gen ptVGen overflow into a trailing bin (or drop it). + N_gen_hist = ( + _append_axis_overflow(hg_gen, "ptVGen") + if ptVGen_overflow + else hg_gen.values(flow=False).astype(np.float64) + ) + + # Out from the with-block: hist is materialized. Keep the gen ptVGen overflow + # as a trailing gen bin so the model can supply σ_gen(qT>44) — else the high- + # ptll reco bins (fed by true qT>44 migrating down) come out low. + R = ( + _append_axis_overflow(h_proj, "ptVGen") + if ptVGen_overflow + else h_proj.values(flow=False).astype(np.float64) + ) + + def _gen_edges(name): + e = np.asarray(h_proj.axes[name].edges, dtype=np.float64) + if ptVGen_overflow and name == "ptVGen": + e = np.concatenate([e, [PTVGEN_OVERFLOW_EDGE]]) # (44, 100] overflow bin + return e reco_meta = [(name, h_proj.axes[name].edges) for name in reco_axes] - gen_meta = [(name, h_proj.axes[name].edges) for name in gen_axes] - reco_shape = tuple(h_proj.axes[name].size for name in reco_axes) - gen_shape = tuple(h_proj.axes[name].size for name in gen_axes) + gen_meta = [(name, _gen_edges(name)) for name in gen_axes] + # Derive shapes from the (possibly overflow-extended) arrays, not the hist. + reco_shape = R.shape[: len(reco_axes)] + gen_shape = R.shape[len(reco_axes) :] + + if N_gen_hist is not None and N_gen_hist.shape != gen_shape: + raise ValueError( + f"gen-total {gen_total_name!r} shape {N_gen_hist.shape} != R gen " + f"shape {gen_shape}; gen binning mismatch." + ) return dict( R=R, + N_gen=N_gen_hist, reco_axes=reco_meta, gen_axes=gen_meta, reco_shape=reco_shape, From 5b062958a86fe4ab253d76b8ead4dca5d3e63fc9 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Mon, 8 Jun 2026 15:49:49 -0400 Subject: [PATCH 05/31] better logs --- wremnants/postprocessing/scetlib_np/btgrid_cache.py | 5 ++++- wremnants/postprocessing/scetlib_np/param_model.py | 3 ++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_cache.py b/wremnants/postprocessing/scetlib_np/btgrid_cache.py index db836e42f..675ee5cc4 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_cache.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_cache.py @@ -64,7 +64,10 @@ def load(submitdir, rebuild=False, verbose=True): with open(combined, "rb") as f: grid = pickle.load(f) if verbose: - print(f"[btgrid_cache] loaded combined pickle in {time.time()-t0:.1f}s") + print( + f"[btgrid_cache] loaded combined pickle in {time.time()-t0:.1f}s", + flush=True, + ) return grid if not shards: diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 725b01cc8..fe1998468 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -784,7 +784,8 @@ def __init__( i = self._param_order.index(name) defaults[i] = float(val) print( - f"[SCETlibNPParamModel] xparamdefault overridden: {dict(zip(self._param_order, defaults))}" + f"[SCETlibNPParamModel] xparamdefault overridden: {dict(zip(self._param_order, defaults))}", + flush=True, ) # rabbit's set_param_default expects an internal-storage convention # where POIs (npoi entries) are SQRT(value) if not allowNegativeParam. From 34e40c24face196606f165d91ea092e7466c7189 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Tue, 9 Jun 2026 11:03:51 -0400 Subject: [PATCH 06/31] set up default inputs --- .../postprocessing/scetlib_np/param_model.py | 227 ++++++++++++++++-- 1 file changed, 210 insertions(+), 17 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index fe1998468..54a4f031c 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -169,6 +169,61 @@ SCETlib correction pkl (see :mod:`scetlib_lambda_central`). The np_model and np_model_nu strings are fixed at construction (from λ_central). All λ values are TF Variables — differentiable in the fit. + +═════════════════════════════════════════════════════════════════════════════ +Getting the postfit Hessian / covariance (uncertainties on λ) +═════════════════════════════════════════════════════════════════════════════ + +The fit floats λ fine, but rabbit's postfit covariance step +``loss_val_grad_hess`` → ``t2.jacobian(grad, x)`` differentiates through the bT +fold once per fit parameter (~3754). pfor cannot see that the fold depends on +the parameters only through the ≤8 λ, so it re-materializes the internal +(Ng × Nbt) ≈ 8.75 GB bT slab for EVERY parameter → ~33 TB → OOM. + +Fix (this module): a "straight-through" surrogate that keeps the exact ratio +VALUE but exposes only a compact quadratic in the ≤8 λ to autodiff: + + ratio~(λ) = stop_gradient(ratio) + J·d + ½ dᵀ K d , d = λ − stop_gradient(λ) + +with J = dratio/dλ ([Nreco, nλ]) and K = d²ratio/dλ² ([Nreco, nλ, nλ]) computed +by forward-mode AD (≤8 / ≤64 fold passes, NOT tiled). d is identically 0 so the +value is unchanged, but ∂d/∂λ = I, so ∂ratio~/∂λ = J and ∂²ratio~/∂λ² = K: +rabbit's jacobian gets the exact derivatives while the big slab stays inside +stop_gradient and never enters the differentiated graph (33 TB → a few MB). +Implemented in ``_ratio_straightthrough`` (+ ``_ratio_compact_jac``, +``_ratio_compact_hess``); selected in ``compute()`` by env flags. Full +derivation + validation in ``HESSIAN_PLAN.md``. + +Two-pass recipe (rabbit still computes the Hessian; NO rabbit changes): + + 1. Fit, no Hessian → postfit: + rabbit_fit.py --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel \ + --noHessian -o fit/ + + 2. Covariance at that postfit (no refit), straight-through ON: + SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH=1 SCETLIB_NP_HESSIAN_GN=1 \ + rabbit_fit.py --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel \ + --externalPostfit fit/fitresults.hdf5 --externalPostfitResult nominal \ + --noFit -t 0 --pseudoData nominal -o cov/ + (do NOT pass --noHessian; do NOT pass --eager.) ~5 min, no OOM. + Uncertainties are sqrt(diag(cov)) in cov/fitresults.hdf5 (results_nominal). + +Env flags (both OFF by default → the fit path is unchanged): + SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH=1 use the straight-through path in compute() + SCETLIB_NP_HESSIAN_GN=1 Gauss-Newton: keep J only, drop the K term + +GN vs full-K. The Poisson Hessian is + H_ij = Σ_b [ (n_b/μ_b²) J_bi J_bj + (1 − n_b/μ_b) K_bij ]. +For ASIMOV data the residual (1 − n/μ) = 0, so the K term vanishes and GN (J +only) is EXACT. For real/toy data K is needed for the exact Hessian — but the +nested-forward-mode K (``_ratio_compact_hess``) currently CRASHES under rabbit's +@tf.function (the tf.where in ``_safe_div``) and is impractically slow under +--eager, so GN (the recipe above) is the working production path today. See +HESSIAN_PLAN.md §9 for the open K item and its fixes. + +WARNING: keep SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH UNSET during the FIT. The +surrogate recomputes J(/K) on every compute() call — fine for the one-shot +covariance pass, but it would cripple the minimizer (many gradient/HVP evals). """ import json @@ -204,6 +259,51 @@ "results_z-2d-nnlo-vj-CT18ZNNLO-{scale}-scetlibmatch.txt", ) +# Default fit inputs, so collaborators can run +# ``--paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel`` with no +# extra positional args. Passing the two positional args explicitly still wins. +# +# The unfolding response lives in wremnants-data (a skim of only the two hists +# load_R reads — ``nominal_prefsr_yieldsUnfolding`` + ``prefsr`` — 6.8 GB → 211 MB; +# made with scripts/inspect/open_narf_h5py.py). Resolved via wrem_common.data_dir. +_UNFOLDING_HDF5_DEFAULT = os.path.join( + wrem_common.data_dir, + "TheoryCorrections", + "scetlib_np", + "mz_dilepton_unfolding_R_skim.hdf5", +) +# The bT-grid combined pickle is ~17.5 GB (two [Nbins×Nbt] float64 slabs: I_pert +# and C_nu) — too large for wremnants-data/git-LFS — so it lives on the shared +# data area next to NanoAOD, as a sibling of NanoAOD under the data base. We +# resolve the NanoAOD base by REPLICATING dataset_tools.getDataPath()'s hostname +# logic rather than importing it: that module does ``import ROOT`` / ``import narf`` +# at module scope (and wremnants.production.__init__ runs ROOT.gInterpreter + +# narf.clingutils.Load), which segfaults when imported mid-fit (after TF is up). +# Only the subMIT copy exists today — pass btgrid_dir explicitly at other sites. +_BTGRID_SUBDIR = ("scetlib_np", "Z_COM13_CT18Z_N3p0LL_btgrid_fineall") + + +def _nano_data_base(): + """NanoAOD base dir per host — mirrors dataset_tools.getDataPath() WITHOUT + importing it (that import pulls ROOT/narf; see note above). Falls back to the + subMIT path on unknown hosts (the only site with the grid today).""" + import socket + + hostname = socket.gethostname() + if hostname.endswith(".cern.ch"): + return "/scratch/shared/NanoAOD" + if hostname == "cmsanalysis.pi.infn.it": + return "/scratchnvme/wmass/NANOV9/postVFP" + if hostname == "cmsasymow.pi.infn.it": + return "/scratch/wmass/y2016" + # .mit.edu and any unknown host -> subMIT shared scratch + return "/scratch/submit/cms/wmass/NanoAOD" + + +def _default_btgrid_dir(): + return os.path.join(os.path.dirname(_nano_data_base()), *_BTGRID_SUBDIR) + + # Ordered list of the v1 continuous λ. CS-side first, then TMD-effective. GNU_PARAMS = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") @@ -419,8 +519,8 @@ class SCETlibNPParamModel(ParamModel): def __init__( self, indata, - unfolding_hdf5_path: str, - btgrid_dir: str, + unfolding_hdf5_path: Optional[str] = None, + btgrid_dir: Optional[str] = None, lambda_central: Optional[Mapping] = None, signal_proc: str = "Zmumu", Q_lo: float = 60.0, @@ -443,8 +543,15 @@ def __init__( Path to the upstream histmaker output containing ``nominal_prefsr_yieldsUnfolding`` for R (and the ``prefsr`` xnorm hist for N_gen) — see :mod:`response_matrix` for the defaults. + Defaults (when None) to the skim shipped in wremnants-data + (``_UNFOLDING_HDF5_DEFAULT``); pass explicitly to override. btgrid_dir - Directory of the SCETlib bT-grid shards (fineall). + Directory holding the SCETlib bT-grid ``combined_btgrid.pkl``. + Defaults (when None) to the shared data-area copy next to NanoAOD + (``_default_btgrid_dir()``, which mirrors + ``dataset_tools.getDataPath()``'s per-host logic without importing it — + that import pulls ROOT/narf and segfaults mid-fit); pass explicitly at + non-subMIT sites. lambda_central Dict with two sub-dicts ``eff_params`` and ``gnu_params`` (same shape as returned by :func:`scetlib_lambda_central.read_lambda_central`). @@ -482,6 +589,15 @@ def __init__( self.indata = indata # ---- Double-counting guard + # Resolve default inputs so collaborators can run with no extra + # --paramModel positional args (explicit args still override). The + # unfolding skim ships in wremnants-data; the big bT-grid lives on the + # shared data area next to NanoAOD (see _default_btgrid_dir). + if unfolding_hdf5_path is None: + unfolding_hdf5_path = _UNFOLDING_HDF5_DEFAULT + if btgrid_dir is None: + btgrid_dir = _default_btgrid_dir() + # If the histmaker baked discrete NP κ-template variations into the # input HDF5, those systs describe the same physics as our continuous # λ POUs. Running both → double-counting (the discrete syst absorbs @@ -1031,6 +1147,82 @@ def _unpack_params(self, param): raise KeyError(name) return eff_params, gnu_params + # ========================================================================= + # ratio(λ) and the straight-through compact-derivative path (Hessian Phase B) + # ========================================================================= + + def _ratio_from_param(self, param): + """λ (full param vector) → floored per-reco-bin ratio, shape (N_reco,). + + The differentiable map that the straight-through Hessian path wraps. The + soft positivity floor (see ``compute``) lives here so both the normal and + straight-through paths apply it identically. + """ + eff_params, gnu_params = self._unpack_params(param) + sigma_gen = self._sigma_gen_at(eff_params, gnu_params) + gen_flat = tf.reshape(sigma_gen, [-1]) + sigma_reco = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) + ratio = sigma_reco / self.sigma_reco_central # (N_reco,) + scale = tf.constant(RATIO_FLOOR_SCALE, dtype=ratio.dtype) + ratio = tf.maximum( + scale * tf.math.softplus(ratio / scale), + tf.constant(RATIO_FLOOR_MIN, dtype=ratio.dtype), + ) + return ratio + + def _ratio_compact_jac(self, param): + """J = d(ratio)/d(param), shape (N_reco, nparam), via forward-mode AD. + + One JVP per parameter (nparam ≤ 8), each a single bT-fold pass — NOT + tiled over params. This is the compact object the Hessian actually needs + from the fold (see HESSIAN_PLAN.md §2).""" + n = int(param.shape[0]) + cols = [] + for i in range(n): + tangent = tf.one_hot(i, n, dtype=param.dtype) + with tf.autodiff.ForwardAccumulator(param, tangent) as acc: + r = self._ratio_from_param(param) + cols.append(acc.jvp(r)) # (N_reco,) = dratio/dparam_i + return tf.stack(cols, axis=1) # (N_reco, nparam) + + def _ratio_compact_hess(self, param): + """K = d²(ratio)/d(param)², shape (N_reco, nparam, nparam), forward-over-forward. + + nparam² JVP-of-JVP passes (≤ 64), each one bT-fold pass — never tiled.""" + n = int(param.shape[0]) + rows = [] + for i in range(n): + ti = tf.one_hot(i, n, dtype=param.dtype) + cols = [] + for j in range(n): + tj = tf.one_hot(j, n, dtype=param.dtype) + with tf.autodiff.ForwardAccumulator(param, tj) as acc_j: + with tf.autodiff.ForwardAccumulator(param, ti) as acc_i: + r = self._ratio_from_param(param) + di = acc_i.jvp(r) # (N_reco,) + cols.append(acc_j.jvp(di)) # (N_reco,) = d²ratio/dparam_i dparam_j + rows.append(tf.stack(cols, axis=1)) # (N_reco, nparam) over j + return tf.stack(rows, axis=2) # (N_reco, j, i) — symmetric in (i, j) + + def _ratio_straightthrough(self, param, use_curvature=True): + """Exact ratio value, but autodiff sees only a compact quadratic in the + ≤8 λ (J, and optionally K) — so the (N_grid, N_bt) bT slab never enters + the differentiated graph and rabbit's covariance jacobian does not OOM. + + At the evaluation point (d = 0) the value is exact, the 1st derivative is + J, and the 2nd derivative is K. ``use_curvature=False`` keeps only J + (Gauss-Newton / Fisher — exact for Asimov data, 8 vs 72 fold passes). + """ + val = self._ratio_from_param(param) + J = tf.stop_gradient(self._ratio_compact_jac(param)) # (N_reco, nparam) + d = param - tf.stop_gradient(param) # value 0, unit gradient + d = tf.cast(d, J.dtype) + out = tf.stop_gradient(val) + tf.linalg.matvec(J, d) + if use_curvature: + K = tf.stop_gradient(self._ratio_compact_hess(param)) # (N_reco, n, n) + out = out + 0.5 * tf.einsum("rij,i,j->r", K, d, d) + return out + def compute(self, param, full=False): """Return per-(bin, proc) scaling tensor. @@ -1049,20 +1241,21 @@ def compute(self, param, full=False): is a numerical safety net, not a physics constraint — it does not stop the fit from exploring negative-λ; it just keeps that exploration finite. """ - eff_params, gnu_params = self._unpack_params(param) - sigma_gen = self._sigma_gen_at(eff_params, gnu_params) - # Fold σ_gen(λ) through the normalized migration response P (self.R). - gen_flat = tf.reshape(sigma_gen, [-1]) - sigma_reco = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) - ratio = sigma_reco / self.sigma_reco_central # (N_reco,) - # Soft positivity floor (see docstring): scale·softplus(r/scale) ≈ r for - # healthy r and → 0⁺ smoothly for r ≤ 0; the max() with a tiny ground - # keeps it strictly positive even when softplus underflows (r ~ -1e43). - scale = tf.constant(RATIO_FLOOR_SCALE, dtype=ratio.dtype) - ratio = tf.maximum( - scale * tf.math.softplus(ratio / scale), - tf.constant(RATIO_FLOOR_MIN, dtype=ratio.dtype), - ) + # ratio(λ) per reco bin. Normal path = exact fold (used for the fit). + # Straight-through path (Hessian-only Phase B) keeps the bT slab off the + # autodiff graph so rabbit's covariance jacobian doesn't OOM. Toggled by + # SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH; SCETLIB_NP_HESSIAN_GN=1 drops the + # curvature term (Gauss-Newton/Fisher — exact for Asimov, 8 vs 72 passes). + # Do NOT enable during the fit: it recomputes J(/K) every call. + if not hasattr(self, "_hess_st"): + self._hess_st = bool( + os.environ.get("SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH", "").strip() + ) + if self._hess_st: + gn = bool(os.environ.get("SCETLIB_NP_HESSIAN_GN", "").strip()) + ratio = self._ratio_straightthrough(param, use_curvature=not gn) + else: + ratio = self._ratio_from_param(param) # Build (N_reco, N_proc) scaling: 1 everywhere except the signal column, # which carries the per-bin ratio. Broadcasting the precomputed (1, N_proc) From d4f73dac78bf11eb3d159f9b68158c1aec841868 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Thu, 11 Jun 2026 09:44:16 -0400 Subject: [PATCH 07/31] bump rabbit --- rabbit | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rabbit b/rabbit index 843fbbd0a..2fa3ae52a 160000 --- a/rabbit +++ b/rabbit @@ -1 +1 @@ -Subproject commit 843fbbd0a9bbf48ab7726e0cc9cb0a9e5629616a +Subproject commit 2fa3ae52a4f67bad6d8e8bb9a585df7e2beb4244 From c31025a1d81576c1a3155869a2685eefe5d66c29 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Thu, 11 Jun 2026 09:48:17 -0400 Subject: [PATCH 08/31] add NP to alphaS impacts --- wremnants/utilities/styles/styles.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wremnants/utilities/styles/styles.py b/wremnants/utilities/styles/styles.py index e0a11ba33..3052c0170 100644 --- a/wremnants/utilities/styles/styles.py +++ b/wremnants/utilities/styles/styles.py @@ -272,6 +272,7 @@ def translate_html_to_latex(n): "ZmassAndWidth", "massAndWidth", "normXsecZ", + "resumNonpert", ] nuisance_grouping = { "super": [ From c898f544f4913f1bed36f8e8ea0530f43a3f717e Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Thu, 11 Jun 2026 11:06:42 -0400 Subject: [PATCH 09/31] bump submodule --- wremnants-data | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wremnants-data b/wremnants-data index 3d2b2b214..81ff2c21c 160000 --- a/wremnants-data +++ b/wremnants-data @@ -1 +1 @@ -Subproject commit 3d2b2b2140751b6212898a923ede103793d89ee2 +Subproject commit 81ff2c21cffae7cb0c853c17484ed6c745ae36ae From 8deafe0365c042362e66667cf72da4a0c8d3cf3d Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 09:50:21 -0400 Subject: [PATCH 10/31] udpates on SCETlib NP param model --- .../postprocessing/scetlib_np/btgrid_tf.py | 270 ++++++- .../scetlib_np/lambda_central.py | 78 +- .../postprocessing/scetlib_np/param_model.py | 714 +++++++++--------- .../production/datasets/dataset_tools.py | 181 +---- wremnants/utilities/data_paths.py | 204 +++++ 5 files changed, 935 insertions(+), 512 deletions(-) create mode 100644 wremnants/utilities/data_paths.py diff --git a/wremnants/postprocessing/scetlib_np/btgrid_tf.py b/wremnants/postprocessing/scetlib_np/btgrid_tf.py index 00edf32f8..10672ab65 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_tf.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_tf.py @@ -117,10 +117,31 @@ def simpson_tf(y, weights): } +def _frozen_eq_zero(x): + """``x == 0`` as a gradient-frozen boolean condition for the NP-factor masks. + + The comparison is a non-differentiable, measure-zero boundary that the + surrounding ``tf.where`` never differentiates through, so freezing its input + leaves both the value and every derivative unchanged. But it is REQUIRED for + the full-K Hessian: the straight-through ``K`` path nests two + ``ForwardAccumulator``s (forward-over-forward AD), and building the JVP of an + ``Equal`` op that receives a tangent-carrying input raises + ``IndexError: list index out of range`` under ``@tf.function`` (a TF + nested-forward-mode bug). With the input frozen no tangent ever reaches the + comparison, so the bug never triggers; the ``tf.where``'s own JVP with a + constant condition is fine. (The earlier GN/J-only path never hit this — it + uses a single ``ForwardAccumulator``, for which ``Equal``'s JVP is fine.)""" + return tf.equal(tf.stop_gradient(x), 0) + + def _safe_div(num, den): - """``num / den`` with the denominator clamped away from zero, masked by - ``tf.where`` at the call site. Keeps gradients finite.""" - den_safe = tf.where(tf.equal(den, 0), tf.ones_like(den), den) + """``num / den`` with the denominator clamped to 1 where it is exactly zero. + + Equivalent to ``num / tf.where(den == 0, 1, den)``; the comparison input is + frozen (see :func:`_frozen_eq_zero`) so the full-K nested forward-mode + Hessian does not crash under ``@tf.function``. Keeps gradients finite; the + den==0 result is masked away by the caller's final ``tf.where``.""" + den_safe = tf.where(_frozen_eq_zero(den), tf.ones_like(den), den) return num / den_safe @@ -175,7 +196,9 @@ def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_ raise ValueError(f"F_eff_tf: unsupported np_model {np_model!r}") full = tf.exp(-2.0 * lambda_inf * bT * func) - return tf.where(tf.equal(lambda_inf, 0), tf.ones_like(full), full) + # lambda_inf == 0 -> 1 (NP off). Frozen comparison input so the full-K nested + # forward-mode Hessian doesn't crash under @tf.function (see _frozen_eq_zero). + return tf.where(_frozen_eq_zero(lambda_inf), tf.ones_like(full), full) def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): @@ -216,7 +239,9 @@ def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): raise ValueError(f"gamma_nu_NP_tf: unsupported np_model_nu {np_model_nu!r}") full = -lambda_inf_nu * func - return tf.where(tf.equal(lambda_inf_nu, 0), tf.zeros_like(full), full) + # lambda_inf_nu == 0 -> 0 (NP off). Frozen comparison input so the full-K nested + # forward-mode Hessian doesn't crash under @tf.function (see _frozen_eq_zero). + return tf.where(_frozen_eq_zero(lambda_inf_nu), tf.zeros_like(full), full) # ============================================================================= @@ -316,8 +341,243 @@ def build_bT_J0_kernel(qT_per_bin, bT): λ-independent — call once at ParamModel construction and pass into :func:`reconstruct_batch_tf` as ``bT_J0_kernel``. + + Note: in the factorized path (:func:`reconstruct_batch_factorized_tf`) + this is instead called with ``qT_unique`` (NqT distinct values, not the + per-bin expansion), giving a (NqT, Nbt) kernel — same numbers, ~4000× + smaller. """ qT_per_bin = _as_dtype(qT_per_bin) bT = _as_dtype(bT) arg = qT_per_bin[:, tf.newaxis] * bT[tf.newaxis, :] return bT[tf.newaxis, :] * tf.math.special.bessel_j0(arg) + + +# ============================================================================= +# Factorized reconstruction (GPU-memory-safe; exact) +# ============================================================================= +# +# The (Nbins, Nbt) layout of reconstruct_batch_tf needs several ~9 GB fp64 +# tensors (Nbins=546840, Nbt=2000) and OOMs a 32 GB GPU at construction. Two +# exact observations shrink it (see scetlib_np/FACTORIZED_RECON.md): +# +# 1. qT enters the λ-dependent integrand ONLY through the bT·J0(qT·bT) +# kernel: the kernel needs the NqT *unique* qT values, not Nbins rows. +# 2. SCETlib's profile scales are piecewise in x = qT/Q and exactly +# canonical (qT-independent) below the first transition point x1·Q, so +# the cached I_pert / C_nu rows are BIT-IDENTICAL across qT there. The +# dedup below discovers identical rows dynamically (byte-wise hashing + +# full verification) — no assumption about profiles or qT ranges. +# +# The bT-Simpson reduction then becomes a (Nu, Nbt) @ (Nbt, NqT) matmul plus +# a per-bin gather. Same integrand, same Simpson weights, same sampling; only +# the floating-point grouping/summation order changes (≲1e-14 relative). + + +def dedup_grid_rows(I_pert, C_nu, feff_idx_per_bin, verbose=True): + """Find bit-identical (I_pert, C_nu, F_eff-index) row triples. + + Construction-time numpy helper (runs once, CPU). Rows are keyed by the + raw bytes of the I_pert row, the C_nu row and the per-bin F_eff Y-index, + so two bins share a unique id iff their λ-dependent integrand columns are + bit-for-bit identical for EVERY λ. The grouping is verified by direct + array comparison afterwards (no reliance on hash collision odds). + + Parameters + ---------- + I_pert, C_nu : (Nbins, Nbt) float64 ndarrays + feff_idx_per_bin : (Nbins,) int ndarray + Index into the unique-Y table used for the F_eff gather (Y enters + F_eff via delta_lambda2·Y²; keying on it keeps the per-unique-row + F_eff well-defined even when delta_lambda2 floats). + + Returns + ------- + dict with: + I_u, C_u : (Nu, Nbt) deduplicated rows (copies, C-contiguous) + row_uid : (Nbins,) int32, bin -> unique-row index + feff_idx_u : (Nu,) int32, unique row -> unique-Y index + n_unique : int + C_uu : (Ncu, Nbt) second-level dedup of the C_u rows. C_nu is the + ν-evolution coefficient — it depends on (Q, profile-qT) + only, not Y, so its standalone unique-row count is ~150x + smaller than Nu (1888 vs 284605 on the fineall grid). The + exp(C·g) transcendentals run on these rows and are + gathered back — bit-identical, ~150x fewer exp() calls, + and the (Nu, Nbt) C constant never needs to exist on + device. + c_of_u : (Nu,) int32, unique row -> C_uu row index + n_unique_C : int + """ + import hashlib + + I_pert = np.ascontiguousarray(I_pert) + C_nu = np.ascontiguousarray(C_nu) + n_bins = I_pert.shape[0] + feff_idx_per_bin = np.asarray(feff_idx_per_bin).reshape(-1) + if feff_idx_per_bin.shape[0] != n_bins: + raise ValueError("feff_idx_per_bin length != Nbins") + + seen = {} + row_uid = np.empty(n_bins, dtype=np.int32) + rep_rows = [] # representative bin index per unique id + for k in range(n_bins): + h = hashlib.blake2b(I_pert[k].tobytes(), digest_size=16) + h.update(C_nu[k].tobytes()) + h.update(int(feff_idx_per_bin[k]).to_bytes(4, "little", signed=True)) + key = h.digest() + uid = seen.get(key) + if uid is None: + uid = len(rep_rows) + seen[key] = uid + rep_rows.append(k) + row_uid[k] = uid + + rep_rows = np.asarray(rep_rows, dtype=np.int64) + I_u = np.ascontiguousarray(I_pert[rep_rows]) + C_u = np.ascontiguousarray(C_nu[rep_rows]) + feff_idx_u = feff_idx_per_bin[rep_rows].astype(np.int32) + + # Verify the grouping bit-exactly: every bin's rows must equal its + # representative's. Chunked to bound the temporary gather copies. + chunk = 20000 + for k0 in range(0, n_bins, chunk): + k1 = min(k0 + chunk, n_bins) + sel = row_uid[k0:k1] + if not ( + np.array_equal(I_pert[k0:k1], I_u[sel]) + and np.array_equal(C_nu[k0:k1], C_u[sel]) + and np.array_equal(feff_idx_per_bin[k0:k1], feff_idx_u[sel]) + ): + raise AssertionError( + f"dedup_grid_rows: hash grouping failed verification in bins " + f"[{k0}, {k1}) — this should be impossible; grid corrupt?" + ) + + # Second-level dedup of the C rows (qT-and-Y-independent below the + # profile transition AND Y-independent everywhere → ~150x smaller). + n_u = len(rep_rows) + seen_c = {} + c_of_u = np.empty(n_u, dtype=np.int32) + rep_c = [] + for k in range(n_u): + key = hashlib.blake2b(C_u[k].tobytes(), digest_size=16).digest() + cid = seen_c.get(key) + if cid is None: + cid = len(rep_c) + seen_c[key] = cid + rep_c.append(k) + c_of_u[k] = cid + C_uu = np.ascontiguousarray(C_u[np.asarray(rep_c, dtype=np.int64)]) + for k0 in range(0, n_u, chunk): + k1 = min(k0 + chunk, n_u) + if not np.array_equal(C_u[k0:k1], C_uu[c_of_u[k0:k1]]): + raise AssertionError( + f"dedup_grid_rows: C-row sub-dedup failed verification in rows " + f"[{k0}, {k1}) — this should be impossible; grid corrupt?" + ) + + if verbose: + print( + f"[dedup_grid_rows] {n_bins} bins -> {n_u} unique rows " + f"({n_bins / n_u:.2f}x dedup, verified bit-exact); " + f"per-tensor {n_bins * I_pert.shape[1] * 8 / 1e9:.2f} GB -> " + f"{n_u * I_pert.shape[1] * 8 / 1e9:.2f} GB; " + f"C_nu sub-dedup {n_u} -> {len(rep_c)} rows " + f"({n_u / len(rep_c):.0f}x, verified bit-exact)", + flush=True, + ) + return dict( + I_u=I_u, + C_u=C_u, + row_uid=row_uid, + feff_idx_u=feff_idx_u, + n_unique=n_u, + C_uu=C_uu, + c_of_u=c_of_u, + n_unique_C=len(rep_c), + ) + + +def reconstruct_batch_factorized_tf( + b_bar, + I_pert_u, + C_nu_u=None, + eff_params: Mapping = None, + gnu_params: Mapping = None, + *, + np_model: str, + np_model_nu: str, + KwqT, + gather_idx, + Y_unique, + feff_idx_u, + C_nu_uu=None, + c_of_u=None, +): + """Memory-factorized, numerically-equivalent form of + :func:`reconstruct_batch_tf`. + + Evaluates σ_i = qT_i Σ_b w_b·bT_b·J0(qT_i bT_b)·I_{u(i),b}·exp(C_{u(i),b} + g_b)·F_{y(u(i)),b} as a (Nu, Nbt) elementwise block, a matmul against the + weighted J0 kernel on the unique-qT grid, and a per-bin gather — no + (Nbins, Nbt) tensor is ever materialized. Same integrand, weights and + sampling as :func:`reconstruct_batch_tf`; only the floating-point + multiplication grouping and summation order differ (≲1e-14 relative — + see the parity script :mod:`scetlib_np_factorized_parity`). + + Parameters + ---------- + b_bar : (Nbt,) — b*(bT), the NP-factor argument + I_pert_u : (Nu, Nbt) — deduplicated grid rows + (from :func:`dedup_grid_rows`) + C_nu_u : (Nu, Nbt), optional + Per-unique-row C_ν. Either this OR (``C_nu_uu``, ``c_of_u``) must be + given; the latter is preferred (~150x fewer exp() calls and no + (Nu, Nbt) C constant on device — bit-identical results, since exp of + identical rows is identical and the gather only replicates rows). + KwqT : (NqT, Nbt) + ``qT_u · bT · J0(qT_u·bT) · w_simpson`` on the unique-qT grid; the + per-bin qT prefactor of reconstruct_batch_tf is folded in here. + gather_idx : (Nbins, 2) int32 — per bin ``[u(i), qT_index(i)]`` + Y_unique : (NY,) — unique Y values for the F_eff evaluation + feff_idx_u : (Nu,) int32 — unique row -> Y_unique index + C_nu_uu : (Ncu, Nbt), optional — second-level deduplicated C_ν rows + c_of_u : (Nu,) int32, optional — unique row -> C_uu row index + """ + b_bar = _as_dtype(b_bar) + I_pert_u = _as_dtype(I_pert_u) + KwqT = _as_dtype(KwqT) + + g_NP = gamma_nu_NP_tf(b_bar, **gnu_params, np_model_nu=np_model_nu) # (Nbt,) + if C_nu_uu is not None and c_of_u is not None: + # exp on the ~150x smaller C-row table, replicated by the gather. + exp_g_uu = tf.exp(_as_dtype(C_nu_uu) * g_NP[tf.newaxis, :]) # (Ncu, Nbt) + exp_g_u = tf.gather(exp_g_uu, c_of_u) # (Nu, Nbt) + elif C_nu_u is not None: + exp_g_u = tf.exp(_as_dtype(C_nu_u) * g_NP[tf.newaxis, :]) # (Nu, Nbt) + else: + raise ValueError( + "reconstruct_batch_factorized_tf: pass either C_nu_u or " + "(C_nu_uu, c_of_u)" + ) + + delta_l2_in = eff_params.get("delta_lambda2", 0.0) + if isinstance(delta_l2_in, (int, float)) and float(delta_l2_in) == 0.0: + # static fast path: F_eff has no Y dependence (matches the + # reconstruct_batch_tf fast path bit-for-bit on the unique rows) + Feff = F_eff_tf(0.0, b_bar, **eff_params, np_model=np_model) # (Nbt,) + Feff_u = Feff[tf.newaxis, :] # broadcast over Nu + else: + # F_eff on the unique-Y rows, gathered to the unique grid rows. Same + # unique-Y trick as reconstruct_batch_tf — the gather replicates rows + # bit-exactly and scatter-adds cotangents in the backward pass. + Y_u = _as_dtype(Y_unique)[:, tf.newaxis] # (NY, 1) + Feff_rows = F_eff_tf( + Y_u, b_bar[tf.newaxis, :], **eff_params, np_model=np_model + ) # (NY, Nbt) + Feff_u = tf.gather(Feff_rows, feff_idx_u) # (Nu, Nbt) + + M = I_pert_u * exp_g_u * Feff_u # (Nu, Nbt) + S = tf.matmul(M, KwqT, transpose_b=True) # (Nu, NqT) + return tf.gather_nd(S, gather_idx) # (Nbins,) diff --git a/wremnants/postprocessing/scetlib_np/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py index d094eeec4..c5b638c59 100644 --- a/wremnants/postprocessing/scetlib_np/lambda_central.py +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -11,6 +11,7 @@ read_lambda_central(hdf5_path, proc="Z") -> dict """ +import json import os import pickle import sys @@ -117,9 +118,59 @@ def read_lambda_central_from_meta(meta, proc="Z", _source=""): return _resolve_tag_to_lambda(tag, proc) +def load_lambda_central_file(path): + """Load a λ_central override from a JSON or YAML file. + + The file must decode to a dict with ``eff_params`` and ``gnu_params`` + sub-dicts (same shape as :func:`read_lambda_central`). Format is chosen + by extension (``.yaml``/``.yml`` → YAML, else JSON); YAML's loader also + accepts JSON, so this is forgiving either way. + """ + if not os.path.exists(path): + raise FileNotFoundError(f"λ_central file missing: {path!r}") + with open(path) as f: + text = f.read() + if path.lower().endswith((".yaml", ".yml")): + import yaml + + data = yaml.safe_load(text) + else: + try: + data = json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError( + f"λ_central file {path!r} is not valid JSON; got {exc}" + ) from exc + if ( + not isinstance(data, dict) + or "eff_params" not in data + or "gnu_params" not in data + ): + raise ValueError( + f"λ_central file {path!r} must decode to a dict with " + f"'eff_params' and 'gnu_params' keys; got {type(data).__name__} " + f"with keys {list(data) if isinstance(data, dict) else ''}." + ) + return data + + def read_lambda_central(hdf5_path, proc="Z"): """Extract λ_central from the SCETlib correction referenced by the hdf5. + Accepts either a fit-input datacard (setupRabbit output) OR a rabbit + ``fitresults*.hdf5``: rabbit propagates the whole datacard meta into the + fitresults as ``meta_info_input``, so for fitresults the lookup is the + same after unwrapping one level. Handy to answer "which λ_central was + this fit's ParamModel initialized with?" from the fit output alone. + + λ_central overrides: the override route is the ``lambda_central=`` + token in the ``--paramModel`` spec — the fit command (including the + token) is stored in the fitresults ``meta_info.args``, so this function + recovers the override automatically. A programmatic constructor-arg dict + is NOT visible in the fit output — for such fits the theoryCorr-tag + route below silently reports the auto-detect values (check the fit log + for "λ_central from" lines). + Returns a dict with: tag : the central theoryCorr tag (first entry of meta args) @@ -137,7 +188,32 @@ def read_lambda_central(hdf5_path, proc="Z"): if "meta" not in f: raise KeyError(f"{hdf5_path}: no 'meta' group — wrong file type?") meta = wums_io.pickle_load_h5py(f["meta"]) - return read_lambda_central_from_meta(meta, proc=proc, _source=hdf5_path) + + # Preference 1 (fitresults): a lambda_central= token in the stored + # --paramModel spec — the fit command records the override for free. + args_meta = (meta.get("meta_info") or {}).get("args") or {} + for spec in args_meta.get("paramModel") or []: + for tok in spec: + if isinstance(tok, str) and tok.startswith("lambda_central="): + path = tok.split("=", 1)[1] + lc = load_lambda_central_file(path) + lc["source"] = f"cli-file:{path}" + return lc + + # Fallback: resolve the datacard's theoryCorr tag. For fitresults this + # ASSUMES no env-var/constructor override was active (those are not + # visible in the output; check the fit log). fitresults nest the + # datacard meta (which itself carries meta_info_input) one level down — + # unwrap until the histmaker args are in view. + while ( + isinstance(meta.get("meta_info_input"), dict) + and "args" not in meta.get("meta_info_input", {}) + and "meta_info_input" in meta["meta_info_input"] + ): + meta = meta["meta_info_input"] + lc = read_lambda_central_from_meta(meta, proc=proc, _source=hdf5_path) + lc["source"] = "theoryCorr-tag (assumes no runtime override)" + return lc def _resolve_tag_to_lambda(tag, proc): diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 54a4f031c..478515f65 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -62,6 +62,13 @@ λ-independent and precomputed once — the bT·J₀ kernel, the bT Simpson weights, and the arctan_Q² Q-integration weights are all built at construction. +The default evaluation uses the memory-factorized layout (deduplicated +(I_pert, C_ν) rows + J₀ kernel on the unique-qT grid + Simpson-as-matmul), +which is numerically equivalent to the per-bin (Nbins, Nbt) layout (≲1e-14 +rel., floating-point summation order only) but ~6× smaller — required to fit +a 32 GB GPU. See FACTORIZED_RECON.md in this directory; the legacy_recon=1 +spec token restores the legacy layout. + (1b) Integrate over Q, then rebin onto the gen grid: σ_resum(λ; g) = rebin_{qT→ptVGen, |Y|→absYVGen} [ ∫_{Q_lo}^{Q_hi} dQ σ(Q, Y, qT; λ) ] @@ -89,7 +96,9 @@ σ_ns is NP-INDEPENDENT (the same for every λ) and is added at GEN level, so it folds through the same response R as σ_resum in Step 3. Because the fit uses a ratio (Step 3), σ_ns correctly DILUTES the NP variation where the FO dominates -(high qT). Set include_nonsingular=False for a resum-only model (σ_ns = 0). +(high qT). σ_ns is ALWAYS included (it is what the histmaker nominal carries); +for resum-only diagnostics subtract the exposed ``sigma_ns`` from +``sigma_gen_central`` (and re-fold with ``R`` for reco level). ═════════════════════════════════════════════════════════════════════════════ Step 3 — σ_reco(λ; b): fold gen → reco through the response matrix @@ -201,34 +210,48 @@ --noHessian -o fit/ 2. Covariance at that postfit (no refit), straight-through ON: - SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH=1 SCETLIB_NP_HESSIAN_GN=1 \ rabbit_fit.py --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel \ + hessian_straightthrough=1 hessian_gn=1 \ --externalPostfit fit/fitresults.hdf5 --externalPostfitResult nominal \ --noFit -t 0 --pseudoData nominal -o cov/ (do NOT pass --noHessian; do NOT pass --eager.) ~5 min, no OOM. Uncertainties are sqrt(diag(cov)) in cov/fitresults.hdf5 (results_nominal). -Env flags (both OFF by default → the fit path is unchanged): - SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH=1 use the straight-through path in compute() - SCETLIB_NP_HESSIAN_GN=1 Gauss-Newton: keep J only, drop the K term +Switches (both OFF by default → the fit path is unchanged). They are spec +tokens inside the --paramModel spec (shown above) — the spec is stored in the +fitresults meta_info.args, so the configuration is recorded in the output +(env vars are NOT supported; all model knobs go through the spec): + hessian_straightthrough=1 use the straight-through path in compute() + hessian_gn=1 Gauss-Newton: keep J only, drop the K term +WARNING: hessian_straightthrough=1 WITHOUT hessian_gn=1 is full-K mode — +correct in principle (needed for real/toy data) but currently INFEASIBLE at +full grid scale (the 64 nested-FA passes unroll into one graph, ~TB peak → +OOM-kill). Until HESSIAN_PLAN.md §9 (precomputed chunked K) is implemented, +always pass hessian_gn=1 (exact for Asimov). GN vs full-K. The Poisson Hessian is H_ij = Σ_b [ (n_b/μ_b²) J_bi J_bj + (1 − n_b/μ_b) K_bij ]. For ASIMOV data the residual (1 − n/μ) = 0, so the K term vanishes and GN (J -only) is EXACT. For real/toy data K is needed for the exact Hessian — but the -nested-forward-mode K (``_ratio_compact_hess``) currently CRASHES under rabbit's -@tf.function (the tf.where in ``_safe_div``) and is impractically slow under ---eager, so GN (the recipe above) is the working production path today. See -HESSIAN_PLAN.md §9 for the open K item and its fixes. - -WARNING: keep SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH UNSET during the FIT. The +only) is EXACT — drop the K term with hessian_gn=1. For real/toy data +the K term is needed for the exact Hessian. The nested-forward-mode K +(``_ratio_compact_hess``) USED to crash under rabbit's @tf.function: building the +JVP of an ``Equal`` op that carries a tangent raised ``IndexError`` in TF's +nested-forward-mode autodiff. The culprit was the λ_inf==0 / den==0 masks in +:mod:`btgrid_tf` comparing a differentiated tensor; freezing the comparison input +(``btgrid_tf._frozen_eq_zero`` = ``tf.equal(stop_gradient(x), 0)``) removes the +tangent into ``Equal`` without changing any value or derivative (the comparison +is a measure-zero boundary ``tf.where`` never differentiates). Full-K now runs +under @tf.function and matches the exact reverse-mode Hessian to machine +precision (≤3e-16 rel; see HESSIAN_PLAN.md §7 + the isolation validation). So +both GN and full-K are available; GN remains the default for Asimov (exact and +cheaper — 8 vs 72 fold passes). + +WARNING: do NOT pass hessian_straightthrough=1 during the FIT. The surrogate recomputes J(/K) on every compute() call — fine for the one-shot covariance pass, but it would cripple the minimizer (many gradient/HVP evals). """ -import json import os -import re from typing import Mapping, Optional import numpy as np @@ -243,11 +266,8 @@ from wremnants.postprocessing.scetlib_np import lambda_central as scetlib_lambda_central from wremnants.postprocessing.scetlib_np import response_matrix as fz_R from wremnants.utilities import common as wrem_common +from wremnants.utilities.data_paths import getDataPath -# Default fixed-order inputs for the NP-independent nonsingular term -# σ_ns = DYTurbo − SCETlib_singular, resolved relative to this package via -# wrem_common.data_dir (= /wremnants-data/data). These are the CT18Z -# N3+0LL fixed-order pieces; both are NP-independent (same for any λ tune). _NONSING_FO_SING_DEFAULT = os.path.join( wrem_common.data_dir, "TheoryCorrections", @@ -258,51 +278,14 @@ "TheoryCorrections", "results_z-2d-nnlo-vj-CT18ZNNLO-{scale}-scetlibmatch.txt", ) - -# Default fit inputs, so collaborators can run -# ``--paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel`` with no -# extra positional args. Passing the two positional args explicitly still wins. -# -# The unfolding response lives in wremnants-data (a skim of only the two hists -# load_R reads — ``nominal_prefsr_yieldsUnfolding`` + ``prefsr`` — 6.8 GB → 211 MB; -# made with scripts/inspect/open_narf_h5py.py). Resolved via wrem_common.data_dir. _UNFOLDING_HDF5_DEFAULT = os.path.join( wrem_common.data_dir, "TheoryCorrections", "scetlib_np", "mz_dilepton_unfolding_R_skim.hdf5", ) -# The bT-grid combined pickle is ~17.5 GB (two [Nbins×Nbt] float64 slabs: I_pert -# and C_nu) — too large for wremnants-data/git-LFS — so it lives on the shared -# data area next to NanoAOD, as a sibling of NanoAOD under the data base. We -# resolve the NanoAOD base by REPLICATING dataset_tools.getDataPath()'s hostname -# logic rather than importing it: that module does ``import ROOT`` / ``import narf`` -# at module scope (and wremnants.production.__init__ runs ROOT.gInterpreter + -# narf.clingutils.Load), which segfaults when imported mid-fit (after TF is up). -# Only the subMIT copy exists today — pass btgrid_dir explicitly at other sites. _BTGRID_SUBDIR = ("scetlib_np", "Z_COM13_CT18Z_N3p0LL_btgrid_fineall") - - -def _nano_data_base(): - """NanoAOD base dir per host — mirrors dataset_tools.getDataPath() WITHOUT - importing it (that import pulls ROOT/narf; see note above). Falls back to the - subMIT path on unknown hosts (the only site with the grid today).""" - import socket - - hostname = socket.gethostname() - if hostname.endswith(".cern.ch"): - return "/scratch/shared/NanoAOD" - if hostname == "cmsanalysis.pi.infn.it": - return "/scratchnvme/wmass/NANOV9/postVFP" - if hostname == "cmsasymow.pi.infn.it": - return "/scratch/wmass/y2016" - # .mit.edu and any unknown host -> subMIT shared scratch - return "/scratch/submit/cms/wmass/NanoAOD" - - -def _default_btgrid_dir(): - return os.path.join(os.path.dirname(_nano_data_base()), *_BTGRID_SUBDIR) - +_DISCRETE_NP_SUBSTRING = "scetlibnp" # Ordered list of the v1 continuous λ. CS-side first, then TMD-effective. GNU_PARAMS = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") @@ -323,27 +306,8 @@ def _default_btgrid_dir(): RATIO_FLOOR_SCALE = 1.0e-4 RATIO_FLOOR_MIN = 1.0e-9 - -# Theorist-recommended Gaussian prior widths for the SCETlib NP λ parameters. -# Source: NP-NP discussion slide, central values 2026-05, plus a wide -# in-house default for delta_lambda2 (the theorist hasn't quoted a width -# for it; 0 ± 0.2 is comfortably wider than its expected scale). -# -# λ₂^ν = 0.15 ± 0.10 -# Λ₂ = 0.40 ⁺⁰·⁶₋₀.₄ (asymmetric) -# Λ₄ = 0.40 ⁺⁰·⁶₋₀.₄ (asymmetric) -# δ Λ₂ = 0.00 ± 0.20 (in-house wide default; not from the slide) -# -# Asymmetric uncertainties (Λ₂, Λ₄) are approximated by a symmetric Gaussian -# with σ = (σ⁺ + σ⁻) / 2. Slightly conservative on the upper side, slightly -# loose on the lower side. A future patch can implement a split-Gaussian if -# needed. -# -# Any λ not listed here gets σ = NaN by default → no prior, floats free. -# Currently those are: lambda_inf, lambda_inf_nu, lambda4_nu, lambda6. -# They are expected to be FROZEN via rabbit's --freezeParameters until -# the theorist provides priors for them. -THEORIST_PRIOR_SIGMAS = { +# Recommended Gaussian prior widths for the SCETlib NP λ parameters +DEFAULT_PRIOR_SIGMAS = { "lambda2_nu": 0.10, "lambda2": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average "lambda4": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average @@ -351,43 +315,16 @@ def _default_btgrid_dir(): } -def _load_lambda_central_file(path): - """Load a λ_central override from a JSON or YAML file. +def _default_btgrid_dir(): + base = getDataPath(fallback="/scratch/submit/cms/wmass/NanoAOD") + return os.path.join(os.path.dirname(base), *_BTGRID_SUBDIR) - The file must decode to a dict with ``eff_params`` and ``gnu_params`` - sub-dicts (same shape as :func:`scetlib_lambda_central.read_lambda_central`). - Format is chosen by extension (``.yaml``/``.yml`` → YAML, else JSON); - YAML's loader also accepts JSON, so this is forgiving either way. - """ - if not os.path.exists(path): - raise FileNotFoundError( - f"SCETLIB_NP_LAMBDA_CENTRAL_FILE points to a missing file: {path!r}" - ) - with open(path) as f: - text = f.read() - if path.lower().endswith((".yaml", ".yml")): - import yaml - - data = yaml.safe_load(text) - else: - try: - data = json.loads(text) - except json.JSONDecodeError as exc: - raise ValueError( - f"SCETLIB_NP_LAMBDA_CENTRAL_FILE {path!r} is not valid JSON; got {exc}" - ) from exc - if ( - not isinstance(data, dict) - or "eff_params" not in data - or "gnu_params" not in data - ): - raise ValueError( - f"SCETLIB_NP_LAMBDA_CENTRAL_FILE {path!r} must decode to a dict with " - f"'eff_params' and 'gnu_params' keys; got {type(data).__name__} " - f"with keys {list(data) if isinstance(data, dict) else ''}." - ) - return data +def _load_lambda_central_file(path): + """ + Load a λ_central override from a JSON or YAML file. + """ + return scetlib_lambda_central.load_lambda_central_file(path) def _crop_R_to_fit(R, R_reco_axes, fit_reco_axes, tol=1e-9): @@ -516,25 +453,55 @@ def _central(h): class SCETlibNPParamModel(ParamModel): + @classmethod + def parse_args(cls, indata, *args, **kwargs): + import inspect + + sig = inspect.signature(cls.__init__) + valid = {n: p for n, p in sig.parameters.items() if n not in ("self", "indata")} + positional = [] + for tok in args: + key = tok.split("=", 1)[0] if isinstance(tok, str) and "=" in tok else None + if key in valid: + 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, unfolding_hdf5_path: Optional[str] = None, btgrid_dir: Optional[str] = None, - lambda_central: Optional[Mapping] = None, + lambda_central=None, signal_proc: str = "Zmumu", Q_lo: float = 60.0, Q_hi: float = 120.0, poi_params: Optional[tuple] = (), + priors: bool = False, prior_sigmas: Optional[Mapping] = None, - include_nonsingular: bool = True, nonsingular_fo_sing: str = _NONSING_FO_SING_DEFAULT, nonsingular_dyturbo: str = _NONSING_DYTURBO_DEFAULT, nonsingular_qt_cutoff: float = 1.0, + legacy_recon: bool = False, + xparam_default: Optional[str] = None, + hessian_straightthrough: bool = False, + hessian_gn: bool = False, **kwargs, ): """Construct the ParamModel. + Usage:: + --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel [key=value ...] + Parameters ---------- indata @@ -548,10 +515,9 @@ def __init__( btgrid_dir Directory holding the SCETlib bT-grid ``combined_btgrid.pkl``. Defaults (when None) to the shared data-area copy next to NanoAOD - (``_default_btgrid_dir()``, which mirrors - ``dataset_tools.getDataPath()``'s per-host logic without importing it — - that import pulls ROOT/narf and segfaults mid-fit); pass explicitly at - non-subMIT sites. + (``_default_btgrid_dir()``, built on the ROOT/narf-free + ``wremnants.utilities.data_paths.getDataPath()``); pass explicitly + at non-subMIT sites. lambda_central Dict with two sub-dicts ``eff_params`` and ``gnu_params`` (same shape as returned by :func:`scetlib_lambda_central.read_lambda_central`). @@ -568,9 +534,19 @@ def __init__( POIs (reported as POIs in the fit output). The rest are reported as model nuisances (npou). The POI vs POU split is independent of the prior assignment (see ``prior_sigmas``). + priors + Enable Gaussian priors on the λ parameters (spec token + ``priors=1``). Rabbit applies priors whenever a ParamModel + *declares* ``prior_sigmas`` — there is no rabbit-side CLA (the + old ``--paramModelPriors`` flag was dropped in WMass/rabbit#133); + the model itself decides. This token IS that decision: only when + set does the model declare ``prior_sigmas``. Default off → + everything floats free. prior_sigmas - Per-name override dict for the Gaussian prior σ on each - parameter. Defaults come from ``THEORIST_PRIOR_SIGMAS``: + Per-name override for the Gaussian prior σ on each parameter — + a Mapping, or as spec token the comma-separated string form + ``prior_sigmas=lambda2=0.3,delta_lambda2=nan`` (same format as + ``xparam_default``). Defaults come from ``DEFAULT_PRIOR_SIGMAS``: lambda2_nu : 0.10 lambda2 : 0.50 (symmetric approx of +0.6/-0.4) @@ -581,43 +557,64 @@ def __init__( ``--freezeParameters`` until the theorist provides priors for them. Pass ``np.nan`` here to free a constrained param, or a finite value to add a prior on one that defaults to NaN. - Only consumed when the fitter is invoked with - ``--paramModelPriors``; otherwise everything floats free. + Only meaningful together with ``priors=1``; ignored (with a + warning) otherwise. Prior mean for each param is ``self.xparamdefault`` (the runcard's λ_central). + nonsingular_fo_sing, nonsingular_dyturbo + Paths to the σ_ns inputs (SCETlib singular pkl / DYTurbo + scetlibmatch txt); default to the wremnants-data + TheoryCorrections copies. σ_ns is always included — the matched + σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns is what the histmaker + nominal carries; resum-only diagnostics subtract ``sigma_ns``. + nonsingular_qt_cutoff + Low-qT cutoff (GeV) below which σ_ns is zeroed (the FO−singular + difference is numerically unreliable at tiny qT). + legacy_recon + Use the legacy per-bin (Nbins, Nbt) reconstruction layout instead + of the default memory-factorized one (numerically equivalent to + ≲1e-14 rel; for parity checks only). See FACTORIZED_RECON.md. + xparam_default + Comma-separated ``name=value,...`` string shifting the fit START + (and the prior mean) off the runcard's λ_central — for closure / + injection tests. The truth (ratio denominator) is NOT moved. + hessian_straightthrough + Expose compact λ-derivatives (J, optionally K) to autodiff while + keeping the exact value — for the one-shot two-pass covariance + recipe ONLY (see the module docstring); never set during a fit. + hessian_gn + With ``hessian_straightthrough``: Gauss-Newton, keep J and drop + the K term (exact for Asimov, where the residual vanishes). """ self.indata = indata - # ---- Double-counting guard - # Resolve default inputs so collaborators can run with no extra - # --paramModel positional args (explicit args still override). The - # unfolding skim ships in wremnants-data; the big bT-grid lives on the - # shared data area next to NanoAOD (see _default_btgrid_dir). if unfolding_hdf5_path is None: unfolding_hdf5_path = _UNFOLDING_HDF5_DEFAULT if btgrid_dir is None: btgrid_dir = _default_btgrid_dir() - # If the histmaker baked discrete NP κ-template variations into the - # input HDF5, those systs describe the same physics as our continuous - # λ POUs. Running both → double-counting (the discrete syst absorbs - # whatever shape variation our ParamModel should describe). Warn - # loudly if any such systs are present and unfrozen. - self._check_discrete_np_double_counting(kwargs.get("freezeParameters")) + self._check_discrete_np_double_counting() # ---- λ_central # Three sources of λ_central, in priority order: - # 1. ``lambda_central`` constructor arg (explicit dict). - # 2. ``SCETLIB_NP_LAMBDA_CENTRAL_FILE`` env var — path to a JSON or - # YAML file with ``eff_params`` and ``gnu_params``. Overrides the + # 1. ``lambda_central=`` spec token — path to a JSON or YAML + # file with ``eff_params`` and ``gnu_params``. Overrides the # metadata auto-detect; useful when the upstream SCETlib pkl isn't # accessible (e.g. a colleague's input). + # 2. ``lambda_central`` constructor arg (explicit dict, programmatic). # 3. Auto-detect from the fit hdf5's theoryCorr → upstream pkl. - env_lc_file = os.environ.get("SCETLIB_NP_LAMBDA_CENTRAL_FILE", "").strip() - if lambda_central is None and env_lc_file: - lambda_central = _load_lambda_central_file(env_lc_file) + lambda_central_source = ( + "constructor-arg" if lambda_central is not None else None + ) + if isinstance(lambda_central, str): + # CLI token (lambda_central=): RECOMMENDED override route — + # the --paramModel spec is stored in the fitresults meta, so the + # override is recorded in the output (env var/dict are not). + lc_path = lambda_central + lambda_central = _load_lambda_central_file(lc_path) + lambda_central_source = f"cli-file:{lc_path}" print( - f"[SCETlibNPParamModel] λ_central from file {env_lc_file!r}", + f"[SCETlibNPParamModel] λ_central from CLI file {lc_path!r}", flush=True, ) if lambda_central is None: @@ -633,6 +630,7 @@ def __init__( lambda_central = scetlib_lambda_central.read_lambda_central_from_meta( indata_meta, _source="indata.metadata" ) + lambda_central_source = "auto-detect:indata.metadata theoryCorr" print( f"[SCETlibNPParamModel] λ_central auto-detected from indata.metadata", flush=True, @@ -641,6 +639,8 @@ def __init__( for key, value in lambda_central.items(): print(f" {key} = {value!r}", flush=True) + self.lambda_central_source = lambda_central_source + self.eff_central = dict(lambda_central["eff_params"]) self.gnu_central = dict(lambda_central["gnu_params"]) self.np_model = self.eff_central["np_model"] @@ -657,32 +657,81 @@ def __init__( # Cache btgrid arrays as TF constants. self.bT = tf.constant(grid["bT"], dtype=fz_tf.DTYPE) self.b_bar = tf.constant(grid["b_bar"], dtype=fz_tf.DTYPE) - self.I_pert = tf.constant(grid["I_pert"][0], dtype=fz_tf.DTYPE) # (Nbins, Nbt) - self.C_nu = tf.constant(grid["C_nu"][0], dtype=fz_tf.DTYPE) - # Per-bin qT and Y (from the bin tuple), for reconstruct_batch_tf. + # Per-bin qT and Y (from the bin tuple). bins = grid["bins"] - self.qT_per_bin = tf.constant( - np.array([b[2] for b in bins], dtype=np.float64), dtype=fz_tf.DTYPE - ) + qT_pb_np = np.array([b[2] for b in bins], dtype=np.float64) + self.qT_per_bin = tf.constant(qT_pb_np, dtype=fz_tf.DTYPE) Y_pb_np = np.array([b[1] for b in bins], dtype=np.float64) self.Y_per_bin = tf.constant(Y_pb_np, dtype=fz_tf.DTYPE) # F_eff depends on the bin only through Y (not Q or qT), and Y takes few - # distinct values across the grid. Precompute the unique-Y map so - # reconstruct_batch_tf evaluates the NP transcendentals on NY rows and + # distinct values across the grid. Precompute the unique-Y map so the + # reconstruction evaluates the NP transcendentals on NY rows and # gathers, instead of recomputing identical rows for every (Q, qT). Y_feff_unique_np, Y_feff_inv_np = np.unique(Y_pb_np, return_inverse=True) + Y_feff_inv_np = Y_feff_inv_np.reshape(-1).astype(np.int32) self.Y_feff_unique = tf.constant(Y_feff_unique_np, dtype=fz_tf.DTYPE) - self.Y_feff_inverse_idx = tf.constant( - Y_feff_inv_np.reshape(-1).astype(np.int32), dtype=tf.int32 - ) - - # Precompute the bT·J0(qT·bT) kernel (λ-independent). - self.bT_J0_kernel = fz_tf.build_bT_J0_kernel(self.qT_per_bin, self.bT) - self.bT_simpson_w = tf.constant( - fz_tf.simpson_weights(np.asarray(self.bT)), dtype=fz_tf.DTYPE - ) + self.Y_feff_inverse_idx = tf.constant(Y_feff_inv_np, dtype=tf.int32) + + bT_simpson_w_np = fz_tf.simpson_weights(np.asarray(grid["bT"])) + self.bT_simpson_w = tf.constant(bT_simpson_w_np, dtype=fz_tf.DTYPE) + + # Reconstruction layout. Default: factorized (deduplicated rows + + # unique-qT J0 kernel + Simpson-as-matmul) — numerically equivalent to + # the legacy (Nbins, Nbt) layout (≲1e-14 rel., summation order only) + # but ~6x smaller, which is what lets the fit run on a 32 GB GPU. + # The legacy_recon=1 spec token selects the legacy path (parity + # checks). + self.factorized = not bool(legacy_recon) + + # Hessian straight-through switches (see the module docstring's + # two-pass recipe). Spec tokens hessian_straightthrough=1 / + # hessian_gn=1, recorded in the fitresults meta via the stored + # --paramModel spec. + self._hess_st = bool(hessian_straightthrough) + self._hess_gn = bool(hessian_gn) + if self.factorized: + dd = fz_tf.dedup_grid_rows( + grid["I_pert"][0], grid["C_nu"][0], Y_feff_inv_np + ) + self.I_pert_u = tf.constant(dd["I_u"], dtype=fz_tf.DTYPE) # (Nu, Nbt) + # C_ν via the second-level dedup: exp(C·g) runs on the small + # (Ncu, Nbt) table and is gathered — bit-identical, ~150x fewer + # transcendentals, no (Nu, Nbt) C constant on device. + self.C_nu_uu = tf.constant(dd["C_uu"], dtype=fz_tf.DTYPE) # (Ncu, Nbt) + self.c_of_u = tf.constant(dd["c_of_u"], dtype=tf.int32) + self.feff_idx_u = tf.constant(dd["feff_idx_u"], dtype=tf.int32) + # Per-bin index into the unique-qT axis. The bin qT values are by + # construction members of qT_unique, so searchsorted is an exact + # lookup (asserted). + qT_idx_np = np.searchsorted(idx_map["qT_unique"], qT_pb_np) + assert np.array_equal(idx_map["qT_unique"][qT_idx_np], qT_pb_np) + self.gather_idx = tf.constant( + np.stack([dd["row_uid"].astype(np.int64), qT_idx_np], axis=1), + dtype=tf.int32, + ) + # Drop the host-side dedup copies (the tf.constants own the data now). + del dd + # Weighted J0 kernel on the unique-qT grid, with the per-bin qT + # prefactor and the Simpson weights folded in: (NqT, Nbt). + K_u = fz_tf.build_bT_J0_kernel( + tf.constant(idx_map["qT_unique"], dtype=fz_tf.DTYPE), self.bT + ) + self.KwqT = ( + tf.constant(idx_map["qT_unique"], dtype=fz_tf.DTYPE)[:, tf.newaxis] + * K_u + * self.bT_simpson_w[tf.newaxis, :] + ) + else: + self.I_pert = tf.constant( + grid["I_pert"][0], dtype=fz_tf.DTYPE + ) # (Nbins, Nbt) + self.C_nu = tf.constant(grid["C_nu"][0], dtype=fz_tf.DTYPE) + # Precompute the bT·J0(qT·bT) kernel (λ-independent). + self.bT_J0_kernel = fz_tf.build_bT_J0_kernel(self.qT_per_bin, self.bT) + # Drop the ~17.5 GB host-side grid reference before TF graph building. + del grid # ---- Q-integration weights (arctan_Q² Simpson on Z mass window). self.Q_weights = tf.constant( @@ -782,47 +831,45 @@ def __init__( self.sigma_YqT_central = self._sigma_YqT_native_at( self.eff_central, self.gnu_central ) - # ---- Optional fixed-order/DYTurbo nonsingular term (NP-independent). + # ---- Fixed-order/DYTurbo nonsingular term (NP-independent). # σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns. Added at GEN level, so it folds # through the same response R as the resummed piece. Because rnorm is a # ratio, this correctly DILUTES the NP variation where the FO dominates - # (high qT). σ_ns is a constant (no λ dependence). - self.include_nonsingular = bool(include_nonsingular) - if self.include_nonsingular: - _dy0 = ( - nonsingular_dyturbo.format(scale="mur1-muf1") - if (nonsingular_dyturbo and "{scale}" in nonsingular_dyturbo) - else nonsingular_dyturbo + # (high qT). σ_ns is a constant (no λ dependence), always included — + # the matched σ_gen is what the histmaker nominal carries; resum-only + # diagnostics subtract self.sigma_ns instead of rebuilding the model. + _dy0 = ( + nonsingular_dyturbo.format(scale="mur1-muf1") + if (nonsingular_dyturbo and "{scale}" in nonsingular_dyturbo) + else nonsingular_dyturbo + ) + missing = [ + p for p in (nonsingular_fo_sing, _dy0) if not (p and os.path.exists(p)) + ] + if missing: + raise FileNotFoundError( + "The matched model needs the fixed-order inputs for " + "σ_ns = DYTurbo − SCETlib_singular, but these are missing:\n " + + "\n ".join(missing) + + "\nThey live under wremnants-data/data/TheoryCorrections (the " + "SCETlib singular …_nnlo_sing…combined.pkl and the DYTurbo " + "results_…scetlibmatch.txt). Pass nonsingular_fo_sing / " + "nonsingular_dyturbo to point at them." ) - missing = [ - p for p in (nonsingular_fo_sing, _dy0) if not (p and os.path.exists(p)) - ] - if missing: - raise FileNotFoundError( - "include_nonsingular=True needs the fixed-order inputs for " - "σ_ns = DYTurbo − SCETlib_singular, but these are missing:\n " - + "\n ".join(missing) - + "\nThey live under wremnants-data/data/TheoryCorrections (the " - "SCETlib singular …_nnlo_sing…combined.pkl and the DYTurbo " - "results_…scetlibmatch.txt). Pass nonsingular_fo_sing / " - "nonsingular_dyturbo, or set include_nonsingular=False for resum-only." - ) - sigma_ns_np = compute_nonsingular_gen( - nonsingular_fo_sing, - nonsingular_dyturbo, - self._gen_axes_meta, - q_lo=Q_lo, - q_hi=Q_hi, - qt_cutoff=nonsingular_qt_cutoff, + sigma_ns_np = compute_nonsingular_gen( + nonsingular_fo_sing, + nonsingular_dyturbo, + self._gen_axes_meta, + q_lo=Q_lo, + q_hi=Q_hi, + qt_cutoff=nonsingular_qt_cutoff, + ) + if sigma_ns_np.shape != tuple(self.gen_shape): + raise ValueError( + f"nonsingular gen shape {sigma_ns_np.shape} != model gen shape " + f"{tuple(self.gen_shape)}" ) - if sigma_ns_np.shape != tuple(self.gen_shape): - raise ValueError( - f"nonsingular gen shape {sigma_ns_np.shape} != model gen shape " - f"{tuple(self.gen_shape)}" - ) - self.sigma_ns = tf.constant(sigma_ns_np, dtype=fz_tf.DTYPE) - else: - self.sigma_ns = tf.zeros(self.gen_shape, dtype=fz_tf.DTYPE) + self.sigma_ns = tf.constant(sigma_ns_np, dtype=fz_tf.DTYPE) # Reuse the native (NY, NqT) integral already computed above for # sigma_YqT_central — no need to run the bT reconstruction at λ_central twice. sigma_gen_central = self._sigma_gen_at( @@ -879,24 +926,45 @@ def __init__( self.npou = len(nou_params) self.params = np.array([p.encode() for p in self._param_order]) + # Impact groups over our own parameters, consumed by the Fitter's + # traditional impacts. rabbit's built-in systgroup machinery can't + # represent these: its group indices are syst-relative, but our λ are + # POUs (model nuisances), which have no syst index. The Fitter resolves + # these labels -> floating full-x indices and computes the conditional + # group impact from the covariance. Split into the two NP sectors: + # CS-side γ_ν (lambda*_nu) vs TMD-effective F_eff. + # + # ``resumNonpert`` is the SAME group name setupRabbit assigns to the + # discrete scetlibNP* template variations in the old-style datacard + # (there resumNonpert == exactly those 4 nuisances). Emitting it here + # (= all our lambda) makes the grouped-impact bar directly comparable + # between the new param model and the old NP variations. It does NOT + # collide with a syst group: the new-model datacard excludes scetlibNP, + # so resumNonpert is absent from indata.systgroups. + self.param_impact_groups = { + "resumNonpert": tuple(ALL_PARAMS), + "scetlibNPgammaNu": tuple(GNU_PARAMS), + "scetlibNPFeff": tuple(EFF_PARAMS), + } + # Defaults: λ_central values per parameter. Optionally overridden by - # the ``SCETLIB_NP_XPARAMDEFAULT`` env var — comma-separated - # ``name=value`` pairs (for closure tests where the data-generating - # / fit-start point should differ from the card's λ_central). + # the ``xparam_default=name=value,...`` spec token — comma-separated + # pairs (for closure tests where the data-generating / fit-start + # point should differ from the card's λ_central). central_lookup = {**self.eff_central, **self.gnu_central} defaults = np.array( [central_lookup[p] for p in self._param_order], dtype=np.float64 ) - env_override = os.environ.get("SCETLIB_NP_XPARAMDEFAULT", "").strip() - if env_override: + start_override = (xparam_default or "").strip() + if start_override: overrides = dict( - tuple(s.split("=")) for s in env_override.split(",") if s.strip() + tuple(s.split("=")) for s in start_override.split(",") if s.strip() ) for name, val in overrides.items(): name = name.strip() if name not in self._param_order: - raise KeyError(f"SCETLIB_NP_XPARAMDEFAULT: unknown param {name!r}") + raise KeyError(f"xparam_default: unknown param {name!r}") i = self._param_order.index(name) defaults[i] = float(val) print( @@ -911,70 +979,65 @@ def __init__( self.is_linear = False self.xparamdefault = tf.constant(defaults, dtype=indata.dtype) - # Gaussian priors (consumed by rabbit's Fitter when --paramModelPriors - # is set; ignored otherwise). Default σ values come from - # ``THEORIST_PRIOR_SIGMAS`` (lambda2_nu, lambda2, lambda4 — the only - # params the theorist provides widths for as of 2026-05). - # All other params default to σ = NaN → no prior, float free; in - # practice those should be frozen via rabbit's --freezeParameters - # until the theorist gives priors for them. - # The ``prior_sigmas`` kwarg is a per-name override dict; pass NaN to - # force a parameter free, or a finite value to add / change a prior. - # The mean of each prior is λ_central (i.e. self.xparamdefault). - prior_sigmas = dict(prior_sigmas or {}) - sigmas_arr = np.empty(self.nparams, dtype=np.float64) - for i, p in enumerate(self._param_order): - if p in prior_sigmas: - sigmas_arr[i] = float(prior_sigmas[p]) # explicit override (may be NaN) - elif p in THEORIST_PRIOR_SIGMAS: - sigmas_arr[i] = THEORIST_PRIOR_SIGMAS[p] # theorist recommendation - else: - sigmas_arr[i] = np.nan # free (expected to be frozen) - self.prior_sigmas = sigmas_arr - # prior_means defaults to xparamdefault if not set, so don't store - # redundantly — Fitter will fall back to xparamdefault. - - # ========================================================================= - # Helpers - # ========================================================================= + # Gaussian priors (semantics documented on the constructor args). + # Rabbit's Fitter applies them whenever the model DECLARES + # ``prior_sigmas`` (no rabbit-side CLA — WMass/rabbit#133), so the + # declaration is gated behind ``priors``: off → no attribute → + # everything floats free. The Fitter takes the prior means from + # xparamdefault, so an xparam_default shift moves start AND prior + # mean together (to centre priors on truth while starting shifted, + # prior_means would have to be decoupled from xparamdefault). + self._use_priors = bool(priors) + # ``prior_sigmas`` may be a Mapping (programmatic) or the spec-token + # string ``prior_sigmas=lambda2=0.3,delta_lambda2=nan`` — the same + # comma-separated name=value format as xparam_default; value ``nan`` + # frees the param. + if isinstance(prior_sigmas, str): + prior_sigmas = dict( + tuple(s.split("=")) for s in prior_sigmas.split(",") if s.strip() + ) + prior_sigmas = {k.strip(): v for k, v in dict(prior_sigmas or {}).items()} + for name in prior_sigmas: + if name not in self._param_order: + raise KeyError(f"prior_sigmas: unknown param {name!r}") + if self._use_priors: + sigmas_arr = np.empty(self.nparams, dtype=np.float64) + for i, p in enumerate(self._param_order): + if p in prior_sigmas: + sigmas_arr[i] = float( + prior_sigmas[p] + ) # explicit override (may be NaN) + elif p in DEFAULT_PRIOR_SIGMAS: + sigmas_arr[i] = DEFAULT_PRIOR_SIGMAS[p] # theorist recommendation + else: + sigmas_arr[i] = np.nan # free (expected to be frozen) + self.prior_sigmas = sigmas_arr + # prior_means defaults to xparamdefault if not set, so don't store + # redundantly — Fitter will fall back to xparamdefault. + print( + "[SCETlibNPParamModel] Gaussian priors ENABLED (priors=1); " + "applied by rabbit's Fitter (pre-#133 rabbit additionally " + "needs --paramModelPriors):", + flush=True, + ) + for i, p in enumerate(self._param_order): + if np.isfinite(sigmas_arr[i]) and sigmas_arr[i] > 0: + print(f" {p}: σ = {sigmas_arr[i]:.4g}", flush=True) + elif prior_sigmas: + print( + "[SCETlibNPParamModel] WARNING: prior_sigmas overrides given " + "but priors are not enabled (pass priors=1); ignoring them — " + "all λ float free.", + flush=True, + ) - # Substrings (case-insensitive) that mark indata.systs as discrete - # NP-template variations of one of our 8 continuous λ parameters. - # Matching is case-insensitive because the histmaker uses inconsistent - # casing (canonical names are uppercase Lambda, but some configurations - # serialize them lowercase). - # - # Canonical names (see theory_variation_labels.py): - # chargeVgenNP0scetlibNPZLambda2 → catches "scetlibnpzlambda" - # chargeVgenNP0scetlibNPZLambda4 → catches "scetlibnpzlambda" - # chargeVgenNP0scetlibNPZDelta_Lambda2 → catches "scetlibnpzdelta" - # chargeVgenNP0scetlibNPLambda2 (W-side) → catches "scetlibnplambda" - # chargeVgenNP0scetlibNPLambda4 (W-side) → catches "scetlibnplambda" - # chargeVgenNP0scetlibNPDelta_Lambda2 → catches "scetlibnpdelta" - # scetlibNPgamma → catches "scetlibnpgamma" - # scetlibNPgammaEigvar{1,2,3} → catches "scetlibnpgamma" - # scetlibNPgammaLambda{2,4,Inf} → catches "scetlibnpgamma" - _DISCRETE_NP_PATTERNS = ( - "scetlibnpzlambda", # Z-side Lambda2 / Lambda4 templates - "scetlibnpzdelta", # Z-side Delta_Lambda2 template - "scetlibnplambda", # W-side Lambda2 / Lambda4 templates - "scetlibnpdelta", # W-side Delta_Lambda2 template - "scetlibnpgamma", # all γ_ν templates (Lambda2/4/Inf, Eigvar1/2/3, "gamma") - ) + def _check_discrete_np_double_counting(self): + """Refuse to run on a datacard containing discrete scetlibNP systs. - def _check_discrete_np_double_counting(self, freeze_patterns): - """Detect indata systs that overlap with our continuous λ POUs. - - Three outcomes: - - The discrete NP syst is **absent** from ``indata.systs`` entirely - (histmaker didn't include it): nothing to do, silent return. - - The discrete NP syst is **present and matched** by - ``freeze_patterns``: it's frozen at central → no double-counting, - silent return. - - The discrete NP syst is **present and unfrozen**: it overlaps - with one of our continuous λ POUs and will absorb shape variation - the ParamModel should describe → print a loud banner with the - exact freeze args to add. + They describe the same physics as this ParamModel's continuous λ; + running both double-counts (the discrete syst absorbs shape variation + the ParamModel should describe: spurious pull on the indata syst, + postfit λ not what the data prefers). """ systs = getattr(self.indata, "systs", None) @@ -982,59 +1045,19 @@ def _check_discrete_np_double_counting(self, freeze_patterns): return syst_names = [s.decode() if isinstance(s, bytes) else str(s) for s in systs] - # Case-insensitive substring match: canonical names use uppercase - # Lambda but some histmaker outputs lowercase the names. conflicting = [ - s - for s in syst_names - if any(pat in s.lower() for pat in self._DISCRETE_NP_PATTERNS) + s for s in syst_names if self._DISCRETE_NP_SUBSTRING in s.lower() ] if not conflicting: - return # not in indata.systs at all → nothing to warn about - - # Which of those are NOT already covered by a user-supplied freeze - # pattern (exact match or anchored regex)? - patterns = list(freeze_patterns or []) - unfrozen = [] - for s in conflicting: - covered = False - for pat in patterns: - if pat == s: - covered = True - break - try: - if re.fullmatch(pat, s): - covered = True - break - except re.error: - continue - if not covered: - unfrozen.append(s) - - if not unfrozen: - return # all conflicting systs are already frozen by the user - - print( - "\n" - "===================================================================\n" - "[SCETlibNPParamModel] DOUBLE-COUNTING WARNING\n" - "===================================================================\n" - f"Detected {len(unfrozen)} discrete NP κ-template syst(s) in the\n" - "input HDF5 that describe the same physics as this ParamModel's\n" - "continuous λ parameters. Running both leads to double-counting:\n" - "the discrete syst absorbs shape variation that the ParamModel\n" - "should describe (the indata syst will show a spurious pull, and\n" - "the postfit λ values are not what the data actually prefers).\n\n" - "Unfrozen conflicting systs:\n" - + "\n".join(f" {s}" for s in unfrozen) - + "\n\n" - "Fix by adding to --freezeParameters, e.g.:\n" - " --freezeParameters '.*scetlibNPZ.*lambda.*' " - "'.*scetlibNPgammaLambda.*' ...\n" - "or list them explicitly:\n" - " --freezeParameters " + " ".join(repr(s) for s in unfrozen) + "\n" - "===================================================================", - flush=True, + return + + raise ValueError( + f"[SCETlibNPParamModel] {len(conflicting)} discrete scetlibNP " + "κ-template syst(s) found in the input HDF5; they describe the " + "same physics as this ParamModel's continuous λ parameters and " + "running both double-counts. Remake the datacard without them " + "(setupRabbit --excludeNuisances '.*scetlibNP.*'). Conflicting " + "systs:\n" + "\n".join(f" {s}" for s in conflicting) ) def _fit_reco_axes(self, indata): @@ -1068,23 +1091,43 @@ def _sigma_YqT_native_at(self, eff_params, gnu_params): qT grid (Y_unique, qT_unique), BEFORE the |Y|-fold and qT-rebin. This is the object that the native-binning validation compares against the SCETlib spectrum reference (curve 1) and the numpy `factorize` (curve 2).""" - # 1. Reconstruct σ on the btgrid's flat (Nbins,) layout. - sigma_flat = fz_tf.reconstruct_batch_tf( - qT_per_bin=self.qT_per_bin, - bT=self.bT, - I_pert=self.I_pert, - C_nu=self.C_nu, - b_bar=self.b_bar, - Y_per_bin=self.Y_per_bin, - eff_params={k: v for k, v in eff_params.items() if k != "np_model"}, - gnu_params={k: v for k, v in gnu_params.items() if k != "np_model_nu"}, - np_model=self.np_model, - np_model_nu=self.np_model_nu, - bT_J0_kernel=self.bT_J0_kernel, - bT_simpson_weights=self.bT_simpson_w, - Y_unique=self.Y_feff_unique, - Y_inverse_idx=self.Y_feff_inverse_idx, - ) + # 1. Reconstruct σ on the btgrid's flat (Nbins,) layout. Factorized + # (default) and legacy layouts are numerically equivalent (≲1e-14 + # rel.; summation order only — see FACTORIZED_RECON.md). + eff = {k: v for k, v in eff_params.items() if k != "np_model"} + gnu = {k: v for k, v in gnu_params.items() if k != "np_model_nu"} + if self.factorized: + sigma_flat = fz_tf.reconstruct_batch_factorized_tf( + b_bar=self.b_bar, + I_pert_u=self.I_pert_u, + C_nu_uu=self.C_nu_uu, + c_of_u=self.c_of_u, + eff_params=eff, + gnu_params=gnu, + np_model=self.np_model, + np_model_nu=self.np_model_nu, + KwqT=self.KwqT, + gather_idx=self.gather_idx, + Y_unique=self.Y_feff_unique, + feff_idx_u=self.feff_idx_u, + ) + else: + sigma_flat = fz_tf.reconstruct_batch_tf( + qT_per_bin=self.qT_per_bin, + bT=self.bT, + I_pert=self.I_pert, + C_nu=self.C_nu, + b_bar=self.b_bar, + Y_per_bin=self.Y_per_bin, + eff_params=eff, + gnu_params=gnu, + np_model=self.np_model, + np_model_nu=self.np_model_nu, + bT_J0_kernel=self.bT_J0_kernel, + bT_simpson_weights=self.bT_simpson_w, + Y_unique=self.Y_feff_unique, + Y_inverse_idx=self.Y_feff_inverse_idx, + ) # 2. Sparse → dense (NQ, NY, NqT). Missing cells get 0. sigma_dense = fz_int.sparse_to_dense_tf(sigma_flat, self.flat_idx) # 3. Integrate over Q (arctan_Q² Simpson) → (NY, NqT). @@ -1244,16 +1287,13 @@ def compute(self, param, full=False): # ratio(λ) per reco bin. Normal path = exact fold (used for the fit). # Straight-through path (Hessian-only Phase B) keeps the bT slab off the # autodiff graph so rabbit's covariance jacobian doesn't OOM. Toggled by - # SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH; SCETLIB_NP_HESSIAN_GN=1 drops the - # curvature term (Gauss-Newton/Fisher — exact for Asimov, 8 vs 72 passes). + # hessian_straightthrough=1 spec token; + # hessian_gn=1 drops the curvature term + # (Gauss-Newton/Fisher — exact for Asimov, 8 vs 72 passes). Resolved at + # construction (self._hess_st / self._hess_gn). # Do NOT enable during the fit: it recomputes J(/K) every call. - if not hasattr(self, "_hess_st"): - self._hess_st = bool( - os.environ.get("SCETLIB_NP_HESSIAN_STRAIGHTTHROUGH", "").strip() - ) if self._hess_st: - gn = bool(os.environ.get("SCETLIB_NP_HESSIAN_GN", "").strip()) - ratio = self._ratio_straightthrough(param, use_curvature=not gn) + ratio = self._ratio_straightthrough(param, use_curvature=not self._hess_gn) else: ratio = self._ratio_from_param(param) diff --git a/wremnants/production/datasets/dataset_tools.py b/wremnants/production/datasets/dataset_tools.py index a9fdd5126..d5c2275eb 100644 --- a/wremnants/production/datasets/dataset_tools.py +++ b/wremnants/production/datasets/dataset_tools.py @@ -3,13 +3,22 @@ """ import importlib -import os -import random import ROOT -import XRootD.client import narf + +# Path / file-list helpers live in a ROOT/narf-free module so they can be +# imported without pulling in ROOT (this package's __init__ imports ROOT and +# narf). Re-exported here for backward compatibility. +from wremnants.utilities.data_paths import ( # noqa: F401 + appendFilesXrd, + buildFileList, + buildFileListPosix, + buildFileListXrd, + getDataPath, + makeFilelist, +) from wums import logging logger = logging.child_logger(__name__) @@ -24,172 +33,6 @@ } -def buildFileListPosix(path): - outfiles = [] - for root, dirs, fnames in os.walk(path): - for fname in fnames: - if fname.lower().endswith(".root"): - outfiles.append(f"{root}/{fname}") - - return outfiles - - -def appendFilesXrd( - filelist, xrdfs, path, suffixes=[".root"], recurse=False, num_clients=16 -): - status, dirlist = xrdfs.dirlist(path, flags=XRootD.client.flags.DirListFlags.STAT) - - if not status.ok: - if status.code == 400 and status.errno == 3011: - logger.warning(f"XRootD directory not found: {path}") - else: - raise RuntimeError( - f"Error in XRootD.client.FileSystem.dirlist: {status.message}, {status.code}, {status.errno}" - ) - - return - - for diritem in dirlist: - is_dir = diritem.statinfo.flags & XRootD.client.flags.StatInfoFlags.IS_DIR - is_other = diritem.statinfo.flags & XRootD.client.flags.StatInfoFlags.OTHER - is_file = not (is_dir or is_other) - - if is_dir and recurse: - childpath = f"{path}/{diritem.name}" - appendFilesXrd( - filelist, - xrdfs, - childpath, - suffixes=suffixes, - recurse=recurse, - num_clients=num_clients, - ) - elif is_file: - lowername = diritem.name.lower() - matchsuffix = False - for suffix in suffixes: - if lowername.endswith(suffix): - matchsuffix = True - break - - if matchsuffix: - if num_clients > 0: - # construct client string if necessary to force multiple xrootd connections - # (needed for good performance when a single or small number of xrootd servers is used) - client = f"user_{random.randrange(num_clients)}" - outname = f"{xrdfs.url.protocol}://{client}@{xrdfs.url.hostname}:{xrdfs.url.port}/{path}/{diritem.name}" - else: - outname = f"{xrdfs.url.protocol}://{xrdfs.url.hostid}/{path}/{diritem.name}" - - filelist.append(outname) - - -def buildFileListXrd(path, num_clients=16): - xrdurl = XRootD.client.URL(path) - - if not xrdurl.is_valid(): - raise ValueError(f"Invalid xrootd path {path}") - - xrdfs = XRootD.client.FileSystem(xrdurl.hostid) - xrdpath = xrdurl.path - - outfiles = [] - appendFilesXrd(outfiles, xrdfs, xrdpath, recurse=True, num_clients=num_clients) - - return outfiles - - -def buildFileList(path): - xrdprefix = "root://" - return ( - buildFileListXrd(path) - if path.startswith(xrdprefix) - else buildFileListPosix(path) - ) - - -# TODO add the rest of the samples! -def makeFilelist( - paths, - maxFiles=-1, - base_path=None, - nano_prod_tags=None, - is_data=False, - oneMCfileEveryN=None, - era=None, -): - filelist = [] - expandedPaths = [] - for orig_path in paths: - if maxFiles > 0 and len(filelist) >= maxFiles: - break - # try each tag in order until files are found - fallback = False - for prod_tag in nano_prod_tags: - format_args = dict(BASE_PATH=base_path, NANO_PROD_TAG=prod_tag, ERA=era) - - path = orig_path.format(**format_args) - expandedPaths.append(path) - logger.debug(f"Reading files from path {path}") - - files = buildFileList(path) - if maxFiles > 0 and len(files) >= maxFiles: - logger.info( - f"Booking {maxFiles} of {len(files)} files with tag {prod_tag} with path {path}" - ) - break - - if len(files) == 0: - fallback = True - logger.warning( - f"Did not find any files for tag {prod_tag} matching path {path}!" - ) - else: - if fallback: - logger.warning(f"Falling back to tag {prod_tag} with path {path}") - else: - logger.info( - f"Booking {maxFiles} of {len(files)} files with tag {prod_tag} with path {path}" - ) - break - - filelist.extend(files) - - toreturn = ( - filelist - if maxFiles < 0 or len(filelist) < maxFiles - else random.Random(1).sample(filelist, maxFiles) - ) - - if oneMCfileEveryN != None and not is_data: - tmplist = [] - for i, f in enumerate(toreturn): - if i % oneMCfileEveryN == 0: - tmplist.append(f) - logger.warning(f"Using {len(tmplist)} files instead of {len(toreturn)}") - toreturn = tmplist - - logger.debug(f"Length of list is {len(toreturn)} for paths {expandedPaths}") - return toreturn - - -def getDataPath(): - import socket - - hostname = socket.gethostname() - - if hostname.endswith(".cern.ch"): - base_path = "/scratch/shared/NanoAOD" - elif hostname.endswith(".mit.edu"): - base_path = "/scratch/submit/cms/wmass/NanoAOD" - elif hostname == "cmsanalysis.pi.infn.it": - # NOTE: If anyone wants to run lowpu analysis at Pisa they'd probably want a different path - base_path = "/scratchnvme/wmass/NANOV9/postVFP" - elif hostname == "cmsasymow.pi.infn.it": - base_path = "/scratch/wmass/y2016" - return base_path - - def is_zombie(file_path): # Try opening the ROOT file and check if it's a zombie file file = ROOT.TFile.Open(file_path) diff --git a/wremnants/utilities/data_paths.py b/wremnants/utilities/data_paths.py new file mode 100644 index 000000000..0e4183413 --- /dev/null +++ b/wremnants/utilities/data_paths.py @@ -0,0 +1,204 @@ +""" +Site-dependent data paths and dataset file-list construction. + +This module is importable WITHOUT ROOT, narf, or XRootD. It deliberately lives +under ``wremnants.utilities`` rather than ``wremnants.production.datasets``: +the ``wremnants.production`` package ``__init__`` imports ROOT and narf, so any +module below it pulls them in at import time. Code that only needs the +per-host data path (e.g. the SCETlib NP ParamModel, which runs inside a fit +where a mid-fit ROOT import can segfault) should import from here. + +``wremnants.production.datasets.dataset_tools`` re-exports everything defined +here, so existing call sites keep working. + +XRootD is imported lazily, only when a ``root://`` path is actually listed. +""" + +import os +import random +import socket + +from wums import logging + +logger = logging.child_logger(__name__) + + +def getDataPath(fallback=None): + """NanoAOD base directory for the current host. + + Raises ValueError on an unknown host unless ``fallback`` is given, in + which case that path is returned instead. + """ + hostname = socket.gethostname() + + if hostname.endswith(".cern.ch"): + return "/scratch/shared/NanoAOD" + elif hostname.endswith(".mit.edu"): + return "/scratch/submit/cms/wmass/NanoAOD" + elif hostname == "cmsanalysis.pi.infn.it": + # NOTE: If anyone wants to run lowpu analysis at Pisa they'd probably want a different path + return "/scratchnvme/wmass/NANOV9/postVFP" + elif hostname == "cmsasymow.pi.infn.it": + return "/scratch/wmass/y2016" + + if fallback is not None: + logger.warning( + f"No data path known for host {hostname}, falling back to {fallback}" + ) + return fallback + raise ValueError( + f"No data path known for host {hostname}; pass an explicit base path" + ) + + +def buildFileListPosix(path): + outfiles = [] + for root, dirs, fnames in os.walk(path): + for fname in fnames: + if fname.lower().endswith(".root"): + outfiles.append(f"{root}/{fname}") + + return outfiles + + +def appendFilesXrd( + filelist, xrdfs, path, suffixes=[".root"], recurse=False, num_clients=16 +): + import XRootD.client + + status, dirlist = xrdfs.dirlist(path, flags=XRootD.client.flags.DirListFlags.STAT) + + if not status.ok: + if status.code == 400 and status.errno == 3011: + logger.warning(f"XRootD directory not found: {path}") + else: + raise RuntimeError( + f"Error in XRootD.client.FileSystem.dirlist: {status.message}, {status.code}, {status.errno}" + ) + + return + + for diritem in dirlist: + is_dir = diritem.statinfo.flags & XRootD.client.flags.StatInfoFlags.IS_DIR + is_other = diritem.statinfo.flags & XRootD.client.flags.StatInfoFlags.OTHER + is_file = not (is_dir or is_other) + + if is_dir and recurse: + childpath = f"{path}/{diritem.name}" + appendFilesXrd( + filelist, + xrdfs, + childpath, + suffixes=suffixes, + recurse=recurse, + num_clients=num_clients, + ) + elif is_file: + lowername = diritem.name.lower() + matchsuffix = False + for suffix in suffixes: + if lowername.endswith(suffix): + matchsuffix = True + break + + if matchsuffix: + if num_clients > 0: + # construct client string if necessary to force multiple xrootd connections + # (needed for good performance when a single or small number of xrootd servers is used) + client = f"user_{random.randrange(num_clients)}" + outname = f"{xrdfs.url.protocol}://{client}@{xrdfs.url.hostname}:{xrdfs.url.port}/{path}/{diritem.name}" + else: + outname = f"{xrdfs.url.protocol}://{xrdfs.url.hostid}/{path}/{diritem.name}" + + filelist.append(outname) + + +def buildFileListXrd(path, num_clients=16): + import XRootD.client + + xrdurl = XRootD.client.URL(path) + + if not xrdurl.is_valid(): + raise ValueError(f"Invalid xrootd path {path}") + + xrdfs = XRootD.client.FileSystem(xrdurl.hostid) + xrdpath = xrdurl.path + + outfiles = [] + appendFilesXrd(outfiles, xrdfs, xrdpath, recurse=True, num_clients=num_clients) + + return outfiles + + +def buildFileList(path): + xrdprefix = "root://" + return ( + buildFileListXrd(path) + if path.startswith(xrdprefix) + else buildFileListPosix(path) + ) + + +# TODO add the rest of the samples! +def makeFilelist( + paths, + maxFiles=-1, + base_path=None, + nano_prod_tags=None, + is_data=False, + oneMCfileEveryN=None, + era=None, +): + filelist = [] + expandedPaths = [] + for orig_path in paths: + if maxFiles > 0 and len(filelist) >= maxFiles: + break + # try each tag in order until files are found + fallback = False + for prod_tag in nano_prod_tags: + format_args = dict(BASE_PATH=base_path, NANO_PROD_TAG=prod_tag, ERA=era) + + path = orig_path.format(**format_args) + expandedPaths.append(path) + logger.debug(f"Reading files from path {path}") + + files = buildFileList(path) + if maxFiles > 0 and len(files) >= maxFiles: + logger.info( + f"Booking {maxFiles} of {len(files)} files with tag {prod_tag} with path {path}" + ) + break + + if len(files) == 0: + fallback = True + logger.warning( + f"Did not find any files for tag {prod_tag} matching path {path}!" + ) + else: + if fallback: + logger.warning(f"Falling back to tag {prod_tag} with path {path}") + else: + logger.info( + f"Booking {maxFiles} of {len(files)} files with tag {prod_tag} with path {path}" + ) + break + + filelist.extend(files) + + toreturn = ( + filelist + if maxFiles < 0 or len(filelist) < maxFiles + else random.Random(1).sample(filelist, maxFiles) + ) + + if oneMCfileEveryN != None and not is_data: + tmplist = [] + for i, f in enumerate(toreturn): + if i % oneMCfileEveryN == 0: + tmplist.append(f) + logger.warning(f"Using {len(tmplist)} files instead of {len(toreturn)}") + toreturn = tmplist + + logger.debug(f"Length of list is {len(toreturn)} for paths {expandedPaths}") + return toreturn From 966e4931e40783b9c834ceb64c8fb4ecb232bae5 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 09:58:09 -0400 Subject: [PATCH 11/31] bump rabbit Co-Authored-By: Claude Fable 5 --- rabbit | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rabbit b/rabbit index 2fa3ae52a..6bfa4ffcb 160000 --- a/rabbit +++ b/rabbit @@ -1 +1 @@ -Subproject commit 2fa3ae52a4f67bad6d8e8bb9a585df7e2beb4244 +Subproject commit 6bfa4ffcbad0e80d58423e1ecbca41654131b2c8 From 2a1550466f7312179b01756a0da3f4068ec7c2e4 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 10:38:30 -0400 Subject: [PATCH 12/31] fix module-level constant reference in discrete-NP guard _DISCRETE_NP_SUBSTRING lives at module scope; reading it via self raised AttributeError on any datacard with systs. Co-Authored-By: Claude Fable 5 --- wremnants/postprocessing/scetlib_np/param_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 478515f65..fd0acb779 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -1046,7 +1046,7 @@ def _check_discrete_np_double_counting(self): syst_names = [s.decode() if isinstance(s, bytes) else str(s) for s in systs] conflicting = [ - s for s in syst_names if self._DISCRETE_NP_SUBSTRING in s.lower() + s for s in syst_names if _DISCRETE_NP_SUBSTRING in s.lower() ] if not conflicting: return From e33ad7112a0cfe5340fbc364d9cbf9cf04d73d85 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 10:42:14 -0400 Subject: [PATCH 13/31] drop the numpy reference implementation from the shipped tree btgrid_numpy was only consumed at runtime through load_btgrid_shards; move that into btgrid_cache and keep the numpy reference + parity tests in the development tree only. Co-Authored-By: Claude Fable 5 --- .../postprocessing/scetlib_np/btgrid_cache.py | 109 ++- .../scetlib_np/btgrid_integrate.py | 9 +- .../postprocessing/scetlib_np/btgrid_numpy.py | 714 ------------------ .../postprocessing/scetlib_np/btgrid_tf.py | 15 +- .../postprocessing/scetlib_np/param_model.py | 12 +- 5 files changed, 122 insertions(+), 737 deletions(-) delete mode 100644 wremnants/postprocessing/scetlib_np/btgrid_numpy.py diff --git a/wremnants/postprocessing/scetlib_np/btgrid_cache.py b/wremnants/postprocessing/scetlib_np/btgrid_cache.py index 675ee5cc4..c3eea7faa 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_cache.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_cache.py @@ -14,11 +14,112 @@ import pickle import time -from wremnants.postprocessing.scetlib_np import btgrid_numpy as fz +import numpy as np _COMBINED_BASENAME = "combined_btgrid.pkl" +def load_btgrid_shards(submitdir_or_glob, runcard_basename=None): + """Combine bT-grid shards produced by --bt-grid mode. + + `submitdir_or_glob` may be: + - a directory: we look for ``*_btgrid.pkl`` inside (recursively only + one level via scetlib_outputs/). + - a glob pattern: used directly. + + Returns dict with: + bT : (Nbt,) + b_bar : (Nbt,) + bins : list of (Q, Y, qT, lep) bin centres, length Nbins + vars : dict of variation index -> setting dict (copied from the + first shard; all shards are expected to carry the same set) + I_pert : (Nvars, Nbins, Nbt) + C_nu : (Nvars, Nbins, Nbt) + config : dict from the first shard (perturbative configuration the + grid was generated against) + n_shards: int + """ + if os.path.isdir(submitdir_or_glob): + candidates = [ + os.path.join(submitdir_or_glob, "scetlib_outputs", "*_btgrid.pkl"), + os.path.join(submitdir_or_glob, "*_btgrid.pkl"), + ] + else: + candidates = [submitdir_or_glob] + + files = [] + for pat in candidates: + files = sorted(glob.glob(pat)) + if files: + break + if not files: + raise FileNotFoundError(f"No btgrid shards found under {submitdir_or_glob!r}") + + # First shard sets the schema; later shards must match. + with open(files[0], "rb") as f: + first = pickle.load(f) + if first.get("schema_version") != "bt_grid_v1": + raise ValueError( + f"Unexpected schema {first.get('schema_version')!r} in {files[0]}" + ) + bT = np.asarray(first["bT"], dtype=float) + b_bar = np.asarray(first["b_bar"], dtype=float) + varis = first["vars"] + config = first["config"] + n_vars = len(varis) + n_bt = bT.size + + # We don't know Nbins ahead of time without scanning all shards. Walk them + # once: build a dict of bin -> (var_idx -> (I_pert_row, C_nu_row)). + bin_to_data = {} + for path in files: + with open(path, "rb") as f: + d = pickle.load(f) + if d.get("schema_version") != "bt_grid_v1": + raise ValueError( + f"Mixed schema versions: {path} has {d.get('schema_version')}" + ) + if d["bT"].shape != bT.shape or not np.allclose(d["bT"], bT): + raise ValueError(f"bT grid mismatch in {path}") + bins_local = d["bins"] + I_local = np.asarray( + d["I_pert"], dtype=float + ) # (Nvars_local, Nbins_local, Nbt) + C_local = np.asarray(d["C_nu"], dtype=float) + # Map local variation indices to the union variation order. We assume + # all shards share the same vars dict (true when produced by the same + # condor submission). + var_order_local = list(d["vars"].keys()) + for b_idx, b_tup in enumerate(bins_local): + slot = bin_to_data.setdefault(tuple(b_tup), {}) + for v_pos, v_idx in enumerate(var_order_local): + slot[v_idx] = (I_local[v_pos, b_idx], C_local[v_pos, b_idx]) + + var_order = list(varis.keys()) + bins_sorted = sorted(bin_to_data.keys(), key=lambda t: (t[0], t[1], t[2])) + n_bins = len(bins_sorted) + I_pert = np.full((n_vars, n_bins, n_bt), np.nan, dtype=float) + C_nu = np.full((n_vars, n_bins, n_bt), np.nan, dtype=float) + for b_pos, b_tup in enumerate(bins_sorted): + per_var = bin_to_data[b_tup] + for v_pos, v_idx in enumerate(var_order): + if v_idx in per_var: + I_pert[v_pos, b_pos] = per_var[v_idx][0] + C_nu[v_pos, b_pos] = per_var[v_idx][1] + + return { + "bT": bT, + "b_bar": b_bar, + "bins": bins_sorted, + "vars": varis, + "var_order": var_order, + "I_pert": I_pert, + "C_nu": C_nu, + "config": config, + "n_shards": len(files), + } + + def _shard_glob(submitdir): for pat in ( os.path.join(submitdir, "scetlib_outputs", "*_btgrid.pkl"), @@ -48,8 +149,8 @@ def load(submitdir, rebuild=False, verbose=True): On first call (or when ``rebuild=True``, or when any shard is newer than the cached combined file), assembles the shards via - :func:`wremnants.postprocessing.scetlib_np.btgrid_numpy.load_btgrid_shards`, - writes ``combined_btgrid.pkl``, and returns the dict. + :func:`load_btgrid_shards`, writes ``combined_btgrid.pkl``, and returns + the dict. On subsequent calls, loads the pickle directly. """ @@ -74,7 +175,7 @@ def load(submitdir, rebuild=False, verbose=True): raise FileNotFoundError(f"No btgrid shards found under {submitdir!r}") t0 = time.time() - grid = fz.load_btgrid_shards(submitdir) + grid = load_btgrid_shards(submitdir) if verbose: print( f"[btgrid_cache] assembled {grid['n_shards']} shards in " diff --git a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py index 57acb50db..5f37d919e 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py @@ -18,8 +18,7 @@ edge list, builds a ``(N_target, N_source)`` Simpson weight matrix. Apply via ``tf.tensordot``. -Parity tests against the numpy reference in :mod:`scetlib_btgrid_numpy` are in -:mod:`scetlib_btgrid_tf_parity` (added in Phase 3). +Parity-tested against the numpy reference implementation (development tree). """ import numpy as np @@ -30,7 +29,7 @@ simpson_weights, ) -# Z resonance defaults (matches integrate_over_Q in btgrid_numpy). +# Z resonance defaults (matches the numpy-reference integrate_over_Q). MZ_PDG = 91.1876 GAMMAZ_PDG = 2.4952 @@ -110,7 +109,7 @@ def q_integrate_weights(Q_grid, Q_lo, Q_hi, q0=MZ_PDG, Gamma=GAMMAZ_PDG): """Simpson weights for integrating over Q ∈ [Q_lo, Q_hi] in arctan-Q² space. Implements the same change of variable as - :func:`scetlib_btgrid_numpy.integrate_over_Q` with ``method="arctan_Q2"``: + the numpy-reference ``integrate_over_Q`` with ``method="arctan_Q2"``: x = arctan((Q² - q0²) / (q0 Γ)) flattens the Breit-Wigner peak, then Simpson on x with the Jacobian dQ/dx. @@ -154,7 +153,7 @@ def rebin_weights(source_grid, target_edges, name="axis", tol=1e-9): target bin. Mirrors the per-bin call pattern of - :func:`scetlib_btgrid_numpy.integrate_over_axis_bin`. + the numpy-reference ``integrate_over_axis_bin``. """ source_grid = np.asarray(source_grid, dtype=np.float64) target_edges = np.asarray(target_edges, dtype=np.float64) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_numpy.py b/wremnants/postprocessing/scetlib_np/btgrid_numpy.py deleted file mode 100644 index 6c0a18ae1..000000000 --- a/wremnants/postprocessing/scetlib_np/btgrid_numpy.py +++ /dev/null @@ -1,714 +0,0 @@ -# ------------------------------------------------------------------------------- -# NP factorization library for the bT-grid workflow. -# -# Vendored copy of: -# /work/submit/lavezzo/alphaS/scetlib-cms-newnp-lambda4fix/prod/scetlib_run/ -# scetlib_run/factorize.py -# Numpy-only; serves as the reference implementation against which the TF port -# (scetlib_btgrid_tf.py, Phase 2) is parity-tested. Kept in sync manually with -# the upstream scetlib repo — when the upstream changes, recopy this file and -# rerun the parity tests. -# -# Provides: -# - Pure-numpy transcriptions of NP_model_effective (F_eff) and -# NP_model_gammanu (gamma_nu^NP) that match the C++ code byte-for-byte. -# - A vectorised Hankel reconstruction of sigma(qT) from a cached bT-grid. -# The full integrand — every factor and its bare-bT / b*(bT) / (Q,Y,qT) / -# lambda dependence — is written out ONCE in the module docstring of -# wremnants/postprocessing/scetlib_np/param_model.py (single source of -# truth). The reconstruct_* functions below implement it. -# - Loaders for the bT-grid pickle shards produced by --bt-grid and for the -# prior-art spectrum-mode "combined" pickles. -# -# Self-contained: depends on numpy only, no scipy / SCETlib runtime at import. -# ------------------------------------------------------------------------------- - -import glob -import os -import pickle -import sys - -import numpy as np - - -# Compat shim: pickles produced with numpy >= 2.0 reference `numpy._core`, -# which does not exist in numpy < 2.0. Alias the old `numpy.core` under the -# new name (and a few of its submodules) so unpickling succeeds. -def _ensure_numpy_core_alias(): - if "numpy._core" in sys.modules: - return - try: - import numpy._core # noqa: F401 - except ImportError: - try: - from numpy import core as _np_core - except ImportError: - return - sys.modules["numpy._core"] = _np_core - for name in ( - "multiarray", - "numeric", - "fromnumeric", - "umath", - "shape_base", - "_methods", - ): - sub = getattr(_np_core, name, None) - if sub is not None: - sys.modules[f"numpy._core.{name}"] = sub - - -_ensure_numpy_core_alias() - - -# ============================================================================= -# Numerics: bare-bones bessel J0 and Simpson, both numpy-only (the container -# in which the fit runs has no scipy). -# ============================================================================= - - -def bessel_j0(x): - """J_0(x) via Abramowitz & Stegun 9.4.1 / 9.4.3. Accurate to ~1.6e-8.""" - x = np.asarray(x, dtype=float) - out = np.empty_like(x) - small = np.abs(x) < 3.0 - xs = x[small] / 3.0 - y = xs * xs - out[small] = ( - 1.0 - - 2.2499997 * y - + 1.2656208 * y**2 - - 0.3163866 * y**3 - + 0.0444479 * y**4 - - 0.0039444 * y**5 - + 0.0002100 * y**6 - ) - xl = np.abs(x[~small]) - z = 3.0 / xl - f0 = ( - 0.79788456 - - 0.00000077 * z - - 0.00552740 * z**2 - - 0.00009512 * z**3 - + 0.00137237 * z**4 - - 0.00072805 * z**5 - + 0.00014476 * z**6 - ) - theta0 = ( - xl - - 0.78539816 - - 0.04166397 * z - - 0.00003954 * z**2 - + 0.00262573 * z**3 - - 0.00054125 * z**4 - - 0.00029333 * z**5 - + 0.00013558 * z**6 - ) - out[~small] = f0 * np.cos(theta0) / np.sqrt(xl) - return out - - -def simpson(y, x, axis=-1): - """Composite Simpson on a (possibly non-uniform) 1-D grid, vectorised - along the given axis. Falls back to trapezoid for the last segment when - the number of intervals is odd.""" - y = np.asarray(y, dtype=float) - x = np.asarray(x, dtype=float) - if x.ndim != 1: - raise ValueError("simpson expects 1-D x") - n = x.size - 1 - if n < 1: - return np.zeros( - y.shape[:-1] if axis == -1 else y.shape[:axis] + y.shape[axis + 1 :] - ) - # move integration axis to the end for simpler slicing - y_moved = np.moveaxis(y, axis, -1) - if n % 2 == 1: - lead = simpson(np.moveaxis(y_moved[..., :-1], -1, axis), x[:-1], axis=axis) - tail = 0.5 * (y_moved[..., -1] + y_moved[..., -2]) * (x[-1] - x[-2]) - return lead + tail - h = np.diff(x) - h0 = h[0::2] - h1 = h[1::2] - s = ( - (h0 + h1) - / 6.0 - * ( - y_moved[..., 0:-1:2] * (2.0 - h1 / h0) - + y_moved[..., 1::2] * (h0 + h1) ** 2 / (h0 * h1) - + y_moved[..., 2::2] * (2.0 - h0 / h1) - ) - ) - return np.sum(s, axis=-1) - - -# ============================================================================= -# NP model transcriptions. Mirror the C++ implementations one-to-one; expect -# bT to be the b̄T (= b_star) at which the NP factor enters. -# ============================================================================= - -# NP_model_effective.np_model enum values supported here: -EFF_MODELS = { - "identity", - "tanh_2", - "tanh_6", - "tanh_4", - "frac_2", - "frac_4", - "exp_2", - "exp_4", - "signed_lambda", - "hyp_tangent", - "square_root", -} - -# NP_model_gammanu.np_model_nu enum values supported here: -GNU_MODELS = { - "tanh_1", - "tanh_2", - "tanh_6", - "frac_1", - "frac_2", - "exp_1", - "exp_2", - "hyp_tangent", - "linear", -} - - -def F_eff(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_model): - """F_eff(Y, b̄T) — NP_model_effective::operator() from NP_models.hpp.""" - bT = np.asarray(bT, dtype=float) - - if np_model == "signed_lambda": - lambda2_Y = lambda2 + delta_lambda2 * Y * Y - if lambda4 <= 0.0 and (lambda2 != 0.0 or delta_lambda2 != 0.0): - raise ValueError( - "signed_lambda requires lambda4 > 0 when lambda2 or delta_lambda2 != 0" - ) - return (1.0 + lambda2_Y * bT**2) ** 2 * np.exp(-2.0 * lambda4 * bT**4) - - lambda2_Y = lambda2 + delta_lambda2 * Y * Y - arg = (lambda2_Y + lambda4 * bT**2) * bT - - if np_model == "identity": - return np.exp(-2.0 * bT * arg) - - if lambda_inf == 0.0: - return np.ones_like(bT) - - arg = arg / lambda_inf - # alias support - model = {"hyp_tangent": "tanh_2", "square_root": "frac_2"}.get(np_model, np_model) - - if model == "tanh_2": - arg = arg + (1.0 / 3.0) * (lambda2_Y * bT / lambda_inf) ** 3 - func = np.tanh(arg) - elif model == "tanh_6": - arg = arg + lambda6 * bT**5 / lambda_inf - arg = arg + (1.0 / 3.0) * (lambda2_Y * bT / lambda_inf) ** 3 - func = np.tanh(arg) - elif model == "tanh_4": - func = np.sqrt(np.tanh(arg**2)) - elif model == "frac_2": - arg = arg + 0.5 * (lambda2_Y * bT / lambda_inf) ** 3 - func = arg / np.sqrt(1.0 + arg**2) - elif model == "frac_4": - func = arg / np.sqrt(np.sqrt(1.0 + arg**4)) - elif model == "exp_2": - arg = arg + 0.25 * (lambda2_Y * bT / lambda_inf) ** 3 - func = np.sqrt(-np.expm1(-(arg**2))) - elif model == "exp_4": - func = np.sqrt(np.sqrt(-np.expm1(-(arg**4)))) - else: - raise ValueError(f"F_eff: unsupported np_model {np_model!r}") - - return np.exp(-2.0 * lambda_inf * bT * func) - - -def gamma_nu_NP(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): - """gamma_nu^NP(b̄T) — NP_model_gammanu::model_gammanu() from Gamma_nu.hpp. - - Note: NP_model_gammanu has its own b_star() (b0_bmax_nu), but the C++ - code calls model_gammanu(bT) with the raw bT passed into Gamma_nu — which - is b_star_global by the time it reaches us. So pass b̄T = b_star_global(bT) - here; b0_bmax_nu does NOT enter model_gammanu directly. - """ - bT = np.asarray(bT, dtype=float) - if lambda_inf_nu == 0.0: - return np.zeros_like(bT) - - bT2 = bT * bT - arg = (lambda2_nu + lambda4_nu * bT2) * bT2 / lambda_inf_nu - - model = {"hyp_tangent": "tanh_2", "linear": "frac_1"}.get(np_model_nu, np_model_nu) - - if model == "tanh_1": - arg = arg + (2.0 / 3.0) * (lambda2_nu * bT2 / lambda_inf_nu) ** 2 - func = np.tanh(np.sqrt(arg)) ** 2 - elif model == "tanh_2": - func = np.tanh(arg) - elif model == "tanh_6": - # NP_model_gammanu hardcodes lambda6_nu = 0.0007 (Gamma_nu.hpp:102) - arg = arg + 0.0007 * bT2**3 / lambda_inf_nu - func = np.tanh(arg) - elif model == "frac_1": - arg = arg + (lambda2_nu * bT2 / lambda_inf_nu) ** 2 - func = arg / (1.0 + arg) - elif model == "frac_2": - func = arg / np.sqrt(1.0 + arg**2) - elif model == "exp_1": - arg = arg + 0.5 * (lambda2_nu * bT2 / lambda_inf_nu) ** 2 - func = -np.expm1(-arg) - elif model == "exp_2": - func = np.sqrt(-np.expm1(-(arg**2))) - else: - raise ValueError(f"gamma_nu_NP: unsupported np_model_nu {np_model_nu!r}") - - return -lambda_inf_nu * func - - -# Convenience defaults: "all NP knobs off". These reproduce the SCETlib -# configuration that --bt-grid mode caches. -NP_ZERO_EFF = dict( - lambda_inf=0.0, - lambda2=0.0, - lambda4=0.0, - lambda6=0.0, - delta_lambda2=0.0, - np_model="identity", -) -NP_ZERO_GNU = dict( - lambda_inf_nu=0.0, lambda2_nu=0.0, lambda4_nu=0.0, np_model_nu="tanh_2" -) - - -def eff_params_from_conf(conf): - """Read NP_model_effective parameters from a configparser [Nonperturbative] - section (defaults to NP_ZERO_EFF if a field is absent).""" - if "Nonperturbative" not in conf: - return dict(NP_ZERO_EFF) - sec = conf["Nonperturbative"] - return dict( - lambda_inf=sec.getfloat("lambda_inf", fallback=NP_ZERO_EFF["lambda_inf"]), - lambda2=sec.getfloat("lambda2", fallback=NP_ZERO_EFF["lambda2"]), - lambda4=sec.getfloat("lambda4", fallback=NP_ZERO_EFF["lambda4"]), - lambda6=sec.getfloat("lambda6", fallback=NP_ZERO_EFF["lambda6"]), - delta_lambda2=sec.getfloat( - "delta_lambda2", fallback=NP_ZERO_EFF["delta_lambda2"] - ), - np_model=sec.get("np_model", fallback=NP_ZERO_EFF["np_model"]), - ) - - -def gnu_params_from_conf(conf): - """Read NP_model_gammanu parameters from a configparser [Nonperturbative] - section (defaults to NP_ZERO_GNU if a field is absent).""" - if "Nonperturbative" not in conf: - return dict(NP_ZERO_GNU) - sec = conf["Nonperturbative"] - return dict( - lambda_inf_nu=sec.getfloat( - "lambda_inf_nu", fallback=NP_ZERO_GNU["lambda_inf_nu"] - ), - lambda2_nu=sec.getfloat("lambda2_nu", fallback=NP_ZERO_GNU["lambda2_nu"]), - lambda4_nu=sec.getfloat("lambda4_nu", fallback=NP_ZERO_GNU["lambda4_nu"]), - np_model_nu=sec.get("np_model_nu", fallback=NP_ZERO_GNU["np_model_nu"]), - ) - - -# ============================================================================= -# Hankel reconstruction -# ============================================================================= - - -def reconstruct_one(qT, bT, I_pert, C_nu, b_bar, Y, eff_params, gnu_params): - """Hankel-reconstruct one sigma(qT) from cached arrays at a single (Q,Y,qT). - - Returns the differential structure-function value at the point (Q,Y,qT), - matching SCETlib's spectrum-mode ang.c convention (i.e. including the - qT factor that arises from the integration-variable choice in SCETlib's - integrator_de_oscillatory). - - qT : scalar - bT : (Nb,) integration variable (raw bT) - I_pert : (Nb,) cached SCETlib integrand at NP off - C_nu : (Nb,) coefficient of gamma_nu^NP in log evolution - b_bar : (Nb,) b_star_global(bT) where NP factors evaluate - Y : scalar rapidity - eff_params : dict for NP_model_effective parameters - gnu_params : dict for NP_model_gammanu parameters - """ - g_NP = gamma_nu_NP(b_bar, **gnu_params) - Feff = F_eff(Y, b_bar, **eff_params) - integrand = bT * bessel_j0(qT * bT) * I_pert * np.exp(C_nu * g_NP) * Feff - return qT * simpson(integrand, bT) - - -def reconstruct_grid_QYqT( - Q_grid, - Y_grid, - qT_grid, - bT, - I_pert, - C_nu, - b_bar, - eff_params, - gnu_params, - verbose=True, -): - """Reconstruct sigma at every (Q, Y, qT) sample point of a regularly-indexed - grid. Computes J0(qT*bT) once on (NqT, Nbt) (rather than the redundant - (NQ*NY*NqT, Nbt) the naive batch would build) and processes one Q-slice at - a time to bound peak memory and emit progress. - - Q_grid : (NQ,) point values in Q - Y_grid : (NY,) point values in Y - qT_grid : (NqT,) point values in qT - bT : (Nbt,) bT integration variable shared by all points - I_pert : (Npts, Nbt) cached perturbative integrand, Npts = NQ*NY*NqT, - indexed as flat list in the same order load_btgrid_shards - returns its `bins` list (sorted lexicographically by Q, Y, qT). - C_nu : (Npts, Nbt) rapidity-evolution-log coefficient - b_bar : (Nbt,) b_star_global(bT) - - Returns: (NQ, NY, NqT) ndarray of sigma_factorised values. - """ - Q_grid = np.asarray(Q_grid, dtype=float) - Y_grid = np.asarray(Y_grid, dtype=float) - qT_grid = np.asarray(qT_grid, dtype=float) - bT = np.asarray(bT, dtype=float) - NQ, NY, NqT = Q_grid.size, Y_grid.size, qT_grid.size - Nbt = bT.size - Npts = NQ * NY * NqT - if I_pert.shape != (Npts, Nbt): - raise ValueError( - f"I_pert shape {I_pert.shape} doesn't match expected " - f"({Npts}, {Nbt}) from {NQ} Q x {NY} Y x {NqT} qT" - ) - - # --- shared (Y-independent) factors over bT --- - g_NP = gamma_nu_NP(b_bar, **gnu_params) # (Nbt,) - delta_l2 = eff_params.get("delta_lambda2", 0.0) - if delta_l2 == 0.0: - Feff_bT = F_eff(0.0, b_bar, **eff_params) # (Nbt,) - else: - Feff_bT = None # per-Y below - - # --- bT*J0(qT*bT) cached once over (NqT, Nbt) --- - bT_J0 = bT[np.newaxis, :] * bessel_j0(qT_grid[:, np.newaxis] * bT[np.newaxis, :]) - # shape (NqT, Nbt); used in every Q-slice below. - - # I_pert and C_nu are stored as (NQ*NY*NqT, Nbt); reshape view to - # (NQ, NY, NqT, Nbt) and process Q-by-Q. - I_pert_r = I_pert.reshape(NQ, NY, NqT, Nbt) - C_nu_r = C_nu.reshape(NQ, NY, NqT, Nbt) - - out = np.empty((NQ, NY, NqT), dtype=float) - - import time - - t0 = time.time() - for iQ in range(NQ): - # exp(C_nu * g_NP) on the (NY, NqT, Nbt) Q-slice -- the only piece - # that doesn't factor across the slice - exp_g_factor = np.exp( - C_nu_r[iQ] * g_NP[np.newaxis, np.newaxis, :] - ) # (NY, NqT, Nbt) - if delta_l2 == 0.0: - # Feff_bT shape (Nbt,); broadcast over (NY, NqT) - integrand = ( - bT_J0[np.newaxis, :, :] - * I_pert_r[iQ] - * exp_g_factor - * Feff_bT[np.newaxis, np.newaxis, :] - ) - else: - # Feff depends on Y (per-Y row), shape (NY, Nbt) broadcast over qT - Feff = np.stack( - [F_eff(Y_i, b_bar, **eff_params) for Y_i in Y_grid] - ) # (NY, Nbt) - integrand = ( - bT_J0[np.newaxis, :, :] - * I_pert_r[iQ] - * exp_g_factor - * Feff[:, np.newaxis, :] - ) - # Simpson over bT (last axis) - sigma_Q = simpson(integrand, bT, axis=-1) # (NY, NqT) - # qT factor (SCETlib's x = qT*bT integration convention) - out[iQ] = qT_grid[np.newaxis, :] * sigma_Q - - if verbose: - elapsed = time.time() - t0 - print( - f" [hankel] Q-slice {iQ+1}/{NQ} done " - f"({elapsed:.1f}s elapsed, ETA {elapsed*(NQ-iQ-1)/(iQ+1):.1f}s)", - flush=True, - ) - return out - - -def integrate_over_Q( - sigma_QYqT, Q_grid, Q_lo, Q_hi, method="arctan_Q2", q0=91.1876, Gamma=2.4952 -): - """Integrate sigma(Q, Y, qT) over Q in [Q_lo, Q_hi]. Q-samples outside - [Q_lo, Q_hi] are dropped. - - sigma_QYqT : (NQ, NY, NqT) point values - Q_grid : (NQ,) Q sample positions (need not be uniform) - Q_lo, Q_hi : integration limits - method : "arctan_Q2" (default; uses x = arctan((Q²-q0²)/(q0*Gamma)) so - the Breit-Wigner Z resonance becomes smooth) or "simpson" - (Simpson directly in Q) or "trapz" (trapezoid in Q). - q0, Gamma : resonance mass and width for the arctan_Q2 transform. - Defaults match Z-boson values (mZ = 91.1876, ΓZ = 2.4952). - - Returns: (NY, NqT) integrated values. - """ - Q_grid = np.asarray(Q_grid, dtype=float) - mask = (Q_grid >= Q_lo) & (Q_grid <= Q_hi) - if mask.sum() < 2: - raise ValueError(f"Need >= 2 Q samples in [{Q_lo}, {Q_hi}]; got {mask.sum()}") - Q_sub = Q_grid[mask] - s_sub = sigma_QYqT[mask] # (NQ_sub, NY, NqT) - s_moved = np.moveaxis(s_sub, 0, -1) # (..., NQ_sub) - - if method == "simpson": - return simpson(s_moved, Q_sub, axis=-1) - if method == "trapz": - return np.trapz(s_moved, Q_sub, axis=-1) - if method == "arctan_Q2": - # x = arctan((Q² - q0²) / (q0 * Gamma)) → flattens the Breit-Wigner peak - x = np.arctan((Q_sub**2 - q0**2) / (q0 * Gamma)) - # dQ/dx = ( q0*Gamma + (Q² - q0²)² / (q0*Gamma) ) / (2 Q) - jac = (q0 * Gamma + (Q_sub**2 - q0**2) ** 2 / (q0 * Gamma)) / (2.0 * Q_sub) - return simpson(s_moved * jac, x, axis=-1) - raise ValueError(f"integrate_over_Q: unknown method {method!r}") - - -def integrate_over_axis_bin(values, axis_grid, axis_lo, axis_hi, name="axis"): - """Integrate a 1-D array of sample values over a single bin [axis_lo, axis_hi] - using Simpson's rule on whatever sample points fall in (and on) the bin edges. - - Designed for use with a btgrid sampled at bin-edges + bin-centres (3 samples - per bin minimum). When the grid carries the bin's 2 edges + central point, - this is a 3-point Simpson (4th-order accurate per bin); when more samples - fall inside the bin, simpson naturally extends to higher order via composite - rule. - - values : (..., N_axis) array values at axis_grid sample points - axis_grid : (N_axis,) sample positions (sorted, not necessarily uniform) - axis_lo, axis_hi : integration limits (bin edges) - - Returns: integrated value with the axis collapsed. - """ - axis_grid = np.asarray(axis_grid, dtype=float) - values = np.asarray(values, dtype=float) - # tolerance for "on the edge" lookups (FP-imprecise bin centres show up as - # e.g. -2.3499999999999996; bin widths are at least ~0.025 so 1e-9 is safe) - tol = 1e-9 - mask = (axis_grid >= axis_lo - tol) & (axis_grid <= axis_hi + tol) - if mask.sum() < 2: - raise ValueError( - f"integrate_over_axis_bin({name}): need >= 2 samples in " - f"[{axis_lo}, {axis_hi}]; got {mask.sum()}" - ) - sub = values[..., mask] - g = axis_grid[mask] - return simpson(sub, g, axis=-1) - - -def integrate_over_Y_bin(sigma_YqT, Y_grid, Y_lo, Y_hi): - """Integrate sigma(Y, qT) over Y in [Y_lo, Y_hi] using sample points of - Y_grid that fall in or on the bin edges. Returns (NqT,) array. - - Expects Y_grid to include Y_lo and Y_hi as samples (the edges) and ideally - the bin centre too — Simpson uses all available samples in [Y_lo, Y_hi].""" - # move Y axis (axis 0) to the end for integrate_over_axis_bin which uses last axis - return integrate_over_axis_bin( - np.moveaxis(sigma_YqT, 0, -1), Y_grid, Y_lo, Y_hi, name="Y" - ) - - -def integrate_over_qT_bin(sigma_YqT, qT_grid, qT_lo, qT_hi): - """Integrate sigma(Y, qT) over qT in [qT_lo, qT_hi]. Returns (NY,) array. - - Expects qT_grid to include qT_lo and qT_hi as samples (the bin edges).""" - return integrate_over_axis_bin(sigma_YqT, qT_grid, qT_lo, qT_hi, name="qT") - - -def reconstruct_batch( - qT_per_bin, bT, I_pert, C_nu, b_bar, Y_per_bin, eff_params, gnu_params -): - """Vectorised reconstruction for many bins at once. - - qT_per_bin : (Nbins,) qT for each bin - Y_per_bin : (Nbins,) Y for each bin (used by F_eff's δλ_2·Y² term) - bT : (Nbt,) integration variable (raw bT), shared by all bins - I_pert : (Nbins, Nbt) cached perturbative bT integrand - C_nu : (Nbins, Nbt) rapidity log coefficient - b_bar : (Nbt,) b_star_global(bT), shared by all bins - - Returns: (Nbins,) sigma_factorized for each bin - """ - qT_per_bin = np.asarray(qT_per_bin, dtype=float) - Y_per_bin = np.asarray(Y_per_bin, dtype=float) - bT = np.asarray(bT, dtype=float) - I_pert = np.asarray(I_pert, dtype=float) - C_nu = np.asarray(C_nu, dtype=float) - b_bar = np.asarray(b_bar, dtype=float) - - # gamma_nu^NP and the gamma_nu exponential factor are bin-shared - # (bT-dependent only); F_eff depends on Y through δλ_2·Y², so we evaluate - # it per bin if δλ_2 != 0. - g_NP = gamma_nu_NP(b_bar, **gnu_params) # (Nbt,) - exp_g_factor = np.exp(C_nu * g_NP[np.newaxis, :]) # (Nbins, Nbt) - - delta_l2 = eff_params.get("delta_lambda2", 0.0) - if delta_l2 == 0.0: - # F_eff has no Y dependence beyond a constant Y² factor, so identical - # across bins -> compute once - Feff = F_eff(0.0, b_bar, **eff_params) # (Nbt,) - bT_J0 = bT * bessel_j0( - qT_per_bin[:, np.newaxis] * bT[np.newaxis, :] - ) # (Nbins, Nbt) - integrand = bT_J0 * I_pert * exp_g_factor * Feff[np.newaxis, :] - else: - # need per-bin F_eff because of Y dependence in lambda2_Y - Feff = np.empty_like(I_pert) - for i, Y_i in enumerate(Y_per_bin): - Feff[i] = F_eff(Y_i, b_bar, **eff_params) - bT_J0 = bT * bessel_j0(qT_per_bin[:, np.newaxis] * bT[np.newaxis, :]) - integrand = bT_J0 * I_pert * exp_g_factor * Feff - - # Multiply by qT to match SCETlib's spectrum-mode integration convention - # (SCETlib's _int_bT integrates in x = qT*bT, picking up an explicit qT - # factor via the Jacobian; we integrate in bT directly so we must add it - # back). - return qT_per_bin * simpson(integrand, bT, axis=-1) - - -# ============================================================================= -# Loaders for the artefacts produced by the bT-grid condor run and by the -# spectrum-mode "combined" pickles. -# ============================================================================= - - -def load_btgrid_shards(submitdir_or_glob, runcard_basename=None): - """Combine bT-grid shards produced by --bt-grid mode. - - `submitdir_or_glob` may be: - - a directory: we look for ``*_btgrid.pkl`` inside (recursively only - one level via scetlib_outputs/). - - a glob pattern: used directly. - - Returns dict with: - bT : (Nbt,) - b_bar : (Nbt,) - bins : list of (Q, Y, qT, lep) bin centres, length Nbins - vars : dict of variation index -> setting dict (copied from the - first shard; all shards are expected to carry the same set) - I_pert : (Nvars, Nbins, Nbt) - C_nu : (Nvars, Nbins, Nbt) - config : dict from the first shard (perturbative configuration the - grid was generated against) - n_shards: int - """ - if os.path.isdir(submitdir_or_glob): - candidates = [ - os.path.join(submitdir_or_glob, "scetlib_outputs", "*_btgrid.pkl"), - os.path.join(submitdir_or_glob, "*_btgrid.pkl"), - ] - else: - candidates = [submitdir_or_glob] - - files = [] - for pat in candidates: - files = sorted(glob.glob(pat)) - if files: - break - if not files: - raise FileNotFoundError(f"No btgrid shards found under {submitdir_or_glob!r}") - - # First shard sets the schema; later shards must match. - with open(files[0], "rb") as f: - first = pickle.load(f) - if first.get("schema_version") != "bt_grid_v1": - raise ValueError( - f"Unexpected schema {first.get('schema_version')!r} in {files[0]}" - ) - bT = np.asarray(first["bT"], dtype=float) - b_bar = np.asarray(first["b_bar"], dtype=float) - varis = first["vars"] - config = first["config"] - n_vars = len(varis) - n_bt = bT.size - - # We don't know Nbins ahead of time without scanning all shards. Walk them - # once: build a dict of bin -> (var_idx -> (I_pert_row, C_nu_row)). - bin_to_data = {} - for path in files: - with open(path, "rb") as f: - d = pickle.load(f) - if d.get("schema_version") != "bt_grid_v1": - raise ValueError( - f"Mixed schema versions: {path} has {d.get('schema_version')}" - ) - if d["bT"].shape != bT.shape or not np.allclose(d["bT"], bT): - raise ValueError(f"bT grid mismatch in {path}") - bins_local = d["bins"] - I_local = np.asarray( - d["I_pert"], dtype=float - ) # (Nvars_local, Nbins_local, Nbt) - C_local = np.asarray(d["C_nu"], dtype=float) - # Map local variation indices to the union variation order. We assume - # all shards share the same vars dict (true when produced by the same - # condor submission). - var_order_local = list(d["vars"].keys()) - for b_idx, b_tup in enumerate(bins_local): - slot = bin_to_data.setdefault(tuple(b_tup), {}) - for v_pos, v_idx in enumerate(var_order_local): - slot[v_idx] = (I_local[v_pos, b_idx], C_local[v_pos, b_idx]) - - var_order = list(varis.keys()) - bins_sorted = sorted(bin_to_data.keys(), key=lambda t: (t[0], t[1], t[2])) - n_bins = len(bins_sorted) - I_pert = np.full((n_vars, n_bins, n_bt), np.nan, dtype=float) - C_nu = np.full((n_vars, n_bins, n_bt), np.nan, dtype=float) - for b_pos, b_tup in enumerate(bins_sorted): - per_var = bin_to_data[b_tup] - for v_pos, v_idx in enumerate(var_order): - if v_idx in per_var: - I_pert[v_pos, b_pos] = per_var[v_idx][0] - C_nu[v_pos, b_pos] = per_var[v_idx][1] - - return { - "bT": bT, - "b_bar": b_bar, - "bins": bins_sorted, - "vars": varis, - "var_order": var_order, - "I_pert": I_pert, - "C_nu": C_nu, - "config": config, - "n_shards": len(files), - } - - -def load_spectrum_reference(combined_pkl): - """Load a 'combined' spectrum-mode pickle written by scetlib-run-qT.py. - - Returns a dict - hist : the hist.Hist object stored in the pickle - config : the [section]->dict mapping the run used - meta_data: the meta-data dict - Use the returned hist directly (axes are typically Q, Y, qT, lep, vars). - """ - with open(combined_pkl, "rb") as f: - d = pickle.load(f) - return { - "hist": d.get("hist"), - "config": d.get("config", {}), - "meta_data": d.get("meta_data", {}), - } diff --git a/wremnants/postprocessing/scetlib_np/btgrid_tf.py b/wremnants/postprocessing/scetlib_np/btgrid_tf.py index 10672ab65..35917bfb7 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_tf.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_tf.py @@ -1,8 +1,9 @@ """TensorFlow port of the bT-grid factorization library. -Mirrors :mod:`scetlib_btgrid_numpy` function-by-function. The numpy module is -the byte-for-byte transcription of SCETlib C++ and the parity test in -:mod:`scetlib_btgrid_tf_parity` keeps the two in sync. +Function-by-function TF port of a numpy reference implementation (itself a +byte-for-byte transcription of SCETlib C++). The numpy reference and the +parity tests that keep the two in sync live in the development tree; only +the TF port is shipped here. Design choices: * ``np_model`` / ``np_model_nu`` strings are fixed at trace time (the SCETlib @@ -49,7 +50,7 @@ def simpson_weights(x): """Return weights ``w`` such that Simpson(y, x) == sum(w * y, axis=-1). ``x`` is a numpy array with size ``N``. Implementation mirrors the numpy - ``simpson`` in :mod:`scetlib_btgrid_numpy` (composite Simpson with + ``simpson`` of the numpy reference (composite Simpson with trapezoid fallback on the last segment when N-1 is odd). """ x = np.asarray(x, dtype=np.float64) @@ -146,7 +147,7 @@ def _safe_div(num, den): def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_model): - """TF port of :func:`scetlib_btgrid_numpy.F_eff` for a fixed ``np_model``.""" + """TF port of the numpy-reference ``F_eff`` for a fixed ``np_model``.""" if np_model not in EFF_MODELS: raise ValueError(f"F_eff_tf: unsupported np_model {np_model!r}") @@ -202,7 +203,7 @@ def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_ def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): - """TF port of :func:`scetlib_btgrid_numpy.gamma_nu_NP` for fixed ``np_model_nu``.""" + """TF port of the numpy-reference ``gamma_nu_NP`` for fixed ``np_model_nu``.""" if np_model_nu not in GNU_MODELS: raise ValueError(f"gamma_nu_NP_tf: unsupported np_model_nu {np_model_nu!r}") @@ -266,7 +267,7 @@ def reconstruct_batch_tf( Y_unique=None, Y_inverse_idx=None, ): - """TF port of :func:`scetlib_btgrid_numpy.reconstruct_batch`. + """TF port of the numpy-reference ``reconstruct_batch``. Implements the per-(Q, Y, qT) bT-space integrand. The full formula with every factor and its bare-bT / b*(bT) / (Q,Y,qT) / λ dependence is written diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index fd0acb779..3e6306ea0 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -4,9 +4,9 @@ SCETlib nonperturbative (NP) prediction at the fitted λ vs. at λ_central. The prediction is built in THREE STEPS, written out in order below — read top to bottom for the full maths. This module docstring is the SINGLE SOURCE OF TRUTH: -:mod:`response_matrix`, :func:`btgrid_tf.reconstruct_batch_tf`, the numpy -reference :mod:`btgrid_numpy`, ``sigma_reco_central.md``, and the validation -scripts all point here rather than restate it. +:mod:`response_matrix`, :func:`btgrid_tf.reconstruct_batch_tf`, +``sigma_reco_central.md``, and the development-tree numpy reference and +validation scripts all point here rather than restate it. Pipeline at a glance (everything is a function of the NP parameters λ): @@ -48,7 +48,7 @@ C_ν coefficient of γ_ν^NP in the rapidity (CS) log evolution. Cached per (Q, Y, qT) AND bT — a full (Nbins, Nbt) array; λ-independent. Taken at the BARE bT: the lone NP-exponent factor NOT frozen to - b* (per the SCETlib convention; see btgrid_numpy NP-model notes). + b* (per the SCETlib convention). Because it varies with bT it sits inside the bT integral (not a constant out front); because it varies with Q the Q integral cannot be collapsed ahead of the fit — the λ-dependence does not @@ -1045,9 +1045,7 @@ def _check_discrete_np_double_counting(self): return syst_names = [s.decode() if isinstance(s, bytes) else str(s) for s in systs] - conflicting = [ - s for s in syst_names if _DISCRETE_NP_SUBSTRING in s.lower() - ] + conflicting = [s for s in syst_names if _DISCRETE_NP_SUBSTRING in s.lower()] if not conflicting: return From 8e113ea1b4ced68910da6898e0e8eb7be21dc2ba Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Mon, 22 Jun 2026 15:51:49 -0400 Subject: [PATCH 14/31] Address PR review (kdlong): s-dep-width MZ const, propagate NP runcard to histmaker meta, de-verbose docstring - Move Z mass/width to common.MZ_S_DEP_WIDTH / GAMMAZ_S_DEP_WIDTH (s-dependent-width scheme 91.1535/2.4932, not the PDG pole values); btgrid_integrate uses them for the arctan-Q^2 change of variable. The choice only sets the transform centre/scale: <0.001% effect on the integral on the real btgrid Q grid. - Propagate the SCETlib NP runcard from the theoryCorr pkl into the histmaker output meta_info (scetlib_np_lambda_central); the fit reads lambda_central from metadata instead of re-opening the pkl by filename. Non-backward- compatible: fit inputs produced before this must be remade. - De-verbose the param_model module docstring (drop "single source of truth" framing, convert Unicode box dividers to ASCII). (The "drop the numpy reference implementation" item was already on this branch.) Co-Authored-By: Claude Opus 4.8 (1M context) --- .../scetlib_np/btgrid_integrate.py | 11 +- .../scetlib_np/lambda_central.py | 332 ++++++++++-------- .../postprocessing/scetlib_np/param_model.py | 40 +-- wremnants/production/histmaker_tools.py | 33 +- wremnants/utilities/common.py | 7 + 5 files changed, 241 insertions(+), 182 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py index 5f37d919e..273d694d0 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py @@ -28,10 +28,13 @@ _as_dtype, simpson_weights, ) +from wremnants.utilities import common as wrem_common -# Z resonance defaults (matches the numpy-reference integrate_over_Q). -MZ_PDG = 91.1876 -GAMMAZ_PDG = 2.4952 +# Z resonance parameters for the Q-integration change of variable, in the +# s-dependent-width scheme (see wremnants.utilities.common). Only set the centre +# and scale of the arctan-Q^2 transform below; they do not change the physics. +MZ_S_DEP_WIDTH = wrem_common.MZ_S_DEP_WIDTH +GAMMAZ_S_DEP_WIDTH = wrem_common.GAMMAZ_S_DEP_WIDTH # ============================================================================= @@ -105,7 +108,7 @@ def sparse_to_dense_tf(sigma_flat, flat_idx): # ============================================================================= -def q_integrate_weights(Q_grid, Q_lo, Q_hi, q0=MZ_PDG, Gamma=GAMMAZ_PDG): +def q_integrate_weights(Q_grid, Q_lo, Q_hi, q0=MZ_S_DEP_WIDTH, Gamma=GAMMAZ_S_DEP_WIDTH): """Simpson weights for integrating over Q ∈ [Q_lo, Q_hi] in arctan-Q² space. Implements the same change of variable as diff --git a/wremnants/postprocessing/scetlib_np/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py index c5b638c59..abba79d80 100644 --- a/wremnants/postprocessing/scetlib_np/lambda_central.py +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -1,20 +1,28 @@ -"""λ_central auto-detect from a fit-input hdf5. +"""Central NP (lambda) parameters for the SCETlib ParamModel. -The SCETlib correction's NP runcard is preserved in the upstream -``*_Corr.pkl.lz4`` file (under ``file_meta_data..config.Nonperturbative``), -but it is **not** propagated through the histmaker into the fit-input hdf5. We -read the correction tag from the hdf5 (``meta_info_input.args.theoryCorr``), -resolve to the upstream pkl, and extract the Nonperturbative section. +The SCETlib correction's Nonperturbative runcard lives in the upstream +``*_Corr.pkl.lz4`` file under +``file_meta_data..config.Nonperturbative``. The histmaker reads that +section when it applies the correction and writes the parsed values into its +output metadata (key ``scetlib_np_lambda_central``); see +:func:`build_lambda_central_meta`. The fit then reads them back from the +metadata that rabbit propagates into the datacard / fitresults -- it never +re-opens the upstream pkl. -Single entry point: +Write side (histmaker): + build_lambda_central_meta(theory_corr_tags, procs) -> {proc: lambda_central} +Read side (fit / postprocessing): read_lambda_central(hdf5_path, proc="Z") -> dict + read_lambda_central_from_meta(meta, proc="Z") -> dict + +A ``lambda_central`` dict has keys ``tag``, ``basename``, ``eff_params`` (for +NP_model_effective / F_eff) and ``gnu_params`` (for NP_model_gammanu). """ import json import os import pickle -import sys import h5py import lz4.frame @@ -22,47 +30,35 @@ from wremnants.utilities import common as wrem_common from wums import ioutils as wums_io -# Names of the parameters the ParamModel cares about, split by which scetlib -# C++ struct consumes them. Values in the Nonperturbative section are strings; -# numeric ones get parsed to float, model names stay as strings. +# Metadata key under which the histmaker stores the parsed central runcard. +META_KEY = "scetlib_np_lambda_central" + +# Parameter names the ParamModel needs, split by the scetlib C++ struct that +# consumes them. Nonperturbative values are strings; numeric ones get floated, +# model names stay strings. GNU_NUMERIC = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") GNU_STRING = ("np_model_nu",) EFF_NUMERIC = ("lambda_inf", "lambda2", "lambda4", "lambda6", "delta_lambda2") EFF_STRING = ("np_model",) -def _correction_pkl_path(tag, proc): - """Resolve a theoryCorr tag to its upstream pkl.lz4 path.""" - return os.path.join( - wrem_common.data_dir, "TheoryCorrections", f"{tag}_Corr{proc}.pkl.lz4" - ) - - -def _load_correction_pkl(tag, proc): - path = _correction_pkl_path(tag, proc) - if not os.path.exists(path): - raise FileNotFoundError( - f"SCETlib correction pkl not found: {path!r}. " - f"Cannot extract λ_central for theoryCorr={tag!r}." - ) - with lz4.frame.open(path, "rb") as f: - return pickle.load(f) +# ============================================================================= +# Parsing the Nonperturbative section out of an upstream correction pkl +# (write side -- only the histmaker runs this, with the pkl already in hand). +# ============================================================================= def _find_nonperturbative(corr_dict): - """Search the upstream correction dict for ``Nonperturbative`` configs. + """Return [(basename, Nonperturbative dict)] for every basename in the pkl. - Returns a list of (basename, Nonperturbative dict) tuples. There are - typically multiple basenames (resummed-singular, fixed-order, etc.) — all - are expected to share the same Nonperturbative section since SCETlib runs - them with one runcard. + A correction pkl usually has several basenames (resummed-singular, + fixed-order, ...); they share one runcard so the Nonperturbative section is + the same, but we keep them all and pick below. """ out = [] meta = corr_dict.get("file_meta_data") if not isinstance(meta, dict): - raise KeyError( - "Correction pkl has no 'file_meta_data' entry — schema mismatch." - ) + raise KeyError("Correction pkl has no 'file_meta_data' entry.") for basename, file_meta in meta.items(): if not isinstance(file_meta, dict): continue @@ -72,20 +68,14 @@ def _find_nonperturbative(corr_dict): npert = cfg.get("Nonperturbative") if isinstance(npert, dict): out.append((basename, npert)) - if not out: - raise KeyError( - "No Nonperturbative section found in any basename of " - "file_meta_data — λ_central undefined." - ) return out def _parse_section(npert): - """Parse one Nonperturbative dict into the two parameter groups. + """Split one Nonperturbative dict into the eff / gnu parameter groups. - Numeric params absent from the runcard default to 0 — SCETlib runcards - only set the keys relevant to the chosen np_model. e.g. tanh_2 setups - typically omit ``lambda6`` (the bT⁵ coefficient only used by tanh_6). + Numeric params absent from the runcard default to 0 -- runcards only set the + keys their np_model uses (e.g. tanh_2 omits ``lambda6``). """ eff_params = {"np_model": npert[EFF_STRING[0]]} gnu_params = {"np_model_nu": npert[GNU_STRING[0]]} @@ -96,38 +86,144 @@ def _parse_section(npert): return eff_params, gnu_params -def read_lambda_central_from_meta(meta, proc="Z", _source=""): - """Same as :func:`read_lambda_central`, but takes the already-loaded - metadata dict (e.g. ``indata.metadata``) instead of an hdf5 path. +def _select_basename(sections, tag): + """Pick the basename whose runcard to use when a pkl bundles several. + + Some pkls carry multiple NP variants (e.g. a lattice central + a FranksVals + variant). Prefer a basename whose name shares a keyword with the tag, then + the resummed-singular file (it carries the full NP set). + """ + tag_lower = tag.lower() + KEYWORDS = ("franksvals", "lattice", "newvars", "lambda6") + matched_kw = next((k for k in KEYWORDS if k in tag_lower), None) + + def _score(name): + name_lower = name.lower() + score = 0 + if matched_kw and matched_kw in name_lower: + score += 10 + if ( + "nnlo_sing" in name_lower + or "_sing_" in name_lower + or name_lower.endswith("sing.pkl") + ): + score += 1 + return score - Avoids re-opening the input HDF5 when the caller already has the meta - in hand. ``_source`` is used only in error messages. + return sorted(sections, key=lambda item: -_score(item[0]))[0] + + +def extract_lambda_central(corr_dict, tag, proc): + """Parse the central lambda parameters out of a loaded correction pkl dict. + + Returns ``{tag, basename, eff_params, gnu_params}``. Raises if the pkl has + no Nonperturbative section. """ - try: - theory_corr = meta["meta_info_input"]["args"]["theoryCorr"] - except (KeyError, TypeError) as exc: + sections = _find_nonperturbative(corr_dict) + if not sections: raise KeyError( - f"{_source}: meta_info_input.args.theoryCorr missing — " - "cannot identify the central SCETlib correction." - ) from exc + f"No Nonperturbative section in correction pkl for tag={tag!r}, " + f"proc={proc!r}." + ) + basename, npert = _select_basename(sections, tag) + eff_params, gnu_params = _parse_section(npert) + return dict(tag=tag, basename=basename, eff_params=eff_params, gnu_params=gnu_params) + + +def _correction_pkl_path(tag, proc, data_dir=None): + data_dir = data_dir if data_dir is not None else wrem_common.data_dir + return os.path.join(data_dir, "TheoryCorrections", f"{tag}_Corr{proc}.pkl.lz4") + + +def build_lambda_central_meta(theory_corr_tags, procs=("Z", "W"), data_dir=None): + """Build the ``scetlib_np_lambda_central`` metadata for the histmaker output. - if not theory_corr: - raise ValueError(f"{_source}: theoryCorr list is empty.") - tag = theory_corr[0] # first entry = central; rest are pdfvars/pdfas + Opens the central correction pkl (``theory_corr_tags[0]``) for each proc and + extracts its Nonperturbative runcard. Returns ``{proc: lambda_central}`` for + the procs whose pkl exists and carries an NP section; procs without one are + skipped silently (most analyses have no SCETlib NP correction). Returns an + empty dict if there are no tags. - return _resolve_tag_to_lambda(tag, proc) + This is the ONLY place the upstream pkl is read; the fit reads the result + back from metadata. + """ + if not theory_corr_tags: + return {} + tag = theory_corr_tags[0] # first entry = central; rest are pdfvars/pdfas + out = {} + for proc in procs: + path = _correction_pkl_path(tag, proc, data_dir=data_dir) + if not os.path.exists(path): + continue + try: + with lz4.frame.open(path, "rb") as f: + corr_dict = pickle.load(f) + out[proc] = extract_lambda_central(corr_dict, tag, proc) + except KeyError: + # pkl present but no Nonperturbative section -- not an NP correction. + continue + return out + + +# ============================================================================= +# Reading the propagated metadata (read side -- fit / postprocessing). +# ============================================================================= + + +def _iter_meta_levels(meta, max_depth=8): + """Yield ``meta``, then ``meta['meta_info_input']``, recursively. + + The histmaker writes the key into ``meta_info``; rabbit nests that under + ``meta_info_input`` in the datacard, and again in the fitresults. Walking the + chain finds the key regardless of which file we were handed. + """ + cur = meta + for _ in range(max_depth): + if not isinstance(cur, dict): + return + yield cur + nxt = cur.get("meta_info_input") + if not isinstance(nxt, dict) or nxt is cur: + return + cur = nxt + + +def read_lambda_central_from_meta(meta, proc="Z", _source=""): + """Fetch the central lambda parameters from an already-loaded metadata dict. + + Searches ``meta`` and any nested ``meta_info_input`` for the propagated + ``scetlib_np_lambda_central`` entry. Raises with a clear message if it is + absent -- old inputs produced before metadata propagation must be remade + (resolving the upstream pkl by filename is no longer supported). + """ + for level in _iter_meta_levels(meta): + lc_all = level.get(META_KEY) + if not isinstance(lc_all, dict) or not lc_all: + continue + if proc in lc_all: + return dict(lc_all[proc]) + if len(lc_all) == 1: + # single proc stored -- use it whatever its label + return dict(next(iter(lc_all.values()))) + raise KeyError( + f"{_source}: {META_KEY!r} has no proc {proc!r} (have {sorted(lc_all)})." + ) + raise KeyError( + f"{_source}: no {META_KEY!r} in metadata. The SCETlib NP runcard is " + f"propagated into the histmaker output since this version; remake the " + f"histmaker output (upstream-pkl resolution by filename was removed)." + ) def load_lambda_central_file(path): - """Load a λ_central override from a JSON or YAML file. + """Load a lambda_central override from a JSON or YAML file. The file must decode to a dict with ``eff_params`` and ``gnu_params`` - sub-dicts (same shape as :func:`read_lambda_central`). Format is chosen - by extension (``.yaml``/``.yml`` → YAML, else JSON); YAML's loader also - accepts JSON, so this is forgiving either way. + sub-dicts. Format is chosen by extension (``.yaml``/``.yml`` -> YAML, else + JSON; YAML also accepts JSON). """ if not os.path.exists(path): - raise FileNotFoundError(f"λ_central file missing: {path!r}") + raise FileNotFoundError(f"lambda_central file missing: {path!r}") with open(path) as f: text = f.read() if path.lower().endswith((".yaml", ".yml")): @@ -139,7 +235,7 @@ def load_lambda_central_file(path): data = json.loads(text) except json.JSONDecodeError as exc: raise ValueError( - f"λ_central file {path!r} is not valid JSON; got {exc}" + f"lambda_central file {path!r} is not valid JSON; got {exc}" ) from exc if ( not isinstance(data, dict) @@ -147,50 +243,32 @@ def load_lambda_central_file(path): or "gnu_params" not in data ): raise ValueError( - f"λ_central file {path!r} must decode to a dict with " - f"'eff_params' and 'gnu_params' keys; got {type(data).__name__} " - f"with keys {list(data) if isinstance(data, dict) else ''}." + f"lambda_central file {path!r} must decode to a dict with " + f"'eff_params' and 'gnu_params' keys; got {type(data).__name__}." ) return data def read_lambda_central(hdf5_path, proc="Z"): - """Extract λ_central from the SCETlib correction referenced by the hdf5. - - Accepts either a fit-input datacard (setupRabbit output) OR a rabbit - ``fitresults*.hdf5``: rabbit propagates the whole datacard meta into the - fitresults as ``meta_info_input``, so for fitresults the lookup is the - same after unwrapping one level. Handy to answer "which λ_central was - this fit's ParamModel initialized with?" from the fit output alone. - - λ_central overrides: the override route is the ``lambda_central=`` - token in the ``--paramModel`` spec — the fit command (including the - token) is stored in the fitresults ``meta_info.args``, so this function - recovers the override automatically. A programmatic constructor-arg dict - is NOT visible in the fit output — for such fits the theoryCorr-tag - route below silently reports the auto-detect values (check the fit log - for "λ_central from" lines). - - Returns a dict with: - - tag : the central theoryCorr tag (first entry of meta args) - pkl_path : path to the upstream correction pkl - basename : the file_meta basename whose runcard was used - eff_params : dict for NP_model_effective (F_eff). Keys: - np_model, lambda_inf, lambda2, lambda4, lambda6, - delta_lambda2. - gnu_params : dict for NP_model_gammanu (γ_ν^NP). Keys: - np_model_nu, lambda_inf_nu, lambda2_nu, lambda4_nu. - - Raises with a clear message if any link in the chain is missing. + """Read the central lambda parameters referenced by an hdf5. + + Accepts a setupRabbit datacard or a rabbit ``fitresults*.hdf5`` (rabbit + nests the datacard meta one level down; the lookup handles both). + + Order of preference: + 1. a ``lambda_central=`` token in the stored ``--paramModel`` spec + (an explicit override the fit command recorded); + 2. the ``scetlib_np_lambda_central`` metadata propagated by the histmaker. + + Returns a dict with ``tag``, ``basename``, ``eff_params``, ``gnu_params`` + and ``source``. Raises with a clear message if neither route resolves. """ with h5py.File(hdf5_path, "r") as f: if "meta" not in f: - raise KeyError(f"{hdf5_path}: no 'meta' group — wrong file type?") + raise KeyError(f"{hdf5_path}: no 'meta' group -- wrong file type?") meta = wums_io.pickle_load_h5py(f["meta"]) - # Preference 1 (fitresults): a lambda_central= token in the stored - # --paramModel spec — the fit command records the override for free. + # Preference 1: an explicit lambda_central= override. args_meta = (meta.get("meta_info") or {}).get("args") or {} for spec in args_meta.get("paramModel") or []: for tok in spec: @@ -200,62 +278,12 @@ def read_lambda_central(hdf5_path, proc="Z"): lc["source"] = f"cli-file:{path}" return lc - # Fallback: resolve the datacard's theoryCorr tag. For fitresults this - # ASSUMES no env-var/constructor override was active (those are not - # visible in the output; check the fit log). fitresults nest the - # datacard meta (which itself carries meta_info_input) one level down — - # unwrap until the histmaker args are in view. - while ( - isinstance(meta.get("meta_info_input"), dict) - and "args" not in meta.get("meta_info_input", {}) - and "meta_info_input" in meta["meta_info_input"] - ): - meta = meta["meta_info_input"] + # Preference 2: the propagated histmaker metadata. lc = read_lambda_central_from_meta(meta, proc=proc, _source=hdf5_path) - lc["source"] = "theoryCorr-tag (assumes no runtime override)" + lc["source"] = "histmaker-metadata" return lc -def _resolve_tag_to_lambda(tag, proc): - """Internal: given a theoryCorr tag, load the pkl and parse λ_central.""" - - corr_dict = _load_correction_pkl(tag, proc) - sections = _find_nonperturbative(corr_dict) - - # Some correction pkls bundle multiple NP variants in one file (e.g. a - # "lattice" central + a "FranksVals" variant). We need to pick the - # basename that matches the analysis's NP tag. Heuristic: look for a - # substring in the basename that also appears in the tag (case-insensitive). - tag_lower = tag.lower() - KEYWORDS = ("franksvals", "lattice", "newvars", "lambda6") - matched_kw = next((k for k in KEYWORDS if k in tag_lower), None) - - def _basename_score(name): - name_lower = name.lower() - score = 0 - if matched_kw and matched_kw in name_lower: - score += 10 # strong preference: matches the analysis variant - if ( - "nnlo_sing" in name_lower - or "_sing_" in name_lower - or name_lower.endswith("sing.pkl") - ): - score += 1 # weak preference: resummed-singular carries the full NP set - return score - - sections_sorted = sorted(sections, key=lambda item: -_basename_score(item[0])) - basename, npert = sections_sorted[0] - eff_params, gnu_params = _parse_section(npert) - - return dict( - tag=tag, - pkl_path=_correction_pkl_path(tag, proc), - basename=basename, - eff_params=eff_params, - gnu_params=gnu_params, - ) - - if __name__ == "__main__": import sys @@ -267,7 +295,7 @@ def _basename_score(name): sys.exit(2) out = read_lambda_central(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else "Z") print(f"tag : {out['tag']}") - print(f"pkl_path : {out['pkl_path']}") print(f"basename : {out['basename']}") + print(f"source : {out.get('source')}") print(f"eff_params: {out['eff_params']}") print(f"gnu_params: {out['gnu_params']}") diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 3e6306ea0..82f870863 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -2,11 +2,9 @@ This ParamModel scales the signal reco template by a per-bin ratio of the SCETlib nonperturbative (NP) prediction at the fitted λ vs. at λ_central. The -prediction is built in THREE STEPS, written out in order below — read top to -bottom for the full maths. This module docstring is the SINGLE SOURCE OF TRUTH: -:mod:`response_matrix`, :func:`btgrid_tf.reconstruct_batch_tf`, -``sigma_reco_central.md``, and the development-tree numpy reference and -validation scripts all point here rather than restate it. +prediction is built in four steps, written out below; the related modules +(:mod:`response_matrix`, :func:`btgrid_tf.reconstruct_batch`, the validation +scripts, ``sigma_reco_central.md``) refer here for the derivation. Pipeline at a glance (everything is a function of the NP parameters λ): @@ -23,9 +21,9 @@ cosThetaStarll_quantile, phiStarll_quantile). λ splits into λ_eff (for F_eff) and λ_ν (for γ_ν^NP) — the 8 differentiable parameters listed at the end. -═════════════════════════════════════════════════════════════════════════════ +============================================================================= Step 1 — σ_resum(λ; g): the resummed prediction from the bT grid -═════════════════════════════════════════════════════════════════════════════ +============================================================================= (1a) Per-(Q, Y, qT) bT-space (Hankel) integral — the NP parameters enter HERE: @@ -79,9 +77,9 @@ Y axis folded into |Y| → absYVGen. The |Y| fold is valid because NP is Y-symmetric (F_eff depends on Y², γ_ν^NP doesn't depend on Y at all). -═════════════════════════════════════════════════════════════════════════════ +============================================================================= Step 2 — σ_gen(λ; g): add the NP-independent fixed-order nonsingular -═════════════════════════════════════════════════════════════════════════════ +============================================================================= σ_gen(λ; g) = σ_resum(λ; g) + σ_ns(g) @@ -100,9 +98,9 @@ for resum-only diagnostics subtract the exposed ``sigma_ns`` from ``sigma_gen_central`` (and re-fold with ``R`` for reco level). -═════════════════════════════════════════════════════════════════════════════ +============================================================================= Step 3 — σ_reco(λ; b): fold gen → reco through the response matrix -═════════════════════════════════════════════════════════════════════════════ +============================================================================= P(b | g) = R_raw(b, g) / N_gen(g) (efficiency × migration) σ_reco(λ; b) = Σ_g P(b | g) · σ_gen(λ; g) @@ -147,9 +145,9 @@ the same QCD/boson gen level (the postfsr variants in the file close ~1% worse — FSR mismatch). -═════════════════════════════════════════════════════════════════════════════ +============================================================================= Step 4 — rnorm(b, proc): the per-reco-bin variation handed to rabbit -═════════════════════════════════════════════════════════════════════════════ +============================================================================= ratio(b) = σ_reco(λ; b) / σ_reco(λ_central; b) rnorm(b, proc) = 1 + (ratio(b) − 1) · [proc is signal] (1 in every other proc) @@ -165,23 +163,23 @@ then σ_gen cancels in σ_reco(λ_c) = R_raw·1 and the λ_central closure can't test the integral.) -───────────────────────────────────────────────────────────────────────────── +----------------------------------------------------------------------------- Parameters and inputs -───────────────────────────────────────────────────────────────────────────── +----------------------------------------------------------------------------- The 8 v1 parameters λ (all factorisable through the current btgrid): γ_ν^NP (CS-side): lambda2_nu, lambda4_nu, lambda_inf_nu F_eff (TMD-effective): lambda2, lambda4, lambda6, delta_lambda2, lambda_inf -λ_central is read from the fit-tensor's meta_info_input via the upstream -SCETlib correction pkl (see :mod:`scetlib_lambda_central`). The np_model and -np_model_nu strings are fixed at construction (from λ_central). All λ values -are TF Variables — differentiable in the fit. +λ_central is read from the fit-tensor's metadata, where the histmaker stored +the SCETlib correction's NP runcard (see :mod:`scetlib_lambda_central`). The +np_model and np_model_nu strings are fixed at construction (from λ_central). +All λ values are TF Variables, differentiable in the fit. -═════════════════════════════════════════════════════════════════════════════ +============================================================================= Getting the postfit Hessian / covariance (uncertainties on λ) -═════════════════════════════════════════════════════════════════════════════ +============================================================================= The fit floats λ fine, but rabbit's postfit covariance step ``loss_val_grad_hess`` → ``t2.jacobian(grad, x)`` differentiates through the bT diff --git a/wremnants/production/histmaker_tools.py b/wremnants/production/histmaker_tools.py index 331fb305b..7a5d384c4 100644 --- a/wremnants/production/histmaker_tools.py +++ b/wremnants/production/histmaker_tools.py @@ -154,6 +154,31 @@ def analysis_debug_output(results): logger.debug("") +def _add_scetlib_np_lambda_central(meta_info, args): + """Propagate the central SCETlib NP runcard into the output metadata. + + The SCETlib ParamModel fit needs the central lambda parameters the + correction was generated with. We read them from the upstream theoryCorr pkl + here (where the correction is applied) and store them under + ``scetlib_np_lambda_central`` so the fit never re-opens the pkl by filename. + No-op for analyses whose correction has no Nonperturbative section. + """ + theory_corr = getattr(args, "theoryCorr", None) + if not theory_corr: + return + try: + from wremnants.postprocessing.scetlib_np import lambda_central + + lc_meta = lambda_central.build_lambda_central_meta(theory_corr) + if lc_meta: + meta_info[lambda_central.META_KEY] = lc_meta + logger.info( + f"Stored scetlib_np_lambda_central for procs {sorted(lc_meta)}" + ) + except Exception as exc: + logger.warning(f"Could not extract scetlib_np_lambda_central: {exc}") + + def write_analysis_output(results, outfile, args, name_append=[]): analysis_debug_output(results) @@ -203,11 +228,9 @@ def write_analysis_output(results, outfile, args, name_append=[]): ioutils.pickle_dump_h5py(k, v, f, override=open_as != "w") if "meta_info" not in f.keys(): - ioutils.pickle_dump_h5py( - "meta_info", - output_tools.make_meta_info_dict(args=args, wd=common.base_dir), - f, - ) + meta_info = output_tools.make_meta_info_dict(args=args, wd=common.base_dir) + _add_scetlib_np_lambda_central(meta_info, args) + ioutils.pickle_dump_h5py("meta_info", meta_info, f) logger.info(f"Writing output: {time.time()-time0}") logger.info(f"Output saved in {outfile}") diff --git a/wremnants/utilities/common.py b/wremnants/utilities/common.py index 50866349c..ddacc5123 100644 --- a/wremnants/utilities/common.py +++ b/wremnants/utilities/common.py @@ -16,6 +16,13 @@ BR_TAUToE = 0.1782 Z_TAU_TO_LEP_RATIO = 1.0 - (1.0 - BR_TAUToMU - BR_TAUToE) ** 2 +# Z mass/width in the s-dependent-width ("running width") scheme used for the +# gen Breit-Wigner lineshape and by SCETlib. These are NOT the PDG pole values +# (mZ = 91.1876, GammaZ = 2.4952); they are the values fed to make_bw_binning in +# the gen histmakers (w_z_gen_dists.py, mz_dilepton.py). +MZ_S_DEP_WIDTH = 91.1535 +GAMMAZ_S_DEP_WIDTH = 2.4932 + # cross sections in pb at sqrt(s)=13TeV (TODO: add source information) xsec_DYJetsToLL = 2001.9 xsec_WplusJetsToLNu = 11765.9 From 782d8110bc123bb5eb49b8bc777cd533a446b16d Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Mon, 22 Jun 2026 16:55:05 -0400 Subject: [PATCH 15/31] Apply black formatting (fix CI linting) Co-Authored-By: Claude Opus 4.8 (1M context) --- wremnants/postprocessing/scetlib_np/btgrid_integrate.py | 4 +++- wremnants/postprocessing/scetlib_np/lambda_central.py | 4 +++- wremnants/production/histmaker_tools.py | 4 +--- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py index 273d694d0..76894223c 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py @@ -108,7 +108,9 @@ def sparse_to_dense_tf(sigma_flat, flat_idx): # ============================================================================= -def q_integrate_weights(Q_grid, Q_lo, Q_hi, q0=MZ_S_DEP_WIDTH, Gamma=GAMMAZ_S_DEP_WIDTH): +def q_integrate_weights( + Q_grid, Q_lo, Q_hi, q0=MZ_S_DEP_WIDTH, Gamma=GAMMAZ_S_DEP_WIDTH +): """Simpson weights for integrating over Q ∈ [Q_lo, Q_hi] in arctan-Q² space. Implements the same change of variable as diff --git a/wremnants/postprocessing/scetlib_np/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py index abba79d80..41b15779c 100644 --- a/wremnants/postprocessing/scetlib_np/lambda_central.py +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -127,7 +127,9 @@ def extract_lambda_central(corr_dict, tag, proc): ) basename, npert = _select_basename(sections, tag) eff_params, gnu_params = _parse_section(npert) - return dict(tag=tag, basename=basename, eff_params=eff_params, gnu_params=gnu_params) + return dict( + tag=tag, basename=basename, eff_params=eff_params, gnu_params=gnu_params + ) def _correction_pkl_path(tag, proc, data_dir=None): diff --git a/wremnants/production/histmaker_tools.py b/wremnants/production/histmaker_tools.py index 7a5d384c4..36079c904 100644 --- a/wremnants/production/histmaker_tools.py +++ b/wremnants/production/histmaker_tools.py @@ -172,9 +172,7 @@ def _add_scetlib_np_lambda_central(meta_info, args): lc_meta = lambda_central.build_lambda_central_meta(theory_corr) if lc_meta: meta_info[lambda_central.META_KEY] = lc_meta - logger.info( - f"Stored scetlib_np_lambda_central for procs {sorted(lc_meta)}" - ) + logger.info(f"Stored scetlib_np_lambda_central for procs {sorted(lc_meta)}") except Exception as exc: logger.warning(f"Could not extract scetlib_np_lambda_central: {exc}") From 7201d52984c16da8e35aa8194051b346f1368fa5 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Tue, 23 Jun 2026 09:53:13 -0400 Subject: [PATCH 16/31] SCETlib-NP: read the response matrix R from the datacard Carry the reco x gen response matrix R (and the gen-total N_gen) through the setupRabbit datacard instead of a separate file, so the fit is self-contained and R is always consistent with the run that built the card (review comment #6). - setupRabbit: after the input loop, extract R + N_gen once from the unfolding histmaker output and embed them in the fit input via rabbit add_auxiliary (the "scetlib_np" group). Presence-based lenient guard (response_matrix.has_response): embed only when an input carries both the response hist and the prefsr gen-total; raise if more than one qualifies. - response_matrix: add has_response() (cheap both-present check, never raises). - param_model: read R only from indata.auxiliary["scetlib_np"]; drop the unfolding_hdf5_path argument and the sigma_gen(lambda_c) N_gen fallback (raise if N_gen is absent -- the proxy made the central closure circular). - scetlib_np package __init__: import SCETlibNPParamModel lazily (PEP 562) so response_matrix can be imported by setupRabbit without pulling in TensorFlow. Requires a rabbit with TensorWriter.add_auxiliary (WMass/rabbit#145); the submodule bump is a follow-up. Co-Authored-By: Claude Opus 4.8 --- scripts/rabbit/setupRabbit.py | 40 +++++++ .../postprocessing/scetlib_np/__init__.py | 26 ++-- .../postprocessing/scetlib_np/param_model.py | 111 ++++++++++-------- .../scetlib_np/response_matrix.py | 26 ++++ 4 files changed, 148 insertions(+), 55 deletions(-) diff --git a/scripts/rabbit/setupRabbit.py b/scripts/rabbit/setupRabbit.py index 88a3fc3f5..6c199eafe 100644 --- a/scripts/rabbit/setupRabbit.py +++ b/scripts/rabbit/setupRabbit.py @@ -18,6 +18,7 @@ from wremnants.postprocessing.datagroups.datagroups import Datagroups from wremnants.postprocessing.histselections import FakeSelectorSimpleABCD from wremnants.postprocessing.regression import Regressor +from wremnants.postprocessing.scetlib_np import response_matrix as scetlib_np_response from wremnants.postprocessing.syst_tools import ( fake_nonclosure_byAxis, fake_transferFactor_ptSyst, @@ -3585,6 +3586,45 @@ def outputFolderName(outfolder, datagroups, doStatOnly, postfix): outfile = "Combination" logger.info(f"Writing output to {outfile}") + # ---- SCETlib-NP response matrix R: embed it in the datacard so the + # SCETlibNPParamModel reads R (and the gen-total N_gen) from the fit input, + # consistent with the run that produced the card, rather than from a + # separate, independently-versioned file. Presence-based *lenient* guard + # (see response_matrix.has_response): embed only when an input carries BOTH + # the response hist and the gen-total, so generic unfolding runs that lack + # the gen-total are a no-op. A genuine NP card missing the gen-total simply + # won't embed and the SCETlibNPParamModel will then error clearly at fit + # time. One source, one path: the ParamModel reads R only from the datacard. + resp_inputs = [f for f in args.inputFile if scetlib_np_response.has_response(f)] + if len(resp_inputs) > 1: + raise RuntimeError( + "Multiple inputs carry the SCETlib-NP response (hist + gen-total): " + f"{resp_inputs}; expected at most one (the Z dilepton --unfolding run)." + ) + if resp_inputs: + if not hasattr(writer, "add_auxiliary"): + raise RuntimeError( + "TensorWriter has no 'add_auxiliary'; update the rabbit submodule " + "to a revision including WMass/rabbit#145 to embed the SCETlib-NP " + "response matrix in the datacard." + ) + logger.info(f"Embedding SCETlib-NP response matrix from {resp_inputs[0]}") + R_info = scetlib_np_response.load_R(resp_inputs[0]) + writer.add_auxiliary( + "scetlib_np", + { + "R": R_info["R"], + "N_gen": R_info["N_gen"], + "reco_axes": [n for n, _ in R_info["reco_axes"]], + "gen_axes": [n for n, _ in R_info["gen_axes"]], + # one edges dataset per reco/gen axis (variable length) + **{ + f"edges__{n}": e + for n, e in R_info["reco_axes"] + R_info["gen_axes"] + }, + }, + ) + # propagate meta info into result file meta = { "meta_info": output_tools.make_meta_info_dict( diff --git a/wremnants/postprocessing/scetlib_np/__init__.py b/wremnants/postprocessing/scetlib_np/__init__.py index 439644204..b12044044 100644 --- a/wremnants/postprocessing/scetlib_np/__init__.py +++ b/wremnants/postprocessing/scetlib_np/__init__.py @@ -1,13 +1,21 @@ -"""SCETlib NP continuous-λ param model and its bT-grid factorisation port. +"""SCETlib-NP postprocessing package. -The fit-time rabbit ParamModel lives in :mod:`.param_model`; the supporting -bT-grid numpy/TF code, Q-integration, caching, response-matrix loader and -λ_central reader are the sibling modules. ``SCETlibNPParamModel`` is -re-exported here so it can be referenced by the short dotted path: - - --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel +``SCETlibNPParamModel`` (and its TensorFlow / btgrid dependencies) is imported +lazily so that lightweight submodules — e.g. :mod:`response_matrix`, used by +setupRabbit to embed the response matrix in the datacard — can be imported +without pulling in TensorFlow. The package-level re-export +``wremnants.postprocessing.scetlib_np.SCETlibNPParamModel`` (used by rabbit's +``--paramModel`` loader) still works, resolved on first access via PEP 562. """ -from wremnants.postprocessing.scetlib_np.param_model import SCETlibNPParamModel - __all__ = ["SCETlibNPParamModel"] + + +def __getattr__(name): + if name == "SCETlibNPParamModel": + from wremnants.postprocessing.scetlib_np.param_model import ( + SCETlibNPParamModel, + ) + + return SCETlibNPParamModel + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 82f870863..64cab9859 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -4,7 +4,7 @@ SCETlib nonperturbative (NP) prediction at the fitted λ vs. at λ_central. The prediction is built in four steps, written out below; the related modules (:mod:`response_matrix`, :func:`btgrid_tf.reconstruct_batch`, the validation -scripts, ``sigma_reco_central.md``) refer here for the derivation. +scripts) refer here for the derivation. Pipeline at a glance (everything is a function of the NP parameters λ): @@ -64,8 +64,7 @@ (I_pert, C_ν) rows + J₀ kernel on the unique-qT grid + Simpson-as-matmul), which is numerically equivalent to the per-bin (Nbins, Nbt) layout (≲1e-14 rel., floating-point summation order only) but ~6× smaller — required to fit -a 32 GB GPU. See FACTORIZED_RECON.md in this directory; the legacy_recon=1 -spec token restores the legacy layout. +a 32 GB GPU. The legacy_recon=1 spec token restores the legacy layout. (1b) Integrate over Q, then rebin onto the gen grid: @@ -159,9 +158,6 @@ purely the SHAPE of the NP variation per reco bin — Steps 1–3 build the absolute σ_reco(λ; b), and Step 4 reduces it to the bin-by-bin template scaling the fit needs. σ_reco(λ_central) is precomputed once at construction as the denominator. -(If the gen-total hist is absent, σ_gen(λ_c) is used as a proxy for N_gen, but -then σ_gen cancels in σ_reco(λ_c) = R_raw·1 and the λ_central closure can't test -the integral.) ----------------------------------------------------------------------------- Parameters and inputs @@ -198,8 +194,7 @@ rabbit's jacobian gets the exact derivatives while the big slab stays inside stop_gradient and never enters the differentiated graph (33 TB → a few MB). Implemented in ``_ratio_straightthrough`` (+ ``_ratio_compact_jac``, -``_ratio_compact_hess``); selected in ``compute()`` by env flags. Full -derivation + validation in ``HESSIAN_PLAN.md``. +``_ratio_compact_hess``); selected in ``compute()`` by env flags. Two-pass recipe (rabbit still computes the Hessian; NO rabbit changes): @@ -224,8 +219,8 @@ WARNING: hessian_straightthrough=1 WITHOUT hessian_gn=1 is full-K mode — correct in principle (needed for real/toy data) but currently INFEASIBLE at full grid scale (the 64 nested-FA passes unroll into one graph, ~TB peak → -OOM-kill). Until HESSIAN_PLAN.md §9 (precomputed chunked K) is implemented, -always pass hessian_gn=1 (exact for Asimov). +OOM-kill). Until a precomputed chunked-K path is implemented, always pass +hessian_gn=1 (exact for Asimov). GN vs full-K. The Poisson Hessian is H_ij = Σ_b [ (n_b/μ_b²) J_bi J_bj + (1 − n_b/μ_b) K_bij ]. @@ -240,7 +235,7 @@ tangent into ``Equal`` without changing any value or derivative (the comparison is a measure-zero boundary ``tf.where`` never differentiates). Full-K now runs under @tf.function and matches the exact reverse-mode Hessian to machine -precision (≤3e-16 rel; see HESSIAN_PLAN.md §7 + the isolation validation). So +precision (≤3e-16 rel; verified by the isolation validation). So both GN and full-K are available; GN remains the default for Asimov (exact and cheaper — 8 vs 72 fold passes). @@ -262,7 +257,6 @@ from wremnants.postprocessing.scetlib_np import btgrid_integrate as fz_int from wremnants.postprocessing.scetlib_np import btgrid_tf as fz_tf from wremnants.postprocessing.scetlib_np import lambda_central as scetlib_lambda_central -from wremnants.postprocessing.scetlib_np import response_matrix as fz_R from wremnants.utilities import common as wrem_common from wremnants.utilities.data_paths import getDataPath @@ -276,12 +270,6 @@ "TheoryCorrections", "results_z-2d-nnlo-vj-CT18ZNNLO-{scale}-scetlibmatch.txt", ) -_UNFOLDING_HDF5_DEFAULT = os.path.join( - wrem_common.data_dir, - "TheoryCorrections", - "scetlib_np", - "mz_dilepton_unfolding_R_skim.hdf5", -) _BTGRID_SUBDIR = ("scetlib_np", "Z_COM13_CT18Z_N3p0LL_btgrid_fineall") _DISCRETE_NP_SUBSTRING = "scetlibnp" @@ -358,6 +346,43 @@ def _crop_R_to_fit(R, R_reco_axes, fit_reco_axes, tol=1e-9): return R[slices] +def _R_info_from_auxiliary(indata): + """Reconstruct the response-matrix dict from the datacard's ``scetlib_np`` + auxiliary bundle. + + setupRabbit extracts R (and the gen-total N_gen, reco/gen axis names + edges) + once from the unfolding histmaker output and embeds it in the fit input via + rabbit's ``add_auxiliary``; rabbit exposes it as ``FitInputData.auxiliary``. + The model reads R ONLY from there — one source, one path — so R is always + consistent with the run that produced the datacard. The returned dict matches + :func:`response_matrix.load_R`'s shape for the keys this model consumes + (``R``, ``N_gen``, and ``reco_axes`` / ``gen_axes`` as ordered + ``(name, edges)`` lists). + """ + aux = getattr(indata, "auxiliary", None) or {} + if "scetlib_np" not in aux: + raise ValueError( + "SCETlibNPParamModel: the datacard has no 'scetlib_np' auxiliary (the " + "reco×gen response matrix R). Rebuild the datacard with a setupRabbit " + "that embeds it from a mz_dilepton --unfolding input (it must carry " + "'nominal_prefsr_yieldsUnfolding' and the 'prefsr' gen-total)." + ) + bundle = aux["scetlib_np"] + 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"] + ], + ) + + def _bin_sum_matrix(src_centers, target_edges, tol=1e-6): """(N_target, N_src) 0/1 matrix that SUMS bin-integrated source bins whose centre falls in each target bin. Source bins outside all target bins are @@ -477,7 +502,6 @@ def parse_args(cls, indata, *args, **kwargs): def __init__( self, indata, - unfolding_hdf5_path: Optional[str] = None, btgrid_dir: Optional[str] = None, lambda_central=None, signal_proc: str = "Zmumu", @@ -503,13 +527,13 @@ def __init__( Parameters ---------- indata - rabbit's input-data structure (passed by ``ph.load_models``). - unfolding_hdf5_path - Path to the upstream histmaker output containing - ``nominal_prefsr_yieldsUnfolding`` for R (and the ``prefsr`` xnorm - hist for N_gen) — see :mod:`response_matrix` for the defaults. - Defaults (when None) to the skim shipped in wremnants-data - (``_UNFOLDING_HDF5_DEFAULT``); pass explicitly to override. + rabbit's input-data structure (passed by ``ph.load_models``). The + reco×gen response matrix R (and the gen-total N_gen, axis names + + edges) is read from ``indata.auxiliary["scetlib_np"]`` — embedded in + the datacard by setupRabbit from a ``mz_dilepton.py --unfolding`` + histmaker output (see :func:`_R_info_from_auxiliary` and + :mod:`response_matrix`). One source, one path: there is no file-path + argument; R always comes from the fit input it is consistent with. btgrid_dir Directory holding the SCETlib bT-grid ``combined_btgrid.pkl``. Defaults (when None) to the shared data-area copy next to NanoAOD @@ -571,7 +595,7 @@ def __init__( legacy_recon Use the legacy per-bin (Nbins, Nbt) reconstruction layout instead of the default memory-factorized one (numerically equivalent to - ≲1e-14 rel; for parity checks only). See FACTORIZED_RECON.md. + ≲1e-14 rel; for parity checks only). xparam_default Comma-separated ``name=value,...`` string shifting the fit START (and the prior mean) off the runcard's λ_central — for closure / @@ -586,8 +610,6 @@ def __init__( """ self.indata = indata - if unfolding_hdf5_path is None: - unfolding_hdf5_path = _UNFOLDING_HDF5_DEFAULT if btgrid_dir is None: btgrid_dir = _default_btgrid_dir() @@ -737,8 +759,8 @@ def __init__( dtype=fz_tf.DTYPE, ) - # ---- R matrix - R_info = fz_R.load_R(unfolding_hdf5_path) + # ---- R matrix (read from the datacard's scetlib_np auxiliary) + R_info = _R_info_from_auxiliary(indata) # The fit-tensor's reco binning may differ from R's by trailing # overflow bins (e.g. R has ptll [0, …, 44, 100] while the fit ends # at 44). Crop R's trailing bins so the reco shape matches. @@ -754,19 +776,16 @@ def __init__( # Gen-total denominator N_gen(g) from the xnorm hist ("prefsr"): the # generated fiducial yield per gen bin (pre-reco-selection). Dividing R # by this gives the theory-independent efficiency×migration response. - # Falls back to the σ_gen(λ_c) proxy if the gen-total isn't in the file. - if R_info.get("N_gen") is not None: - self._N_gen_flat = tf.constant( - R_info["N_gen"].reshape(-1), dtype=fz_tf.DTYPE - ) - else: - self._N_gen_flat = None - print( - "[SCETlibNPParamModel] WARNING: no gen-total hist in unfolding " - "output; falling back to σ_gen(λ_c) as the response normalizer " - "(circular closure — see param_model docstring).", - flush=True, + # REQUIRED: setupRabbit only embeds the response when the gen-total is + # present, so N_gen should always be here; raise if not (a σ_gen(λ_c) + # proxy would make the central closure circular — see module docstring). + if R_info.get("N_gen") is None: + raise ValueError( + "SCETlibNPParamModel: the 'scetlib_np' auxiliary has no N_gen " + "(gen-total). Rebuild the datacard from a histmaker output that " + "carries the 'prefsr' xnorm hist." ) + self._N_gen_flat = tf.constant(R_info["N_gen"].reshape(-1), dtype=fz_tf.DTYPE) self._reco_axes_meta = [ (name, fit_axes[1]) for (name, fit_axes) in zip( @@ -884,7 +903,7 @@ def __init__( f"SCETlibNPParamModel: {n_bad} gen bins have non-positive " f"σ_gen(λ_central); cannot normalize / fold the response." ) - N_gen = self._N_gen_flat if self._N_gen_flat is not None else gen_flat + N_gen = self._N_gen_flat # Guard empty gen bins (no generated events): leave column at 0. safe_N_gen = tf.where(N_gen > 0, N_gen, tf.ones_like(N_gen)) self.R = self._R_raw / safe_N_gen[tf.newaxis, :] # P(b|g) = R_raw / N_gen @@ -1089,7 +1108,7 @@ def _sigma_YqT_native_at(self, eff_params, gnu_params): SCETlib spectrum reference (curve 1) and the numpy `factorize` (curve 2).""" # 1. Reconstruct σ on the btgrid's flat (Nbins,) layout. Factorized # (default) and legacy layouts are numerically equivalent (≲1e-14 - # rel.; summation order only — see FACTORIZED_RECON.md). + # rel.; summation order only). eff = {k: v for k, v in eff_params.items() if k != "np_model"} gnu = {k: v for k, v in gnu_params.items() if k != "np_model_nu"} if self.factorized: @@ -1214,7 +1233,7 @@ def _ratio_compact_jac(self, param): One JVP per parameter (nparam ≤ 8), each a single bT-fold pass — NOT tiled over params. This is the compact object the Hessian actually needs - from the fold (see HESSIAN_PLAN.md §2).""" + from the fold.""" n = int(param.shape[0]) cols = [] for i in range(n): diff --git a/wremnants/postprocessing/scetlib_np/response_matrix.py b/wremnants/postprocessing/scetlib_np/response_matrix.py index c80a43226..d2983abc3 100644 --- a/wremnants/postprocessing/scetlib_np/response_matrix.py +++ b/wremnants/postprocessing/scetlib_np/response_matrix.py @@ -95,6 +95,32 @@ def _append_axis_overflow(h, axis_name): return np.concatenate([inr, over], axis=pos) +def has_response( + unfolding_hdf5_path, + sample_key=DEFAULT_SAMPLE, + hist_name=DEFAULT_HIST, + gen_total_name=DEFAULT_GENTOTAL, +): + """Cheap guard: does this histmaker output carry BOTH the reco x gen + response hist and the gen-total xnorm hist needed to build R *and* N_gen? + + Used by setupRabbit to decide whether to embed the SCETlib-NP response in + the datacard (presence-based, *lenient* guard): returns True only when both + are present, so a generic unfolding run that has the response hist but not + the gen-total is a silent no-op rather than an error. Never raises (any + structural problem -> False); does not materialize any histogram. + """ + try: + with h5py.File(unfolding_hdf5_path, "r") as f: + if sample_key not in f: + return False + sample = wums_io.pickle_load_h5py(f[sample_key]) + output = sample["output"] + return hist_name in output and gen_total_name in output + except (OSError, KeyError, TypeError): + return False + + def load_R( unfolding_hdf5_path, sample_key=DEFAULT_SAMPLE, From 072693154484824524cb0a60227072f43679c1bb Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Tue, 23 Jun 2026 11:16:37 -0400 Subject: [PATCH 17/31] clean up docstrings --- .../postprocessing/scetlib_np/btgrid_cache.py | 6 +-- .../scetlib_np/btgrid_integrate.py | 2 - .../postprocessing/scetlib_np/btgrid_tf.py | 37 ++++++------- .../scetlib_np/lambda_central.py | 53 +++++++++---------- 4 files changed, 45 insertions(+), 53 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_cache.py b/wremnants/postprocessing/scetlib_np/btgrid_cache.py index c3eea7faa..e220e0235 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_cache.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_cache.py @@ -1,8 +1,8 @@ """One-shot pickle cache for the combined SCETlib bT-grid. -Loading the fineall btgrid as 1519 individual shards takes ~110s. After the -first call, this module writes a single ``combined_btgrid.pkl`` in the btgrid -directory; subsequent calls load that in a few seconds. +Assembling the bT-grid from its individual shards is slow. On the first call +this module writes a single ``combined_btgrid.pkl`` in the btgrid directory; +subsequent calls load that combined file directly, which is much faster. Usage: from wremnants.postprocessing.scetlib_np import btgrid_cache diff --git a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py index 76894223c..73d325280 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py @@ -17,8 +17,6 @@ 3. :func:`rebin_weights` — given a fine source grid and a coarser target-bin edge list, builds a ``(N_target, N_source)`` Simpson weight matrix. Apply via ``tf.tensordot``. - -Parity-tested against the numpy reference implementation (development tree). """ import numpy as np diff --git a/wremnants/postprocessing/scetlib_np/btgrid_tf.py b/wremnants/postprocessing/scetlib_np/btgrid_tf.py index 35917bfb7..1a3ff88fb 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_tf.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_tf.py @@ -1,16 +1,14 @@ -"""TensorFlow port of the bT-grid factorization library. +"""TensorFlow bT-grid factorization library. -Function-by-function TF port of a numpy reference implementation (itself a -byte-for-byte transcription of SCETlib C++). The numpy reference and the -parity tests that keep the two in sync live in the development tree; only -the TF port is shipped here. +Differentiable TF implementation of the SCETlib bT-space form factors and Hankel +reconstruction (transcribed from the SCETlib C++). Design choices: * ``np_model`` / ``np_model_nu`` strings are fixed at trace time (the SCETlib runcard sets them once per fit). The TF functions dispatch on the string at Python level — no ``tf.cond``. * λ parameters are TF tensors (typically scalars, but broadcasting follows - the same rules as numpy). + the usual array rules). * All ops are differentiable in λ. Branches on λ values use ``tf.where`` with a safe denominator to avoid NaN gradients. * ``b_star_global`` is not ported — the cached ``b_bar`` array in the bT-grid @@ -24,8 +22,7 @@ import numpy as np import tensorflow as tf -# Set the dtype used for all ops in this module. Match the numpy reference -# (which uses ``float`` ≡ float64) to keep parity tight. +# Set the dtype used for all ops in this module (float64 throughout). DTYPE = tf.float64 @@ -49,9 +46,8 @@ def _as_dtype(x, dtype=DTYPE): def simpson_weights(x): """Return weights ``w`` such that Simpson(y, x) == sum(w * y, axis=-1). - ``x`` is a numpy array with size ``N``. Implementation mirrors the numpy - ``simpson`` of the numpy reference (composite Simpson with - trapezoid fallback on the last segment when N-1 is odd). + ``x`` is a numpy array with size ``N``. Composite Simpson with a trapezoid + fallback on the last segment when N-1 is odd. """ x = np.asarray(x, dtype=np.float64) n_intervals = x.size - 1 @@ -147,7 +143,7 @@ def _safe_div(num, den): def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_model): - """TF port of the numpy-reference ``F_eff`` for a fixed ``np_model``.""" + """TMD-effective NP form factor F_eff(Y, bT) for a fixed ``np_model``.""" if np_model not in EFF_MODELS: raise ValueError(f"F_eff_tf: unsupported np_model {np_model!r}") @@ -169,8 +165,8 @@ def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_ if np_model == "identity": return tf.exp(-2.0 * bT * arg) - # lambda_inf == 0 returns ones (matches numpy short-circuit). We compute - # the full formula with a safe denominator and mask at the end. + # lambda_inf == 0 returns ones. We compute the full formula with a safe + # denominator and mask at the end. arg_inf = _safe_div(arg, lambda_inf) model = {"hyp_tangent": "tanh_2", "square_root": "frac_2"}.get(np_model, np_model) @@ -203,7 +199,7 @@ def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_ def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): - """TF port of the numpy-reference ``gamma_nu_NP`` for fixed ``np_model_nu``.""" + """CS-side NP rapidity anomalous dimension γ_ν^NP(bT) for fixed ``np_model_nu``.""" if np_model_nu not in GNU_MODELS: raise ValueError(f"gamma_nu_NP_tf: unsupported np_model_nu {np_model_nu!r}") @@ -267,12 +263,12 @@ def reconstruct_batch_tf( Y_unique=None, Y_inverse_idx=None, ): - """TF port of the numpy-reference ``reconstruct_batch``. + """Reconstruct σ on a batch of (Q, Y, qT) grid points from the bT integrand. Implements the per-(Q, Y, qT) bT-space integrand. The full formula with every factor and its bare-bT / b*(bT) / (Q,Y,qT) / λ dependence is written - out once in the :mod:`param_model` module docstring (single source of - truth) — consult it rather than re-deriving the factors from this code. + out in the :mod:`param_model` module docstring — consult it rather than + re-deriving the factors from this code. All array-shape arguments are TF tensors or numpy arrays (will be cast). The λ values inside ``eff_params`` / ``gnu_params`` are the differentiable @@ -360,7 +356,7 @@ def build_bT_J0_kernel(qT_per_bin, bT): # # The (Nbins, Nbt) layout of reconstruct_batch_tf needs several ~9 GB fp64 # tensors (Nbins=546840, Nbt=2000) and OOMs a 32 GB GPU at construction. Two -# exact observations shrink it (see scetlib_np/FACTORIZED_RECON.md): +# exact observations shrink it: # # 1. qT enters the λ-dependent integrand ONLY through the bT·J0(qT·bT) # kernel: the kernel needs the NqT *unique* qT values, not Nbins rows. @@ -524,8 +520,7 @@ def reconstruct_batch_factorized_tf( weighted J0 kernel on the unique-qT grid, and a per-bin gather — no (Nbins, Nbt) tensor is ever materialized. Same integrand, weights and sampling as :func:`reconstruct_batch_tf`; only the floating-point - multiplication grouping and summation order differ (≲1e-14 relative — - see the parity script :mod:`scetlib_np_factorized_parity`). + multiplication grouping and summation order differ (≲1e-14 relative). Parameters ---------- diff --git a/wremnants/postprocessing/scetlib_np/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py index 41b15779c..61eb83e1f 100644 --- a/wremnants/postprocessing/scetlib_np/lambda_central.py +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -51,9 +51,11 @@ def _find_nonperturbative(corr_dict): """Return [(basename, Nonperturbative dict)] for every basename in the pkl. - A correction pkl usually has several basenames (resummed-singular, - fixed-order, ...); they share one runcard so the Nonperturbative section is - the same, but we keep them all and pick below. + A correction pkl carries several basenames (the resummed SCETlib file, the + fixed-order singular file, the gen hist, ...). The resummed and singular + files can have DIFFERENT Nonperturbative runcards (e.g. the FranksVals + correction), so we keep all NP-bearing basenames and let + :func:`_select_resummed` pick the central one. """ out = [] meta = corr_dict.get("file_meta_data") @@ -86,31 +88,28 @@ def _parse_section(npert): return eff_params, gnu_params -def _select_basename(sections, tag): - """Pick the basename whose runcard to use when a pkl bundles several. +def _select_resummed(sections): + """Pick the resummed prediction's runcard from the NP-bearing basenames. - Some pkls carry multiple NP variants (e.g. a lattice central + a FranksVals - variant). Prefer a basename whose name shares a keyword with the tag, then - the resummed-singular file (it carries the full NP set). + A scetlib_dyturbo correction is built (in ``make_theory_corr.py``) from a + resummed SCETlib file plus a fixed-order *singular* file that is subtracted + in the matching. Only the resummed file's Nonperturbative runcard is the + central NP; the singular file's is not (it can differ, e.g. FranksVals). + ``make_theory_corr.py`` tells the two apart by the ``"sing"`` substring in + the filename, so we do the same: keep the basename without ``"sing"``. """ - tag_lower = tag.lower() - KEYWORDS = ("franksvals", "lattice", "newvars", "lambda6") - matched_kw = next((k for k in KEYWORDS if k in tag_lower), None) - - def _score(name): - name_lower = name.lower() - score = 0 - if matched_kw and matched_kw in name_lower: - score += 10 - if ( - "nnlo_sing" in name_lower - or "_sing_" in name_lower - or name_lower.endswith("sing.pkl") - ): - score += 1 - return score - - return sorted(sections, key=lambda item: -_score(item[0]))[0] + resummed = [item for item in sections if "sing" not in item[0]] + if len(resummed) == 1: + return resummed[0] + if not resummed: + raise KeyError( + "No resummed (non-'sing') basename carries a Nonperturbative " + f"section; basenames seen: {[bn for bn, _ in sections]}." + ) + raise KeyError( + "Multiple resummed basenames carry a Nonperturbative section " + f"({[bn for bn, _ in resummed]}); cannot pick the central runcard." + ) def extract_lambda_central(corr_dict, tag, proc): @@ -125,7 +124,7 @@ def extract_lambda_central(corr_dict, tag, proc): f"No Nonperturbative section in correction pkl for tag={tag!r}, " f"proc={proc!r}." ) - basename, npert = _select_basename(sections, tag) + basename, npert = _select_resummed(sections) eff_params, gnu_params = _parse_section(npert) return dict( tag=tag, basename=basename, eff_params=eff_params, gnu_params=gnu_params From 29ee8e217e459dafffa90fbdd519d1d547ab2f34 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Tue, 23 Jun 2026 11:17:20 -0400 Subject: [PATCH 18/31] clean up imports --- wremnants/production/datasets/dataset_tools.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/wremnants/production/datasets/dataset_tools.py b/wremnants/production/datasets/dataset_tools.py index d5c2275eb..a75982f36 100644 --- a/wremnants/production/datasets/dataset_tools.py +++ b/wremnants/production/datasets/dataset_tools.py @@ -10,15 +10,8 @@ # Path / file-list helpers live in a ROOT/narf-free module so they can be # imported without pulling in ROOT (this package's __init__ imports ROOT and -# narf). Re-exported here for backward compatibility. -from wremnants.utilities.data_paths import ( # noqa: F401 - appendFilesXrd, - buildFileList, - buildFileListPosix, - buildFileListXrd, - getDataPath, - makeFilelist, -) +# narf). +from wremnants.utilities.data_paths import getDataPath, makeFilelist from wums import logging logger = logging.child_logger(__name__) From ae0c210a28db63e1219032375dfde02f2447659e Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Tue, 23 Jun 2026 13:07:35 -0400 Subject: [PATCH 19/31] fix scales naming --- wremnants/utilities/io_tools/input_tools.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/wremnants/utilities/io_tools/input_tools.py b/wremnants/utilities/io_tools/input_tools.py index 1f6e0e64d..e45d467a2 100644 --- a/wremnants/utilities/io_tools/input_tools.py +++ b/wremnants/utilities/io_tools/input_tools.py @@ -215,7 +215,9 @@ def read_dyturbo_vars_hist(base_name, var_axis=None, axes=("Y", "qT"), charge=No f"Scale variation {var} found for fo_sing piece but no corresponding variation for dyturbo" ) dyturbo_scale = scales_map.get(var, "mur1-muf1") + print(var, dyturbo_scale) dyturbo_name = base_name.format(i=pdf_member, scale=dyturbo_scale) + print(dyturbo_name) h = read_dyturbo_hist([dyturbo_name], axes=axes, charge=charge) if not var_hist: var_hist = hist.Hist(*h.axes, var_axis, storage=h.storage_type()) @@ -333,10 +335,11 @@ def read_dyturbo_hist( hists = [] for fn in filenames: - if "-mur0p5-" in fn.split("/")[-1]: - fn = fn.replace("-mur0p5-", "-murH-") - if "-muf0p5-" in fn.split("/")[-1]: - fn = fn.replace("-muf0p5-", "-mufH-") + # TODO the naming convention is unclear and inconsistent. These may be needed in the future. + # if "-mur0p5-" in fn.split("/")[-1]: + # fn = fn.replace("-mur0p5-", "-murH-") + # if "-muf0p5-" in fn.split("/")[-1]: + # fn = fn.replace("-muf0p5-", "-mufH-") expandedf = fn.split("+") From 563a5b2bc8a9ce22fc6c8704f793a573575f813b Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Tue, 23 Jun 2026 13:22:03 -0400 Subject: [PATCH 20/31] storing response matrix is optional --- scripts/rabbit/setupRabbit.py | 52 +++++++++++++++++------------------ 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/scripts/rabbit/setupRabbit.py b/scripts/rabbit/setupRabbit.py index 6c199eafe..71d2b2807 100644 --- a/scripts/rabbit/setupRabbit.py +++ b/scripts/rabbit/setupRabbit.py @@ -641,6 +641,11 @@ def make_parser(parser=None, argv=None): action="store_true", help="Add custom recoil systematic uncertainties from smearing met pt/phi and scaling met pt", ) + parser.add_argument( + "--storeResponseMatrix", + action="store_true", + help="Store response matrix for SCETlib-NP parameter model", + ) parser.add_argument( "--ABCDedgesByAxis", @@ -3595,35 +3600,30 @@ def outputFolderName(outfolder, datagroups, doStatOnly, postfix): # the gen-total are a no-op. A genuine NP card missing the gen-total simply # won't embed and the SCETlibNPParamModel will then error clearly at fit # time. One source, one path: the ParamModel reads R only from the datacard. - resp_inputs = [f for f in args.inputFile if scetlib_np_response.has_response(f)] - if len(resp_inputs) > 1: - raise RuntimeError( - "Multiple inputs carry the SCETlib-NP response (hist + gen-total): " - f"{resp_inputs}; expected at most one (the Z dilepton --unfolding run)." - ) - if resp_inputs: - if not hasattr(writer, "add_auxiliary"): + if args.storeResponseMatrix: + resp_inputs = [f for f in args.inputFile if scetlib_np_response.has_response(f)] + if len(resp_inputs) > 1: raise RuntimeError( - "TensorWriter has no 'add_auxiliary'; update the rabbit submodule " - "to a revision including WMass/rabbit#145 to embed the SCETlib-NP " - "response matrix in the datacard." + "Multiple inputs carry the SCETlib-NP response (hist + gen-total): " + f"{resp_inputs}; expected at most one (the Z dilepton --unfolding run)." ) - logger.info(f"Embedding SCETlib-NP response matrix from {resp_inputs[0]}") - R_info = scetlib_np_response.load_R(resp_inputs[0]) - writer.add_auxiliary( - "scetlib_np", - { - "R": R_info["R"], - "N_gen": R_info["N_gen"], - "reco_axes": [n for n, _ in R_info["reco_axes"]], - "gen_axes": [n for n, _ in R_info["gen_axes"]], - # one edges dataset per reco/gen axis (variable length) - **{ - f"edges__{n}": e - for n, e in R_info["reco_axes"] + R_info["gen_axes"] + if resp_inputs: + logger.info(f"Embedding SCETlib-NP response matrix from {resp_inputs[0]}") + R_info = scetlib_np_response.load_R(resp_inputs[0]) + writer.add_auxiliary( + "scetlib_np", + { + "R": R_info["R"], + "N_gen": R_info["N_gen"], + "reco_axes": [n for n, _ in R_info["reco_axes"]], + "gen_axes": [n for n, _ in R_info["gen_axes"]], + # one edges dataset per reco/gen axis (variable length) + **{ + f"edges__{n}": e + for n, e in R_info["reco_axes"] + R_info["gen_axes"] + }, }, - }, - ) + ) # propagate meta info into result file meta = { From 108b50d384126a5b1ad8c30c03eb1f0bc6e71d02 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 24 Jun 2026 11:07:24 -0400 Subject: [PATCH 21/31] Bump rabbit submodule to WMass/rabbit main (948a94a) The PR previously pinned an off-main rabbit feature commit (6bfa4ff). All required rabbit-side changes (ParamModel priors/blinding, auxiliary array group, composite POU [POIs|POUs] layout fix, asym impacts) are now merged into rabbit main, so track it directly. Co-Authored-By: Claude Opus 4.8 (1M context) --- rabbit | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rabbit b/rabbit index 6bfa4ffcb..948a94acf 160000 --- a/rabbit +++ b/rabbit @@ -1 +1 @@ -Subproject commit 6bfa4ffcbad0e80d58423e1ecbca41654131b2c8 +Subproject commit 948a94acfb4add5807f901013fd5fc886b4eaee4 From fbd20903a101bcdab62f2f9694ab48db112fc98f Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 24 Jun 2026 11:07:26 -0400 Subject: [PATCH 22/31] Restore wremnants-data submodule pin to main (3d2b2b2) Revert the inadvertent wremnants-data bump so the PR matches upstream main. Co-Authored-By: Claude Opus 4.8 (1M context) --- wremnants-data | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wremnants-data b/wremnants-data index 81ff2c21c..3d2b2b214 160000 --- a/wremnants-data +++ b/wremnants-data @@ -1 +1 @@ -Subproject commit 81ff2c21cffae7cb0c853c17484ed6c745ae36ae +Subproject commit 3d2b2b2140751b6212898a923ede103793d89ee2 From b7fc284ca21ea15e84f88a1339d7ca851b004370 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 24 Jun 2026 11:30:57 -0400 Subject: [PATCH 23/31] Review cleanup: drop debug prints, clarify external refs, trim docstring - input_tools.read_dyturbo_vars_hist: remove two leftover debug print()s. - btgrid_integrate / param_model: the "numpy reference" comments named functions that live in the external scetlib_run.factorize library, not in this package; qualify them so a maintainer grepping the package isn't lost. - param_model module docstring: condense the GN-vs-full-K implementation history (kept the Hessian formula, the Asimov/GN rationale, and the _frozen_eq_zero footgun; dropped the blow-by-blow narration). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../scetlib_np/btgrid_integrate.py | 10 +++---- .../postprocessing/scetlib_np/param_model.py | 28 ++++++++----------- wremnants/utilities/io_tools/input_tools.py | 2 -- 3 files changed, 17 insertions(+), 23 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py index 73d325280..ed7fcfaec 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py @@ -102,7 +102,7 @@ def sparse_to_dense_tf(sigma_flat, flat_idx): # ============================================================================= -# Q-integration weights (arctan_Q² method, matches numpy integrate_over_Q) +# Q-integration weights (arctan_Q² method, matches scetlib_run.factorize.integrate_over_Q) # ============================================================================= @@ -111,8 +111,8 @@ def q_integrate_weights( ): """Simpson weights for integrating over Q ∈ [Q_lo, Q_hi] in arctan-Q² space. - Implements the same change of variable as - the numpy-reference ``integrate_over_Q`` with ``method="arctan_Q2"``: + Implements the same change of variable as the external numpy reference + ``scetlib_run.factorize.integrate_over_Q`` with ``method="arctan_Q2"``: x = arctan((Q² - q0²) / (q0 Γ)) flattens the Breit-Wigner peak, then Simpson on x with the Jacobian dQ/dx. @@ -155,8 +155,8 @@ def rebin_weights(source_grid, target_edges, name="axis", tol=1e-9): matrix; entries are 0 for source samples not contributing to a given target bin. - Mirrors the per-bin call pattern of - the numpy-reference ``integrate_over_axis_bin``. + Mirrors the per-bin call pattern of the external numpy reference + ``scetlib_run.factorize`` (``integrate_over_{qT,Y}_bin``). """ source_grid = np.asarray(source_grid, dtype=np.float64) target_edges = np.asarray(target_edges, dtype=np.float64) diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 64cab9859..fef637eaa 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -225,19 +225,14 @@ GN vs full-K. The Poisson Hessian is H_ij = Σ_b [ (n_b/μ_b²) J_bi J_bj + (1 − n_b/μ_b) K_bij ]. For ASIMOV data the residual (1 − n/μ) = 0, so the K term vanishes and GN (J -only) is EXACT — drop the K term with hessian_gn=1. For real/toy data -the K term is needed for the exact Hessian. The nested-forward-mode K -(``_ratio_compact_hess``) USED to crash under rabbit's @tf.function: building the -JVP of an ``Equal`` op that carries a tangent raised ``IndexError`` in TF's -nested-forward-mode autodiff. The culprit was the λ_inf==0 / den==0 masks in -:mod:`btgrid_tf` comparing a differentiated tensor; freezing the comparison input -(``btgrid_tf._frozen_eq_zero`` = ``tf.equal(stop_gradient(x), 0)``) removes the -tangent into ``Equal`` without changing any value or derivative (the comparison -is a measure-zero boundary ``tf.where`` never differentiates). Full-K now runs -under @tf.function and matches the exact reverse-mode Hessian to machine -precision (≤3e-16 rel; verified by the isolation validation). So -both GN and full-K are available; GN remains the default for Asimov (exact and -cheaper — 8 vs 72 fold passes). +only) is EXACT — drop the K term with hessian_gn=1. Real/toy data need the K +term; full-K (``_ratio_compact_hess``) matches the exact reverse-mode Hessian. +Full-K requires the ``btgrid_tf._frozen_eq_zero`` guard: the λ_inf==0 / den==0 +masks compare a differentiated tensor, which makes TF's nested-forward-mode AD +raise ``IndexError`` on the ``Equal`` op; freezing the comparison input +(``tf.equal(stop_gradient(x), 0)``) drops the tangent into ``Equal`` without +changing any value or derivative (it is a measure-zero ``tf.where`` boundary). +GN remains the default for Asimov (exact and cheaper — 8 vs 72 fold passes). WARNING: do NOT pass hessian_straightthrough=1 during the FIT. The surrogate recomputes J(/K) on every compute() call — fine for the one-shot @@ -843,8 +838,8 @@ def __init__( # circular (σ_gen cancels to R_raw·1). # Native-binning Q-integrated reconstruction (NY, NqT) on the signed-Y / # qT grid, BEFORE the |Y|-fold and qT-rebin — exposed so the native-binning - # validation can compare it to the SCETlib reference / numpy factorize - # without the projection layer. + # validation can compare it to the SCETlib reference / external + # scetlib_run.factorize without the projection layer. self.sigma_YqT_central = self._sigma_YqT_native_at( self.eff_central, self.gnu_central ) @@ -1105,7 +1100,8 @@ def _sigma_YqT_native_at(self, eff_params, gnu_params): in the btgrid's *native* binning: shape (NY, NqT) on the signed-Y / qT grid (Y_unique, qT_unique), BEFORE the |Y|-fold and qT-rebin. This is the object that the native-binning validation compares against the - SCETlib spectrum reference (curve 1) and the numpy `factorize` (curve 2).""" + SCETlib spectrum reference (curve 1) and the external scetlib_run.factorize + (curve 2).""" # 1. Reconstruct σ on the btgrid's flat (Nbins,) layout. Factorized # (default) and legacy layouts are numerically equivalent (≲1e-14 # rel.; summation order only). diff --git a/wremnants/utilities/io_tools/input_tools.py b/wremnants/utilities/io_tools/input_tools.py index e45d467a2..d2052e649 100644 --- a/wremnants/utilities/io_tools/input_tools.py +++ b/wremnants/utilities/io_tools/input_tools.py @@ -215,9 +215,7 @@ def read_dyturbo_vars_hist(base_name, var_axis=None, axes=("Y", "qT"), charge=No f"Scale variation {var} found for fo_sing piece but no corresponding variation for dyturbo" ) dyturbo_scale = scales_map.get(var, "mur1-muf1") - print(var, dyturbo_scale) dyturbo_name = base_name.format(i=pdf_member, scale=dyturbo_scale) - print(dyturbo_name) h = read_dyturbo_hist([dyturbo_name], axes=axes, charge=charge) if not var_hist: var_hist = hist.Hist(*h.axes, var_axis, storage=h.storage_type()) From 33fdde88f020316946438478211afbd7cd2f7a3e Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 24 Jun 2026 19:19:28 -0400 Subject: [PATCH 24/31] =?UTF-8?q?SCETlib-NP:=20add=20gen=5Flevel=3D1=20mod?= =?UTF-8?q?e=20(gen-level=20=CF=83UL=20fit,=20no=20response=20fold)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit gen_level=1 makes the fit channel the gen (ptVGen, absY) binning and skips Step 3 (the gen→reco fold): compute() returns the per-gen-bin ratio σ_gen(λ)/σ_gen(λ_central) from Steps 1–2. No scetlib_np auxiliary / R / N_gen is required — the gen binning is read from the single fit channel. For the direct-theory σUL closure. Recovered from an uncommitted stash where it had been lost during the PR701 sync. param_model.py: gen_level constructor flag + docstring; branch the R / gen-axes setup, the central-denominator caching, and _ratio_from_param so the normal reco path is byte-for-byte unchanged when gen_level is off. Also folds in the pending non-finite btgrid-cell sanitization (NaN/inf -> 0 for kinematically-forbidden grid points) that was sitting in the working tree. Co-Authored-By: Claude Opus 4.8 --- .../postprocessing/scetlib_np/param_model.py | 167 ++++++++++++------ 1 file changed, 116 insertions(+), 51 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index fef637eaa..9aee3044e 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -512,6 +512,7 @@ def __init__( xparam_default: Optional[str] = None, hessian_straightthrough: bool = False, hessian_gn: bool = False, + gen_level: bool = False, **kwargs, ): """Construct the ParamModel. @@ -602,6 +603,15 @@ def __init__( hessian_gn With ``hessian_straightthrough``: Gauss-Newton, keep J and drop the K term (exact for Asimov, where the residual vanishes). + gen_level + Gen-level σUL fit mode (spec token ``gen_level=1``). The fit channel + IS the gen (ptVGen, |Y|) binning, so there is NO response matrix and + NO gen→reco fold (Step 3 is skipped): compute() returns the + per-GEN-bin ratio σ_gen(λ) / σ_gen(λ_central) from Steps 1–2. The + gen binning is read from the single fit channel's axes, so no + ``scetlib_np`` auxiliary / R / N_gen is needed. Used for the + direct-theory σUL closure (this ParamModel as the λ model, fit + against an injected gen-level σUL pseudodata). """ self.indata = indata @@ -663,6 +673,24 @@ def __init__( # ---- btgrid + dense layout grid = btgrid_cache.load(btgrid_dir) + # Sanitize non-finite bt-grid cells. Kinematically-forbidden points + # (x = (Q/Ecm)·e^|Y| ≥ 1, e.g. extreme forward Y near the Z peak) can + # come back as NaN from SCETlib instead of the physical 0. Their true + # cross section is zero, so replace NaN/inf with 0 → harmless zero-rows. + # Without this, dedup_grid_rows' hash-group verification fails (NaN != + # NaN). (For grids produced as condor shards the forbidden cells are + # simply absent and dense_index_map 0-fills them; a single-process local + # run instead writes them in as NaN, which is what this handles.) + for _key in ("I_pert", "C_nu"): + _arr = grid[_key] + if not np.isfinite(_arr).all(): + _nbad = int((~np.isfinite(_arr)).any(axis=-1).sum()) + np.nan_to_num(_arr, copy=False, nan=0.0, posinf=0.0, neginf=0.0) + print( + f"[SCETlibNPParamModel] sanitized {_nbad} non-finite " + f"{_key} bt-grid rows -> 0 (kinematically-forbidden cells)", + flush=True, + ) idx_map = fz_int.dense_index_map(grid["bins"]) self.Q_unique = idx_map["Q_unique"] self.Y_unique = idx_map["Y_unique"] @@ -754,41 +782,65 @@ def __init__( dtype=fz_tf.DTYPE, ) - # ---- R matrix (read from the datacard's scetlib_np auxiliary) - R_info = _R_info_from_auxiliary(indata) - # The fit-tensor's reco binning may differ from R's by trailing - # overflow bins (e.g. R has ptll [0, …, 44, 100] while the fit ends - # at 44). Crop R's trailing bins so the reco shape matches. - fit_reco_axes = self._fit_reco_axes(indata) - R_arr = _crop_R_to_fit(R_info["R"], R_info["reco_axes"], fit_reco_axes) - # Tighten the metadata to match the cropped R. - self.reco_shape = R_arr.shape[: len(fit_reco_axes)] - self.gen_shape = R_arr.shape[len(fit_reco_axes) :] - N_reco = int(np.prod(self.reco_shape)) - N_gen = int(np.prod(self.gen_shape)) - # Raw response counts; normalized to a response below. - self._R_raw = tf.constant(R_arr.reshape(N_reco, N_gen), dtype=fz_tf.DTYPE) - # Gen-total denominator N_gen(g) from the xnorm hist ("prefsr"): the - # generated fiducial yield per gen bin (pre-reco-selection). Dividing R - # by this gives the theory-independent efficiency×migration response. - # REQUIRED: setupRabbit only embeds the response when the gen-total is - # present, so N_gen should always be here; raise if not (a σ_gen(λ_c) - # proxy would make the central closure circular — see module docstring). - if R_info.get("N_gen") is None: - raise ValueError( - "SCETlibNPParamModel: the 'scetlib_np' auxiliary has no N_gen " - "(gen-total). Rebuild the datacard from a histmaker output that " - "carries the 'prefsr' xnorm hist." - ) - self._N_gen_flat = tf.constant(R_info["N_gen"].reshape(-1), dtype=fz_tf.DTYPE) - self._reco_axes_meta = [ - (name, fit_axes[1]) - for (name, fit_axes) in zip( - [a[0] for a in R_info["reco_axes"]], - fit_reco_axes, + # ---- Gen/reco binning + (reco mode) the response matrix R. + # gen_level=1: the fit channel IS the gen (ptVGen, absY) binning, so + # there is NO response matrix and NO gen→reco fold — compute() returns + # the per-GEN-bin ratio σ_gen(λ)/σ_gen(λ_central) (Steps 1–2 only), and + # the scetlib_np auxiliary / N_gen are not required. + self.gen_level = bool(gen_level) + if self.gen_level: + gen_axes = self._fit_reco_axes(indata) + if len(gen_axes) != 2: + raise NotImplementedError( + "gen_level SCETlibNPParamModel expects a single fit channel " + "with 2 gen axes (ptVGen, absY); got " + f"{[n for n, _ in gen_axes]}" + ) + self.R = None + self._R_raw = None + self._N_gen_flat = None + self.reco_shape = None + self._reco_axes_meta = None + self._gen_axes_meta = gen_axes + self.gen_shape = tuple(len(e) - 1 for (_, e) in gen_axes) + else: + # ---- R matrix (read from the datacard's scetlib_np auxiliary) + R_info = _R_info_from_auxiliary(indata) + # The fit-tensor's reco binning may differ from R's by trailing + # overflow bins (e.g. R has ptll [0, …, 44, 100] while the fit ends + # at 44). Crop R's trailing bins so the reco shape matches. + fit_reco_axes = self._fit_reco_axes(indata) + R_arr = _crop_R_to_fit(R_info["R"], R_info["reco_axes"], fit_reco_axes) + # Tighten the metadata to match the cropped R. + self.reco_shape = R_arr.shape[: len(fit_reco_axes)] + self.gen_shape = R_arr.shape[len(fit_reco_axes) :] + N_reco = int(np.prod(self.reco_shape)) + N_gen = int(np.prod(self.gen_shape)) + # Raw response counts; normalized to a response below. + self._R_raw = tf.constant(R_arr.reshape(N_reco, N_gen), dtype=fz_tf.DTYPE) + # Gen-total denominator N_gen(g) from the xnorm hist ("prefsr"): the + # generated fiducial yield per gen bin (pre-reco-selection). Dividing R + # by this gives the theory-independent efficiency×migration response. + # REQUIRED: setupRabbit only embeds the response when the gen-total is + # present, so N_gen should always be here; raise if not (a σ_gen(λ_c) + # proxy would make the central closure circular — see module docstring). + if R_info.get("N_gen") is None: + raise ValueError( + "SCETlibNPParamModel: the 'scetlib_np' auxiliary has no N_gen " + "(gen-total). Rebuild the datacard from a histmaker output that " + "carries the 'prefsr' xnorm hist." + ) + self._N_gen_flat = tf.constant( + R_info["N_gen"].reshape(-1), dtype=fz_tf.DTYPE ) - ] - self._gen_axes_meta = R_info["gen_axes"] + self._reco_axes_meta = [ + (name, fit_axes[1]) + for (name, fit_axes) in zip( + [a[0] for a in R_info["reco_axes"]], + fit_reco_axes, + ) + ] + self._gen_axes_meta = R_info["gen_axes"] # ---- Rebin weights: btgrid (NY signed) → (NabsYVGen) via |Y| folding # and (NqT) → (NptVGen). @@ -898,21 +950,30 @@ def __init__( f"SCETlibNPParamModel: {n_bad} gen bins have non-positive " f"σ_gen(λ_central); cannot normalize / fold the response." ) - N_gen = self._N_gen_flat - # Guard empty gen bins (no generated events): leave column at 0. - safe_N_gen = tf.where(N_gen > 0, N_gen, tf.ones_like(N_gen)) - self.R = self._R_raw / safe_N_gen[tf.newaxis, :] # P(b|g) = R_raw / N_gen - # Free the raw counts: only the normalized response self.R is used from - # here on (compute() never touches _R_raw) — no need to hold both. - del self._R_raw - self.sigma_reco_central = tf.linalg.matvec(self.R, gen_flat) # Σ_g P·σ_gen(λ_c) - if tf.reduce_any(self.sigma_reco_central <= 0).numpy(): - n_bad = int(tf.reduce_sum(tf.cast(self.sigma_reco_central <= 0, tf.int32))) - raise ValueError( - f"SCETlibNPParamModel: {n_bad} reco bins have non-positive " - f"σ_reco(λ_central). Likely a binning mismatch between R and " - f"the fit-tensor reco axes." - ) + if self.gen_level: + # Gen-level σUL fit: the fit bins ARE the gen bins, so the per-bin + # ratio denominator is σ_gen(λ_central) directly (no reco fold). + self.sigma_gen_central_flat = gen_flat + else: + N_gen = self._N_gen_flat + # Guard empty gen bins (no generated events): leave column at 0. + safe_N_gen = tf.where(N_gen > 0, N_gen, tf.ones_like(N_gen)) + self.R = self._R_raw / safe_N_gen[tf.newaxis, :] # P(b|g) = R_raw / N_gen + # Free the raw counts: only the normalized response self.R is used from + # here on (compute() never touches _R_raw) — no need to hold both. + del self._R_raw + self.sigma_reco_central = tf.linalg.matvec( + self.R, gen_flat + ) # Σ_g P·σ_gen(λ_c) + if tf.reduce_any(self.sigma_reco_central <= 0).numpy(): + n_bad = int( + tf.reduce_sum(tf.cast(self.sigma_reco_central <= 0, tf.int32)) + ) + raise ValueError( + f"SCETlibNPParamModel: {n_bad} reco bins have non-positive " + f"σ_reco(λ_central). Likely a binning mismatch between R and " + f"the fit-tensor reco axes." + ) # ---- Process index: signal column gets the ratio, others get 1. procs = [p.decode() if isinstance(p, bytes) else str(p) for p in indata.procs] @@ -1215,8 +1276,12 @@ def _ratio_from_param(self, param): eff_params, gnu_params = self._unpack_params(param) sigma_gen = self._sigma_gen_at(eff_params, gnu_params) gen_flat = tf.reshape(sigma_gen, [-1]) - sigma_reco = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) - ratio = sigma_reco / self.sigma_reco_central # (N_reco,) + if self.gen_level: + # Gen-level σUL fit: the fit bins ARE the gen bins — no reco fold. + ratio = gen_flat / self.sigma_gen_central_flat # (N_gen,) + else: + sigma_reco = tf.linalg.matvec(self.R, gen_flat) # (N_reco,) + ratio = sigma_reco / self.sigma_reco_central # (N_reco,) scale = tf.constant(RATIO_FLOOR_SCALE, dtype=ratio.dtype) ratio = tf.maximum( scale * tf.math.softplus(ratio / scale), From ccab63c2810d1628ec42a5e5541a057fa4fb8be3 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 24 Jun 2026 19:19:28 -0400 Subject: [PATCH 25/31] =?UTF-8?q?SCETlib-NP:=20add=20NP=20form-factor=20pl?= =?UTF-8?q?otting=20+=20fitresult-=CE=BB=20reader?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two companion modules for visualizing the SCETlib NP form factors from a fit: - np_function_plots.py: a PURE plotter for the CS γ_ν^NP(b_T) and TMD F_eff(b_T, y) form factors. Calls the same btgrid_tf.F_eff_tf / gamma_nu_NP_tf the fit integrates (driven by the np_model strings), so a drawn curve is exactly the fitted model. Takes physical λ sets + optional toys for a percentile band; knows nothing about fitresults. Raw-λ and --from-fitresult CLI modes. - fitresult_lambdas.py: the OUTPUT-side reader that turns a rabbit fitresults HDF5 into the λ sets / toy ensembles the plotter consumes. Reads λ from EITHER fit flavour behind one NPLambdas interface: the new continuous-λ param model (physical λ in parms + covariance) and the old template-based fit (discrete-nuisance pulls mapped to physical λ). Also prints the prefit/postfit/constraint/frozen λ table (CLI). Co-Authored-By: Claude Opus 4.8 --- .../scetlib_np/fitresult_lambdas.py | 513 ++++++++++++++++++ .../scetlib_np/np_function_plots.py | 357 ++++++++++++ 2 files changed, 870 insertions(+) create mode 100644 wremnants/postprocessing/scetlib_np/fitresult_lambdas.py create mode 100644 wremnants/postprocessing/scetlib_np/np_function_plots.py diff --git a/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py b/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py new file mode 100644 index 000000000..250904ed7 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py @@ -0,0 +1,513 @@ +"""Read SCETlib NP λ out of a rabbit fitresults HDF5 (table / curves / toys). + +The OUTPUT-side companion to :mod:`lambda_central` (which reads the INPUT-side +λ_central from the upstream correction pkl). Two responsibilities, kept apart +from the plotting: + + * Tabulate the λ — prefit/postfit value, prefit/postfit 1σ constraint, and + whether each was frozen — and print it (``read_lambdas`` + the CLI here). + * Turn a fitresults into the λ sets / toy ensembles the pure plotter + :mod:`np_function_plots` consumes (``lambdas_from_fitresult``, + ``sample_lambda_toys``, ``plot_series_from_fitresult``). + +Two fit flavours are supported through SEPARATE readers, both emitting the same +:class:`~np_function_plots.NPLambdas` interface so the plotter never learns +which it came from: + + * NEW continuous-λ param model — λ stored PHYSICALLY in ``parms`` + (``allowNegativeParam=True``; no sqrt convention), covariance over the + floating λ in ``cov``. ``lambdas_from_fitresult`` / ``sample_lambda_toys``. + * OLD template-based fit — discrete NP nuisances (``scetlibNPgamma*`` …) + whose pulls map to physical λ through a template/piecewise map. + ``lambdas_from_template_fit`` (param-map driven). + +Units: the new-model λ are already physical (see ``param_model`` / +``allowNegativeParam``); no conversion is applied. The ``np_model`` / +``np_model_nu`` strings the curves need are recovered via +:func:`lambda_central.read_lambda_central` (which reads them straight from a +fitresults), with a CLI/argument override. + +CLI (print the table):: + + python -m wremnants.postprocessing.scetlib_np.fitresult_lambdas +""" + +import argparse +import math +import re + +import numpy as np + +from rabbit import io_tools +from wremnants.postprocessing.scetlib_np import lambda_central as _lc +from wremnants.postprocessing.scetlib_np.np_function_plots import ( + EFF_PARAMS, + GNU_PARAMS, + NPLambdas, + Series, +) + +ALL_PARAMS = GNU_PARAMS + EFF_PARAMS +SECTOR = { + **{p: "gamma_nu (CS)" for p in GNU_PARAMS}, + **{p: "F_eff (TMD)" for p in EFF_PARAMS}, +} +# Fallback model strings when lambda_central can't reach the upstream pkl. +DEFAULT_NP_MODEL = "tanh_6" +DEFAULT_NP_MODEL_NU = "tanh_2" + + +# --------------------------------------------------------------------------- +# low-level helpers +# --------------------------------------------------------------------------- + + +def _frozen_matcher(freeze_exprs): + """rabbit's freeze semantics: exact name OR anchored re.match. None if no list.""" + if not freeze_exprs: + return None + exact = set(freeze_exprs) + compiled = [re.compile(e) for e in freeze_exprs] + return lambda name: (name in exact) or any(r.match(name) for r in compiled) + + +def _parse_param_model_spec(args): + """(model_name, {token: value}) from the recorded --paramModel arg.""" + pm = args.get("paramModel") + if not pm: + return None, {} + tokens = pm[0] if isinstance(pm[0], (list, tuple)) else pm + tokens = [t.decode() if isinstance(t, bytes) else str(t) for t in tokens] + if not tokens: + return None, {} + spec = {} + for tok in tokens[1:]: + if "=" in tok: + k, v = tok.split("=", 1) + spec[k] = v + return tokens[0], spec + + +def _names(parms_hist): + return [ + n.decode() if isinstance(n, bytes) else str(n) for n in list(parms_hist.axes[0]) + ] + + +# --------------------------------------------------------------------------- +# table +# --------------------------------------------------------------------------- + + +def read_lambdas(fitresult_path, result=None, params=None): + """Structured λ readout from a (new-model) fitresults. + + Returns a dict:: + + { + "context": {file, result, model, freeze, spec, ...}, + "params": { name: {prefit, postfit, prefit_sigma, postfit_sigma, + frozen, sector, present} , ... } + } + + ``params`` defaults to all known λ; pass a list to restrict/extend. + """ + fitresult, meta = io_tools.get_fitresult(fitresult_path, result, meta=True) + post = fitresult["parms"].get() + pre = fitresult["parms_prefit"].get() + names = _names(post) + idx = {n: i for i, n in enumerate(names)} + + pv, prv = post.values(), pre.values() + post_var, pre_var = post.variances(), pre.variances() + + fit_args = ( + (meta.get("meta_info", {}) or {}).get("args", {}) + if isinstance(meta, dict) + else {} + ) + freeze = fit_args.get("freezeParameters") or [] + is_frozen = _frozen_matcher(freeze) + model_name, spec = _parse_param_model_spec(fit_args) + + want = list(params) if params else list(ALL_PARAMS) + + out = {} + for name in want: + if name not in idx: + out[name] = dict(present=False, sector=SECTOR.get(name, "")) + continue + i = idx[name] + post_sigma = ( + math.sqrt(post_var[i]) + if post_var is not None and post_var[i] >= 0 + else float("nan") + ) + pre_sigma = ( + math.sqrt(pre_var[i]) + if pre_var is not None and pre_var[i] >= 0 + else float("nan") + ) + if is_frozen is not None: + frozen = bool(is_frozen(name)) + else: + frozen = (post_var is not None) and (post_var[i] == 0.0) + out[name] = dict( + present=True, + sector=SECTOR.get(name, ""), + prefit=float(prv[i]), + postfit=float(pv[i]), + prefit_sigma=float(pre_sigma), + postfit_sigma=float(post_sigma), + frozen=frozen, + ) + return { + "context": dict( + file=fitresult_path, + result=result or "(default)", + model=model_name, + freeze=freeze, + spec=spec, + ), + "params": out, + } + + +def format_lambda_table(readout): + """Render the :func:`read_lambdas` result as the printable table string.""" + ctx, params = readout["context"], readout["params"] + + def fmt(x, nd=5): + if x is None or (isinstance(x, float) and math.isnan(x)): + return "n/a" + return f"{x:.{nd}g}" + + lines = [] + lines.append(f"File : {ctx['file']}") + lines.append(f"Result : {ctx['result']}") + if ctx.get("model"): + lines.append(f"Model : {ctx['model']}") + spec = ctx.get("spec") or {} + if spec.get("xparam_default"): + lines.append(f" xparam_default (start shift) : {spec['xparam_default']}") + if spec.get("priors") in ("1", "true", "True", "yes", "on"): + extra = f" ({spec['prior_sigmas']})" if spec.get("prior_sigmas") else "" + lines.append(f" priors : ENABLED{extra}") + if ctx.get("freeze"): + lines.append(f"Freeze : {ctx['freeze']}") + lines.append("Values are PHYSICAL λ (allowNegativeParam=True; no sqrt conversion).") + lines.append("") + + cols = ("Parameter", "Sector", "Fixed", "Prefit", "Postfit", "Postfit±", "Prefit±") + w = (16, 14, 6, 12, 12, 12, 10) + header = " ".join( + f"{c:>{wi}}" if i else f"{c:<{wi}}" for i, (c, wi) in enumerate(zip(cols, w)) + ) + lines.append(header) + lines.append("-" * len(header)) + for name, d in params.items(): + sector = d.get("sector", "") + if not d.get("present"): + lines.append( + " ".join( + [ + f"{name:<{w[0]}}", + f"{sector:>{w[1]}}", + f"{'--':>{w[2]}}", + f"{'absent':>{w[3]}}", + f"{'absent':>{w[4]}}", + f"{'--':>{w[5]}}", + f"{'--':>{w[6]}}", + ] + ) + ) + continue + lines.append( + " ".join( + [ + f"{name:<{w[0]}}", + f"{sector:>{w[1]}}", + f"{('YES' if d['frozen'] else 'no'):>{w[2]}}", + f"{fmt(d['prefit']):>{w[3]}}", + f"{fmt(d['postfit']):>{w[4]}}", + f"{fmt(d['postfit_sigma']):>{w[5]}}", + f"{fmt(d['prefit_sigma']):>{w[6]}}", + ] + ) + ) + lines.append("") + lines.append( + "Postfit± / Prefit± are the 1σ constraints (sqrt of the stored variance)." + ) + lines.append( + "Prefit± = 0 means unconstrained (no Gaussian prior); >0 is the prior σ." + ) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# curve / toy feeding (new continuous-λ model) +# --------------------------------------------------------------------------- + + +def _flat_values(fitresult_path, which="postfit", result=None): + """{λ name: value} for all present λ, from ``parms`` (postfit) or + ``parms_prefit`` (prefit).""" + fitresult = io_tools.get_fitresult(fitresult_path, result) + h = fitresult["parms" if which == "postfit" else "parms_prefit"].get() + names = _names(h) + vals = h.values() + return {p: float(vals[names.index(p)]) for p in ALL_PARAMS if p in names} + + +def _resolve_models(fitresult_path, np_model=None, np_model_nu=None): + """np_model strings from lambda_central (reads them off the fitresults), + falling back to defaults if the upstream pkl is unreachable. Explicit + arguments win.""" + if np_model and np_model_nu: + return np_model, np_model_nu + eff_model, gnu_model = DEFAULT_NP_MODEL, DEFAULT_NP_MODEL_NU + try: + lc = _lc.read_lambda_central(fitresult_path) + eff_model = lc["eff_params"].get("np_model", eff_model) + gnu_model = lc["gnu_params"].get("np_model_nu", gnu_model) + except Exception as exc: # pkl unreachable / non-NP fit: warn, use defaults + print( + f"[fitresult_lambdas] could not read np_model from {fitresult_path} " + f"({exc}); using defaults {eff_model!r}/{gnu_model!r}. " + f"Pass --np-model/--np-model-nu to override." + ) + return np_model or eff_model, np_model_nu or gnu_model + + +def lambdas_from_fitresult( + fitresult_path, which="postfit", result=None, np_model=None, np_model_nu=None +): + """:class:`NPLambdas` for the prefit or postfit point of a new-model fit.""" + np_model, np_model_nu = _resolve_models(fitresult_path, np_model, np_model_nu) + vals = _flat_values(fitresult_path, which=which, result=result) + return NPLambdas.from_flat(vals, np_model, np_model_nu) + + +def read_lambda_covariance(fitresult_path, result=None, names=ALL_PARAMS): + """(floating_names, mean, cov) over the FLOATING λ (variance > 0). + + Frozen λ (variance 0, absent, or pinned) are excluded — they don't enter + the toy band. The submatrix keeps the full correlations among the floating + λ (these are strong: e.g. lambda2_nu↔lambda2 ≈ −1).""" + fitresult = io_tools.get_fitresult(fitresult_path, result) + parms = fitresult["parms"].get() + pnames = _names(parms) + pvals = parms.values() + cov_h = fitresult["cov"].get() + cov_names = _names(cov_h) + cov = cov_h.values() + cidx = {n: i for i, n in enumerate(cov_names)} + + floating, mean, sel = [], [], [] + for p in names: + if p in cidx and cov[cidx[p], cidx[p]] > 0: + floating.append(p) + sel.append(cidx[p]) + mean.append(float(pvals[pnames.index(p)])) + if not floating: + return [], np.zeros(0), np.zeros((0, 0)) + sub = cov[np.ix_(sel, sel)] + return floating, np.asarray(mean), np.asarray(sub) + + +def sample_lambda_toys( + fitresult_path, n_toys=500, seed=0, result=None, np_model=None, np_model_nu=None +): + """List of :class:`NPLambdas` toys sampled from the postfit MVN. + + Floating λ are drawn jointly from their postfit covariance; frozen λ are + held at their postfit value. (For real/Asimov data the postfit point is the + correct band centre.)""" + np_model, np_model_nu = _resolve_models(fitresult_path, np_model, np_model_nu) + base = _flat_values(fitresult_path, which="postfit", result=result) + floating, mean, cov = read_lambda_covariance(fitresult_path, result=result) + if not floating: + return [] + rng = np.random.default_rng(seed) + draws = rng.multivariate_normal(mean, cov, size=n_toys) + toys = [] + for d in draws: + vals = dict(base) + for k, name in enumerate(floating): + vals[name] = float(d[k]) + toys.append(NPLambdas.from_flat(vals, np_model, np_model_nu)) + return toys + + +def plot_series_from_fitresult( + fitresult_path, result=None, n_toys=500, seed=0, np_model=None, np_model_nu=None +): + """Build the [prefit dashed, postfit solid + band] series for the plotter.""" + np_model, np_model_nu = _resolve_models(fitresult_path, np_model, np_model_nu) + pre = lambdas_from_fitresult( + fitresult_path, "prefit", result, np_model, np_model_nu + ) + post = lambdas_from_fitresult( + fitresult_path, "postfit", result, np_model, np_model_nu + ) + toys = sample_lambda_toys( + fitresult_path, n_toys, seed, result, np_model, np_model_nu + ) + return [ + Series(label="prefit (λ_central)", lam=pre, color="C0", linestyle="--", lw=1.8), + Series(label="postfit", lam=post, color="C3", linestyle="-", lw=2.0, toys=toys), + ] + + +# --------------------------------------------------------------------------- +# old template-based fit reader (legacy adapter) +# --------------------------------------------------------------------------- + +# AN parameter name (per side) -> param-model name. +_AN_TO_BTGRID = { + ("CS", "lambda_2"): "lambda2_nu", + ("CS", "lambda_4"): "lambda4_nu", + ("CS", "lambda_inf"): "lambda_inf_nu", + ("TMD", "Lambda_2"): "lambda2", + ("TMD", "Lambda_4"): "lambda4", + ("TMD", "Delta_Lambda_2"): "delta_lambda2", + ("TMD", "Lambda_inf"): "lambda_inf", + ("TMD", "Lambda_6"): "lambda6", +} + + +def _template_theta_to_physical(theta, entry, kfactor): + """Piecewise linearization param(θ) used by the old discrete NP nuisances: + nominal + max(θ,0)·(Up−nom)·kf − max(−θ,0)·(nom−Down)·kf.""" + nom = entry["nominal"] + d_up = (entry["Up_template_value"] - nom) * kfactor + d_dn = (nom - entry["Down_template_value"]) * kfactor + return nom + max(theta, 0.0) * d_up - max(-theta, 0.0) * d_dn + + +def _template_base_eff_gnu(param_map): + """eff/gnu dicts seeded with the param-map's fixed parameters + models.""" + fixed = param_map.get("fixed_parameters", {}) + eff = dict( + lambda_inf=fixed.get("Lambda_inf_TMD", {}).get("value", 1.0), + lambda2=0.0, + lambda4=0.0, + lambda6=fixed.get("Lambda_6", {}).get("value", 0.016), + delta_lambda2=0.0, + np_model="tanh_6", + ) + gnu = dict(lambda_inf_nu=0.0, lambda2_nu=0.0, lambda4_nu=0.0, np_model_nu="tanh_6") + return eff, gnu + + +def _template_apply(eff, gnu, param_map, theta_by_nuis, kfactors): + """Fill eff/gnu (copies) with physical λ from per-nuisance θ.""" + eff, gnu = dict(eff), dict(gnu) + for nuis, entry in param_map["nuisances"].items(): + key = (entry["side"], entry["param_AN"]) + bt = _AN_TO_BTGRID.get(key) + if bt is None: + continue + kf = kfactors.get(nuis, kfactors.get(entry["param_AN"], 1.0)) + val = _template_theta_to_physical(theta_by_nuis.get(nuis, 0.0), entry, kf) + (gnu if entry["side"] == "CS" else eff)[bt] = val + return eff, gnu + + +def lambdas_from_template_fit( + fitresult_path, np_param_map, result=None, kfactors=None, n_toys=0, seed=0 +): + """Read an OLD template-based fit → (central :class:`NPLambdas`, toys list). + + ``np_param_map`` is the JSON path (or already-loaded dict) describing each + discrete NP nuisance's template Up/Down and its physical AN parameter. The + discrete-nuisance pulls are mapped to physical λ via the same piecewise + linearization the template histograms were built with. ``kfactors`` scales + a nuisance's template delta (use when the workspace was built with + ``--scaleParams``); keyed by rabbit nuisance name or AN param name. + + With ``n_toys>0`` the floating nuisances are sampled from their postfit + covariance (in NUISANCE space — the linearization is applied per toy), so + the band correctly reflects the nonlinear θ→λ map. + """ + import json + + if isinstance(np_param_map, str): + with open(np_param_map) as f: + np_param_map = json.load(f) + kfactors = kfactors or {} + + fitresult = io_tools.get_fitresult(fitresult_path, result) + parms = fitresult["parms"].get() + pnames = _names(parms) + pvals = parms.values() + nuis_list = list(np_param_map["nuisances"].keys()) + theta = { + nm: (float(pvals[pnames.index(nm)]) if nm in pnames else 0.0) + for nm in nuis_list + } + + base_eff, base_gnu = _template_base_eff_gnu(np_param_map) + c_eff, c_gnu = _template_apply(base_eff, base_gnu, np_param_map, theta, kfactors) + central = NPLambdas(eff=c_eff, gnu=c_gnu) + + toys = [] + if n_toys > 0: + cov_h = fitresult["cov"].get() + cov_names = _names(cov_h) + cov = cov_h.values() + cidx = {n: i for i, n in enumerate(cov_names)} + floating = [ + nm for nm in nuis_list if nm in cidx and cov[cidx[nm], cidx[nm]] > 0 + ] + if floating: + sel = [cidx[nm] for nm in floating] + sub = cov[np.ix_(sel, sel)] + mean = np.array([theta[nm] for nm in floating]) + rng = np.random.default_rng(seed) + draws = rng.multivariate_normal(mean, sub, size=n_toys) + for d in draws: + th = dict(theta) + for k, nm in enumerate(floating): + th[nm] = float(d[k]) + e, g = _template_apply(base_eff, base_gnu, np_param_map, th, kfactors) + toys.append(NPLambdas(eff=e, gnu=g)) + return central, toys + + +# --------------------------------------------------------------------------- +# CLI (print the table) +# --------------------------------------------------------------------------- + + +def make_parser(): + p = argparse.ArgumentParser( + description="Print SCETlib NP λ (prefit/postfit/constraints/fixed) from a rabbit fitresults HDF5.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument("infile", help="path to fitresults*.hdf5") + p.add_argument( + "--result", + default=None, + help="results group suffix (e.g. 'nominal'); default 'results'.", + ) + p.add_argument( + "--params", + nargs="+", + default=None, + help="restrict to these parameter names (default: all known λ).", + ) + return p + + +def main(argv=None): + args = make_parser().parse_args(argv) + readout = read_lambdas(args.infile, result=args.result, params=args.params) + print() + print(format_lambda_table(readout)) + print() + + +if __name__ == "__main__": + main() diff --git a/wremnants/postprocessing/scetlib_np/np_function_plots.py b/wremnants/postprocessing/scetlib_np/np_function_plots.py new file mode 100644 index 000000000..f90af5508 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/np_function_plots.py @@ -0,0 +1,357 @@ +"""Plot the SCETlib NP form factors — CS γ_ν^NP(b_T) and TMD F_eff(b_T, y). + +This is a PURE plotting library: it takes physical λ values (the two parameter +dicts the model uses) and draws the two NP functions. It knows nothing about +fitresults — where the λ come from (a new continuous-λ fit, an old +template-based fit, or hand-picked values to test) is the caller's job. The +companion reader :mod:`fitresult_lambdas` turns a fitresults HDF5 into the λ +sets / toy ensembles this module consumes; ``main()`` below glues the two for +convenience but the plot functions stay reader-agnostic. + +The curves call the same form factors the fit integrates +(:func:`btgrid_tf.F_eff_tf` / :func:`btgrid_tf.gamma_nu_NP_tf`), driven by the +``np_model`` / ``np_model_nu`` strings, so a plotted curve is exactly the model +the fit used. + +A "λ set" is the pair of dicts :class:`NPLambdas` carries: + + eff = {lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_model} (F_eff / TMD) + gnu = {lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu} (γ_ν / CS) + +These map 1:1 onto the form-factor keyword arguments. A band is drawn from +a list of λ-set "toys" handed in by the caller (this module just takes +percentiles of the resulting curves); it never samples and never sees a +covariance. + +CLI (plot a raw λ set, no fit involved):: + + python -m wremnants.postprocessing.scetlib_np.np_function_plots \\ + --lambda2 0.4 --lambda4 0.4 --lambda2_nu 0.15 \\ + --np-model tanh_6 --np-model-nu tanh_2 -o /tmp/np.png + +CLI (from a fitresults: prefit dashed, postfit solid + 68% band):: + + python -m wremnants.postprocessing.scetlib_np.np_function_plots \\ + --from-fitresult -o /tmp/np.png +""" + +import argparse +import os +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple + +import numpy as np + +from wremnants.postprocessing.scetlib_np import btgrid_tf + +# λ split across the two NP sectors, with sensible "all knobs off" defaults so a +# bare CLI call still draws something. Mirrors param_model.{GNU,EFF}_PARAMS. +GNU_PARAMS = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") +EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") + + +@dataclass +class NPLambdas: + """One physical λ point: the two form-factor parameter dicts. + + ``eff`` / ``gnu`` hold exactly the kwargs ``btgrid_tf.F_eff_tf`` / + ``gamma_nu_NP_tf`` expect (numeric λ + the ``np_model`` / ``np_model_nu`` + string). Build one from a fit via :mod:`fitresult_lambdas`, or by hand. + """ + + eff: dict + gnu: dict + + @classmethod + def from_flat(cls, values, np_model, np_model_nu): + """Build from a flat name->value mapping (the param-model λ names).""" + eff = {k: float(values[k]) for k in EFF_PARAMS if k in values} + gnu = {k: float(values[k]) for k in GNU_PARAMS if k in values} + eff["np_model"] = np_model + gnu["np_model_nu"] = np_model_nu + return cls(eff=eff, gnu=gnu) + + +@dataclass +class Series: + """A curve to draw: a λ set, styling, and optional toys for an error band.""" + + label: str + lam: NPLambdas + color: Optional[str] = None + linestyle: str = "-" + lw: float = 2.0 + toys: Optional[List[NPLambdas]] = None # band drawn from these if present + band_pct: Tuple[float, float] = (16.0, 84.0) + + +def gamma_nu_curve(bT, gnu): + """γ_ν^NP(b_T) for one gnu dict (CS sector).""" + return np.asarray(btgrid_tf.gamma_nu_NP_tf(bT, **gnu), dtype=float) + + +def f_eff_curve(bT, y, eff): + """F_eff(y, b_T) for one eff dict at rapidity ``y`` (TMD sector).""" + return np.asarray(btgrid_tf.F_eff_tf(y, bT, **eff), dtype=float) + + +def _band(curves, pct): + """(lo, hi) percentile envelope across a stack of curves, or None if empty.""" + if not len(curves): + return None + stack = np.asarray(curves) + return np.percentile(stack, pct[0], axis=0), np.percentile(stack, pct[1], axis=0) + + +_CORNERS = { + "upper right": (0.97, 0.97, "top", "right"), + "upper left": (0.03, 0.97, "top", "left"), + "lower left": (0.03, 0.03, "bottom", "left"), + "lower right": (0.97, 0.03, "bottom", "right"), +} + + +def _param_inset(ax, lam, sector, corner="upper right"): + """Small text box listing the λ that drive the panel.""" + if sector == "gnu": + lines = [ + rf"$\lambda_2^\nu = {lam.gnu.get('lambda2_nu', 0):+.4f}$", + rf"$\lambda_4^\nu = {lam.gnu.get('lambda4_nu', 0):+.4f}$", + rf"$\lambda_\infty^\nu = {lam.gnu.get('lambda_inf_nu', 0):+.4f}$", + rf"model: {lam.gnu.get('np_model_nu', '?')}", + ] + else: + lines = [ + rf"$\lambda_2 = {lam.eff.get('lambda2', 0):+.4f}$", + rf"$\lambda_4 = {lam.eff.get('lambda4', 0):+.4f}$", + rf"$\delta\lambda_2 = {lam.eff.get('delta_lambda2', 0):+.4f}$", + rf"$\lambda_6 = {lam.eff.get('lambda6', 0):+.4f}$", + rf"$\lambda_\infty = {lam.eff.get('lambda_inf', 0):+.4f}$", + rf"model: {lam.eff.get('np_model', '?')}", + ] + box = dict(boxstyle="round,pad=0.35", fc="white", ec="0.6", alpha=0.85) + x, y, va, ha = _CORNERS[corner] + ax.text( + x, + y, + "\n".join(lines), + transform=ax.transAxes, + fontsize=8, + va=va, + ha=ha, + bbox=box, + ) + + +def plot_np_functions( + series: Sequence[Series], + *, + y_values: Sequence[float] = (0.0, 2.5), + bT_max: float = 4.0, + n_points: int = 401, + outpath: str, + inset_from: Optional[Series] = None, + f_ymax: Optional[float] = None, +): + """Draw the two NP form factors for one or more λ sets. + + Parameters + ---------- + series + Curves to overlay. The first series with ``toys`` set draws a + percentile band on each panel. Pure: the caller supplies the toys. + y_values + Rapidity values for the TMD panel (F_eff depends on y; γ_ν does not). + bT_max, n_points + b_T grid for the curves [GeV^-1]. + inset_from + Series whose λ are written into the per-panel parameter box (defaults + to the last series — typically the postfit point). + """ + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + bT = np.linspace(0.0, bT_max, n_points) + fig, (axL, axR) = plt.subplots(1, 2, figsize=(13, 4.5)) + + auto_colors = [c for c in plt.rcParams["axes.prop_cycle"].by_key()["color"]] + cmap_tmd = plt.cm.viridis + + # NP factors are evaluated at the bare b_T grid (this grid's b* prescription + # is the identity, b_bar == b_T; see param_model / base.conf b0_over_bmax=0). + # For a b*-frozen grid the caller would need to map b_T -> b_bar first. + # F_eff can run away at large b_T for λ4 < 0 toys, so we scale the TMD panel + # to the line curves rather than let a runaway band tail set the autoscale. + line_fmax = 0.0 + + for si, s in enumerate(series): + color = s.color or auto_colors[si % len(auto_colors)] + + # ---- CS panel: γ_ν^NP(b_T) (no y dependence) ---- + axL.plot( + bT, + gamma_nu_curve(bT, s.lam.gnu), + label=s.label, + color=color, + lw=s.lw, + ls=s.linestyle, + ) + if s.toys: + band = _band([gamma_nu_curve(bT, t.gnu) for t in s.toys], s.band_pct) + if band is not None: + axL.fill_between( + bT, + band[0], + band[1], + color=color, + alpha=0.22, + label=f"{s.label} {int(s.band_pct[1]-s.band_pct[0])}% band", + ) + + # ---- TMD panel: F_eff(b_T, y) per requested y ---- + n_y = len(y_values) + for yi, y in enumerate(y_values): + shade = 0.25 + 0.6 * (yi / max(n_y - 1, 1)) + yc = color if n_y == 1 else cmap_tmd(shade) + line = f_eff_curve(bT, y, s.lam.eff) + line_fmax = max(line_fmax, float(np.nanmax(line))) + axR.plot( + bT, line, color=yc, lw=s.lw, ls=s.linestyle, label=f"{s.label}, y={y:g}" + ) + if s.toys: + band = _band([f_eff_curve(bT, y, t.eff) for t in s.toys], s.band_pct) + if band is not None: + axR.fill_between(bT, band[0], band[1], color=yc, alpha=0.18) + + axL.axhline(0, color="k", lw=0.5) + axL.set_xlabel(r"$b_T$ [GeV$^{-1}$]") + axL.set_ylabel(r"$\tilde\gamma_\nu^{\rm NP}(b_T)$") + axL.set_title( + r"CS rapidity anomalous dimension $\tilde\gamma_\nu^{\rm NP}(b_T)$", fontsize=11 + ) + axL.legend(loc="lower left", fontsize=8) + axL.grid(alpha=0.3) + + axR.set_xlabel(r"$b_T$ [GeV$^{-1}$]") + axR.set_ylabel(r"$F_{\rm eff}(b_T, y)$") + axR.set_title(r"TMD-effective NP factor $F_{\rm eff}(b_T, y)$", fontsize=11) + axR.legend(loc="upper right", fontsize=8) + axR.grid(alpha=0.3) + + # Scale the TMD panel to the line curves so a runaway band tail (bare F_eff + # for λ4 < 0) doesn't dominate the autoscale. + top = f_ymax if f_ymax is not None else max(1.1, 1.2 * line_fmax) + axR.set_ylim(0.0, top) + + # Param boxes diagonally opposite each panel's legend so they don't collide: + # CS legend lower-left -> box upper-right; TMD legend upper-right -> box lower-left. + inset = inset_from if inset_from is not None else (series[-1] if series else None) + if inset is not None: + _param_inset(axL, inset.lam, "gnu", corner="upper right") + _param_inset(axR, inset.lam, "eff", corner="lower left") + + os.makedirs(os.path.dirname(os.path.abspath(outpath)) or ".", exist_ok=True) + fig.tight_layout() + fig.savefig(outpath, dpi=140) + plt.close(fig) + print(f"Wrote {outpath}") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +# Defaults for a bare CLI call: SCETlib "knobs off" plus the conventional model +# strings (override with --np-model / --np-model-nu and the --lambda* flags). +_CLI_DEFAULTS = dict( + lambda2=0.0, + lambda4=0.0, + lambda6=0.0, + delta_lambda2=0.0, + lambda_inf=1.0, + lambda2_nu=0.0, + lambda4_nu=0.0, + lambda_inf_nu=1.0, +) + + +def make_parser(): + p = argparse.ArgumentParser( + description="Plot SCETlib NP form factors γ_ν^NP(b_T) and F_eff(b_T,y).", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + src = p.add_argument_group("input (pick one mode)") + src.add_argument( + "--from-fitresult", + default=None, + help="fitresults HDF5: plot prefit (dashed) + postfit (solid + band). " + "Uses fitresult_lambdas to read λ / sample the band.", + ) + src.add_argument( + "--result", default=None, help="results group suffix for --from-fitresult." + ) + src.add_argument( + "--n-toys", type=int, default=500, help="band toys (--from-fitresult)." + ) + src.add_argument("--seed", type=int, default=0, help="band RNG seed.") + + raw = p.add_argument_group("raw λ mode (no fit)") + for k, v in _CLI_DEFAULTS.items(): + raw.add_argument(f"--{k}", type=float, default=None, help=f"default {v}") + raw.add_argument("--np-model", default="tanh_6", help="F_eff model string.") + raw.add_argument("--np-model-nu", default="tanh_2", help="γ_ν model string.") + + p.add_argument( + "--y", + type=float, + nargs="+", + default=[0.0, 2.5], + help="rapidity values for the TMD panel.", + ) + p.add_argument("--bT-max", type=float, default=4.0) + p.add_argument( + "--f-ymax", + type=float, + default=None, + help="fixed upper y-limit for the TMD panel (default: auto " + "from the line curves; clips runaway negative-λ band tails).", + ) + p.add_argument("--label", default="input") + p.add_argument("--outpath", "-o", required=True) + return p + + +def main(argv=None): + args = make_parser().parse_args(argv) + + if args.from_fitresult: + # Reading lives in the reader module; the plotter stays pure. + from wremnants.postprocessing.scetlib_np import fitresult_lambdas as frl + + series = frl.plot_series_from_fitresult( + args.from_fitresult, + result=args.result, + n_toys=args.n_toys, + seed=args.seed, + ) + else: + vals = { + k: (getattr(args, k) if getattr(args, k) is not None else v) + for k, v in _CLI_DEFAULTS.items() + } + lam = NPLambdas.from_flat(vals, args.np_model, args.np_model_nu) + series = [Series(label=args.label, lam=lam, color="C3")] + + plot_np_functions( + series, + y_values=args.y, + bT_max=args.bT_max, + outpath=args.outpath, + f_ymax=args.f_ymax, + ) + + +if __name__ == "__main__": + main() From d1a26d2446a739a4e0fffd1fb92ecd183af725d4 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 1 Jul 2026 11:40:43 -0400 Subject: [PATCH 26/31] datacard-free SigmaGen; clean up docstrings; some utility scripts fro SCETlib --- .../postprocessing/scetlib_np/__init__.py | 20 +- .../postprocessing/scetlib_np/btgrid_cache.py | 51 +- .../scetlib_np/btgrid_integrate.py | 63 +- .../postprocessing/scetlib_np/btgrid_tf.py | 247 ++-- .../scetlib_np/fitresult_lambdas.py | 77 +- .../scetlib_np/lambda_central.py | 189 ++- .../scetlib_np/np_damping_wall.py | 337 +++++ .../scetlib_np/np_function_plots.py | 182 +-- .../postprocessing/scetlib_np/param_model.py | 1111 ++++++----------- .../scetlib_np/param_model_diagnostics.py | 525 ++++++++ wremnants/postprocessing/scetlib_np/params.py | 122 ++ .../scetlib_np/response_matrix.py | 150 ++- .../postprocessing/scetlib_np/sigma_gen.py | 576 +++++++++ .../scetlib_np/sigma_gen_at_lambda.py | 661 ++++++++++ wremnants/utilities/styles/styles.py | 4 +- 15 files changed, 3137 insertions(+), 1178 deletions(-) create mode 100644 wremnants/postprocessing/scetlib_np/np_damping_wall.py create mode 100644 wremnants/postprocessing/scetlib_np/param_model_diagnostics.py create mode 100644 wremnants/postprocessing/scetlib_np/params.py create mode 100644 wremnants/postprocessing/scetlib_np/sigma_gen.py create mode 100644 wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py diff --git a/wremnants/postprocessing/scetlib_np/__init__.py b/wremnants/postprocessing/scetlib_np/__init__.py index b12044044..6a060bab0 100644 --- a/wremnants/postprocessing/scetlib_np/__init__.py +++ b/wremnants/postprocessing/scetlib_np/__init__.py @@ -1,14 +1,16 @@ """SCETlib-NP postprocessing package. -``SCETlibNPParamModel`` (and its TensorFlow / btgrid dependencies) is imported -lazily so that lightweight submodules — e.g. :mod:`response_matrix`, used by -setupRabbit to embed the response matrix in the datacard — can be imported -without pulling in TensorFlow. The package-level re-export -``wremnants.postprocessing.scetlib_np.SCETlibNPParamModel`` (used by rabbit's -``--paramModel`` loader) still works, resolved on first access via PEP 562. +``SCETlibNPParamModel`` (rabbit adapter) and ``SigmaGenModel`` (datacard-free +σ_gen physics core), with their TensorFlow / btgrid dependencies, are imported +lazily so lightweight submodules (e.g. :mod:`response_matrix`, used by setupRabbit +to embed the response matrix in the datacard) import without pulling in +TensorFlow. The package-level re-exports +``wremnants.postprocessing.scetlib_np.SCETlibNPParamModel`` (rabbit's +``--paramModel`` loader) and ``…​.SigmaGenModel`` still work, resolved on first +access via PEP 562. """ -__all__ = ["SCETlibNPParamModel"] +__all__ = ["SCETlibNPParamModel", "SigmaGenModel"] def __getattr__(name): @@ -18,4 +20,8 @@ def __getattr__(name): ) return SCETlibNPParamModel + if name == "SigmaGenModel": + from wremnants.postprocessing.scetlib_np.sigma_gen import SigmaGenModel + + return SigmaGenModel raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/wremnants/postprocessing/scetlib_np/btgrid_cache.py b/wremnants/postprocessing/scetlib_np/btgrid_cache.py index e220e0235..800ebc35b 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_cache.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_cache.py @@ -1,8 +1,8 @@ """One-shot pickle cache for the combined SCETlib bT-grid. -Assembling the bT-grid from its individual shards is slow. On the first call -this module writes a single ``combined_btgrid.pkl`` in the btgrid directory; -subsequent calls load that combined file directly, which is much faster. +Assembling the bT-grid from its shards is slow. The first call writes a single +``combined_btgrid.pkl`` in the btgrid directory; later calls load it directly, +much faster. Usage: from wremnants.postprocessing.scetlib_np import btgrid_cache @@ -23,20 +23,20 @@ def load_btgrid_shards(submitdir_or_glob, runcard_basename=None): """Combine bT-grid shards produced by --bt-grid mode. `submitdir_or_glob` may be: - - a directory: we look for ``*_btgrid.pkl`` inside (recursively only - one level via scetlib_outputs/). + - a directory: looks for ``*_btgrid.pkl`` inside (one level deep via + scetlib_outputs/). - a glob pattern: used directly. Returns dict with: bT : (Nbt,) b_bar : (Nbt,) bins : list of (Q, Y, qT, lep) bin centres, length Nbins - vars : dict of variation index -> setting dict (copied from the - first shard; all shards are expected to carry the same set) + vars : dict variation index -> setting dict (from the first shard; + all shards expected to carry the same set) I_pert : (Nvars, Nbins, Nbt) C_nu : (Nvars, Nbins, Nbt) - config : dict from the first shard (perturbative configuration the - grid was generated against) + config : dict from the first shard (perturbative config the grid was + generated against) n_shards: int """ if os.path.isdir(submitdir_or_glob): @@ -69,8 +69,8 @@ def load_btgrid_shards(submitdir_or_glob, runcard_basename=None): n_vars = len(varis) n_bt = bT.size - # We don't know Nbins ahead of time without scanning all shards. Walk them - # once: build a dict of bin -> (var_idx -> (I_pert_row, C_nu_row)). + # Nbins is unknown without scanning all shards. Walk once: dict of + # bin -> (var_idx -> (I_pert_row, C_nu_row)). bin_to_data = {} for path in files: with open(path, "rb") as f: @@ -86,9 +86,8 @@ def load_btgrid_shards(submitdir_or_glob, runcard_basename=None): d["I_pert"], dtype=float ) # (Nvars_local, Nbins_local, Nbt) C_local = np.asarray(d["C_nu"], dtype=float) - # Map local variation indices to the union variation order. We assume - # all shards share the same vars dict (true when produced by the same - # condor submission). + # Map local variation indices to the union order. Assumes all shards + # share the same vars dict (true within one condor submission). var_order_local = list(d["vars"].keys()) for b_idx, b_tup in enumerate(bins_local): slot = bin_to_data.setdefault(tuple(b_tup), {}) @@ -147,12 +146,10 @@ def _cache_is_fresh(combined, shards): def load(submitdir, rebuild=False, verbose=True): """Load the combined bT-grid for ``submitdir``. - On first call (or when ``rebuild=True``, or when any shard is newer than - the cached combined file), assembles the shards via - :func:`load_btgrid_shards`, writes ``combined_btgrid.pkl``, and returns - the dict. - - On subsequent calls, loads the pickle directly. + On first call (or ``rebuild=True``, or any shard newer than the cached + combined file), assembles the shards via :func:`load_btgrid_shards`, writes + ``combined_btgrid.pkl``, returns the dict. Otherwise loads the pickle + directly. """ if not os.path.isdir(submitdir): raise ValueError(f"{submitdir!r} is not a directory") @@ -187,3 +184,17 @@ def load(submitdir, rebuild=False, verbose=True): pickle.dump(grid, f, protocol=pickle.HIGHEST_PROTOCOL) os.replace(tmp, combined) return grid + + +def combined_path(submitdir): + """Path to the combined bT-grid pickle for ``submitdir`` (may not exist yet).""" + return _combined_path(submitdir) + + +def is_combined_fresh(submitdir): + """True if ``combined_btgrid.pkl`` exists and is at least as new as every + shard, i.e. :func:`load` would read it directly rather than reassemble. + + Exposed so a derived cache (e.g. the factorized layout in :mod:`sigma_gen`) + can key its own freshness on the combined pickle without loading it.""" + return _cache_is_fresh(_combined_path(submitdir), _shard_glob(submitdir)) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py index ed7fcfaec..a62bec23f 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_integrate.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py @@ -1,22 +1,20 @@ """Q integration and Y/qT rebin helpers for the SCETlib bT-grid ParamModel. -All weight-construction is numpy (runs once at construction time); runtime -contractions are simple ``tf.tensordot`` / ``tf.einsum`` calls. +Weight construction is numpy (once at construction time); runtime contractions +are ``tf.tensordot`` / ``tf.einsum``. -Three pieces: +1. :func:`dense_index_map` — ``(NQ, NY, NqT)`` int array mapping each + rectangular grid cell to a flat sparse-btgrid bin index; ``-1`` marks + missing combos. ``tf.gather`` with a sentinel pads a sparse ``(Nbins,)`` σ + into a dense ``(NQ, NY, NqT)``. -1. :func:`dense_index_map` — build a ``(NQ, NY, NqT)`` int array mapping each - rectangular grid cell to a flat bin index in the sparse btgrid; ``-1`` - marks missing combos. Use ``tf.gather`` with a sentinel to pad a sparse - ``(Nbins,)`` σ tensor into a dense ``(NQ, NY, NqT)``. - -2. :func:`q_integrate_weights` — produces a ``(NQ,)`` weight vector for - ``arctan_Q²``-method Simpson integration over the Z mass window. Apply - via ``tf.einsum('q, qyz -> yz', w, sigma)``. +2. :func:`q_integrate_weights` — ``(NQ,)`` weight vector for ``arctan_Q²`` + Simpson integration over the Z mass window. Apply via + ``tf.einsum('q, qyz -> yz', w, sigma)``. 3. :func:`rebin_weights` — given a fine source grid and a coarser target-bin - edge list, builds a ``(N_target, N_source)`` Simpson weight matrix. - Apply via ``tf.tensordot``. + edge list, a ``(N_target, N_source)`` Simpson weight matrix. Apply via + ``tf.tensordot``. """ import numpy as np @@ -28,9 +26,9 @@ ) from wremnants.utilities import common as wrem_common -# Z resonance parameters for the Q-integration change of variable, in the -# s-dependent-width scheme (see wremnants.utilities.common). Only set the centre -# and scale of the arctan-Q^2 transform below; they do not change the physics. +# Z resonance parameters for the Q-integration change of variable, +# s-dependent-width scheme (see wremnants.utilities.common). Set only the centre +# and scale of the arctan-Q^2 transform; do not change the physics. MZ_S_DEP_WIDTH = wrem_common.MZ_S_DEP_WIDTH GAMMAZ_S_DEP_WIDTH = wrem_common.GAMMAZ_S_DEP_WIDTH @@ -46,8 +44,8 @@ def dense_index_map(bins, Q_unique=None, Y_unique=None, qT_unique=None): Parameters ---------- bins : list of tuples - From ``load_btgrid_shards``: each element is ``(Q, Y, qT, lep)``, - sorted lexicographically. + From ``load_btgrid_shards``: each ``(Q, Y, qT, lep)``, sorted + lexicographically. Returns ------- @@ -90,11 +88,11 @@ def dense_index_map(bins, Q_unique=None, Y_unique=None, qT_unique=None): def sparse_to_dense_tf(sigma_flat, flat_idx): """Reshape a sparse ``(Nbins,)`` σ tensor to dense ``(NQ, NY, NqT)``. - Missing cells (``flat_idx == -1``) are padded with 0. Implemented via - ``tf.gather`` with a 0-padded sentinel row. + Missing cells (``flat_idx == -1``) padded with 0, via ``tf.gather`` with a + 0-padded sentinel row. """ sigma_flat = _as_dtype(sigma_flat) - # Append one extra "zero" entry that the -1 indices will gather. + # Append one "zero" entry that the -1 indices gather. extended = tf.concat([sigma_flat, tf.zeros([1], dtype=sigma_flat.dtype)], axis=0) sentinel = tf.cast(tf.shape(sigma_flat)[0], tf.int64) # index of the appended zero idx_safe = tf.where(tf.equal(flat_idx, -1), sentinel, flat_idx) @@ -111,12 +109,12 @@ def q_integrate_weights( ): """Simpson weights for integrating over Q ∈ [Q_lo, Q_hi] in arctan-Q² space. - Implements the same change of variable as the external numpy reference + Same change of variable as the numpy reference ``scetlib_run.factorize.integrate_over_Q`` with ``method="arctan_Q2"``: - x = arctan((Q² - q0²) / (q0 Γ)) flattens the Breit-Wigner peak, then - Simpson on x with the Jacobian dQ/dx. + x = arctan((Q² - q0²) / (q0 Γ)) flattens the Breit-Wigner peak, then Simpson + on x with the Jacobian dQ/dx. - Returns a ``(NQ,)`` weight vector with zeros outside ``[Q_lo, Q_hi]``. + Returns a ``(NQ,)`` weight vector, zero outside ``[Q_lo, Q_hi]``. """ Q_grid = np.asarray(Q_grid, dtype=np.float64) mask = (Q_grid >= Q_lo) & (Q_grid <= Q_hi) @@ -149,13 +147,12 @@ def integrate_over_Q_tf(sigma_QYqT, Q_weights): def rebin_weights(source_grid, target_edges, name="axis", tol=1e-9): """Build a ``(N_target, N_source)`` Simpson rebin matrix. - For each target bin ``[target_edges[i], target_edges[i+1]]``, find the - source samples that fall in (with ``tol`` slack) the bin's interior + - edges, then compute Simpson weights for those samples. Returns a dense - matrix; entries are 0 for source samples not contributing to a given - target bin. + For each target bin ``[target_edges[i], target_edges[i+1]]``, find the source + samples within ``tol`` slack of the bin's interior + edges and compute their + Simpson weights. Dense matrix; entries are 0 for source samples not + contributing to a target bin. - Mirrors the per-bin call pattern of the external numpy reference + Mirrors the per-bin call pattern of the numpy reference ``scetlib_run.factorize`` (``integrate_over_{qT,Y}_bin``). """ source_grid = np.asarray(source_grid, dtype=np.float64) @@ -192,8 +189,8 @@ def rebin_axis_tf(values, axis, weights): rank = len(values.shape) if axis < 0: axis += rank - # tensordot contracts values[axis] with weights[1]; target axis ends up - # at the END of the result. Permute it back to position ``axis``. + # tensordot contracts values[axis] with weights[1]; target axis lands at the + # END of the result. Permute it back to position ``axis``. out = tf.tensordot(values, weights, axes=[[axis], [1]]) perm = list(range(rank - 1)) perm.insert(axis, rank - 1) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_tf.py b/wremnants/postprocessing/scetlib_np/btgrid_tf.py index 1a3ff88fb..cb7aac502 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_tf.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_tf.py @@ -1,20 +1,17 @@ """TensorFlow bT-grid factorization library. -Differentiable TF implementation of the SCETlib bT-space form factors and Hankel -reconstruction (transcribed from the SCETlib C++). - -Design choices: - * ``np_model`` / ``np_model_nu`` strings are fixed at trace time (the SCETlib - runcard sets them once per fit). The TF functions dispatch on the string at - Python level — no ``tf.cond``. - * λ parameters are TF tensors (typically scalars, but broadcasting follows - the usual array rules). - * All ops are differentiable in λ. Branches on λ values use ``tf.where`` - with a safe denominator to avoid NaN gradients. - * ``b_star_global`` is not ported — the cached ``b_bar`` array in the bT-grid - shards is precomputed and travels as a ``tf.constant``. - * Simpson weights are precomputed at trace time from the (static) bT, Y, qT - grids; the runtime cost is just ``tf.reduce_sum(w * y)``. +Differentiable TF transcription of the SCETlib bT-space form factors and Hankel +reconstruction (from the SCETlib C++). + + * ``np_model`` / ``np_model_nu`` fixed at trace time (set once per fit by the + runcard); functions dispatch on the string at Python level, not ``tf.cond``. + * λ parameters are TF tensors (typically scalars; usual broadcasting). + * All ops differentiable in λ; branches on λ use ``tf.where`` with a safe + denominator to avoid NaN gradients. + * ``b_star_global`` not ported: the precomputed ``b_bar`` shard array travels + as a ``tf.constant``. + * Simpson weights precomputed at trace time from the static bT, Y, qT grids; + runtime cost is ``tf.reduce_sum(w * y)``. """ from typing import Mapping @@ -22,16 +19,16 @@ import numpy as np import tensorflow as tf -# Set the dtype used for all ops in this module (float64 throughout). +# float64 throughout this module. DTYPE = tf.float64 def _as_dtype(x, dtype=DTYPE): """Coerce ``x`` to ``dtype`` without losing precision on Python scalars. - ``tf.cast(0.4, tf.float64)`` round-trips through float32 (returning - ``0.4000000059604645``); ``tf.constant(0.4, dtype=tf.float64)`` does not. - Use this helper everywhere a possibly-Python-float input enters the graph. + ``tf.cast(0.4, tf.float64)`` round-trips through float32 + (``0.4000000059604645``); ``tf.constant(0.4, dtype=tf.float64)`` does not. + Use wherever a possibly-Python-float input enters the graph. """ if isinstance(x, (int, float)): return tf.constant(x, dtype=dtype) @@ -44,10 +41,10 @@ def _as_dtype(x, dtype=DTYPE): def simpson_weights(x): - """Return weights ``w`` such that Simpson(y, x) == sum(w * y, axis=-1). + """Weights ``w`` such that Simpson(y, x) == sum(w * y, axis=-1). - ``x`` is a numpy array with size ``N``. Composite Simpson with a trapezoid - fallback on the last segment when N-1 is odd. + ``x`` is a size-``N`` numpy array. Composite Simpson, trapezoid fallback on + the last segment when N-1 is odd. """ x = np.asarray(x, dtype=np.float64) n_intervals = x.size - 1 @@ -55,7 +52,7 @@ def simpson_weights(x): return np.zeros_like(x) if n_intervals % 2 == 1: - # leading n-1 intervals get Simpson, last segment gets trapezoid + # leading n-1 intervals Simpson, last segment trapezoid w_lead = simpson_weights(x[:-1]) w = np.concatenate([w_lead, [0.0]]) h_last = x[-1] - x[-2] @@ -115,29 +112,28 @@ def simpson_tf(y, weights): def _frozen_eq_zero(x): - """``x == 0`` as a gradient-frozen boolean condition for the NP-factor masks. + """``x == 0`` as a gradient-frozen condition for the NP-factor masks. - The comparison is a non-differentiable, measure-zero boundary that the + The comparison is a non-differentiable, measure-zero boundary the surrounding ``tf.where`` never differentiates through, so freezing its input - leaves both the value and every derivative unchanged. But it is REQUIRED for - the full-K Hessian: the straight-through ``K`` path nests two - ``ForwardAccumulator``s (forward-over-forward AD), and building the JVP of an - ``Equal`` op that receives a tangent-carrying input raises - ``IndexError: list index out of range`` under ``@tf.function`` (a TF - nested-forward-mode bug). With the input frozen no tangent ever reaches the - comparison, so the bug never triggers; the ``tf.where``'s own JVP with a - constant condition is fine. (The earlier GN/J-only path never hit this — it - uses a single ``ForwardAccumulator``, for which ``Equal``'s JVP is fine.)""" + changes no value or derivative. REQUIRED for the full-K Hessian: the + straight-through ``K`` path nests two ``ForwardAccumulator``s + (forward-over-forward AD), and the JVP of an ``Equal`` op fed a + tangent-carrying input raises ``IndexError: list index out of range`` under + ``@tf.function`` (a TF nested-forward-mode bug). Frozen input → no tangent + reaches the comparison; the ``tf.where`` JVP with a constant condition is + fine. (The GN/J-only path uses one ``ForwardAccumulator``, where ``Equal``'s + JVP is fine.)""" return tf.equal(tf.stop_gradient(x), 0) def _safe_div(num, den): - """``num / den`` with the denominator clamped to 1 where it is exactly zero. + """``num / den`` with the denominator clamped to 1 where exactly zero. - Equivalent to ``num / tf.where(den == 0, 1, den)``; the comparison input is - frozen (see :func:`_frozen_eq_zero`) so the full-K nested forward-mode - Hessian does not crash under ``@tf.function``. Keeps gradients finite; the - den==0 result is masked away by the caller's final ``tf.where``.""" + Equivalent to ``num / tf.where(den == 0, 1, den)`` with a frozen comparison + input (see :func:`_frozen_eq_zero`) so the full-K nested forward-mode Hessian + doesn't crash under ``@tf.function``. Gradients stay finite; the den==0 + result is masked away by the caller's final ``tf.where``.""" den_safe = tf.where(_frozen_eq_zero(den), tf.ones_like(den), den) return num / den_safe @@ -165,8 +161,8 @@ def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_ if np_model == "identity": return tf.exp(-2.0 * bT * arg) - # lambda_inf == 0 returns ones. We compute the full formula with a safe - # denominator and mask at the end. + # lambda_inf == 0 returns ones: compute the full formula with a safe + # denominator, mask at the end. arg_inf = _safe_div(arg, lambda_inf) model = {"hyp_tangent": "tanh_2", "square_root": "frac_2"}.get(np_model, np_model) @@ -193,13 +189,21 @@ def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_ raise ValueError(f"F_eff_tf: unsupported np_model {np_model!r}") full = tf.exp(-2.0 * lambda_inf * bT * func) - # lambda_inf == 0 -> 1 (NP off). Frozen comparison input so the full-K nested - # forward-mode Hessian doesn't crash under @tf.function (see _frozen_eq_zero). + # lambda_inf == 0 -> 1 (NP off); frozen comparison input (see + # _frozen_eq_zero) for the full-K @tf.function Hessian. return tf.where(_frozen_eq_zero(lambda_inf), tf.ones_like(full), full) -def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): - """CS-side NP rapidity anomalous dimension γ_ν^NP(bT) for fixed ``np_model_nu``.""" +def gamma_nu_NP_tf( + bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, lambda6_nu=0.0, np_model_nu +): + """CS-side NP rapidity anomalous dimension γ_ν^NP(bT) for fixed ``np_model_nu``. + + ``lambda6_nu`` is the b⁶ coefficient used only by the ``tanh_6`` model; other + models ignore it. It defaults to 0 (then tanh_6 reduces to tanh_2). SCETlib's + own ``NP_model_gammanu`` uses 0.0007 (Gamma_nu.hpp:102) — pass that to + reproduce SCETlib's regulated CS kernel. + """ if np_model_nu not in GNU_MODELS: raise ValueError(f"gamma_nu_NP_tf: unsupported np_model_nu {np_model_nu!r}") @@ -207,6 +211,7 @@ def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): lambda_inf_nu = _as_dtype(lambda_inf_nu) lambda2_nu = _as_dtype(lambda2_nu) lambda4_nu = _as_dtype(lambda4_nu) + lambda6_nu = _as_dtype(lambda6_nu) bT2 = bT * bT arg = _safe_div((lambda2_nu + lambda4_nu * bT2) * bT2, lambda_inf_nu) @@ -219,8 +224,9 @@ def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): elif model == "tanh_2": func = tf.tanh(arg) elif model == "tanh_6": - # NP_model_gammanu hardcodes lambda6_nu = 0.0007 (Gamma_nu.hpp:102) - a = arg + _safe_div(0.0007 * bT2**3, lambda_inf_nu) + # b⁶ term (SCETlib NP_model_gammanu uses lambda6_nu = 0.0007, + # Gamma_nu.hpp:102); here it is the fittable lambda6_nu (default 0). + a = arg + _safe_div(lambda6_nu * bT2**3, lambda_inf_nu) func = tf.tanh(a) elif model == "frac_1": a = arg + _safe_div(lambda2_nu * bT2, lambda_inf_nu) ** 2 @@ -236,8 +242,8 @@ def gamma_nu_NP_tf(bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu): raise ValueError(f"gamma_nu_NP_tf: unsupported np_model_nu {np_model_nu!r}") full = -lambda_inf_nu * func - # lambda_inf_nu == 0 -> 0 (NP off). Frozen comparison input so the full-K nested - # forward-mode Hessian doesn't crash under @tf.function (see _frozen_eq_zero). + # lambda_inf_nu == 0 -> 0 (NP off); frozen comparison input (see + # _frozen_eq_zero) for the full-K @tf.function Hessian. return tf.where(_frozen_eq_zero(lambda_inf_nu), tf.zeros_like(full), full) @@ -266,25 +272,23 @@ def reconstruct_batch_tf( """Reconstruct σ on a batch of (Q, Y, qT) grid points from the bT integrand. Implements the per-(Q, Y, qT) bT-space integrand. The full formula with - every factor and its bare-bT / b*(bT) / (Q,Y,qT) / λ dependence is written - out in the :mod:`param_model` module docstring — consult it rather than - re-deriving the factors from this code. + every factor and its bare-bT / b*(bT) / (Q,Y,qT) / λ dependence is in the + :mod:`param_model` module docstring. - All array-shape arguments are TF tensors or numpy arrays (will be cast). - The λ values inside ``eff_params`` / ``gnu_params`` are the differentiable - parameters — pass them as TF scalars (Variables or constants). + Array-shape arguments are TF tensors or numpy arrays (cast on entry). The λ + values in ``eff_params`` / ``gnu_params`` are the differentiable parameters; + pass them as TF scalars (Variables or constants). ``bT_simpson_weights`` and ``bT_J0_kernel`` are optional precomputed - constants. Pass them in the ParamModel to avoid recomputing per-step. - - ``Y_unique`` / ``Y_inverse_idx`` are an optional precomputed unique-Y map - (``Y_unique`` = sorted distinct Y values, shape ``(NY,)``; ``Y_inverse_idx`` - = per-bin index into ``Y_unique``, shape ``(Nbins,)``). ``F_eff`` depends on - the bin only through Y, so when this map is supplied the NP transcendentals - are evaluated on the ``NY`` unique rows and gathered back to ``(Nbins, Nbt)`` - — bit-for-bit identical to the per-bin evaluation, but the expensive ops and - their λ-gradients run on ``NY`` rows instead of ``Nbins`` (Q and qT don't - enter ``F_eff``). Without the map it falls back to the full per-bin path. + constants; pass them from the ParamModel to avoid recomputing per-step. + + ``Y_unique`` / ``Y_inverse_idx`` are an optional unique-Y map (``Y_unique`` = + sorted distinct Y, shape ``(NY,)``; ``Y_inverse_idx`` = per-bin index into + it, shape ``(Nbins,)``). ``F_eff`` depends on the bin only through Y, so the + NP transcendentals run on the ``NY`` unique rows and gather back to + ``(Nbins, Nbt)`` — bit-for-bit identical to per-bin, but with the expensive + ops and their λ-gradients on ``NY`` rows not ``Nbins`` (Q, qT don't enter + ``F_eff``). Without the map, falls back to the full per-bin path. """ qT_per_bin = _as_dtype(qT_per_bin) # (Nbins,) Y_per_bin = _as_dtype(Y_per_bin) # (Nbins,) @@ -298,7 +302,7 @@ def reconstruct_batch_tf( bT_J0_kernel = _as_dtype(bT_J0_kernel) # (Nbins, Nbt) if bT_simpson_weights is None: - # bT is a tf.Tensor here; convert to numpy for the Python-side weights + # bT is a tf.Tensor; to numpy for the Python-side weights bT_simpson_weights = simpson_weights(np.asarray(bT)) bT_simpson_weights = _as_dtype(bT_simpson_weights) # (Nbt,) @@ -311,17 +315,17 @@ def reconstruct_batch_tf( Feff = F_eff_tf(0.0, b_bar, **eff_params, np_model=np_model) # (Nbt,) Feff_b = Feff[tf.newaxis, :] elif Y_unique is not None and Y_inverse_idx is not None: - # per-bin F_eff, but evaluated on the unique Y rows then gathered. - # Exact: identical Y -> identical F_eff row for any λ; the gather only - # replicates rows (its backward scatter-adds the cotangents, so λ-grads - # are unchanged). Transcendentals run on (NY, Nbt), not (Nbins, Nbt). + # per-bin F_eff on the unique-Y rows, then gathered. Exact: identical Y + # -> identical F_eff row for any λ; the gather only replicates rows + # (backward scatter-adds cotangents, λ-grads unchanged). Transcendentals + # run on (NY, Nbt), not (Nbins, Nbt). Y_u = _as_dtype(Y_unique)[:, tf.newaxis] # (NY, 1) b_b = b_bar[tf.newaxis, :] Feff_u = F_eff_tf(Y_u, b_b, **eff_params, np_model=np_model) # (NY, Nbt) Feff_b = tf.gather(Feff_u, Y_inverse_idx) # (Nbins, Nbt) else: - # per-bin F_eff due to Y dependence (delta_lambda2 * Y^2) - # build (Nbins, Nbt) by broadcasting Y_per_bin over bT + # per-bin F_eff (Y dependence via delta_lambda2 * Y^2): build + # (Nbins, Nbt) by broadcasting Y_per_bin over bT Y_b = Y_per_bin[:, tf.newaxis] b_b = b_bar[tf.newaxis, :] Feff_b = F_eff_tf(Y_b, b_b, **eff_params, np_model=np_model) # (Nbins, Nbt) @@ -336,13 +340,12 @@ def reconstruct_batch_tf( def build_bT_J0_kernel(qT_per_bin, bT): """Precompute ``bT * J_0(qT*bT)`` on the (Nbins, Nbt) grid. - λ-independent — call once at ParamModel construction and pass into + λ-independent: call once at ParamModel construction, pass into :func:`reconstruct_batch_tf` as ``bT_J0_kernel``. - Note: in the factorized path (:func:`reconstruct_batch_factorized_tf`) - this is instead called with ``qT_unique`` (NqT distinct values, not the - per-bin expansion), giving a (NqT, Nbt) kernel — same numbers, ~4000× - smaller. + The factorized path (:func:`reconstruct_batch_factorized_tf`) instead calls + it with ``qT_unique`` (NqT distinct values, not the per-bin expansion), + giving a (NqT, Nbt) kernel: same numbers, ~4000× smaller. """ qT_per_bin = _as_dtype(qT_per_bin) bT = _as_dtype(bT) @@ -354,39 +357,39 @@ def build_bT_J0_kernel(qT_per_bin, bT): # Factorized reconstruction (GPU-memory-safe; exact) # ============================================================================= # -# The (Nbins, Nbt) layout of reconstruct_batch_tf needs several ~9 GB fp64 -# tensors (Nbins=546840, Nbt=2000) and OOMs a 32 GB GPU at construction. Two -# exact observations shrink it: +# reconstruct_batch_tf's (Nbins, Nbt) layout needs several ~9 GB fp64 tensors +# (Nbins=546840, Nbt=2000) and OOMs a 32 GB GPU at construction. Two exact +# observations shrink it: # -# 1. qT enters the λ-dependent integrand ONLY through the bT·J0(qT·bT) -# kernel: the kernel needs the NqT *unique* qT values, not Nbins rows. -# 2. SCETlib's profile scales are piecewise in x = qT/Q and exactly -# canonical (qT-independent) below the first transition point x1·Q, so -# the cached I_pert / C_nu rows are BIT-IDENTICAL across qT there. The -# dedup below discovers identical rows dynamically (byte-wise hashing + -# full verification) — no assumption about profiles or qT ranges. +# 1. qT enters the λ-dependent integrand ONLY via the bT·J0(qT·bT) kernel, +# which needs the NqT *unique* qT values, not Nbins rows. +# 2. SCETlib's profile scales are piecewise in x = qT/Q and exactly canonical +# (qT-independent) below the first transition x1·Q, so the cached +# I_pert / C_nu rows are BIT-IDENTICAL across qT there. The dedup below +# discovers identical rows dynamically (byte-wise hashing + full +# verification) — no assumption about profiles or qT ranges. # -# The bT-Simpson reduction then becomes a (Nu, Nbt) @ (Nbt, NqT) matmul plus -# a per-bin gather. Same integrand, same Simpson weights, same sampling; only -# the floating-point grouping/summation order changes (≲1e-14 relative). +# The bT-Simpson reduction then becomes a (Nu, Nbt) @ (Nbt, NqT) matmul plus a +# per-bin gather. Same integrand, weights and sampling; only the floating-point +# grouping/summation order changes (≲1e-14 relative). def dedup_grid_rows(I_pert, C_nu, feff_idx_per_bin, verbose=True): """Find bit-identical (I_pert, C_nu, F_eff-index) row triples. - Construction-time numpy helper (runs once, CPU). Rows are keyed by the - raw bytes of the I_pert row, the C_nu row and the per-bin F_eff Y-index, - so two bins share a unique id iff their λ-dependent integrand columns are - bit-for-bit identical for EVERY λ. The grouping is verified by direct - array comparison afterwards (no reliance on hash collision odds). + Construction-time numpy helper (runs once, CPU). Rows are keyed by the raw + bytes of the I_pert row, the C_nu row and the per-bin F_eff Y-index, so two + bins share a unique id iff their λ-dependent integrand columns are + bit-for-bit identical for EVERY λ. Grouping verified by direct array + comparison afterward, not hash-collision odds. Parameters ---------- I_pert, C_nu : (Nbins, Nbt) float64 ndarrays feff_idx_per_bin : (Nbins,) int ndarray - Index into the unique-Y table used for the F_eff gather (Y enters - F_eff via delta_lambda2·Y²; keying on it keeps the per-unique-row - F_eff well-defined even when delta_lambda2 floats). + Index into the unique-Y table for the F_eff gather (Y enters F_eff via + delta_lambda2·Y²; keying on it keeps per-unique-row F_eff well-defined + even when delta_lambda2 floats). Returns ------- @@ -395,14 +398,12 @@ def dedup_grid_rows(I_pert, C_nu, feff_idx_per_bin, verbose=True): row_uid : (Nbins,) int32, bin -> unique-row index feff_idx_u : (Nu,) int32, unique row -> unique-Y index n_unique : int - C_uu : (Ncu, Nbt) second-level dedup of the C_u rows. C_nu is the - ν-evolution coefficient — it depends on (Q, profile-qT) - only, not Y, so its standalone unique-row count is ~150x - smaller than Nu (1888 vs 284605 on the fineall grid). The - exp(C·g) transcendentals run on these rows and are - gathered back — bit-identical, ~150x fewer exp() calls, - and the (Nu, Nbt) C constant never needs to exist on - device. + C_uu : (Ncu, Nbt) second-level dedup of C_u. C_nu depends on + (Q, profile-qT) only, not Y, so its standalone unique-row + count is ~150x smaller than Nu (1888 vs 284605 on fineall). + The exp(C·g) transcendentals run on these rows, gathered + back: bit-identical, ~150x fewer exp() calls, and the + (Nu, Nbt) C constant never exists on device. c_of_u : (Nu,) int32, unique row -> C_uu row index n_unique_C : int """ @@ -435,8 +436,8 @@ def dedup_grid_rows(I_pert, C_nu, feff_idx_per_bin, verbose=True): C_u = np.ascontiguousarray(C_nu[rep_rows]) feff_idx_u = feff_idx_per_bin[rep_rows].astype(np.int32) - # Verify the grouping bit-exactly: every bin's rows must equal its - # representative's. Chunked to bound the temporary gather copies. + # Verify bit-exactly: every bin's rows must equal its representative's. + # Chunked to bound the temporary gather copies. chunk = 20000 for k0 in range(0, n_bins, chunk): k1 = min(k0 + chunk, n_bins) @@ -451,8 +452,8 @@ def dedup_grid_rows(I_pert, C_nu, feff_idx_per_bin, verbose=True): f"[{k0}, {k1}) — this should be impossible; grid corrupt?" ) - # Second-level dedup of the C rows (qT-and-Y-independent below the - # profile transition AND Y-independent everywhere → ~150x smaller). + # Second-level dedup of the C rows (qT-independent below the profile + # transition AND Y-independent everywhere → ~150x smaller). n_u = len(rep_rows) seen_c = {} c_of_u = np.empty(n_u, dtype=np.int32) @@ -517,10 +518,10 @@ def reconstruct_batch_factorized_tf( Evaluates σ_i = qT_i Σ_b w_b·bT_b·J0(qT_i bT_b)·I_{u(i),b}·exp(C_{u(i),b} g_b)·F_{y(u(i)),b} as a (Nu, Nbt) elementwise block, a matmul against the - weighted J0 kernel on the unique-qT grid, and a per-bin gather — no - (Nbins, Nbt) tensor is ever materialized. Same integrand, weights and - sampling as :func:`reconstruct_batch_tf`; only the floating-point - multiplication grouping and summation order differ (≲1e-14 relative). + weighted J0 kernel on the unique-qT grid, and a per-bin gather; no + (Nbins, Nbt) tensor is materialized. Same integrand, weights and sampling as + :func:`reconstruct_batch_tf`; only the floating-point multiplication grouping + and summation order differ (≲1e-14 relative). Parameters ---------- @@ -528,13 +529,13 @@ def reconstruct_batch_factorized_tf( I_pert_u : (Nu, Nbt) — deduplicated grid rows (from :func:`dedup_grid_rows`) C_nu_u : (Nu, Nbt), optional - Per-unique-row C_ν. Either this OR (``C_nu_uu``, ``c_of_u``) must be - given; the latter is preferred (~150x fewer exp() calls and no - (Nu, Nbt) C constant on device — bit-identical results, since exp of - identical rows is identical and the gather only replicates rows). + Per-unique-row C_ν. Pass this OR (``C_nu_uu``, ``c_of_u``); the latter + is preferred (~150x fewer exp() calls, no (Nu, Nbt) C constant on + device). Bit-identical: exp of identical rows is identical, gather only + replicates rows. KwqT : (NqT, Nbt) - ``qT_u · bT · J0(qT_u·bT) · w_simpson`` on the unique-qT grid; the - per-bin qT prefactor of reconstruct_batch_tf is folded in here. + ``qT_u · bT · J0(qT_u·bT) · w_simpson`` on the unique-qT grid; folds in + reconstruct_batch_tf's per-bin qT prefactor. gather_idx : (Nbins, 2) int32 — per bin ``[u(i), qT_index(i)]`` Y_unique : (NY,) — unique Y values for the F_eff evaluation feff_idx_u : (Nu,) int32 — unique row -> Y_unique index @@ -565,9 +566,9 @@ def reconstruct_batch_factorized_tf( Feff = F_eff_tf(0.0, b_bar, **eff_params, np_model=np_model) # (Nbt,) Feff_u = Feff[tf.newaxis, :] # broadcast over Nu else: - # F_eff on the unique-Y rows, gathered to the unique grid rows. Same - # unique-Y trick as reconstruct_batch_tf — the gather replicates rows - # bit-exactly and scatter-adds cotangents in the backward pass. + # F_eff on the unique-Y rows, gathered to the unique grid rows (same + # unique-Y trick as reconstruct_batch_tf: gather replicates rows + # bit-exactly, scatter-adds cotangents on the backward pass). Y_u = _as_dtype(Y_unique)[:, tf.newaxis] # (NY, 1) Feff_rows = F_eff_tf( Y_u, b_bar[tf.newaxis, :], **eff_params, np_model=np_model diff --git a/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py b/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py index 250904ed7..74487e5e4 100644 --- a/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py +++ b/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py @@ -1,30 +1,28 @@ """Read SCETlib NP λ out of a rabbit fitresults HDF5 (table / curves / toys). -The OUTPUT-side companion to :mod:`lambda_central` (which reads the INPUT-side +OUTPUT-side companion to :mod:`lambda_central` (which reads the INPUT-side λ_central from the upstream correction pkl). Two responsibilities, kept apart -from the plotting: +from plotting: - * Tabulate the λ — prefit/postfit value, prefit/postfit 1σ constraint, and - whether each was frozen — and print it (``read_lambdas`` + the CLI here). + * Tabulate the λ (prefit/postfit value, prefit/postfit 1σ constraint, frozen + flag) and print it (``read_lambdas`` + the CLI). * Turn a fitresults into the λ sets / toy ensembles the pure plotter :mod:`np_function_plots` consumes (``lambdas_from_fitresult``, ``sample_lambda_toys``, ``plot_series_from_fitresult``). -Two fit flavours are supported through SEPARATE readers, both emitting the same -:class:`~np_function_plots.NPLambdas` interface so the plotter never learns -which it came from: +Two fit flavours, SEPARATE readers, both emitting the same +:class:`~np_function_plots.NPLambdas` so the plotter never learns the source: - * NEW continuous-λ param model — λ stored PHYSICALLY in ``parms`` + * NEW continuous-λ param model: λ stored PHYSICALLY in ``parms`` (``allowNegativeParam=True``; no sqrt convention), covariance over the floating λ in ``cov``. ``lambdas_from_fitresult`` / ``sample_lambda_toys``. - * OLD template-based fit — discrete NP nuisances (``scetlibNPgamma*`` …) - whose pulls map to physical λ through a template/piecewise map. + * OLD template-based fit: discrete NP nuisances (``scetlibNPgamma*`` …) whose + pulls map to physical λ via a template/piecewise map. ``lambdas_from_template_fit`` (param-map driven). Units: the new-model λ are already physical (see ``param_model`` / -``allowNegativeParam``); no conversion is applied. The ``np_model`` / -``np_model_nu`` strings the curves need are recovered via -:func:`lambda_central.read_lambda_central` (which reads them straight from a +``allowNegativeParam``); no conversion. The ``np_model`` / ``np_model_nu`` strings +the curves need come from :func:`lambda_central.read_lambda_central` (read off the fitresults), with a CLI/argument override. CLI (print the table):: @@ -40,14 +38,8 @@ from rabbit import io_tools from wremnants.postprocessing.scetlib_np import lambda_central as _lc -from wremnants.postprocessing.scetlib_np.np_function_plots import ( - EFF_PARAMS, - GNU_PARAMS, - NPLambdas, - Series, -) - -ALL_PARAMS = GNU_PARAMS + EFF_PARAMS +from wremnants.postprocessing.scetlib_np.np_function_plots import NPLambdas, Series +from wremnants.postprocessing.scetlib_np.params import ALL_PARAMS, EFF_PARAMS, GNU_PARAMS SECTOR = { **{p: "gamma_nu (CS)" for p in GNU_PARAMS}, **{p: "F_eff (TMD)" for p in EFF_PARAMS}, @@ -261,9 +253,8 @@ def _flat_values(fitresult_path, which="postfit", result=None): def _resolve_models(fitresult_path, np_model=None, np_model_nu=None): - """np_model strings from lambda_central (reads them off the fitresults), - falling back to defaults if the upstream pkl is unreachable. Explicit - arguments win.""" + """np_model strings from lambda_central (read off the fitresults), falling + back to defaults if the upstream pkl is unreachable. Explicit arguments win.""" if np_model and np_model_nu: return np_model, np_model_nu eff_model, gnu_model = DEFAULT_NP_MODEL, DEFAULT_NP_MODEL_NU @@ -292,10 +283,15 @@ def lambdas_from_fitresult( def read_lambda_covariance(fitresult_path, result=None, names=ALL_PARAMS): """(floating_names, mean, cov) over the FLOATING λ (variance > 0). - Frozen λ (variance 0, absent, or pinned) are excluded — they don't enter - the toy band. The submatrix keeps the full correlations among the floating - λ (these are strong: e.g. lambda2_nu↔lambda2 ≈ −1).""" + Frozen λ (variance 0, absent, or pinned) are excluded from the toy band. The + submatrix keeps the full correlations among the floating λ (strong: e.g. + lambda2_nu↔lambda2 ≈ −1). + + Returns empties if the fitresults has no ``cov`` (e.g. a ``--noFit`` / + no-Hessian run); callers then skip the toy band.""" fitresult = io_tools.get_fitresult(fitresult_path, result) + if "cov" not in fitresult.keys(): + return [], np.zeros(0), np.zeros((0, 0)) parms = fitresult["parms"].get() pnames = _names(parms) pvals = parms.values() @@ -321,9 +317,8 @@ def sample_lambda_toys( ): """List of :class:`NPLambdas` toys sampled from the postfit MVN. - Floating λ are drawn jointly from their postfit covariance; frozen λ are - held at their postfit value. (For real/Asimov data the postfit point is the - correct band centre.)""" + Floating λ drawn jointly from their postfit covariance; frozen λ held at their + postfit value. (For real/Asimov data the postfit point is the band centre.)""" np_model, np_model_nu = _resolve_models(fitresult_path, np_model, np_model_nu) base = _flat_values(fitresult_path, which="postfit", result=result) floating, mean, cov = read_lambda_covariance(fitresult_path, result=result) @@ -397,7 +392,11 @@ def _template_base_eff_gnu(param_map): delta_lambda2=0.0, np_model="tanh_6", ) - gnu = dict(lambda_inf_nu=0.0, lambda2_nu=0.0, lambda4_nu=0.0, np_model_nu="tanh_6") + gnu = dict( + lambda_inf_nu=0.0, lambda2_nu=0.0, lambda4_nu=0.0, + lambda6_nu=0.0007, # SCETlib NP_model_gammanu b⁶ coeff (Gamma_nu.hpp:102) + np_model_nu="tanh_6", + ) return eff, gnu @@ -420,16 +419,16 @@ def lambdas_from_template_fit( ): """Read an OLD template-based fit → (central :class:`NPLambdas`, toys list). - ``np_param_map`` is the JSON path (or already-loaded dict) describing each - discrete NP nuisance's template Up/Down and its physical AN parameter. The - discrete-nuisance pulls are mapped to physical λ via the same piecewise - linearization the template histograms were built with. ``kfactors`` scales - a nuisance's template delta (use when the workspace was built with - ``--scaleParams``); keyed by rabbit nuisance name or AN param name. + ``np_param_map`` is the JSON path (or loaded dict) describing each discrete NP + nuisance's template Up/Down and its physical AN parameter. The nuisance pulls + map to physical λ via the same piecewise linearization the template histograms + were built with. ``kfactors`` scales a nuisance's template delta (use when the + workspace was built with ``--scaleParams``); keyed by rabbit nuisance name or + AN param name. With ``n_toys>0`` the floating nuisances are sampled from their postfit - covariance (in NUISANCE space — the linearization is applied per toy), so - the band correctly reflects the nonlinear θ→λ map. + covariance in NUISANCE space (linearization applied per toy), so the band + reflects the nonlinear θ→λ map. """ import json diff --git a/wremnants/postprocessing/scetlib_np/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py index 61eb83e1f..39a04a1d4 100644 --- a/wremnants/postprocessing/scetlib_np/lambda_central.py +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -1,13 +1,12 @@ """Central NP (lambda) parameters for the SCETlib ParamModel. The SCETlib correction's Nonperturbative runcard lives in the upstream -``*_Corr.pkl.lz4`` file under -``file_meta_data..config.Nonperturbative``. The histmaker reads that -section when it applies the correction and writes the parsed values into its -output metadata (key ``scetlib_np_lambda_central``); see -:func:`build_lambda_central_meta`. The fit then reads them back from the -metadata that rabbit propagates into the datacard / fitresults -- it never -re-opens the upstream pkl. +``*_Corr.pkl.lz4`` under +``file_meta_data..config.Nonperturbative``. The histmaker parses that +section when it applies the correction and writes the values to its output +metadata (key ``scetlib_np_lambda_central``); see +:func:`build_lambda_central_meta`. The fit reads them back from the metadata +rabbit propagates into the datacard / fitresults; it never re-opens the pkl. Write side (histmaker): build_lambda_central_meta(theory_corr_tags, procs) -> {proc: lambda_central} @@ -20,42 +19,39 @@ NP_model_effective / F_eff) and ``gnu_params`` (for NP_model_gammanu). """ -import json import os import pickle import h5py import lz4.frame +from wremnants.postprocessing.scetlib_np.params import ( + EFF_MODEL_KEY, + EFF_PARAMS, + GNU_MODEL_KEY, + GNU_PARAMS, + active_params, +) from wremnants.utilities import common as wrem_common from wums import ioutils as wums_io # Metadata key under which the histmaker stores the parsed central runcard. META_KEY = "scetlib_np_lambda_central" -# Parameter names the ParamModel needs, split by the scetlib C++ struct that -# consumes them. Nonperturbative values are strings; numeric ones get floated, -# model names stay strings. -GNU_NUMERIC = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") -GNU_STRING = ("np_model_nu",) -EFF_NUMERIC = ("lambda_inf", "lambda2", "lambda4", "lambda6", "delta_lambda2") -EFF_STRING = ("np_model",) - # ============================================================================= -# Parsing the Nonperturbative section out of an upstream correction pkl -# (write side -- only the histmaker runs this, with the pkl already in hand). +# Parse the Nonperturbative section out of an upstream correction pkl +# (write side -- only the histmaker runs this, pkl already in hand). # ============================================================================= def _find_nonperturbative(corr_dict): """Return [(basename, Nonperturbative dict)] for every basename in the pkl. - A correction pkl carries several basenames (the resummed SCETlib file, the - fixed-order singular file, the gen hist, ...). The resummed and singular - files can have DIFFERENT Nonperturbative runcards (e.g. the FranksVals - correction), so we keep all NP-bearing basenames and let - :func:`_select_resummed` pick the central one. + A correction pkl carries several basenames (resummed SCETlib file, fixed-order + singular file, gen hist, ...). Resummed and singular files can have DIFFERENT + Nonperturbative runcards (e.g. FranksVals), so keep all NP-bearing basenames + and let :func:`_select_resummed` pick the central one. """ out = [] meta = corr_dict.get("file_meta_data") @@ -76,14 +72,14 @@ def _find_nonperturbative(corr_dict): def _parse_section(npert): """Split one Nonperturbative dict into the eff / gnu parameter groups. - Numeric params absent from the runcard default to 0 -- runcards only set the - keys their np_model uses (e.g. tanh_2 omits ``lambda6``). + Numeric params absent from the runcard default to 0; a runcard only sets the + keys its np_model uses (e.g. tanh_2 omits ``lambda6``). """ - eff_params = {"np_model": npert[EFF_STRING[0]]} - gnu_params = {"np_model_nu": npert[GNU_STRING[0]]} - for k in EFF_NUMERIC: + eff_params = {EFF_MODEL_KEY: npert[EFF_MODEL_KEY]} + gnu_params = {GNU_MODEL_KEY: npert[GNU_MODEL_KEY]} + for k in EFF_PARAMS: eff_params[k] = float(npert.get(k, 0.0)) - for k in GNU_NUMERIC: + for k in GNU_PARAMS: gnu_params[k] = float(npert.get(k, 0.0)) return eff_params, gnu_params @@ -91,12 +87,12 @@ def _parse_section(npert): def _select_resummed(sections): """Pick the resummed prediction's runcard from the NP-bearing basenames. - A scetlib_dyturbo correction is built (in ``make_theory_corr.py``) from a - resummed SCETlib file plus a fixed-order *singular* file that is subtracted - in the matching. Only the resummed file's Nonperturbative runcard is the - central NP; the singular file's is not (it can differ, e.g. FranksVals). - ``make_theory_corr.py`` tells the two apart by the ``"sing"`` substring in - the filename, so we do the same: keep the basename without ``"sing"``. + A scetlib_dyturbo correction is built (``make_theory_corr.py``) from a + resummed SCETlib file plus a fixed-order *singular* file subtracted in the + matching. Only the resummed file's runcard is the central NP; the singular + file's can differ (e.g. FranksVals). ``make_theory_corr.py`` distinguishes + them by the ``"sing"`` substring in the filename, so keep the basename + without ``"sing"``. """ resummed = [item for item in sections if "sing" not in item[0]] if len(resummed) == 1: @@ -113,7 +109,7 @@ def _select_resummed(sections): def extract_lambda_central(corr_dict, tag, proc): - """Parse the central lambda parameters out of a loaded correction pkl dict. + """Parse the central lambda parameters from a loaded correction pkl dict. Returns ``{tag, basename, eff_params, gnu_params}``. Raises if the pkl has no Nonperturbative section. @@ -139,14 +135,13 @@ def _correction_pkl_path(tag, proc, data_dir=None): def build_lambda_central_meta(theory_corr_tags, procs=("Z", "W"), data_dir=None): """Build the ``scetlib_np_lambda_central`` metadata for the histmaker output. - Opens the central correction pkl (``theory_corr_tags[0]``) for each proc and + Opens the central correction pkl (``theory_corr_tags[0]``) per proc and extracts its Nonperturbative runcard. Returns ``{proc: lambda_central}`` for - the procs whose pkl exists and carries an NP section; procs without one are - skipped silently (most analyses have no SCETlib NP correction). Returns an - empty dict if there are no tags. + procs whose pkl exists and carries an NP section; procs without one are + skipped (most analyses have no SCETlib NP correction). Empty dict if no tags. - This is the ONLY place the upstream pkl is read; the fit reads the result - back from metadata. + The ONLY place the upstream pkl is read; the fit reads the result back from + metadata. """ if not theory_corr_tags: return {} @@ -161,13 +156,13 @@ def build_lambda_central_meta(theory_corr_tags, procs=("Z", "W"), data_dir=None) corr_dict = pickle.load(f) out[proc] = extract_lambda_central(corr_dict, tag, proc) except KeyError: - # pkl present but no Nonperturbative section -- not an NP correction. + # pkl present, no Nonperturbative section -- not an NP correction. continue return out # ============================================================================= -# Reading the propagated metadata (read side -- fit / postprocessing). +# Read the propagated metadata (read side -- fit / postprocessing). # ============================================================================= @@ -176,7 +171,7 @@ def _iter_meta_levels(meta, max_depth=8): The histmaker writes the key into ``meta_info``; rabbit nests that under ``meta_info_input`` in the datacard, and again in the fitresults. Walking the - chain finds the key regardless of which file we were handed. + chain finds the key whichever file was handed in. """ cur = meta for _ in range(max_depth): @@ -189,23 +184,56 @@ def _iter_meta_levels(meta, max_depth=8): cur = nxt +def _fill_missing_params(lc): + """Complete ``lc``'s eff/gnu sub-dicts for the model's full λ-vector — but + HARD-FAIL if the card omits a λ its OWN np_model uses. + + A λ the card's np_model does NOT use (e.g. ``lambda6`` / ``lambda6_nu`` under + tanh_2) is inert; the model still needs a vector slot for it, so it is filled + with 0.0 (correct, not a guess — the form ignores it). A λ the np_model DOES + use (per :func:`params.active_params`) but the metadata lacks means a + stale/corrupt card that cannot describe its own model → raise rather than + silently default it. Cards written before an *inert* param was added (e.g. + pre-``lambda6_nu`` tanh_2 cards) therefore still load.""" + lc = dict(lc) + eff = dict(lc.get("eff_params", {})) + gnu = dict(lc.get("gnu_params", {})) + needed = active_params( + np_model=eff.get(EFF_MODEL_KEY), np_model_nu=gnu.get(GNU_MODEL_KEY) + ) + missing_used = sorted(needed - (set(eff) | set(gnu))) + if missing_used: + raise KeyError( + f"lambda_central metadata is missing λ {missing_used} that its np_model " + f"({eff.get(EFF_MODEL_KEY)} / {gnu.get(GNU_MODEL_KEY)}) USES — the card " + f"cannot describe its own model; remake the histmaker output." + ) + for k in EFF_PARAMS: + eff.setdefault(k, 0.0) # inert under this np_model -> 0 (vector slot only) + for k in GNU_PARAMS: + gnu.setdefault(k, 0.0) + lc["eff_params"], lc["gnu_params"] = eff, gnu + return lc + + def read_lambda_central_from_meta(meta, proc="Z", _source=""): - """Fetch the central lambda parameters from an already-loaded metadata dict. + """Fetch the central lambda parameters from a loaded metadata dict. Searches ``meta`` and any nested ``meta_info_input`` for the propagated - ``scetlib_np_lambda_central`` entry. Raises with a clear message if it is - absent -- old inputs produced before metadata propagation must be remade - (resolving the upstream pkl by filename is no longer supported). + ``scetlib_np_lambda_central`` entry. Raises if absent: inputs produced before + metadata propagation must be remade (upstream-pkl resolution by filename is + no longer supported). Params added after the card was written are filled with + 0.0 (see :func:`_fill_missing_params`). """ for level in _iter_meta_levels(meta): lc_all = level.get(META_KEY) if not isinstance(lc_all, dict) or not lc_all: continue if proc in lc_all: - return dict(lc_all[proc]) + return _fill_missing_params(lc_all[proc]) if len(lc_all) == 1: # single proc stored -- use it whatever its label - return dict(next(iter(lc_all.values()))) + return _fill_missing_params(next(iter(lc_all.values()))) raise KeyError( f"{_source}: {META_KEY!r} has no proc {proc!r} (have {sorted(lc_all)})." ) @@ -216,70 +244,21 @@ def read_lambda_central_from_meta(meta, proc="Z", _source=""): ) -def load_lambda_central_file(path): - """Load a lambda_central override from a JSON or YAML file. - - The file must decode to a dict with ``eff_params`` and ``gnu_params`` - sub-dicts. Format is chosen by extension (``.yaml``/``.yml`` -> YAML, else - JSON; YAML also accepts JSON). - """ - if not os.path.exists(path): - raise FileNotFoundError(f"lambda_central file missing: {path!r}") - with open(path) as f: - text = f.read() - if path.lower().endswith((".yaml", ".yml")): - import yaml - - data = yaml.safe_load(text) - else: - try: - data = json.loads(text) - except json.JSONDecodeError as exc: - raise ValueError( - f"lambda_central file {path!r} is not valid JSON; got {exc}" - ) from exc - if ( - not isinstance(data, dict) - or "eff_params" not in data - or "gnu_params" not in data - ): - raise ValueError( - f"lambda_central file {path!r} must decode to a dict with " - f"'eff_params' and 'gnu_params' keys; got {type(data).__name__}." - ) - return data - - def read_lambda_central(hdf5_path, proc="Z"): """Read the central lambda parameters referenced by an hdf5. - Accepts a setupRabbit datacard or a rabbit ``fitresults*.hdf5`` (rabbit - nests the datacard meta one level down; the lookup handles both). - - Order of preference: - 1. a ``lambda_central=`` token in the stored ``--paramModel`` spec - (an explicit override the fit command recorded); - 2. the ``scetlib_np_lambda_central`` metadata propagated by the histmaker. + Accepts a setupRabbit datacard or a rabbit ``fitresults*.hdf5`` (rabbit nests + the datacard meta one level down; the lookup handles both). Reads the + ``scetlib_np_lambda_central`` metadata the histmaker propagates. - Returns a dict with ``tag``, ``basename``, ``eff_params``, ``gnu_params`` - and ``source``. Raises with a clear message if neither route resolves. + Returns ``{tag, basename, eff_params, gnu_params, source}``. Raises if the + metadata is absent. """ with h5py.File(hdf5_path, "r") as f: if "meta" not in f: raise KeyError(f"{hdf5_path}: no 'meta' group -- wrong file type?") meta = wums_io.pickle_load_h5py(f["meta"]) - # Preference 1: an explicit lambda_central= override. - args_meta = (meta.get("meta_info") or {}).get("args") or {} - for spec in args_meta.get("paramModel") or []: - for tok in spec: - if isinstance(tok, str) and tok.startswith("lambda_central="): - path = tok.split("=", 1)[1] - lc = load_lambda_central_file(path) - lc["source"] = f"cli-file:{path}" - return lc - - # Preference 2: the propagated histmaker metadata. lc = read_lambda_central_from_meta(meta, proc=proc, _source=hdf5_path) lc["source"] = "histmaker-metadata" return lc diff --git a/wremnants/postprocessing/scetlib_np/np_damping_wall.py b/wremnants/postprocessing/scetlib_np/np_damping_wall.py new file mode 100644 index 000000000..09641be2f --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/np_damping_wall.py @@ -0,0 +1,337 @@ +"""Physical-damping wall for the continuous-λ SCETlib NP param model. + +Companion to ``wremnants/postprocessing/np_monotonicity.py``: same rabbit +``Regularizer`` mechanism and hinge-loss (relu²) wall structure, for the +:class:`~wremnants.postprocessing.scetlib_np.param_model.SCETlibNPParamModel`, +whose λ are direct fit parameters (the model-param block of ``x``), not discrete +template nuisances. No ``PARAM_MAP`` / θ-interpolation: the regularizer reads the +physical λ from ``x[:nparams]`` in the model's canonical parameter order. + +The model FORMS and λ ORDER are DERIVED from the model, not re-declared here. +``SCETlibNPParamModel.__init__`` publishes itself on the shared ``indata`` +(``indata.scetlib_np_param_model``); this regularizer — built afterward with the +same ``indata`` — reads ``model.fit_forms`` (the F_eff / γ_ν NUMERATOR forms the +fit integrates) and ``model._param_order`` (POI-reordering and all). So you pass +the form ONCE, to the model; nothing on the ``-r`` line repeats it, and a +``poi_params`` reorder is tracked automatically. The wall RAISES if no model has +registered, or if a side's form is one it has no walls for (see dispatch below). + +Why a wall (see ``param_model.py``): a wrong-sign λ anti-damps the NP form +factors, so the bT-integral diverges / the differential σ(qT) oscillates negative +— a genuine unphysical region. The fit's qT→ptVGen rebin launders it by averaging +the negative differential σ away, so σ_gen / σ_reco / the NLL stay finite and the +minimizer cannot see the unphysical-ness. These λ must lie in the physical +(damping) region; this regularizer adds the fit-time penalty keeping them there. + +Walls, per side and per FORM, read off the TF forms the fit integrates +(``btgrid_tf.gamma_nu_NP_tf`` / ``btgrid_tf.F_eff_tf``). The damping criterion is +"the tanh argument is ≥ 0 ∀b" (γ_ν ≤ 0 / F_eff decays), i.e. a polynomial in +u ≡ b² is non-negative on u ≥ 0: + + CS side γ_ν^NP(b) = −λ∞_ν · tanh( P(u)/λ∞_ν ) , + P(u) = λ2_ν·u + λ4_ν·u² + λ6_ν·u³ (λ6_ν ≡ 0 for tanh_2) + tanh_2: λ∞_ν > 0 , λ2_ν ≥ 0 , λ4_ν ≥ 0 + tanh_6: λ∞_ν > 0 , λ2_ν ≥ 0 , λ6_ν ≥ 0 (leading) , + and λ4_ν ≥ 0 OR λ4_ν² ≤ 4·λ2_ν·λ6_ν (interior; λ4_ν may dip + negative while the b⁶ term keeps P ≥ 0) + + TMD side F_eff(Y,b) = exp(−2·λ∞·b·tanh(a)) , + a·λ∞ = b·Q(u) , Q(u) = λ2_Y + B·u + λ6·u² , + B = λ4 + λ2_Y³/(3λ∞²) , λ2_Y = λ2 + δλ2·Y² (λ6 ≡ 0 for tanh_2) + tanh_2: λ∞ > 0 , λ2_Y ≥ 0 , B ≥ 0 [≡ 3·λ∞²·λ4 + λ2_Y³ ≥ 0] + tanh_6: λ∞ > 0 , λ2_Y ≥ 0 , λ6 ≥ 0 (leading) , + and B ≥ 0 OR B² ≤ 4·λ2_Y·λ6 (interior) + evaluated at Y=0 and Y=Y_MAX (covers δλ2 of either sign — λ2_Y is + monotonic in Y², so the binding |Y| is one of the two extremes). + +The tanh_2 walls are EXACT in λ∞ (read off the actual TF forms, not the AN +λ∞-normalised parametrisation). They are the λ6 = 0 reduction of the tanh_6 walls +(then the leading coefficient is λ4_ν / B, hence the simpler ≥0 limit), and of the +``np_monotonicity.py`` monotonicity walls. NOTE this file uses the DAMPING +criterion (P ≥ 0), which lets the interior coefficient dip negative within the +discriminant bound; ``np_monotonicity.py`` uses the stricter MONOTONICITY +criterion (√3 in place of √4). The interior relu²(relu(−λ4_ν)² − 4·λ2_ν·λ6_ν) +penalty form is self-gating: it vanishes for λ4_ν ≥ 0 (no division by λ6). + +Wall hardness is set at fit time by rabbit's ``--regularizationStrength`` (the +penalty × ``exp(2·tau)`` in ``fitter.py``; ``tau`` is a fixed multiplier, NOT a +minimised parameter). A large strength makes this a BARRIER: ≈0 inside the +physical region, steeply rising outside. Free (small strength) vs walled (large +strength) Δχ² is the data–model tension diagnostic: railing against a wall with +large Δχ² is genuine tension, not a masked pathology. + +The small-b turn-on walls (λ2_ν ≥ 0, λ2_Y ≥ 0) are a stronger condition than the +large-b limit (they forbid an anti-damping bump near b→0, not just the wrong +asymptote). They can be switched off with the mapping flag ``smallb=0`` — then +ONLY the limiting/interior behaviour and the λ∞ floors are enforced, and the +leading b² coefficient floats either sign (use the postfit σ(qT)≥0 check in +``param_model_diagnostics`` as the real guard then). + +Invoke (the model form is taken from ``--paramModel``; nothing repeats it here): + + rabbit_fit.py ... \\ + --regularizationStrength 3 \\ + -r wremnants.postprocessing.scetlib_np.np_damping_wall.NPDampingWall \\ + wremnants.postprocessing.scetlib_np.np_damping_wall.NPDampingMapping \\ + [smallb=0] + +(``Y_MAX`` = 5 — the kinematic ceiling / btgrid Y reach — the λ∞ floor, and the +damping margin are fixed module constants, not -r options; both Y_MAX=5 and a +nonzero margin are UNDER TEST, edit the constants to change them.) + +References: + AN-25-085 theory.tex Eqs. eq:npgamma, eq:npf; + param_model.py / sigma_gen.py docstrings (the σ_gen pipeline & the binning- + launders-the-pathology discussion). +""" + +# rabbit / TF imports deferred to the lazy class factories so the module stays +# importable without rabbit/TF (mirrors np_monotonicity.py). + +# Forms this wall has damping conditions for. Anything else (frac_*, exp_*, +# tanh_1, tanh_4, signed_lambda, identity, …) raises rather than silently +# applying the wrong tanh_2/tanh_6 walls. +SUPPORTED_FORMS = ("tanh_2", "tanh_6") + +# btgrid_tf form aliases that resolve to a SUPPORTED_FORMS entry (mirrors the +# alias maps in gamma_nu_NP_tf / F_eff_tf). Other aliases ("linear"->frac_1, +# "square_root"->frac_2) resolve to unsupported forms and so fall through to the +# raise. Applied to both sides; only "hyp_tangent" reaches a supported form. +_FORM_ALIASES = {"hyp_tangent": "tanh_2"} + +# λ that the walls reference (always present in the model's _param_order, which +# is built from params.ALL_PARAMS); cross-checked at construction. +_REQUIRED_PARAMS = ( + "lambda2_nu", + "lambda4_nu", + "lambda6_nu", + "lambda_inf_nu", + "lambda2", + "lambda4", + "lambda6", + "delta_lambda2", + "lambda_inf", +) + +# Fixed knobs (NOT CLI options — edit here to test). Only ``smallb`` is exposed +# on the -r line; these are deliberately constants to keep that line minimal. +Y_MAX = 5.0 # binding |y| for the F_eff Y-evaluation: the kinematic ceiling +# ln(√s/Q) ≈ the btgrid Y reach (±5). UNDER TEST (was 2.5 = the |y| +# acceptance); 5 demands physical damping over the full phase space. +LAMBDA_INF_FLOOR = 1e-3 # positive floor on the λ∞ saturation scales +NP_DAMPING_MARGIN = 5e-3 # positive cushion: enforce each damping coeff ≥ this, +# not just ≥ 0, to keep the soft-wall equilibrium off the boundary. +# UNDER TEST; default 0 (off) — set > 0 here to try the cushion. + + +def _make_mapping_class(): + from rabbit.mappings.mapping import BaseMapping + + class NPDampingMapping(BaseMapping): + """Vestigial BaseMapping carrying the wall's option to the regularizer. + + Only option (``key=value`` token, optional): + smallb=<0|1> enforce the small-b turn-on walls λ2_ν≥0 and λ2_Y≥0 + (default 1). smallb=0 drops them, keeping ONLY the + large-b limit/interior walls and the λ∞ floors — i.e. + constrain the limiting behaviour but let the leading + small-b coefficient float either sign. + + The model forms and λ order are derived from the registered + ``SCETlibNPParamModel`` (see module docstring). The binding |Y| (``Y_MAX``), + the λ∞ floor (``LAMBDA_INF_FLOOR``), and the damping cushion + (``NP_DAMPING_MARGIN``) are FIXED module constants, not CLI options — edit + them in this file to test, kept off the -r line on purpose. + """ + + def __init__(self, indata, key, smallb=True): + super().__init__(indata, key) + self.indata = indata + self.smallb = bool(smallb) + + @classmethod + def parse_args(cls, indata, *args): + smallb = True + for a in args: + if "=" not in a: + raise ValueError( + f"NPDampingMapping: arg must be 'smallb=<0|1>', got '{a}'" + ) + k, v = a.split("=", 1) + if k == "smallb": + smallb = v.strip().lower() not in ("0", "false", "no", "off") + else: + raise ValueError( + f"NPDampingMapping: unknown key '{k}'; only 'smallb' is " + f"supported (ymax/eps/margin are fixed module constants)." + ) + return cls(indata, f"{cls.__name__} smallb={int(smallb)}", smallb=smallb) + + return NPDampingMapping + + +def _make_regularizer_class(): + import tensorflow as tf + from rabbit.regularization.regularizer import Regularizer + + class NPDampingWall(Regularizer): + """Hinge-loss penalty enforcing NP damping, per-side and per-form + (tanh_2 / tanh_6); see the module docstring.""" + + def __init__(self, mapping, dtype): + super().__init__(mapping, dtype) + self.dtype = dtype + self.mapping = mapping + self.indata = mapping.indata + # ymax / eps / margin are fixed module constants (not CLI options); + # only smallb comes from the mapping. + self.ymax = Y_MAX + self.eps = LAMBDA_INF_FLOOR + self.margin = NP_DAMPING_MARGIN + self.enforce_small_b = bool(getattr(mapping, "smallb", True)) + + # Forms + λ order are DERIVED from the SCETlibNPParamModel that + # published itself on the shared indata (built before this regularizer + # in rabbit_fit; see param_model.py). No -r-line repetition. + model = getattr(self.indata, "scetlib_np_param_model", None) + if model is None: + raise ValueError( + "NPDampingWall: no SCETlibNPParamModel registered on indata. " + "This wall derives the NP form and λ order from the param " + "model, so the fit must use " + "--paramModel ...scetlib_np.param_model.SCETlibNPParamModel " + "(which publishes indata.scetlib_np_param_model)." + ) + + self._order = tuple(model._param_order) + self._pidx = {name: i for i, name in enumerate(self._order)} + missing = [p for p in _REQUIRED_PARAMS if p not in self._pidx] + if missing: + raise ValueError( + f"NPDampingWall: model param order {self._order} is missing " + f"required λ {missing}." + ) + + # FIT (numerator) forms — the ones the fit integrates, which the wall + # must constrain (NOT the card/denominator form). Resolve aliases and + # fail on any form we have no walls for. + forms = model.fit_forms + self._np_model_nu = self._resolve_form( + forms["np_model_nu"], side="CS (γ_ν, np_model_nu)" + ) + self._np_model = self._resolve_form( + forms["np_model"], side="TMD (F_eff, np_model)" + ) + + self._cast = lambda v: tf.constant(v, dtype=self.dtype) + # Model-param block is x[:nparams]; nparams resolved at set_expectations. + self._nparams = None + + @staticmethod + def _resolve_form(form, side): + resolved = _FORM_ALIASES.get(form, form) + if resolved not in SUPPORTED_FORMS: + raise NotImplementedError( + f"NPDampingWall: no damping walls for {side} form {form!r}" + + (f" (resolves to {resolved!r})" if resolved != form else "") + + f"; supported: {sorted(SUPPORTED_FORMS)}. Add walls for it " + "or run that side with a supported form." + ) + return resolved + + def set_expectations(self, initial_params, initial_observables): + nsyst = len(self.indata.systs) + self._nparams = int(initial_params.shape[0]) - nsyst + if self._nparams != len(self._order): + raise ValueError( + f"NPDampingWall: the fit's model-param block is {self._nparams} " + f"wide but the model param order has {len(self._order)} entries " + f"{self._order}. A wrapping/composite param model (e.g. the " + "saturated goodness-of-fit path) reorders/resizes the block in " + "a way this wall's flat indexing cannot follow." + ) + + def _lam(self, params, name): + # λ stored directly in the model-param block (allowNegativeParam=True), + # so x[index] IS the physical λ — no theta interpolation. + return params[self._pidx[name]] + + def compute_nll_penalty(self, params, observables): + zero = self._cast(0.0) + eps = self._cast(self.eps) + m = self._cast(self.margin) # positive cushion: enforce coeff ≥ margin + three = self._cast(3.0) + four = self._cast(4.0) + thirtysix = self._cast(36.0) + + def relu2(x): # hinge: 0 if x ≤ 0 else x² + return tf.square(tf.maximum(zero, x)) + + # Each damping condition "coeff ≥ 0" is enforced as "coeff ≥ margin" + # (margin=0 → bare ≥0): relu2(margin - coeff). The soft wall's gradient + # vanishes at the knee, so a weakly-constrained coeff settles a hair + # past it; the margin moves the knee so the equilibrium stays damping. + def wall(coeff): # coeff ≥ margin + return relu2(m - coeff) + + # ---- CS-side γ_ν^NP damping: P(u)=λ2_ν·u+λ4_ν·u²+λ6_ν·u³ ≥ 0 ∀u≥0. + l2nu = self._lam(params, "lambda2_nu") + l4nu = self._lam(params, "lambda4_nu") + l6nu = self._lam(params, "lambda6_nu") + linfnu = self._lam(params, "lambda_inf_nu") + pens = [relu2(eps - linfnu)] # λ∞_ν > 0 (saturation-scale regime) + if self._np_model_nu == "tanh_2": + pens.append(wall(l4nu)) # λ4_ν ≥ margin (large-b leading) + else: # tanh_6 + pens.append(wall(l6nu)) # λ6_ν ≥ margin (large-b leading) + # interior: λ4_ν ≥ 0 OR λ4_ν² ≤ 4·λ2_ν·λ6_ν. Self-gating — the + # relu2(-l4nu) vanishes for λ4_ν ≥ 0, so no penalty there; and no + # division by λ6_ν. (No margin on the interior discriminant.) + pens.append(relu2(relu2(-l4nu) - four * l2nu * l6nu)) + if self.enforce_small_b: + pens.append(wall(l2nu)) # λ2_ν ≥ margin (small-b turn-on) + + # ---- TMD-side F_eff damping: Q(u)=λ2_Y+B·u+λ6·u² ≥ 0 ∀u≥0, evaluated + # at the binding |Y| extremes. cubic ≡ 3·λ∞²·B = 3·λ∞²·λ4 + λ2_Y³, so + # all conditions stay division-free (multiply through by 3·λ∞² > 0). + l2 = self._lam(params, "lambda2") + l4 = self._lam(params, "lambda4") + l6 = self._lam(params, "lambda6") + dl2 = self._lam(params, "delta_lambda2") + linf = self._lam(params, "lambda_inf") + pens.append(relu2(eps - linf)) # λ∞ > 0 + linf2 = linf * linf + for y_sq in (0.0, self.ymax * self.ymax): + l2Y = l2 + dl2 * self._cast(y_sq) + if self.enforce_small_b: + pens.append(wall(l2Y)) # λ2_Y ≥ margin (small-b turn-on) + cubic = three * linf2 * l4 + l2Y**3 # 3·λ∞²·B + if self._np_model == "tanh_2": + pens.append(wall(cubic)) # B ≥ margin (large-b leading) + else: # tanh_6 + pens.append(wall(l6)) # λ6 ≥ margin (large-b leading) + # interior: B ≥ 0 OR B² ≤ 4·λ2_Y·λ6. In cubic-space (cubic = + # 3·λ∞²·B): cubic ≥ 0 OR cubic² ≤ 36·λ∞⁴·λ2_Y·λ6. Self-gating, + # division-free. (No margin on the interior discriminant.) + bound = thirtysix * linf2 * linf2 * l2Y * l6 + pens.append(relu2(relu2(-cubic) - bound)) + + return tf.add_n(pens) + + return NPDampingWall + + +# PEP-562 lazy class resolution: rabbit's loader does +# module = importlib.import_module(...); cls = getattr(module, class_name) +# so the classes are synthesised on first attribute access, keeping the module +# importable without rabbit / TF (matches np_monotonicity.py). +def __getattr__(name): + if name == "NPDampingMapping": + cls = _make_mapping_class() + globals()["NPDampingMapping"] = cls + return cls + if name == "NPDampingWall": + cls = _make_regularizer_class() + globals()["NPDampingWall"] = cls + return cls + raise AttributeError(name) diff --git a/wremnants/postprocessing/scetlib_np/np_function_plots.py b/wremnants/postprocessing/scetlib_np/np_function_plots.py index f90af5508..5348b0453 100644 --- a/wremnants/postprocessing/scetlib_np/np_function_plots.py +++ b/wremnants/postprocessing/scetlib_np/np_function_plots.py @@ -1,38 +1,39 @@ -"""Plot the SCETlib NP form factors — CS γ_ν^NP(b_T) and TMD F_eff(b_T, y). +"""Plot the SCETlib NP form factors: CS γ_ν^NP(b_T) and TMD F_eff(b_T, y). -This is a PURE plotting library: it takes physical λ values (the two parameter -dicts the model uses) and draws the two NP functions. It knows nothing about -fitresults — where the λ come from (a new continuous-λ fit, an old -template-based fit, or hand-picked values to test) is the caller's job. The -companion reader :mod:`fitresult_lambdas` turns a fitresults HDF5 into the λ -sets / toy ensembles this module consumes; ``main()`` below glues the two for -convenience but the plot functions stay reader-agnostic. +A pure plotting library: takes physical λ values (the two parameter dicts the +model uses) and draws the two NP functions. Where the λ come from (a new +continuous-λ fit, an old template-based fit, or hand-picked values) is the +caller's job. The companion reader :mod:`fitresult_lambdas` turns a fitresults +HDF5 into the λ sets / toy ensembles this module consumes; ``main()`` glues the +two together, but the plot functions stay reader-agnostic. The curves call the same form factors the fit integrates (:func:`btgrid_tf.F_eff_tf` / :func:`btgrid_tf.gamma_nu_NP_tf`), driven by the -``np_model`` / ``np_model_nu`` strings, so a plotted curve is exactly the model -the fit used. +``np_model`` / ``np_model_nu`` strings, so a plotted curve is the fit's model. A "λ set" is the pair of dicts :class:`NPLambdas` carries: eff = {lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_model} (F_eff / TMD) gnu = {lambda_inf_nu, lambda2_nu, lambda4_nu, np_model_nu} (γ_ν / CS) -These map 1:1 onto the form-factor keyword arguments. A band is drawn from -a list of λ-set "toys" handed in by the caller (this module just takes -percentiles of the resulting curves); it never samples and never sees a -covariance. +These map 1:1 onto the form-factor keyword arguments. A band is drawn from a +caller-supplied list of λ-set "toys" (this module takes percentiles of the +resulting curves); it never samples and never sees a covariance. CLI (plot a raw λ set, no fit involved):: python -m wremnants.postprocessing.scetlib_np.np_function_plots \\ - --lambda2 0.4 --lambda4 0.4 --lambda2_nu 0.15 \\ + --lambdas lambda2=0.4,lambda4=0.4,lambda2_nu=0.15 \\ --np-model tanh_6 --np-model-nu tanh_2 -o /tmp/np.png +The ``--lambdas`` names must be λ the chosen models actually use (e.g. ``lambda6`` +needs ``--np-model tanh_6``); naming an inert λ is a hard error rather than a +silently-ignored value. Unset λ stay at their defaults (NP-unit point). + CLI (from a fitresults: prefit dashed, postfit solid + 68% band):: python -m wremnants.postprocessing.scetlib_np.np_function_plots \\ - --from-fitresult -o /tmp/np.png + --fitresult -o /tmp/np.png """ import argparse @@ -43,20 +44,22 @@ import numpy as np from wremnants.postprocessing.scetlib_np import btgrid_tf - -# λ split across the two NP sectors, with sensible "all knobs off" defaults so a -# bare CLI call still draws something. Mirrors param_model.{GNU,EFF}_PARAMS. -GNU_PARAMS = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") -EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") +from wremnants.postprocessing.scetlib_np.params import ( + EFF_PARAMS, + GNU_PARAMS, + active_params, + parse_lambda_overrides, + split_eff_gnu, +) @dataclass class NPLambdas: """One physical λ point: the two form-factor parameter dicts. - ``eff`` / ``gnu`` hold exactly the kwargs ``btgrid_tf.F_eff_tf`` / - ``gamma_nu_NP_tf`` expect (numeric λ + the ``np_model`` / ``np_model_nu`` - string). Build one from a fit via :mod:`fitresult_lambdas`, or by hand. + ``eff`` / ``gnu`` hold the kwargs ``btgrid_tf.F_eff_tf`` / ``gamma_nu_NP_tf`` + expect (numeric λ + ``np_model`` / ``np_model_nu``). Build via + :mod:`fitresult_lambdas` or by hand. """ eff: dict @@ -65,8 +68,7 @@ class NPLambdas: @classmethod def from_flat(cls, values, np_model, np_model_nu): """Build from a flat name->value mapping (the param-model λ names).""" - eff = {k: float(values[k]) for k in EFF_PARAMS if k in values} - gnu = {k: float(values[k]) for k in GNU_PARAMS if k in values} + eff, gnu = split_eff_gnu(values) eff["np_model"] = np_model gnu["np_model_nu"] = np_model_nu return cls(eff=eff, gnu=gnu) @@ -111,24 +113,36 @@ def _band(curves, pct): } +# LaTeX labels per λ name, for the parameter inset. +_EFF_LABELS = { + "lambda2": r"\lambda_2", + "lambda4": r"\lambda_4", + "delta_lambda2": r"\delta\lambda_2", + "lambda6": r"\lambda_6", + "lambda_inf": r"\lambda_\infty", +} +_GNU_LABELS = { + "lambda2_nu": r"\lambda_2^\nu", + "lambda4_nu": r"\lambda_4^\nu", + "lambda6_nu": r"\lambda_6^\nu", + "lambda_inf_nu": r"\lambda_\infty^\nu", +} + + def _param_inset(ax, lam, sector, corner="upper right"): - """Small text box listing the λ that drive the panel.""" + """Small text box listing only the λ that drive the panel for its model.""" if sector == "gnu": - lines = [ - rf"$\lambda_2^\nu = {lam.gnu.get('lambda2_nu', 0):+.4f}$", - rf"$\lambda_4^\nu = {lam.gnu.get('lambda4_nu', 0):+.4f}$", - rf"$\lambda_\infty^\nu = {lam.gnu.get('lambda_inf_nu', 0):+.4f}$", - rf"model: {lam.gnu.get('np_model_nu', '?')}", - ] + model = lam.gnu.get("np_model_nu", "?") + active = active_params(np_model_nu=model) + labels, src, order = _GNU_LABELS, lam.gnu, GNU_PARAMS else: - lines = [ - rf"$\lambda_2 = {lam.eff.get('lambda2', 0):+.4f}$", - rf"$\lambda_4 = {lam.eff.get('lambda4', 0):+.4f}$", - rf"$\delta\lambda_2 = {lam.eff.get('delta_lambda2', 0):+.4f}$", - rf"$\lambda_6 = {lam.eff.get('lambda6', 0):+.4f}$", - rf"$\lambda_\infty = {lam.eff.get('lambda_inf', 0):+.4f}$", - rf"model: {lam.eff.get('np_model', '?')}", - ] + model = lam.eff.get("np_model", "?") + active = active_params(np_model=model) + labels, src, order = _EFF_LABELS, lam.eff, EFF_PARAMS + lines = [ + rf"${labels[k]} = {src.get(k, 0):+.4f}$" for k in order if k in active + ] + lines.append(rf"model: {model}") box = dict(boxstyle="round,pad=0.35", fc="white", ec="0.6", alpha=0.85) x, y, va, ha = _CORNERS[corner] ax.text( @@ -146,7 +160,7 @@ def _param_inset(ax, lam, sector, corner="upper right"): def plot_np_functions( series: Sequence[Series], *, - y_values: Sequence[float] = (0.0, 2.5), + y_values: Sequence[float] = (0.0, 2.5, 5.0), bT_max: float = 4.0, n_points: int = 401, outpath: str, @@ -158,15 +172,15 @@ def plot_np_functions( Parameters ---------- series - Curves to overlay. The first series with ``toys`` set draws a - percentile band on each panel. Pure: the caller supplies the toys. + Curves to overlay. Each series with ``toys`` draws a percentile band on + its panels; the caller supplies the toys. y_values Rapidity values for the TMD panel (F_eff depends on y; γ_ν does not). bT_max, n_points b_T grid for the curves [GeV^-1]. inset_from - Series whose λ are written into the per-panel parameter box (defaults - to the last series — typically the postfit point). + Series whose λ fill the per-panel parameter box (default: last series, + typically the postfit point). """ import matplotlib @@ -179,11 +193,11 @@ def plot_np_functions( auto_colors = [c for c in plt.rcParams["axes.prop_cycle"].by_key()["color"]] cmap_tmd = plt.cm.viridis - # NP factors are evaluated at the bare b_T grid (this grid's b* prescription - # is the identity, b_bar == b_T; see param_model / base.conf b0_over_bmax=0). - # For a b*-frozen grid the caller would need to map b_T -> b_bar first. - # F_eff can run away at large b_T for λ4 < 0 toys, so we scale the TMD panel - # to the line curves rather than let a runaway band tail set the autoscale. + # NP factors evaluated at the bare b_T grid (this grid's b* prescription is + # the identity, b_bar == b_T; see param_model / base.conf b0_over_bmax=0); a + # b*-frozen grid would need b_T -> b_bar mapped first. F_eff runs away at large + # b_T for λ4 < 0 toys, so the TMD panel scales to the line curves (below), not + # a runaway band tail. line_fmax = 0.0 for si, s in enumerate(series): @@ -240,18 +254,22 @@ def plot_np_functions( axR.legend(loc="upper right", fontsize=8) axR.grid(alpha=0.3) - # Scale the TMD panel to the line curves so a runaway band tail (bare F_eff - # for λ4 < 0) doesn't dominate the autoscale. + # Scale to the line curves so a runaway band tail (bare F_eff, λ4 < 0) can't + # dominate the autoscale. top = f_ymax if f_ymax is not None else max(1.1, 1.2 * line_fmax) axR.set_ylim(0.0, top) - # Param boxes diagonally opposite each panel's legend so they don't collide: - # CS legend lower-left -> box upper-right; TMD legend upper-right -> box lower-left. + # Param boxes diagonally opposite each panel's legend to avoid collisions: + # CS legend lower-left -> box upper-right; TMD upper-right -> box lower-left. inset = inset_from if inset_from is not None else (series[-1] if series else None) if inset is not None: _param_inset(axL, inset.lam, "gnu", corner="upper right") _param_inset(axR, inset.lam, "eff", corner="lower left") + # Allow --outpath to be a directory (trailing slash or no extension): append + # a default filename rather than erroring on a bare ".png". + if outpath.endswith(("/", os.sep)) or os.path.isdir(outpath) or not os.path.splitext(outpath)[1]: + outpath = os.path.join(outpath, "np_functions.png") os.makedirs(os.path.dirname(os.path.abspath(outpath)) or ".", exist_ok=True) fig.tight_layout() fig.savefig(outpath, dpi=140) @@ -263,9 +281,10 @@ def plot_np_functions( # CLI # --------------------------------------------------------------------------- -# Defaults for a bare CLI call: SCETlib "knobs off" plus the conventional model -# strings (override with --np-model / --np-model-nu and the --lambda* flags). -_CLI_DEFAULTS = dict( +# Baseline λ for the raw mode: SCETlib "knobs off" (NP-unit point). --lambdas +# overrides any subset; unset λ stay here. Carries all ALL_PARAMS so the +# form-factor kwargs are complete regardless of the chosen model. +_DEFAULT_LAMBDAS = dict( lambda2=0.0, lambda4=0.0, lambda6=0.0, @@ -273,6 +292,7 @@ def plot_np_functions( lambda_inf=1.0, lambda2_nu=0.0, lambda4_nu=0.0, + lambda6_nu=0.0, lambda_inf_nu=1.0, ) @@ -284,30 +304,35 @@ def make_parser(): ) src = p.add_argument_group("input (pick one mode)") src.add_argument( - "--from-fitresult", + "--fitresult", default=None, help="fitresults HDF5: plot prefit (dashed) + postfit (solid + band). " "Uses fitresult_lambdas to read λ / sample the band.", ) src.add_argument( - "--result", default=None, help="results group suffix for --from-fitresult." + "--result", default=None, help="results group suffix for --fitresult." ) src.add_argument( - "--n-toys", type=int, default=500, help="band toys (--from-fitresult)." + "--n-toys", type=int, default=500, help="band toys (--fitresult)." ) src.add_argument("--seed", type=int, default=0, help="band RNG seed.") raw = p.add_argument_group("raw λ mode (no fit)") - for k, v in _CLI_DEFAULTS.items(): - raw.add_argument(f"--{k}", type=float, default=None, help=f"default {v}") - raw.add_argument("--np-model", default="tanh_6", help="F_eff model string.") + raw.add_argument( + "--lambdas", + default=None, + help="λ overrides 'name=val,...' (e.g. lambda2=0.4,lambda2_nu=0.15); unset " + "λ stay at the NP-unit defaults. Names must be λ the chosen models use " + "(lambda6/lambda6_nu need tanh_6); an inert λ is a hard error.", + ) + raw.add_argument("--np-model", default="tanh_2", help="F_eff model string.") raw.add_argument("--np-model-nu", default="tanh_2", help="γ_ν model string.") p.add_argument( "--y", type=float, nargs="+", - default=[0.0, 2.5], + default=[0.0, 2.5, 5.0], help="rapidity values for the TMD panel.", ) p.add_argument("--bT-max", type=float, default=4.0) @@ -324,23 +349,36 @@ def make_parser(): def main(argv=None): - args = make_parser().parse_args(argv) + parser = make_parser() + args = parser.parse_args(argv) - if args.from_fitresult: + if args.fitresult: # Reading lives in the reader module; the plotter stays pure. from wremnants.postprocessing.scetlib_np import fitresult_lambdas as frl series = frl.plot_series_from_fitresult( - args.from_fitresult, + args.fitresult, result=args.result, n_toys=args.n_toys, seed=args.seed, ) else: - vals = { - k: (getattr(args, k) if getattr(args, k) is not None else v) - for k, v in _CLI_DEFAULTS.items() - } + try: + overrides = parse_lambda_overrides(args.lambdas) + except ValueError as e: + parser.error(str(e)) + active = active_params(args.np_model, args.np_model_nu) + inert = [k for k in overrides if k not in active] + if inert: + parser.error( + "--lambdas: " + + ", ".join(inert) + + f" not used by np_model={args.np_model} / " + + f"np_model_nu={args.np_model_nu} (active: " + + ", ".join(k for k in (*EFF_PARAMS, *GNU_PARAMS) if k in active) + + ")" + ) + vals = {**_DEFAULT_LAMBDAS, **overrides} lam = NPLambdas.from_flat(vals, args.np_model, args.np_model_nu) series = [Series(label=args.label, lam=lam, color="C3")] diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 9aee3044e..7ff44e884 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -1,25 +1,41 @@ """SCETlibNPParamModel — continuous-λ rabbit ParamModel for SCETlib NP. -This ParamModel scales the signal reco template by a per-bin ratio of the -SCETlib nonperturbative (NP) prediction at the fitted λ vs. at λ_central. The -prediction is built in four steps, written out below; the related modules -(:mod:`response_matrix`, :func:`btgrid_tf.reconstruct_batch`, the validation -scripts) refer here for the derivation. - -Pipeline at a glance (everything is a function of the NP parameters λ): - - Step 1 btgrid Hankel + Q integral → σ_resum(λ; g) resummed, on the gen grid - Step 2 + fixed-order matching → σ_gen(λ; g) = σ_resum(λ; g) + σ_ns(g) - Step 3 fold through response R → σ_reco(λ; b) gen → reco - Step 4 ratio vs λ_central → rnorm(b, proc) the shape handed to rabbit - -Steps 1–3 build the absolute physical cross section; Step 4 alone produces the -per-bin variation the fit consumes — they are kept separate on purpose. - -Indices: Q, Y, qT are the SCETlib btgrid axes (boson mass / rapidity / qT); -g = flattened gen bin (ptVGen, absYVGen); b = flattened reco bin (ptll, yll, -cosThetaStarll_quantile, phiStarll_quantile). λ splits into λ_eff (for F_eff) -and λ_ν (for γ_ν^NP) — the 8 differentiable parameters listed at the end. +Scales the signal reco template by a per-bin ratio of the SCETlib +nonperturbative (NP) prediction at the fitted λ vs at λ_central. Built in four +steps; :mod:`response_matrix`, :mod:`btgrid_tf` and the validation scripts refer +here for the derivation. + + Step 1 btgrid Hankel + Q integral → σ_resum(λ; g) resummed, gen grid + Step 2 + fixed-order matching → σ_gen(λ; g) = σ_resum + σ_ns + Step 3 fold through response R → σ_reco(λ; b) gen → reco + Step 4 ratio vs λ_central → rnorm(b, proc) handed to rabbit + +Steps 1–3 build the absolute cross section; Step 4 forms the per-bin variation +the fit consumes. + +Code split: Steps 1–2 (the datacard-free physics — btgrid integral → matched +σ_gen on a gen grid, from btgrid + λ_central + gen edges) live in +:class:`~wremnants.postprocessing.scetlib_np.sigma_gen.SigmaGenModel`. This class +is the loader/rabbit adapter: it resolves λ_central, the gen grid, and (in the +reco path) R / N_gen from the datacard, holds a ``SigmaGenModel`` as +``self.core``, and adds Steps 3–4. ``_sigma_gen_at`` / ``_sigma_YqT_native_at`` +alias the core; its physics attributes (``eff_central``, ``np_model``, +``sigma_ns``, …) are reachable as ``model.`` via ``__getattr__``. + +NP physical validity: a wrong-sign point (λ2_ν < 0, or the λ2_eff < 0 +divergence) makes the form factors anti-damping and the differential σ(qT) +oscillate negative, but the qT-rebin into the coarse gen grid averages that away +— so the binned σ_gen, σ_reco and likelihood stay smooth and positive, and the +fit gets no signal from the likelihood to avoid the unphysical region. +Enforcement is the rabbit ``Regularizer`` ``np_damping_wall.NPDampingWall`` (a +one-sided hinge on the λ, via ``-r``; hardness via ``--regularizationStrength``), +encoding the tanh_2 damping conditions (CS: λ2_ν ≥ 0, λ4_ν ≥ 0; TMD: λ2_Y ≥ 0 +and 3·λ∞²·λ4 + λ2_Y³ ≥ 0). Postfit-only validity checks (``np_damping_ok``, +``spectrum_negativity``) live in ``param_model_diagnostics``. + +Indices: Q, Y, qT are the btgrid axes (mass / rapidity / qT); g = gen bin +(ptVGen, absYVGen); b = reco bin (ptll, yll, cosThetaStarll_quantile, +phiStarll_quantile). λ splits into λ_eff (F_eff) and λ_ν (γ_ν^NP). ============================================================================= Step 1 — σ_resum(λ; g): the resummed prediction from the bT grid @@ -32,7 +48,7 @@ · exp[ C_ν(Q, Y, qT; bT) · γ_ν^NP(b*(bT); λ_ν) ] · F_eff(Y; b*(bT); λ_eff) -Where each factor lives — bare bT vs the b*-frozen b̄T: +Each factor, bare bT vs the b*-frozen b̄T: bT bare bT, the Hankel integration variable. Enters ONLY the measure factor bT and the Fourier kernel J₀(qT·bT). The leading @@ -57,21 +73,19 @@ (no Q/qT). λ-dependent. Only γ_ν^NP and F_eff carry λ; everything else (bT, J₀, I_pert, C_ν, b*) is -λ-independent and precomputed once — the bT·J₀ kernel, the bT Simpson weights, -and the arctan_Q² Q-integration weights are all built at construction. - -The default evaluation uses the memory-factorized layout (deduplicated -(I_pert, C_ν) rows + J₀ kernel on the unique-qT grid + Simpson-as-matmul), -which is numerically equivalent to the per-bin (Nbins, Nbt) layout (≲1e-14 -rel., floating-point summation order only) but ~6× smaller — required to fit -a 32 GB GPU. The legacy_recon=1 spec token restores the legacy layout. +λ-independent and precomputed at construction (the bT·J₀ kernel, bT Simpson +weights, arctan_Q² Q-weights). The reconstruction uses a memory-factorized layout +(deduplicated (I_pert, C_ν) rows + J₀ on the unique-qT grid + Simpson-as-matmul), +~6× smaller than the dense (Nbins, Nbt) and small enough for a 32 GB GPU; it is +memoized in a ``.npz`` next to ``combined_btgrid.pkl`` (staleness auto-detected) +so repeat constructions skip the ~18 GB raw load. (1b) Integrate over Q, then rebin onto the gen grid: σ_resum(λ; g) = rebin_{qT→ptVGen, |Y|→absYVGen} [ ∫_{Q_lo}^{Q_hi} dQ σ(Q, Y, qT; λ) ] The Q integral uses an arctan_Q² Simpson rule (the x = arctan((Q²−q0²)/(q0·Γ)) -transform flattens the Breit-Wigner Z peak). The result (Y, qT) is then rebinned +transform flattens the Breit-Wigner Z peak). The (Y, qT) result is rebinned (Simpson) onto the unfolding hist's gen edges: qT → ptVGen, and the signed btgrid Y axis folded into |Y| → absYVGen. The |Y| fold is valid because NP is Y-symmetric (F_eff depends on Y², γ_ν^NP doesn't depend on Y at all). @@ -85,17 +99,27 @@ σ_ns(g) = rebin_{qT→ptVGen, |Y|→absYVGen} [ σ_DYTurbo^FO − σ_SCETlib-sing^FO ] The fixed-order matching adds the nonsingular piece σ_ns = (DYTurbo fixed order) -− (SCETlib singular fixed order), read straight from the original FO inputs -(the ``…_nnlo_sing…combined.pkl`` and the DYTurbo ``results_…scetlibmatch.txt``), +− (SCETlib singular fixed order), from the original FO inputs (the +``…_nnlo_sing…combined.pkl`` and the DYTurbo ``results_…scetlibmatch.txt``), Q-windowed to [Q_lo, Q_hi], |Y|-folded, zeroed below qt_cutoff, and summed onto the (ptVGen, absYVGen) gen bins (see :func:`compute_nonsingular_gen`). -σ_ns is NP-INDEPENDENT (the same for every λ) and is added at GEN level, so it -folds through the same response R as σ_resum in Step 3. Because the fit uses a -ratio (Step 3), σ_ns correctly DILUTES the NP variation where the FO dominates -(high qT). σ_ns is ALWAYS included (it is what the histmaker nominal carries); -for resum-only diagnostics subtract the exposed ``sigma_ns`` from -``sigma_gen_central`` (and re-fold with ``R`` for reco level). +σ_ns is NP-INDEPENDENT (the same for every λ) and added at GEN level, so it folds +through the same response R as σ_resum in Step 3. Because the fit uses a ratio +(Step 3), σ_ns DILUTES the NP variation where the FO dominates (high qT). σ_ns is +ALWAYS included (it is what the histmaker nominal carries); for resum-only +diagnostics subtract the exposed ``sigma_ns`` from ``sigma_gen_central`` (and +re-fold with ``R`` for reco level). + +Known limitation — qT > 100 GeV truncation. Both σ_gen ingredients stop at +qT = 100 (the bT grid and the DYTurbo FO input), so σ_gen(λ; g) ≡ 0 above 100, +while R and N_gen lump the full unbounded qT > last_edge into the (last_edge, 100] +overflow (PTVGEN_OVERFLOW_EDGE, see :mod:`response_matrix`). The under-fed overflow +makes the top in-range reco bin (ptll [37, 44]) ~3% low vs the histmaker nominal +and trips the agreement guard, but does NOT bias the fit: the missing piece is +fixed-order, NP-independent, so it cancels in the Step-4 ratio. A high-ptll +extension would need DYTurbo FO past 100, higher bT-grid/overflow ceilings, and +the histmaker correction remade (it shares the truncated inputs). ============================================================================= Step 3 — σ_reco(λ; b): fold gen → reco through the response matrix @@ -104,15 +128,11 @@ P(b | g) = R_raw(b, g) / N_gen(g) (efficiency × migration) σ_reco(λ; b) = Σ_g P(b | g) · σ_gen(λ; g) -where, in these expressions, - g = the flattened GEN bin (ptVGen, absYVGen) — the grid σ_gen from Steps - 1–2 lives on (boson qT and |Y|), summed over by Σ_g; - b = the flattened RECO bin (ptll, yll, cosThetaStarll_quantile, - phiStarll_quantile) — the measured dilepton observables σ_reco lives on. -P(b | g) is the gen→reco mapping (one reco column per gen bin), so σ_reco(λ; b) -is just σ_gen pushed through the detector. The Σ_g is the matvec -``tf.linalg.matvec(self.R, σ_gen_flat)``. This step is pure detector folding — -no λ_central and no ratio enter here; that is Step 4. +g = gen bin (ptVGen, |Y|), summed over by Σ_g; b = reco bin (ptll, yll, +cosThetaStarll_quantile, phiStarll_quantile). P(b | g) is the gen→reco map (one +reco column per gen bin), so σ_reco is σ_gen pushed through the detector; Σ_g is +``tf.linalg.matvec(self.R, σ_gen_flat)``. Pure detector folding — no λ_central or +ratio here (that is Step 4). Factors (all loaded by :mod:`response_matrix`; see it for the hist mechanics): @@ -135,14 +155,12 @@ Taking UL of R instead would discard the angular partition and inflate the closure ~15×. -Why N_gen and not the reco-passing marginal Σ_b R_raw(b, g): R is -post-reco-selection so that marginal already carries efficiency, and dividing -by it cancels efficiency (migration-only) — closes far worse, since efficiency -is strongly gen-dependent (ε ≈ 0.07–0.54 across gen bins on the current file). -N_gen(g) is the true generated total, so P = R_raw/N_gen is the -theory-independent gen→reco map. Pre-FSR: σ_gen, R, and N_gen must all sit at -the same QCD/boson gen level (the postfsr variants in the file close ~1% worse -— FSR mismatch). +Why N_gen and not the reco-passing marginal Σ_b R_raw(b, g): that marginal +already carries efficiency, so dividing by it cancels efficiency (migration-only) +and closes far worse — efficiency is strongly gen-dependent (ε ≈ 0.07–0.54 here). +N_gen(g), the true generated total, makes P = R_raw/N_gen the theory-independent +gen→reco map. σ_gen, R, N_gen must all sit at the same (pre-FSR) gen level; the +postfsr variants close ~1% worse. ============================================================================= Step 4 — rnorm(b, proc): the per-reco-bin variation handed to rabbit @@ -151,13 +169,11 @@ ratio(b) = σ_reco(λ; b) / σ_reco(λ_central; b) rnorm(b, proc) = 1 + (ratio(b) − 1) · [proc is signal] (1 in every other proc) -This is the only object that leaves the model: compute() returns rnorm(b, proc), -and rabbit multiplies the signal process's reco template (reco bin b) by it, -leaving every other process at 1. Dividing by σ_reco(λ_central) cancels the -event-count↔cross-section scale and any overall normalization, so rnorm carries -purely the SHAPE of the NP variation per reco bin — Steps 1–3 build the absolute -σ_reco(λ; b), and Step 4 reduces it to the bin-by-bin template scaling the fit -needs. σ_reco(λ_central) is precomputed once at construction as the denominator. +The only object that leaves the model: compute() returns rnorm(b, proc), and +rabbit multiplies the signal reco template by it (other processes stay at 1). +Dividing by σ_reco(λ_central) (precomputed at construction) cancels the +event-count↔σ scale and overall normalization, so rnorm carries purely the SHAPE +of the per-reco-bin NP variation. ----------------------------------------------------------------------------- Parameters and inputs @@ -169,74 +185,26 @@ F_eff (TMD-effective): lambda2, lambda4, lambda6, delta_lambda2, lambda_inf λ_central is read from the fit-tensor's metadata, where the histmaker stored -the SCETlib correction's NP runcard (see :mod:`scetlib_lambda_central`). The -np_model and np_model_nu strings are fixed at construction (from λ_central). -All λ values are TF Variables, differentiable in the fit. +the SCETlib correction's NP runcard (see :mod:`lambda_central`). The np_model and +np_model_nu strings are fixed at construction (from λ_central). All λ values are +TF Variables, differentiable in the fit. ============================================================================= -Getting the postfit Hessian / covariance (uncertainties on λ) +Postfit Hessian / covariance (uncertainties on λ) ============================================================================= -The fit floats λ fine, but rabbit's postfit covariance step -``loss_val_grad_hess`` → ``t2.jacobian(grad, x)`` differentiates through the bT -fold once per fit parameter (~3754). pfor cannot see that the fold depends on -the parameters only through the ≤8 λ, so it re-materializes the internal -(Ng × Nbt) ≈ 8.75 GB bT slab for EVERY parameter → ~33 TB → OOM. - -Fix (this module): a "straight-through" surrogate that keeps the exact ratio -VALUE but exposes only a compact quadratic in the ≤8 λ to autodiff: - - ratio~(λ) = stop_gradient(ratio) + J·d + ½ dᵀ K d , d = λ − stop_gradient(λ) - -with J = dratio/dλ ([Nreco, nλ]) and K = d²ratio/dλ² ([Nreco, nλ, nλ]) computed -by forward-mode AD (≤8 / ≤64 fold passes, NOT tiled). d is identically 0 so the -value is unchanged, but ∂d/∂λ = I, so ∂ratio~/∂λ = J and ∂²ratio~/∂λ² = K: -rabbit's jacobian gets the exact derivatives while the big slab stays inside -stop_gradient and never enters the differentiated graph (33 TB → a few MB). -Implemented in ``_ratio_straightthrough`` (+ ``_ratio_compact_jac``, -``_ratio_compact_hess``); selected in ``compute()`` by env flags. - -Two-pass recipe (rabbit still computes the Hessian; NO rabbit changes): - - 1. Fit, no Hessian → postfit: - rabbit_fit.py --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel \ - --noHessian -o fit/ - - 2. Covariance at that postfit (no refit), straight-through ON: - rabbit_fit.py --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel \ - hessian_straightthrough=1 hessian_gn=1 \ - --externalPostfit fit/fitresults.hdf5 --externalPostfitResult nominal \ - --noFit -t 0 --pseudoData nominal -o cov/ - (do NOT pass --noHessian; do NOT pass --eager.) ~5 min, no OOM. - Uncertainties are sqrt(diag(cov)) in cov/fitresults.hdf5 (results_nominal). - -Switches (both OFF by default → the fit path is unchanged). They are spec -tokens inside the --paramModel spec (shown above) — the spec is stored in the -fitresults meta_info.args, so the configuration is recorded in the output -(env vars are NOT supported; all model knobs go through the spec): - hessian_straightthrough=1 use the straight-through path in compute() - hessian_gn=1 Gauss-Newton: keep J only, drop the K term -WARNING: hessian_straightthrough=1 WITHOUT hessian_gn=1 is full-K mode — -correct in principle (needed for real/toy data) but currently INFEASIBLE at -full grid scale (the 64 nested-FA passes unroll into one graph, ~TB peak → -OOM-kill). Until a precomputed chunked-K path is implemented, always pass -hessian_gn=1 (exact for Asimov). - -GN vs full-K. The Poisson Hessian is - H_ij = Σ_b [ (n_b/μ_b²) J_bi J_bj + (1 − n_b/μ_b) K_bij ]. -For ASIMOV data the residual (1 − n/μ) = 0, so the K term vanishes and GN (J -only) is EXACT — drop the K term with hessian_gn=1. Real/toy data need the K -term; full-K (``_ratio_compact_hess``) matches the exact reverse-mode Hessian. -Full-K requires the ``btgrid_tf._frozen_eq_zero`` guard: the λ_inf==0 / den==0 -masks compare a differentiated tensor, which makes TF's nested-forward-mode AD -raise ``IndexError`` on the ``Equal`` op; freezing the comparison input -(``tf.equal(stop_gradient(x), 0)``) drops the tangent into ``Equal`` without -changing any value or derivative (it is a measure-zero ``tf.where`` boundary). -GN remains the default for Asimov (exact and cheaper — 8 vs 72 fold passes). - -WARNING: do NOT pass hessian_straightthrough=1 during the FIT. The -surrogate recomputes J(/K) on every compute() call — fine for the one-shot -covariance pass, but it would cripple the minimizer (many gradient/HVP evals). +rabbit's covariance step differentiates the bT fold once per fit parameter +(~3754) and can't see it depends only on the ≤8 λ, re-materializing the +(Ng × Nbt) ≈ 8.75 GB slab per parameter → ~33 TB → OOM. Fix: a straight-through +surrogate (``_ratio_straightthrough`` + ``_ratio_compact_jac`` / +``_ratio_compact_hess``) exposing only a compact quadratic in λ to autodiff +(J = dratio/dλ, K = d²ratio/dλ² by forward-mode AD) while the exact ratio value +stays under stop_gradient, so the big slab never enters the differentiated graph. +Enable via the ``hessian_straightthrough=1`` / ``hessian_gn=1`` spec tokens and a +two-pass run (fit ``--noHessian``, then covariance at the postfit with +``--externalPostfit --noFit``). GN (``hessian_gn=1``, drops the K term) is exact +for Asimov and the default; full-K is for real/toy data. Do NOT enable during the +fit. Full derivation, GN-vs-full-K, and the exact commands: ``docs/HESSIAN_PLAN.md``. """ import os @@ -246,74 +214,47 @@ import tensorflow as tf from rabbit.param_models.param_model import ParamModel -from wremnants.postprocessing.scetlib_np import ( - btgrid_cache, -) -from wremnants.postprocessing.scetlib_np import btgrid_integrate as fz_int from wremnants.postprocessing.scetlib_np import btgrid_tf as fz_tf from wremnants.postprocessing.scetlib_np import lambda_central as scetlib_lambda_central -from wremnants.utilities import common as wrem_common -from wremnants.utilities.data_paths import getDataPath -_NONSING_FO_SING_DEFAULT = os.path.join( - wrem_common.data_dir, - "TheoryCorrections", - "inclusive_Z_COM13_CT18Z_N3+0LL_lattice_lambda4bugfix_fine_nnlo_sing_combined.pkl", -) -_NONSING_DYTURBO_DEFAULT = os.path.join( - wrem_common.data_dir, - "TheoryCorrections", - "results_z-2d-nnlo-vj-CT18ZNNLO-{scale}-scetlibmatch.txt", +# Physics core (Steps 1–2) lives in :mod:`sigma_gen`; this module is the +# datacard/rabbit adapter (Steps 3–4), holding a :class:`SigmaGenModel` as +# ``self.core``. The λ-name tuples, σ_ns builder, and default-btgrid helper are +# re-exported here for backward compat (older imports referenced them on this +# module). ``fz_tf`` is still needed for the reco-side tensor dtype (``fz_tf.DTYPE``). +from wremnants.postprocessing.scetlib_np.params import DEFAULT_PRIOR_SIGMAS +from wremnants.postprocessing.scetlib_np.sigma_gen import ( # noqa: F401 + ALL_PARAMS, + EFF_PARAMS, + GNU_PARAMS, + SigmaGenModel, + _NONSING_DYTURBO_DEFAULT, + _NONSING_FO_SING_DEFAULT, + _default_btgrid_dir, + compute_nonsingular_gen, ) -_BTGRID_SUBDIR = ("scetlib_np", "Z_COM13_CT18Z_N3p0LL_btgrid_fineall") -_DISCRETE_NP_SUBSTRING = "scetlibnp" -# Ordered list of the v1 continuous λ. CS-side first, then TMD-effective. -GNU_PARAMS = ("lambda2_nu", "lambda4_nu", "lambda_inf_nu") -EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") -ALL_PARAMS = GNU_PARAMS + EFF_PARAMS +_DISCRETE_NP_SUBSTRING = "scetlibnp" # Positivity floor for the per-bin reco ratio in compute() (see its docstring). # A pathological λ (e.g. λ4 < 0 with the bounded-tanh model) can drive σ_reco — -# and hence the predicted signal yield — negative, which makes the Poisson NLL -# NaN and stalls the minimizer. We soft-floor the ratio to a small positive value -# so a bad point becomes a LARGE-BUT-FINITE penalty the fit can back off from, -# with a non-zero gradient through the transition (softplus, not a hard clamp). -# RATIO_FLOOR_SCALE — softplus transition width. Chosen FAR below any physical -# response so healthy ratios (~0.9–1.1, and every validated λ-variation) pass -# through to machine precision: scale·softplus(r/scale) == r for r ≫ scale. +# and the predicted signal yield — negative, giving NaN Poisson NLL and stalling +# the minimizer. Soft-floor the ratio to a small positive value (softplus, not a +# hard clamp), so a bad point is a LARGE-BUT-FINITE penalty with non-zero gradient. +# RATIO_FLOOR_SCALE — softplus transition width. FAR below any physical response +# so healthy ratios (~0.9–1.1, every validated λ-variation) pass to machine +# precision: scale·softplus(r/scale) == r for r ≫ scale. # RATIO_FLOOR_MIN — hard positive ground: softplus underflows to exactly 0 for -# the extreme (r ~ -1e43) case, so this keeps the yield strictly > 0 (no NaN). +# the extreme (r ~ -1e43) case, keeping the yield strictly > 0 (no NaN). RATIO_FLOOR_SCALE = 1.0e-4 RATIO_FLOOR_MIN = 1.0e-9 -# Recommended Gaussian prior widths for the SCETlib NP λ parameters -DEFAULT_PRIOR_SIGMAS = { - "lambda2_nu": 0.10, - "lambda2": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average - "lambda4": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average - "delta_lambda2": 0.20, # 0 ± 0.20 wide default (no theorist value yet) -} - - -def _default_btgrid_dir(): - base = getDataPath(fallback="/scratch/submit/cms/wmass/NanoAOD") - return os.path.join(os.path.dirname(base), *_BTGRID_SUBDIR) - - -def _load_lambda_central_file(path): - """ - Load a λ_central override from a JSON or YAML file. - """ - return scetlib_lambda_central.load_lambda_central_file(path) - - def _crop_R_to_fit(R, R_reco_axes, fit_reco_axes, tol=1e-9): """Crop R's trailing reco bins so its reco shape matches the fit. R's reco binning is typically a superset of the fit's (e.g. R has one extra overflow ptll bin past the fit's last edge). For each reco axis, - require that R's leading edges match the fit's edges; crop R along that + require R's leading edges to match the fit's edges and crop R along that axis to keep only the matching bins. """ if len(R_reco_axes) != len(fit_reco_axes): @@ -348,7 +289,7 @@ def _R_info_from_auxiliary(indata): setupRabbit extracts R (and the gen-total N_gen, reco/gen axis names + edges) once from the unfolding histmaker output and embeds it in the fit input via rabbit's ``add_auxiliary``; rabbit exposes it as ``FitInputData.auxiliary``. - The model reads R ONLY from there — one source, one path — so R is always + The model reads R ONLY from there (one source, one path), so R is always consistent with the run that produced the datacard. The returned dict matches :func:`response_matrix.load_R`'s shape for the keys this model consumes (``R``, ``N_gen``, and ``reco_axes`` / ``gen_axes`` as ordered @@ -378,97 +319,6 @@ def _R_info_from_auxiliary(indata): ) -def _bin_sum_matrix(src_centers, target_edges, tol=1e-6): - """(N_target, N_src) 0/1 matrix that SUMS bin-integrated source bins whose - centre falls in each target bin. Source bins outside all target bins are - dropped — a natural truncation to the target range (qT>ptVGen_max, |Y|>absY_max).""" - src = np.asarray(src_centers, dtype=np.float64) - edges = np.asarray(target_edges, dtype=np.float64) - W = np.zeros((edges.size - 1, src.size), dtype=np.float64) - for i in range(edges.size - 1): - m = (src >= edges[i] - tol) & (src <= edges[i + 1] + tol) - W[i, m] = 1.0 - return W - - -def compute_nonsingular_gen( - fo_sing_path, - dyturbo_path, - gen_axes_meta, - charge=0, - q_lo=60.0, - q_hi=120.0, - qt_cutoff=1.0, - dyturbo_axes=("Q", "Y", "qT"), -): - """Nonsingular FO term on the model gen grid (NptVGen, NabsYVGen). - - The fixed-order/DYTurbo matching adds a NP-INDEPENDENT piece to σ_gen: - σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns , - where the nonsingular is read straight from the original fixed-order inputs: - σ_ns = (DYTurbo fixed order) − (SCETlib singular fixed order) - — exactly the ``-hfo_sing + hfo`` that ``read_matched_scetlib_hist`` forms. - ``fo_sing_path`` is the SCETlib singular ``…_nnlo_sing…combined.pkl``; - ``dyturbo_path`` is the DYTurbo fixed-order ``results_…scetlibmatch.txt`` - (use ``{scale}`` → mur1-muf1 for the central). The nonsingular is zeroed below - ``qt_cutoff`` (as make_theory_corr does), Q-windowed to [q_lo, q_hi], |Y|-folded, - and projected onto the coarse (ptVGen, absYVGen) gen bins by SUMMING the - bin-integrated native bins. - """ - from wremnants.utilities.io_tools import input_tools - from wums import boostHistHelpers as hh - - def _central(h): - if "vars" in h.axes.name: - names = list(h.axes["vars"]) - idx = 0 - for c in ("central", "pdf0", "nominal"): - if c in names: - idx = names.index(c) - break - h = h[{"vars": idx}] - return h - - # SCETlib singular fixed order, and DYTurbo fixed order, from their own files. - dyturbo_path = ( - dyturbo_path.format(scale="mur1-muf1") - if "{scale}" in dyturbo_path - else dyturbo_path - ) - hfo_sing = _central(input_tools.read_scetlib_hist(fo_sing_path, charge=charge)) - hfo = input_tools.read_dyturbo_hist( - [dyturbo_path], axes=list(dyturbo_axes), charge=charge - ) - if "vars" in hfo.axes.name: - hfo = _central(hfo) - - # Align shared physics axes (DYTurbo is coarser), then σ_ns = DYTurbo − singular. - for ax in ("Y", "Q", "qT"): - if ax in set(hfo.axes.name) & set(hfo_sing.axes.name): - hfo, hfo_sing = hh.rebinHistsToCommon([hfo, hfo_sing], ax) - nonsing_h = hh.addHists(-1.0 * hfo_sing, hfo, flow=False, by_ax_name=False) - - if "charge" in nonsing_h.axes.name: - nonsing_h = nonsing_h[{"charge": sum}] - # Q-window: slice(...,sum) sums ONLY the in-range Q bins (no underflow leak). - Qe = np.asarray(nonsing_h.axes["Q"].edges, dtype=np.float64) - qi = int(np.argmin(np.abs(Qe - q_lo))) - qj = int(np.argmin(np.abs(Qe - q_hi))) - nonsing_h = nonsing_h[{"Q": slice(qi, qj, sum)}] - nonsing_h = hh.makeAbsHist(nonsing_h, "Y") # signed Y -> |Y| - - qT_c = np.asarray(nonsing_h.axes["qT"].centers, dtype=np.float64) - absY_c = np.asarray(nonsing_h.axes["absY"].centers, dtype=np.float64) - v = nonsing_h.project("qT", "absY").values(flow=False) # (qT, absY) - v[qT_c < qt_cutoff, :] = 0.0 # zero the nonsingular below the cutoff - - ptV_edges = np.asarray(gen_axes_meta[0][1], dtype=np.float64) - absY_edges = np.asarray(gen_axes_meta[1][1], dtype=np.float64) - Wp = _bin_sum_matrix(qT_c, ptV_edges) # (NptVGen, NqT) - Wa = _bin_sum_matrix(absY_c, absY_edges) # (NabsYVGen, NabsYsrc) - return Wp @ v @ Wa.T # (NptVGen, NabsYVGen) - - class SCETlibNPParamModel(ParamModel): @classmethod @@ -508,11 +358,16 @@ def __init__( nonsingular_fo_sing: str = _NONSING_FO_SING_DEFAULT, nonsingular_dyturbo: str = _NONSING_DYTURBO_DEFAULT, nonsingular_qt_cutoff: float = 1.0, - legacy_recon: bool = False, xparam_default: Optional[str] = None, hessian_straightthrough: bool = False, hessian_gn: bool = False, gen_level: bool = False, + check_agreement: bool = True, + check_agreement_threshold: float = 0.005, + check_agreement_strict: bool = False, + check_agreement_min_yield: float = 0.0, + np_model_fit: Optional[str] = None, + np_model_nu_fit: Optional[str] = None, **kwargs, ): """Construct the ParamModel. @@ -525,11 +380,11 @@ def __init__( indata rabbit's input-data structure (passed by ``ph.load_models``). The reco×gen response matrix R (and the gen-total N_gen, axis names + - edges) is read from ``indata.auxiliary["scetlib_np"]`` — embedded in + edges) is read from ``indata.auxiliary["scetlib_np"]``, embedded in the datacard by setupRabbit from a ``mz_dilepton.py --unfolding`` histmaker output (see :func:`_R_info_from_auxiliary` and - :mod:`response_matrix`). One source, one path: there is no file-path - argument; R always comes from the fit input it is consistent with. + :mod:`response_matrix`). There is no file-path argument; R always + comes from the fit input it is consistent with. btgrid_dir Directory holding the SCETlib bT-grid ``combined_btgrid.pkl``. Defaults (when None) to the shared data-area copy next to NanoAOD @@ -538,7 +393,7 @@ def __init__( at non-subMIT sites. lambda_central Dict with two sub-dicts ``eff_params`` and ``gnu_params`` (same - shape as returned by :func:`scetlib_lambda_central.read_lambda_central`). + shape as returned by :func:`lambda_central.read_lambda_central`). By default λ_central is auto-detected from ``indata.metadata``; pass this only to override or to support hand-built indata that lacks metadata. @@ -555,50 +410,32 @@ def __init__( priors Enable Gaussian priors on the λ parameters (spec token ``priors=1``). Rabbit applies priors whenever a ParamModel - *declares* ``prior_sigmas`` — there is no rabbit-side CLA (the - old ``--paramModelPriors`` flag was dropped in WMass/rabbit#133); + *declares* ``prior_sigmas``; there is no rabbit-side CLA (the + old ``--paramModelPriors`` flag was dropped in WMass/rabbit#133), the model itself decides. This token IS that decision: only when set does the model declare ``prior_sigmas``. Default off → everything floats free. prior_sigmas - Per-name override for the Gaussian prior σ on each parameter — + Per-name override for the Gaussian prior σ on each parameter: a Mapping, or as spec token the comma-separated string form ``prior_sigmas=lambda2=0.3,delta_lambda2=nan`` (same format as - ``xparam_default``). Defaults come from ``DEFAULT_PRIOR_SIGMAS``: - - lambda2_nu : 0.10 - lambda2 : 0.50 (symmetric approx of +0.6/-0.4) - lambda4 : 0.50 (symmetric approx of +0.6/-0.4) - - All other params default to ``NaN`` → no prior, float free; in - practice they are expected to be frozen with rabbit's - ``--freezeParameters`` until the theorist provides priors for - them. Pass ``np.nan`` here to free a constrained param, or a - finite value to add a prior on one that defaults to NaN. - Only meaningful together with ``priors=1``; ignored (with a - warning) otherwise. - Prior mean for each param is ``self.xparamdefault`` (the - runcard's λ_central). + ``xparam_default``). Defaults come from ``DEFAULT_PRIOR_SIGMAS`` nonsingular_fo_sing, nonsingular_dyturbo Paths to the σ_ns inputs (SCETlib singular pkl / DYTurbo scetlibmatch txt); default to the wremnants-data - TheoryCorrections copies. σ_ns is always included — the matched + TheoryCorrections copies. σ_ns is always included: the matched σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns is what the histmaker nominal carries; resum-only diagnostics subtract ``sigma_ns``. nonsingular_qt_cutoff Low-qT cutoff (GeV) below which σ_ns is zeroed (the FO−singular difference is numerically unreliable at tiny qT). - legacy_recon - Use the legacy per-bin (Nbins, Nbt) reconstruction layout instead - of the default memory-factorized one (numerically equivalent to - ≲1e-14 rel; for parity checks only). xparam_default Comma-separated ``name=value,...`` string shifting the fit START - (and the prior mean) off the runcard's λ_central — for closure / + (and the prior mean) off the runcard's λ_central, for closure / injection tests. The truth (ratio denominator) is NOT moved. hessian_straightthrough Expose compact λ-derivatives (J, optionally K) to autodiff while - keeping the exact value — for the one-shot two-pass covariance + keeping the exact value, for the one-shot two-pass covariance recipe ONLY (see the module docstring); never set during a fit. hessian_gn With ``hessian_straightthrough``: Gauss-Newton, keep J and drop @@ -611,7 +448,38 @@ def __init__( gen binning is read from the single fit channel's axes, so no ``scetlib_np`` auxiliary / R / N_gen is needed. Used for the direct-theory σUL closure (this ParamModel as the λ model, fit - against an injected gen-level σUL pseudodata). + against injected gen-level σUL pseudodata). + check_agreement + Run the in-fit reco agreement guard at construction (default ON; + spec token ``check_agreement=0`` disables). It compares the model's + σ_reco(λ_central) SHAPE to the card's signal nominal template + and trips if ANY bin's |σ_reco(λ_c)/card_nominal − 1| exceeds + ``check_agreement_threshold``, catching the misuse where the + model's λ_central baseline doesn't match the card. + check_agreement_threshold + Per-bin trip threshold for ``check_agreement`` (fractional; + default 0.005 = 0.5%). + check_agreement_strict + Raise instead of warn when the guard trips (spec token + ``check_agreement_strict=1``). Default warns (lists worst bins). + check_agreement_min_yield + Optional reference-yield floor (fraction of the max reco bin) below + which bins are ignored by the guard, suppressing sparse/near-empty + corner bins where a shape residual is meaningless. Default 0 (every + bin counts). + np_model_fit, np_model_nu_fit + Optional override of the F_eff / γ_ν functional form used for the + NUMERATOR σ(λ) (spec tokens ``np_model_fit=...`` / ``np_model_nu_fit= + tanh_6``). The DENOMINATOR (central σ(λ_c), the ratio's reference) is + ALWAYS the card's form — fixed by λ_central, immutable — so the model + stays consistent with the histmaker template. Default (None) → the + numerator uses the card form too (rnorm(λ_c)=1, unchanged behaviour). + Setting e.g. ``np_model_nu_fit=tanh_6`` makes the fit predict in + tanh_6 while transporting from the tanh_2 baseline: + rnorm = σ^(fit)(λ) / σ^(card)(λ_c). NOTE this is a model change — + validate the fit form against a matching SCETlib reference before + trusting results, and ``lambda6_nu`` (the tanh_6 b⁶ coefficient) is a + normal fittable λ (default 0, inert under tanh_2). """ self.indata = indata @@ -621,27 +489,18 @@ def __init__( self._check_discrete_np_double_counting() # ---- λ_central - # Three sources of λ_central, in priority order: - # 1. ``lambda_central=`` spec token — path to a JSON or YAML - # file with ``eff_params`` and ``gnu_params``. Overrides the - # metadata auto-detect; useful when the upstream SCETlib pkl isn't - # accessible (e.g. a colleague's input). - # 2. ``lambda_central`` constructor arg (explicit dict, programmatic). - # 3. Auto-detect from the fit hdf5's theoryCorr → upstream pkl. + # Anchor point of the model; must match the SCETlib NP runcard the input's + # theory correction was built with (so rnorm(λ_central) == 1). Two sources, + # priority order: + # 1. ``lambda_central`` constructor arg (explicit dict) — standalone + # diagnostic scripts; not reachable from the rabbit CLI. + # 2. Auto-detect from the fit hdf5's propagated histmaker metadata — the + # production path. No CLI override by design: a mismatched anchor + # silently biases the fit, so an input lacking the metadata must be + # remade, not overridden. lambda_central_source = ( "constructor-arg" if lambda_central is not None else None ) - if isinstance(lambda_central, str): - # CLI token (lambda_central=): RECOMMENDED override route — - # the --paramModel spec is stored in the fitresults meta, so the - # override is recorded in the output (env var/dict are not). - lc_path = lambda_central - lambda_central = _load_lambda_central_file(lc_path) - lambda_central_source = f"cli-file:{lc_path}" - print( - f"[SCETlibNPParamModel] λ_central from CLI file {lc_path!r}", - flush=True, - ) if lambda_central is None: # Auto-detect from indata.metadata (loaded by rabbit's # FitInputData from the input HDF5's "meta" group). @@ -666,127 +525,19 @@ def __init__( self.lambda_central_source = lambda_central_source - self.eff_central = dict(lambda_central["eff_params"]) - self.gnu_central = dict(lambda_central["gnu_params"]) - self.np_model = self.eff_central["np_model"] - self.np_model_nu = self.gnu_central["np_model_nu"] - - # ---- btgrid + dense layout - grid = btgrid_cache.load(btgrid_dir) - # Sanitize non-finite bt-grid cells. Kinematically-forbidden points - # (x = (Q/Ecm)·e^|Y| ≥ 1, e.g. extreme forward Y near the Z peak) can - # come back as NaN from SCETlib instead of the physical 0. Their true - # cross section is zero, so replace NaN/inf with 0 → harmless zero-rows. - # Without this, dedup_grid_rows' hash-group verification fails (NaN != - # NaN). (For grids produced as condor shards the forbidden cells are - # simply absent and dense_index_map 0-fills them; a single-process local - # run instead writes them in as NaN, which is what this handles.) - for _key in ("I_pert", "C_nu"): - _arr = grid[_key] - if not np.isfinite(_arr).all(): - _nbad = int((~np.isfinite(_arr)).any(axis=-1).sum()) - np.nan_to_num(_arr, copy=False, nan=0.0, posinf=0.0, neginf=0.0) - print( - f"[SCETlibNPParamModel] sanitized {_nbad} non-finite " - f"{_key} bt-grid rows -> 0 (kinematically-forbidden cells)", - flush=True, - ) - idx_map = fz_int.dense_index_map(grid["bins"]) - self.Q_unique = idx_map["Q_unique"] - self.Y_unique = idx_map["Y_unique"] - self.qT_unique = idx_map["qT_unique"] - self.flat_idx = tf.constant(idx_map["flat_idx"], dtype=tf.int64) - - # Cache btgrid arrays as TF constants. - self.bT = tf.constant(grid["bT"], dtype=fz_tf.DTYPE) - self.b_bar = tf.constant(grid["b_bar"], dtype=fz_tf.DTYPE) - - # Per-bin qT and Y (from the bin tuple). - bins = grid["bins"] - qT_pb_np = np.array([b[2] for b in bins], dtype=np.float64) - self.qT_per_bin = tf.constant(qT_pb_np, dtype=fz_tf.DTYPE) - Y_pb_np = np.array([b[1] for b in bins], dtype=np.float64) - self.Y_per_bin = tf.constant(Y_pb_np, dtype=fz_tf.DTYPE) - - # F_eff depends on the bin only through Y (not Q or qT), and Y takes few - # distinct values across the grid. Precompute the unique-Y map so the - # reconstruction evaluates the NP transcendentals on NY rows and - # gathers, instead of recomputing identical rows for every (Q, qT). - Y_feff_unique_np, Y_feff_inv_np = np.unique(Y_pb_np, return_inverse=True) - Y_feff_inv_np = Y_feff_inv_np.reshape(-1).astype(np.int32) - self.Y_feff_unique = tf.constant(Y_feff_unique_np, dtype=fz_tf.DTYPE) - self.Y_feff_inverse_idx = tf.constant(Y_feff_inv_np, dtype=tf.int32) - - bT_simpson_w_np = fz_tf.simpson_weights(np.asarray(grid["bT"])) - self.bT_simpson_w = tf.constant(bT_simpson_w_np, dtype=fz_tf.DTYPE) - - # Reconstruction layout. Default: factorized (deduplicated rows + - # unique-qT J0 kernel + Simpson-as-matmul) — numerically equivalent to - # the legacy (Nbins, Nbt) layout (≲1e-14 rel., summation order only) - # but ~6x smaller, which is what lets the fit run on a 32 GB GPU. - # The legacy_recon=1 spec token selects the legacy path (parity - # checks). - self.factorized = not bool(legacy_recon) - - # Hessian straight-through switches (see the module docstring's - # two-pass recipe). Spec tokens hessian_straightthrough=1 / - # hessian_gn=1, recorded in the fitresults meta via the stored - # --paramModel spec. + # ---- Hessian straight-through switches (module docstring's two-pass + # recipe). Spec tokens hessian_straightthrough=1 / hessian_gn=1, recorded + # in the fitresults meta via the stored --paramModel spec. self._hess_st = bool(hessian_straightthrough) self._hess_gn = bool(hessian_gn) - if self.factorized: - dd = fz_tf.dedup_grid_rows( - grid["I_pert"][0], grid["C_nu"][0], Y_feff_inv_np - ) - self.I_pert_u = tf.constant(dd["I_u"], dtype=fz_tf.DTYPE) # (Nu, Nbt) - # C_ν via the second-level dedup: exp(C·g) runs on the small - # (Ncu, Nbt) table and is gathered — bit-identical, ~150x fewer - # transcendentals, no (Nu, Nbt) C constant on device. - self.C_nu_uu = tf.constant(dd["C_uu"], dtype=fz_tf.DTYPE) # (Ncu, Nbt) - self.c_of_u = tf.constant(dd["c_of_u"], dtype=tf.int32) - self.feff_idx_u = tf.constant(dd["feff_idx_u"], dtype=tf.int32) - # Per-bin index into the unique-qT axis. The bin qT values are by - # construction members of qT_unique, so searchsorted is an exact - # lookup (asserted). - qT_idx_np = np.searchsorted(idx_map["qT_unique"], qT_pb_np) - assert np.array_equal(idx_map["qT_unique"][qT_idx_np], qT_pb_np) - self.gather_idx = tf.constant( - np.stack([dd["row_uid"].astype(np.int64), qT_idx_np], axis=1), - dtype=tf.int32, - ) - # Drop the host-side dedup copies (the tf.constants own the data now). - del dd - # Weighted J0 kernel on the unique-qT grid, with the per-bin qT - # prefactor and the Simpson weights folded in: (NqT, Nbt). - K_u = fz_tf.build_bT_J0_kernel( - tf.constant(idx_map["qT_unique"], dtype=fz_tf.DTYPE), self.bT - ) - self.KwqT = ( - tf.constant(idx_map["qT_unique"], dtype=fz_tf.DTYPE)[:, tf.newaxis] - * K_u - * self.bT_simpson_w[tf.newaxis, :] - ) - else: - self.I_pert = tf.constant( - grid["I_pert"][0], dtype=fz_tf.DTYPE - ) # (Nbins, Nbt) - self.C_nu = tf.constant(grid["C_nu"][0], dtype=fz_tf.DTYPE) - # Precompute the bT·J0(qT·bT) kernel (λ-independent). - self.bT_J0_kernel = fz_tf.build_bT_J0_kernel(self.qT_per_bin, self.bT) - # Drop the ~17.5 GB host-side grid reference before TF graph building. - del grid - - # ---- Q-integration weights (arctan_Q² Simpson on Z mass window). - self.Q_weights = tf.constant( - fz_int.q_integrate_weights(self.Q_unique, Q_lo, Q_hi), - dtype=fz_tf.DTYPE, - ) - # ---- Gen/reco binning + (reco mode) the response matrix R. - # gen_level=1: the fit channel IS the gen (ptVGen, absY) binning, so - # there is NO response matrix and NO gen→reco fold — compute() returns - # the per-GEN-bin ratio σ_gen(λ)/σ_gen(λ_central) (Steps 1–2 only), and - # the scetlib_np auxiliary / N_gen are not required. + # ---- Gen/reco binning. Resolve the gen grid — and, in the reco path, the + # response matrix R (and gen-total N_gen) — from the datacard, then hand the + # gen edges to the physics core (Steps 1–2: btgrid integral → matched σ_gen + # on that grid). gen_level=1: the fit channel IS the gen (ptVGen, absY) + # binning, so NO response matrix and NO gen→reco fold — compute() returns the + # per-GEN-bin ratio σ_gen(λ)/σ_gen(λ_central) — and the scetlib_np auxiliary + # / N_gen are not required. self.gen_level = bool(gen_level) if self.gen_level: gen_axes = self._fit_reco_axes(indata) @@ -796,31 +547,23 @@ def __init__( "with 2 gen axes (ptVGen, absY); got " f"{[n for n, _ in gen_axes]}" ) - self.R = None - self._R_raw = None - self._N_gen_flat = None + R_arr = None + N_gen_arr = None self.reco_shape = None self._reco_axes_meta = None - self._gen_axes_meta = gen_axes - self.gen_shape = tuple(len(e) - 1 for (_, e) in gen_axes) else: # ---- R matrix (read from the datacard's scetlib_np auxiliary) R_info = _R_info_from_auxiliary(indata) - # The fit-tensor's reco binning may differ from R's by trailing - # overflow bins (e.g. R has ptll [0, …, 44, 100] while the fit ends - # at 44). Crop R's trailing bins so the reco shape matches. + # Fit-tensor reco binning may differ from R's by trailing overflow bins + # (e.g. R has ptll [0, …, 44, 100] while the fit ends at 44). Crop R's + # trailing bins so the reco shape matches. fit_reco_axes = self._fit_reco_axes(indata) R_arr = _crop_R_to_fit(R_info["R"], R_info["reco_axes"], fit_reco_axes) - # Tighten the metadata to match the cropped R. self.reco_shape = R_arr.shape[: len(fit_reco_axes)] - self.gen_shape = R_arr.shape[len(fit_reco_axes) :] - N_reco = int(np.prod(self.reco_shape)) - N_gen = int(np.prod(self.gen_shape)) - # Raw response counts; normalized to a response below. - self._R_raw = tf.constant(R_arr.reshape(N_reco, N_gen), dtype=fz_tf.DTYPE) - # Gen-total denominator N_gen(g) from the xnorm hist ("prefsr"): the + gen_axes = R_info["gen_axes"] + # Gen-total denominator N_gen(g) from the xnorm hist ("prefsr"): # generated fiducial yield per gen bin (pre-reco-selection). Dividing R - # by this gives the theory-independent efficiency×migration response. + # by it gives the theory-independent efficiency×migration response. # REQUIRED: setupRabbit only embeds the response when the gen-total is # present, so N_gen should always be here; raise if not (a σ_gen(λ_c) # proxy would make the central closure circular — see module docstring). @@ -830,9 +573,7 @@ def __init__( "(gen-total). Rebuild the datacard from a histmaker output that " "carries the 'prefsr' xnorm hist." ) - self._N_gen_flat = tf.constant( - R_info["N_gen"].reshape(-1), dtype=fz_tf.DTYPE - ) + N_gen_arr = R_info["N_gen"] self._reco_axes_meta = [ (name, fit_axes[1]) for (name, fit_axes) in zip( @@ -840,128 +581,86 @@ def __init__( fit_reco_axes, ) ] - self._gen_axes_meta = R_info["gen_axes"] - - # ---- Rebin weights: btgrid (NY signed) → (NabsYVGen) via |Y| folding - # and (NqT) → (NptVGen). - absY_edges = self._gen_axes_meta[1][1] # absYVGen edges - ptVGen_edges = self._gen_axes_meta[0][1] # ptVGen edges - - # |Y| folding: σ(Y) is symmetric in Y so the absY-bin integral is - # 2·∫_{absY_lo}^{absY_hi} σ(Y) dY. Use Y >= 0 source samples and - # multiply by 2. - Y_pos_mask = self.Y_unique >= 0 - Y_pos = self.Y_unique[Y_pos_mask] - absY_rebin_pos = fz_int.rebin_weights(Y_pos, absY_edges, name="absY") - # Pad to full NY: zero on negative-Y columns. - W_absY = np.zeros((absY_edges.size - 1, self.Y_unique.size), dtype=np.float64) - W_absY[:, Y_pos_mask] = 2.0 * absY_rebin_pos - self.W_absY = tf.constant(W_absY, dtype=fz_tf.DTYPE) - - # qT rebin: btgrid qT (signed nonneg, NqT=141) → ptVGen edges. When - # load_R was built with ptVGen_overflow=True (default), ptVGen_edges ends - # in the overflow bin [last_gen_edge, PTVGEN_OVERFLOW_EDGE] (e.g. [44, 100]), - # so rebin_weights' last row Simpson-integrates the btgrid tail qT∈(44,100] - # into that overflow gen bin — matching R's gen-overflow column (true qT>44 - # migrating into the high-ptll reco bins). btgrid qT past the last edge - # (>100, beyond the grid) is dropped; negligible. Without the overflow - # column ptVGen_edges ends at 44 and that tail is simply truncated. - self.W_ptVGen = tf.constant( - fz_int.rebin_weights(self.qT_unique, ptVGen_edges, name="ptVGen"), - dtype=fz_tf.DTYPE, - ) - - # ---- Normalize the response, then cache σ_reco(λ_central). - # A response matrix must encode only the gen→reco *mapping*, not the - # MC's absolute gen spectrum. Normalize each gen column by the gen-total - # N_gen(g) (the xnorm "prefsr" hist — generated fiducial yield before - # reco selection) → P(b|g) = eff×migration (theory-independent): - # P(b|g) = R_raw(b,g) / N_gen(g) - # σ_reco(λ;b) = Σ_g P(b|g) · σ_gen(λ;g) - # σ_reco(λ_c;b) = Σ_g P(b|g) · σ_gen(λ_c;g) - # NB the reco-passing marginal Σ_b R_raw(b,g) is the WRONG normalizer: - # it already includes efficiency (R is post-reco-selection), so dividing - # by it cancels efficiency (migration-only) — closes far worse, ε is not - # flat in gen bin. We use the true gen-total N_gen instead. Because - # σ_reco_central then depends on σ_gen(λ_c) (it does NOT collapse to - # R_raw·1), the λ_central closure is a genuine test of the integral. - # Fallback: if the gen-total hist is absent, use σ_gen(λ_c) as a proxy - # for N_gen (∝ σ_gen^MC) — keeps efficiency but makes the closure - # circular (σ_gen cancels to R_raw·1). - # Native-binning Q-integrated reconstruction (NY, NqT) on the signed-Y / - # qT grid, BEFORE the |Y|-fold and qT-rebin — exposed so the native-binning - # validation can compare it to the SCETlib reference / external - # scetlib_run.factorize without the projection layer. - self.sigma_YqT_central = self._sigma_YqT_native_at( - self.eff_central, self.gnu_central - ) - # ---- Fixed-order/DYTurbo nonsingular term (NP-independent). - # σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns. Added at GEN level, so it folds - # through the same response R as the resummed piece. Because rnorm is a - # ratio, this correctly DILUTES the NP variation where the FO dominates - # (high qT). σ_ns is a constant (no λ dependence), always included — - # the matched σ_gen is what the histmaker nominal carries; resum-only - # diagnostics subtract self.sigma_ns instead of rebuilding the model. - _dy0 = ( - nonsingular_dyturbo.format(scale="mur1-muf1") - if (nonsingular_dyturbo and "{scale}" in nonsingular_dyturbo) - else nonsingular_dyturbo - ) - missing = [ - p for p in (nonsingular_fo_sing, _dy0) if not (p and os.path.exists(p)) - ] - if missing: - raise FileNotFoundError( - "The matched model needs the fixed-order inputs for " - "σ_ns = DYTurbo − SCETlib_singular, but these are missing:\n " - + "\n ".join(missing) - + "\nThey live under wremnants-data/data/TheoryCorrections (the " - "SCETlib singular …_nnlo_sing…combined.pkl and the DYTurbo " - "results_…scetlibmatch.txt). Pass nonsingular_fo_sing / " - "nonsingular_dyturbo to point at them." - ) - sigma_ns_np = compute_nonsingular_gen( - nonsingular_fo_sing, - nonsingular_dyturbo, - self._gen_axes_meta, - q_lo=Q_lo, - q_hi=Q_hi, - qt_cutoff=nonsingular_qt_cutoff, + self._gen_axes_meta = gen_axes + + # ---- Physics core (Steps 1–2): btgrid Hankel + arctan-Q² Simpson + + # |Y|/qT rebin + NP-independent fixed-order nonsingular → matched σ_gen(λ) + # on the (ptVGen, absYVGen) gen grid. Datacard-free, rebuildable standalone + # from (btgrid, λ_central, gen edges); the model delegates all σ_gen + # evaluation to it (see _sigma_gen_at / _sigma_YqT_native_at and __getattr__ + # forwarding of its physics attributes — eff_central, np_model, Y_unique, + # sigma_ns, …). + self.core = SigmaGenModel( + btgrid_dir=btgrid_dir, + lambda_central=lambda_central, + gen_axes=gen_axes, + Q_lo=Q_lo, + Q_hi=Q_hi, + nonsingular_fo_sing=nonsingular_fo_sing, + nonsingular_dyturbo=nonsingular_dyturbo, + nonsingular_qt_cutoff=nonsingular_qt_cutoff, ) - if sigma_ns_np.shape != tuple(self.gen_shape): + self.gen_shape = self.core.gen_shape + + # ---- Numerator form (default = card form). The denominator (central + # σ(λ_c) below) ALWAYS uses the card's form via the core, so consistency + # with the histmaker template is fixed; only the numerator σ(λ) may use a + # different form when np_model_(nu_)fit is given (see the constructor doc). + self._np_model_fit = np_model_fit or self.core.np_model + self._np_model_nu_fit = np_model_nu_fit or self.core.np_model_nu + if self._np_model_fit not in fz_tf.EFF_MODELS: raise ValueError( - f"nonsingular gen shape {sigma_ns_np.shape} != model gen shape " - f"{tuple(self.gen_shape)}" + f"np_model_fit={self._np_model_fit!r} not in {sorted(fz_tf.EFF_MODELS)}" ) - self.sigma_ns = tf.constant(sigma_ns_np, dtype=fz_tf.DTYPE) - # Reuse the native (NY, NqT) integral already computed above for - # sigma_YqT_central — no need to run the bT reconstruction at λ_central twice. - sigma_gen_central = self._sigma_gen_at( - self.eff_central, self.gnu_central, sigma_YqT=self.sigma_YqT_central - ) - # The pure gen-level integral (NptVGen, NabsYVGen), BEFORE folding through - # the response — used by the gen-level validation to test the integral - # in isolation (no R). - self.sigma_gen_central = sigma_gen_central - gen_flat = tf.reshape(sigma_gen_central, [-1]) - if tf.reduce_any(gen_flat <= 0).numpy(): - n_bad = int(tf.reduce_sum(tf.cast(gen_flat <= 0, tf.int32))) + if self._np_model_nu_fit not in fz_tf.GNU_MODELS: raise ValueError( - f"SCETlibNPParamModel: {n_bad} gen bins have non-positive " - f"σ_gen(λ_central); cannot normalize / fold the response." + f"np_model_nu_fit={self._np_model_nu_fit!r} not in " + f"{sorted(fz_tf.GNU_MODELS)}" ) + if (self._np_model_fit, self._np_model_nu_fit) != ( + self.core.np_model, + self.core.np_model_nu, + ): + print( + f"[SCETlibNPParamModel] NUMERATOR form overridden: " + f"F_eff {self.core.np_model}->{self._np_model_fit}, " + f"γ_ν {self.core.np_model_nu}->{self._np_model_nu_fit}; " + f"denominator (central) stays card form " + f"({self.core.np_model}/{self.core.np_model_nu}). " + f"VALIDATE against a matching SCETlib reference before trusting.", + flush=True, + ) + + # ---- Reco fold (Step 3) / gen-level baseline (Step 4 denominator). + # σ_gen(λ_central) comes from the core. R must encode only the gen→reco + # *mapping*, not the MC's absolute gen spectrum, so normalize each gen + # column by the gen-total N_gen(g): + # P(b|g) = R_raw(b,g) / N_gen(g) (eff × migration) + # σ_reco(λ_c;b) = Σ_g P(b|g) · σ_gen(λ_c;g) + # The reco-passing marginal Σ_b R_raw(b,g) is the WRONG normalizer (R is + # post-reco-selection, already carrying efficiency; dividing by it cancels + # efficiency and closes far worse). Because σ_reco_central then depends on + # σ_gen(λ_c) — it does NOT collapse to R_raw·1 — the λ_central closure is a + # genuine test of the integral. + gen_flat = tf.reshape(self.core.sigma_gen_central, [-1]) if self.gen_level: - # Gen-level σUL fit: the fit bins ARE the gen bins, so the per-bin - # ratio denominator is σ_gen(λ_central) directly (no reco fold). + # Gen-level σUL fit: fit bins ARE the gen bins, so the per-bin ratio + # denominator is σ_gen(λ_central) directly (no reco fold). + self.R = None + self._N_gen_flat = None self.sigma_gen_central_flat = gen_flat else: - N_gen = self._N_gen_flat + N_reco = int(np.prod(self.reco_shape)) + N_gen = int(np.prod(self.gen_shape)) + R_raw = tf.constant(R_arr.reshape(N_reco, N_gen), dtype=fz_tf.DTYPE) + self._N_gen_flat = tf.constant(N_gen_arr.reshape(-1), dtype=fz_tf.DTYPE) # Guard empty gen bins (no generated events): leave column at 0. - safe_N_gen = tf.where(N_gen > 0, N_gen, tf.ones_like(N_gen)) - self.R = self._R_raw / safe_N_gen[tf.newaxis, :] # P(b|g) = R_raw / N_gen - # Free the raw counts: only the normalized response self.R is used from - # here on (compute() never touches _R_raw) — no need to hold both. - del self._R_raw + safe_N_gen = tf.where( + self._N_gen_flat > 0, + self._N_gen_flat, + tf.ones_like(self._N_gen_flat), + ) + self.R = R_raw / safe_N_gen[tf.newaxis, :] # P(b|g) = R_raw / N_gen self.sigma_reco_central = tf.linalg.matvec( self.R, gen_flat ) # Σ_g P·σ_gen(λ_c) @@ -984,8 +683,8 @@ def __init__( ) self.signal_proc_idx = procs.index(signal_proc) self.nproc = indata.nproc - # Precompute the (1, N_proc) one-hot selecting the signal column, reused - # in every compute() to place the per-bin ratio (others stay at 1). + # Precompute the (1, N_proc) one-hot selecting the signal column, reused in + # every compute() to place the per-bin ratio (others stay at 1). self._signal_col_mask = tf.reshape( tf.one_hot(self.signal_proc_idx, self.nproc, dtype=indata.dtype), [1, self.nproc], @@ -1000,30 +699,29 @@ def __init__( self.params = np.array([p.encode() for p in self._param_order]) # Impact groups over our own parameters, consumed by the Fitter's - # traditional impacts. rabbit's built-in systgroup machinery can't - # represent these: its group indices are syst-relative, but our λ are - # POUs (model nuisances), which have no syst index. The Fitter resolves - # these labels -> floating full-x indices and computes the conditional - # group impact from the covariance. Split into the two NP sectors: - # CS-side γ_ν (lambda*_nu) vs TMD-effective F_eff. + # traditional impacts. rabbit's built-in systgroup machinery can't represent + # these: its group indices are syst-relative, but our λ are POUs (model + # nuisances) with no syst index. The Fitter resolves these labels -> floating + # full-x indices and computes the conditional group impact from the + # covariance. Split into the two NP sectors: CS-side γ_ν (lambda*_nu) vs + # TMD-effective F_eff. # - # ``resumNonpert`` is the SAME group name setupRabbit assigns to the - # discrete scetlibNP* template variations in the old-style datacard - # (there resumNonpert == exactly those 4 nuisances). Emitting it here - # (= all our lambda) makes the grouped-impact bar directly comparable - # between the new param model and the old NP variations. It does NOT - # collide with a syst group: the new-model datacard excludes scetlibNP, - # so resumNonpert is absent from indata.systgroups. + # ``resumNonpert`` is the SAME group name setupRabbit assigns to the discrete + # scetlibNP* template variations in the old-style datacard (there == exactly + # those 4 nuisances). Emitting it here (= all our lambda) makes the + # grouped-impact bar directly comparable between the new param model and the + # old NP variations. No collision with a syst group: the new-model datacard + # excludes scetlibNP, so resumNonpert is absent from indata.systgroups. self.param_impact_groups = { "resumNonpert": tuple(ALL_PARAMS), "scetlibNPgammaNu": tuple(GNU_PARAMS), "scetlibNPFeff": tuple(EFF_PARAMS), } - # Defaults: λ_central values per parameter. Optionally overridden by - # the ``xparam_default=name=value,...`` spec token — comma-separated - # pairs (for closure tests where the data-generating / fit-start - # point should differ from the card's λ_central). + # Defaults: λ_central values per parameter. Optionally overridden by the + # ``xparam_default=name=value,...`` spec token — comma-separated pairs (for + # closure tests where the data-generating / fit-start point should differ + # from the card's λ_central). central_lookup = {**self.eff_central, **self.gnu_central} defaults = np.array( [central_lookup[p] for p in self._param_order], dtype=np.float64 @@ -1044,27 +742,24 @@ def __init__( f"[SCETlibNPParamModel] xparamdefault overridden: {dict(zip(self._param_order, defaults))}", flush=True, ) - # rabbit's set_param_default expects an internal-storage convention - # where POIs (npoi entries) are SQRT(value) if not allowNegativeParam. - # For our λ which can in principle be tiny / zero (delta_lambda2), - # default to allowNegativeParam=True so the stored value == λ directly. + # rabbit's set_param_default stores POIs (npoi entries) as SQRT(value) if not + # allowNegativeParam. Our λ can be tiny/zero (delta_lambda2), so default to + # allowNegativeParam=True: stored value == λ directly. self.allowNegativeParam = True self.is_linear = False self.xparamdefault = tf.constant(defaults, dtype=indata.dtype) - # Gaussian priors (semantics documented on the constructor args). - # Rabbit's Fitter applies them whenever the model DECLARES - # ``prior_sigmas`` (no rabbit-side CLA — WMass/rabbit#133), so the - # declaration is gated behind ``priors``: off → no attribute → - # everything floats free. The Fitter takes the prior means from - # xparamdefault, so an xparam_default shift moves start AND prior - # mean together (to centre priors on truth while starting shifted, - # prior_means would have to be decoupled from xparamdefault). + # Gaussian priors (semantics on the constructor args). Rabbit's Fitter + # applies them whenever the model DECLARES ``prior_sigmas`` (no rabbit-side + # CLA — WMass/rabbit#133), so the declaration is gated behind ``priors``: off + # → no attribute → everything floats free. The Fitter takes prior means from + # xparamdefault, so an xparam_default shift moves start AND prior mean + # together (centring priors on truth while starting shifted would need + # prior_means decoupled from xparamdefault). self._use_priors = bool(priors) - # ``prior_sigmas`` may be a Mapping (programmatic) or the spec-token - # string ``prior_sigmas=lambda2=0.3,delta_lambda2=nan`` — the same - # comma-separated name=value format as xparam_default; value ``nan`` - # frees the param. + # ``prior_sigmas`` may be a Mapping (programmatic) or the spec-token string + # ``prior_sigmas=lambda2=0.3,delta_lambda2=nan`` — same comma-separated + # name=value format as xparam_default; value ``nan`` frees the param. if isinstance(prior_sigmas, str): prior_sigmas = dict( tuple(s.split("=")) for s in prior_sigmas.split(",") if s.strip() @@ -1086,7 +781,7 @@ def __init__( sigmas_arr[i] = np.nan # free (expected to be frozen) self.prior_sigmas = sigmas_arr # prior_means defaults to xparamdefault if not set, so don't store - # redundantly — Fitter will fall back to xparamdefault. + # redundantly — Fitter falls back to xparamdefault. print( "[SCETlibNPParamModel] Gaussian priors ENABLED (priors=1); " "applied by rabbit's Fitter (pre-#133 rabbit additionally " @@ -1104,12 +799,57 @@ def __init__( flush=True, ) + # ---- In-fit reco agreement guard (default ON; check_agreement=0 off). + # Compares the model's σ_reco(λ_central) SHAPE to the card's signal nominal + # template (what rnorm multiplies) and trips on any bin over + # check_agreement_threshold. Pure-numpy, runs once, never crashes the fit. + # Skipped in gen_level mode (no reco fold). Full reco+gen report and plots in + # param_model_diagnostics.run_card_diagnostics. + self.check_agreement = bool(check_agreement) + self.check_agreement_threshold = float(check_agreement_threshold) + self.check_agreement_strict = bool(check_agreement_strict) + self.check_agreement_min_yield = float(check_agreement_min_yield) + if self.check_agreement and not self.gen_level: + from wremnants.postprocessing.scetlib_np import ( + param_model_diagnostics as _diag, + ) + + _diag.run_reco_guard( + self, + indata, + threshold=self.check_agreement_threshold, + strict=self.check_agreement_strict, + min_yield_frac=self.check_agreement_min_yield, + ) + + # ---- Publish this fully-built model on the shared indata so the + # NPDampingWall regularizer can derive the FIT (numerator) forms and the + # canonical λ order from it, instead of being told them a second time on + # the -r line. indata is the only object both this model (built first in + # rabbit_fit.load_models) and the regularizer's mapping (built afterward + # with the same indata) share. The card form in indata.metadata is the + # WRONG one for the wall — the wall must constrain the numerator σ(λ), + # which np_model_(nu_)fit may override. See np_damping_wall.py. + indata.scetlib_np_param_model = self + + @property + def fit_forms(self): + """The NUMERATOR (fit) NP forms the wall must constrain — NOT the card + form (``self.np_model`` / ``self.np_model_nu``, forwarded from the core, + is the immutable denominator form). Returns ``{"np_model", "np_model_nu"}`` + with the F_eff / γ_ν forms actually integrated for σ(λ). The canonical λ + order is ``self._param_order`` (poi_params first, then the rest).""" + return { + "np_model": self._np_model_fit, + "np_model_nu": self._np_model_nu_fit, + } + def _check_discrete_np_double_counting(self): """Refuse to run on a datacard containing discrete scetlibNP systs. - They describe the same physics as this ParamModel's continuous λ; - running both double-counts (the discrete syst absorbs shape variation - the ParamModel should describe: spurious pull on the indata syst, + They describe the same physics as this ParamModel's continuous λ, so + running both double-counts: the discrete syst absorbs shape variation + the ParamModel should describe (spurious pull on the indata syst, postfit λ not what the data prefers). """ @@ -1153,77 +893,47 @@ def _fit_reco_axes(self, indata): ] # ========================================================================= - # σ_gen evaluation + # σ_gen evaluation — delegated to the physics core (Steps 1–2) # ========================================================================= - def _sigma_YqT_native_at(self, eff_params, gnu_params): - """Reconstruct σ(λ) on the btgrid and Q-integrate, returning the result - in the btgrid's *native* binning: shape (NY, NqT) on the signed-Y / - qT grid (Y_unique, qT_unique), BEFORE the |Y|-fold and qT-rebin. This - is the object that the native-binning validation compares against the - SCETlib spectrum reference (curve 1) and the external scetlib_run.factorize - (curve 2).""" - # 1. Reconstruct σ on the btgrid's flat (Nbins,) layout. Factorized - # (default) and legacy layouts are numerically equivalent (≲1e-14 - # rel.; summation order only). - eff = {k: v for k, v in eff_params.items() if k != "np_model"} - gnu = {k: v for k, v in gnu_params.items() if k != "np_model_nu"} - if self.factorized: - sigma_flat = fz_tf.reconstruct_batch_factorized_tf( - b_bar=self.b_bar, - I_pert_u=self.I_pert_u, - C_nu_uu=self.C_nu_uu, - c_of_u=self.c_of_u, - eff_params=eff, - gnu_params=gnu, - np_model=self.np_model, - np_model_nu=self.np_model_nu, - KwqT=self.KwqT, - gather_idx=self.gather_idx, - Y_unique=self.Y_feff_unique, - feff_idx_u=self.feff_idx_u, - ) - else: - sigma_flat = fz_tf.reconstruct_batch_tf( - qT_per_bin=self.qT_per_bin, - bT=self.bT, - I_pert=self.I_pert, - C_nu=self.C_nu, - b_bar=self.b_bar, - Y_per_bin=self.Y_per_bin, - eff_params=eff, - gnu_params=gnu, - np_model=self.np_model, - np_model_nu=self.np_model_nu, - bT_J0_kernel=self.bT_J0_kernel, - bT_simpson_weights=self.bT_simpson_w, - Y_unique=self.Y_feff_unique, - Y_inverse_idx=self.Y_feff_inverse_idx, - ) - # 2. Sparse → dense (NQ, NY, NqT). Missing cells get 0. - sigma_dense = fz_int.sparse_to_dense_tf(sigma_flat, self.flat_idx) - # 3. Integrate over Q (arctan_Q² Simpson) → (NY, NqT). - return fz_int.integrate_over_Q_tf(sigma_dense, self.Q_weights) + def __getattr__(self, name): + """Forward physics attributes (σ_gen tensors, btgrid axes, λ_central) to + ``self.core``. + + Python calls ``__getattr__`` only when normal attribute lookup fails, so + this never shadows a real ParamModel attribute; it makes the physics + surface the core owns (``eff_central``, ``gnu_central``, ``np_model``, + ``Y_unique``, ``qT_unique``, ``sigma_ns``, ``sigma_YqT_central``, + ``sigma_gen_central``, …) reachable as ``model.`` for backward + compatibility. Guarded against the ``core`` name itself and the pre-core + construction window to avoid recursion. + """ + if name != "core": + core = self.__dict__.get("core") + if core is not None and hasattr(core, name): + return getattr(core, name) + raise AttributeError( + f"{type(self).__name__!r} object has no attribute {name!r}" + ) - def _sigma_gen_at(self, eff_params, gnu_params, sigma_YqT=None): - """Evaluate σ_gen(λ) on R's gen binning. Returns shape (NptVGen, NabsYVGen). + def _sigma_YqT_native_at(self, eff_params, gnu_params, np_model=None, np_model_nu=None): + """Backward-compat alias for :meth:`SigmaGenModel.sigma_YqT_native` — the + native (NY, NqT) Q-integrated σ(λ) before the |Y|-fold / qT-rebin. + ``np_model`` / ``np_model_nu`` override the form (default: card form).""" + return self.core.sigma_YqT_native( + eff_params, gnu_params, np_model=np_model, np_model_nu=np_model_nu + ) - ``sigma_YqT`` lets a caller pass an already-computed native (NY, NqT) - integral to skip the (expensive) bT reconstruction — used at construction - to reuse ``self.sigma_YqT_central`` instead of integrating λ_central twice. - """ - if sigma_YqT is None: - sigma_YqT = self._sigma_YqT_native_at(eff_params, gnu_params) - # 4. Rebin Y (signed) → absYVGen (|Y|-folded): (NabsYVGen, NqT). - sigma_absY_qT = fz_int.rebin_axis_tf(sigma_YqT, axis=0, weights=self.W_absY) - # 5. Rebin qT → ptVGen: (NabsYVGen, NptVGen). - sigma_absY_ptV = fz_int.rebin_axis_tf( - sigma_absY_qT, axis=1, weights=self.W_ptVGen + def _sigma_gen_at( + self, eff_params, gnu_params, sigma_YqT=None, np_model=None, np_model_nu=None + ): + """Backward-compat alias for :meth:`SigmaGenModel.sigma_gen` — the matched + σ_gen(λ) on the (NptVGen, NabsYVGen) gen grid (Steps 1–2). + ``np_model`` / ``np_model_nu`` override the form (default: card form).""" + return self.core.sigma_gen( + eff_params, gnu_params, sigma_YqT=sigma_YqT, + np_model=np_model, np_model_nu=np_model_nu, ) - # 6. Reorder to (NptVGen, NabsYVGen) to match R's gen axis order. - sigma_resum = tf.transpose(sigma_absY_ptV, perm=[1, 0]) - # 7. Add the NP-independent fixed-order/DYTurbo nonsingular (zeros if off). - return sigma_resum + self.sigma_ns # ========================================================================= # λ-vector helpers @@ -1269,12 +979,20 @@ def _unpack_params(self, param): def _ratio_from_param(self, param): """λ (full param vector) → floored per-reco-bin ratio, shape (N_reco,). - The differentiable map that the straight-through Hessian path wraps. The - soft positivity floor (see ``compute``) lives here so both the normal and + The differentiable map the straight-through Hessian path wraps. The soft + positivity floor (see ``compute``) lives here so the normal and straight-through paths apply it identically. """ eff_params, gnu_params = self._unpack_params(param) - sigma_gen = self._sigma_gen_at(eff_params, gnu_params) + # Numerator uses the fit form (default = card form); the denominator + # (sigma_reco_central / sigma_gen_central_flat) was built from the core's + # card-form central, so it always stays the histmaker-consistent baseline. + sigma_gen = self._sigma_gen_at( + eff_params, + gnu_params, + np_model=self._np_model_fit, + np_model_nu=self._np_model_nu_fit, + ) gen_flat = tf.reshape(sigma_gen, [-1]) if self.gen_level: # Gen-level σUL fit: the fit bins ARE the gen bins — no reco fold. @@ -1292,9 +1010,9 @@ def _ratio_from_param(self, param): def _ratio_compact_jac(self, param): """J = d(ratio)/d(param), shape (N_reco, nparam), via forward-mode AD. - One JVP per parameter (nparam ≤ 8), each a single bT-fold pass — NOT - tiled over params. This is the compact object the Hessian actually needs - from the fold.""" + One JVP per parameter (nparam ≤ 8), each a single bT-fold pass, NOT + tiled over params. The compact object the Hessian needs from the + fold.""" n = int(param.shape[0]) cols = [] for i in range(n): @@ -1307,7 +1025,7 @@ def _ratio_compact_jac(self, param): def _ratio_compact_hess(self, param): """K = d²(ratio)/d(param)², shape (N_reco, nparam, nparam), forward-over-forward. - nparam² JVP-of-JVP passes (≤ 64), each one bT-fold pass — never tiled.""" + nparam² JVP-of-JVP passes (≤ 64), each one bT-fold pass, never tiled.""" n = int(param.shape[0]) rows = [] for i in range(n): @@ -1325,12 +1043,12 @@ def _ratio_compact_hess(self, param): def _ratio_straightthrough(self, param, use_curvature=True): """Exact ratio value, but autodiff sees only a compact quadratic in the - ≤8 λ (J, and optionally K) — so the (N_grid, N_bt) bT slab never enters + ≤8 λ (J, and optionally K), so the (N_grid, N_bt) bT slab never enters the differentiated graph and rabbit's covariance jacobian does not OOM. At the evaluation point (d = 0) the value is exact, the 1st derivative is J, and the 2nd derivative is K. ``use_curvature=False`` keeps only J - (Gauss-Newton / Fisher — exact for Asimov data, 8 vs 72 fold passes). + (Gauss-Newton / Fisher, exact for Asimov data, 8 vs 72 fold passes). """ val = self._ratio_from_param(param) J = tf.stop_gradient(self._ratio_compact_jac(param)) # (N_reco, nparam) @@ -1349,25 +1067,24 @@ def compute(self, param, full=False): ratio σ_reco(λ; b) / σ_reco(λ_central; b); other columns are 1. Positivity floor: the bT integral has no hard wall against pathological λ - (e.g. λ4 < 0 with the bounded-tanh models), which — via the b*-saturated, - hugely enhanced large-b region — can make σ_reco, and thus the predicted - signal yield, NEGATIVE. That would give a NaN Poisson NLL and a flat - gradient (tanh saturates), trapping the minimizer. We soft-floor the ratio - to a small positive value (RATIO_FLOOR_SCALE / RATIO_FLOOR_MIN) so a bad - point is a large-but-finite penalty with a usable gradient, NOT a crash. - The scale is far below any physical response, so the validated central and + (e.g. λ4 < 0 with the bounded-tanh models), which, via the b*-saturated, + hugely enhanced large-b region, can make σ_reco, and thus the predicted + signal yield, NEGATIVE. That gives a NaN Poisson NLL and a flat gradient + (tanh saturates), trapping the minimizer. We soft-floor the ratio to a + small positive value (RATIO_FLOOR_SCALE / RATIO_FLOOR_MIN) so a bad point + is a large-but-finite penalty with a usable gradient, NOT a crash. The + scale is far below any physical response, so the validated central and λ-variation closures are unchanged (ratio == 1 at λ_central, to fp). This - is a numerical safety net, not a physics constraint — it does not stop the - fit from exploring negative-λ; it just keeps that exploration finite. + is a numerical safety net, not a physics constraint: it does not stop the + fit from exploring negative-λ, it keeps that exploration finite. """ # ratio(λ) per reco bin. Normal path = exact fold (used for the fit). # Straight-through path (Hessian-only Phase B) keeps the bT slab off the - # autodiff graph so rabbit's covariance jacobian doesn't OOM. Toggled by - # hessian_straightthrough=1 spec token; - # hessian_gn=1 drops the curvature term + # autodiff graph so rabbit's covariance jacobian doesn't OOM. Toggled by the + # hessian_straightthrough=1 spec token; hessian_gn=1 drops the curvature term # (Gauss-Newton/Fisher — exact for Asimov, 8 vs 72 passes). Resolved at - # construction (self._hess_st / self._hess_gn). - # Do NOT enable during the fit: it recomputes J(/K) every call. + # construction (self._hess_st / self._hess_gn). Do NOT enable during the fit: + # it recomputes J(/K) every call. if self._hess_st: ratio = self._ratio_straightthrough(param, use_curvature=not self._hess_gn) else: @@ -1375,7 +1092,7 @@ def compute(self, param, full=False): # Build (N_reco, N_proc) scaling: 1 everywhere except the signal column, # which carries the per-bin ratio. Broadcasting the precomputed (1, N_proc) - # one-hot avoids materializing a separate ones tensor / rebuilding one_hot. + # one-hot avoids materializing a ones tensor / rebuilding one_hot. ratio_col = tf.cast( tf.reshape(ratio, [-1, 1]), self.indata.dtype ) # (N_reco, 1) diff --git a/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py b/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py new file mode 100644 index 000000000..acb8623c1 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py @@ -0,0 +1,525 @@ +"""Agreement diagnostics for :class:`SCETlibNPParamModel`: does the model +reproduce, at λ_central, the SHAPE present in the datacard? + +The model builds its λ_central prediction two ways: + + σ_gen(λ_c) the bt-grid Hankel + Q integral, matched (Steps 1-2) + σ_reco(λ_c) = R · σ_gen(λ_c) folded through the response (Step 3) + +and the fit applies ``rnorm = σ_reco(λ)/σ_reco(λ_c)`` on the datacard's signal +nominal template. For that transport to be valid the model's λ_central shape must +match the shape baked into the card (see :mod:`param_model`). This module checks +against the references that live IN THE CARD, so it needs no external histmaker: + + reco : σ_reco(λ_c) vs indata.norm[:, signal] (the template the + ratio multiplies) + gen : σ_gen(λ_c) vs N_gen (the gen-total in the scetlib_np + auxiliary, == the histmaker's + NP-corrected gen σ at λ_central) + +Both are SHAPE comparisons: σ_reco/σ_gen are folded theory cross sections, +norm/N_gen are weighted yields, so they differ by an overall normalization (it +cancels in the fit's ratio). Density plots unit-normalize both curves, so the +comparison carries no ad-hoc scale. + +Entry points: + * :func:`run_reco_guard` — PURE-NUMPY per-bin reco check for the in-fit + auto-guard (warns, or raises with ``strict``; never imports plotting). + * :func:`run_card_diagnostics` — full reco + gen comparison (+ optional plots) + for interactive use; returns the per-bin residuals. + * ``python -m wremnants.postprocessing.scetlib_np.param_model_diagnostics + --datacard [--outdir ]`` — construct from a card, write the + reco + gen agreement plots. + +Heavy deps (``hist``, ``wums.plot_tools``, and the plotting/projection helpers +from ``scetlib_np.validation_plots``) are imported LAZILY in the print/plot paths +only, so this module — and the in-fit guard — stays numpy-only. +""" + +import os + +import numpy as np + +# Per-reco/gen axis projection order for the standalone plots (shared names). +from wremnants.postprocessing.scetlib_np.params import ( + GEN_AXES as GEN_PROJ_AXES, + RECO_AXES as RECO_PROJ_AXES, +) + + +# ============================================================================= +# Postfit NP physical-validity detectors (standalone — NOT part of the fit). +# ============================================================================= +# A wrong-sign NP point anti-damps the form factors and makes the *differential* +# σ(qT) oscillate negative; the qT→ptVGen rebin AVERAGES that away, so it is +# invisible in the binned σ_gen / σ_reco / NLL the fit sees. FIT-TIME enforcement +# is ``np_damping_wall.NPDampingWall`` (the exact tanh_2/tanh_6 damping walls); these are +# cheap POSTFIT cross-checks on a constructed :class:`SigmaGenModel` ``core`` at a +# λ point (eff/gnu dicts, same shape as ``core.eff_central`` / ``core.gnu_central``). +# ``np_damping_ok`` probes the CAUSE (the forms must damp); ``spectrum_negativity`` +# measures the EFFECT (σ(qT) ≥ 0). +NP_PROBE_BT = (0.3, 1.0, 2.0, 5.0, 10.0, 20.0) + + +def np_damping_ok(core, eff_params, gnu_params, b_probe=NP_PROBE_BT, gamma_tol=1e-3): + """Probe the NP form factors (no bT integral) for the physical DAMPING sign. + + Evaluates the actual ``btgrid_tf`` forms the fit integrates at a few bT: + γ_ν^NP(b) ≤ 0 — CS Sudakov damping; γ_ν^NP > 0 is the anti-damping wrong + sign (λ2_ν < 0 / λ4_ν < 0) that makes σ(qT) oscillate neg. + F_eff(b) decays — TMD damping / bT-integral convergence; F_eff growing with + bT is the λ2_eff < 0 divergence trap. + Empirical proxy for the exact ``np_damping_wall.NPDampingWall`` walls: the γ_ν + probes test the SAME damping condition the CS walls enforce (form-agnostically + — the b⁶ term is in the evaluated form, so this also covers tanh_6); the F_eff + endpoint test at Y=0 is CRUDER than the wall's exact a≥0 ∀b at Y=0 and Y_max. + NOTE this probes the CARD form (``core.np_model``/``np_model_nu``), so it does + not see a numerator-form override (``np_model_(nu_)fit``). Cheap (1-D evals).""" + from wremnants.postprocessing.scetlib_np import btgrid_tf as fz_tf + + b = np.asarray(b_probe, dtype=np.float64) + eff = {k: v for k, v in eff_params.items() if k != "np_model"} + gnu = {k: v for k, v in gnu_params.items() if k != "np_model_nu"} + g = fz_tf.gamma_nu_NP_tf(b, np_model_nu=core.np_model_nu, **gnu).numpy() + F = fz_tf.F_eff_tf(0.0, b, np_model=core.np_model, **eff).numpy() + gamma_max = float(np.max(g)) + feff_growing = bool(F[-1] > F[0]) + return { + "probe_b": b, + "gamma_nu": g, + "F_eff": F, + "gamma_nu_max": gamma_max, + "gamma_nu_wrong_sign": bool(gamma_max > gamma_tol), + "F_eff_growing": feff_growing, + "ok": (gamma_max <= gamma_tol) and (not feff_growing), + } + + +def spectrum_negativity(core, eff_params, gnu_params, sigma_YqT=None, locate=True): + """Negativity of the native (Y, qT) resummed spectrum at a λ point — the + σ(qT) < 0 pathology the gen-binning averages away. Scale-free metrics: + neg_area_frac = Σ|min(σ,0)| / Σ|σ| (0 physical; → O(1) pathological) + min_over_peak = min(σ) / max(σ) (≈0 physical; ≤ −O(1) pathological) + Pass ``sigma_YqT`` (e.g. ``core.sigma_YqT_central``) to skip recomputation, + else reconstructed via ``core.sigma_YqT_native``. Judge relative to the + λ_central baseline (``np_physical_report`` does this): the singular-only + spectrum carries a tiny benign qT→0 dip. + + With ``locate`` (default True) and a 2-D spectrum on the core's native grids + (``core.Y_unique`` × ``core.qT_unique``), also returns WHERE the negativity + sits — the key discriminator for whether it touches the fit region or is + laundered/out-of-acceptance: + ``worst`` {iY, iqT, Y, qT, value, frac_of_peak} of the most-negative cell + ``neg_bins`` per-cell {Y, qT, value, frac_of_peak}, most-negative first + ``neg_qT_range`` / ``neg_absY_max`` extent of the negative region + (omitted if the grids are absent or the shape doesn't match).""" + s = ( + core.sigma_YqT_native(eff_params, gnu_params) + if sigma_YqT is None + else sigma_YqT + ) + s = np.asarray(s) + peak = float(np.max(s)) + neg = float(np.sum(np.abs(np.minimum(s, 0.0)))) + tot = float(np.sum(np.abs(s))) + out = { + "neg_area_frac": neg / tot, + "min_over_peak": float(np.min(s)) / peak, + "n_neg_bins": int(np.sum(s < 0)), + } + Yg = np.asarray(getattr(core, "Y_unique", None)) if locate else None + qTg = np.asarray(getattr(core, "qT_unique", None)) if locate else None + if ( + locate + and s.ndim == 2 + and Yg is not None and qTg is not None + and Yg.ndim == 1 and qTg.ndim == 1 + and s.shape == (Yg.size, qTg.size) + ): + jmin = np.unravel_index(int(np.argmin(s)), s.shape) + out["worst"] = dict( + iY=int(jmin[0]), iqT=int(jmin[1]), + Y=float(Yg[jmin[0]]), qT=float(qTg[jmin[1]]), + value=float(s[jmin]), frac_of_peak=float(s[jmin] / peak), + ) + negidx = np.argwhere(s < 0) + if negidx.size: + order = np.argsort(s[negidx[:, 0], negidx[:, 1]]) # most negative first + out["neg_bins"] = [ + dict( + Y=float(Yg[negidx[k, 0]]), qT=float(qTg[negidx[k, 1]]), + value=float(s[negidx[k, 0], negidx[k, 1]]), + frac_of_peak=float(s[negidx[k, 0], negidx[k, 1]] / peak), + ) + for k in order + ] + out["neg_qT_range"] = ( + float(qTg[negidx[:, 1]].min()), float(qTg[negidx[:, 1]].max()) + ) + out["neg_absY_max"] = float(np.abs(Yg[negidx[:, 0]]).max()) + else: + out["neg_bins"] = [] + out["neg_qT_range"] = None + out["neg_absY_max"] = None + return out + + +def np_physical_report( + core, eff_params, gnu_params, sigma_YqT=None, central_neg_area=None +): + """Combine both detectors into a postfit verdict (no printing, no raising). + + Returns ``{ok, issues, damp, neg, central_neg_area}``: ``ok`` the overall + verdict, ``issues`` human-readable problems, ``damp``/``neg`` the raw + sub-results. ``central_neg_area`` anchors the relative negativity threshold + (default: from ``core`` at λ_central). Callers format/act.""" + damp = np_damping_ok(core, eff_params, gnu_params) + neg = spectrum_negativity(core, eff_params, gnu_params, sigma_YqT=sigma_YqT) + if central_neg_area is None: + central_neg_area = spectrum_negativity( + core, + core.eff_central, + core.gnu_central, + sigma_YqT=getattr(core, "sigma_YqT_central", None), + locate=False, + )["neg_area_frac"] + neg_bad = neg["neg_area_frac"] > max(0.01, 5.0 * central_neg_area) + issues = [] + if damp["gamma_nu_wrong_sign"]: + issues.append( + f"γ_ν^NP > 0 (anti-damping, wrong CS sign; max={damp['gamma_nu_max']:+.3g}) " + f"on probe bT={list(damp['probe_b'])}" + ) + if damp["F_eff_growing"]: + issues.append( + "F_eff grows with bT (TMD divergence sign — λ2_eff likely < 0; the bT " + "integral is then finite only by grid truncation)" + ) + if neg_bad: + issues.append( + f"native σ(qT) significantly negative: neg_area_frac=" + f"{neg['neg_area_frac']:.3g} (λ_central {central_neg_area:.3g}), " + f"min/peak={neg['min_over_peak']:+.3g}, n_neg_bins={neg['n_neg_bins']}" + ) + return { + "ok": damp["ok"] and not neg_bad, + "issues": issues, + "damp": damp, + "neg": neg, + "central_neg_area": central_neg_area, + } + + +# ============================================================================= +# Pure-numpy core (no heavy imports) — shared by the guard and the full report. +# ============================================================================= +def _shape_residual(model_vals, ref_vals): + """Per-bin shape residual after matching the overall integral. + + Returns ``(resid, stats)``: ``resid`` has the input shape and is + ``scale·model/ref - 1`` (NaN where ref<=0), ``scale = Σref/Σmodel``. ``stats`` + carries what a threshold guard keys on: yield-weighted mean |residual| and the + worst bin (value + flat index).""" + m = np.asarray(model_vals, dtype=np.float64) + n = np.asarray(ref_vals, dtype=np.float64) + if m.shape != n.shape: + raise ValueError(f"shape mismatch: model {m.shape} vs ref {n.shape}") + msum, nsum = m.sum(), n.sum() + scale = nsum / msum if msum != 0 else np.nan + good = n > 0 + resid = np.full(m.shape, np.nan) + resid[good] = scale * m[good] / n[good] - 1.0 + rg = np.abs(resid[good]) + wmad = float(np.average(rg, weights=n[good])) if good.any() else np.nan + imax = int(np.nanargmax(np.abs(resid))) if good.any() else -1 + stats = dict( + scale=float(scale), + n_bins=int(good.sum()), + max_abs=float(rg.max()) if rg.size else np.nan, + yield_weighted_mean_abs=wmad, + worst_flat_idx=imax, + worst_value=float(resid.flat[imax]) if imax >= 0 else np.nan, + ) + return resid, stats + + +def card_reco_reference(model, indata): + """Signal nominal reco template from the card, on the model's reco binning. + + ``indata.norm`` is ``(nbins, nproc)``; the signal column is the histmaker + nominal yield ``rnorm`` multiplies. Reshaped to ``model.reco_shape`` (same + axis order as ``sigma_reco_central``).""" + norm = np.asarray(indata.norm, dtype=np.float64) + n_reco = int(np.prod(model.reco_shape)) + if norm.shape[0] < n_reco: + raise ValueError( + f"indata.norm has {norm.shape[0]} bins < reco bins {n_reco}" + ) + # v1 single non-masked channel: the reco bins are the leading rows of norm. + col = norm[:n_reco, model.signal_proc_idx] + return col.reshape(model.reco_shape) + + +def card_gen_reference(model): + """Gen-total ``N_gen`` from the scetlib_np auxiliary, on the model gen grid. + + ``N_gen`` is the prefsr xnorm UL gen-total in the card — the histmaker's + NP-corrected gen σ at λ_central. It normalizes R, and ``P = R_raw/N_gen`` is + built theory-independent, so N_gen carries the same nominal correction as + R_raw. Shape == ``model.gen_shape``.""" + return np.asarray(model._N_gen_flat, dtype=np.float64).reshape(model.gen_shape) + + +def reco_offending_bins(model, indata, threshold, min_yield_frac=0.0): + """Reco bins whose |shape residual| exceeds ``threshold`` (fractional). + + Pure numpy (no plotting/hist imports), so the in-fit guard can call it. + ``min_yield_frac`` (of the max reference bin) optionally drops near-empty bins, + where a shape residual is meaningless. Returns ``(offenders, stats)``, + offenders sorted by |residual| descending.""" + ref = card_reco_reference(model, indata) + mdl = np.asarray(model.sigma_reco_central, np.float64).reshape(model.reco_shape) + resid, stats = _shape_residual(mdl, ref) + absr = np.abs(resid) + if min_yield_frac > 0: + floor = float(min_yield_frac) * float(np.nanmax(ref)) + absr = np.where(ref >= floor, absr, np.nan) + names = [n for n, _ in model._reco_axes_meta] + offenders = [] + for idx in np.argwhere(absr > threshold): + idx = tuple(int(i) for i in idx) + offenders.append( + dict( + coord=dict(zip(names, idx)), + residual=float(resid[idx]), + ref_yield=float(ref[idx]), + ) + ) + offenders.sort(key=lambda d: -abs(d["residual"])) + return offenders, stats + + +def run_reco_guard( + model, indata, threshold=0.005, strict=False, min_yield_frac=0.0, max_list=12 +): + """In-fit reco agreement guard (pure numpy). Trips on ANY bin exceeding + ``threshold`` (|σ_reco(λ_c)/card_nominal − 1|). + + Warns and lists the worst offenders by default; raises iff ``strict``. A + diagnostic failure never takes down the fit — any error is caught and logged. + """ + tag = "[SCETlibNPParamModel]" + try: + offenders, stats = reco_offending_bins(model, indata, threshold, min_yield_frac) + except Exception as e: # a diagnostic must never crash the minimizer + print(f"{tag} reco agreement check SKIPPED (error: {e})", flush=True) + return None + wmean = stats["yield_weighted_mean_abs"] * 100 + if not offenders: + print( + f"{tag} reco agreement OK: all {stats['n_bins']} bins within " + f"{threshold*100:.2f}% (yield-weighted mean |shape−1| = {wmean:.3f}%).", + flush=True, + ) + return offenders + head = "\n".join( + f" {d['residual']*100:+.2f}% (" + + ", ".join(f"{k}={v}" for k, v in d["coord"].items()) + + f") ref_yield={d['ref_yield']:.3g}" + for d in offenders[:max_list] + ) + more = "" if len(offenders) <= max_list else f"\n … +{len(offenders)-max_list} more" + msg = ( + f"{tag} reco agreement: {len(offenders)} bin(s) exceed {threshold*100:.2f}% " + f"|σ_reco(λ_c)/card_nominal − 1| (yield-weighted mean = {wmean:.3f}%). " + f"The model's λ_central shape does not match the card's signal template " + f"in these bins — the rnorm it applies there is transported off the wrong " + f"baseline. (Trips concentrated in the top ptll bin [37,44] are the KNOWN " + f"qT>100 grid truncation — see the param_model module docstring 'Known " + f"limitation' — not misuse; they do not bias the fit.) Worst:\n{head}{more}" + ) + if strict: + raise ValueError( + msg + + "\n(check_agreement_strict=1 → raising. Pass check_agreement=0 to " + "disable, raise check_agreement_threshold, or set check_agreement_min_yield " + "to ignore sparse bins.)" + ) + print("WARNING: " + msg, flush=True) + return offenders + + +# ============================================================================= +# Full report (+ optional plots) — heavy deps imported lazily here. +# ============================================================================= +def compare_level(model_vals, ref_vals, axes_meta, label): + """Print detailed per-bin diagnostics + return ``(resid, stats)``. + + ``model_vals`` / ``ref_vals`` are ndarrays on the same binning; ``axes_meta`` + the ``(name, edges)`` list for that level (reco or gen).""" + from wremnants.postprocessing.scetlib_np.validation_plots import summarize + + print(f"\n{'='*70}\n{label}\n{'='*70}") + summarize(np.asarray(model_vals, np.float64), np.asarray(ref_vals, np.float64), axes_meta) + resid, stats = _shape_residual(model_vals, ref_vals) + names = [n for n, _ in axes_meta] + shape = tuple(len(e) - 1 for _, e in axes_meta) + coord = ( + np.unravel_index(stats["worst_flat_idx"], shape) + if stats["worst_flat_idx"] >= 0 + else None + ) + coord_str = ( + ", ".join(f"{nm}={c}" for nm, c in zip(names, coord)) if coord is not None else "n/a" + ) + print( + f"\n >> {label}: yield-weighted mean|shape−1| = " + f"{stats['yield_weighted_mean_abs']*100:.3f}% | " + f"worst bin = {stats['worst_value']*100:+.2f}% at ({coord_str})" + ) + return resid, stats + + +def run_card_diagnostics( + model, indata, outdir=None, do_plots=True, ref_label_reco=None, + gen_exclude_overflow=True, +): + """Full reco + gen comparison of the model's λ_central shape to the card's + references, with optional per-axis shape plots. + + The gen `ptVGen` axis has a known-TRUNCATED overflow bin: the model integrates + qT only to the bt-grid ceiling (PTVGEN_OVERFLOW_EDGE), while N_gen's overflow + holds all qT beyond the last edge (unbounded). A global Σ/Σ shape + normalization forces the totals to match, smearing that one large-yield + deficit into a flat pedestal across every other bin — an artifact masquerading + as a constant offset (verified 260625). With ``gen_exclude_overflow`` (default + True) the gen shape is normalized on the RESOLVED bins only, so the bulk + closure is faithful and the truncation shows as its own step in the `ptVGen` + ratio; the `absYVGen` plot then projects resolved-qT only (overflow zeroed) + for a clean rapidity-shape comparison. + + Returns ``{'reco': (resid, stats), 'gen': (resid, stats)}`` (stats are the + global-norm ones from compare_level; the resolved-norm bulk figure is + printed). Reco is unaffected: its overflow ptll bin was cropped to the fit + binning, so there is no truncated bin to smear.""" + out = {} + reco_ref = card_reco_reference(model, indata) + reco_model = np.asarray(model.sigma_reco_central, np.float64).reshape(model.reco_shape) + out["reco"] = compare_level( + reco_model, reco_ref, model._reco_axes_meta, "RECO σ_reco(λ_c) vs card norm[signal]" + ) + + gen_ref = card_gen_reference(model) + gen_model = np.asarray(model.sigma_gen_central, np.float64).reshape(model.gen_shape) + out["gen"] = compare_level( + gen_model, gen_ref, model._gen_axes_meta, "GEN σ_gen(λ_c) vs card N_gen" + ) + + # Resolved-qT (overflow-excluded) gen normalization — the faithful bulk view. + gen_names = [n for n, _ in model._gen_axes_meta] + ptv_ax = gen_names.index("ptVGen") if "ptVGen" in gen_names else 0 + drop = 1 if gen_exclude_overflow else 0 + resolved = [slice(None)] * gen_model.ndim + resolved[ptv_ax] = slice(0, gen_model.shape[ptv_ax] - drop) + resolved = tuple(resolved) + gscale = gen_ref[resolved].sum() / gen_model[resolved].sum() + if gen_exclude_overflow: + rres, wres = gscale * gen_model[resolved] / gen_ref[resolved], gen_ref[resolved] + print( + f"\n >> GEN resolved-qT norm (overflow bin excluded): bulk " + f"yield-weighted mean|shape−1| = " + f"{np.average(np.abs(rres - 1.0), weights=wres) * 100:.3f}% " + f"(the global-norm {out['gen'][1]['yield_weighted_mean_abs']*100:.2f}% is that " + f"truncation smeared into a pedestal — see docstring)" + ) + + if do_plots and outdir: + from wremnants.postprocessing.scetlib_np.validation_plots import ( + plot_ptll_ratio, + tf_to_hist, + ) + + os.makedirs(outdir, exist_ok=True) + rlabel_reco = ref_label_reco or "card nominal (signal)" + h_reco_m = tf_to_hist(reco_model, model._reco_axes_meta) + h_reco_n = tf_to_hist(reco_ref, model._reco_axes_meta) + for ax in RECO_PROJ_AXES: + plot_ptll_ratio( + h_reco_m, h_reco_n, axis=ax, + out_path=os.path.join(outdir, f"reco_{ax}.png"), + ref_label=rlabel_reco, model_label=r"ParamModel $\sigma_{reco}(\lambda_c)$", + rlabel="model / card", density=True, + title=f"σ_reco(λ_c) vs card nominal — {ax}", + ) + # gen ptVGen: normalize on resolved bins but keep all bins, so the + # truncated overflow shows as a step (not a pedestal on the bulk). + # Pre-scale the model and pass scale=1.0 so the legend has no "(×scale)"; + # the normalization choice lives in the title. + norm_tag = "resolved-qT norm" if gen_exclude_overflow else "global norm" + plot_ptll_ratio( + tf_to_hist(gen_model * gscale, model._gen_axes_meta), + tf_to_hist(gen_ref, model._gen_axes_meta), + axis="ptVGen", density=False, + out_path=os.path.join(outdir, "gen_ptVGen.png"), + ref_label=r"card $N_{gen}$", + model_label=r"ParamModel $\sigma_{gen}(\lambda_c)$", + rlabel="model / $N_{gen}$", rrange=(0.78, 1.05), + title=f"σ_gen(λ_c) vs card N_gen — ptVGen ({norm_tag})", + ) + # gen absYVGen: project resolved-qT only (overflow zeroed in both) so the + # rapidity shape isn't contaminated by the truncated bin. + gm_res, gn_res = gen_model.copy(), gen_ref.copy() + if gen_exclude_overflow: + ov = [slice(None)] * gen_model.ndim + ov[ptv_ax] = gen_model.shape[ptv_ax] - 1 + gm_res[tuple(ov)] = 0.0 + gn_res[tuple(ov)] = 0.0 + plot_ptll_ratio( + tf_to_hist(gm_res, model._gen_axes_meta), + tf_to_hist(gn_res, model._gen_axes_meta), + axis="absYVGen", density=True, + out_path=os.path.join(outdir, "gen_absYVGen.png"), + ref_label=r"card $N_{gen}$", + model_label=r"ParamModel $\sigma_{gen}(\lambda_c)$", + rlabel="model / $N_{gen}$", rrange=(0.95, 1.05), + title="σ_gen(λ_c) vs card N_gen — |y|" + (", resolved qT" if gen_exclude_overflow else ""), + ) + print(f"\n plots written under: {outdir}") + return out + + +def main(argv=None): + import argparse + import time + + from rabbit.inputdata import FitInputData + + from wremnants.postprocessing.scetlib_np.param_model import SCETlibNPParamModel + + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--datacard", required=True, help="fit-input hdf5 (must carry the scetlib_np auxiliary)") + p.add_argument("--btgrid", default=None, help="SCETlib bt-grid dir (default: model's data-area copy)") + p.add_argument("--signal-proc", default="Zmumu") + p.add_argument("--outdir", default=None, help="plot output dir ('' / unset to skip plotting)") + args = p.parse_args(argv) + + print("Loading FitInputData …", flush=True) + t0 = time.time() + indata = FitInputData(args.datacard) + print(f" loaded in {time.time()-t0:.1f}s; nproc={indata.nproc}", flush=True) + + print("Constructing SCETlibNPParamModel (runs the bt integral at λ_central) …", flush=True) + t0 = time.time() + kw = dict(signal_proc=args.signal_proc, check_agreement=False) # report below, not the guard + if args.btgrid: + kw["btgrid_dir"] = args.btgrid + model = SCETlibNPParamModel(indata, **kw) + print(f" constructed in {time.time()-t0:.1f}s", flush=True) + + run_card_diagnostics( + model, indata, outdir=(args.outdir or None), do_plots=bool(args.outdir) + ) + + +if __name__ == "__main__": + main() diff --git a/wremnants/postprocessing/scetlib_np/params.py b/wremnants/postprocessing/scetlib_np/params.py new file mode 100644 index 000000000..438eb1a7e --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/params.py @@ -0,0 +1,122 @@ +"""Shared TF-free vocabulary and helpers for the SCETlib-NP param model. + +Single source of truth for the λ parameter names, the np_model selector keys, +the reco/gen axis names, and the few numpy helpers used by both the TF core +(:mod:`sigma_gen`, :mod:`param_model`) and the TF-free tools +(:mod:`sigma_gen_at_lambda`, :mod:`np_function_plots`, :mod:`fitresult_lambdas`, +:mod:`lambda_central`). Import-light (numpy only) so the lightweight tools reach +these without pulling in the TF core. +""" + +import numpy as np + +# λ parameter names: CS-side γ_ν^NP first, then TMD-effective F_eff. +# lambda6_nu is the CS b⁶ coefficient — only the tanh_6 γ_ν model uses it; tanh_2 +# ignores it, so it is inert there and defaults to 0 from the card runcard. +GNU_PARAMS = ("lambda2_nu", "lambda4_nu", "lambda6_nu", "lambda_inf_nu") +EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") +ALL_PARAMS = GNU_PARAMS + EFF_PARAMS + +# Recommended Gaussian prior widths per λ (theorist recommendations), applied by +# SCETlibNPParamModel only when priors are enabled (priors=1); a λ absent here +# floats free (NaN width). Lives in this config home keyed by the names above. +DEFAULT_PRIOR_SIGMAS = { + "lambda2_nu": 0.10, + "lambda4_nu": 0.50, + "lambda6_nu": 0.10, + "lambda2": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average + "lambda4": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average + "delta_lambda2": 0.20, # 0 ± 0.20 wide default (no theorist value yet) + "lambda6": 0.1, +} + +# np_model selector keys carried alongside the numeric λ in a tune dict. +EFF_MODEL_KEY = "np_model" +GNU_MODEL_KEY = "np_model_nu" + +# Fit observable axis names. +RECO_AXES = ("ptll", "yll", "cosThetaStarll_quantile", "phiStarll_quantile") +GEN_AXES = ("ptVGen", "absYVGen") + +# Which λ each NP model actually uses. The form factors in btgrid_tf +# (F_eff_tf / gamma_nu_NP_tf) read a fixed subset of the λ per model string; a λ +# outside that subset is inert (e.g. lambda6 under tanh_2). Mirrored here as +# plain data so this module stays TF-free — btgrid_tf is the source of truth. +_EFF_MODEL_ALIASES = {"hyp_tangent": "tanh_2", "square_root": "frac_2"} +_GNU_MODEL_ALIASES = {"hyp_tangent": "tanh_2", "linear": "frac_1"} +# eff models with no lambda_inf damping (plain polynomial form factors). +_EFF_NO_LAMBDA_INF = {"signed_lambda", "identity"} + + +def active_params(np_model=None, np_model_nu=None): + """Names of the λ the chosen NP model(s) actually use. + + Pass ``np_model`` for the F_eff (TMD) set, ``np_model_nu`` for the γ_ν (CS) + set, or both for the union. A λ outside the returned set is inert for that + model (``lambda6``/``lambda6_nu`` are used only by ``tanh_6``); callers can + reject such λ instead of silently ignoring them. Source of truth: + ``btgrid_tf.F_eff_tf`` / ``btgrid_tf.gamma_nu_NP_tf``.""" + out = set() + if np_model is not None: + m = _EFF_MODEL_ALIASES.get(np_model, np_model) + out |= {"lambda2", "lambda4", "delta_lambda2"} + if m not in _EFF_NO_LAMBDA_INF: + out.add("lambda_inf") + if m == "tanh_6": + out.add("lambda6") + if np_model_nu is not None: + m = _GNU_MODEL_ALIASES.get(np_model_nu, np_model_nu) + out |= {"lambda2_nu", "lambda4_nu", "lambda_inf_nu"} + if m == "tanh_6": + out.add("lambda6_nu") + return out + + +def parse_lambda_overrides(spec): + """Parse a ``"name=val,name=val"`` λ-override string into ``{name: float}``. + + Hard-errors (``ValueError``) on a malformed token, an unknown parameter name + (not in :data:`ALL_PARAMS`), or a non-float value. An empty / ``None`` spec + yields ``{}``. Model-awareness (whether a *known* λ is used by the chosen + model) is the caller's job via :func:`active_params`.""" + out = {} + for tok in (spec or "").split(","): + tok = tok.strip() + if not tok: + continue + if "=" not in tok: + raise ValueError(f"--lambdas: expected 'name=value', got {tok!r}") + k, v = tok.split("=", 1) + k = k.strip() + if k not in ALL_PARAMS: + raise ValueError( + f"--lambdas: unknown NP parameter {k!r} " + f"(known: {', '.join(ALL_PARAMS)})" + ) + try: + out[k] = float(v) + except ValueError: + raise ValueError(f"--lambdas: {k}={v.strip()!r} is not a float") + return out + + +def split_eff_gnu(values): + """Split a ``{name: value}`` mapping into ``(eff_params, gnu_params)`` dicts + by membership in EFF_PARAMS / GNU_PARAMS (values floated; names in neither, + e.g. the model-name keys, are dropped).""" + eff = {k: float(values[k]) for k in EFF_PARAMS if k in values} + gnu = {k: float(values[k]) for k in GNU_PARAMS if k in values} + return eff, gnu + + +def bin_sum_matrix(src_centers, target_edges, tol=1e-6): + """(N_target, N_src) 0/1 matrix summing source bins whose centre falls in + each target bin. Source bins outside every target bin get 0, truncating to + the target range (e.g. qT > ptVGen_max, |Y| > absY_max).""" + src = np.asarray(src_centers, dtype=np.float64) + edges = np.asarray(target_edges, dtype=np.float64) + W = np.zeros((edges.size - 1, src.size), dtype=np.float64) + for i in range(edges.size - 1): + m = (src >= edges[i] - tol) & (src <= edges[i + 1] + tol) + W[i, m] = 1.0 + return W diff --git a/wremnants/postprocessing/scetlib_np/response_matrix.py b/wremnants/postprocessing/scetlib_np/response_matrix.py index d2983abc3..31b51d96c 100644 --- a/wremnants/postprocessing/scetlib_np/response_matrix.py +++ b/wremnants/postprocessing/scetlib_np/response_matrix.py @@ -1,16 +1,14 @@ """Load the (reco × gen) response matrix R for the ParamModel. R lives in the unfolding histmaker output (a separate hdf5 from the fit tensor) -as ``nominal_prefsr_yieldsUnfolding`` under the Z sample group. We slice -``acceptance=True`` (gen-fiducial) and project to the reco × (ptVGen, absYVGen) -axes — which SUMS the helicitySig axis: R is filled with the weight PARTITION -``nominal_weight_helicity`` whose 8 pieces add back up, so the physical yield -is the helicitySig SUM (NOT UL). The gen-total normalizer N_gen, by contrast, -is filled with ``csAngularMoments`` and takes the UL component -(``helicitySig=-1``); see ``_select_ul_helicity`` and the inline comment in -``load_R``. The full response-fold formula and this SUM-vs-UL subtlety are -documented once in the :mod:`param_model` module docstring (single source of -truth) — consult it rather than re-deriving here. +as ``nominal_prefsr_yieldsUnfolding`` under the Z sample group. Slice +``acceptance=True`` (gen-fiducial), project to reco × (ptVGen, absYVGen): this +SUMS the helicitySig axis. R is filled with the weight PARTITION +``nominal_weight_helicity`` whose 8 pieces add back up, so the physical yield is +the helicitySig SUM (NOT UL). The gen-total normalizer N_gen is filled with +``csAngularMoments`` and takes the UL component (``helicitySig=-1``); see +``_select_ul_helicity`` and the inline comment in ``load_R``. Full response-fold +formula and this SUM-vs-UL subtlety: :mod:`param_model` module docstring. Single entry point: @@ -24,47 +22,41 @@ from wums import ioutils as wums_io -# Pre-FSR gen level: the SCETlib btgrid σ_gen is the resummed *boson* qT/Y -# (QCD, before QED FSR), so the response and gen-total must also be pre-FSR for -# σ_gen, R, and N_gen to live at the same gen level. (postfsr variants exist in -# the same file — nominal_postfsr_yieldsUnfolding / "postfsr" — for comparison.) +# Axes kept, in canonical order: reco first, then gen. +from wremnants.postprocessing.scetlib_np.params import GEN_AXES, RECO_AXES + +# Pre-FSR: the btgrid σ_gen is resummed *boson* qT/Y (QCD, pre-QED-FSR), so R +# and N_gen must also be pre-FSR for σ_gen, R, N_gen to share a gen level. +# (postfsr variants — nominal_postfsr_yieldsUnfolding / "postfsr" — also exist.) DEFAULT_HIST = "nominal_prefsr_yieldsUnfolding" DEFAULT_GENTOTAL = "prefsr" # xnorm gen-total denominator (pre-reco-selection) DEFAULT_SAMPLE = "Zmumu_2016PostVFP" - -# Axes we keep, in canonical order: reco first, then gen. -RECO_AXES = ("ptll", "yll", "cosThetaStarll_quantile", "phiStarll_quantile") -GEN_AXES = ("ptVGen", "absYVGen") -# helicitySig is an angular-moment axis; we take the UL component (value -1), -# the angular-integrated total — see _select_ul_helicity. acceptance is sliced -# to True (gen-fiducial) at use. +# helicitySig: angular-moment axis; take UL (value -1), the angular-integrated +# total (see _select_ul_helicity). acceptance sliced True (gen-fiducial) at use. HELICITY_AXIS = "helicitySig" -# Gen ptVGen overflow handling. The unfolding hist's ptVGen axis ends at 44 (the -# last reco-ptll edge), but ~3.6% of the Z yield has true gen qT > 44 and -# resolution-migrates into the high-ptll reco bins (6.2% of the last reco bin). -# Dropping that gen-overflow column makes σ_reco low there. We instead fold the -# ptVGen overflow into an extra gen bin so the model can supply a σ_gen for it -# (the btgrid integral over qT ∈ (44, PTVGEN_OVERFLOW_EDGE]). The edge must be -# ≤ the btgrid qT max (the fineall grid runs to 100) and should coincide with a -# gen-histmaker ptVgen edge so the gen-level cross-check's _merge_matrix is exact -# — 100 satisfies both. (absYVGen has zero overflow: |Y| ≤ 2.5 is fully contained, -# so only ptVGen needs this.) +# Gen ptVGen overflow. The ptVGen axis ends at 44 (last reco-ptll edge), but +# ~3.6% of the Z yield has true gen qT > 44 and resolution-migrates into the +# high-ptll reco bins (6.2% of the last reco bin); dropping that column makes +# σ_reco low there. Instead fold the overflow into an extra gen bin so the model +# can supply a σ_gen for it (btgrid integral over qT ∈ (44, PTVGEN_OVERFLOW_EDGE]). +# The edge must be ≤ the btgrid qT max (fineall runs to 100) and should coincide +# with a gen-histmaker ptVgen edge so the cross-check's _merge_matrix is exact; +# 100 satisfies both. (absYVGen has zero overflow: |Y| ≤ 2.5 is fully contained.) PTVGEN_OVERFLOW_EDGE = 100.0 def _select_ul_helicity(h): """Select the UL angular component (helicitySig = -1), if that axis exists. - Applied to the gen-total denominator N_gen ONLY — NOT to R. N_gen is filled - with ``csAngularMoments``, a moment expansion whose A_i bins (0..7) are - signed and do NOT sum to σ; only UL (value -1) is the angular-integrated - total, so N_gen takes UL. R is the opposite: filled with the weight - PARTITION ``nominal_weight_helicity``, it is recovered by SUMMING - helicitySig (``project``), not by taking UL — see the inline comment in - ``load_R``. Same axis, opposite reduction; taking UL of R would inflate the - closure ~15×. Returning h unchanged when the axis is absent keeps - non-helicity inputs working. + For the gen-total denominator N_gen ONLY, NOT R. N_gen is filled with + ``csAngularMoments``, a moment expansion whose A_i bins (0..7) are signed and + do NOT sum to σ; only UL (value -1) is the angular-integrated total, so N_gen + takes UL. R is filled with the weight PARTITION ``nominal_weight_helicity`` + and is recovered by SUMMING helicitySig (``project``), not by UL (see the + inline comment in ``load_R``). Same axis, opposite reduction; taking UL of R + would inflate the closure ~15×. Returning h unchanged when the axis is absent + keeps non-helicity inputs working. """ if HELICITY_AXIS in [a.name for a in h.axes]: ul_idx = h.axes[HELICITY_AXIS].index(-1) @@ -74,9 +66,9 @@ def _select_ul_helicity(h): def _append_axis_overflow(h, axis_name): """Values array for ``h`` with ``axis_name``'s OVERFLOW bin appended as one - extra in-range bin along that axis; every other axis is taken in-range (no - flow). I.e. ``flow=False`` everywhere except that we keep ``axis_name``'s - overflow as a genuine trailing bin.""" + extra in-range bin along that axis; every other axis in-range (no flow). + I.e. ``flow=False`` everywhere except ``axis_name``'s overflow kept as a + trailing bin.""" full = h.values(flow=True) inr = h.values(flow=False).astype(np.float64) idx, pos = [], None @@ -101,14 +93,13 @@ def has_response( hist_name=DEFAULT_HIST, gen_total_name=DEFAULT_GENTOTAL, ): - """Cheap guard: does this histmaker output carry BOTH the reco x gen - response hist and the gen-total xnorm hist needed to build R *and* N_gen? + """True iff this histmaker output carries BOTH the reco x gen response hist + and the gen-total xnorm hist (needed for R *and* N_gen). - Used by setupRabbit to decide whether to embed the SCETlib-NP response in - the datacard (presence-based, *lenient* guard): returns True only when both - are present, so a generic unfolding run that has the response hist but not - the gen-total is a silent no-op rather than an error. Never raises (any - structural problem -> False); does not materialize any histogram. + setupRabbit uses this to decide whether to embed the SCETlib-NP response in + the datacard. Requiring both makes a generic unfolding run (response hist but + no gen-total) a silent no-op, not an error. Never raises (any structural + problem -> False); materializes no histogram. """ try: with h5py.File(unfolding_hdf5_path, "r") as f: @@ -141,10 +132,10 @@ def load_R( source : (path, sample_key, hist_name) for traceability ``ptVGen_overflow`` (default True): append the gen ptVGen overflow (true - qT > last gen edge) as a trailing gen bin in R and N_gen, with edge - ``PTVGEN_OVERFLOW_EDGE``. This lets the model fold a σ_gen(qT>44) through the - migration into the high-ptll reco bins (see the PTVGEN_OVERFLOW_EDGE note); - set False for the legacy in-range-only response. + qT > last gen edge) as a trailing gen bin in R and N_gen, edge + ``PTVGEN_OVERFLOW_EDGE``, so the model can fold σ_gen(qT>44) through the + migration into the high-ptll reco bins (see the PTVGEN_OVERFLOW_EDGE note). + False = legacy in-range-only response. """ with h5py.File(unfolding_hdf5_path, "r") as f: if sample_key not in f: @@ -180,38 +171,37 @@ def load_R( f"{hist_name}: missing expected axes {missing}. " f"Got: {ax_names}" ) - # Select acceptance=True (gen events in fiducial), then keep the reco + - # gen axes. project() SUMS helicitySig out — which is correct *for R*: - # the joint yield is filled with `nominal_weight_helicity` - # (= nominal_weight × helWeight_tensor, see helicity_utils), a PARTITION - # of the event weight into the 8 helicity pieces g_i(cosθ,φ) that ADD - # BACK UP to the full angular weight. So Σ_helicitySig R is the physical - # angular-resolved reco×gen yield (the angular dependence lives in the - # cosThetaStar*/phiStar* reco bins). NB this is the OPPOSITE reduction - # from N_gen below: N_gen is filled with `csAngularMoments` (a moment - # expansion whose 0..7 bins do NOT sum to σ), so it takes UL (-1). Same - # axis, different fill tensor → different recovery. Taking UL of R would - # discard the angular partition and inflate the closure (~15×). + # Select acceptance=True (fiducial gen), keep reco + gen axes. project() + # SUMS helicitySig — correct *for R*: the joint yield is filled with + # `nominal_weight_helicity` (= nominal_weight × helWeight_tensor, see + # helicity_utils), a PARTITION of the event weight into the 8 helicity + # pieces g_i(cosθ,φ) that ADD BACK UP to the full angular weight. So + # Σ_helicitySig R is the physical angular-resolved reco×gen yield (angular + # dependence lives in the cosThetaStar*/phiStar* reco bins). OPPOSITE + # reduction from N_gen below: N_gen is filled with `csAngularMoments` (a + # moment expansion whose 0..7 bins do NOT sum to σ), so it takes UL (-1). + # Same axis, different fill tensor → different recovery. Taking UL of R + # would discard the angular partition and inflate the closure (~15×). h_sel = h[{"acceptance": True}] h_proj = h_sel.project(*reco_axes, *gen_axes) # Gen-total denominator N_gen(g): the xnorm histogram (e.g. "postfsr"), # filled on fiducial gen events BEFORE reco selection. Its gen marginal - # is the generated total per gen bin (efficiency NOT yet applied), so - # N_reco(b,g)/N_gen(g) = efficiency × migration — the theory-independent - # gen→reco response. (The gen marginal of R itself is reco-passing, i.e. - # already × efficiency, which is the wrong normalizer.) + # is the generated total per gen bin (no efficiency yet), so + # N_reco(b,g)/N_gen(g) = efficiency × migration, the theory-independent + # gen→reco response. (R's own gen marginal is reco-passing, i.e. already + # × efficiency — the wrong normalizer.) N_gen_hist = None if gen_total_name is not None and gen_total_name in output: gp = output[gen_total_name] hg = gp.get() if hasattr(gp, "get") else gp - # postfsr/prefsr axes: (count, ptVGen, absYVGen, helicitySig). This - # one is filled with `csAngularMoments`: a moment expansion whose UL - # bin (-1) is the angular-integrated total σ, while bins 0..7 are the - # A_i moments (orthogonal, signed) that do NOT sum to σ — summing them - # overcounts by ~19% here. So take UL. (R above is the opposite: a - # weight partition, recovered by SUMMING helicitySig.) Then project - # to the gen axes (sums the trivial 'count' axis) to match R's binning. + # postfsr/prefsr axes: (count, ptVGen, absYVGen, helicitySig). Filled + # with `csAngularMoments`: a moment expansion whose UL bin (-1) is the + # angular-integrated total σ, while bins 0..7 are the A_i moments + # (orthogonal, signed) that do NOT sum to σ (summing overcounts ~19% + # here). So take UL. (R above is the opposite: a weight partition, + # recovered by SUMMING helicitySig.) Then project to the gen axes + # (sums the trivial 'count' axis) to match R's binning. hg = _select_ul_helicity(hg) hg_gen = hg.project(*gen_axes) # Fold the gen ptVGen overflow into a trailing bin (or drop it). @@ -221,9 +211,9 @@ def load_R( else hg_gen.values(flow=False).astype(np.float64) ) - # Out from the with-block: hist is materialized. Keep the gen ptVGen overflow - # as a trailing gen bin so the model can supply σ_gen(qT>44) — else the high- - # ptll reco bins (fed by true qT>44 migrating down) come out low. + # Out of the with-block: hist materialized. Keep the gen ptVGen overflow as a + # trailing gen bin so the model can supply σ_gen(qT>44); else the high-ptll + # reco bins (fed by true qT>44 migrating down) come out low. R = ( _append_axis_overflow(h_proj, "ptVGen") if ptVGen_overflow diff --git a/wremnants/postprocessing/scetlib_np/sigma_gen.py b/wremnants/postprocessing/scetlib_np/sigma_gen.py new file mode 100644 index 000000000..dff831d6c --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/sigma_gen.py @@ -0,0 +1,576 @@ +"""SigmaGenModel — the datacard-free SCETlib NP σ_gen(λ) core (Steps 1–2). + +The physics half of the SCETlib NP prediction, factored out of +:class:`~wremnants.postprocessing.scetlib_np.param_model.SCETlibNPParamModel` so +it runs from a bT-grid directory, λ_central (with the np_model strings), and the +gen-bin edges — no rabbit / datacard / fit input. Owns Steps 1–2: + + Step 1 btgrid Hankel + Q integral → σ_resum(λ; g) resummed, on the gen grid + Step 2 + fixed-order matching → σ_gen(λ; g) = σ_resum(λ; g) + σ_ns(g) + +Every factor (b*, I_pert, C_ν, γ_ν^NP, F_eff, the arctan-Q² Q integral, the +|Y|-fold and qT→ptVGen rebin, σ_ns matching) is derived in the ``param_model.py`` +module docstring; this class is that arithmetic without loader/fit-interface +concerns. Steps 3–4 (gen→reco fold, per-bin ratio) stay in +``SCETlibNPParamModel``, which holds a ``SigmaGenModel`` as ``self.core``. + +Public surface (used by the validation scripts and the σ_gen-at-λ tool): + + eff_central, gnu_central, np_model, np_model_nu λ_central + functional forms + gen_axes, gen_shape the (ptVGen, absYVGen) gen grid + Y_unique, qT_unique, Q_unique btgrid native axes + sigma_ns NP-independent FO nonsingular + sigma_YqT_central native (NY, NqT) resum-only σ at λ_c + sigma_gen_central matched σ_gen on the gen grid at λ_c + sigma_YqT_native(eff, gnu) native (NY, NqT) σ(λ), pre-fold + sigma_gen(eff, gnu[, sigma_YqT]) matched σ_gen(λ) on the gen grid + +``eff``/``gnu`` are dicts of the λ values plus the ``np_model`` / +``np_model_nu`` form strings, same shape as ``eff_central`` / ``gnu_central``. +Start from those and override the λ you want; un-supplied params stay at +λ_central. All tensors are TF, so ``sigma_gen`` is differentiable in λ. +""" + +import os +from typing import Optional + +import numpy as np +import tensorflow as tf + +from wremnants.postprocessing.scetlib_np import btgrid_cache +from wremnants.postprocessing.scetlib_np import btgrid_integrate as fz_int +from wremnants.postprocessing.scetlib_np import btgrid_tf as fz_tf +from wremnants.postprocessing.scetlib_np.params import ( # noqa: F401 (re-export) + ALL_PARAMS, + EFF_PARAMS, + GNU_PARAMS, + bin_sum_matrix, +) +from wremnants.utilities import common as wrem_common +from wremnants.utilities.data_paths import getDataPath + +_NONSING_FO_SING_DEFAULT = os.path.join( + wrem_common.data_dir, + "TheoryCorrections", + "inclusive_Z_COM13_CT18Z_N3+0LL_lattice_lambda4bugfix_fine_nnlo_sing_combined.pkl", +) +_NONSING_DYTURBO_DEFAULT = os.path.join( + wrem_common.data_dir, + "TheoryCorrections", + "results_z-2d-nnlo-vj-CT18ZNNLO-{scale}-scetlibmatch.txt", +) +_BTGRID_SUBDIR = ("scetlib_np", "Z_COM13_CT18Z_N3p0LL_btgrid_fineall") + +# λ name tuples (GNU_PARAMS / EFF_PARAMS / ALL_PARAMS) and bin_sum_matrix are +# imported above from :mod:`params` and re-exported here for back-compat. + + +def _default_btgrid_dir(): + base = getDataPath(fallback="/scratch/submit/cms/wmass/NanoAOD") + return os.path.join(os.path.dirname(base), *_BTGRID_SUBDIR) + + +def compute_nonsingular_gen( + fo_sing_path, + dyturbo_path, + gen_axes_meta, + charge=0, + q_lo=60.0, + q_hi=120.0, + qt_cutoff=1.0, + dyturbo_axes=("Q", "Y", "qT"), +): + """Nonsingular FO term on the model gen grid (NptVGen, NabsYVGen). + + The fixed-order/DYTurbo matching adds a NP-INDEPENDENT piece to σ_gen: + σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns , + σ_ns = (DYTurbo fixed order) − (SCETlib singular fixed order) , + the ``-hfo_sing + hfo`` that ``read_matched_scetlib_hist`` forms. + ``fo_sing_path`` is the SCETlib singular ``…_nnlo_sing…combined.pkl``; + ``dyturbo_path`` is the DYTurbo FO ``results_…scetlibmatch.txt`` (``{scale}`` + → mur1-muf1 for the central). σ_ns is zeroed below ``qt_cutoff`` (as + make_theory_corr does), Q-windowed to [q_lo, q_hi], |Y|-folded, and projected + onto the coarse (ptVGen, absYVGen) gen bins by SUMMING the bin-integrated + native bins. + """ + from wremnants.utilities.io_tools import input_tools + from wums import boostHistHelpers as hh + + def _central(h): + if "vars" in h.axes.name: + names = list(h.axes["vars"]) + idx = 0 + for c in ("central", "pdf0", "nominal"): + if c in names: + idx = names.index(c) + break + h = h[{"vars": idx}] + return h + + # SCETlib singular FO and DYTurbo FO, from their own files. + dyturbo_path = ( + dyturbo_path.format(scale="mur1-muf1") + if "{scale}" in dyturbo_path + else dyturbo_path + ) + hfo_sing = _central(input_tools.read_scetlib_hist(fo_sing_path, charge=charge)) + hfo = input_tools.read_dyturbo_hist( + [dyturbo_path], axes=list(dyturbo_axes), charge=charge + ) + if "vars" in hfo.axes.name: + hfo = _central(hfo) + + # Align shared physics axes (DYTurbo is coarser), then σ_ns = DYTurbo − singular. + for ax in ("Y", "Q", "qT"): + if ax in set(hfo.axes.name) & set(hfo_sing.axes.name): + hfo, hfo_sing = hh.rebinHistsToCommon([hfo, hfo_sing], ax) + nonsing_h = hh.addHists(-1.0 * hfo_sing, hfo, flow=False, by_ax_name=False) + + if "charge" in nonsing_h.axes.name: + nonsing_h = nonsing_h[{"charge": sum}] + # Q-window: slice(...,sum) sums ONLY the in-range Q bins (no underflow leak). + Qe = np.asarray(nonsing_h.axes["Q"].edges, dtype=np.float64) + qi = int(np.argmin(np.abs(Qe - q_lo))) + qj = int(np.argmin(np.abs(Qe - q_hi))) + nonsing_h = nonsing_h[{"Q": slice(qi, qj, sum)}] + nonsing_h = hh.makeAbsHist(nonsing_h, "Y") # signed Y -> |Y| + + qT_c = np.asarray(nonsing_h.axes["qT"].centers, dtype=np.float64) + absY_c = np.asarray(nonsing_h.axes["absY"].centers, dtype=np.float64) + v = nonsing_h.project("qT", "absY").values(flow=False) # (qT, absY) + v[qT_c < qt_cutoff, :] = 0.0 # zero the nonsingular below the cutoff + + ptV_edges = np.asarray(gen_axes_meta[0][1], dtype=np.float64) + absY_edges = np.asarray(gen_axes_meta[1][1], dtype=np.float64) + Wp = bin_sum_matrix(qT_c, ptV_edges) # (NptVGen, NqT) + Wa = bin_sum_matrix(absY_c, absY_edges) # (NabsYVGen, NabsYsrc) + return Wp @ v @ Wa.T # (NptVGen, NabsYVGen) + + +# ============================================================================ +# Factorized derived cache +# ============================================================================ +# ``combined_btgrid.pkl`` (:mod:`btgrid_cache`) holds the RAW grid. Deriving the +# reconstruction layout from it (sanitize non-finite cells → dense index map → +# ``dedup_grid_rows`` → weighted J0 kernel) is the slow part of construction +# (~18 GB load + dedup) and a PURE function of that grid, so we memoize the +# derived arrays in an .npz next to the pickle; repeat constructions skip the raw +# load and dedup. Invalidated on combined-pickle change (mtime+size) OR derivation +# code change (bump _FACTORIZED_SCHEMA_VERSION). The combined pickle is REQUIRED: +# it is what freshness is verified against; absent → rebuild, not trust the .npz. +# The .npz lives in the btgrid dir, shared across all users of that grid. +_FACTORIZED_CACHE_BASENAME = "combined_btgrid.factorized.npz" +_FACTORIZED_SCHEMA_VERSION = "factorized_v1" +# The arrays that fully populate the factorized layout (see _assign_factorized). +_FACTORIZED_KEYS = ( + "flat_idx", "Q_unique", "Y_unique", "qT_unique", "bT", "b_bar", + "Y_feff_unique", "bT_simpson_w", "I_pert_u", "C_nu_uu", "c_of_u", + "feff_idx_u", "gather_idx", "KwqT", +) + + +def _build_factorized_arrays(grid): + """Derive the factorized-layout numpy arrays from a raw combined grid. + + Sanitize non-finite cells, build the dense index map, dedup the (I_pert, + C_nu) rows (bit-exact-verified inside ``dedup_grid_rows``), fold the per-qT + prefactor + bT Simpson weights into the unique-qT J0 kernel. A pure function + of ``grid``; the returned dict is what :meth:`_assign_factorized` turns into + TF constants and what the derived cache stores.""" + # Sanitize non-finite bt-grid cells. Kinematically-forbidden points + # (x = (Q/Ecm)·e^|Y| ≥ 1, e.g. extreme forward Y near the Z peak) come back + # as NaN from SCETlib instead of the physical 0; their true σ is 0, and + # dedup_grid_rows' hash-group verification needs finite cells (NaN != NaN). + for _key in ("I_pert", "C_nu"): + _arr = grid[_key] + if not np.isfinite(_arr).all(): + _nbad = int((~np.isfinite(_arr)).any(axis=-1).sum()) + np.nan_to_num(_arr, copy=False, nan=0.0, posinf=0.0, neginf=0.0) + print( + f"[SigmaGenModel] sanitized {_nbad} non-finite {_key} bt-grid " + f"rows -> 0 (kinematically-forbidden cells)", + flush=True, + ) + idx_map = fz_int.dense_index_map(grid["bins"]) + bins = grid["bins"] + qT_pb = np.array([b[2] for b in bins], dtype=np.float64) + Y_pb = np.array([b[1] for b in bins], dtype=np.float64) + # F_eff depends on the bin only through Y (few distinct values): map to + # unique Y so the NP transcendentals run on NY rows and gather. + Y_feff_unique, Y_feff_inv = np.unique(Y_pb, return_inverse=True) + Y_feff_inv = Y_feff_inv.reshape(-1).astype(np.int32) + bT = np.asarray(grid["bT"], dtype=np.float64) + b_bar = np.asarray(grid["b_bar"], dtype=np.float64) + bT_simpson = np.asarray(fz_tf.simpson_weights(bT), dtype=np.float64) + # Dedup the (I_pert, C_nu) rows (~2x), plus a 2nd-level C_nu dedup — both + # verified bit-exact inside dedup_grid_rows. + dd = fz_tf.dedup_grid_rows(grid["I_pert"][0], grid["C_nu"][0], Y_feff_inv) + # Per-bin index into the unique-qT axis (exact lookup, asserted). + qT_idx = np.searchsorted(idx_map["qT_unique"], qT_pb) + assert np.array_equal(idx_map["qT_unique"][qT_idx], qT_pb) + gather_idx = np.stack( + [dd["row_uid"].astype(np.int64), qT_idx.astype(np.int64)], axis=1 + ).astype(np.int32) + # Weighted J0 kernel on the unique-qT grid (per-qT prefactor + bT Simpson + # weights folded in): (NqT, Nbt). + qTu = np.asarray(idx_map["qT_unique"], dtype=np.float64) + K_u = fz_tf.build_bT_J0_kernel( + tf.constant(qTu, dtype=fz_tf.DTYPE), tf.constant(bT, dtype=fz_tf.DTYPE) + ) + KwqT = ( + tf.constant(qTu, dtype=fz_tf.DTYPE)[:, tf.newaxis] + * K_u + * tf.constant(bT_simpson, dtype=fz_tf.DTYPE)[tf.newaxis, :] + ).numpy() + return { + "flat_idx": np.asarray(idx_map["flat_idx"], dtype=np.int64), + "Q_unique": np.asarray(idx_map["Q_unique"], dtype=np.float64), + "Y_unique": np.asarray(idx_map["Y_unique"], dtype=np.float64), + "qT_unique": qTu, + "bT": bT, + "b_bar": b_bar, + "Y_feff_unique": np.asarray(Y_feff_unique, dtype=np.float64), + "bT_simpson_w": bT_simpson, + "I_pert_u": np.asarray(dd["I_u"], dtype=np.float64), + "C_nu_uu": np.asarray(dd["C_uu"], dtype=np.float64), + "c_of_u": np.asarray(dd["c_of_u"], dtype=np.int32), + "feff_idx_u": np.asarray(dd["feff_idx_u"], dtype=np.int32), + "gather_idx": gather_idx, + "KwqT": KwqT, + } + + +def _read_factorized_cache(cache_path, btgrid_dir): + """Return the cached factorized arrays (an open ``NpzFile``) if present AND + consistent with the current combined pickle, else ``None``. + + The combined pickle is REQUIRED so freshness can be VERIFIED: it must exist + and be fresh vs its shards, the schema tag must match, and the combined + (mtime+size) must match what was recorded at cache-write. Absent pickle → + cannot verify → rebuild (from the shards) rather than trust a stale .npz. Any + read error → ``None`` (rebuild).""" + if not os.path.exists(cache_path): + return None + combined = btgrid_cache.combined_path(btgrid_dir) + if not (os.path.exists(combined) and btgrid_cache.is_combined_fresh(btgrid_dir)): + return None + try: + z = np.load(cache_path, allow_pickle=False) + st = os.stat(combined) + if ( + str(z["schema_version"].item()) != _FACTORIZED_SCHEMA_VERSION + or int(z["combined_mtime_ns"].item()) != int(st.st_mtime_ns) + or int(z["combined_size"].item()) != int(st.st_size) + or any(k not in z.files for k in _FACTORIZED_KEYS) + ): + return None + except Exception as exc: # corrupt / partial cache -> rebuild + print( + f"[SigmaGenModel] ignoring unreadable factorized cache " + f"{cache_path} ({exc})", + flush=True, + ) + return None + return z + + +def _write_factorized_cache(cache_path, arr, btgrid_dir): + """Atomically write the derived arrays + a freshness header next to the + combined pickle. PID-unique temp name so concurrent builders (e.g. parallel + fit jobs) don't clobber each other; ``os.replace`` is atomic.""" + st = os.stat(btgrid_cache.combined_path(btgrid_dir)) + header = { + "schema_version": np.array(_FACTORIZED_SCHEMA_VERSION), + "combined_mtime_ns": np.int64(st.st_mtime_ns), + "combined_size": np.int64(st.st_size), + } + tmp = f"{cache_path}.tmp.{os.getpid()}.npz" + np.savez(tmp, **header, **arr) + os.replace(tmp, cache_path) + + +class SigmaGenModel: + """Btgrid → σ_gen(λ) on a fixed (ptVGen, absYVGen) gen grid. See module docstring.""" + + def __init__( + self, + btgrid_dir: Optional[str] = None, + lambda_central: Optional[dict] = None, + gen_axes=None, + Q_lo: float = 60.0, + Q_hi: float = 120.0, + nonsingular_fo_sing: str = _NONSING_FO_SING_DEFAULT, + nonsingular_dyturbo: str = _NONSING_DYTURBO_DEFAULT, + nonsingular_qt_cutoff: float = 1.0, + include_nonsingular: bool = True, + ): + """Build the σ_gen core. + + Parameters + ---------- + btgrid_dir + Directory holding the SCETlib bT-grid ``combined_btgrid.pkl``. + Defaults (when None) to the shared data-area copy next to NanoAOD. + lambda_central + Dict with ``eff_params`` and ``gnu_params`` sub-dicts (same shape as + :func:`lambda_central.read_lambda_central`). Carries both the central + λ values and the ``np_model`` / ``np_model_nu`` functional-form strings. + gen_axes + Ordered ``[("ptVGen", edges), ("absYVGen", edges)]`` — the gen grid + σ_gen is rebinned onto. ``edges`` are 1-D arrays of bin edges. + Q_lo, Q_hi + Z mass window for the Q-integration on the btgrid. + nonsingular_fo_sing, nonsingular_dyturbo, nonsingular_qt_cutoff + σ_ns = DYTurbo − SCETlib_singular inputs / low-qT cutoff (see + :func:`compute_nonsingular_gen`). + include_nonsingular + Add the matched FO nonsingular σ_ns (default True — the matched σ_gen + the histmaker nominal carries). False → resum-only (σ_ns = 0), FO + inputs not read. + + The derived factorized btgrid layout is always memoized in an .npz next to + ``combined_btgrid.pkl`` (see "Factorized derived cache" above). No on/off + knob: staleness is auto-detected, and a fresh cache lets construction skip + the ~18 GB raw load and the row dedup. + """ + if lambda_central is None: + raise ValueError("SigmaGenModel requires lambda_central (eff/gnu params).") + if gen_axes is None or len(gen_axes) != 2: + raise ValueError( + "SigmaGenModel requires gen_axes = [(ptVGen, edges), (absYVGen, edges)]." + ) + if btgrid_dir is None: + btgrid_dir = _default_btgrid_dir() + + # ---- λ_central + functional forms. + self.eff_central = dict(lambda_central["eff_params"]) + self.gnu_central = dict(lambda_central["gnu_params"]) + self.np_model = self.eff_central["np_model"] + self.np_model_nu = self.gnu_central["np_model_nu"] + + # ---- gen grid. + 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) + + # ---- btgrid factorized reconstruction layout (Step-1 tensors). Loaded + # from the derived .npz cache when fresh, else built from the raw grid + # (sanitize → dense index → row dedup → weighted J0 kernel) and cached. + # Sets: flat_idx, Q/Y/qT_unique, bT, b_bar, Y_feff_unique, bT_simpson_w, + # I_pert_u, C_nu_uu, c_of_u, feff_idx_u, gather_idx, KwqT. + self._setup_btgrid(btgrid_dir) + + # ---- Q-integration weights (arctan_Q² Simpson on Z mass window). + self.Q_weights = tf.constant( + fz_int.q_integrate_weights(self.Q_unique, Q_lo, Q_hi), + dtype=fz_tf.DTYPE, + ) + + # ---- Rebin weights: btgrid (NY signed) → (NabsYVGen) via |Y| folding, + # (NqT) → (NptVGen). + ptVGen_edges = self.gen_axes[0][1] + absY_edges = self.gen_axes[1][1] + # |Y| folding: σ(Y) symmetric in Y, so the absY-bin integral is + # 2·∫_{absY_lo}^{absY_hi} σ(Y) dY. Use Y >= 0 samples, multiply by 2. + Y_pos_mask = self.Y_unique >= 0 + Y_pos = self.Y_unique[Y_pos_mask] + absY_rebin_pos = fz_int.rebin_weights(Y_pos, absY_edges, name="absY") + # Pad to full NY: zero on negative-Y columns. + W_absY = np.zeros((absY_edges.size - 1, self.Y_unique.size), dtype=np.float64) + W_absY[:, Y_pos_mask] = 2.0 * absY_rebin_pos + self.W_absY = tf.constant(W_absY, dtype=fz_tf.DTYPE) + # qT rebin: btgrid qT (signed nonneg, NqT=141) → ptVGen edges. With a + # ptVGen overflow bin [last_gen_edge, OVERFLOW_EDGE] (e.g. [44, 100]), + # rebin_weights' last row Simpson-integrates the btgrid tail qT∈(44,100] + # into it; btgrid qT past the last edge (>100, off-grid) is dropped + # (negligible). + self.W_ptVGen = tf.constant( + fz_int.rebin_weights(self.qT_unique, ptVGen_edges, name="ptVGen"), + dtype=fz_tf.DTYPE, + ) + + # ---- Native (NY, NqT) Q-integrated reconstruction at λ_central, BEFORE + # the |Y|-fold and qT-rebin — exposed so the native-binning validation + # compares it to the SCETlib reference without the projection layer. + self.sigma_YqT_central = self.sigma_YqT_native(self.eff_central, self.gnu_central) + + # ---- Fixed-order/DYTurbo nonsingular term (NP-independent). + # σ_gen^matched(λ) = σ_gen^resum(λ) + σ_ns, added at GEN level so it folds + # through the same response R as the resummed piece. σ_ns is constant (no + # λ dependence); included by default (the matched σ_gen the histmaker + # nominal carries). include_nonsingular=False → resum-only. + if include_nonsingular: + _dy0 = ( + nonsingular_dyturbo.format(scale="mur1-muf1") + if (nonsingular_dyturbo and "{scale}" in nonsingular_dyturbo) + else nonsingular_dyturbo + ) + missing = [ + p for p in (nonsingular_fo_sing, _dy0) if not (p and os.path.exists(p)) + ] + if missing: + raise FileNotFoundError( + "The matched model needs the fixed-order inputs for " + "σ_ns = DYTurbo − SCETlib_singular, but these are missing:\n " + + "\n ".join(missing) + + "\nThey live under wremnants-data/data/TheoryCorrections (the " + "SCETlib singular …_nnlo_sing…combined.pkl and the DYTurbo " + "results_…scetlibmatch.txt). Pass nonsingular_fo_sing / " + "nonsingular_dyturbo to point at them (or include_nonsingular=" + "False for resum-only)." + ) + sigma_ns_np = compute_nonsingular_gen( + nonsingular_fo_sing, + nonsingular_dyturbo, + self.gen_axes, + q_lo=Q_lo, + q_hi=Q_hi, + qt_cutoff=nonsingular_qt_cutoff, + ) + if sigma_ns_np.shape != tuple(self.gen_shape): + raise ValueError( + f"nonsingular gen shape {sigma_ns_np.shape} != model gen shape " + f"{tuple(self.gen_shape)}" + ) + self.sigma_ns = tf.constant(sigma_ns_np, dtype=fz_tf.DTYPE) + else: + self.sigma_ns = tf.zeros(self.gen_shape, dtype=fz_tf.DTYPE) + + # ---- Matched σ_gen(λ_central). Reuse the native (NY, NqT) integral + # already computed for sigma_YqT_central (no 2nd bT reconstruction). + sigma_gen_central = self.sigma_gen( + self.eff_central, self.gnu_central, sigma_YqT=self.sigma_YqT_central + ) + self.sigma_gen_central = sigma_gen_central + gen_flat = tf.reshape(sigma_gen_central, [-1]) + if tf.reduce_any(gen_flat <= 0).numpy(): + n_bad = int(tf.reduce_sum(tf.cast(gen_flat <= 0, tf.int32))) + raise ValueError( + f"SigmaGenModel: {n_bad} gen bins have non-positive " + f"σ_gen(λ_central); cannot normalize / fold the response." + ) + + # ========================================================================= + # btgrid factorized layout (derived-cache aware) + # ========================================================================= + + def _setup_btgrid(self, btgrid_dir): + """Populate the factorized-layout tensors: from the derived cache when + usable (see :func:`_read_factorized_cache`), else built from the combined + grid and cached next to it for reuse. Staleness handled automatically + (combined mtime+size + schema).""" + cache_path = os.path.join(btgrid_dir, _FACTORIZED_CACHE_BASENAME) + z = _read_factorized_cache(cache_path, btgrid_dir) + if z is not None: + self._assign_factorized(z) + print( + f"[SigmaGenModel] loaded factorized btgrid cache {cache_path}", + flush=True, + ) + return + grid = btgrid_cache.load(btgrid_dir) + arr = _build_factorized_arrays(grid) + del grid # free the ~18 GB host grid before TF graph build + self._assign_factorized(arr) + try: + _write_factorized_cache(cache_path, arr, btgrid_dir) + print( + f"[SigmaGenModel] wrote factorized btgrid cache {cache_path}", + flush=True, + ) + except OSError as exc: # read-only area etc. — still usable, just slow + print( + f"[SigmaGenModel] WARNING: could not write factorized cache " + f"{cache_path} ({exc}); continuing without it", + flush=True, + ) + + def _assign_factorized(self, a): + """Set the factorized-layout attributes from a mapping of numpy arrays + (:func:`_build_factorized_arrays` output or a loaded npz). Q/Y/qT_unique + stay numpy (used in numpy rebin-weight construction); the rest become TF + constants.""" + D = fz_tf.DTYPE + self.flat_idx = tf.constant(a["flat_idx"], dtype=tf.int64) + self.Q_unique = np.asarray(a["Q_unique"], dtype=np.float64) + self.Y_unique = np.asarray(a["Y_unique"], dtype=np.float64) + self.qT_unique = np.asarray(a["qT_unique"], dtype=np.float64) + self.bT = tf.constant(a["bT"], dtype=D) + self.b_bar = tf.constant(a["b_bar"], dtype=D) + self.Y_feff_unique = tf.constant(a["Y_feff_unique"], dtype=D) + self.bT_simpson_w = tf.constant(a["bT_simpson_w"], dtype=D) + self.I_pert_u = tf.constant(a["I_pert_u"], dtype=D) # (Nu, Nbt) + self.C_nu_uu = tf.constant(a["C_nu_uu"], dtype=D) # (Ncu, Nbt) + self.c_of_u = tf.constant(a["c_of_u"], dtype=tf.int32) + self.feff_idx_u = tf.constant(a["feff_idx_u"], dtype=tf.int32) + self.gather_idx = tf.constant(a["gather_idx"], dtype=tf.int32) + self.KwqT = tf.constant(a["KwqT"], dtype=D) + + # ========================================================================= + # σ_gen evaluation + # ========================================================================= + + def sigma_YqT_native(self, eff_params, gnu_params, np_model=None, np_model_nu=None): + """Reconstruct σ(λ) on the btgrid and Q-integrate, in the btgrid's + *native* binning: shape (NY, NqT) on the signed-Y / qT grid (Y_unique, + qT_unique), BEFORE the |Y|-fold and qT-rebin. The object the native-binning + validation compares against the SCETlib spectrum reference (curve 1) and + the external scetlib_run.factorize (curve 2). + + ``np_model`` / ``np_model_nu`` override the functional form applied to the + λ for THIS evaluation (default: the construction forms ``self.np_model`` / + ``self.np_model_nu``). The bt-grid (I_pert, C_nu) is NP-model-independent, + so a different form is just a different analytic factor on the same grid — + used to evaluate a numerator in one form while the denominator (central) + stays in the card's form (see ``SCETlibNPParamModel``).""" + # 1. Reconstruct σ on the btgrid's flat (Nbins,) layout via the + # memory-factorized path (deduplicated rows + unique-qT J0 kernel + + # Simpson-as-matmul) — ~6x smaller than dense (Nbins, Nbt), which is what + # lets the fit run on a 32 GB GPU. + eff = {k: v for k, v in eff_params.items() if k != "np_model"} + gnu = {k: v for k, v in gnu_params.items() if k != "np_model_nu"} + sigma_flat = fz_tf.reconstruct_batch_factorized_tf( + b_bar=self.b_bar, + I_pert_u=self.I_pert_u, + C_nu_uu=self.C_nu_uu, + c_of_u=self.c_of_u, + eff_params=eff, + gnu_params=gnu, + np_model=np_model or self.np_model, + np_model_nu=np_model_nu or self.np_model_nu, + KwqT=self.KwqT, + gather_idx=self.gather_idx, + Y_unique=self.Y_feff_unique, + feff_idx_u=self.feff_idx_u, + ) + # 2. Sparse → dense (NQ, NY, NqT). Missing cells get 0. + sigma_dense = fz_int.sparse_to_dense_tf(sigma_flat, self.flat_idx) + # 3. Integrate over Q (arctan_Q² Simpson) → (NY, NqT). + return fz_int.integrate_over_Q_tf(sigma_dense, self.Q_weights) + + def sigma_gen( + self, eff_params, gnu_params, sigma_YqT=None, np_model=None, np_model_nu=None + ): + """Evaluate matched σ_gen(λ) on the gen binning. Returns (NptVGen, NabsYVGen). + + ``sigma_YqT`` passes an already-computed native (NY, NqT) integral to skip + the expensive bT reconstruction; used at construction to reuse + ``sigma_YqT_central`` instead of integrating λ_central twice. + ``np_model`` / ``np_model_nu`` override the functional form for this + evaluation (default: the construction forms) — see ``sigma_YqT_native``. + """ + if sigma_YqT is None: + sigma_YqT = self.sigma_YqT_native( + eff_params, gnu_params, np_model=np_model, np_model_nu=np_model_nu + ) + # 4. Rebin Y (signed) → absYVGen (|Y|-folded): (NabsYVGen, NqT). + sigma_absY_qT = fz_int.rebin_axis_tf(sigma_YqT, axis=0, weights=self.W_absY) + # 5. Rebin qT → ptVGen: (NabsYVGen, NptVGen). + sigma_absY_ptV = fz_int.rebin_axis_tf( + sigma_absY_qT, axis=1, weights=self.W_ptVGen + ) + # 6. Reorder to (NptVGen, NabsYVGen) to match R's gen axis order. + sigma_resum = tf.transpose(sigma_absY_ptV, perm=[1, 0]) + # 7. Add the NP-independent fixed-order/DYTurbo nonsingular (zeros if off). + return sigma_resum + self.sigma_ns diff --git a/wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py b/wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py new file mode 100644 index 000000000..8dd70f777 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py @@ -0,0 +1,661 @@ +"""Evaluate the SCETlib NP matched σ_gen at a given λ tune, via the core. + +Builds the datacard-free core ``SigmaGenModel`` and runs its matched gen-level +prediction σ_gen(λ; g) = σ_resum(λ; g) + σ_ns on the (ptVGen, absYVGen) gen grid +— Steps 1–2 of ``param_model.py`` (the object that folds through R into the fit, +BEFORE the gen→reco fold and ratio). Prints σ_gen on the gen grid and can plot a +1-D projection (e.g. the ptZ = ptVGen distribution). + +The λ to evaluate at = a physical BASE tune + overrides: + * base: ``--meta-from HDF5`` / the ``--theory-corr`` file's Nonperturbative + runcard / else the canonical FranksVals tanh_2 default. The model is BUILT at + this base (positive σ_gen, which the constructor requires); λ are evaluated on + top, so params not set stay at the BASE value, not 0; + * ``--fitresult HDF5`` postfit λ (optional); + * ``--lambdas name=val,...`` explicit values (optional; win over the rest). +Common use is ``--lambdas lambda2=0.5``, the rest staying at FranksVals; no +λ_central source needed. Evaluating, unlike constructing, has no positivity +guard, so a weak-NP tune can be inspected; non-positive bins are warned, not +rejected. + +Can ALSO overlay the gen distribution in an official ``TheoryCorrection`` hist +(``--theory-corr``). Those ``.pkl.lz4`` files carry the ``{generator}_hist`` +object — the official SCETlib+DYTurbo prediction on a (Q, absY, qT, charge, vars) +grid — so its central (``pdf0``) entry is the same physical object the param-model +σ_gen reconstructs. Overlaying is a direct end-to-end check that the bt-grid + +on-the-fly reconstruction reproduces the official run; pick any ``vars`` label, +e.g. ``lambda21.0``, to overlay a λ-shifted official run against the model at the +matching λ. + +The bT-grid is required (``--btgrid``). The gen-bin edges (ptVGen, absYVGen) are +chosen per axis, in order: explicit ``--ptv-edges`` / ``--absy-edges``, then a +``--gen-edges-from`` / ``--datacard`` hdf5, then a built-in default (1-GeV ptVGen +bins over [0, 40]; a single rapidity-inclusive absYVGen bin [0, 5]) — so the +script runs with no gen-edge input. ``--datacard`` feeds ONLY the gen edges; λ +are not sourced from it (use ``--meta-from``). + +Run inside a container that binds the inputs (same as the validation scripts): + + export APPTAINER_BIND="/scratch,/cvmfs,/work,/ceph,/home" + singularity run --cleanenv bash -c \\ + "source main/WRemnants/setup.sh; \\ + python3 -m wremnants.postprocessing.scetlib_np.sigma_gen_at_lambda \\ + --lambdas lambda2=0.4,lambda4=0.1,lambda2_nu=0.15 \\ + --theory-corr /data/TheoryCorrections/scetlib_dyturbo_..._CorrZ.pkl.lz4 \\ + --plot ~/public_html/alphaS/YYMMDD_sigmagen/ptZ.png" +""" + +import argparse +import os +import sys +import time + +import numpy as np + +from wremnants.postprocessing.scetlib_np.params import ( + EFF_PARAMS, + GNU_PARAMS, + parse_lambda_overrides, +) + +# btgrid default mirrors the validation scripts; the datacard is intentionally +# NOT defaulted — pass --datacard (or explicit --*-edges) for the gen edges. +BTGRID_DIR = "/scratch/submit/cms/wmass/scetlib_np/Z_COM13_CT18Z_N3p0LL_btgrid_fineall/" +Q_LO, Q_HI = 60.0, 120.0 +# Canonical FranksVals (CT18Z N3+0LL lattice λ4-bugfix) tanh_2 runcard — the +# production λ_central. Construction BASE when no base λ is sourced: the model +# must be built at a PHYSICAL tune (positive σ_gen, so the constructor's response +# guard passes), and the requested λ evaluated on top. Source of truth: a +# correction file's Nonperturbative section (file_meta_data → config → +# Nonperturbative); the LatticeNPLambda4Bugfix_FranksVals_CT18Z values. +CANONICAL_BASE = { + "eff_params": { + "np_model": "tanh_2", "lambda2": 0.4, "lambda4": 0.4, + "lambda6": 0.0, "delta_lambda2": 0.0, "lambda_inf": 1.0, + }, + "gnu_params": { + "np_model_nu": "tanh_2", "lambda2_nu": 0.15, + "lambda4_nu": 0.0, "lambda6_nu": 0.0, "lambda_inf_nu": 2.0, + }, +} + +# Built-in gen grid used when neither explicit edges nor a datacard/hdf5 source is +# given: 1-GeV ptVGen bins over [0, 40], rapidity-inclusive in a single absYVGen +# bin [0, 5] (5.0 is a TheoryCorrection absY edge, so the overlay still aligns). +DEFAULT_PTV_EDGES = np.arange(0.0, 41.0, 1.0) +DEFAULT_ABSY_EDGES = np.array([0.0, 5.0]) + +# corr-hist axis ↔ model gen axis (the TheoryCorrection _hist uses SCETlib names). +_CORR_AXIS = {"ptVGen": "qT", "absYVGen": "absY"} + + +def _parse_edges(s): + """``a,b,c,...`` -> float ndarray of bin edges.""" + return np.array([float(x) for x in s.split(",") if x.strip()], dtype=np.float64) + + +def _merge_matrix(fine_edges, coarse_edges, name="axis", tol=1e-6): + """(N_coarse, N_fine) 0/1 matrix summing fine bins into coarse bins. + + Requires every coarse edge to coincide with a fine edge (coarse is a + sub-binning of fine): exact merge, no interpolation. Fine bins whose centre + lies outside every coarse bin (e.g. qT beyond the model's ptVGen overflow + edge) get weight 0 and are dropped, matching the model. + """ + fine_edges = np.asarray(fine_edges, dtype=np.float64) + coarse_edges = np.asarray(coarse_edges, dtype=np.float64) + for e in coarse_edges: + if not np.any(np.isclose(fine_edges, e, atol=tol)): + raise SystemExit( + f"_merge_matrix[{name}]: model edge {e} is not a TheoryCorrection " + f"bin edge (its binning is not a refinement of the model grid on " + f"this axis). corr edges: {fine_edges}" + ) + centers = 0.5 * (fine_edges[:-1] + fine_edges[1:]) + W = np.zeros((coarse_edges.size - 1, fine_edges.size - 1), dtype=np.float64) + for i in range(coarse_edges.size - 1): + m = (centers >= coarse_edges[i]) & (centers <= coarse_edges[i + 1]) + W[i, m] = 1.0 + return W + + +def resolve_base_lambda(args): + """Physical BASE λ tune (eff_params/gnu_params) the model is CONSTRUCTED at. + + Priority: ``--meta-from HDF5`` > the ``--theory-corr`` file's embedded + Nonperturbative runcard > the canonical FranksVals tanh_2 default. Always a + complete physical tune (never None), so construction lands on a positive-σ_gen + point (the constructor's response guard); the requested λ are evaluated on top. + """ + from wremnants.postprocessing.scetlib_np import lambda_central as lc + + if args.meta_from: + print(f"[λ base] from hdf5 metadata {args.meta_from}") + return lc.read_lambda_central(args.meta_from) + if args.theory_corr: + import pickle + + import lz4.frame + + with lz4.frame.open(args.theory_corr) as fh: + corr = pickle.load(fh) + base = lc.extract_lambda_central( + corr, tag=os.path.basename(args.theory_corr), + proc=args.theory_corr_proc or "Z", + ) + print(f"[λ base] from the --theory-corr Nonperturbative runcard " + f"({base.get('basename')})") + return {"eff_params": base["eff_params"], "gnu_params": base["gnu_params"]} + print("[λ base] none given -> canonical FranksVals tanh_2 default") + return CANONICAL_BASE + + +def assemble_tune(base, overrides): + """Full (eff_params, gnu_params) for the EVAL point = base tune + overrides, + plus the explicitly-set names. + + ``base`` is a physical lambda_central dict (with the np_model form strings); + params not in ``overrides`` stay at the base value (NOT 0). Each override is + routed to eff or gnu by membership.""" + eff = dict(base["eff_params"]) + gnu = dict(base["gnu_params"]) + explicit = {} + for name, val in overrides.items(): + if name in EFF_PARAMS: + eff[name] = val + elif name in GNU_PARAMS: + gnu[name] = val + else: + raise SystemExit( + f"unknown λ {name!r}; valid: {list(GNU_PARAMS) + list(EFF_PARAMS)}" + ) + explicit[name] = val + return eff, gnu, explicit + + +def resolve_gen_axes(args): + """gen_axes = [(ptVGen, edges), (absYVGen, edges)], chosen per axis in order: + explicit --ptv-edges/--absy-edges, then a --gen-edges-from/--datacard hdf5, + then the built-in defaults (1-GeV ptVGen [0,40]; single absYVGen [0,5]).""" + ptv = _parse_edges(args.ptv_edges) if args.ptv_edges else None + absy = _parse_edges(args.absy_edges) if args.absy_edges else None + src = args.gen_edges_from or args.datacard + + src_axes = None + if (ptv is None or absy is None) and src: + print(f"[gen-axes] reading the scetlib_np auxiliary of {src}") + from rabbit.inputdata import FitInputData + + from wremnants.postprocessing.scetlib_np.param_model import ( + _R_info_from_auxiliary, + ) + + indata = FitInputData(src) + src_axes = { + n: np.asarray(e, dtype=np.float64) + for n, e in _R_info_from_auxiliary(indata)["gen_axes"] + } + + def pick(name, explicit, default): + if explicit is not None: + print(f"[gen-axes] {name}: explicit ({explicit.size - 1} bins)") + return explicit + if src_axes is not None and name in src_axes: + print(f"[gen-axes] {name}: from {src} ({src_axes[name].size - 1} bins)") + return src_axes[name] + print(f"[gen-axes] {name}: built-in default ({default.size - 1} bins, " + f"[{default[0]:g}, {default[-1]:g}])") + return default + + return [ + ("ptVGen", pick("ptVGen", ptv, DEFAULT_PTV_EDGES)), + ("absYVGen", pick("absYVGen", absy, DEFAULT_ABSY_EDGES)), + ] + + +def load_theory_corr_hist(path, proc=None): + """Load the ``{generator}_hist`` (SCETlib+DYTurbo) Hist from a TheoryCorrection + ``.pkl.lz4``. + + The file maps ``corr[proc][histname]``. ``proc`` defaults to the single + physics key (``meta_data`` / ``file_meta_data`` excluded); the hist is the + lone ``*_hist`` that is not ``minnlo_ref_hist`` (the prediction, not the + MiNNLO reference or the ratio). + """ + import pickle + + import lz4.frame + + with lz4.frame.open(path) as fh: + corr = pickle.load(fh) + + meta_keys = {"meta_data", "file_meta_data"} + procs = [k for k in corr.keys() if k not in meta_keys] + if proc is None: + if len(procs) != 1: + raise SystemExit( + f"--theory-corr-proc needed: {os.path.basename(path)} has procs {procs}" + ) + proc = procs[0] + elif proc not in corr: + raise SystemExit(f"proc {proc!r} not in {list(corr.keys())}") + + entry = corr[proc] + cands = [k for k in entry if k.endswith("_hist") and k != "minnlo_ref_hist"] + if len(cands) != 1: + raise SystemExit( + f"expected one {{generator}}_hist in {proc}; found {list(entry.keys())}" + ) + histname = cands[0] + + print(f"[theory-corr] {os.path.basename(path)} :: {proc} / {histname}") + return entry[histname] + + +def theory_corr_projection(h, gen_axes, plot_axis, var="pdf0", q_window=(Q_LO, Q_HI), + tol=1e-6): + """Project a TheoryCorrection ``_hist`` onto the model's ``plot_axis`` gen bins. + + Reduces the (Q, absY, qT, charge, vars) Hist to a 1-D bin-integrated σ on the + model's ``plot_axis`` edges, restricted to the model's gen-grid extent on the + OTHER axis so it covers the same phase space the model σ_gen projection does: + + 1. select the ``vars`` entry (default ``pdf0`` = central tune); + 2. sum the Q bins whose centre falls in ``q_window`` (in-range only); + 3. sum the charge axis (in-range), if present; + 4. sum the OTHER gen axis over the model's extent [0, other_max]; + 5. rebin the projection axis onto the model's ``plot_axis`` edges (model + edges must be a sub-binning of the corr hist's: qT is fine enough that + ptVGen always aligns; absY uses SCETlib's binning so absYVGen may not). + + Returns an ndarray of length ``len(plot_axis edges) - 1`` (bin-integrated σ). + """ + names = [n for n, _ in gen_axes] + if plot_axis not in names or plot_axis not in _CORR_AXIS: + raise SystemExit(f"--plot-axis {plot_axis!r} not a model gen axis {names}") + edges_by_name = {n: np.asarray(e, dtype=np.float64) for n, e in gen_axes} + other_model = names[1] if plot_axis == names[0] else names[0] + proj_corr = _CORR_AXIS[plot_axis] + other_corr = _CORR_AXIS[other_model] + + have = [a.name for a in h.axes] + for need in ("Q", proj_corr, other_corr, "vars"): + if need not in have: + raise SystemExit( + f"theory-corr hist missing {need!r} axis; has {have}" + ) + + # 1. vars selection. + vlist = list(h.axes["vars"]) + if var not in vlist: + raise SystemExit( + f"--theory-corr-var {var!r} not in corr hist vars; have {vlist}" + ) + h = h[{"vars": vlist.index(var)}] + + # 2. Q window (sum in-range bins whose centre is inside the window). + qe = np.asarray(h.axes["Q"].edges, dtype=np.float64) + qc = 0.5 * (qe[:-1] + qe[1:]) + qsel = np.where((qc >= q_window[0] - tol) & (qc <= q_window[1] + tol))[0] + if not qsel.size: + raise SystemExit(f"no Q bins in window {q_window}; corr Q edges {qe}") + h = h[{"Q": slice(int(qsel[0]), int(qsel[-1]) + 1, sum)}] + + # 3. charge sum (in-range), if a charge axis is present. + if "charge" in [a.name for a in h.axes]: + h = h[{"charge": slice(0, h.axes["charge"].size, sum)}] + + # 4. sum the OTHER axis over the model's extent [0, other_max]. Non-coinciding + # upper edge (SCETlib's absY binning): cut at the nearest corr edge, warn. + other_max = edges_by_name[other_model][-1] + oe = np.asarray(h.axes[other_corr].edges, dtype=np.float64) + oc = 0.5 * (oe[:-1] + oe[1:]) + osel = np.where(oc <= other_max + tol)[0] + if not osel.size: + raise SystemExit( + f"theory-corr {other_corr} has no bins below the model {other_model} " + f"max {other_max}; corr edges {oe}" + ) + cut_idx = int(osel[-1]) + 1 + actual_edge = oe[cut_idx] + if abs(actual_edge - other_max) > tol: + print( + f"[theory-corr] WARNING: model {other_model} max {other_max} does not " + f"coincide with a corr {other_corr} edge; summing corr up to " + f"{actual_edge} ({abs(actual_edge - other_max):.3g} off)." + ) + h = h[{other_corr: slice(0, cut_idx, sum)}] + + # 5. rebin the projection axis onto the model's plot_axis edges. + W = _merge_matrix( + np.asarray(h.axes[proj_corr].edges, dtype=np.float64), + edges_by_name[plot_axis], + name=plot_axis, + tol=tol, + ) + return W @ np.asarray(h.values(flow=False), dtype=np.float64) + + +def _lambda_box_text(eff, gnu): + """Compact multi-line λ-tune annotation: np_model form(s) + non-zero params.""" + lines = [f"np_model = {eff.get('np_model')}"] + if gnu.get("np_model_nu") != eff.get("np_model"): + lines.append(f"np_model_nu = {gnu.get('np_model_nu')}") + for p in GNU_PARAMS: + if abs(gnu.get(p, 0.0)) > 0: + lines.append(f"{p} = {gnu[p]:.4g}") + for p in EFF_PARAMS: + if abs(eff.get(p, 0.0)) > 0: + lines.append(f"{p} = {eff[p]:.4g}") + return "λ tune:\n" + "\n".join(lines) + + +def make_projection_plot(sigma_gen, gen_axes, axis, out_path, eff, gnu, + s_corr=None, corr_label=None): + """Step histogram of the matched σ_gen(λ) projection onto one gen axis + (default ptVGen = ptZ), summing over the other. + + With a TheoryCorrection projection ``s_corr`` (bin-integrated σ on the SAME + ``axis`` edges), it is overlaid and TWO residual panels added: a ratio (param + model ÷ SCETlib+DYTurbo) and a DIFFERENTIAL difference Δ(dσ/dx) = + (model − corr)/width. The diff is in the top panel's units, so an additive + pedestal in the density reads as a horizontal line while a multiplicative bias + slopes with the spectrum — discriminating a constant offset from a fractional + one. Without ``s_corr`` the figure is a single panel. Values are plotted as + DIFFERENTIAL dσ/dx (bin-integrated σ ÷ bin width) so the variable binning, + notably the wide ptVGen overflow bin, reads correctly; the ratio is + width-independent. + """ + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + names = [n for n, _ in gen_axes] + if axis not in names: + raise SystemExit(f"--plot-axis {axis!r} not in gen axes {names}") + ai = names.index(axis) + other = 1 - ai # exactly 2 gen axes (ptVGen, absYVGen) + edges = np.asarray(gen_axes[ai][1], dtype=np.float64) + widths = np.diff(edges) + s = sigma_gen.sum(axis=other) + ds = s / widths + corr_label = corr_label or "SCETlib+DYTurbo" + show_ratio = s_corr is not None + if show_ratio: + ratio = np.divide(s, s_corr, out=np.ones_like(s), where=s_corr != 0) + diff = (s - s_corr) / widths # differential difference Δ(dσ/dx) + + if show_ratio: + fig, (ax, axr, axd) = plt.subplots( + 3, 1, sharex=True, figsize=(7, 7.2), + gridspec_kw={"height_ratios": [3, 1, 1], "hspace": 0.06}, + ) + else: + fig, ax = plt.subplots(figsize=(7, 5)) + axr = axd = None + + lab = "p$_T^Z$ (ptVGen) [GeV]" if axis == "ptVGen" else axis + ax.stairs(ds, edges, color="C3", lw=1.6, label="σ_gen(λ) (param model)") + if s_corr is not None: + ax.stairs(s_corr / widths, edges, color="C0", lw=1.6, ls=(0, (4, 2)), + label=corr_label) + ax.set_ylabel(r"d$\sigma_{\mathrm{gen}}$/d(" + axis + ") [a.u.]") + ax.margins(x=0) + ax.legend(loc="upper right", fontsize=9) + ax.text( + 0.975, 0.60, _lambda_box_text(eff, gnu), transform=ax.transAxes, + ha="right", va="top", fontsize=7.5, family="monospace", + bbox=dict(boxstyle="round", facecolor="white", edgecolor="0.7", alpha=0.9), + ) + + if show_ratio: + axr.stairs(ratio, edges, color="k", lw=1.4) + axr.axhline(1.0, color="0.5", lw=0.8, ls="--") + axr.set_ylabel("param model /\nSCETlib+DYTurbo") + axr.margins(x=0) + # Zoom around 1 to show the (often sub-%) residual, keeping 1.0 in frame; + # small window if the ratio is flat. + rlo, rhi = float(np.min(ratio)), float(np.max(ratio)) + pad = max((rhi - rlo) * 0.25, 0.003) + axr.set_ylim(min(rlo, 1.0) - pad, max(rhi, 1.0) + pad) + # Differential difference: horizontal ⇒ additive pedestal in density; + # tracks the spectrum shape ⇒ multiplicative (fractional) offset. + axd.stairs(diff, edges, color="C2", lw=1.4) + axd.axhline(0.0, color="0.5", lw=0.8, ls="--") + axd.set_xlabel(lab) + axd.set_ylabel("(model − corr)\n/ d(" + axis + ")") + axd.margins(x=0) + rng = (f"; model/corr [{rlo:.4f}, {rhi:.4f}]" + f"; Δ(dσ/dx) [{float(np.min(diff)):.3g}, {float(np.max(diff)):.3g}]") + else: + ax.set_xlabel(lab) + rng = "" + + out_dir = os.path.dirname(out_path) + if out_dir and not os.path.exists(out_dir): + os.makedirs(out_dir, exist_ok=True) + fig.savefig(out_path, dpi=130, bbox_inches="tight") + plt.close(fig) + print(f"[plot] wrote {out_path} (axis={axis}, summed over {names[other]}{rng})") + + +def main(argv=None): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--btgrid", default=BTGRID_DIR, help="SCETlib bT-grid directory") + p.add_argument("--datacard", default=None, + help="hdf5 fallback for the GEN EDGES (no default). λ are NOT " + "sourced from it — use --meta-from for that") + # λ tune: base source (optional) + functional form + p.add_argument("--meta-from", default=None, + help="hdf5 to read the base λ tune from (datacard/fitresults metadata); optional") + p.add_argument("--np-model", default=None, + help="F_eff functional-form override (default: the base tune's form — " + "tanh_2 for the canonical / --theory-corr base)") + p.add_argument("--np-model-nu", default=None, + help="γ_ν^NP functional-form override (default: the base tune's form)") + # λ values to evaluate at (applied on top of the base, in this order) + p.add_argument("--fitresult", default=None, + help="fitresults hdf5 to read the POSTFIT λ from (optional)") + p.add_argument("--result", default=None, help="fitresult group suffix (e.g. 'nominal')") + p.add_argument("--lambdas", default=None, + help="λ values 'name=val,...' evaluated on top of the base tune " + "(e.g. lambda2=0.5); unset params stay at the base (FranksVals " + "by default), NOT 0") + # gen-edge source + p.add_argument("--ptv-edges", default=None, + help="ptVGen edges 'a,b,c,...' (default: 1-GeV bins over [0,40])") + p.add_argument("--absy-edges", default=None, + help="absYVGen edges 'a,b,c,...' (default: single bin [0,5])") + p.add_argument("--gen-edges-from", default=None, + help="hdf5 whose scetlib_np auxiliary gives the gen edges " + "(default: --datacard, else the built-in defaults)") + # TheoryCorrection overlay + p.add_argument("--theory-corr", default=None, + help="TheoryCorrection .pkl.lz4 to overlay the official " + "SCETlib+DYTurbo gen distribution (its {generator}_hist)") + p.add_argument("--theory-corr-proc", default=None, + help="proc key in the corr file (default: the single physics key)") + p.add_argument("--theory-corr-var", default="pdf0", + help="vars label to read from the corr hist (default pdf0 = central)") + p.add_argument("--theory-corr-normalize", action="store_true", + help="rescale the corr curve to the model σ_gen integral " + "(shape-only comparison; default off = absolute overlay)") + # model / output + p.add_argument("--q-lo", type=float, default=Q_LO) + p.add_argument("--q-hi", type=float, default=Q_HI) + p.add_argument("--no-nonsingular", action="store_true", + help="resum-only σ_gen (σ_ns = 0; skips the FO inputs)") + p.add_argument("--plot", default=None, + help="optional path (e.g. .png/.pdf) to write a 1-D projection plot " + "of σ_gen(λ) [+ theory-corr overlay/ratio]; see --plot-axis") + p.add_argument("--plot-axis", default="ptVGen", choices=["ptVGen", "absYVGen"], + help="gen axis to project onto for --plot (default ptVGen = ptZ)") + args = p.parse_args(argv) + + # ---- λ: build a PHYSICAL base tune (model CONSTRUCTED there so the + # positive-σ_gen guard passes), then EVALUATE at base + the requested + # overrides (--fitresult postfit, then --lambdas). Params not set stay at the + # base, NOT at 0. + import copy + + overrides = {} + if args.fitresult: + from wremnants.postprocessing.scetlib_np.fitresult_lambdas import _flat_values + + pf = _flat_values(args.fitresult, which="postfit", result=args.result) + overrides.update(pf) + print(f"[λ] postfit from {args.fitresult}: {pf}") + try: + overrides.update(parse_lambda_overrides(args.lambdas)) + except ValueError as e: + p.error(str(e)) + + base = copy.deepcopy(resolve_base_lambda(args)) + if args.np_model: + base["eff_params"]["np_model"] = args.np_model + if args.np_model_nu: + base["gnu_params"]["np_model_nu"] = args.np_model_nu + eff, gnu, explicit = assemble_tune(base, overrides) + + gen_axes = resolve_gen_axes(args) + + from wremnants.postprocessing.scetlib_np.sigma_gen import SigmaGenModel + + print("\n[core] constructing SigmaGenModel at the base tune (bt-grid integral) …", + flush=True) + t0 = time.time() + core = SigmaGenModel( + btgrid_dir=args.btgrid, + lambda_central=base, + gen_axes=gen_axes, + Q_lo=args.q_lo, + Q_hi=args.q_hi, + include_nonsingular=not args.no_nonsingular, + ) + print(f" constructed in {time.time()-t0:.1f}s; gen grid {core.gen_shape} " + f"({[n for n, _ in core.gen_axes]})") + + print(f"\n[λ] evaluating matched σ_gen at:") + print(f" F_eff : {eff}") + print(f" γ_ν^NP : {gnu}") + if explicit: + print(f" (set via --lambdas/--fitresult: {explicit}; the rest stay at the base)") + else: + print(" (no overrides — evaluating at the base tune itself)") + + # ---- σ_gen on the (ptVGen, absYVGen) gen grid. Reuse the construction central + # when no overrides; else evaluate the tune (no positivity guard, so a weak-NP + # tune can be inspected — it can dip negative at the lowest qT). + t0 = time.time() + if explicit: + sigma_gen = np.asarray(core.sigma_gen(eff, gnu).numpy(), dtype=np.float64) + else: + sigma_gen = np.asarray(core.sigma_gen_central.numpy(), dtype=np.float64) + print(f" σ_gen computed in {time.time()-t0:.1f}s; shape {sigma_gen.shape}") + + n_bad = int(np.sum(sigma_gen <= 0)) + if n_bad: + print(f" [warning] {n_bad}/{sigma_gen.size} σ_gen bins are non-positive at " + f"this tune (expected where the NP damping is weak, esp. low qT).") + + print(f"\n Σ σ_gen = {sigma_gen.sum():.6g}") + print(f" per-bin σ_gen : min {sigma_gen.min():.4g} max {sigma_gen.max():.4g}") + print("\n σ_gen(λ) per (ptVGen × absY) bin:") + with np.printoptions(precision=4, suppress=True, linewidth=140): + print(sigma_gen) + + # ---- NP physical-validity detectors at THIS λ. A wrong-sign tune's pathology + # (anti-damping NP → oscillating, negative native σ(qT)) is AVERAGED AWAY in + # the binned σ_gen above, so check the native spectrum and form factors + # directly. Detectors only — change nothing; fit-time enforcement is the + # np_damping_wall.NPDampingWall regularizer. + from wremnants.postprocessing.scetlib_np import param_model_diagnostics as ppd + + rep = ppd.np_physical_report(core, eff, gnu) + damp, neg = rep["damp"], rep["neg"] + print("\n NP physical-validity:") + print(f" γ_ν^NP damping : {'OK' if not damp['gamma_nu_wrong_sign'] else 'WRONG SIGN'}" + f" (max γ_ν over probe bT = {damp['gamma_nu_max']:+.3g}; must be ≤ 0)") + print(f" F_eff decays : {'OK' if not damp['F_eff_growing'] else 'GROWING (bT-integral divergence sign)'}") + print(f" native σ(qT)≥0 : neg_area_frac={neg['neg_area_frac']:.3g} " + f"(λ_central {rep['central_neg_area']:.3g}), min/peak={neg['min_over_peak']:+.3g}, " + f"n_neg_bins={neg['n_neg_bins']}") + + # WHERE the negativity sits (native Y × qT). The σ(qT)<0 dip is laundered by + # the gen-binning, so locating it is the discriminator: negative cells inside + # the |Y|≤2.5 acceptance at accessible qT matter for interpretation; cells only + # at |Y|>2.5 or beyond the fit's qT reach are doubly invisible. ACCEPT_ABSY is + # the Z dilepton |yll| acceptance edge (boson-Y proxy). + ACCEPT_ABSY = 2.5 + if neg.get("n_neg_bins") and neg.get("worst") is not None: + w = neg["worst"] + print(f" worst neg cell : σ={w['value']:.4g} ({w['frac_of_peak']*100:+.1f}% of peak) " + f"at Y={w['Y']:+.3g}, qT={w['qT']:.3g} GeV") + cells = neg.get("neg_bins") or [] + in_cells = [c for c in cells if abs(c["Y"]) <= ACCEPT_ABSY] + n_in, n_out = len(in_cells), len(cells) - len(in_cells) + qr = neg.get("neg_qT_range") + print(f" neg-cell split : {n_in} inside |Y|≤{ACCEPT_ABSY}, {n_out} outside " + f"(|Y|≤{neg.get('neg_absY_max', float('nan')):.2g} reached); " + f"qT∈[{qr[0]:.3g}, {qr[1]:.3g}] GeV" if qr else "") + if in_cells: + wi = min(in_cells, key=lambda c: c["value"]) + print(f" worst in-acc : σ={wi['value']:.4g} ({wi['frac_of_peak']*100:+.1f}% of peak) " + f"at Y={wi['Y']:+.3g}, qT={wi['qT']:.3g} GeV " + f"[the |Y|≤{ACCEPT_ABSY} pathology the fit region could see]") + n_show = min(10, len(cells)) + if n_show: + print(f" most-negative {n_show} cells (Y, qT[GeV], σ, %peak):") + for c in cells[:n_show]: + inacc = "in " if abs(c["Y"]) <= ACCEPT_ABSY else "out" + print(f" [{inacc}] Y={c['Y']:+.3g} qT={c['qT']:7.3g} " + f"σ={c['value']:+.4g} ({c['frac_of_peak']*100:+.1f}%)") + + if not rep["ok"]: + print(" ⚠ UNPHYSICAL NP TUNE — the differential σ(qT) is negative / the NP is " + "anti-damping.\n This is hidden in the binned σ_gen above; do not treat " + "this point as a physical prediction.") + + # ---- optional: project an official TheoryCorrection run onto the same axis. + s_corr = None + corr_label = None + if args.theory_corr: + h_corr = load_theory_corr_hist(args.theory_corr, args.theory_corr_proc) + s_corr = theory_corr_projection( + h_corr, core.gen_axes, args.plot_axis, var=args.theory_corr_var, + q_window=(args.q_lo, args.q_hi), + ) + other = 1 - [n for n, _ in core.gen_axes].index(args.plot_axis) + s_model = sigma_gen.sum(axis=other) + sum_corr, sum_model = float(s_corr.sum()), float(s_model.sum()) + norm_label = "" + if args.theory_corr_normalize and sum_corr != 0: + scale = sum_model / sum_corr + s_corr = s_corr * scale + norm_label = f", ×{scale:.4f} (shape-normalized)" + print(f"\n[theory-corr] shape-normalized to the model integral (×{scale:.5f})") + corr_label = f"SCETlib+DYTurbo ({args.theory_corr_var}){norm_label}" + rcorr = np.divide(s_model, s_corr, out=np.ones_like(s_model), where=s_corr != 0) + print(f"\n[theory-corr] projection onto {args.plot_axis} (var={args.theory_corr_var}):") + print(f" Σ corr = {sum_corr:.6g}") + print(f" Σ model / Σ corr = {sum_model/sum_corr:.5f}" + + (" (== 1 after --theory-corr-normalize)" if args.theory_corr_normalize else "")) + print(f" model / corr per bin : min {rcorr.min():.4f} max {rcorr.max():.4f}") + with np.printoptions(precision=4, suppress=True, linewidth=140): + print(f" model/corr: {rcorr}") + if not args.plot: + print("[theory-corr] (pass --plot to also write the overlay figure)") + + if args.plot: + make_projection_plot( + sigma_gen, core.gen_axes, args.plot_axis, args.plot, eff, gnu, + s_corr=s_corr, corr_label=corr_label, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/wremnants/utilities/styles/styles.py b/wremnants/utilities/styles/styles.py index 3052c0170..0f3f5736e 100644 --- a/wremnants/utilities/styles/styles.py +++ b/wremnants/utilities/styles/styles.py @@ -1,6 +1,6 @@ import copy -import matplotlib.cm as cm +import matplotlib from wums import boostHistHelpers as hh from wums import logging @@ -819,7 +819,7 @@ def get_labels_colors_procs_sorted(procs): "Rare", ][::-1] - cmap = cm.get_cmap("tab10") + cmap = matplotlib.colormaps["tab10"] procs = sorted( procs, key=lambda x: procs_sort.index(x) if x in procs_sort else len(procs_sort) From 30e5d3a6c500a71b35e5c415a5ed975e9991afa5 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 1 Jul 2026 15:11:30 -0400 Subject: [PATCH 27/31] =?UTF-8?q?SCETlib-NP:=20model=E2=86=92=CE=BB=20regi?= =?UTF-8?q?stry=20in=20params.py=20(single=20source)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One nested registry {model: {param: {value, sigma}}}, split by family, storing each model's λ set AND its fit defaults inline (value = neutral start fallback, sigma = default prior width, None = free). Replaces the hardcoded active_params branches, the flat DEFAULT_PRIOR_SIGMAS, and btgrid_tf's EFF_MODELS/GNU_MODELS name-sets. active_params()/param_defaults()/valid-name sets derive from it. Co-Authored-By: Claude Opus 4.8 --- wremnants/postprocessing/scetlib_np/params.py | 160 ++++++++++++++---- 1 file changed, 123 insertions(+), 37 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/params.py b/wremnants/postprocessing/scetlib_np/params.py index 438eb1a7e..dc9cf913a 100644 --- a/wremnants/postprocessing/scetlib_np/params.py +++ b/wremnants/postprocessing/scetlib_np/params.py @@ -17,19 +17,6 @@ EFF_PARAMS = ("lambda2", "lambda4", "lambda6", "delta_lambda2", "lambda_inf") ALL_PARAMS = GNU_PARAMS + EFF_PARAMS -# Recommended Gaussian prior widths per λ (theorist recommendations), applied by -# SCETlibNPParamModel only when priors are enabled (priors=1); a λ absent here -# floats free (NaN width). Lives in this config home keyed by the names above. -DEFAULT_PRIOR_SIGMAS = { - "lambda2_nu": 0.10, - "lambda4_nu": 0.50, - "lambda6_nu": 0.10, - "lambda2": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average - "lambda4": 0.50, # 0.4 ⁺⁰·⁶₋₀.₄ -> symmetric average - "delta_lambda2": 0.20, # 0 ± 0.20 wide default (no theorist value yet) - "lambda6": 0.1, -} - # np_model selector keys carried alongside the numeric λ in a tune dict. EFF_MODEL_KEY = "np_model" GNU_MODEL_KEY = "np_model_nu" @@ -38,40 +25,139 @@ RECO_AXES = ("ptll", "yll", "cosThetaStarll_quantile", "phiStarll_quantile") GEN_AXES = ("ptVGen", "absYVGen") -# Which λ each NP model actually uses. The form factors in btgrid_tf -# (F_eff_tf / gamma_nu_NP_tf) read a fixed subset of the λ per model string; a λ -# outside that subset is inert (e.g. lambda6 under tanh_2). Mirrored here as -# plain data so this module stays TF-free — btgrid_tf is the source of truth. +# np_model selector aliases -> canonical model name. The form-factor branches in +# btgrid_tf and the registry below key on the canonical name. _EFF_MODEL_ALIASES = {"hyp_tangent": "tanh_2", "square_root": "frac_2"} _GNU_MODEL_ALIASES = {"hyp_tangent": "tanh_2", "linear": "frac_1"} -# eff models with no lambda_inf damping (plain polynomial form factors). -_EFF_NO_LAMBDA_INF = {"signed_lambda", "identity"} +# ---- Model → λ registry (single source of truth) ----------------------------- +# Per model: the λ it uses and their fit defaults (value = neutral start fallback, +# sigma = default prior width, None = free). Must stay in sync with the btgrid_tf +# form branches, which read exactly these λ by name. +EFF_MODEL_PARAMS = { + "identity": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + }, + "signed_lambda": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + }, + "tanh_2": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + "lambda_inf": {"value": 0.0, "sigma": None}, + }, + "tanh_4": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + "lambda_inf": {"value": 0.0, "sigma": None}, + }, + "tanh_6": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + "lambda_inf": {"value": 0.0, "sigma": None}, + "lambda6": {"value": 0.0, "sigma": 0.10}, + }, + "frac_2": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + "lambda_inf": {"value": 0.0, "sigma": None}, + }, + "frac_4": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + "lambda_inf": {"value": 0.0, "sigma": None}, + }, + "exp_2": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + "lambda_inf": {"value": 0.0, "sigma": None}, + }, + "exp_4": { + "lambda2": {"value": 0.0, "sigma": 0.50}, + "lambda4": {"value": 0.0, "sigma": 0.50}, + "delta_lambda2": {"value": 0.0, "sigma": 0.20}, + "lambda_inf": {"value": 0.0, "sigma": None}, + }, +} +GNU_MODEL_PARAMS = { + "tanh_1": { + "lambda2_nu": {"value": 0.0, "sigma": 0.10}, + "lambda4_nu": {"value": 0.0, "sigma": 0.50}, + "lambda_inf_nu": {"value": 0.0, "sigma": None}, + }, + "tanh_2": { + "lambda2_nu": {"value": 0.0, "sigma": 0.10}, + "lambda4_nu": {"value": 0.0, "sigma": 0.50}, + "lambda_inf_nu": {"value": 0.0, "sigma": None}, + }, + "tanh_6": { + "lambda2_nu": {"value": 0.0, "sigma": 0.10}, + "lambda4_nu": {"value": 0.0, "sigma": 0.50}, + "lambda_inf_nu": {"value": 0.0, "sigma": None}, + "lambda6_nu": {"value": 0.0, "sigma": 0.10}, + }, + "frac_1": { + "lambda2_nu": {"value": 0.0, "sigma": 0.10}, + "lambda4_nu": {"value": 0.0, "sigma": 0.50}, + "lambda_inf_nu": {"value": 0.0, "sigma": None}, + }, + "frac_2": { + "lambda2_nu": {"value": 0.0, "sigma": 0.10}, + "lambda4_nu": {"value": 0.0, "sigma": 0.50}, + "lambda_inf_nu": {"value": 0.0, "sigma": None}, + }, + "exp_1": { + "lambda2_nu": {"value": 0.0, "sigma": 0.10}, + "lambda4_nu": {"value": 0.0, "sigma": 0.50}, + "lambda_inf_nu": {"value": 0.0, "sigma": None}, + }, + "exp_2": { + "lambda2_nu": {"value": 0.0, "sigma": 0.10}, + "lambda4_nu": {"value": 0.0, "sigma": 0.50}, + "lambda_inf_nu": {"value": 0.0, "sigma": None}, + }, +} + +# Valid model names (canonical + aliases). The single validation set for both +# btgrid_tf (``np_model not in EFF_MODELS``) and the param model. Re-exported by +# btgrid_tf so it need not keep its own copy. +EFF_MODELS = frozenset(EFF_MODEL_PARAMS) | frozenset(_EFF_MODEL_ALIASES) +GNU_MODELS = frozenset(GNU_MODEL_PARAMS) | frozenset(_GNU_MODEL_ALIASES) -def active_params(np_model=None, np_model_nu=None): - """Names of the λ the chosen NP model(s) actually use. - Pass ``np_model`` for the F_eff (TMD) set, ``np_model_nu`` for the γ_ν (CS) - set, or both for the union. A λ outside the returned set is inert for that - model (``lambda6``/``lambda6_nu`` are used only by ``tanh_6``); callers can - reject such λ instead of silently ignoring them. Source of truth: - ``btgrid_tf.F_eff_tf`` / ``btgrid_tf.gamma_nu_NP_tf``.""" - out = set() +def param_defaults(np_model=None, np_model_nu=None): + """``{name: {"value", "sigma"}}`` for the λ the given model(s) use. + + Union of the F_eff (``np_model``) and γ_ν (``np_model_nu``) registry rows, + aliases resolved. Raises ``KeyError`` on an unknown model name.""" + out = {} if np_model is not None: - m = _EFF_MODEL_ALIASES.get(np_model, np_model) - out |= {"lambda2", "lambda4", "delta_lambda2"} - if m not in _EFF_NO_LAMBDA_INF: - out.add("lambda_inf") - if m == "tanh_6": - out.add("lambda6") + out.update(EFF_MODEL_PARAMS[_EFF_MODEL_ALIASES.get(np_model, np_model)]) if np_model_nu is not None: - m = _GNU_MODEL_ALIASES.get(np_model_nu, np_model_nu) - out |= {"lambda2_nu", "lambda4_nu", "lambda_inf_nu"} - if m == "tanh_6": - out.add("lambda6_nu") + out.update(GNU_MODEL_PARAMS[_GNU_MODEL_ALIASES.get(np_model_nu, np_model_nu)]) return out +def active_params(np_model=None, np_model_nu=None): + """Names of the λ the chosen NP model(s) actually use (registry keys). + + Pass ``np_model`` for the F_eff (TMD) set, ``np_model_nu`` for the γ_ν (CS) + set, or both for the union. A λ outside the returned set is inert for that + model and is NOT a fit parameter. Source of truth: :data:`EFF_MODEL_PARAMS` / + :data:`GNU_MODEL_PARAMS`, audited against the ``btgrid_tf`` form branches.""" + return set(param_defaults(np_model=np_model, np_model_nu=np_model_nu)) + + def parse_lambda_overrides(spec): """Parse a ``"name=val,name=val"`` λ-override string into ``{name: float}``. From e71ab13a2190837ce4711d6e6a5d150d3fdf0c3d Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 1 Jul 2026 15:11:30 -0400 Subject: [PATCH 28/31] =?UTF-8?q?SCETlib-NP:=20de-hardcode=20F=5Feff=5Ftf?= =?UTF-8?q?=20/=20gamma=5Fnu=5FNP=5Ftf=20(read=20=CE=BB=20by=20name)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The form factors take a values dict + np_model selector; each branch reads ONLY the λ its formula uses (missing -> KeyError = fail out). Adding a model with new λ is a new branch + a registry row, no signature churn. EFF_MODELS/GNU_MODELS + alias maps imported from params. Updated the 7 internal call sites and the np_function_plots / param_model_diagnostics callers (**dict -> dict). Co-Authored-By: Claude Opus 4.8 --- .../postprocessing/scetlib_np/btgrid_tf.py | 99 +++++++++---------- .../scetlib_np/np_function_plots.py | 13 ++- .../scetlib_np/param_model_diagnostics.py | 4 +- 3 files changed, 56 insertions(+), 60 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_tf.py b/wremnants/postprocessing/scetlib_np/btgrid_tf.py index cb7aac502..397c56f3c 100644 --- a/wremnants/postprocessing/scetlib_np/btgrid_tf.py +++ b/wremnants/postprocessing/scetlib_np/btgrid_tf.py @@ -19,6 +19,16 @@ import numpy as np import tensorflow as tf +# Valid-name sets + alias maps from the numpy-only params module (single source). +# The form branches below must read exactly the λ each model lists in the params +# registry (EFF_MODEL_PARAMS / GNU_MODEL_PARAMS). +from wremnants.postprocessing.scetlib_np.params import ( + EFF_MODELS, + GNU_MODELS, + _EFF_MODEL_ALIASES, + _GNU_MODEL_ALIASES, +) + # float64 throughout this module. DTYPE = tf.float64 @@ -85,30 +95,7 @@ def simpson_tf(y, weights): # F_eff and gamma_nu^NP — TF transcriptions # ============================================================================= -EFF_MODELS = { - "identity", - "tanh_2", - "tanh_6", - "tanh_4", - "frac_2", - "frac_4", - "exp_2", - "exp_4", - "signed_lambda", - "hyp_tangent", - "square_root", -} -GNU_MODELS = { - "tanh_1", - "tanh_2", - "tanh_6", - "frac_1", - "frac_2", - "exp_1", - "exp_2", - "hyp_tangent", - "linear", -} +# EFF_MODELS / GNU_MODELS (valid np_model names) come from params (imported above). def _frozen_eq_zero(x): @@ -138,24 +125,27 @@ def _safe_div(num, den): return num / den_safe -def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_model): - """TMD-effective NP form factor F_eff(Y, bT) for a fixed ``np_model``.""" +def F_eff_tf(Y, bT, values, *, np_model): + """TMD-effective NP form factor F_eff(Y, bT) for a fixed ``np_model``. + + ``values`` maps λ name -> value (TF scalar / Variable / constant, or python + float). Each ``np_model`` branch reads ONLY the λ its formula uses — a missing + one raises ``KeyError`` (fail out; no fabricated default). Extra keys (e.g. + ``np_model``) are ignored. The λ each model reads is declared in + :data:`params.EFF_MODEL_PARAMS`; keep the two in sync.""" if np_model not in EFF_MODELS: raise ValueError(f"F_eff_tf: unsupported np_model {np_model!r}") bT = _as_dtype(bT) Y = _as_dtype(Y) - lambda_inf = _as_dtype(lambda_inf) - lambda2 = _as_dtype(lambda2) - lambda4 = _as_dtype(lambda4) - lambda6 = _as_dtype(lambda6) - delta_lambda2 = _as_dtype(delta_lambda2) + lambda2 = _as_dtype(values["lambda2"]) + lambda4 = _as_dtype(values["lambda4"]) + delta_lambda2 = _as_dtype(values["delta_lambda2"]) + lambda2_Y = lambda2 + delta_lambda2 * Y * Y if np_model == "signed_lambda": - lambda2_Y = lambda2 + delta_lambda2 * Y * Y return (1.0 + lambda2_Y * bT**2) ** 2 * tf.exp(-2.0 * lambda4 * bT**4) - lambda2_Y = lambda2 + delta_lambda2 * Y * Y arg = (lambda2_Y + lambda4 * bT**2) * bT if np_model == "identity": @@ -163,13 +153,15 @@ def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_ # lambda_inf == 0 returns ones: compute the full formula with a safe # denominator, mask at the end. + lambda_inf = _as_dtype(values["lambda_inf"]) arg_inf = _safe_div(arg, lambda_inf) - model = {"hyp_tangent": "tanh_2", "square_root": "frac_2"}.get(np_model, np_model) + model = _EFF_MODEL_ALIASES.get(np_model, np_model) if model == "tanh_2": a = arg_inf + (1.0 / 3.0) * _safe_div(lambda2_Y * bT, lambda_inf) ** 3 func = tf.tanh(a) elif model == "tanh_6": + lambda6 = _as_dtype(values["lambda6"]) a = arg_inf + _safe_div(lambda6 * bT**5, lambda_inf) a = a + (1.0 / 3.0) * _safe_div(lambda2_Y * bT, lambda_inf) ** 3 func = tf.tanh(a) @@ -194,29 +186,26 @@ def F_eff_tf(Y, bT, *, lambda_inf, lambda2, lambda4, lambda6, delta_lambda2, np_ return tf.where(_frozen_eq_zero(lambda_inf), tf.ones_like(full), full) -def gamma_nu_NP_tf( - bT, *, lambda_inf_nu, lambda2_nu, lambda4_nu, lambda6_nu=0.0, np_model_nu -): +def gamma_nu_NP_tf(bT, values, *, np_model_nu): """CS-side NP rapidity anomalous dimension γ_ν^NP(bT) for fixed ``np_model_nu``. - ``lambda6_nu`` is the b⁶ coefficient used only by the ``tanh_6`` model; other - models ignore it. It defaults to 0 (then tanh_6 reduces to tanh_2). SCETlib's - own ``NP_model_gammanu`` uses 0.0007 (Gamma_nu.hpp:102) — pass that to - reproduce SCETlib's regulated CS kernel. + ``values`` maps λ name -> value; each branch reads ONLY the λ its formula uses + (a missing one raises ``KeyError`` — fail out). Extra keys (e.g. + ``np_model_nu``) are ignored. The λ each model reads is declared in + :data:`params.GNU_MODEL_PARAMS`; keep the two in sync. """ if np_model_nu not in GNU_MODELS: raise ValueError(f"gamma_nu_NP_tf: unsupported np_model_nu {np_model_nu!r}") bT = _as_dtype(bT) - lambda_inf_nu = _as_dtype(lambda_inf_nu) - lambda2_nu = _as_dtype(lambda2_nu) - lambda4_nu = _as_dtype(lambda4_nu) - lambda6_nu = _as_dtype(lambda6_nu) + lambda_inf_nu = _as_dtype(values["lambda_inf_nu"]) + lambda2_nu = _as_dtype(values["lambda2_nu"]) + lambda4_nu = _as_dtype(values["lambda4_nu"]) bT2 = bT * bT arg = _safe_div((lambda2_nu + lambda4_nu * bT2) * bT2, lambda_inf_nu) - model = {"hyp_tangent": "tanh_2", "linear": "frac_1"}.get(np_model_nu, np_model_nu) + model = _GNU_MODEL_ALIASES.get(np_model_nu, np_model_nu) if model == "tanh_1": a = arg + (2.0 / 3.0) * _safe_div(lambda2_nu * bT2, lambda_inf_nu) ** 2 @@ -224,8 +213,8 @@ def gamma_nu_NP_tf( elif model == "tanh_2": func = tf.tanh(arg) elif model == "tanh_6": - # b⁶ term (SCETlib NP_model_gammanu uses lambda6_nu = 0.0007, - # Gamma_nu.hpp:102); here it is the fittable lambda6_nu (default 0). + # tanh_2 plus a b⁶ term with fittable coefficient lambda6_nu. + lambda6_nu = _as_dtype(values["lambda6_nu"]) a = arg + _safe_div(lambda6_nu * bT2**3, lambda_inf_nu) func = tf.tanh(a) elif model == "frac_1": @@ -306,13 +295,13 @@ def reconstruct_batch_tf( bT_simpson_weights = simpson_weights(np.asarray(bT)) bT_simpson_weights = _as_dtype(bT_simpson_weights) # (Nbt,) - g_NP = gamma_nu_NP_tf(b_bar, **gnu_params, np_model_nu=np_model_nu) # (Nbt,) + g_NP = gamma_nu_NP_tf(b_bar, gnu_params, np_model_nu=np_model_nu) # (Nbt,) exp_g_factor = tf.exp(C_nu * g_NP[tf.newaxis, :]) # (Nbins, Nbt) delta_l2_in = eff_params.get("delta_lambda2", 0.0) if isinstance(delta_l2_in, (int, float)) and float(delta_l2_in) == 0.0: # static fast path: F_eff has no Y dependence - Feff = F_eff_tf(0.0, b_bar, **eff_params, np_model=np_model) # (Nbt,) + Feff = F_eff_tf(0.0, b_bar, eff_params, np_model=np_model) # (Nbt,) Feff_b = Feff[tf.newaxis, :] elif Y_unique is not None and Y_inverse_idx is not None: # per-bin F_eff on the unique-Y rows, then gathered. Exact: identical Y @@ -321,14 +310,14 @@ def reconstruct_batch_tf( # run on (NY, Nbt), not (Nbins, Nbt). Y_u = _as_dtype(Y_unique)[:, tf.newaxis] # (NY, 1) b_b = b_bar[tf.newaxis, :] - Feff_u = F_eff_tf(Y_u, b_b, **eff_params, np_model=np_model) # (NY, Nbt) + Feff_u = F_eff_tf(Y_u, b_b, eff_params, np_model=np_model) # (NY, Nbt) Feff_b = tf.gather(Feff_u, Y_inverse_idx) # (Nbins, Nbt) else: # per-bin F_eff (Y dependence via delta_lambda2 * Y^2): build # (Nbins, Nbt) by broadcasting Y_per_bin over bT Y_b = Y_per_bin[:, tf.newaxis] b_b = b_bar[tf.newaxis, :] - Feff_b = F_eff_tf(Y_b, b_b, **eff_params, np_model=np_model) # (Nbins, Nbt) + Feff_b = F_eff_tf(Y_b, b_b, eff_params, np_model=np_model) # (Nbins, Nbt) integrand = bT_J0_kernel * I_pert * exp_g_factor * Feff_b # (Nbins, Nbt) sigma = simpson_tf(integrand, bT_simpson_weights) # (Nbins,) @@ -546,7 +535,7 @@ def reconstruct_batch_factorized_tf( I_pert_u = _as_dtype(I_pert_u) KwqT = _as_dtype(KwqT) - g_NP = gamma_nu_NP_tf(b_bar, **gnu_params, np_model_nu=np_model_nu) # (Nbt,) + g_NP = gamma_nu_NP_tf(b_bar, gnu_params, np_model_nu=np_model_nu) # (Nbt,) if C_nu_uu is not None and c_of_u is not None: # exp on the ~150x smaller C-row table, replicated by the gather. exp_g_uu = tf.exp(_as_dtype(C_nu_uu) * g_NP[tf.newaxis, :]) # (Ncu, Nbt) @@ -563,7 +552,7 @@ def reconstruct_batch_factorized_tf( if isinstance(delta_l2_in, (int, float)) and float(delta_l2_in) == 0.0: # static fast path: F_eff has no Y dependence (matches the # reconstruct_batch_tf fast path bit-for-bit on the unique rows) - Feff = F_eff_tf(0.0, b_bar, **eff_params, np_model=np_model) # (Nbt,) + Feff = F_eff_tf(0.0, b_bar, eff_params, np_model=np_model) # (Nbt,) Feff_u = Feff[tf.newaxis, :] # broadcast over Nu else: # F_eff on the unique-Y rows, gathered to the unique grid rows (same @@ -571,7 +560,7 @@ def reconstruct_batch_factorized_tf( # bit-exactly, scatter-adds cotangents on the backward pass). Y_u = _as_dtype(Y_unique)[:, tf.newaxis] # (NY, 1) Feff_rows = F_eff_tf( - Y_u, b_bar[tf.newaxis, :], **eff_params, np_model=np_model + Y_u, b_bar[tf.newaxis, :], eff_params, np_model=np_model ) # (NY, Nbt) Feff_u = tf.gather(Feff_rows, feff_idx_u) # (Nu, Nbt) diff --git a/wremnants/postprocessing/scetlib_np/np_function_plots.py b/wremnants/postprocessing/scetlib_np/np_function_plots.py index 5348b0453..54db262e6 100644 --- a/wremnants/postprocessing/scetlib_np/np_function_plots.py +++ b/wremnants/postprocessing/scetlib_np/np_function_plots.py @@ -88,13 +88,20 @@ class Series: def gamma_nu_curve(bT, gnu): - """γ_ν^NP(b_T) for one gnu dict (CS sector).""" - return np.asarray(btgrid_tf.gamma_nu_NP_tf(bT, **gnu), dtype=float) + """γ_ν^NP(b_T) for one gnu dict (CS sector). + + ``gnu`` carries the numeric λ plus the ``np_model_nu`` key; the form reads the + λ it needs and ignores the model key (passed explicitly as the selector).""" + return np.asarray( + btgrid_tf.gamma_nu_NP_tf(bT, gnu, np_model_nu=gnu["np_model_nu"]), dtype=float + ) def f_eff_curve(bT, y, eff): """F_eff(y, b_T) for one eff dict at rapidity ``y`` (TMD sector).""" - return np.asarray(btgrid_tf.F_eff_tf(y, bT, **eff), dtype=float) + return np.asarray( + btgrid_tf.F_eff_tf(y, bT, eff, np_model=eff["np_model"]), dtype=float + ) def _band(curves, pct): diff --git a/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py b/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py index acb8623c1..f248aabe9 100644 --- a/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py +++ b/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py @@ -80,8 +80,8 @@ def np_damping_ok(core, eff_params, gnu_params, b_probe=NP_PROBE_BT, gamma_tol=1 b = np.asarray(b_probe, dtype=np.float64) eff = {k: v for k, v in eff_params.items() if k != "np_model"} gnu = {k: v for k, v in gnu_params.items() if k != "np_model_nu"} - g = fz_tf.gamma_nu_NP_tf(b, np_model_nu=core.np_model_nu, **gnu).numpy() - F = fz_tf.F_eff_tf(0.0, b, np_model=core.np_model, **eff).numpy() + g = fz_tf.gamma_nu_NP_tf(b, gnu, np_model_nu=core.np_model_nu).numpy() + F = fz_tf.F_eff_tf(0.0, b, eff, np_model=core.np_model).numpy() gamma_max = float(np.max(g)) feff_growing = bool(F[-1] > F[0]) return { From 31556b982417f9769c079c3b263df6fc8a1ae342 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 1 Jul 2026 15:11:30 -0400 Subject: [PATCH 29/31] =?UTF-8?q?SCETlib-NP:=20fit=20only=20the=20active?= =?UTF-8?q?=20=CE=BB;=20central=20path=20validate-only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit param_model: _param_order = active_params(fit forms) — inert λ (e.g. lambda6/ lambda6_nu under tanh_2) are no longer registered, so they can't add a zero-derivative Hessian row (fixes the postfit 'Cholesky failed, Hessian not positive-definite'). Start precedence xparam_default > card λ_central > registry neutral; prior sigmas from the registry; impact groups intersected; poi_params validated; prior_sigmas/xparam_default naming a non-fitted λ warn-and-ignore. lambda_central: _fill_missing_params validate-only (no 0.0 fill). Co-Authored-By: Claude Opus 4.8 --- .../scetlib_np/lambda_central.py | 23 +++--- .../postprocessing/scetlib_np/param_model.py | 79 +++++++++++++++---- 2 files changed, 71 insertions(+), 31 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py index 39a04a1d4..3f4fa4341 100644 --- a/wremnants/postprocessing/scetlib_np/lambda_central.py +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -185,17 +185,17 @@ def _iter_meta_levels(meta, max_depth=8): def _fill_missing_params(lc): - """Complete ``lc``'s eff/gnu sub-dicts for the model's full λ-vector — but - HARD-FAIL if the card omits a λ its OWN np_model uses. + """Validate that the card carries every λ its OWN np_model USES; return ``lc`` + unchanged otherwise (the name is historical — it no longer fills). A λ the card's np_model does NOT use (e.g. ``lambda6`` / ``lambda6_nu`` under - tanh_2) is inert; the model still needs a vector slot for it, so it is filled - with 0.0 (correct, not a guess — the form ignores it). A λ the np_model DOES - use (per :func:`params.active_params`) but the metadata lacks means a - stale/corrupt card that cannot describe its own model → raise rather than - silently default it. Cards written before an *inert* param was added (e.g. - pre-``lambda6_nu`` tanh_2 cards) therefore still load.""" - lc = dict(lc) + tanh_2) is simply absent: the de-hardcoded ``btgrid_tf`` form factors read only + the λ their branch needs, so no placeholder slot is required (previously such λ + were filled with 0.0). A λ the np_model DOES use (per + :func:`params.active_params`) but the metadata lacks means a stale/corrupt card + that cannot describe its own model → raise rather than silently default it. + Cards written before an *inert* λ was added (e.g. pre-``lambda6_nu`` tanh_2 + cards) therefore still load — the λ is neither needed nor filled.""" eff = dict(lc.get("eff_params", {})) gnu = dict(lc.get("gnu_params", {})) needed = active_params( @@ -208,11 +208,6 @@ def _fill_missing_params(lc): f"({eff.get(EFF_MODEL_KEY)} / {gnu.get(GNU_MODEL_KEY)}) USES — the card " f"cannot describe its own model; remake the histmaker output." ) - for k in EFF_PARAMS: - eff.setdefault(k, 0.0) # inert under this np_model -> 0 (vector slot only) - for k in GNU_PARAMS: - gnu.setdefault(k, 0.0) - lc["eff_params"], lc["gnu_params"] = eff, gnu return lc diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index 7ff44e884..da8ff441b 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -222,7 +222,7 @@ # ``self.core``. The λ-name tuples, σ_ns builder, and default-btgrid helper are # re-exported here for backward compat (older imports referenced them on this # module). ``fz_tf`` is still needed for the reco-side tensor dtype (``fz_tf.DTYPE``). -from wremnants.postprocessing.scetlib_np.params import DEFAULT_PRIOR_SIGMAS +from wremnants.postprocessing.scetlib_np.params import active_params, param_defaults from wremnants.postprocessing.scetlib_np.sigma_gen import ( # noqa: F401 ALL_PARAMS, EFF_PARAMS, @@ -419,7 +419,9 @@ def __init__( Per-name override for the Gaussian prior σ on each parameter: a Mapping, or as spec token the comma-separated string form ``prior_sigmas=lambda2=0.3,delta_lambda2=nan`` (same format as - ``xparam_default``). Defaults come from ``DEFAULT_PRIOR_SIGMAS`` + ``xparam_default``). Defaults come from the per-model registry + (``params.EFF_MODEL_PARAMS`` / ``GNU_MODEL_PARAMS``); a σ of ``None`` + there floats the param free. nonsingular_fo_sing, nonsingular_dyturbo Paths to the σ_ns inputs (SCETlib singular pkl / DYTurbo scetlibmatch txt); default to the wremnants-data @@ -691,8 +693,22 @@ def __init__( ) # ---- ParamModel registration (POIs first, then NOUs). + # Fit only the λ the numerator forms use (registry active set); inert λ + # aren't registered, so they can't add a zero-derivative (singular) Hessian row. poi_params = tuple(poi_params or ()) - nou_params = tuple(p for p in ALL_PARAMS if p not in poi_params) + active = active_params( + np_model=self._np_model_fit, np_model_nu=self._np_model_nu_fit + ) + bad_pois = [p for p in poi_params if p not in active] + if bad_pois: + raise ValueError( + f"poi_params {bad_pois} are not used by the fit forms " + f"({self._np_model_fit}/{self._np_model_nu_fit}); " + f"active λ: {sorted(active)}" + ) + nou_params = tuple( + p for p in ALL_PARAMS if p in active and p not in poi_params + ) self._param_order = poi_params + nou_params self.npoi = len(poi_params) self.npou = len(nou_params) @@ -712,21 +728,30 @@ def __init__( # grouped-impact bar directly comparable between the new param model and the # old NP variations. No collision with a syst group: the new-model datacard # excludes scetlibNP, so resumNonpert is absent from indata.systgroups. + # Intersect with the active params — a group must not name a non-fitted λ. + active_set = set(self._param_order) self.param_impact_groups = { - "resumNonpert": tuple(ALL_PARAMS), - "scetlibNPgammaNu": tuple(GNU_PARAMS), - "scetlibNPFeff": tuple(EFF_PARAMS), + "resumNonpert": tuple(p for p in ALL_PARAMS if p in active_set), + "scetlibNPgammaNu": tuple(p for p in GNU_PARAMS if p in active_set), + "scetlibNPFeff": tuple(p for p in EFF_PARAMS if p in active_set), } - # Defaults: λ_central values per parameter. Optionally overridden by the - # ``xparam_default=name=value,...`` spec token — comma-separated pairs (for - # closure tests where the data-generating / fit-start point should differ - # from the card's λ_central). + # Start value / prior mean per fitted λ, precedence: + # xparam_default[user] ▷ card λ_central (if the card carries it) ▷ + # registry neutral value (transform-extension λ the card lacks; = 0). + # pdefs holds {name: {"value","sigma"}} for the fit forms' active λ. + pdefs = param_defaults( + np_model=self._np_model_fit, np_model_nu=self._np_model_nu_fit + ) central_lookup = {**self.eff_central, **self.gnu_central} defaults = np.array( - [central_lookup[p] for p in self._param_order], dtype=np.float64 + [central_lookup.get(p, pdefs[p]["value"]) for p in self._param_order], + dtype=np.float64, ) + # ``xparam_default=name=value,...`` shifts the fit START (and prior mean), + # e.g. for closure / injection / transport tests. A name not used by the fit + # forms is warned-and-ignored; an unknown λ name is an error (typo guard). start_override = (xparam_default or "").strip() if start_override: overrides = dict( @@ -734,8 +759,16 @@ def __init__( ) for name, val in overrides.items(): name = name.strip() - if name not in self._param_order: + if name not in ALL_PARAMS: raise KeyError(f"xparam_default: unknown param {name!r}") + if name not in self._param_order: + print( + f"[SCETlibNPParamModel] WARNING: xparam_default {name!r} is " + f"not used by the fit forms " + f"({self._np_model_fit}/{self._np_model_nu_fit}); ignoring.", + flush=True, + ) + continue i = self._param_order.index(name) defaults[i] = float(val) print( @@ -765,9 +798,21 @@ def __init__( tuple(s.split("=")) for s in prior_sigmas.split(",") if s.strip() ) prior_sigmas = {k.strip(): v for k, v in dict(prior_sigmas or {}).items()} - for name in prior_sigmas: - if name not in self._param_order: + # A prior on a λ not used by the fit forms is warned-and-ignored; an unknown + # λ name is an error (typo guard). + _kept_prior_sigmas = {} + for name, v in prior_sigmas.items(): + if name not in ALL_PARAMS: raise KeyError(f"prior_sigmas: unknown param {name!r}") + if name not in self._param_order: + print( + f"[SCETlibNPParamModel] WARNING: prior_sigmas {name!r} is not used " + f"by the fit forms; ignoring.", + flush=True, + ) + continue + _kept_prior_sigmas[name] = v + prior_sigmas = _kept_prior_sigmas if self._use_priors: sigmas_arr = np.empty(self.nparams, dtype=np.float64) for i, p in enumerate(self._param_order): @@ -775,10 +820,10 @@ def __init__( sigmas_arr[i] = float( prior_sigmas[p] ) # explicit override (may be NaN) - elif p in DEFAULT_PRIOR_SIGMAS: - sigmas_arr[i] = DEFAULT_PRIOR_SIGMAS[p] # theorist recommendation else: - sigmas_arr[i] = np.nan # free (expected to be frozen) + # registry default sigma for this (fit-model, param); None = free + s = pdefs[p]["sigma"] + sigmas_arr[i] = np.nan if s is None else float(s) self.prior_sigmas = sigmas_arr # prior_means defaults to xparamdefault if not set, so don't store # redundantly — Fitter falls back to xparamdefault. From 9ff9eec218f765d99b8cc4747a5af06ddab44619 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Wed, 1 Jul 2026 17:36:44 -0400 Subject: [PATCH 30/31] =?UTF-8?q?SCETlib-NP:=20derive=20NPDampingWall=20re?= =?UTF-8?q?quired=20=CE=BB=20from=20the=20central=20registry?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wall hardcoded _REQUIRED_PARAMS (all 9 λ) and read λ6*/λ6 unconditionally, assuming every model carries the full ALL_PARAMS vocabulary. Since "fit only the active λ" made the model's _param_order form-dependent (tanh_2 drops λ6*), the wall wrongly required λ6* for a tanh_2 model and raised at construction. Now it derives its required λ from active_params(fit_forms) — the same registry (params.EFF/GNU_MODEL_PARAMS) that builds the model's _param_order — and reads λ6*/λ6 only inside their tanh_6 branches. The wall's λ vocabulary tracks the chosen model per-side; it can never require a λ the model omits. Validated: tanh_2 + wall constructs and the penalty evaluates (was the ValueError); tanh_6-CS override requires λ6_ν but not λ6. Co-Authored-By: Claude Opus 4.8 --- .../scetlib_np/np_damping_wall.py | 47 +++++++++---------- 1 file changed, 23 insertions(+), 24 deletions(-) diff --git a/wremnants/postprocessing/scetlib_np/np_damping_wall.py b/wremnants/postprocessing/scetlib_np/np_damping_wall.py index 09641be2f..a8aa00e98 100644 --- a/wremnants/postprocessing/scetlib_np/np_damping_wall.py +++ b/wremnants/postprocessing/scetlib_np/np_damping_wall.py @@ -86,7 +86,9 @@ """ # rabbit / TF imports deferred to the lazy class factories so the module stays -# importable without rabbit/TF (mirrors np_monotonicity.py). +# importable without rabbit/TF (mirrors np_monotonicity.py). The λ registry +# (params) is numpy-only, so it is safe to import at module level. +from wremnants.postprocessing.scetlib_np.params import active_params # Forms this wall has damping conditions for. Anything else (frac_*, exp_*, # tanh_1, tanh_4, signed_lambda, identity, …) raises rather than silently @@ -99,20 +101,6 @@ # raise. Applied to both sides; only "hyp_tangent" reaches a supported form. _FORM_ALIASES = {"hyp_tangent": "tanh_2"} -# λ that the walls reference (always present in the model's _param_order, which -# is built from params.ALL_PARAMS); cross-checked at construction. -_REQUIRED_PARAMS = ( - "lambda2_nu", - "lambda4_nu", - "lambda6_nu", - "lambda_inf_nu", - "lambda2", - "lambda4", - "lambda6", - "delta_lambda2", - "lambda_inf", -) - # Fixed knobs (NOT CLI options — edit here to test). Only ``smallb`` is exposed # on the -r line; these are deliberately constants to keep that line minimal. Y_MAX = 5.0 # binding |y| for the F_eff Y-evaluation: the kinematic ceiling @@ -172,6 +160,7 @@ def parse_args(cls, indata, *args): def _make_regularizer_class(): import tensorflow as tf + from rabbit.regularization.regularizer import Regularizer class NPDampingWall(Regularizer): @@ -205,12 +194,6 @@ def __init__(self, mapping, dtype): self._order = tuple(model._param_order) self._pidx = {name: i for i, name in enumerate(self._order)} - missing = [p for p in _REQUIRED_PARAMS if p not in self._pidx] - if missing: - raise ValueError( - f"NPDampingWall: model param order {self._order} is missing " - f"required λ {missing}." - ) # FIT (numerator) forms — the ones the fit integrates, which the wall # must constrain (NOT the card/denominator form). Resolve aliases and @@ -223,6 +206,21 @@ def __init__(self, mapping, dtype): forms["np_model"], side="TMD (F_eff, np_model)" ) + # Required λ come from the CENTRAL REGISTRY, keyed on the resolved + # forms — the SAME source (active_params) the model uses to build + # _param_order. Each model has its own λ vocabulary (e.g. tanh_2 has + # no λ6*), so the wall can never require a λ the chosen model omits. + required = active_params( + np_model=self._np_model, np_model_nu=self._np_model_nu + ) + missing = [p for p in sorted(required) if p not in self._pidx] + if missing: + raise ValueError( + f"NPDampingWall: model param order {self._order} is missing " + f"λ {missing} required by the fit forms " + f"(np_model={self._np_model!r}, np_model_nu={self._np_model_nu!r})." + ) + self._cast = lambda v: tf.constant(v, dtype=self.dtype) # Model-param block is x[:nparams]; nparams resolved at set_expectations. self._nparams = None @@ -277,12 +275,12 @@ def wall(coeff): # coeff ≥ margin # ---- CS-side γ_ν^NP damping: P(u)=λ2_ν·u+λ4_ν·u²+λ6_ν·u³ ≥ 0 ∀u≥0. l2nu = self._lam(params, "lambda2_nu") l4nu = self._lam(params, "lambda4_nu") - l6nu = self._lam(params, "lambda6_nu") linfnu = self._lam(params, "lambda_inf_nu") pens = [relu2(eps - linfnu)] # λ∞_ν > 0 (saturation-scale regime) if self._np_model_nu == "tanh_2": pens.append(wall(l4nu)) # λ4_ν ≥ margin (large-b leading) - else: # tanh_6 + else: # tanh_6 — λ6_ν only exists in the tanh_6 vocabulary + l6nu = self._lam(params, "lambda6_nu") pens.append(wall(l6nu)) # λ6_ν ≥ margin (large-b leading) # interior: λ4_ν ≥ 0 OR λ4_ν² ≤ 4·λ2_ν·λ6_ν. Self-gating — the # relu2(-l4nu) vanishes for λ4_ν ≥ 0, so no penalty there; and no @@ -296,9 +294,10 @@ def wall(coeff): # coeff ≥ margin # all conditions stay division-free (multiply through by 3·λ∞² > 0). l2 = self._lam(params, "lambda2") l4 = self._lam(params, "lambda4") - l6 = self._lam(params, "lambda6") dl2 = self._lam(params, "delta_lambda2") linf = self._lam(params, "lambda_inf") + # λ6 only exists in the tanh_6 vocabulary; read it only when used. + l6 = self._lam(params, "lambda6") if self._np_model == "tanh_6" else None pens.append(relu2(eps - linf)) # λ∞ > 0 linf2 = linf * linf for y_sq in (0.0, self.ymax * self.ymax): From ffbe5b46f291e085123c625f3b3e1b484934cb84 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 10 Jul 2026 14:51:32 -0400 Subject: [PATCH 31/31] updates --- wremnants/postprocessing/scetlib_np/README.md | 45 ++ .../scetlib_np/fitresult_lambdas.py | 38 +- .../scetlib_np/lambda_central.py | 57 +- .../scetlib_np/np_function_plots.py | 49 +- .../postprocessing/scetlib_np/param_model.py | 75 ++- .../scetlib_np/param_model_diagnostics.py | 226 ++++--- .../postprocessing/scetlib_np/plot_output.py | 74 ++ .../scetlib_np/sigma_gen_at_lambda.py | 630 ++++++++---------- .../scetlib_np/validate_agreement.py | 339 ++++++++++ 9 files changed, 1015 insertions(+), 518 deletions(-) create mode 100644 wremnants/postprocessing/scetlib_np/README.md create mode 100644 wremnants/postprocessing/scetlib_np/plot_output.py create mode 100644 wremnants/postprocessing/scetlib_np/validate_agreement.py diff --git a/wremnants/postprocessing/scetlib_np/README.md b/wremnants/postprocessing/scetlib_np/README.md new file mode 100644 index 000000000..f5b273bb7 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/README.md @@ -0,0 +1,45 @@ +# `scetlib_np` — SCETlib non-perturbative param model + +The fit-time NP model (`SCETlibNPParamModel`) that reconstructs σ_gen from a +cached SCETlib bt-grid, folds it through the response matrix R, and applies +`rnorm = σ_reco(λ)/σ_reco(λ_c)` in the rabbit fit — plus the tools to validate it. + +Run everything inside the wmass singularity with the venv + `setup.sh` sourced. + +## User entry points + +### Validation / agreement — the two you'll normally use +| command | what it checks | +|---|---| +| `python -m …scetlib_np.validate_agreement --reference card --datacard [--outdir ]` | model σ_reco **and** σ_gen at λ_central vs the **datacard itself** (`norm[signal]` + `N_gen`). No external inputs. | +| `python -m …scetlib_np.validate_agreement --reference histmaker --datacard --histmaker [--plot-out ] [--gen-histmaker ] [--variation lambda21.0 …]` | the same vs an external **histmaker** `nominal` / gen MC, plus the reco λ-**variations** (`Corr[var]/Corr[pdf0]`). | +| `python -m …scetlib_np.sigma_gen_at_lambda --theory-corr [--datacard ] [--lambdas lambda2=0.5 …] [--fitresult ] [--plot ]` | σ_gen at an **arbitrary λ tune** (cardless, gen-only) vs the official **TheoryCorrection**. Distinct: no datacard/response needed, any λ. | + +### Analysis / inspection +| command | what it does | +|---|---| +| `python -m …scetlib_np.fitresult_lambdas …` | read fitted NP λ out of a rabbit fitresults (table / form-factor curves / toys). | +| `python -m …scetlib_np.np_function_plots …` | plot the NP form factors: CS γ_ν^NP(b_T) and TMD F_eff(b_T, y). | +| `python -m …scetlib_np.point_to_binned …` | convert a POINT-spectrum SCETlib pickle → a binned `{hist: hist.Hist}`. | +| `python -m …scetlib_np.lambda_central ` | inspect the central NP (λ) tune carried by a datacard/correction. | +| `python -m …scetlib_np.response_matrix ` | load / inspect the (reco × gen) response matrix R. | + +### Developer validation, smoke & timing (`…scetlib_np.validation.`) +Not everyday tools — deeper cross-checks of the bt-grid factorization: +`native_validation` (native-binning vs SCETlib spectrum-mode ref), +`resum_validation` (resummed σ_gen from two sources), `export_spectrum` +(export σ onto a SCETlib run's grid), `gen_level_smoke` (gen_level=1 fold-free), +`factorized_parity` (legacy vs factorized reconstruction), `truth_start_grid` +(random-truth POI recovery), `damping_wall_dispatch` (NPDampingWall dispatch), +`timing` (one-σ_gen cost). + +## Library modules (imported, not run) +`param_model.py` (the rabbit fit model) · `sigma_gen.py` (datacard-free σ_gen core) · +`btgrid_{cache,integrate,tf}.py` (bt-grid load + Hankel/TF integration) · +`response_matrix.py` (R) · `params.py`, `lambda_central.py` (axes / central λ) · +`np_damping_wall.py` (fit-time physical-NP regularizer) · +`validation/agreement.py` (shared reference loaders/aligners for the CLIs) · +`param_model_diagnostics.py` (**Layer 0**: the numpy in-fit `run_reco_guard`, the +postfit pathology detectors, and `run_card_diagnostics` — the impl behind +`validate_agreement --reference card`) · `validation_plots.py`, `plot_output.py` +(shared plotting/array helpers + the one save entry point). diff --git a/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py b/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py index 74487e5e4..50d00b9b6 100644 --- a/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py +++ b/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py @@ -22,8 +22,10 @@ Units: the new-model λ are already physical (see ``param_model`` / ``allowNegativeParam``); no conversion. The ``np_model`` / ``np_model_nu`` strings -the curves need come from :func:`lambda_central.read_lambda_central` (read off the -fitresults), with a CLI/argument override. +the curves need are the FIT (numerator) forms: the ``np_model_fit`` / +``np_model_nu_fit`` override tokens rabbit stored in the fitresults meta when +given, else the card (denominator) form from +:func:`lambda_central.read_lambda_central`, with a CLI/argument override on top. CLI (print the table):: @@ -39,7 +41,12 @@ from rabbit import io_tools from wremnants.postprocessing.scetlib_np import lambda_central as _lc from wremnants.postprocessing.scetlib_np.np_function_plots import NPLambdas, Series -from wremnants.postprocessing.scetlib_np.params import ALL_PARAMS, EFF_PARAMS, GNU_PARAMS +from wremnants.postprocessing.scetlib_np.params import ( + ALL_PARAMS, + EFF_PARAMS, + GNU_PARAMS, +) + SECTOR = { **{p: "gamma_nu (CS)" for p in GNU_PARAMS}, **{p: "F_eff (TMD)" for p in EFF_PARAMS}, @@ -253,22 +260,25 @@ def _flat_values(fitresult_path, which="postfit", result=None): def _resolve_models(fitresult_path, np_model=None, np_model_nu=None): - """np_model strings from lambda_central (read off the fitresults), falling - back to defaults if the upstream pkl is unreachable. Explicit arguments win.""" + """The FIT (numerator) np_model strings. Explicit arguments win, then + :func:`lambda_central.read_np_models` (the central resolver: card form + overridden by the fit's ``np_model_(nu_)fit`` tokens), then the defaults.""" if np_model and np_model_nu: return np_model, np_model_nu eff_model, gnu_model = DEFAULT_NP_MODEL, DEFAULT_NP_MODEL_NU try: - lc = _lc.read_lambda_central(fitresult_path) - eff_model = lc["eff_params"].get("np_model", eff_model) - gnu_model = lc["gnu_params"].get("np_model_nu", gnu_model) - except Exception as exc: # pkl unreachable / non-NP fit: warn, use defaults + eff_model, gnu_model = _lc.read_np_models(fitresult_path) + except Exception as exc: # metadata unreadable / non-NP fit: warn, use defaults print( - f"[fitresult_lambdas] could not read np_model from {fitresult_path} " - f"({exc}); using defaults {eff_model!r}/{gnu_model!r}. " + f"[fitresult_lambdas] could not read the np_model forms from " + f"{fitresult_path} ({exc}); using defaults " + f"{DEFAULT_NP_MODEL!r}/{DEFAULT_NP_MODEL_NU!r}. " f"Pass --np-model/--np-model-nu to override." ) - return np_model or eff_model, np_model_nu or gnu_model + return ( + np_model or eff_model or DEFAULT_NP_MODEL, + np_model_nu or gnu_model or DEFAULT_NP_MODEL_NU, + ) def lambdas_from_fitresult( @@ -393,7 +403,9 @@ def _template_base_eff_gnu(param_map): np_model="tanh_6", ) gnu = dict( - lambda_inf_nu=0.0, lambda2_nu=0.0, lambda4_nu=0.0, + lambda_inf_nu=0.0, + lambda2_nu=0.0, + lambda4_nu=0.0, lambda6_nu=0.0007, # SCETlib NP_model_gammanu b⁶ coeff (Gamma_nu.hpp:102) np_model_nu="tanh_6", ) diff --git a/wremnants/postprocessing/scetlib_np/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py index 3f4fa4341..3d1cd9038 100644 --- a/wremnants/postprocessing/scetlib_np/lambda_central.py +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -239,6 +239,14 @@ def read_lambda_central_from_meta(meta, proc="Z", _source=""): ) +def _read_meta(hdf5_path): + """Load the pickled ``meta`` group of a datacard / fitresults hdf5.""" + with h5py.File(hdf5_path, "r") as f: + if "meta" not in f: + raise KeyError(f"{hdf5_path}: no 'meta' group -- wrong file type?") + return wums_io.pickle_load_h5py(f["meta"]) + + def read_lambda_central(hdf5_path, proc="Z"): """Read the central lambda parameters referenced by an hdf5. @@ -249,16 +257,53 @@ def read_lambda_central(hdf5_path, proc="Z"): Returns ``{tag, basename, eff_params, gnu_params, source}``. Raises if the metadata is absent. """ - with h5py.File(hdf5_path, "r") as f: - if "meta" not in f: - raise KeyError(f"{hdf5_path}: no 'meta' group -- wrong file type?") - meta = wums_io.pickle_load_h5py(f["meta"]) - - lc = read_lambda_central_from_meta(meta, proc=proc, _source=hdf5_path) + lc = read_lambda_central_from_meta( + _read_meta(hdf5_path), proc=proc, _source=hdf5_path + ) lc["source"] = "histmaker-metadata" return lc +def read_fit_form_overrides_from_meta(meta): + """``(np_model_fit, np_model_nu_fit)`` from the rabbit ``--paramModel`` spec + stored in a fitresults' ``meta_info.args``; ``None`` per slot when the fit did + not override that form. A datacard (no ``meta_info.args.paramModel``) yields + ``(None, None)``. Keys on the literal ``np_model_(nu_)fit=`` spec tokens — + coupled to the ``param_model`` constructor's argument spelling.""" + args = (meta.get("meta_info") or {}).get("args") or {} + specs = args.get("paramModel") or [] + eff = gnu = None + for spec in specs: + for tok in spec: + tok = tok.decode() if isinstance(tok, bytes) else str(tok) + if tok.startswith("np_model_fit="): + eff = tok.split("=", 1)[1] + elif tok.startswith("np_model_nu_fit="): + gnu = tok.split("=", 1)[1] + return eff, gnu + + +def read_np_models(hdf5_path, proc="Z"): + """The NP functional forms that apply to PREDICTIONS from this hdf5: + ``(np_model, np_model_nu)`` — the single resolver every offline tool + should use. + + The card (denominator) forms from the propagated + ``scetlib_np_lambda_central`` metadata, overridden per sector by the + ``np_model_(nu_)fit`` numerator tokens when the file is a rabbit fitresults + whose ``--paramModel`` spec carried them (mirrors ``param_model``: + ``np_model_fit or card form``). A datacard has no fit override and resolves + to the card forms. Raises (KeyError) if the file carries no lambda_central + metadata at all (non-NP input).""" + meta = _read_meta(hdf5_path) + lc = read_lambda_central_from_meta(meta, proc=proc, _source=hdf5_path) + eff_fit, gnu_fit = read_fit_form_overrides_from_meta(meta) + return ( + eff_fit or lc["eff_params"].get(EFF_MODEL_KEY), + gnu_fit or lc["gnu_params"].get(GNU_MODEL_KEY), + ) + + if __name__ == "__main__": import sys diff --git a/wremnants/postprocessing/scetlib_np/np_function_plots.py b/wremnants/postprocessing/scetlib_np/np_function_plots.py index 54db262e6..0cc0e00b2 100644 --- a/wremnants/postprocessing/scetlib_np/np_function_plots.py +++ b/wremnants/postprocessing/scetlib_np/np_function_plots.py @@ -37,7 +37,6 @@ """ import argparse -import os from dataclasses import dataclass from typing import List, Optional, Sequence, Tuple @@ -146,9 +145,7 @@ def _param_inset(ax, lam, sector, corner="upper right"): model = lam.eff.get("np_model", "?") active = active_params(np_model=model) labels, src, order = _EFF_LABELS, lam.eff, EFF_PARAMS - lines = [ - rf"${labels[k]} = {src.get(k, 0):+.4f}$" for k in order if k in active - ] + lines = [rf"${labels[k]} = {src.get(k, 0):+.4f}$" for k in order if k in active] lines.append(rf"model: {model}") box = dict(boxstyle="round,pad=0.35", fc="white", ec="0.6", alpha=0.85) x, y, va, ha = _CORNERS[corner] @@ -173,6 +170,7 @@ def plot_np_functions( outpath: str, inset_from: Optional[Series] = None, f_ymax: Optional[float] = None, + args=None, ): """Draw the two NP form factors for one or more λ sets. @@ -275,13 +273,15 @@ def plot_np_functions( # Allow --outpath to be a directory (trailing slash or no extension): append # a default filename rather than erroring on a bare ".png". - if outpath.endswith(("/", os.sep)) or os.path.isdir(outpath) or not os.path.splitext(outpath)[1]: - outpath = os.path.join(outpath, "np_functions.png") - os.makedirs(os.path.dirname(os.path.abspath(outpath)) or ".", exist_ok=True) + from wremnants.postprocessing.scetlib_np import plot_output + fig.tight_layout() - fig.savefig(outpath, dpi=140) + outdir, basename = plot_output.split_outpath( + outpath, default_name="np_functions.png" + ) + plot_output.save_plot(outdir, basename, fig=fig, args=args, dpi=140) plt.close(fig) - print(f"Wrote {outpath}") + print(f"Wrote {outdir}/{basename}.png(.pdf) + {basename}.log") # --------------------------------------------------------------------------- @@ -319,9 +319,7 @@ def make_parser(): src.add_argument( "--result", default=None, help="results group suffix for --fitresult." ) - src.add_argument( - "--n-toys", type=int, default=500, help="band toys (--fitresult)." - ) + src.add_argument("--n-toys", type=int, default=500, help="band toys (--fitresult).") src.add_argument("--seed", type=int, default=0, help="band RNG seed.") raw = p.add_argument_group("raw λ mode (no fit)") @@ -332,8 +330,18 @@ def make_parser(): "λ stay at the NP-unit defaults. Names must be λ the chosen models use " "(lambda6/lambda6_nu need tanh_6); an inert λ is a hard error.", ) - raw.add_argument("--np-model", default="tanh_2", help="F_eff model string.") - raw.add_argument("--np-model-nu", default="tanh_2", help="γ_ν model string.") + raw.add_argument( + "--np-model", + default=None, + help="F_eff model string (raw mode default: tanh_2; with --fitresult " + "the fit form is read from the fitresults, this overrides it).", + ) + raw.add_argument( + "--np-model-nu", + default=None, + help="γ_ν model string (raw mode default: tanh_2; with --fitresult " + "the fit form is read from the fitresults, this overrides it).", + ) p.add_argument( "--y", @@ -368,25 +376,29 @@ def main(argv=None): result=args.result, n_toys=args.n_toys, seed=args.seed, + np_model=args.np_model, + np_model_nu=args.np_model_nu, ) else: + np_model = args.np_model or "tanh_2" + np_model_nu = args.np_model_nu or "tanh_2" try: overrides = parse_lambda_overrides(args.lambdas) except ValueError as e: parser.error(str(e)) - active = active_params(args.np_model, args.np_model_nu) + active = active_params(np_model, np_model_nu) inert = [k for k in overrides if k not in active] if inert: parser.error( "--lambdas: " + ", ".join(inert) - + f" not used by np_model={args.np_model} / " - + f"np_model_nu={args.np_model_nu} (active: " + + f" not used by np_model={np_model} / " + + f"np_model_nu={np_model_nu} (active: " + ", ".join(k for k in (*EFF_PARAMS, *GNU_PARAMS) if k in active) + ")" ) vals = {**_DEFAULT_LAMBDAS, **overrides} - lam = NPLambdas.from_flat(vals, args.np_model, args.np_model_nu) + lam = NPLambdas.from_flat(vals, np_model, np_model_nu) series = [Series(label=args.label, lam=lam, color="C3")] plot_np_functions( @@ -395,6 +407,7 @@ def main(argv=None): bT_max=args.bT_max, outpath=args.outpath, f_ymax=args.f_ymax, + args=args, ) diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py index da8ff441b..8ccf92882 100644 --- a/wremnants/postprocessing/scetlib_np/param_model.py +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -128,8 +128,11 @@ P(b | g) = R_raw(b, g) / N_gen(g) (efficiency × migration) σ_reco(λ; b) = Σ_g P(b | g) · σ_gen(λ; g) -g = gen bin (ptVGen, |Y|), summed over by Σ_g; b = reco bin (ptll, yll, -cosThetaStarll_quantile, phiStarll_quantile). P(b | g) is the gen→reco map (one +g = gen bin (ptVGen, |Y|), summed over by Σ_g; b = reco bin — the fit channel's +axes, any prefix subset of the canonical (ptll, yll, cosThetaStarll_quantile, +phiStarll_quantile) order (e.g. a 2D ptll-yll fit; the embedded 4D R is +marginalized over the missing axes, see ``_marginalize_R_reco``). P(b | g) is +the gen→reco map (one reco column per gen bin), so σ_reco is σ_gen pushed through the detector; Σ_g is ``tf.linalg.matvec(self.R, σ_gen_flat)``. Pure detector folding — no λ_central or ratio here (that is Step 4). @@ -207,7 +210,6 @@ fit. Full derivation, GN-vs-full-K, and the exact commands: ``docs/HESSIAN_PLAN.md``. """ -import os from typing import Mapping, Optional import numpy as np @@ -224,12 +226,12 @@ # module). ``fz_tf`` is still needed for the reco-side tensor dtype (``fz_tf.DTYPE``). from wremnants.postprocessing.scetlib_np.params import active_params, param_defaults from wremnants.postprocessing.scetlib_np.sigma_gen import ( # noqa: F401 + _NONSING_DYTURBO_DEFAULT, + _NONSING_FO_SING_DEFAULT, ALL_PARAMS, EFF_PARAMS, GNU_PARAMS, SigmaGenModel, - _NONSING_DYTURBO_DEFAULT, - _NONSING_FO_SING_DEFAULT, _default_btgrid_dir, compute_nonsingular_gen, ) @@ -249,6 +251,7 @@ 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 trailing reco bins so its reco shape matches the fit. @@ -282,6 +285,41 @@ def _crop_R_to_fit(R, R_reco_axes, fit_reco_axes, tol=1e-9): return R[slices] +def _marginalize_R_reco(R, R_reco_axes, fit_axis_names): + """Sum R over the reco axes the fit channel doesn't have. + + The datacard embeds R at the canonical reco binning (the full 4D + ptll/yll/cosThetaStar*/phiStar* grid, see ``params.RECO_AXES``) regardless + of the fit channel's dimensionality. R(b, g) is a counts response, so the + response for a lower-dimensional reco channel (e.g. a 2D ptll-yll fit) is + exactly the marginal over the dropped reco axes — sum them out here. Gen + axes (the trailing axes of R) are untouched. The kept axes must appear in + the fit's order (both follow the canonical ordering, so a mismatch means a + non-canonical fit channel, which the crop couldn't handle either). + """ + 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"canonical 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"[SCETlibNPParamModel] 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): """Reconstruct the response-matrix dict from the datacard's ``scetlib_np`` auxiliary bundle. @@ -556,11 +594,17 @@ def __init__( else: # ---- R matrix (read from the datacard's scetlib_np auxiliary) R_info = _R_info_from_auxiliary(indata) + fit_reco_axes = self._fit_reco_axes(indata) + # The auxiliary R carries the canonical (4D) reco binning whatever + # the fit channel is; a lower-dimensional channel (e.g. 2D ptll-yll) + # uses the marginal response — sum R over the axes the fit lacks. + R_full, R_reco_axes = _marginalize_R_reco( + R_info["R"], R_info["reco_axes"], [n for n, _ in fit_reco_axes] + ) # Fit-tensor reco binning may differ from R's by trailing overflow bins # (e.g. R has ptll [0, …, 44, 100] while the fit ends at 44). Crop R's # trailing bins so the reco shape matches. - fit_reco_axes = self._fit_reco_axes(indata) - R_arr = _crop_R_to_fit(R_info["R"], R_info["reco_axes"], 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"] # Gen-total denominator N_gen(g) from the xnorm hist ("prefsr"): @@ -579,7 +623,7 @@ def __init__( self._reco_axes_meta = [ (name, fit_axes[1]) for (name, fit_axes) in zip( - [a[0] for a in R_info["reco_axes"]], + [a[0] for a in R_reco_axes], fit_reco_axes, ) ] @@ -706,9 +750,7 @@ def __init__( f"({self._np_model_fit}/{self._np_model_nu_fit}); " f"active λ: {sorted(active)}" ) - nou_params = tuple( - p for p in ALL_PARAMS if p in active and p not in poi_params - ) + nou_params = tuple(p for p in ALL_PARAMS if p in active and p not in poi_params) self._param_order = poi_params + nou_params self.npoi = len(poi_params) self.npou = len(nou_params) @@ -961,7 +1003,9 @@ def __getattr__(self, name): f"{type(self).__name__!r} object has no attribute {name!r}" ) - def _sigma_YqT_native_at(self, eff_params, gnu_params, np_model=None, np_model_nu=None): + def _sigma_YqT_native_at( + self, eff_params, gnu_params, np_model=None, np_model_nu=None + ): """Backward-compat alias for :meth:`SigmaGenModel.sigma_YqT_native` — the native (NY, NqT) Q-integrated σ(λ) before the |Y|-fold / qT-rebin. ``np_model`` / ``np_model_nu`` override the form (default: card form).""" @@ -976,8 +1020,11 @@ def _sigma_gen_at( σ_gen(λ) on the (NptVGen, NabsYVGen) gen grid (Steps 1–2). ``np_model`` / ``np_model_nu`` override the form (default: card form).""" return self.core.sigma_gen( - eff_params, gnu_params, sigma_YqT=sigma_YqT, - np_model=np_model, np_model_nu=np_model_nu, + eff_params, + gnu_params, + sigma_YqT=sigma_YqT, + np_model=np_model, + np_model_nu=np_model_nu, ) # ========================================================================= diff --git a/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py b/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py index f248aabe9..fdba702a6 100644 --- a/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py +++ b/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py @@ -22,14 +22,19 @@ cancels in the fit's ratio). Density plots unit-normalize both curves, so the comparison carries no ad-hoc scale. -Entry points: +This is a LIBRARY module (no CLI). The user-facing card agreement check is +``python -m wremnants.postprocessing.scetlib_np.validate_agreement --reference card``, +which calls :func:`run_card_diagnostics` here. + +Entry points (imported, not run): * :func:`run_reco_guard` — PURE-NUMPY per-bin reco check for the in-fit - auto-guard (warns, or raises with ``strict``; never imports plotting). - * :func:`run_card_diagnostics` — full reco + gen comparison (+ optional plots) - for interactive use; returns the per-bin residuals. - * ``python -m wremnants.postprocessing.scetlib_np.param_model_diagnostics - --datacard [--outdir ]`` — construct from a card, write the - reco + gen agreement plots. + auto-guard (warns, or raises with ``strict``; never imports plotting). Called + by the fit (``param_model.py``) at construction time, default-on. + * :func:`run_card_diagnostics` — full reco + gen comparison (+ optional plots); + the implementation behind ``validate_agreement --reference card``. + * the pathology detectors (:func:`np_damping_ok`, :func:`spectrum_negativity`, + :func:`np_physical_report`) — postfit NP-validity cross-checks, also used by + ``sigma_gen_at_lambda``. Heavy deps (``hist``, ``wums.plot_tools``, and the plotting/projection helpers from ``scetlib_np.validation_plots``) are imported LAZILY in the print/plot paths @@ -40,13 +45,6 @@ import numpy as np -# Per-reco/gen axis projection order for the standalone plots (shared names). -from wremnants.postprocessing.scetlib_np.params import ( - GEN_AXES as GEN_PROJ_AXES, - RECO_AXES as RECO_PROJ_AXES, -) - - # ============================================================================= # Postfit NP physical-validity detectors (standalone — NOT part of the fit). # ============================================================================= @@ -61,7 +59,15 @@ NP_PROBE_BT = (0.3, 1.0, 2.0, 5.0, 10.0, 20.0) -def np_damping_ok(core, eff_params, gnu_params, b_probe=NP_PROBE_BT, gamma_tol=1e-3): +def np_damping_ok( + core, + eff_params, + gnu_params, + b_probe=NP_PROBE_BT, + gamma_tol=1e-3, + np_model=None, + np_model_nu=None, +): """Probe the NP form factors (no bT integral) for the physical DAMPING sign. Evaluates the actual ``btgrid_tf`` forms the fit integrates at a few bT: @@ -73,15 +79,19 @@ def np_damping_ok(core, eff_params, gnu_params, b_probe=NP_PROBE_BT, gamma_tol=1 probes test the SAME damping condition the CS walls enforce (form-agnostically — the b⁶ term is in the evaluated form, so this also covers tanh_6); the F_eff endpoint test at Y=0 is CRUDER than the wall's exact a≥0 ∀b at Y=0 and Y_max. - NOTE this probes the CARD form (``core.np_model``/``np_model_nu``), so it does - not see a numerator-form override (``np_model_(nu_)fit``). Cheap (1-D evals).""" + ``np_model`` / ``np_model_nu`` select the forms to probe (default: the card + forms ``core.np_model`` / ``core.np_model_nu``) — for a numerator-form + override fit (``np_model_(nu_)fit``) pass the FIT forms the λ belong to, else + the verdict is about the wrong form. Cheap (1-D evals).""" from wremnants.postprocessing.scetlib_np import btgrid_tf as fz_tf b = np.asarray(b_probe, dtype=np.float64) eff = {k: v for k, v in eff_params.items() if k != "np_model"} gnu = {k: v for k, v in gnu_params.items() if k != "np_model_nu"} - g = fz_tf.gamma_nu_NP_tf(b, gnu, np_model_nu=core.np_model_nu).numpy() - F = fz_tf.F_eff_tf(0.0, b, eff, np_model=core.np_model).numpy() + g = fz_tf.gamma_nu_NP_tf( + b, gnu, np_model_nu=np_model_nu or core.np_model_nu + ).numpy() + F = fz_tf.F_eff_tf(0.0, b, eff, np_model=np_model or core.np_model).numpy() gamma_max = float(np.max(g)) feff_growing = bool(F[-1] > F[0]) return { @@ -95,15 +105,25 @@ def np_damping_ok(core, eff_params, gnu_params, b_probe=NP_PROBE_BT, gamma_tol=1 } -def spectrum_negativity(core, eff_params, gnu_params, sigma_YqT=None, locate=True): +def spectrum_negativity( + core, + eff_params, + gnu_params, + sigma_YqT=None, + locate=True, + np_model=None, + np_model_nu=None, +): """Negativity of the native (Y, qT) resummed spectrum at a λ point — the σ(qT) < 0 pathology the gen-binning averages away. Scale-free metrics: neg_area_frac = Σ|min(σ,0)| / Σ|σ| (0 physical; → O(1) pathological) min_over_peak = min(σ) / max(σ) (≈0 physical; ≤ −O(1) pathological) Pass ``sigma_YqT`` (e.g. ``core.sigma_YqT_central``) to skip recomputation, - else reconstructed via ``core.sigma_YqT_native``. Judge relative to the - λ_central baseline (``np_physical_report`` does this): the singular-only - spectrum carries a tiny benign qT→0 dip. + else reconstructed via ``core.sigma_YqT_native`` with the ``np_model`` / + ``np_model_nu`` forms (default: the card/construction forms — for a + numerator-form override fit pass the FIT forms the λ belong to). Judge + relative to the λ_central baseline (``np_physical_report`` does this): the + singular-only spectrum carries a tiny benign qT→0 dip. With ``locate`` (default True) and a 2-D spectrum on the core's native grids (``core.Y_unique`` × ``core.qT_unique``), also returns WHERE the negativity @@ -114,7 +134,9 @@ def spectrum_negativity(core, eff_params, gnu_params, sigma_YqT=None, locate=Tru ``neg_qT_range`` / ``neg_absY_max`` extent of the negative region (omitted if the grids are absent or the shape doesn't match).""" s = ( - core.sigma_YqT_native(eff_params, gnu_params) + core.sigma_YqT_native( + eff_params, gnu_params, np_model=np_model, np_model_nu=np_model_nu + ) if sigma_YqT is None else sigma_YqT ) @@ -132,29 +154,36 @@ def spectrum_negativity(core, eff_params, gnu_params, sigma_YqT=None, locate=Tru if ( locate and s.ndim == 2 - and Yg is not None and qTg is not None - and Yg.ndim == 1 and qTg.ndim == 1 + and Yg is not None + and qTg is not None + and Yg.ndim == 1 + and qTg.ndim == 1 and s.shape == (Yg.size, qTg.size) ): jmin = np.unravel_index(int(np.argmin(s)), s.shape) out["worst"] = dict( - iY=int(jmin[0]), iqT=int(jmin[1]), - Y=float(Yg[jmin[0]]), qT=float(qTg[jmin[1]]), - value=float(s[jmin]), frac_of_peak=float(s[jmin] / peak), + iY=int(jmin[0]), + iqT=int(jmin[1]), + Y=float(Yg[jmin[0]]), + qT=float(qTg[jmin[1]]), + value=float(s[jmin]), + frac_of_peak=float(s[jmin] / peak), ) negidx = np.argwhere(s < 0) if negidx.size: order = np.argsort(s[negidx[:, 0], negidx[:, 1]]) # most negative first out["neg_bins"] = [ dict( - Y=float(Yg[negidx[k, 0]]), qT=float(qTg[negidx[k, 1]]), + Y=float(Yg[negidx[k, 0]]), + qT=float(qTg[negidx[k, 1]]), value=float(s[negidx[k, 0], negidx[k, 1]]), frac_of_peak=float(s[negidx[k, 0], negidx[k, 1]] / peak), ) for k in order ] out["neg_qT_range"] = ( - float(qTg[negidx[:, 1]].min()), float(qTg[negidx[:, 1]].max()) + float(qTg[negidx[:, 1]].min()), + float(qTg[negidx[:, 1]].max()), ) out["neg_absY_max"] = float(np.abs(Yg[negidx[:, 0]]).max()) else: @@ -165,16 +194,35 @@ def spectrum_negativity(core, eff_params, gnu_params, sigma_YqT=None, locate=Tru def np_physical_report( - core, eff_params, gnu_params, sigma_YqT=None, central_neg_area=None + core, + eff_params, + gnu_params, + sigma_YqT=None, + central_neg_area=None, + np_model=None, + np_model_nu=None, ): """Combine both detectors into a postfit verdict (no printing, no raising). Returns ``{ok, issues, damp, neg, central_neg_area}``: ``ok`` the overall verdict, ``issues`` human-readable problems, ``damp``/``neg`` the raw sub-results. ``central_neg_area`` anchors the relative negativity threshold - (default: from ``core`` at λ_central). Callers format/act.""" - damp = np_damping_ok(core, eff_params, gnu_params) - neg = spectrum_negativity(core, eff_params, gnu_params, sigma_YqT=sigma_YqT) + (default: from ``core`` at λ_central, evaluated at the CARD forms — the + correct baseline regardless of override). ``np_model`` / ``np_model_nu`` + select the forms the λ point is probed under (default: the card forms; pass + the FIT forms for a ``np_model_(nu_)fit`` override fit, e.g. from + ``lambda_central.read_np_models``). Callers format/act.""" + damp = np_damping_ok( + core, eff_params, gnu_params, np_model=np_model, np_model_nu=np_model_nu + ) + neg = spectrum_negativity( + core, + eff_params, + gnu_params, + sigma_YqT=sigma_YqT, + np_model=np_model, + np_model_nu=np_model_nu, + ) if central_neg_area is None: central_neg_area = spectrum_negativity( core, @@ -252,9 +300,7 @@ def card_reco_reference(model, indata): norm = np.asarray(indata.norm, dtype=np.float64) n_reco = int(np.prod(model.reco_shape)) if norm.shape[0] < n_reco: - raise ValueError( - f"indata.norm has {norm.shape[0]} bins < reco bins {n_reco}" - ) + raise ValueError(f"indata.norm has {norm.shape[0]} bins < reco bins {n_reco}") # v1 single non-masked channel: the reco bins are the leading rows of norm. col = norm[:n_reco, model.signal_proc_idx] return col.reshape(model.reco_shape) @@ -328,7 +374,9 @@ def run_reco_guard( + f") ref_yield={d['ref_yield']:.3g}" for d in offenders[:max_list] ) - more = "" if len(offenders) <= max_list else f"\n … +{len(offenders)-max_list} more" + more = ( + "" if len(offenders) <= max_list else f"\n … +{len(offenders)-max_list} more" + ) msg = ( f"{tag} reco agreement: {len(offenders)} bin(s) exceed {threshold*100:.2f}% " f"|σ_reco(λ_c)/card_nominal − 1| (yield-weighted mean = {wmean:.3f}%). " @@ -340,8 +388,7 @@ def run_reco_guard( ) if strict: raise ValueError( - msg - + "\n(check_agreement_strict=1 → raising. Pass check_agreement=0 to " + msg + "\n(check_agreement_strict=1 → raising. Pass check_agreement=0 to " "disable, raise check_agreement_threshold, or set check_agreement_min_yield " "to ignore sparse bins.)" ) @@ -360,7 +407,9 @@ def compare_level(model_vals, ref_vals, axes_meta, label): from wremnants.postprocessing.scetlib_np.validation_plots import summarize print(f"\n{'='*70}\n{label}\n{'='*70}") - summarize(np.asarray(model_vals, np.float64), np.asarray(ref_vals, np.float64), axes_meta) + summarize( + np.asarray(model_vals, np.float64), np.asarray(ref_vals, np.float64), axes_meta + ) resid, stats = _shape_residual(model_vals, ref_vals) names = [n for n, _ in axes_meta] shape = tuple(len(e) - 1 for _, e in axes_meta) @@ -370,7 +419,9 @@ def compare_level(model_vals, ref_vals, axes_meta, label): else None ) coord_str = ( - ", ".join(f"{nm}={c}" for nm, c in zip(names, coord)) if coord is not None else "n/a" + ", ".join(f"{nm}={c}" for nm, c in zip(names, coord)) + if coord is not None + else "n/a" ) print( f"\n >> {label}: yield-weighted mean|shape−1| = " @@ -381,8 +432,13 @@ def compare_level(model_vals, ref_vals, axes_meta, label): def run_card_diagnostics( - model, indata, outdir=None, do_plots=True, ref_label_reco=None, + model, + indata, + outdir=None, + do_plots=True, + ref_label_reco=None, gen_exclude_overflow=True, + args=None, ): """Full reco + gen comparison of the model's λ_central shape to the card's references, with optional per-axis shape plots. @@ -404,9 +460,14 @@ def run_card_diagnostics( binning, so there is no truncated bin to smear.""" out = {} reco_ref = card_reco_reference(model, indata) - reco_model = np.asarray(model.sigma_reco_central, np.float64).reshape(model.reco_shape) + reco_model = np.asarray(model.sigma_reco_central, np.float64).reshape( + model.reco_shape + ) out["reco"] = compare_level( - reco_model, reco_ref, model._reco_axes_meta, "RECO σ_reco(λ_c) vs card norm[signal]" + reco_model, + reco_ref, + model._reco_axes_meta, + "RECO σ_reco(λ_c) vs card norm[signal]", ) gen_ref = card_gen_reference(model) @@ -443,28 +504,34 @@ def run_card_diagnostics( rlabel_reco = ref_label_reco or "card nominal (signal)" h_reco_m = tf_to_hist(reco_model, model._reco_axes_meta) h_reco_n = tf_to_hist(reco_ref, model._reco_axes_meta) - for ax in RECO_PROJ_AXES: + # Project onto the model's own reco axes (the fit channel's — 2D or 4D), + # not the canonical RECO_AXES: a 2D ptll-yll fit has no angular axes. + for ax in [n for n, _ in model._reco_axes_meta]: plot_ptll_ratio( - h_reco_m, h_reco_n, axis=ax, + h_reco_m, + h_reco_n, + axis=ax, out_path=os.path.join(outdir, f"reco_{ax}.png"), - ref_label=rlabel_reco, model_label=r"ParamModel $\sigma_{reco}(\lambda_c)$", - rlabel="model / card", density=True, - title=f"σ_reco(λ_c) vs card nominal — {ax}", + ref_label=rlabel_reco, + model_label=r"ParamModel $\sigma_{reco}(\lambda_c)$", + rlabel="model / card", + density=True, + args=args, ) # gen ptVGen: normalize on resolved bins but keep all bins, so the # truncated overflow shows as a step (not a pedestal on the bulk). - # Pre-scale the model and pass scale=1.0 so the legend has no "(×scale)"; - # the normalization choice lives in the title. - norm_tag = "resolved-qT norm" if gen_exclude_overflow else "global norm" + # Pre-scale the model and pass scale=1.0 so the legend has no "(×scale)". plot_ptll_ratio( tf_to_hist(gen_model * gscale, model._gen_axes_meta), tf_to_hist(gen_ref, model._gen_axes_meta), - axis="ptVGen", density=False, + axis="ptVGen", + density=False, out_path=os.path.join(outdir, "gen_ptVGen.png"), ref_label=r"card $N_{gen}$", model_label=r"ParamModel $\sigma_{gen}(\lambda_c)$", - rlabel="model / $N_{gen}$", rrange=(0.78, 1.05), - title=f"σ_gen(λ_c) vs card N_gen — ptVGen ({norm_tag})", + rlabel="model / $N_{gen}$", + rrange=(0.78, 1.05), + args=args, ) # gen absYVGen: project resolved-qT only (overflow zeroed in both) so the # rapidity shape isn't contaminated by the truncated bin. @@ -477,49 +544,14 @@ def run_card_diagnostics( plot_ptll_ratio( tf_to_hist(gm_res, model._gen_axes_meta), tf_to_hist(gn_res, model._gen_axes_meta), - axis="absYVGen", density=True, + axis="absYVGen", + density=True, out_path=os.path.join(outdir, "gen_absYVGen.png"), ref_label=r"card $N_{gen}$", model_label=r"ParamModel $\sigma_{gen}(\lambda_c)$", - rlabel="model / $N_{gen}$", rrange=(0.95, 1.05), - title="σ_gen(λ_c) vs card N_gen — |y|" + (", resolved qT" if gen_exclude_overflow else ""), + rlabel="model / $N_{gen}$", + rrange=(0.95, 1.05), + args=args, ) print(f"\n plots written under: {outdir}") return out - - -def main(argv=None): - import argparse - import time - - from rabbit.inputdata import FitInputData - - from wremnants.postprocessing.scetlib_np.param_model import SCETlibNPParamModel - - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--datacard", required=True, help="fit-input hdf5 (must carry the scetlib_np auxiliary)") - p.add_argument("--btgrid", default=None, help="SCETlib bt-grid dir (default: model's data-area copy)") - p.add_argument("--signal-proc", default="Zmumu") - p.add_argument("--outdir", default=None, help="plot output dir ('' / unset to skip plotting)") - args = p.parse_args(argv) - - print("Loading FitInputData …", flush=True) - t0 = time.time() - indata = FitInputData(args.datacard) - print(f" loaded in {time.time()-t0:.1f}s; nproc={indata.nproc}", flush=True) - - print("Constructing SCETlibNPParamModel (runs the bt integral at λ_central) …", flush=True) - t0 = time.time() - kw = dict(signal_proc=args.signal_proc, check_agreement=False) # report below, not the guard - if args.btgrid: - kw["btgrid_dir"] = args.btgrid - model = SCETlibNPParamModel(indata, **kw) - print(f" constructed in {time.time()-t0:.1f}s", flush=True) - - run_card_diagnostics( - model, indata, outdir=(args.outdir or None), do_plots=bool(args.outdir) - ) - - -if __name__ == "__main__": - main() diff --git a/wremnants/postprocessing/scetlib_np/plot_output.py b/wremnants/postprocessing/scetlib_np/plot_output.py new file mode 100644 index 000000000..2aea5b88d --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/plot_output.py @@ -0,0 +1,74 @@ +"""One save entry point for every SCETlib-NP plot. + +Wraps the two ``wums`` output helpers so each plot is written in BOTH formats and +with a provenance sidecar, instead of a bare ``fig.savefig`` (single format, no +log): + +* :func:`wums.plot_tools.save_pdf_and_png` -> ``{basename}.pdf`` + ``{basename}.png`` +* :func:`wums.output_tools.write_index_and_log` -> ``{basename}.log`` (the exact + command line, the parsed args, a timestamp, and the git hash/diff) plus the + ``index.php`` web-gallery template. The gallery globs ``*.png`` and links each + plot to its same-basename ``.log``/``.pdf``, so on the webdir "click a plot -> + read the command that made it" just works. + +Import is cheap (only ``os`` at module load); ``wums`` is imported lazily inside +:func:`save_plot`, so packaged modules can import this without paying for it. +""" + +import os + + +def save_plot(outdir, basename, fig=None, args=None, meta_info=None, dpi=None): + """Write ``fig`` to ``{outdir}/{basename}.{pdf,png}`` + a provenance log/index. + + Parameters + ---------- + outdir : str + Output directory (created if missing). Falsy -> current directory. + basename : str + Filename stem, WITHOUT extension. + fig : matplotlib Figure, optional + Figure to save; if ``None`` the current pyplot figure is used. + args : argparse.Namespace, optional + The script's parsed args, recorded in the ``.log``. The command line is + captured from ``sys.argv`` regardless, so ``None`` (e.g. env-var-driven + scripts) still logs the invocation. + meta_info : dict, optional + Extra ``key -> value`` entries appended to the ``.log`` (e.g. config that + does not live in ``args``, like env-var settings or fit inputs). + dpi : int, optional + PNG resolution. The PDF is vector, so this only affects the raster + output. ``None`` keeps the matplotlib default. + """ + from wums import output_tools, plot_tools + + outdir = outdir or "." + os.makedirs(outdir, exist_ok=True) + if dpi is not None and fig is not None: + fig.set_dpi(dpi) + plot_tools.save_pdf_and_png(outdir, basename, fig=fig) + output_tools.write_index_and_log( + outdir, + basename, + analysis_meta_info=meta_info or {}, + args=args, + ) + + +def split_outpath(out_path, default_name=None): + """Split a user ``--out``-style path into ``(outdir, basename)``. + + Accepts a file path (``dir/name.png``), a bare name, or -- when + ``default_name`` is given -- a directory (trailing slash or no suffix), in + which case ``default_name`` is appended. The returned ``basename`` has no + extension, ready for :func:`save_plot`. + """ + if default_name is not None and ( + out_path.endswith(("/", os.sep)) + or os.path.isdir(out_path) + or not os.path.splitext(out_path)[1] + ): + out_path = os.path.join(out_path, default_name) + outdir = os.path.dirname(out_path) or "." + basename = os.path.splitext(os.path.basename(out_path))[0] + return outdir, basename diff --git a/wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py b/wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py index 8dd70f777..7d7330fb9 100644 --- a/wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py +++ b/wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py @@ -11,7 +11,11 @@ runcard / else the canonical FranksVals tanh_2 default. The model is BUILT at this base (positive σ_gen, which the constructor requires); λ are evaluated on top, so params not set stay at the BASE value, not 0; - * ``--fitresult HDF5`` postfit λ (optional); + * ``--fitresult HDF5`` postfit λ (optional). Also sources the base tune from the + fitresults metadata when ``--meta-from`` is not given, and applies the fit's + ``np_model_(nu_)fit`` NUMERATOR-form override (if any) to the EVALUATION — + construction stays at the base (card/denominator) form, mirroring + ``param_model``; * ``--lambdas name=val,...`` explicit values (optional; win over the rest). Common use is ``--lambdas lambda2=0.5``, the rest staying at FranksVals; no λ_central source needed. Evaluating, unlike constructing, has no positivity @@ -46,7 +50,6 @@ """ import argparse -import os import sys import time @@ -57,283 +60,19 @@ GNU_PARAMS, parse_lambda_overrides, ) - -# btgrid default mirrors the validation scripts; the datacard is intentionally -# NOT defaulted — pass --datacard (or explicit --*-edges) for the gen edges. -BTGRID_DIR = "/scratch/submit/cms/wmass/scetlib_np/Z_COM13_CT18Z_N3p0LL_btgrid_fineall/" -Q_LO, Q_HI = 60.0, 120.0 -# Canonical FranksVals (CT18Z N3+0LL lattice λ4-bugfix) tanh_2 runcard — the -# production λ_central. Construction BASE when no base λ is sourced: the model -# must be built at a PHYSICAL tune (positive σ_gen, so the constructor's response -# guard passes), and the requested λ evaluated on top. Source of truth: a -# correction file's Nonperturbative section (file_meta_data → config → -# Nonperturbative); the LatticeNPLambda4Bugfix_FranksVals_CT18Z values. -CANONICAL_BASE = { - "eff_params": { - "np_model": "tanh_2", "lambda2": 0.4, "lambda4": 0.4, - "lambda6": 0.0, "delta_lambda2": 0.0, "lambda_inf": 1.0, - }, - "gnu_params": { - "np_model_nu": "tanh_2", "lambda2_nu": 0.15, - "lambda4_nu": 0.0, "lambda6_nu": 0.0, "lambda_inf_nu": 2.0, - }, -} - -# Built-in gen grid used when neither explicit edges nor a datacard/hdf5 source is -# given: 1-GeV ptVGen bins over [0, 40], rapidity-inclusive in a single absYVGen -# bin [0, 5] (5.0 is a TheoryCorrection absY edge, so the overlay still aligns). -DEFAULT_PTV_EDGES = np.arange(0.0, 41.0, 1.0) -DEFAULT_ABSY_EDGES = np.array([0.0, 5.0]) - -# corr-hist axis ↔ model gen axis (the TheoryCorrection _hist uses SCETlib names). -_CORR_AXIS = {"ptVGen": "qT", "absYVGen": "absY"} - - -def _parse_edges(s): - """``a,b,c,...`` -> float ndarray of bin edges.""" - return np.array([float(x) for x in s.split(",") if x.strip()], dtype=np.float64) - - -def _merge_matrix(fine_edges, coarse_edges, name="axis", tol=1e-6): - """(N_coarse, N_fine) 0/1 matrix summing fine bins into coarse bins. - - Requires every coarse edge to coincide with a fine edge (coarse is a - sub-binning of fine): exact merge, no interpolation. Fine bins whose centre - lies outside every coarse bin (e.g. qT beyond the model's ptVGen overflow - edge) get weight 0 and are dropped, matching the model. - """ - fine_edges = np.asarray(fine_edges, dtype=np.float64) - coarse_edges = np.asarray(coarse_edges, dtype=np.float64) - for e in coarse_edges: - if not np.any(np.isclose(fine_edges, e, atol=tol)): - raise SystemExit( - f"_merge_matrix[{name}]: model edge {e} is not a TheoryCorrection " - f"bin edge (its binning is not a refinement of the model grid on " - f"this axis). corr edges: {fine_edges}" - ) - centers = 0.5 * (fine_edges[:-1] + fine_edges[1:]) - W = np.zeros((coarse_edges.size - 1, fine_edges.size - 1), dtype=np.float64) - for i in range(coarse_edges.size - 1): - m = (centers >= coarse_edges[i]) & (centers <= coarse_edges[i + 1]) - W[i, m] = 1.0 - return W - - -def resolve_base_lambda(args): - """Physical BASE λ tune (eff_params/gnu_params) the model is CONSTRUCTED at. - - Priority: ``--meta-from HDF5`` > the ``--theory-corr`` file's embedded - Nonperturbative runcard > the canonical FranksVals tanh_2 default. Always a - complete physical tune (never None), so construction lands on a positive-σ_gen - point (the constructor's response guard); the requested λ are evaluated on top. - """ - from wremnants.postprocessing.scetlib_np import lambda_central as lc - - if args.meta_from: - print(f"[λ base] from hdf5 metadata {args.meta_from}") - return lc.read_lambda_central(args.meta_from) - if args.theory_corr: - import pickle - - import lz4.frame - - with lz4.frame.open(args.theory_corr) as fh: - corr = pickle.load(fh) - base = lc.extract_lambda_central( - corr, tag=os.path.basename(args.theory_corr), - proc=args.theory_corr_proc or "Z", - ) - print(f"[λ base] from the --theory-corr Nonperturbative runcard " - f"({base.get('basename')})") - return {"eff_params": base["eff_params"], "gnu_params": base["gnu_params"]} - print("[λ base] none given -> canonical FranksVals tanh_2 default") - return CANONICAL_BASE - - -def assemble_tune(base, overrides): - """Full (eff_params, gnu_params) for the EVAL point = base tune + overrides, - plus the explicitly-set names. - - ``base`` is a physical lambda_central dict (with the np_model form strings); - params not in ``overrides`` stay at the base value (NOT 0). Each override is - routed to eff or gnu by membership.""" - eff = dict(base["eff_params"]) - gnu = dict(base["gnu_params"]) - explicit = {} - for name, val in overrides.items(): - if name in EFF_PARAMS: - eff[name] = val - elif name in GNU_PARAMS: - gnu[name] = val - else: - raise SystemExit( - f"unknown λ {name!r}; valid: {list(GNU_PARAMS) + list(EFF_PARAMS)}" - ) - explicit[name] = val - return eff, gnu, explicit - - -def resolve_gen_axes(args): - """gen_axes = [(ptVGen, edges), (absYVGen, edges)], chosen per axis in order: - explicit --ptv-edges/--absy-edges, then a --gen-edges-from/--datacard hdf5, - then the built-in defaults (1-GeV ptVGen [0,40]; single absYVGen [0,5]).""" - ptv = _parse_edges(args.ptv_edges) if args.ptv_edges else None - absy = _parse_edges(args.absy_edges) if args.absy_edges else None - src = args.gen_edges_from or args.datacard - - src_axes = None - if (ptv is None or absy is None) and src: - print(f"[gen-axes] reading the scetlib_np auxiliary of {src}") - from rabbit.inputdata import FitInputData - - from wremnants.postprocessing.scetlib_np.param_model import ( - _R_info_from_auxiliary, - ) - - indata = FitInputData(src) - src_axes = { - n: np.asarray(e, dtype=np.float64) - for n, e in _R_info_from_auxiliary(indata)["gen_axes"] - } - - def pick(name, explicit, default): - if explicit is not None: - print(f"[gen-axes] {name}: explicit ({explicit.size - 1} bins)") - return explicit - if src_axes is not None and name in src_axes: - print(f"[gen-axes] {name}: from {src} ({src_axes[name].size - 1} bins)") - return src_axes[name] - print(f"[gen-axes] {name}: built-in default ({default.size - 1} bins, " - f"[{default[0]:g}, {default[-1]:g}])") - return default - - return [ - ("ptVGen", pick("ptVGen", ptv, DEFAULT_PTV_EDGES)), - ("absYVGen", pick("absYVGen", absy, DEFAULT_ABSY_EDGES)), - ] - - -def load_theory_corr_hist(path, proc=None): - """Load the ``{generator}_hist`` (SCETlib+DYTurbo) Hist from a TheoryCorrection - ``.pkl.lz4``. - - The file maps ``corr[proc][histname]``. ``proc`` defaults to the single - physics key (``meta_data`` / ``file_meta_data`` excluded); the hist is the - lone ``*_hist`` that is not ``minnlo_ref_hist`` (the prediction, not the - MiNNLO reference or the ratio). - """ - import pickle - - import lz4.frame - - with lz4.frame.open(path) as fh: - corr = pickle.load(fh) - - meta_keys = {"meta_data", "file_meta_data"} - procs = [k for k in corr.keys() if k not in meta_keys] - if proc is None: - if len(procs) != 1: - raise SystemExit( - f"--theory-corr-proc needed: {os.path.basename(path)} has procs {procs}" - ) - proc = procs[0] - elif proc not in corr: - raise SystemExit(f"proc {proc!r} not in {list(corr.keys())}") - - entry = corr[proc] - cands = [k for k in entry if k.endswith("_hist") and k != "minnlo_ref_hist"] - if len(cands) != 1: - raise SystemExit( - f"expected one {{generator}}_hist in {proc}; found {list(entry.keys())}" - ) - histname = cands[0] - - print(f"[theory-corr] {os.path.basename(path)} :: {proc} / {histname}") - return entry[histname] - - -def theory_corr_projection(h, gen_axes, plot_axis, var="pdf0", q_window=(Q_LO, Q_HI), - tol=1e-6): - """Project a TheoryCorrection ``_hist`` onto the model's ``plot_axis`` gen bins. - - Reduces the (Q, absY, qT, charge, vars) Hist to a 1-D bin-integrated σ on the - model's ``plot_axis`` edges, restricted to the model's gen-grid extent on the - OTHER axis so it covers the same phase space the model σ_gen projection does: - - 1. select the ``vars`` entry (default ``pdf0`` = central tune); - 2. sum the Q bins whose centre falls in ``q_window`` (in-range only); - 3. sum the charge axis (in-range), if present; - 4. sum the OTHER gen axis over the model's extent [0, other_max]; - 5. rebin the projection axis onto the model's ``plot_axis`` edges (model - edges must be a sub-binning of the corr hist's: qT is fine enough that - ptVGen always aligns; absY uses SCETlib's binning so absYVGen may not). - - Returns an ndarray of length ``len(plot_axis edges) - 1`` (bin-integrated σ). - """ - names = [n for n, _ in gen_axes] - if plot_axis not in names or plot_axis not in _CORR_AXIS: - raise SystemExit(f"--plot-axis {plot_axis!r} not a model gen axis {names}") - edges_by_name = {n: np.asarray(e, dtype=np.float64) for n, e in gen_axes} - other_model = names[1] if plot_axis == names[0] else names[0] - proj_corr = _CORR_AXIS[plot_axis] - other_corr = _CORR_AXIS[other_model] - - have = [a.name for a in h.axes] - for need in ("Q", proj_corr, other_corr, "vars"): - if need not in have: - raise SystemExit( - f"theory-corr hist missing {need!r} axis; has {have}" - ) - - # 1. vars selection. - vlist = list(h.axes["vars"]) - if var not in vlist: - raise SystemExit( - f"--theory-corr-var {var!r} not in corr hist vars; have {vlist}" - ) - h = h[{"vars": vlist.index(var)}] - - # 2. Q window (sum in-range bins whose centre is inside the window). - qe = np.asarray(h.axes["Q"].edges, dtype=np.float64) - qc = 0.5 * (qe[:-1] + qe[1:]) - qsel = np.where((qc >= q_window[0] - tol) & (qc <= q_window[1] + tol))[0] - if not qsel.size: - raise SystemExit(f"no Q bins in window {q_window}; corr Q edges {qe}") - h = h[{"Q": slice(int(qsel[0]), int(qsel[-1]) + 1, sum)}] - - # 3. charge sum (in-range), if a charge axis is present. - if "charge" in [a.name for a in h.axes]: - h = h[{"charge": slice(0, h.axes["charge"].size, sum)}] - - # 4. sum the OTHER axis over the model's extent [0, other_max]. Non-coinciding - # upper edge (SCETlib's absY binning): cut at the nearest corr edge, warn. - other_max = edges_by_name[other_model][-1] - oe = np.asarray(h.axes[other_corr].edges, dtype=np.float64) - oc = 0.5 * (oe[:-1] + oe[1:]) - osel = np.where(oc <= other_max + tol)[0] - if not osel.size: - raise SystemExit( - f"theory-corr {other_corr} has no bins below the model {other_model} " - f"max {other_max}; corr edges {oe}" - ) - cut_idx = int(osel[-1]) + 1 - actual_edge = oe[cut_idx] - if abs(actual_edge - other_max) > tol: - print( - f"[theory-corr] WARNING: model {other_model} max {other_max} does not " - f"coincide with a corr {other_corr} edge; summing corr up to " - f"{actual_edge} ({abs(actual_edge - other_max):.3g} off)." - ) - h = h[{other_corr: slice(0, cut_idx, sum)}] - - # 5. rebin the projection axis onto the model's plot_axis edges. - W = _merge_matrix( - np.asarray(h.axes[proj_corr].edges, dtype=np.float64), - edges_by_name[plot_axis], - name=plot_axis, - tol=tol, - ) - return W @ np.asarray(h.values(flow=False), dtype=np.float64) +from wremnants.postprocessing.scetlib_np.sigma_gen import _default_btgrid_dir + +# The λ-tune resolvers + the theory-correction reference loader/projection now +# live in the shared validation library so the CLIs share one implementation. +from wremnants.postprocessing.scetlib_np.validation.agreement import ( + Q_HI, + Q_LO, + assemble_tune, + load_theory_corr_hist, + resolve_base_lambda, + resolve_gen_axes, + theory_corr_projection, +) def _lambda_box_text(eff, gnu): @@ -350,8 +89,17 @@ def _lambda_box_text(eff, gnu): return "λ tune:\n" + "\n".join(lines) -def make_projection_plot(sigma_gen, gen_axes, axis, out_path, eff, gnu, - s_corr=None, corr_label=None): +def make_projection_plot( + sigma_gen, + gen_axes, + axis, + out_path, + eff, + gnu, + s_corr=None, + corr_label=None, + args=None, +): """Step histogram of the matched σ_gen(λ) projection onto one gen axis (default ptVGen = ptZ), summing over the other. @@ -388,7 +136,10 @@ def make_projection_plot(sigma_gen, gen_axes, axis, out_path, eff, gnu, if show_ratio: fig, (ax, axr, axd) = plt.subplots( - 3, 1, sharex=True, figsize=(7, 7.2), + 3, + 1, + sharex=True, + figsize=(7, 7.2), gridspec_kw={"height_ratios": [3, 1, 1], "hspace": 0.06}, ) else: @@ -398,14 +149,21 @@ def make_projection_plot(sigma_gen, gen_axes, axis, out_path, eff, gnu, lab = "p$_T^Z$ (ptVGen) [GeV]" if axis == "ptVGen" else axis ax.stairs(ds, edges, color="C3", lw=1.6, label="σ_gen(λ) (param model)") if s_corr is not None: - ax.stairs(s_corr / widths, edges, color="C0", lw=1.6, ls=(0, (4, 2)), - label=corr_label) + ax.stairs( + s_corr / widths, edges, color="C0", lw=1.6, ls=(0, (4, 2)), label=corr_label + ) ax.set_ylabel(r"d$\sigma_{\mathrm{gen}}$/d(" + axis + ") [a.u.]") ax.margins(x=0) ax.legend(loc="upper right", fontsize=9) ax.text( - 0.975, 0.60, _lambda_box_text(eff, gnu), transform=ax.transAxes, - ha="right", va="top", fontsize=7.5, family="monospace", + 0.975, + 0.60, + _lambda_box_text(eff, gnu), + transform=ax.transAxes, + ha="right", + va="top", + fontsize=7.5, + family="monospace", bbox=dict(boxstyle="round", facecolor="white", edgecolor="0.7", alpha=0.9), ) @@ -426,73 +184,131 @@ def make_projection_plot(sigma_gen, gen_axes, axis, out_path, eff, gnu, axd.set_xlabel(lab) axd.set_ylabel("(model − corr)\n/ d(" + axis + ")") axd.margins(x=0) - rng = (f"; model/corr [{rlo:.4f}, {rhi:.4f}]" - f"; Δ(dσ/dx) [{float(np.min(diff)):.3g}, {float(np.max(diff)):.3g}]") + rng = ( + f"; model/corr [{rlo:.4f}, {rhi:.4f}]" + f"; Δ(dσ/dx) [{float(np.min(diff)):.3g}, {float(np.max(diff)):.3g}]" + ) else: ax.set_xlabel(lab) rng = "" - out_dir = os.path.dirname(out_path) - if out_dir and not os.path.exists(out_dir): - os.makedirs(out_dir, exist_ok=True) - fig.savefig(out_path, dpi=130, bbox_inches="tight") + from wremnants.postprocessing.scetlib_np import plot_output + + outdir, basename = plot_output.split_outpath(out_path) + plot_output.save_plot(outdir, basename, fig=fig, args=args, dpi=130) plt.close(fig) - print(f"[plot] wrote {out_path} (axis={axis}, summed over {names[other]}{rng})") + print( + f"[plot] wrote {outdir}/{basename}.png(.pdf) + {basename}.log " + f"(axis={axis}, summed over {names[other]}{rng})" + ) def main(argv=None): p = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) - p.add_argument("--btgrid", default=BTGRID_DIR, help="SCETlib bT-grid directory") - p.add_argument("--datacard", default=None, - help="hdf5 fallback for the GEN EDGES (no default). λ are NOT " - "sourced from it — use --meta-from for that") + p.add_argument( + "--btgrid", default=_default_btgrid_dir(), help="SCETlib bT-grid directory" + ) + p.add_argument( + "--datacard", + default=None, + help="hdf5 fallback for the GEN EDGES (no default). λ are NOT " + "sourced from it — use --meta-from for that", + ) # λ tune: base source (optional) + functional form - p.add_argument("--meta-from", default=None, - help="hdf5 to read the base λ tune from (datacard/fitresults metadata); optional") - p.add_argument("--np-model", default=None, - help="F_eff functional-form override (default: the base tune's form — " - "tanh_2 for the canonical / --theory-corr base)") - p.add_argument("--np-model-nu", default=None, - help="γ_ν^NP functional-form override (default: the base tune's form)") + p.add_argument( + "--meta-from", + default=None, + help="hdf5 to read the base λ tune from (datacard/fitresults metadata); optional", + ) + p.add_argument( + "--np-model", + default=None, + help="F_eff functional-form override (default: the base tune's form — " + "tanh_2 for the canonical / --theory-corr base)", + ) + p.add_argument( + "--np-model-nu", + default=None, + help="γ_ν^NP functional-form override (default: the base tune's form)", + ) # λ values to evaluate at (applied on top of the base, in this order) - p.add_argument("--fitresult", default=None, - help="fitresults hdf5 to read the POSTFIT λ from (optional)") - p.add_argument("--result", default=None, help="fitresult group suffix (e.g. 'nominal')") - p.add_argument("--lambdas", default=None, - help="λ values 'name=val,...' evaluated on top of the base tune " - "(e.g. lambda2=0.5); unset params stay at the base (FranksVals " - "by default), NOT 0") + p.add_argument( + "--fitresult", + default=None, + help="fitresults hdf5 to read the POSTFIT λ from (optional)", + ) + p.add_argument( + "--result", default=None, help="fitresult group suffix (e.g. 'nominal')" + ) + p.add_argument( + "--lambdas", + default=None, + help="λ values 'name=val,...' evaluated on top of the base tune " + "(e.g. lambda2=0.5); unset params stay at the base (FranksVals " + "by default), NOT 0", + ) # gen-edge source - p.add_argument("--ptv-edges", default=None, - help="ptVGen edges 'a,b,c,...' (default: 1-GeV bins over [0,40])") - p.add_argument("--absy-edges", default=None, - help="absYVGen edges 'a,b,c,...' (default: single bin [0,5])") - p.add_argument("--gen-edges-from", default=None, - help="hdf5 whose scetlib_np auxiliary gives the gen edges " - "(default: --datacard, else the built-in defaults)") + p.add_argument( + "--ptv-edges", + default=None, + help="ptVGen edges 'a,b,c,...' (default: 1-GeV bins over [0,40])", + ) + p.add_argument( + "--absy-edges", + default=None, + help="absYVGen edges 'a,b,c,...' (default: single bin [0,5])", + ) + p.add_argument( + "--gen-edges-from", + default=None, + help="hdf5 whose scetlib_np auxiliary gives the gen edges " + "(default: --datacard, else the built-in defaults)", + ) # TheoryCorrection overlay - p.add_argument("--theory-corr", default=None, - help="TheoryCorrection .pkl.lz4 to overlay the official " - "SCETlib+DYTurbo gen distribution (its {generator}_hist)") - p.add_argument("--theory-corr-proc", default=None, - help="proc key in the corr file (default: the single physics key)") - p.add_argument("--theory-corr-var", default="pdf0", - help="vars label to read from the corr hist (default pdf0 = central)") - p.add_argument("--theory-corr-normalize", action="store_true", - help="rescale the corr curve to the model σ_gen integral " - "(shape-only comparison; default off = absolute overlay)") + p.add_argument( + "--theory-corr", + default=None, + help="TheoryCorrection .pkl.lz4 to overlay the official " + "SCETlib+DYTurbo gen distribution (its {generator}_hist)", + ) + p.add_argument( + "--theory-corr-proc", + default=None, + help="proc key in the corr file (default: the single physics key)", + ) + p.add_argument( + "--theory-corr-var", + default="pdf0", + help="vars label to read from the corr hist (default pdf0 = central)", + ) + p.add_argument( + "--theory-corr-normalize", + action="store_true", + help="rescale the corr curve to the model σ_gen integral " + "(shape-only comparison; default off = absolute overlay)", + ) # model / output p.add_argument("--q-lo", type=float, default=Q_LO) p.add_argument("--q-hi", type=float, default=Q_HI) - p.add_argument("--no-nonsingular", action="store_true", - help="resum-only σ_gen (σ_ns = 0; skips the FO inputs)") - p.add_argument("--plot", default=None, - help="optional path (e.g. .png/.pdf) to write a 1-D projection plot " - "of σ_gen(λ) [+ theory-corr overlay/ratio]; see --plot-axis") - p.add_argument("--plot-axis", default="ptVGen", choices=["ptVGen", "absYVGen"], - help="gen axis to project onto for --plot (default ptVGen = ptZ)") + p.add_argument( + "--no-nonsingular", + action="store_true", + help="resum-only σ_gen (σ_ns = 0; skips the FO inputs)", + ) + p.add_argument( + "--plot", + default=None, + help="optional path (e.g. .png/.pdf) to write a 1-D projection plot " + "of σ_gen(λ) [+ theory-corr overlay/ratio]; see --plot-axis", + ) + p.add_argument( + "--plot-axis", + default="ptVGen", + choices=["ptVGen", "absYVGen"], + help="gen axis to project onto for --plot (default ptVGen = ptZ)", + ) args = p.parse_args(argv) # ---- λ: build a PHYSICAL base tune (model CONSTRUCTED there so the @@ -502,12 +318,24 @@ def main(argv=None): import copy overrides = {} + fit_eff_form = fit_gnu_form = None if args.fitresult: + from wremnants.postprocessing.scetlib_np import lambda_central as lc from wremnants.postprocessing.scetlib_np.fitresult_lambdas import _flat_values + if not args.meta_from: + # the fitresults carries the card λ_central metadata: construct there + args.meta_from = args.fitresult pf = _flat_values(args.fitresult, which="postfit", result=args.result) overrides.update(pf) print(f"[λ] postfit from {args.fitresult}: {pf}") + # Resolved FIT forms (card form, overridden by np_model_(nu_)fit if the + # fit carried the override) — the forms the postfit λ belong to. + fit_eff_form, fit_gnu_form = lc.read_np_models(args.fitresult) + print( + f"[λ] fit forms from the fitresults: " + f"np_model={fit_eff_form}, np_model_nu={fit_gnu_form}" + ) try: overrides.update(parse_lambda_overrides(args.lambdas)) except ValueError as e: @@ -519,13 +347,24 @@ def main(argv=None): if args.np_model_nu: base["gnu_params"]["np_model_nu"] = args.np_model_nu eff, gnu, explicit = assemble_tune(base, overrides) + # EVALUATION forms: explicit --np-model(-nu) (already folded into the base) > + # the fit's numerator override > the base form. Construction keeps the base + # form (param_model's denominator/numerator split). + eval_np_model = args.np_model or fit_eff_form or base["eff_params"]["np_model"] + eval_np_model_nu = ( + args.np_model_nu or fit_gnu_form or base["gnu_params"]["np_model_nu"] + ) + eff["np_model"] = eval_np_model + gnu["np_model_nu"] = eval_np_model_nu gen_axes = resolve_gen_axes(args) from wremnants.postprocessing.scetlib_np.sigma_gen import SigmaGenModel - print("\n[core] constructing SigmaGenModel at the base tune (bt-grid integral) …", - flush=True) + print( + "\n[core] constructing SigmaGenModel at the base tune (bt-grid integral) …", + flush=True, + ) t0 = time.time() core = SigmaGenModel( btgrid_dir=args.btgrid, @@ -535,14 +374,18 @@ def main(argv=None): Q_hi=args.q_hi, include_nonsingular=not args.no_nonsingular, ) - print(f" constructed in {time.time()-t0:.1f}s; gen grid {core.gen_shape} " - f"({[n for n, _ in core.gen_axes]})") + print( + f" constructed in {time.time()-t0:.1f}s; gen grid {core.gen_shape} " + f"({[n for n, _ in core.gen_axes]})" + ) print(f"\n[λ] evaluating matched σ_gen at:") print(f" F_eff : {eff}") print(f" γ_ν^NP : {gnu}") if explicit: - print(f" (set via --lambdas/--fitresult: {explicit}; the rest stay at the base)") + print( + f" (set via --lambdas/--fitresult: {explicit}; the rest stay at the base)" + ) else: print(" (no overrides — evaluating at the base tune itself)") @@ -551,15 +394,22 @@ def main(argv=None): # tune can be inspected — it can dip negative at the lowest qT). t0 = time.time() if explicit: - sigma_gen = np.asarray(core.sigma_gen(eff, gnu).numpy(), dtype=np.float64) + sigma_gen = np.asarray( + core.sigma_gen( + eff, gnu, np_model=eval_np_model, np_model_nu=eval_np_model_nu + ).numpy(), + dtype=np.float64, + ) else: sigma_gen = np.asarray(core.sigma_gen_central.numpy(), dtype=np.float64) print(f" σ_gen computed in {time.time()-t0:.1f}s; shape {sigma_gen.shape}") n_bad = int(np.sum(sigma_gen <= 0)) if n_bad: - print(f" [warning] {n_bad}/{sigma_gen.size} σ_gen bins are non-positive at " - f"this tune (expected where the NP damping is weak, esp. low qT).") + print( + f" [warning] {n_bad}/{sigma_gen.size} σ_gen bins are non-positive at " + f"this tune (expected where the NP damping is weak, esp. low qT)." + ) print(f"\n Σ σ_gen = {sigma_gen.sum():.6g}") print(f" per-bin σ_gen : min {sigma_gen.min():.4g} max {sigma_gen.max():.4g}") @@ -574,15 +424,23 @@ def main(argv=None): # np_damping_wall.NPDampingWall regularizer. from wremnants.postprocessing.scetlib_np import param_model_diagnostics as ppd - rep = ppd.np_physical_report(core, eff, gnu) + rep = ppd.np_physical_report( + core, eff, gnu, np_model=eval_np_model, np_model_nu=eval_np_model_nu + ) damp, neg = rep["damp"], rep["neg"] print("\n NP physical-validity:") - print(f" γ_ν^NP damping : {'OK' if not damp['gamma_nu_wrong_sign'] else 'WRONG SIGN'}" - f" (max γ_ν over probe bT = {damp['gamma_nu_max']:+.3g}; must be ≤ 0)") - print(f" F_eff decays : {'OK' if not damp['F_eff_growing'] else 'GROWING (bT-integral divergence sign)'}") - print(f" native σ(qT)≥0 : neg_area_frac={neg['neg_area_frac']:.3g} " - f"(λ_central {rep['central_neg_area']:.3g}), min/peak={neg['min_over_peak']:+.3g}, " - f"n_neg_bins={neg['n_neg_bins']}") + print( + f" γ_ν^NP damping : {'OK' if not damp['gamma_nu_wrong_sign'] else 'WRONG SIGN'}" + f" (max γ_ν over probe bT = {damp['gamma_nu_max']:+.3g}; must be ≤ 0)" + ) + print( + f" F_eff decays : {'OK' if not damp['F_eff_growing'] else 'GROWING (bT-integral divergence sign)'}" + ) + print( + f" native σ(qT)≥0 : neg_area_frac={neg['neg_area_frac']:.3g} " + f"(λ_central {rep['central_neg_area']:.3g}), min/peak={neg['min_over_peak']:+.3g}, " + f"n_neg_bins={neg['n_neg_bins']}" + ) # WHERE the negativity sits (native Y × qT). The σ(qT)<0 dip is laundered by # the gen-binning, so locating it is the discriminator: negative cells inside @@ -592,32 +450,44 @@ def main(argv=None): ACCEPT_ABSY = 2.5 if neg.get("n_neg_bins") and neg.get("worst") is not None: w = neg["worst"] - print(f" worst neg cell : σ={w['value']:.4g} ({w['frac_of_peak']*100:+.1f}% of peak) " - f"at Y={w['Y']:+.3g}, qT={w['qT']:.3g} GeV") + print( + f" worst neg cell : σ={w['value']:.4g} ({w['frac_of_peak']*100:+.1f}% of peak) " + f"at Y={w['Y']:+.3g}, qT={w['qT']:.3g} GeV" + ) cells = neg.get("neg_bins") or [] in_cells = [c for c in cells if abs(c["Y"]) <= ACCEPT_ABSY] n_in, n_out = len(in_cells), len(cells) - len(in_cells) qr = neg.get("neg_qT_range") - print(f" neg-cell split : {n_in} inside |Y|≤{ACCEPT_ABSY}, {n_out} outside " - f"(|Y|≤{neg.get('neg_absY_max', float('nan')):.2g} reached); " - f"qT∈[{qr[0]:.3g}, {qr[1]:.3g}] GeV" if qr else "") + print( + f" neg-cell split : {n_in} inside |Y|≤{ACCEPT_ABSY}, {n_out} outside " + f"(|Y|≤{neg.get('neg_absY_max', float('nan')):.2g} reached); " + f"qT∈[{qr[0]:.3g}, {qr[1]:.3g}] GeV" + if qr + else "" + ) if in_cells: wi = min(in_cells, key=lambda c: c["value"]) - print(f" worst in-acc : σ={wi['value']:.4g} ({wi['frac_of_peak']*100:+.1f}% of peak) " - f"at Y={wi['Y']:+.3g}, qT={wi['qT']:.3g} GeV " - f"[the |Y|≤{ACCEPT_ABSY} pathology the fit region could see]") + print( + f" worst in-acc : σ={wi['value']:.4g} ({wi['frac_of_peak']*100:+.1f}% of peak) " + f"at Y={wi['Y']:+.3g}, qT={wi['qT']:.3g} GeV " + f"[the |Y|≤{ACCEPT_ABSY} pathology the fit region could see]" + ) n_show = min(10, len(cells)) if n_show: print(f" most-negative {n_show} cells (Y, qT[GeV], σ, %peak):") for c in cells[:n_show]: inacc = "in " if abs(c["Y"]) <= ACCEPT_ABSY else "out" - print(f" [{inacc}] Y={c['Y']:+.3g} qT={c['qT']:7.3g} " - f"σ={c['value']:+.4g} ({c['frac_of_peak']*100:+.1f}%)") + print( + f" [{inacc}] Y={c['Y']:+.3g} qT={c['qT']:7.3g} " + f"σ={c['value']:+.4g} ({c['frac_of_peak']*100:+.1f}%)" + ) if not rep["ok"]: - print(" ⚠ UNPHYSICAL NP TUNE — the differential σ(qT) is negative / the NP is " - "anti-damping.\n This is hidden in the binned σ_gen above; do not treat " - "this point as a physical prediction.") + print( + " ⚠ UNPHYSICAL NP TUNE — the differential σ(qT) is negative / the NP is " + "anti-damping.\n This is hidden in the binned σ_gen above; do not treat " + "this point as a physical prediction." + ) # ---- optional: project an official TheoryCorrection run onto the same axis. s_corr = None @@ -625,7 +495,10 @@ def main(argv=None): if args.theory_corr: h_corr = load_theory_corr_hist(args.theory_corr, args.theory_corr_proc) s_corr = theory_corr_projection( - h_corr, core.gen_axes, args.plot_axis, var=args.theory_corr_var, + h_corr, + core.gen_axes, + args.plot_axis, + var=args.theory_corr_var, q_window=(args.q_lo, args.q_hi), ) other = 1 - [n for n, _ in core.gen_axes].index(args.plot_axis) @@ -636,13 +509,23 @@ def main(argv=None): scale = sum_model / sum_corr s_corr = s_corr * scale norm_label = f", ×{scale:.4f} (shape-normalized)" - print(f"\n[theory-corr] shape-normalized to the model integral (×{scale:.5f})") + print( + f"\n[theory-corr] shape-normalized to the model integral (×{scale:.5f})" + ) corr_label = f"SCETlib+DYTurbo ({args.theory_corr_var}){norm_label}" rcorr = np.divide(s_model, s_corr, out=np.ones_like(s_model), where=s_corr != 0) - print(f"\n[theory-corr] projection onto {args.plot_axis} (var={args.theory_corr_var}):") + print( + f"\n[theory-corr] projection onto {args.plot_axis} (var={args.theory_corr_var}):" + ) print(f" Σ corr = {sum_corr:.6g}") - print(f" Σ model / Σ corr = {sum_model/sum_corr:.5f}" - + (" (== 1 after --theory-corr-normalize)" if args.theory_corr_normalize else "")) + print( + f" Σ model / Σ corr = {sum_model/sum_corr:.5f}" + + ( + " (== 1 after --theory-corr-normalize)" + if args.theory_corr_normalize + else "" + ) + ) print(f" model / corr per bin : min {rcorr.min():.4f} max {rcorr.max():.4f}") with np.printoptions(precision=4, suppress=True, linewidth=140): print(f" model/corr: {rcorr}") @@ -651,8 +534,15 @@ def main(argv=None): if args.plot: make_projection_plot( - sigma_gen, core.gen_axes, args.plot_axis, args.plot, eff, gnu, - s_corr=s_corr, corr_label=corr_label, + sigma_gen, + core.gen_axes, + args.plot_axis, + args.plot, + eff, + gnu, + s_corr=s_corr, + corr_label=corr_label, + args=args, ) return 0 diff --git a/wremnants/postprocessing/scetlib_np/validate_agreement.py b/wremnants/postprocessing/scetlib_np/validate_agreement.py new file mode 100644 index 000000000..b4d58b3c2 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/validate_agreement.py @@ -0,0 +1,339 @@ +"""Unified entry point for the SCETlib-NP param-model agreement checks (option A). + +One CLI, ``--reference {card,histmaker}``, over the two agreement references that +share the datacard-built ``SCETlibNPParamModel``: + + * ``card`` : the model's λ_central σ_reco / σ_gen vs the datacard itself + (``indata.norm[signal]`` and the ``N_gen`` auxiliary) — no + external inputs. + * ``histmaker`` : the same vs an external histmaker ``nominal`` (+ gen MC and + the Corr[var] λ-variations). + +The third reference — the official ``TheoryCorrection`` at an arbitrary λ tune, +gen-only, cardless — is its own CLI (``sigma_gen_at_lambda``): genuinely distinct +(no datacard/response, arbitrary λ, theory truth). + +The shared reference loaders/aligners live in ``validation.agreement``; the +fit-side guard + pathology detectors and the card comparison +(``run_card_diagnostics``) live in ``param_model_diagnostics`` (numpy Layer 0). + +Run inside the wmass singularity, e.g.: + + python3 -m wremnants.postprocessing.scetlib_np.validate_agreement \\ + --reference card --datacard [--outdir ] + + python3 -m wremnants.postprocessing.scetlib_np.validate_agreement \\ + --reference histmaker --datacard --histmaker \\ + [--plot-out ] [--gen-histmaker ] [--variation lambda21.0 ...] +""" + +import argparse +import sys +import time + +import numpy as np + +from rabbit.inputdata import FitInputData +from wremnants.postprocessing.scetlib_np.param_model import SCETlibNPParamModel +from wremnants.postprocessing.scetlib_np.param_model_diagnostics import ( + run_card_diagnostics, +) +from wremnants.postprocessing.scetlib_np.sigma_gen import _default_btgrid_dir +from wremnants.postprocessing.scetlib_np.validation.agreement import ( + GEN_HIST, + GEN_SAMPLE, + NOMINAL_HIST, + SIGNAL_PROC, + SIGNAL_SAMPLE, + VARIATION_HIST, + align_gen, + align_nominal, + load_gen_hist, + load_nominal, + validate_variation, +) +from wremnants.postprocessing.scetlib_np.validation_plots import ( + plot_ptll_ratio, + summarize, + tf_to_hist, +) + + +def run_card(args): + """``--reference card``: model σ_reco/σ_gen(λ_c) vs the datacard itself + (``indata.norm[signal]`` + the ``N_gen`` auxiliary). No external inputs.""" + print("Loading FitInputData …", flush=True) + t0 = time.time() + indata = FitInputData(args.datacard) + print(f" loaded in {time.time()-t0:.1f}s; nproc={indata.nproc}", flush=True) + + print( + "Constructing SCETlibNPParamModel (runs the bt integral at λ_central) …", + flush=True, + ) + t0 = time.time() + kw = dict( + signal_proc=args.signal_proc, check_agreement=False + ) # report below, not the guard + if args.btgrid: + kw["btgrid_dir"] = args.btgrid + model = SCETlibNPParamModel(indata, **kw) + print(f" constructed in {time.time()-t0:.1f}s", flush=True) + + run_card_diagnostics( + model, + indata, + outdir=(args.outdir or None), + do_plots=bool(args.outdir), + args=args, + ) + return 0 + + +def run_histmaker(args): + """``--reference histmaker``: model σ_reco/σ_gen(λ_c) + λ-variations vs an + external histmaker ``nominal`` / gen MC / Corr[var].""" + import os + + btgrid = args.btgrid or _default_btgrid_dir() + + print("=" * 70) + print("SCETlibNPParamModel vs histmaker nominal — integral validation") + print("=" * 70) + print(f" datacard : {args.datacard}") + print(f" histmaker : {args.histmaker}") + print(f" btgrid : {btgrid}") + + print("\nLoading FitInputData …") + t0 = time.time() + indata = FitInputData(args.datacard) + print(f" loaded in {time.time()-t0:.1f}s; nproc={indata.nproc}") + + print("\nConstructing SCETlibNPParamModel (runs the bT integral at λ_central) …") + t0 = time.time() + model = SCETlibNPParamModel( + indata, + btgrid_dir=btgrid, + signal_proc=args.signal_proc, + ) + print(f" constructed in {time.time()-t0:.1f}s") + print(f" λ_central eff : {model.eff_central}") + print(f" λ_central gnu : {model.gnu_central}") + print(f" reco axes : {[n for n, _ in model._reco_axes_meta]}") + print(f" reco shape : {model.reco_shape}") + + # Everything stays at the hist level: the nominal is reordered/cropped onto + # the model's reco binning (a Hist), and the model's tf output is wrapped + # into a matching Hist via tf_to_hist. + print(f"\nLoading histmaker {args.hist!r} for {args.sample!r} …") + h_nom = align_nominal( + load_nominal(args.histmaker, args.sample, args.hist), model._reco_axes_meta + ) + h_model = tf_to_hist(model.sigma_reco_central, model._reco_axes_meta) + + summarize( + h_model.values(flow=False), h_nom.values(flow=False), model._reco_axes_meta + ) + + # Optional third curve: resum-only σ_reco, derived from the matched model + # by subtracting the exposed σ_ns at gen level and re-folding through R + # (σ_gen = σ_resum + σ_ns, both cached on the model — no second model). + extra_models = None + model_label = r"ParamModel $\sigma_{reco}$" + if args.overlay_resum: + import tensorflow as tf + + print("\nDeriving resum-only σ_reco (matched − σ_ns, folded through R) …") + gen_resum = model.sigma_gen_central - model.sigma_ns + reco_resum = tf.linalg.matvec(model.R, tf.reshape(gen_resum, [-1])) + h_resum = tf_to_hist(reco_resum, model._reco_axes_meta) + scale_r = float(h_nom.values().sum() / h_resum.values().sum()) + extra_models = [(h_resum, "ParamModel resum-only", "blue", scale_r)] + model_label = r"ParamModel matched (resum$\oplus$FO)" + + if args.plot_out: + # σ_reco_central is a theory σ (the response is normalized) while nominal + # is a weighted event yield — density-normalize both for a pure shape + # comparison, so there's no ad-hoc Σ/Σ scale factor (see plot_ptll_ratio). + # Project onto each reco shape axis: ptll (-> args.plot_out) and yll + # (-> _yll); the helper sums out the remaining reco axes. + base, ext = os.path.splitext(args.plot_out) + for ax_name in ("ptll", "yll"): + out = args.plot_out if ax_name == "ptll" else f"{base}_{ax_name}{ext}" + plot_ptll_ratio( + h_model, + h_nom, + axis=ax_name, + out_path=out, + density=True, + model_label=model_label, + title="", + autorrange=0.2, + ratio_legend=False, + no_sci=True, + extra_models=extra_models, + ) + + # ---- Gen-level cross-check: σ_gen(λ_central) [the bT integral, NO response] + # vs a gen-level histmaker at the same NP tune. Tests the integral alone. + if args.gen_histmaker: + print(f"\nGen-level check: σ_gen vs {args.gen_hist!r} from {args.gen_sample!r}") + gen_arr = align_gen( + load_gen_hist(args.gen_histmaker, args.gen_sample, args.gen_hist), + model._gen_axes_meta, + ) + h_gen_mc = tf_to_hist(gen_arr, model._gen_axes_meta) + h_gen_model = tf_to_hist(model.sigma_gen_central, model._gen_axes_meta) + print(" [gen level] σ_gen(integral) vs gen MC, on the model gen grid:") + summarize( + h_gen_model.values(flow=False), + h_gen_mc.values(flow=False), + model._gen_axes_meta, + ) + if args.plot_out: + base, ext = os.path.splitext(args.plot_out) + gen_out = f"{base}_gen{ext}" + # ptVGen-projected unit-normalized (density) shape ratio model/gen MC, + # used only to size the ratio panel — the plot itself density-normalizes. + edges = np.asarray(model._gen_axes_meta[0][1], dtype=np.float64) + centers = 0.5 * (edges[:-1] + edges[1:]) + mp = h_gen_model.project("ptVGen").values(flow=False) + npj = h_gen_mc.project("ptVGen").values(flow=False) + good = npj > 0 + rp = np.where( + good, + (mp / mp.sum()) / np.where(good, npj / npj.sum(), 1.0), + np.nan, + ) + # Focus the view on the fit region (ptll fit stops at 44); the wide + # [44,100] gen-grid tail bin would otherwise squash the spectrum and + # loosen the ratio range. Set the ratio range from the visible bins. + xhi = 44.0 + vis = good & (centers < xhi) + half = float(np.nanmax(np.abs(rp[vis] - 1.0))) + rr = max(half * 1.3, 0.004) + plot_ptll_ratio( + h_gen_model, + h_gen_mc, + axis="ptVGen", + out_path=gen_out, + density=True, + model_label=r"ParamModel $\sigma_{\mathrm{gen}}$", + ref_label="histmaker gen MC", + rlabel="model / gen MC", + title="", + rrange=(1.0 - rr, 1.0 + rr), + xlim=(0.0, xhi), + ratio_legend=False, + no_sci=True, + ) + + # ---- Off-central λ-response: model rnorm vs histmaker Corr[var]/Corr[pdf0] ---- + for var_label in args.variation or []: + out_v = None + if args.plot_out: + base, ext = os.path.splitext(args.plot_out) + out_v = f"{base}_var_{var_label}{ext}" + validate_variation(model, args, var_label, out_path=out_v) + + print("\nDone.") + return 0 + + +def main(argv=None): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument( + "--reference", + required=True, + choices=["card", "histmaker"], + help="agreement reference: 'card' (the datacard itself) or 'histmaker' " + "(an external histmaker nominal + gen MC + Corr variations)", + ) + # ---- shared ---- + p.add_argument( + "--datacard", + required=True, + help="fit-input hdf5 (FitInputData + the scetlib_np auxiliary + λ_central metadata)", + ) + p.add_argument( + "--btgrid", + default=None, + help="SCETlib bT-grid dir (default: the model's data-area copy)", + ) + p.add_argument( + "--signal-proc", default=SIGNAL_PROC, help="indata signal process name" + ) + # ---- card only ---- + p.add_argument( + "--outdir", + default=None, + help="[card] plot output dir ('' / unset to skip plotting)", + ) + # ---- histmaker only ---- + p.add_argument( + "--histmaker", + default=None, + help="[histmaker] histmaker hdf5 holding the 'nominal' hist", + ) + p.add_argument( + "--sample", + default=SIGNAL_SAMPLE, + help="[histmaker] histmaker sample group for the signal", + ) + p.add_argument( + "--hist", + default=NOMINAL_HIST, + help="[histmaker] histogram name to compare against", + ) + p.add_argument( + "--plot-out", + default="", + help="[histmaker] path for the ptll-projection ratio plot ('' to skip)", + ) + p.add_argument( + "--gen-histmaker", + default="", + help="[histmaker] gen-level histmaker hdf5 (same NP tune) for the gen-level " + "cross-check (σ_gen vs gen MC, NO response matrix); '' to skip", + ) + p.add_argument( + "--gen-hist", default=GEN_HIST, help="[histmaker] gen-level hist name" + ) + p.add_argument( + "--gen-sample", + default=GEN_SAMPLE, + help="[histmaker] gen histmaker sample group", + ) + p.add_argument( + "--variation", + nargs="*", + default=[], + metavar="LABEL", + help="[histmaker] off-central λ-response check: one or more 'vars'-axis labels " + "(e.g. lambda21.0 lambda2_nu0.25 lambda41.0). For each, compare the model's " + "rnorm=σ_reco(λ)/σ_reco(λ_c) to the histmaker's Corr[var]/Corr[pdf0].", + ) + p.add_argument( + "--variation-hist", + default=VARIATION_HIST, + help="[histmaker] histmaker hist (with a 'vars' axis) holding the NP λ variations", + ) + p.add_argument( + "--overlay-resum", + action="store_true", + help="[histmaker] overlay the resum-only σ_reco (matched minus the exposed σ_ns, " + "re-folded through R) as a third curve on the ptll plot.", + ) + args = p.parse_args(argv) + + if args.reference == "card": + return run_card(args) + if not args.histmaker: + p.error("--reference histmaker requires --histmaker") + return run_histmaker(args) + + +if __name__ == "__main__": + sys.exit(main())