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/__init__.py

@@ -0,0 +1,21 @@
+"""SCETlib-NP postprocessing package.
+
+``SCETlibNPParamModel`` (and its TensorFlow / btgrid dependencies) is imported
+lazily so that lightweight submodules — e.g. :mod:`response_matrix`, used by
+setupRabbit to embed the response matrix in the datacard — can be imported
+without pulling in TensorFlow. The package-level re-export
+``wremnants.postprocessing.scetlib_np.SCETlibNPParamModel`` (used by rabbit's
+``--paramModel`` loader) still works, resolved on first access via PEP 562.
+"""
+
+__all__ = ["SCETlibNPParamModel"]
+
+
+def __getattr__(name):
+    if name == "SCETlibNPParamModel":
+        from wremnants.postprocessing.scetlib_np.param_model import (
+            SCETlibNPParamModel,
+        )
+
+        return SCETlibNPParamModel
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

wremnants/postprocessing/scetlib_np/btgrid_cache.py

@@ -0,0 +1,189 @@
+"""One-shot pickle cache for the combined SCETlib bT-grid.
+
+Assembling the bT-grid from its individual shards is slow. On the first call
+this module writes a single ``combined_btgrid.pkl`` in the btgrid directory;
+subsequent calls load that combined file directly, which is much faster.
+
+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: we look for ``*_btgrid.pkl`` inside (recursively only
+        one level via scetlib_outputs/).
+      - a glob pattern: used directly.
+
+    Returns dict with:
+        bT      : (Nbt,)
+        b_bar   : (Nbt,)
+        bins    : list of (Q, Y, qT, lep) bin centres, length Nbins
+        vars    : dict of variation index -> setting dict (copied from the
+                  first shard; all shards are expected to carry the same set)
+        I_pert  : (Nvars, Nbins, Nbt)
+        C_nu   : (Nvars, Nbins, Nbt)
+        config  : dict from the first shard (perturbative configuration the
+                  grid was generated against)
+        n_shards: int
+    """
+    if os.path.isdir(submitdir_or_glob):
+        candidates = [
+            os.path.join(submitdir_or_glob, "scetlib_outputs", "*_btgrid.pkl"),
+            os.path.join(submitdir_or_glob, "*_btgrid.pkl"),
+        ]
+    else:
+        candidates = [submitdir_or_glob]
+
+    files = []
+    for pat in candidates:
+        files = sorted(glob.glob(pat))
+        if files:
+            break
+    if not files:
+        raise FileNotFoundError(f"No btgrid shards found under {submitdir_or_glob!r}")
+
+    # First shard sets the schema; later shards must match.
+    with open(files[0], "rb") as f:
+        first = pickle.load(f)
+    if first.get("schema_version") != "bt_grid_v1":
+        raise ValueError(
+            f"Unexpected schema {first.get('schema_version')!r} in {files[0]}"
+        )
+    bT = np.asarray(first["bT"], dtype=float)
+    b_bar = np.asarray(first["b_bar"], dtype=float)
+    varis = first["vars"]
+    config = first["config"]
+    n_vars = len(varis)
+    n_bt = bT.size
+
+    # We don't know Nbins ahead of time without scanning all shards. Walk them
+    # once: build a dict of bin -> (var_idx -> (I_pert_row, C_nu_row)).
+    bin_to_data = {}
+    for path in files:
+        with open(path, "rb") as f:
+            d = pickle.load(f)
+        if d.get("schema_version") != "bt_grid_v1":
+            raise ValueError(
+                f"Mixed schema versions: {path} has {d.get('schema_version')}"
+            )
+        if d["bT"].shape != bT.shape or not np.allclose(d["bT"], bT):
+            raise ValueError(f"bT grid mismatch in {path}")
+        bins_local = d["bins"]
+        I_local = np.asarray(
+            d["I_pert"], dtype=float
+        )  # (Nvars_local, Nbins_local, Nbt)
+        C_local = np.asarray(d["C_nu"], dtype=float)
+        # Map local variation indices to the union variation order. We assume
+        # all shards share the same vars dict (true when produced by the same
+        # condor submission).
+        var_order_local = list(d["vars"].keys())
+        for b_idx, b_tup in enumerate(bins_local):
+            slot = bin_to_data.setdefault(tuple(b_tup), {})
+            for v_pos, v_idx in enumerate(var_order_local):
+                slot[v_idx] = (I_local[v_pos, b_idx], C_local[v_pos, b_idx])
+
+    var_order = list(varis.keys())
+    bins_sorted = sorted(bin_to_data.keys(), key=lambda t: (t[0], t[1], t[2]))
+    n_bins = len(bins_sorted)
+    I_pert = np.full((n_vars, n_bins, n_bt), np.nan, dtype=float)
+    C_nu = np.full((n_vars, n_bins, n_bt), np.nan, dtype=float)
+    for b_pos, b_tup in enumerate(bins_sorted):
+        per_var = bin_to_data[b_tup]
+        for v_pos, v_idx in enumerate(var_order):
+            if v_idx in per_var:
+                I_pert[v_pos, b_pos] = per_var[v_idx][0]
+                C_nu[v_pos, b_pos] = per_var[v_idx][1]
+
+    return {
+        "bT": bT,
+        "b_bar": b_bar,
+        "bins": bins_sorted,
+        "vars": varis,
+        "var_order": var_order,
+        "I_pert": I_pert,
+        "C_nu": C_nu,
+        "config": config,
+        "n_shards": len(files),
+    }
+
+
+def _shard_glob(submitdir):
+    for pat in (
+        os.path.join(submitdir, "scetlib_outputs", "*_btgrid.pkl"),
+        os.path.join(submitdir, "*_btgrid.pkl"),
+    ):
+        files = glob.glob(pat)
+        if files:
+            return files
+    return []
+
+
+def _combined_path(submitdir):
+    return os.path.join(submitdir, _COMBINED_BASENAME)
+
+
+def _cache_is_fresh(combined, shards):
+    if not os.path.exists(combined):
+        return False
+    if not shards:
+        return True  # nothing to compare against; trust the cache
+    mtime = os.path.getmtime(combined)
+    return mtime >= max(os.path.getmtime(s) for s in shards)
+
+
+def load(submitdir, rebuild=False, verbose=True):
+    """Load the combined bT-grid for ``submitdir``.
+
+    On first call (or when ``rebuild=True``, or when any shard is newer than
+    the cached combined file), assembles the shards via
+    :func:`load_btgrid_shards`, writes ``combined_btgrid.pkl``, and returns
+    the dict.
+
+    On subsequent calls, loads the pickle directly.
+    """
+    if not os.path.isdir(submitdir):
+        raise ValueError(f"{submitdir!r} is not a directory")
+
+    combined = _combined_path(submitdir)
+    shards = _shard_glob(submitdir)
+
+    if not rebuild and _cache_is_fresh(combined, shards):
+        t0 = time.time()
+        with open(combined, "rb") as f:
+            grid = pickle.load(f)
+        if verbose:
+            print(
+                f"[btgrid_cache] loaded combined pickle in {time.time()-t0:.1f}s",
+                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