diff --git a/rabbit b/rabbit index 73f346aa2..948a94acf 160000 --- a/rabbit +++ b/rabbit @@ -1 +1 @@ -Subproject commit 73f346aa2fbc05f3f5823bda2588b0539fd355ae +Subproject commit 948a94acfb4add5807f901013fd5fc886b4eaee4 diff --git a/scripts/rabbit/setupRabbit.py b/scripts/rabbit/setupRabbit.py index 88a3fc3f5..71d2b2807 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, @@ -640,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", @@ -3585,6 +3591,40 @@ 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. + 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( + "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: + 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/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/__init__.py b/wremnants/postprocessing/scetlib_np/__init__.py new file mode 100644 index 000000000..6a060bab0 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/__init__.py @@ -0,0 +1,27 @@ +"""SCETlib-NP postprocessing package. + +``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", "SigmaGenModel"] + + +def __getattr__(name): + if name == "SCETlibNPParamModel": + from wremnants.postprocessing.scetlib_np.param_model import ( + SCETlibNPParamModel, + ) + + 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 new file mode 100644 index 000000000..800ebc35b --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/btgrid_cache.py @@ -0,0 +1,200 @@ +"""One-shot pickle cache for the combined SCETlib bT-grid. + +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 + grid = btgrid_cache.load(BTGRID_DIR) +""" + +import glob +import os +import pickle +import time + +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: 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 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 config 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 + + # 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: + 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 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), {}) + 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"), + 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 ``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") + + 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", + flush=True, + ) + return grid + + if not shards: + raise FileNotFoundError(f"No btgrid shards found under {submitdir!r}") + + t0 = time.time() + grid = 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 + + +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 new file mode 100644 index 000000000..a62bec23f --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/btgrid_integrate.py @@ -0,0 +1,197 @@ +"""Q integration and Y/qT rebin helpers for the SCETlib bT-grid ParamModel. + +Weight construction is numpy (once at construction time); runtime contractions +are ``tf.tensordot`` / ``tf.einsum``. + +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)``. + +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, a ``(N_target, N_source)`` Simpson weight matrix. Apply via + ``tf.tensordot``. +""" + +import numpy as np +import tensorflow as tf + +from wremnants.postprocessing.scetlib_np.btgrid_tf import ( + _as_dtype, + simpson_weights, +) +from wremnants.utilities import common as wrem_common + +# 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 + + +# ============================================================================= +# 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 ``(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``) padded with 0, via ``tf.gather`` with a + 0-padded sentinel row. + """ + sigma_flat = _as_dtype(sigma_flat) + # 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) + return tf.gather(extended, idx_safe) + + +# ============================================================================= +# Q-integration weights (arctan_Q² method, matches scetlib_run.factorize.integrate_over_Q) +# ============================================================================= + + +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. + + 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. + + 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) + 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 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 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) + 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 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) + return tf.transpose(out, perm) diff --git a/wremnants/postprocessing/scetlib_np/btgrid_tf.py b/wremnants/postprocessing/scetlib_np/btgrid_tf.py new file mode 100644 index 000000000..397c56f3c --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/btgrid_tf.py @@ -0,0 +1,569 @@ +"""TensorFlow bT-grid factorization library. + +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 + +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 + + +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 + (``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) + return tf.cast(x, dtype) + + +# ============================================================================= +# Simpson on a static 1-D non-uniform grid. +# ============================================================================= + + +def simpson_weights(x): + """Weights ``w`` such that Simpson(y, x) == sum(w * y, axis=-1). + + ``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 + if n_intervals < 1: + return np.zeros_like(x) + + if n_intervals % 2 == 1: + # 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] + 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 / GNU_MODELS (valid np_model names) come from params (imported above). + + +def _frozen_eq_zero(x): + """``x == 0`` as a gradient-frozen condition for the NP-factor masks. + + The comparison is a non-differentiable, measure-zero boundary the + surrounding ``tf.where`` never differentiates through, so freezing its input + 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 exactly zero. + + 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 + + +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) + 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": + return (1.0 + lambda2_Y * bT**2) ** 2 * tf.exp(-2.0 * lambda4 * bT**4) + + arg = (lambda2_Y + lambda4 * bT**2) * bT + + if np_model == "identity": + return tf.exp(-2.0 * bT * arg) + + # 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 = _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) + 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) + # 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, values, *, np_model_nu): + """CS-side NP rapidity anomalous dimension γ_ν^NP(bT) for fixed ``np_model_nu``. + + ``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(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 = _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 + func = tf.tanh(tf.sqrt(a)) ** 2 + elif model == "tanh_2": + func = tf.tanh(arg) + elif model == "tanh_6": + # 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": + 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 + # 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) + + +# ============================================================================= +# 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, +): + """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 in the + :mod:`param_model` module docstring. + + 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 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,) + 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; 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 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 (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) + + 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, pass into + :func:`reconstruct_batch_tf` as ``bT_J0_kernel``. + + 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) + 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) +# ============================================================================= +# +# 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 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, 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 λ. 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 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 + ------- + 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 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 + """ + 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 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-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 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 + ---------- + 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_ν. 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; 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 + 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: 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 + ) # (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/fitresult_lambdas.py b/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py new file mode 100644 index 000000000..50d00b9b6 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/fitresult_lambdas.py @@ -0,0 +1,524 @@ +"""Read SCETlib NP λ out of a rabbit fitresults HDF5 (table / curves / toys). + +OUTPUT-side companion to :mod:`lambda_central` (which reads the INPUT-side +λ_central from the upstream correction pkl). Two responsibilities, kept apart +from plotting: + + * 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, 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`` + (``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 λ 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. The ``np_model`` / ``np_model_nu`` strings +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):: + + 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 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}, +} +# 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): + """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: + 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 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 or DEFAULT_NP_MODEL, + np_model_nu or gnu_model or DEFAULT_NP_MODEL_NU, + ) + + +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 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() + 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 λ 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) + 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, + lambda6_nu=0.0007, # SCETlib NP_model_gammanu b⁶ coeff (Gamma_nu.hpp:102) + 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 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 (linearization applied per toy), so the band + 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/lambda_central.py b/wremnants/postprocessing/scetlib_np/lambda_central.py new file mode 100644 index 000000000..3d1cd9038 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/lambda_central.py @@ -0,0 +1,321 @@ +"""Central NP (lambda) parameters for the SCETlib ParamModel. + +The SCETlib correction's Nonperturbative runcard lives in the upstream +``*_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} + +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 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" + + +# ============================================================================= +# 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 (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") + if not isinstance(meta, dict): + raise KeyError("Correction pkl has no 'file_meta_data' entry.") + 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)) + return out + + +def _parse_section(npert): + """Split one Nonperturbative dict into the eff / gnu parameter groups. + + 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 = {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_PARAMS: + gnu_params[k] = float(npert.get(k, 0.0)) + return eff_params, gnu_params + + +def _select_resummed(sections): + """Pick the resummed prediction's runcard from the NP-bearing basenames. + + 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: + 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): + """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. + """ + sections = _find_nonperturbative(corr_dict) + if not sections: + raise KeyError( + f"No Nonperturbative section in correction pkl for tag={tag!r}, " + f"proc={proc!r}." + ) + 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 + ) + + +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. + + Opens the central correction pkl (``theory_corr_tags[0]``) per proc and + extracts its Nonperturbative runcard. Returns ``{proc: lambda_central}`` for + 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. + + 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, no Nonperturbative section -- not an NP correction. + continue + return out + + +# ============================================================================= +# Read 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 whichever file was handed in. + """ + 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 _fill_missing_params(lc): + """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 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( + 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." + ) + return lc + + +def read_lambda_central_from_meta(meta, proc="Z", _source=""): + """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 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 _fill_missing_params(lc_all[proc]) + if len(lc_all) == 1: + # single proc stored -- use it whatever its label + 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)})." + ) + 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 _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. + + 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 ``{tag, basename, eff_params, gnu_params, source}``. Raises if the + metadata is absent. + """ + 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 + + 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"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/np_damping_wall.py b/wremnants/postprocessing/scetlib_np/np_damping_wall.py new file mode 100644 index 000000000..a8aa00e98 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/np_damping_wall.py @@ -0,0 +1,336 @@ +"""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). 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 +# 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"} + +# 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)} + + # 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)" + ) + + # 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 + + @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") + 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 — λ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 + # 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") + 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): + 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 new file mode 100644 index 000000000..0cc0e00b2 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/np_function_plots.py @@ -0,0 +1,415 @@ +"""Plot the SCETlib NP form factors: CS γ_ν^NP(b_T) and TMD F_eff(b_T, y). + +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 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 +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 \\ + --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 \\ + --fitresult -o /tmp/np.png +""" + +import argparse +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple + +import numpy as np + +from wremnants.postprocessing.scetlib_np import btgrid_tf +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 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 + 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, gnu = split_eff_gnu(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). + + ``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, np_model=eff["np_model"]), 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"), +} + + +# 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 only the λ that drive the panel for its model.""" + if sector == "gnu": + model = lam.gnu.get("np_model_nu", "?") + active = active_params(np_model_nu=model) + labels, src, order = _GNU_LABELS, lam.gnu, GNU_PARAMS + else: + 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( + 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, 5.0), + bT_max: float = 4.0, + n_points: int = 401, + 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. + + Parameters + ---------- + series + 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 λ fill the per-panel parameter box (default: 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 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): + 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 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 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". + from wremnants.postprocessing.scetlib_np import plot_output + + fig.tight_layout() + 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 {outdir}/{basename}.png(.pdf) + {basename}.log") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +# 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, + delta_lambda2=0.0, + lambda_inf=1.0, + lambda2_nu=0.0, + lambda4_nu=0.0, + lambda6_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( + "--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 --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)") + 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=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", + type=float, + nargs="+", + default=[0.0, 2.5, 5.0], + 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): + parser = make_parser() + args = parser.parse_args(argv) + + 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.fitresult, + 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(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={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, np_model, 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, + args=args, + ) + + +if __name__ == "__main__": + main() diff --git a/wremnants/postprocessing/scetlib_np/param_model.py b/wremnants/postprocessing/scetlib_np/param_model.py new file mode 100644 index 000000000..8ccf92882 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/param_model.py @@ -0,0 +1,1193 @@ +"""SCETlibNPParamModel — continuous-λ rabbit ParamModel for SCETlib NP. + +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 +============================================================================= + +(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) + +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 + 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). + 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 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 (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). + +============================================================================= +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), 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 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 +============================================================================= + + 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 — 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). + +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): 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 +============================================================================= + + ratio(b) = σ_reco(λ; b) / σ_reco(λ_central; b) + rnorm(b, proc) = 1 + (ratio(b) − 1) · [proc is signal] (1 in every other proc) + +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 +----------------------------------------------------------------------------- + +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 metadata, where the histmaker stored +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. + +============================================================================= +Postfit Hessian / covariance (uncertainties on λ) +============================================================================= + +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``. +""" + +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_tf as fz_tf +from wremnants.postprocessing.scetlib_np import lambda_central as scetlib_lambda_central + +# 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 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, + _default_btgrid_dir, + compute_nonsingular_gen, +) + +_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 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, keeping the yield strictly > 0 (no NaN). +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. + + 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 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): + 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] + + +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. + + 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"] + ], + ) + + +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, + btgrid_dir: Optional[str] = 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, + nonsingular_fo_sing: str = _NONSING_FO_SING_DEFAULT, + nonsingular_dyturbo: str = _NONSING_DYTURBO_DEFAULT, + nonsingular_qt_cutoff: float = 1.0, + 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. + + Usage:: + --paramModel wremnants.postprocessing.scetlib_np.SCETlibNPParamModel [key=value ...] + + Parameters + ---------- + 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 + the datacard by setupRabbit from a ``mz_dilepton.py --unfolding`` + histmaker output (see :func:`_R_info_from_auxiliary` and + :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 + (``_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:`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``). + 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 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 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 + 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). + 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). + 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 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 + + if btgrid_dir is None: + btgrid_dir = _default_btgrid_dir() + + self._check_discrete_np_double_counting() + + # ---- λ_central + # 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 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" + ) + lambda_central_source = "auto-detect:indata.metadata theoryCorr" + 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.lambda_central_source = lambda_central_source + + # ---- 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) + + # ---- 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) + 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]}" + ) + R_arr = None + N_gen_arr = None + self.reco_shape = None + self._reco_axes_meta = None + 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. + 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"): + # generated fiducial yield per gen bin (pre-reco-selection). Dividing R + # 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). + 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." + ) + N_gen_arr = R_info["N_gen"] + self._reco_axes_meta = [ + (name, fit_axes[1]) + for (name, fit_axes) in zip( + [a[0] for a in R_reco_axes], + fit_reco_axes, + ) + ] + 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, + ) + 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"np_model_fit={self._np_model_fit!r} not in {sorted(fz_tf.EFF_MODELS)}" + ) + if self._np_model_nu_fit not in fz_tf.GNU_MODELS: + raise ValueError( + 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: 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_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( + 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) + 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 + # 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). + # 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 ()) + 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) + 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) 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 == 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. + # 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(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), + } + + # 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.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( + 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 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( + f"[SCETlibNPParamModel] xparamdefault overridden: {dict(zip(self._param_order, defaults))}", + flush=True, + ) + # 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 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`` — 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()} + # 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): + if p in prior_sigmas: + sigmas_arr[i] = float( + prior_sigmas[p] + ) # explicit override (may be NaN) + else: + # 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. + 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, + ) + + # ---- 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 λ, 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). + """ + + 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] + + conflicting = [s for s in syst_names if _DISCRETE_NP_SUBSTRING in s.lower()] + if not conflicting: + 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): + """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 — delegated to the physics core (Steps 1–2) + # ========================================================================= + + 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_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 + ) + + 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, + ) + + # ========================================================================= + # λ-vector helpers + # ========================================================================= + + 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 + + # ========================================================================= + # 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 + + # ========================================================================= + # 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 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) + # 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. + 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), + 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. The compact object the Hessian needs from the + fold.""" + 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. + + Shape: (N_reco, N_proc). Signal-proc column carries the per-reco-bin + 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 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 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 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. + if self._hess_st: + ratio = self._ratio_straightthrough(param, use_curvature=not self._hess_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) + # 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) + rnorm = 1.0 + (ratio_col - 1.0) * self._signal_col_mask + + return rnorm 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..fdba702a6 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/param_model_diagnostics.py @@ -0,0 +1,557 @@ +"""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. + +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). 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 +only, so this module — and the in-fit guard — stays numpy-only. +""" + +import os + +import numpy as np + +# ============================================================================= +# 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, + 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: + γ_ν^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. + ``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=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 { + "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, + 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`` 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 + 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, np_model=np_model, np_model_nu=np_model_nu + ) + 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, + 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, 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, + 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, + args=None, +): + """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) + # 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, + 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, + 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)". + 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), + args=args, + ) + # 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), + args=args, + ) + print(f"\n plots written under: {outdir}") + return out diff --git a/wremnants/postprocessing/scetlib_np/params.py b/wremnants/postprocessing/scetlib_np/params.py new file mode 100644 index 000000000..dc9cf913a --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/params.py @@ -0,0 +1,208 @@ +"""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 + +# 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") + +# 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"} + +# ---- 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 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: + out.update(EFF_MODEL_PARAMS[_EFF_MODEL_ALIASES.get(np_model, np_model)]) + if np_model_nu is not None: + 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}``. + + 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/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/response_matrix.py b/wremnants/postprocessing/scetlib_np/response_matrix.py new file mode 100644 index 000000000..31b51d96c --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/response_matrix.py @@ -0,0 +1,274 @@ +"""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. 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: + + load_R(unfolding_hdf5_path, + sample_key="Zmumu_2016PostVFP", + hist_name="nominal_prefsr_yieldsUnfolding") -> dict +""" + +import h5py +import numpy as np + +from wums import ioutils as wums_io + +# 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" +# 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. 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. + + 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) + 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 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 + 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 has_response( + unfolding_hdf5_path, + sample_key=DEFAULT_SAMPLE, + hist_name=DEFAULT_HIST, + gen_total_name=DEFAULT_GENTOTAL, +): + """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). + + 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: + 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, + 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. + + 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 + + ``ptVGen_overflow`` (default True): append the gen ptVGen overflow (true + 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: + 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", 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 (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 (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). 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). + N_gen_hist = ( + _append_axis_overflow(hg_gen, "ptVGen") + if ptVGen_overflow + else hg_gen.values(flow=False).astype(np.float64) + ) + + # 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 + 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, _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, + 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}]" + ) 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..7d7330fb9 --- /dev/null +++ b/wremnants/postprocessing/scetlib_np/sigma_gen_at_lambda.py @@ -0,0 +1,551 @@ +"""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). 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 +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 sys +import time + +import numpy as np + +from wremnants.postprocessing.scetlib_np.params import ( + EFF_PARAMS, + GNU_PARAMS, + parse_lambda_overrides, +) +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): + """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, + args=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 = "" + + 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 {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=_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 = {} + 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: + 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) + # 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, + ) + 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, 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"\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, 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']}" + ) + + # 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, + args=args, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) 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()) diff --git a/wremnants/production/datasets/dataset_tools.py b/wremnants/production/datasets/dataset_tools.py index 4feaaa07c..a75982f36 100644 --- a/wremnants/production/datasets/dataset_tools.py +++ b/wremnants/production/datasets/dataset_tools.py @@ -3,13 +3,15 @@ """ 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). +from wremnants.utilities.data_paths import getDataPath, makeFilelist from wums import logging logger = logging.child_logger(__name__) @@ -24,172 +26,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}") - - 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/production/histmaker_tools.py b/wremnants/production/histmaker_tools.py index 331fb305b..36079c904 100644 --- a/wremnants/production/histmaker_tools.py +++ b/wremnants/production/histmaker_tools.py @@ -154,6 +154,29 @@ 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 +226,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 e683395d7..70015c433 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 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 diff --git a/wremnants/utilities/io_tools/input_tools.py b/wremnants/utilities/io_tools/input_tools.py index 1f6e0e64d..d2052e649 100644 --- a/wremnants/utilities/io_tools/input_tools.py +++ b/wremnants/utilities/io_tools/input_tools.py @@ -333,10 +333,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("+") diff --git a/wremnants/utilities/styles/styles.py b/wremnants/utilities/styles/styles.py index e0a11ba33..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 @@ -272,6 +272,7 @@ def translate_html_to_latex(n): "ZmassAndWidth", "massAndWidth", "normXsecZ", + "resumNonpert", ] nuisance_grouping = { "super": [ @@ -818,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)