SCETlib NP continuous parameter model

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(

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 <hdf5> [--outdir <dir>]` | 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 <hdf5> --histmaker <hdf5> [--plot-out <path>] [--gen-histmaker <hdf5>] [--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 <pkl.lz4> [--datacard <hdf5>] [--lambdas lambda2=0.5 …] [--fitresult <hdf5>] [--plot <path>]` | σ_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 <fitresults.hdf5> …` | 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 <hdf5>` | inspect the central NP (λ) tune carried by a datacard/correction. |
+| `python -m …scetlib_np.response_matrix <hdf5>` | load / inspect the (reco × gen) response matrix R. |
+
+### Developer validation, smoke & timing (`…scetlib_np.validation.<name>`)
+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).

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}")

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))