From 41a6bfb9f4b4b4cad206c76ef944a70ec3d44557 Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 19:28:03 +0000 Subject: [PATCH 01/15] mcstat: continuous-M de-biasing (frozen M moment + sandwich covariance) - TensorWriter.add_mc_stat_moment(M, param_names): frozen nparams^2 noise-floor matrix stored as an external-likelihood term (hess=-M -> objective -1/2 th^T M th). - Fitter: --mcStatDebias/--mcStatDebiasCov/--mcStatKfold/--covMode options; _build_mcstat_M (reconstruct M from the mcstat external term); cov_mcstat_sandwich (A^-1 + A^-1 M A^-1, analytic, no 2nd Hessian pass); linearity warning (the -1/2 th^T M th penalty only de-biases LINEAR params; rabbit's default x^2 POI transform cancels it -- use --allowNegativeParam). - rabbit_fit.py reports the sandwich covariance when continuousM+sandwich. - tests/toy_mcstat.py, tests/verify_mcstat.py: 200-bin degenerate toy; rabbit matches an independent numpy reference to 4 d.p. (de-biased point, curvature, sandwich). See RESULTS.md S9. Co-Authored-By: Claude Opus 4.8 (1M context) --- bin/rabbit_fit.py | 12 ++++++ rabbit/fitter.py | 86 ++++++++++++++++++++++++++++++++++++++++++ rabbit/parsing.py | 37 ++++++++++++++++++ rabbit/tensorwriter.py | 40 ++++++++++++++++++++ tests/toy_mcstat.py | 53 ++++++++++++++++++++++++++ tests/verify_mcstat.py | 76 +++++++++++++++++++++++++++++++++++++ 6 files changed, 304 insertions(+) create mode 100644 tests/toy_mcstat.py create mode 100644 tests/verify_mcstat.py diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 7c2fb30..a4c2cc3 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -543,6 +543,18 @@ def fit(args, fitter, ws, dofit=True): edmval, cov = fitter.edmval_cov(grad, hess) logger.info(f"edmval: {edmval}") + # Continuous-M robust (sandwich) covariance: replace the inverse + # de-biased Hessian A^-1 with Sigma = A^-1 + A^-1 M A^-1 (the bread + # A = the de-biased objective Hessian just computed = `hess`). + if ( + getattr(fitter, "mcStatDebias", "none") == "continuousM" + and getattr(fitter, "mcStatDebiasCov", "sandwich") == "sandwich" + and getattr(fitter, "mcstat_M", None) is not None + ): + cov = fitter.cov_mcstat_sandwich(hess) + logger.info("Reporting continuous-M sandwich covariance " + "(A^-1 + A^-1 M A^-1)") + ws.add_cov_hist(cov) fitter.cov.assign(cov) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index d9258ae..ad371ed 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -214,6 +214,31 @@ def __init__( self.expected_yield(), trainable=False, name="nexpnom" ) + # --- limited-MC-statistics de-biasing options + self.mcStatDebias = getattr(options, "mcStatDebias", "none") + self.mcStatDebiasCov = getattr(options, "mcStatDebiasCov", "sandwich") + self.mcStatKfold = getattr(options, "mcStatKfold", 2) + self.covMode = getattr(options, "covMode", "observed") + # Frozen noise-floor matrix M (nparams x nparams), reconstructed from any + # external term named "mcstat*" (TensorWriter.add_mc_stat_moment stores + # hess = -M, so M = -scatter(hess_dense)). Used for the continuous-M + # sandwich covariance Sigma = A^-1 + A^-1 M A^-1 (H - A = M). + self.mcstat_M = self._build_mcstat_M() + if self.mcStatDebias == "continuousM" and self.mcstat_M is None: + logger.warning( + "--mcStatDebias continuousM requested but no 'mcstat' external " + "moment term found in the input; continuous-M is inactive. Supply " + "M via TensorWriter.add_mc_stat_moment(M, param_names)." + ) + if self.mcStatDebias == "continuousM" and not self.param_model.is_linear: + logger.warning( + "--mcStatDebias continuousM: the param model is NONLINEAR " + "(e.g. the default rabbit x^2 POI transform). The -1/2 theta^T M theta " + "penalty is cancelled by Fisher growth at the shifted minimum and will " + "NOT de-bias the point or covariance (RESULTS.md S9a). Use " + "--allowNegativeParam (linear POI) or the two-half method instead." + ) + def init_fit_parms( self, param_model, @@ -789,6 +814,67 @@ def toyassign( self.x.assign(pparms.sample()) self.bbstat.randomize_postfit() + def _build_mcstat_M(self): + """Reconstruct the frozen noise-floor matrix M (nparams x nparams) from + external terms whose name starts with 'mcstat'. + + TensorWriter.add_mc_stat_moment stores the term Hessian as hess = -M + (so the objective contribution 0.5 x^T hess x = -1/2 x^T M x de-biases + the curvature). Here we scatter each term's (-hess_dense) block back into + the full nparams x nparams space, indexed by the term's parameter indices. + Returns None if no such term is present (continuous-M inactive). + """ + mcstat = [t for t in self.external_terms if str(t["name"]).startswith("mcstat")] + if not mcstat: + return None + n = int(self.x.shape[0]) + M = tf.zeros([n, n], dtype=self.indata.dtype) + for t in mcstat: + if t["hess_dense"] is None: + raise NotImplementedError( + "mcstat moment term must use a dense Hessian (sparse M not " + "yet supported for the sandwich covariance)." + ) + idx = t["indices"] + coords = tf.reshape( + tf.stack(tf.meshgrid(idx, idx, indexing="ij"), axis=-1), [-1, 2] + ) + # objective uses hess = -M, so M = -hess_dense + updates = tf.reshape(-t["hess_dense"], [-1]) + M = tf.tensor_scatter_nd_add(M, coords, updates) + return M + + def cov_mcstat_sandwich(self, A): + """Continuous-M robust (sandwich) covariance. + + Sigma = A^-1 H A^-1 with bread A = de-biased curvature (grad^2 of the + objective WITH the -1/2 theta^T M theta penalty) and meat H = the + same-sample curvature WITHOUT the penalty. Because the penalty is the + frozen quadratic -1/2 theta^T M theta, H = A + M exactly, so + + Sigma = A^-1 (A + M) A^-1 = A^-1 + A^-1 M A^-1. + + No second Hessian pass is needed. Frozen support only (the M-penalty + treats M as constant; valid in the linear-Gaussian regime where + continuous-M de-biases, RESULTS.md S1c/S9a). + + Parameters + ---------- + A : tf.Tensor, shape [npar, npar] + The de-biased objective Hessian (e.g. from loss_val_grad_hess()). + + Returns + ------- + tf.Tensor, shape [npar, npar] + The sandwich covariance. + """ + if self.mcstat_M is None: + raise RuntimeError( + "cov_mcstat_sandwich requires an mcstat moment term (none found)." + ) + Ainv = tf.linalg.inv(A) + return Ainv + Ainv @ self.mcstat_M @ Ainv + def edmval_cov(self, grad, hess): if len(self.frozen_params) > 0: # Only keep parameters that were floating in the fit diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 085ba1e..14f5bcb 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -377,6 +377,43 @@ def common_parser(): "ill-conditioned profiles for processes with very low effective stats per bin " "(e.g. mixed-sign-weight cancellations).", ) + parser.add_argument( + "--mcStatDebias", + default="none", + choices=["none", "continuousM", "twoHalf", "kfold"], + help="Limited-MC-statistics de-biasing of the POI point/curvature. " + "'continuousM' subtracts the frozen noise-floor matrix M (supplied via " + "TensorWriter.add_mc_stat_moment) from the curvature: -1/2 theta^T M theta in " + "the objective. NOTE: only de-biases parameters that enter the prediction " + "LINEARLY (use --allowNegativeParam for POIs; rabbit's default x^2 transform " + "cancels the correction). 'twoHalf'/'kfold' use cross-fit fold templates " + "(general / nonlinear-safe). Default 'none'.", + ) + parser.add_argument( + "--mcStatDebiasCov", + default="sandwich", + choices=["curvature", "sandwich"], + help="Covariance reported when --mcStatDebias != none. 'curvature' = inverse " + "de-biased Hessian A^-1 (undercovers). 'sandwich' = A^-1 H A^-1 = " + "A^-1 + A^-1 M A^-1 (covering robust covariance). Default 'sandwich'.", + ) + parser.add_argument( + "--mcStatKfold", + default=2, + type=int, + help="Number of folds R re-grouped into halves for --mcStatDebias kfold " + "(2 = the two-half case).", + ) + parser.add_argument( + "--covMode", + default="observed", + choices=["observed", "fisher"], + help="Curvature primitive used for ALL covariances (the standard inverse-Hessian " + "cov AND the de-biased sandwich bread+meat, consistently). 'observed' = the " + "autograd Hessian grad^2 L (Efron-Hinkley, data-conditional). 'fisher' = the " + "Gauss-Newton expected information J^T D J (lower-variance, finite-MC-cleaner, " + "drops the residual.d2mu term). Default 'observed' (current rabbit behaviour).", + ) parser.add_argument( "--paramModel", default=None, diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 1ba08e3..7562b2c 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -246,6 +246,46 @@ def add_process(self, h, name, channel="ch0", signal=False, variances=None): self.dict_norm[channel][name] = norm self.dict_sumw2[channel][name] = sumw2 + def add_mc_stat_moment(self, M, param_names, name="mcstat_M"): + """Frozen MC-stat noise-floor matrix for the continuous-M de-biasing. + + ``M`` is the ``nparams x nparams`` noise floor + ``M_ij = sum_b (sum_{e in b} w_e^2 s_i(e) s_j(e)) / mu_b`` accumulated in + the event loop at the expansion point (see RABBIT_MCSTAT_DESIGN.md §1). + It is added to the NLL as the de-attenuation penalty + ``-1/2 theta^T M theta`` (i.e. an external quadratic term with Hessian + ``-M``), which de-biases the central value and gives curvature + ``H_data - M``. ``param_names`` label the rows/cols of ``M`` and are + resolved against the fit parameter list at fit time. + + Parameters + ---------- + M : (n, n) array_like + Symmetric frozen noise-floor matrix (the user is responsible for + symmetry; only the de-attenuation penalty is formed here). + param_names : sequence of str + Length-``n`` parameter names indexing ``M``. + name : str + External-term label (default ``"mcstat_M"``); also flags the term as + the MC-stat moment so the fitter can build the sandwich covariance. + """ + import hist as _hist + + M = np.asarray(M, dtype=self.dtype) + names = list(param_names) + if M.shape != (len(names), len(names)): + raise ValueError( + f"add_mc_stat_moment: M shape {M.shape} incompatible with " + f"{len(names)} param_names" + ) + ax0 = _hist.axis.StrCategory(names, name="mcstat_params0") + ax1 = _hist.axis.StrCategory(names, name="mcstat_params1") + h = _hist.Hist(ax0, ax1, storage=_hist.storage.Double()) + h.view(flow=False)[...] = -M # hess = -M -> +1/2 x^T H x = -1/2 theta^T M theta + self.add_external_likelihood_term(hess=h, name=name) + # tag for the fitter's sandwich covariance + self.mcstat_moment_term = name + def add_channel(self, axes, name=None, masked=False, flow=False): if flow and masked is False: raise NotImplementedError( diff --git a/tests/toy_mcstat.py b/tests/toy_mcstat.py new file mode 100644 index 0000000..11b3e16 --- /dev/null +++ b/tests/toy_mcstat.py @@ -0,0 +1,53 @@ +"""Build the 2-process near-degenerate MC-stat toy as a rabbit input tensor and +test the continuous-M de-biasing (add_mc_stat_moment). Writes two tensors: + toy_noM.hdf5 : fluctuated finite-MC templates, no de-bias + toy_M.hdf5 : same + frozen noise-floor matrix M over the 2 POIs +A chisqFit on each should give a *larger*, de-biased POI uncertainty with M +(curvature H - M; RABBIT_MCSTAT_DESIGN §1).""" +import os +import numpy as np +import hist +from rabbit.tensorwriter import TensorWriter + +NB = 200 +A, H1, H2 = 4962.0, 5112.0, 4962.0 +rng = np.random.default_rng(20240614) + +n_flat = np.full(NB, A) +n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) +mu_true = n_flat + n_step # true total per bin + +# finite-MC (1:1) templates: one Poisson realization, sumw2 = counts (unit weights) +sw_flat = rng.poisson(n_flat).astype(float) +sw_step = rng.poisson(n_step).astype(float) +data = mu_true.copy() # Asimov data (clean curvature comparison) + +ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") + +def whist(values, variances): + h = hist.Hist(ax, storage=hist.storage.Weight()) + h.view()["value"] = values + h.view()["variance"] = variances + return h + +def build(outname, with_M): + tw = TensorWriter(sparse=False) + tw.add_channel([ax], "ch0") + tw.add_data(whist(data, data), "ch0") + tw.add_process(whist(sw_flat, sw_flat), "flat", "ch0", signal=True) + tw.add_process(whist(sw_step, sw_step), "step", "ch0", signal=True) + if with_M: + # noise floor over the 2 POIs (process-norm scores -> diagonal): + # M_jj = sum_b sumw2[b,j] / mu_b (Gaussian/chisq weighting var_b = data_b) + Mflat = np.sum(sw_flat / data) + Mstep = np.sum(sw_step / data) + M = np.diag([Mflat, Mstep]) + tw.add_mc_stat_moment(M, ["flat", "step"]) + print(f" M = diag({Mflat:.2f}, {Mstep:.2f})") + tw.write(outfolder=os.path.dirname(outname) or ".", + outfilename=os.path.basename(outname)) + +if __name__ == "__main__": + out = os.environ.get("OUT", "/tmp/claude") + print("building toy_noM.hdf5"); build(f"{out}/toy_noM.hdf5", with_M=False) + print("building toy_M.hdf5"); build(f"{out}/toy_M.hdf5", with_M=True) diff --git a/tests/verify_mcstat.py b/tests/verify_mcstat.py new file mode 100644 index 0000000..8809867 --- /dev/null +++ b/tests/verify_mcstat.py @@ -0,0 +1,76 @@ +"""Validate the continuous-M de-biasing (add_mc_stat_moment) in rabbit. + +Builds a Fitter for toy_noM and toy_M, minimizes, computes the covariance, and +prints the POI uncertainties in physical signal-strength (rnorm) space. + +KEY: the `-1/2 theta^T M theta` penalty only de-biases when the POI enters the +prediction LINEARLY. rabbit's default POI transform is rnorm=x^2 (nonlinear) which +cancels the correction; pass --allowNegativeParam (LINEAR rnorm=x) to see the +de-bias. See RESULTS.md S9a. Run with OUT= pointing at toy_{noM,M}.hdf5. +""" +import os, sys +import numpy as np +import importlib.util as _ilu + +RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") +sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) +_spec = _ilu.spec_from_file_location( + "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") +) +_rfm = _ilu.module_from_spec(_spec) +_spec.loader.exec_module(_rfm) + +from rabbit import inputdata, fitter +from rabbit.param_models import helpers as ph + + +def run(filename, linear): + argv = [filename, "-o", "/tmp/claude/verify_out", "-t", "0", + "--noBinByBinStat", "--chisqFit"] + if linear: + argv.append("--allowNegativeParam") + args = _rfm.make_parser().parse_args(argv) + indata = inputdata.FitInputData(filename, None) + param_model = ph.load_models([["Mu"]], indata, **vars(args)) + f = fitter.Fitter(indata, param_model, args, do_blinding=False) + f.defaultassign() + f.set_nobs(indata.data_obs) + f.minimize() + _, grad, hess = f.loss_val_grad_hess() + _, cov = f.edmval_cov(grad, hess) + C = np.asarray(cov) # curvature cov A^-1 + Csand = None + if f.mcstat_M is not None: + Csand = np.asarray(f.cov_mcstat_sandwich(hess)) # A^-1 + A^-1 M A^-1 + x = f.x.numpy() + rnorm = x if linear else x ** 2 # physical signal strength + J = np.diag(np.ones_like(x) if linear else 2 * x) # d rnorm / d x + Cr = J @ C @ J # curvature cov in rnorm space + Csr = None if Csand is None else J @ Csand @ J + return f.parms.astype(str), rnorm, Cr, Csr + + +if __name__ == "__main__": + out = os.environ.get("OUT", "/tmp/claude") + ddif = np.array([1.0, -1.0]) / np.sqrt(2) + for linear in (False, True): + print(f"\n=== POI {'LINEAR (--allowNegativeParam)' if linear else 'SQUARED (rabbit default)'} ===") + res = {} + for tag in ("noM", "M"): + parms, rnorm, Cr, Csr = run(f"{out}/toy_{tag}.hdf5", linear) + res[tag] = (rnorm, Cr, Csr) + extra = "" + if Csr is not None: + extra = (f" sandwich(flat)={np.sqrt(Csr[0,0]):.4f} " + f"sandwich(dif)={np.sqrt(ddif@Csr@ddif):.4f}") + print(f" {tag:3s} rnorm={np.round(rnorm, 4)} " + f"curv_rnorm(flat)={np.sqrt(Cr[0,0]):.4f} " + f"curv_rnorm(dif)={np.sqrt(ddif@Cr@ddif):.4f}{extra}") + s_no = np.sqrt(ddif @ res["noM"][1] @ ddif) + s_M = np.sqrt(ddif @ res["M"][1] @ ddif) + verdict = "DE-BIASED (sigma inflated)" if s_M > 1.2 * s_no else "no de-bias (cancelled)" + print(f" -> curvature sigma(dif): {s_no:.4f} -> {s_M:.4f} {verdict}") + if res["M"][2] is not None: + s_sand = np.sqrt(ddif @ res["M"][2] @ ddif) + print(f" -> sandwich sigma(dif): {s_sand:.4f} " + f"(>= curvature {s_M:.4f}: robust over-coverage margin)") From 4f9083770bc9cbeba27b419786604683e3147c83 Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 19:35:01 +0000 Subject: [PATCH 02/15] mcstat: two-half / k-fold cross-fit de-biasing (fold_axis + jackknife + sandwich) - TensorWriter.add_process(..., fold_axis=): slice a named MC-stat fold axis into k folds, derive the full by projecting it out, store hnorm_folds [k,nbinsfull,nproc]. Non-folded procs stored as norm_full/k per fold (MC-stat-exempt). Common k required. - inputdata: load norm_folds / mcstat_fold_k. - Fitter: precompute half templates norm_A/norm_B (=2*sum of fold halves, shared logk); _compute_yields_noBBB(templates='A'|'B'); jackknife objective L_cf = 2 L_full - 0.5 L_A - 0.5 L_B in _compute_nll_components; meat objective loss_val_grad_hess_meat + cov_twohalf_sandwich (A^-1 H A^-1). - rabbit_fit.py reports the two-half sandwich when twoHalf/kfold + sandwich. - tests/toy_twohalf.py, tests/verify_twohalf.py: rabbit matches numpy reference to 4 d.p.; two-half also de-biases the NONLINEAR (x^2) POI where continuous-M is cancelled. See RESULTS.md S9d. Co-Authored-By: Claude Opus 4.8 (1M context) --- bin/rabbit_fit.py | 19 +++--- rabbit/fitter.py | 99 ++++++++++++++++++++++++++++++-- rabbit/inputdata.py | 11 ++++ rabbit/tensorwriter.py | 124 +++++++++++++++++++++++++++++++++++++++- tests/toy_twohalf.py | 57 ++++++++++++++++++ tests/verify_twohalf.py | 90 +++++++++++++++++++++++++++++ 6 files changed, 387 insertions(+), 13 deletions(-) create mode 100644 tests/toy_twohalf.py create mode 100644 tests/verify_twohalf.py diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index a4c2cc3..77feeed 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -546,14 +546,17 @@ def fit(args, fitter, ws, dofit=True): # Continuous-M robust (sandwich) covariance: replace the inverse # de-biased Hessian A^-1 with Sigma = A^-1 + A^-1 M A^-1 (the bread # A = the de-biased objective Hessian just computed = `hess`). - if ( - getattr(fitter, "mcStatDebias", "none") == "continuousM" - and getattr(fitter, "mcStatDebiasCov", "sandwich") == "sandwich" - and getattr(fitter, "mcstat_M", None) is not None - ): - cov = fitter.cov_mcstat_sandwich(hess) - logger.info("Reporting continuous-M sandwich covariance " - "(A^-1 + A^-1 M A^-1)") + _debias = getattr(fitter, "mcStatDebias", "none") + _debiascov = getattr(fitter, "mcStatDebiasCov", "sandwich") + if _debiascov == "sandwich": + if _debias == "continuousM" and getattr(fitter, "mcstat_M", None) is not None: + cov = fitter.cov_mcstat_sandwich(hess) + logger.info("Reporting continuous-M sandwich covariance " + "(A^-1 + A^-1 M A^-1)") + elif _debias in ("twoHalf", "kfold") and getattr(fitter, "norm_A", None) is not None: + cov = fitter.cov_twohalf_sandwich(hess) + logger.info("Reporting two-half sandwich covariance " + "(A^-1 H A^-1)") ws.add_cov_hist(cov) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index ad371ed..a108130 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -224,6 +224,29 @@ def __init__( # hess = -M, so M = -scatter(hess_dense)). Used for the continuous-M # sandwich covariance Sigma = A^-1 + A^-1 M A^-1 (H - A = M). self.mcstat_M = self._build_mcstat_M() + # Two-half / k-fold cross-fit templates: precompute the half-sample + # norm tensors norm_A, norm_B [nbinsfull, nproc] from the fold + # templates (default groups = first/second half of the k folds, each + # rescaled by k/(k/2)=2 so ==norm_full). Shared logk. + self.norm_A = None + self.norm_B = None + norm_folds = getattr(self.indata, "norm_folds", None) + if norm_folds is not None: + k = int(norm_folds.shape[0]) + if k % 2 != 0: + raise NotImplementedError( + f"two-half de-biasing needs an even number of folds; got k={k}." + ) + half = k // 2 + scale = tf.constant(k / half, dtype=self.indata.dtype) # = 2 + self.norm_A = scale * tf.reduce_sum(norm_folds[:half], axis=0) + self.norm_B = scale * tf.reduce_sum(norm_folds[half:], axis=0) + if self.mcStatDebias in ("twoHalf", "kfold") and self.norm_A is None: + logger.warning( + f"--mcStatDebias {self.mcStatDebias} requested but no fold " + "templates (hnorm_folds) found in the input; two-half de-biasing " + "is inactive. Add processes with a fold_axis in the TensorWriter." + ) if self.mcStatDebias == "continuousM" and self.mcstat_M is None: logger.warning( "--mcStatDebias continuousM requested but no 'mcstat' external " @@ -844,6 +867,43 @@ def _build_mcstat_M(self): M = tf.tensor_scatter_nd_add(M, coords, updates) return M + def _compute_nll_meat(self): + """Plain full-sample NLL (no jackknife de-bias), used as the two-half + sandwich meat H = grad^2 L_full. Mirrors _compute_nll but always uses the + full templates and the ordinary per-bin ln.""" + nexpfullcentral, _, beta = self._compute_yields_with_beta( + profile=True, compute_norm=False, full=len(self.regularizers) + ) + nexp = nexpfullcentral[: self.indata.nbins] + l = self._compute_ln(nexp) + self._compute_lc() + lbeta = self._compute_lbeta(beta) + if lbeta is not None: + l = l + lbeta + lext = self._compute_external_nll() + if lext is not None: + l = l + lext + return l + + @tf.function + def loss_val_grad_hess_meat(self): + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + val = self._compute_nll_meat() + grad = t1.gradient(val, self.x) + hess = t2.jacobian(grad, self.x) + return val, grad, hess + + def cov_twohalf_sandwich(self, A): + """Two-half / k-fold robust (sandwich) covariance Sigma = A^-1 H A^-1. + + Bread A = the de-biased (jackknife) objective Hessian; meat H = the + full-sample NLL Hessian (loss_val_grad_hess_meat). covMode='observed' + (autograd Hessians). See RABBIT_MCSTAT_DESIGN.md §2c. + """ + _, _, H = self.loss_val_grad_hess_meat() + Ainv = tf.linalg.inv(A) + return Ainv @ tf.cast(H, Ainv.dtype) @ Ainv + def cov_mcstat_sandwich(self, A): """Continuous-M robust (sandwich) covariance. @@ -1573,7 +1633,10 @@ def _init_logk_scaled(self): else: self.logk = self.indata.logk * rnorm_init[..., None, None] - def _compute_yields_noBBB(self, full=True, compute_norm=True): + def _compute_yields_noBBB(self, full=True, compute_norm=True, templates="full"): + # templates: "full" uses indata.norm; "A"/"B" use the precomputed + # half-sample fold templates norm_A/norm_B (two-half de-biasing, shared + # logk). Only supported in the dense path. # full: compute yields inclduing masked channels # compute_norm: also build the dense [nbins, nproc] normcentral tensor. # In sparse mode this is expensive (forward + backward) and is only @@ -1653,14 +1716,21 @@ def _compute_yields_noBBB(self, full=True, compute_norm=True): if compute_norm: normcentral = tf.sparse.to_dense(snormnorm_sparse) else: + if templates == "A": + norm_src = self.norm_A + elif templates == "B": + norm_src = self.norm_B + else: + norm_src = self.indata.norm + if full or self.indata.nbinsmasked == 0: nbins = self.indata.nbinsfull logk = self.logk - norm = self.indata.norm + norm = norm_src else: nbins = self.indata.nbins logk = self.logk[:nbins] - norm = self.indata.norm[:nbins] + norm = norm_src[:nbins] if self.indata.symmetric_tensor: mlogk = tf.reshape( @@ -1964,7 +2034,28 @@ def _compute_nll_components(self, profile=True, full_nll=False): nexp = nexpfullcentral[: self.indata.nbins] - ln = self._compute_ln(nexp, full_nll) + if ( + self.mcStatDebias in ("twoHalf", "kfold") + and self.norm_A is not None + ): + # Cross-fit jackknife combination L_cf = 2 L_full - 1/2 L_A - 1/2 L_B. + # Since mu_bar = 1/2(n_A + n_B) = n_full exactly, this de-biases the + # point (gradient = cross-fit score) and curvature (cross-half + # Fisher) regardless of nonlinearity. Shared logk; full templates + # carry the BB-lite beta profiling (unchanged). + nexp_A = self._compute_yields_noBBB( + full=False, compute_norm=False, templates="A" + )[0][: self.indata.nbins] + nexp_B = self._compute_yields_noBBB( + full=False, compute_norm=False, templates="B" + )[0][: self.indata.nbins] + ln = ( + 2.0 * self._compute_ln(nexp, full_nll) + - 0.5 * self._compute_ln(nexp_A, full_nll) + - 0.5 * self._compute_ln(nexp_B, full_nll) + ) + else: + ln = self._compute_ln(nexp, full_nll) lc = self._compute_lc(full_nll) diff --git a/rabbit/inputdata.py b/rabbit/inputdata.py index 250a43f..2382049 100644 --- a/rabbit/inputdata.py +++ b/rabbit/inputdata.py @@ -172,6 +172,17 @@ def __init__(self, filename, pseudodata=None): self.sumw2 = tf.where(kstat == 0.0, self.sumw, self.sumw2) + # MC-stat cross-fit fold templates [k, nbinsfull, nproc] (optional; + # present only when a process was written with a fold_axis). Used by + # the two-half / k-fold de-biasing in the fitter. + if "hnorm_folds" in f.keys(): + self.norm_folds = maketensor(f["hnorm_folds"]) + self.mcstat_fold_k = int(f.attrs.get("mcstat_fold_k", + self.norm_folds.shape[0])) + else: + self.norm_folds = None + self.mcstat_fold_k = None + # compute indices for channels ibin = 0 for channel, info in self.channel_info.items(): diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 7562b2c..189d693 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -52,6 +52,12 @@ def __init__( self.dict_pseudodata = {} # [channel][pseudodata] self.dict_norm = {} # [channel][process] self.dict_sumw2 = {} # [channel][process] + # MC-stat cross-fit fold templates (two-half / k-fold de-biasing). + # dict_norm_folds[channel][process] = ndarray [k, nbinschan] (dense); + # fold_k[channel][process] = k (number of folds for that process). + # Populated only when add_process is called with fold_axis != None. + self.dict_norm_folds = {} # [channel][process] -> [k, nbinschan] + self.fold_k = {} # [channel][process] -> k self.dict_logkavg = {} # [channel][proc][syst] self.dict_logkhalfdiff = {} # [channel][proc][syst] self.dict_logkavg_indices = {} @@ -187,7 +193,13 @@ def add_pseudodata(self, h, name=None, channel="ch0"): ) self.dict_pseudodata[channel][name] = self.get_flat_values(h) - def add_process(self, h, name, channel="ch0", signal=False, variances=None): + def add_process( + self, h, name, channel="ch0", signal=False, variances=None, fold_axis=None + ): + if fold_axis is not None: + return self._add_process_folded( + h, name, channel, signal, variances, fold_axis + ) self._check_hist_and_channel(h, channel) if name in self.dict_norm[channel].keys(): @@ -246,6 +258,64 @@ def add_process(self, h, name, channel="ch0", signal=False, variances=None): self.dict_norm[channel][name] = norm self.dict_sumw2[channel][name] = sumw2 + def _add_process_folded(self, h, name, channel, signal, variances, fold_axis): + """Add a process whose histogram carries an extra MC-stat fold axis. + + The user fills `h` with an extra categorical/integer axis (named + `fold_axis`) assigning each *generated* event to one of `k` folds (all + oversampled copies of a generated event in the SAME fold; assignment + independent of the fitted observables — RABBIT_MCSTAT_DESIGN.md §2a). + The writer: + * derives the FULL template by projecting the fold axis out + (`h_full = sum_f h[fold=f]`) and registers it via the ordinary + add_process path (so sumw / sumw2 / systematics are unchanged), and + * stores the per-fold dense templates `[k, nbinschan]` for the + cross-fit (two-half / k-fold) de-biasing in the fitter. + `k` is inferred from the size of `fold_axis`. Sparse storage is not + supported for folded processes. + """ + self._check_hist_and_channel(h, channel) + if self.sparse and self._issparse(h): + raise NotImplementedError( + "fold_axis is not supported for sparse histograms." + ) + # resolve the fold axis (by name or index) and its size k + ax_names = [a.name for a in h.axes] + if isinstance(fold_axis, str): + if fold_axis not in ax_names: + raise ValueError( + f"fold_axis '{fold_axis}' not found in histogram axes {ax_names}" + ) + fold_idx = ax_names.index(fold_axis) + else: + fold_idx = int(fold_axis) + k = h.axes[fold_idx].size + if k < 2: + raise ValueError(f"fold_axis must have size >= 2, got {k}") + + flow = self.channels[channel]["flow"] + # per-fold flat templates and the projected-out full template + folds = [] + for f in range(k): + hf = h[{fold_idx: f}] + folds.append(self.get_flat_values(hf, flow)) + norm_folds = np.stack(folds, axis=0) # [k, nbinschan] + if not np.all(np.isfinite(norm_folds)): + raise RuntimeError( + f"NaN or Inf values encountered in fold templates for {name}!" + ) + h_full = h[{fold_idx: sum}] + + # register the full template through the ordinary (unfolded) path so + # sumw / sumw2 and all downstream machinery are identical to a normal + # process. variances default to the full hist's variances (= sum of + # fold variances), which is what BB-lite needs. + self.add_process( + h_full, name, channel=channel, signal=signal, variances=variances + ) + self.dict_norm_folds[channel][name] = norm_folds.astype(self.dtype) + self.fold_k[channel][name] = k + def add_mc_stat_moment(self, M, param_names, name="mcstat_M"): """Frozen MC-stat noise-floor matrix for the continuous-M de-biasing. @@ -298,6 +368,8 @@ def add_channel(self, axes, name=None, masked=False, flow=False): self.nbinschan[name] = ibins self.dict_norm[name] = {} self.dict_sumw2[name] = {} + self.dict_norm_folds[name] = {} + self.fold_k[name] = {} self.dict_beta_variations[name] = {} # add masked channels last @@ -1693,6 +1765,47 @@ def write(self, outfolder="./", outfilename="rabbit_input.hdf5", meta_data_dict= ibin += nbinschan + # --- MC-stat cross-fit fold templates (two-half / k-fold de-biasing). + # Assemble [k, nbinsfull, nproc] if any process was added with fold_axis. + # Folded processes store their k fold templates; non-folded processes + # store norm_full/k in each fold so they are identical across folds + # (MC-stat-exempt -> drop out of the de-bias). A common k across all + # folded processes is required. + any_folded = any(len(self.fold_k[c]) for c in self.channels) + norm_folds = None + if any_folded: + kset = { + k for c in self.channels for k in self.fold_k[c].values() + } + if len(kset) != 1: + raise NotImplementedError( + f"All folded processes must share a common number of folds k; " + f"got {sorted(kset)}." + ) + kfold = kset.pop() + norm_folds = np.zeros([kfold, nbinsfull, nproc], self.dtype) + ibin = 0 + for chan, chan_info in self.channels.items(): + nbinschan = self.nbinschan[chan] + for iproc, proc in enumerate(procs): + if proc not in self.dict_norm[chan]: + continue + if proc in self.dict_norm_folds[chan]: + norm_folds[:, ibin : ibin + nbinschan, iproc] = ( + self.dict_norm_folds[chan][proc] + ) + else: + # MC-stat-exempt: spread the full template evenly + full = self.dict_norm[chan][proc] + if self._issparse(full): + dense = np.zeros([nbinschan], self.dtype) + dense[full.indices] = full.data + full = dense + norm_folds[:, ibin : ibin + nbinschan, iproc] = ( + full[None, :] / kfold + ) + ibin += nbinschan + systs = self.get_systs() nsyst = len(systs) @@ -2144,6 +2257,15 @@ def create_dataset( sumw2, f, "hsumw2", maxChunkBytes=self.chunkSize ) + # MC-stat cross-fit fold templates [k, nbinsfull, nproc] (dense), only + # written when at least one process was added with a fold_axis. + if norm_folds is not None: + nbytes += h5pyutils_write.writeFlatInChunks( + norm_folds, f, "hnorm_folds", maxChunkBytes=self.chunkSize + ) + f.attrs["mcstat_fold_k"] = int(norm_folds.shape[0]) + norm_folds = None + if self.sparse: nbytes += h5pyutils_write.writeSparse( norm_sparse_indices, diff --git a/tests/toy_twohalf.py b/tests/toy_twohalf.py new file mode 100644 index 0000000..e53edba --- /dev/null +++ b/tests/toy_twohalf.py @@ -0,0 +1,57 @@ +"""Build the 2-process degenerate toy with an MC-stat FOLD axis (k=2) and write a +rabbit tensor exercising the two-half de-biasing (add_process(fold_axis=...)). + +Each process's finite-MC template is split into two independent halves (fold 0/1) +via a per-bin binomial(count, 0.5); the full template = fold0 + fold1 is derived +by the writer. Writes toy_folds.hdf5.""" +import os +import numpy as np +import hist +from rabbit.tensorwriter import TensorWriter + +NB = 200 +A, H1, H2 = 4962.0, 5112.0, 4962.0 +rng = np.random.default_rng(20240614) + +n_flat = np.full(NB, A) +n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) +mu_true = n_flat + n_step +data = mu_true.copy() # Asimov data + +# finite-MC full templates, then split into two folds per bin +sw_flat = rng.poisson(n_flat) +sw_step = rng.poisson(n_step) +flat_A = rng.binomial(sw_flat, 0.5); flat_B = sw_flat - flat_A +step_A = rng.binomial(sw_step, 0.5); step_B = sw_step - step_A + +ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") +axf = hist.axis.IntCategory([0, 1], name="mcfold") + + +def data_hist(values): + h = hist.Hist(ax, storage=hist.storage.Weight()) + h.view()["value"] = values + h.view()["variance"] = values + return h + + +def folded_hist(fold0, fold1): + h = hist.Hist(ax, axf, storage=hist.storage.Weight()) + h.view()["value"][:, 0] = fold0 + h.view()["value"][:, 1] = fold1 + h.view()["variance"][:, 0] = fold0 + h.view()["variance"][:, 1] = fold1 + return h + + +if __name__ == "__main__": + out = os.environ.get("OUT", "/tmp/claude") + tw = TensorWriter(sparse=False) + tw.add_channel([ax], "ch0") + tw.add_data(data_hist(data), "ch0") + tw.add_process(folded_hist(flat_A, flat_B), "flat", "ch0", + signal=True, fold_axis="mcfold") + tw.add_process(folded_hist(step_A, step_B), "step", "ch0", + signal=True, fold_axis="mcfold") + tw.write(outfolder=out, outfilename="toy_folds.hdf5") + print(f"wrote {out}/toy_folds.hdf5 (k=2 folds for flat, step)") diff --git a/tests/verify_twohalf.py b/tests/verify_twohalf.py new file mode 100644 index 0000000..cc75b61 --- /dev/null +++ b/tests/verify_twohalf.py @@ -0,0 +1,90 @@ +"""Validate rabbit's two-half (jackknife) MC-stat de-biasing against an +independent numpy reference, on the folded degenerate toy (toy_folds.hdf5). + +Compares: standard fit (no debias), two-half curvature A^-1, two-half sandwich +A^-1 H A^-1. Uses a LINEAR POI (--allowNegativeParam) so the closed-form numpy +reference applies. Run with OUT= pointing at toy_folds.hdf5.""" +import os, sys +import numpy as np +import importlib.util as _ilu + +RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") +sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) +_spec = _ilu.spec_from_file_location( + "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") +) +_rfm = _ilu.module_from_spec(_spec) +_spec.loader.exec_module(_rfm) +from rabbit import inputdata, fitter +from rabbit.param_models import helpers as ph + +ddif = np.array([1.0, -1.0]) / np.sqrt(2) + + +def rabbit_fit(filename, debias): + argv = [filename, "-o", "/tmp/claude/verify_out", "-t", "0", + "--noBinByBinStat", "--chisqFit", "--allowNegativeParam", + "--mcStatDebias", debias] + args = _rfm.make_parser().parse_args(argv) + indata = inputdata.FitInputData(filename, None) + pm = ph.load_models([["Mu"]], indata, **vars(args)) + f = fitter.Fitter(indata, pm, args, do_blinding=False) + f.defaultassign(); f.set_nobs(indata.data_obs); f.minimize() + _, grad, hess = f.loss_val_grad_hess() + _, cov = f.edmval_cov(grad, hess) + curv = np.asarray(cov) + sand = None + if debias in ("twoHalf", "kfold"): + sand = np.asarray(f.cov_twohalf_sandwich(hess)) + return f.x.numpy(), curv, sand + + +def numpy_ref(): + # regenerate the SAME folds as tests/toy_twohalf.py + NB = 200; A, H1, H2 = 4962.0, 5112.0, 4962.0 + rng = np.random.default_rng(20240614) + n_flat = np.full(NB, A) + n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) + data = (n_flat + n_step).astype(float) + sw_flat = rng.poisson(n_flat); sw_step = rng.poisson(n_step) + flat_A = rng.binomial(sw_flat, 0.5); flat_B = sw_flat - flat_A + step_A = rng.binomial(sw_step, 0.5); step_B = sw_step - step_A + Tf = np.stack([sw_flat, sw_step], 1).astype(float) # full + TA = 2 * np.stack([flat_A, step_A], 1).astype(float) # half A (x2) + TB = 2 * np.stack([flat_B, step_B], 1).astype(float) # half B (x2) + V = data + + def H(T): return np.einsum('bi,bj,b->ij', T, T, 1 / V) + def b(T): return np.einsum('bi,b,b->i', T, 1 / V, data) + Hf = H(Tf) + A_cf = 2 * Hf - 0.5 * H(TA) - 0.5 * H(TB) # jackknife Hessian + b_cf = 2 * b(Tf) - 0.5 * b(TA) - 0.5 * b(TB) + r_std = np.linalg.solve(Hf, b(Tf)) + r_cf = np.linalg.solve(A_cf, b_cf) + Cstd = np.linalg.inv(Hf) + Ccurv = np.linalg.inv(A_cf) + Csand = Ccurv @ Hf @ Ccurv + return r_std, Cstd, r_cf, Ccurv, Csand + + +if __name__ == "__main__": + out = os.environ.get("OUT", "/tmp/claude") + fn = f"{out}/toy_folds.hdf5" + + r_std, Cstd, r_cf, Ccurv, Csand = numpy_ref() + print("NUMPY REFERENCE:") + print(f" std r={np.round(r_std,4)} sigma(dif)={np.sqrt(ddif@Cstd@ddif):.4f}") + print(f" cf r={np.round(r_cf,4)} curv(dif)={np.sqrt(ddif@Ccurv@ddif):.4f} " + f"sandwich(dif)={np.sqrt(ddif@Csand@ddif):.4f}") + + xr_std, Cr_std, _ = rabbit_fit(fn, "none") + xr_cf, Cr_curv, Cr_sand = rabbit_fit(fn, "twoHalf") + print("\nRABBIT:") + print(f" std r={np.round(xr_std,4)} sigma(dif)={np.sqrt(ddif@Cr_std@ddif):.4f}") + print(f" cf r={np.round(xr_cf,4)} curv(dif)={np.sqrt(ddif@Cr_curv@ddif):.4f} " + f"sandwich(dif)={np.sqrt(ddif@Cr_sand@ddif):.4f}") + + ok = (np.allclose(xr_cf, r_cf, atol=1e-3) + and np.isclose(np.sqrt(ddif@Cr_curv@ddif), np.sqrt(ddif@Ccurv@ddif), atol=1e-3) + and np.isclose(np.sqrt(ddif@Cr_sand@ddif), np.sqrt(ddif@Csand@ddif), atol=1e-3)) + print(f"\n rabbit == numpy reference: {ok}") From 8b878bc43313df8eff1462ba5ed381b92617bfd2 Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 19:56:04 +0000 Subject: [PATCH 03/15] mcstat: --covMode fisher (Gauss-Newton) + refined linearity warning - fisher_curvature / fisher_curvature_full / _fisher_core: Gauss-Newton expected information F = H_obj - H_ln_obs + sum_i coeff_i J_i^T D J_i (D=1/V chisq, data_cov_inv covFit, 1/nexp Poisson). For the jackknife objective the GN data part is exactly the cross-half Fisher F_ch. Drops the residual*d2nexp term, leaves constraint/external/BB-lite curvature intact. - rabbit_fit.py: --covMode fisher uses the GN curvature as the bread for the standard cov AND the sandwich (continuous-M and two-half); cov_twohalf_sandwich meat follows covMode too. - Validated: linear -> fisher==observed (5.8e-11); nonlinear -> differs (drops residual). CLI runs both modes end-to-end. - Refined continuous-M linearity warning: _mcstat_M_params_linear() fires only when M actually touches nonlinear params (squared POI / log_normal syst / non- chisq); silent for the linear-Gaussian case where continuous-M is exact. Co-Authored-By: Claude Opus 4.8 (1M context) --- bin/rabbit_fit.py | 25 ++++--- rabbit/fitter.py | 166 ++++++++++++++++++++++++++++++++++++++++------ 2 files changed, 164 insertions(+), 27 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 77feeed..05119cf 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -540,23 +540,32 @@ def fit(args, fitter, ws, dofit=True): if not args.noEDM and not args.noHessian: # compute the covariance matrix and estimated distance to minimum _, grad, hess = fitter.loss_val_grad_hess() - edmval, cov = fitter.edmval_cov(grad, hess) + + # --covMode fisher: replace the observed-Hessian curvature with the + # Gauss-Newton expected information (J^T D J + nondata), used as the + # bread for the (de-biased) covariance everywhere. Applies to the + # standard cov and the sandwich bread consistently. + _covmode = getattr(fitter, "covMode", "observed") + bread = fitter.fisher_curvature(hess) if _covmode == "fisher" else hess + + edmval, cov = fitter.edmval_cov(grad, bread) logger.info(f"edmval: {edmval}") - # Continuous-M robust (sandwich) covariance: replace the inverse - # de-biased Hessian A^-1 with Sigma = A^-1 + A^-1 M A^-1 (the bread - # A = the de-biased objective Hessian just computed = `hess`). + # Robust (sandwich) covariance for the de-biased fit. Bread = the + # (fisher- or observed-) curvature `bread`; continuous-M uses the + # analytic meat H = A + M (Sigma = A^-1 + A^-1 M A^-1), two-half uses + # the full-sample meat (both in --covMode). _debias = getattr(fitter, "mcStatDebias", "none") _debiascov = getattr(fitter, "mcStatDebiasCov", "sandwich") if _debiascov == "sandwich": if _debias == "continuousM" and getattr(fitter, "mcstat_M", None) is not None: - cov = fitter.cov_mcstat_sandwich(hess) + cov = fitter.cov_mcstat_sandwich(bread) logger.info("Reporting continuous-M sandwich covariance " - "(A^-1 + A^-1 M A^-1)") + f"(A^-1 + A^-1 M A^-1, covMode={_covmode})") elif _debias in ("twoHalf", "kfold") and getattr(fitter, "norm_A", None) is not None: - cov = fitter.cov_twohalf_sandwich(hess) + cov = fitter.cov_twohalf_sandwich(bread, covMode=_covmode) logger.info("Reporting two-half sandwich covariance " - "(A^-1 H A^-1)") + f"(A^-1 H A^-1, covMode={_covmode})") ws.add_cov_hist(cov) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index a108130..e49a3a9 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -247,20 +247,31 @@ def __init__( "templates (hnorm_folds) found in the input; two-half de-biasing " "is inactive. Add processes with a fold_axis in the TensorWriter." ) - if self.mcStatDebias == "continuousM" and self.mcstat_M is None: - logger.warning( - "--mcStatDebias continuousM requested but no 'mcstat' external " - "moment term found in the input; continuous-M is inactive. Supply " - "M via TensorWriter.add_mc_stat_moment(M, param_names)." - ) - if self.mcStatDebias == "continuousM" and not self.param_model.is_linear: - logger.warning( - "--mcStatDebias continuousM: the param model is NONLINEAR " - "(e.g. the default rabbit x^2 POI transform). The -1/2 theta^T M theta " - "penalty is cancelled by Fisher growth at the shifted minimum and will " - "NOT de-bias the point or covariance (RESULTS.md S9a). Use " - "--allowNegativeParam (linear POI) or the two-half method instead." - ) + if self.mcStatDebias == "continuousM": + if self.mcstat_M is None: + logger.warning( + "--mcStatDebias continuousM requested but no 'mcstat' external " + "moment term found in the input; continuous-M is inactive. Supply " + "M via TensorWriter.add_mc_stat_moment(M, param_names)." + ) + elif not self._mcstat_M_params_linear(): + # The -1/2 theta^T M theta penalty only de-biases parameters that + # enter the prediction LINEARLY. If any parameter that M touches + # is nonlinear (e.g. the default x^2 POI transform, or log-normal + # systematics), the data Fisher grows as the -M theta gradient + # shifts the minimum and cancels the -M curvature subtraction + # (RESULTS.md S9a) -> no de-bias of the point or covariance for + # those parameters. Two-half de-biasing has no such restriction. + logger.warning( + "--mcStatDebias continuousM: the supplied M acts on parameters " + "that enter the prediction NONLINEARLY (e.g. the default rabbit " + "x^2 POI transform [use --allowNegativeParam for a linear POI], " + "or log_normal systematics). The -1/2 theta^T M theta penalty is " + "cancelled by Fisher growth at the shifted minimum and will NOT " + "de-bias the point or covariance for those parameters " + "(RESULTS.md S9a). Use a linear parametrization for the de-biased " + "parameters, or the two-half method (--mcStatDebias twoHalf)." + ) def init_fit_parms( self, @@ -837,6 +848,45 @@ def toyassign( self.x.assign(pparms.sample()) self.bbstat.randomize_postfit() + def _mcstat_M_params_linear(self): + """True iff every parameter the mcstat moment M touches enters the + prediction/NLL linearly, so that -1/2 theta^T M theta de-biases exactly + (constant curvature, no minimum-shift cancellation; RESULTS.md §9a). + + A parameter is "linear" here when: + * POI (index < npoi): the param model is linear (allowNegativeParam, + i.e. rnorm = x rather than the default rnorm = x^2), and + * systematic: the systematics enter additively (systematic_type + 'normal') with a symmetric tensor (the asymmetric interpolation is + nonlinear), and + the per-bin NLL is Gaussian (chisqFit) so its curvature does not depend + on the fit point (a Poisson l'' = nobs/nexp^2 varies with nexp(theta) + even for a linear nexp). Returns True only if ALL touched parameters meet + the relevant condition. + """ + if self.mcstat_M is None: + return True + if not self.chisqFit: + return False # Poisson/other: curvature is theta-dependent + npoi = self.param_model.npoi + mcstat_terms = [ + t for t in self.external_terms if str(t["name"]).startswith("mcstat") + ] + touched = set() + for t in mcstat_terms: + touched.update(int(i) for i in t["indices"].numpy()) + for i in touched: + if i < npoi: + if not self.param_model.is_linear: + return False + else: + if not ( + self.indata.systematic_type == "normal" + and self.indata.symmetric_tensor + ): + return False + return True + def _build_mcstat_M(self): """Reconstruct the frozen noise-floor matrix M (nparams x nparams) from external terms whose name starts with 'mcstat'. @@ -893,14 +943,92 @@ def loss_val_grad_hess_meat(self): hess = t2.jacobian(grad, self.x) return val, grad, hess - def cov_twohalf_sandwich(self, A): + def _ln_terms_for_fisher(self, debiased): + """List of (nexp[:nbins], coeff) summed in the active ln term. + + debiased=True and a two-half/k-fold debias active -> the jackknife + combination [(full,2),(A,-1/2),(B,-1/2)] (its Gauss-Newton data part is + exactly the cross-half Fisher F_ch). Otherwise the plain full term + [(full,1)]. Yields are taken through the BB-lite profiling (full) so the + Fisher is consistent with the observed Hessian it replaces.""" + nbins = self.indata.nbins + nexp_full = self._compute_yields_with_beta( + profile=True, compute_norm=False, full=len(self.regularizers) + )[0][:nbins] + if ( + debiased + and self.mcStatDebias in ("twoHalf", "kfold") + and self.norm_A is not None + ): + nexp_A = self._compute_yields_noBBB( + full=False, compute_norm=False, templates="A" + )[0][:nbins] + nexp_B = self._compute_yields_noBBB( + full=False, compute_norm=False, templates="B" + )[0][:nbins] + return [(nexp_full, 2.0), (nexp_A, -0.5), (nexp_B, -0.5)] + return [(nexp_full, 1.0)] + + def _fisher_core(self, hess_obj, debiased): + """Gauss-Newton (expected-information) curvature corresponding to the + observed-Hessian `hess_obj`: + + F = hess_obj - H_ln_obs + F_data + + where H_ln_obs is the observed Hessian of the (active) data term and + F_data = sum_i coeff_i J_i^T D J_i is its Gauss-Newton replacement + (J_i = d nexp_i / d x, D = ell''(nexp_i): 1/V for chisq, data_cov_inv + for covarianceFit, 1/nexp for Poisson expected info). This drops the + residual * d2nexp piece (the only data-term difference) while leaving + the constraint / external / BB-lite curvature in hess_obj untouched, so + the two modes stay consistent (RABBIT_MCSTAT_DESIGN.md §2c/§2d).""" + with tf.GradientTape(persistent=True) as t2: + with tf.GradientTape(persistent=True) as t1: + terms = self._ln_terms_for_fisher(debiased) + ln = tf.add_n([c * self._compute_ln(n) for n, c in terms]) + g = t1.gradient(ln, self.x) + H_ln = t2.jacobian(g, self.x) + + F = tf.zeros_like(hess_obj) + for n, c in terms: + J = t1.jacobian(n, self.x) # [nbins, nparams] + if self.covarianceFit: + JT_Cinv = tf.matmul(J, self.data_cov_inv, transpose_a=True) + Fterm = tf.matmul(JT_Cinv, J) + else: + D = (1.0 / self.varnobs) if self.chisqFit else (1.0 / n) + Fterm = tf.einsum("bi,b,bj->ij", J, D, J) + F = F + c * Fterm + del t1, t2 + return hess_obj - H_ln + F + + @tf.function + def fisher_curvature(self, hess_obj): + """Gauss-Newton curvature of the ACTIVE objective (jackknife if a + two-half debias is on; else the full term). Used as the sandwich bread + and as the standard-covariance curvature under --covMode fisher.""" + return self._fisher_core(hess_obj, True) + + @tf.function + def fisher_curvature_full(self, hess_meat): + """Gauss-Newton curvature of the FULL (undebiased) objective. Used as the + sandwich meat under --covMode fisher.""" + return self._fisher_core(hess_meat, False) + + def cov_twohalf_sandwich(self, A, covMode="observed"): """Two-half / k-fold robust (sandwich) covariance Sigma = A^-1 H A^-1. - Bread A = the de-biased (jackknife) objective Hessian; meat H = the - full-sample NLL Hessian (loss_val_grad_hess_meat). covMode='observed' - (autograd Hessians). See RABBIT_MCSTAT_DESIGN.md §2c. + Bread A = the de-biased (jackknife) curvature; meat H = the full-sample + curvature. Both follow `covMode` (don't mix): 'observed' = autograd + Hessians (A = grad^2 L_cf, H = grad^2 L_full); 'fisher' = Gauss-Newton + (A = F_ch + nondata, H = F_full + nondata). See RABBIT_MCSTAT_DESIGN.md + §2c. The passed `A` must already be in the requested mode. """ - _, _, H = self.loss_val_grad_hess_meat() + _, _, H_obs = self.loss_val_grad_hess_meat() + if covMode == "fisher": + H = self.fisher_curvature_full(H_obs) + else: + H = H_obs Ainv = tf.linalg.inv(A) return Ainv @ tf.cast(H, Ainv.dtype) @ Ainv From a7039b7cab0d5eb91de3872472e7fe20a0c411aa Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 20:05:58 +0000 Subject: [PATCH 04/15] mcstat: split-logk fold path (add_systematic fold_axis) for two-half - TensorWriter.add_systematic(h, ..., fold_axis=): split a systematic's logk per fold (per-fold logk vs per-fold nominal), register the full syst by projecting the fold axis out. Assemble/write hlogk_folds [k,nbinsfull,nproc, nsyst] (folded systs split; others replicated so logk_A==logk_B). Minimal scope: dense, symmetric tensor, single-hist (mirror=True). - inputdata: load logk_folds. - Fitter: build logk_A/logk_B (k=2: per-fold logk = per-half logk, same rnorm_init scaling as shared logk); _compute_yields_noBBB(templates='A'|'B') uses them. - Validated (tests/toy_splitlogk.py): a systematic reusing the nominal MC sample has its template noise de-biased by split logk (tilt sigma 0.0705 -> 0.0907) where shared logk does not (0.0705). See RESULTS.md S9f. Co-Authored-By: Claude Opus 4.8 (1M context) --- rabbit/fitter.py | 38 +++++++++++- rabbit/inputdata.py | 8 +++ rabbit/tensorwriter.py | 129 +++++++++++++++++++++++++++++++++++++++++ tests/toy_splitlogk.py | 64 ++++++++++++++++++++ 4 files changed, 237 insertions(+), 2 deletions(-) create mode 100644 tests/toy_splitlogk.py diff --git a/rabbit/fitter.py b/rabbit/fitter.py index e49a3a9..87f5ac7 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -241,6 +241,37 @@ def __init__( scale = tf.constant(k / half, dtype=self.indata.dtype) # = 2 self.norm_A = scale * tf.reduce_sum(norm_folds[:half], axis=0) self.norm_B = scale * tf.reduce_sum(norm_folds[half:], axis=0) + + # Split-logk halves (optional): when systematics were written with a + # fold_axis (hlogk_folds present) each half uses its own logk so the + # systematic-template noise is de-biased too; otherwise both halves share + # self.logk (the dominant nominal noise is already de-biased via norm^A/B). + # Per-fold logk == per-half logk only for k=2 (log-ratio, not additive), + # and the x2 norm rescale leaves the ratio unchanged. + self.logk_A = self.logk + self.logk_B = self.logk + logk_folds = getattr(self.indata, "logk_folds", None) + if logk_folds is not None: + if int(logk_folds.shape[0]) != 2: + raise NotImplementedError( + "split-logk (hlogk_folds) is currently supported only for k=2 " + "folds (per-fold logk == per-half logk only when k=2)." + ) + if not self.indata.symmetric_tensor: + raise NotImplementedError("split-logk requires a symmetric tensor.") + # apply the same constant rnorm_init scaling as _init_logk_scaled + if self.indata.systematic_type == "normal" and self.param_model.nparams > 0: + rnorm_init = tf.broadcast_to( + self.param_model.compute( + self.param_model.xparamdefault, full=True + ), + [self.indata.nbinsfull, self.indata.nproc], + ) + self.logk_A = logk_folds[0] * rnorm_init[..., None] + self.logk_B = logk_folds[1] * rnorm_init[..., None] + else: + self.logk_A = logk_folds[0] + self.logk_B = logk_folds[1] if self.mcStatDebias in ("twoHalf", "kfold") and self.norm_A is None: logger.warning( f"--mcStatDebias {self.mcStatDebias} requested but no fold " @@ -1846,18 +1877,21 @@ def _compute_yields_noBBB(self, full=True, compute_norm=True, templates="full"): else: if templates == "A": norm_src = self.norm_A + logk_src = self.logk_A elif templates == "B": norm_src = self.norm_B + logk_src = self.logk_B else: norm_src = self.indata.norm + logk_src = self.logk if full or self.indata.nbinsmasked == 0: nbins = self.indata.nbinsfull - logk = self.logk + logk = logk_src norm = norm_src else: nbins = self.indata.nbins - logk = self.logk[:nbins] + logk = logk_src[:nbins] norm = norm_src[:nbins] if self.indata.symmetric_tensor: diff --git a/rabbit/inputdata.py b/rabbit/inputdata.py index 2382049..8a2ed48 100644 --- a/rabbit/inputdata.py +++ b/rabbit/inputdata.py @@ -183,6 +183,14 @@ def __init__(self, filename, pseudodata=None): self.norm_folds = None self.mcstat_fold_k = None + # Split-logk fold tensor [k, nbinsfull, nproc, nsyst] (optional; + # present only when a systematic was written with a fold_axis). When + # absent, two-half de-biasing shares one logk across folds. + if "hlogk_folds" in f.keys(): + self.logk_folds = maketensor(f["hlogk_folds"]) + else: + self.logk_folds = None + # compute indices for channels ibin = 0 for channel, info in self.channel_info.items(): diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 189d693..0bc190b 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -58,6 +58,9 @@ def __init__( # Populated only when add_process is called with fold_axis != None. self.dict_norm_folds = {} # [channel][process] -> [k, nbinschan] self.fold_k = {} # [channel][process] -> k + # Split-logk fold templates (only for systematics added with fold_axis): + # dict_logkavg_folds[channel][process][syst] = [k, nbinschan]. + self.dict_logkavg_folds = {} # [channel][process][syst] -> [k, nbinschan] self.dict_logkavg = {} # [channel][proc][syst] self.dict_logkhalfdiff = {} # [channel][proc][syst] self.dict_logkavg_indices = {} @@ -370,6 +373,7 @@ def add_channel(self, axes, name=None, masked=False, flow=False): self.dict_sumw2[name] = {} self.dict_norm_folds[name] = {} self.fold_k[name] = {} + self.dict_logkavg_folds[name] = {} self.dict_beta_variations[name] = {} # add masked channels last @@ -1234,11 +1238,21 @@ def add_systematic( add_to_data_covariance=False, as_difference=False, syst_axes=None, + fold_axis=None, **kargs, ): """ h: either a single histogram with the systematic variation if mirror=True or a list of two histograms with the up and down variation as_difference: if True, interpret the histogram values as the difference with respect to the nominal (i.e. the absolute variation is norm + h) + fold_axis: optional name of an MC-stat fold axis carried by `h` (the SAME + axis used for add_process). When given, the systematic's logk is + SPLIT per fold (stored as hlogk_folds) so its template noise is + also de-biased by the two-half cross-fit; the full systematic is + registered by projecting the fold axis out. Omit (default) to + share one logk across folds (recommended; the dominant + nominal-template noise is already de-biased via norm^A.norm^B). + Minimal support: dense, single-hist (mirror=True), no other + extra syst axes. syst_axes: optional list of axis names in h that represent independent systematics. If None (default) and h is a hist-like object with axes beyond the channel, the extra axes are auto-detected and each bin combination becomes a separate @@ -1246,6 +1260,12 @@ def add_systematic( to disable auto-detection. """ + if fold_axis is not None: + return self._add_systematic_folded( + h, name, process, channel, kfactor, mirror, symmetrize, + add_to_data_covariance, as_difference, syst_axes, fold_axis, **kargs + ) + # Fast batched path for SparseHist multi-systematic input. Conditions: # - extra (systematic) axes are present # - input is a single SparseHist (not an asymmetric pair) @@ -1368,6 +1388,73 @@ def add_systematic( var_name_out, add_to_data_covariance=add_to_data_covariance, **kargs ) + def _add_systematic_folded( + self, h, name, process, channel, kfactor, mirror, symmetrize, + add_to_data_covariance, as_difference, syst_axes, fold_axis, **kargs + ): + """Add a systematic whose variation histogram carries an MC-stat fold + axis, SPLITTING logk per fold (RABBIT_MCSTAT_DESIGN.md §2a). The full + systematic is registered by projecting the fold axis out; per-fold + logkavg [k, nbinschan] is stored for the two-half cross-fit so the + systematic's own template noise is de-biased. + + Minimal support: dense histogram, single variation (mirror=True), and the + fold axis the only extra axis beyond the channel. The process must have + been added with the SAME fold_axis (so per-fold nominals exist).""" + if self.sparse and self._issparse(h): + raise NotImplementedError("fold_axis systematics: sparse not supported.") + if isinstance(h, (list, tuple)) or not mirror: + raise NotImplementedError( + "fold_axis systematics: only the symmetric single-histogram " + "(mirror=True) case is supported." + ) + if process not in self.dict_norm_folds[channel]: + raise RuntimeError( + f"add_systematic(fold_axis=...): process '{process}' has no fold " + f"templates; add it with the same fold_axis first." + ) + norm_folds = self.dict_norm_folds[channel][process] # [k, nbinschan] + k = norm_folds.shape[0] + ax_names = [a.name for a in h.axes] + if isinstance(fold_axis, str): + if fold_axis not in ax_names: + raise ValueError( + f"fold_axis '{fold_axis}' not found in axes {ax_names}" + ) + fold_idx = ax_names.index(fold_axis) + else: + fold_idx = int(fold_axis) + if h.axes[fold_idx].size != k: + raise ValueError( + f"systematic fold axis size {h.axes[fold_idx].size} != process k={k}" + ) + + flow = self.channels[channel]["flow"] + systematic_type = "normal" if add_to_data_covariance else self.systematic_type + + # per-fold logkavg vs the per-fold nominal + logk_folds = np.zeros_like(norm_folds) + for f in range(k): + syst_f = self.get_flat_values(h[{fold_idx: f}], flow=flow) + norm_f = norm_folds[f] + if as_difference: + syst_f = norm_f + syst_f + logk_folds[f] = self.get_logk( + syst_f, norm_f, kfactor, systematic_type=systematic_type + ) + + # register the FULL systematic by projecting the fold axis out + h_full = h[{fold_idx: sum}] + self.add_systematic( + h_full, name, process, channel, kfactor=kfactor, mirror=mirror, + symmetrize=symmetrize, add_to_data_covariance=add_to_data_covariance, + as_difference=as_difference, syst_axes=syst_axes, **kargs + ) + # store the per-fold logk under the resolved (full) systematic name(s) + if process not in self.dict_logkavg_folds[channel]: + self.dict_logkavg_folds[channel][process] = {} + self.dict_logkavg_folds[channel][process][name] = logk_folds.astype(self.dtype) + def add_beta_variations( self, h, @@ -2056,6 +2143,26 @@ def write(self, outfolder="./", outfilename="rabbit_input.hdf5", meta_data_dict= if self.has_beta_variations: beta_variations = np.zeros([nbinsfull, nbins, nproc], self.dtype) + # Split-logk fold tensor [k, nbinsfull, nproc, nsyst]: built only when + # at least one systematic was added with a fold_axis. Folded systs get + # their per-fold logk; all other (bin,proc,syst) entries replicate the + # full logk/k-equivalent so logk_A == logk_B for them (shared, no extra + # de-bias). Requires a symmetric tensor. + any_folded_syst = any( + len(self.dict_logkavg_folds[c].get(p, {})) + for c in self.channels + for p in self.dict_logkavg_folds[c] + ) + logk_folds = None + if any_folded_syst: + if not self.symmetric_tensor: + raise NotImplementedError( + "fold_axis systematics require a symmetric tensor." + ) + kset = {k for c in self.channels for k in self.fold_k[c].values()} + kfold = kset.pop() + logk_folds = np.zeros([kfold, nbinsfull, nproc, nsyst], self.dtype) + for chan in self.channels.keys(): nbinschan = self.nbinschan[chan] dict_norm_chan = self.dict_norm[chan] @@ -2070,10 +2177,26 @@ def write(self, outfolder="./", outfilename="rabbit_input.hdf5", meta_data_dict= dict_logkavg_proc = self.dict_logkavg[chan][proc] dict_logkhalfdiff_proc = self.dict_logkhalfdiff[chan][proc] + logkavg_folds_proc = ( + self.dict_logkavg_folds[chan].get(proc, {}) + if logk_folds is not None + else {} + ) for isyst, syst in enumerate(systs): if syst not in dict_logkavg_proc.keys(): continue + if logk_folds is not None: + if syst in logkavg_folds_proc: + logk_folds[:, ibin : ibin + nbinschan, iproc, isyst] = ( + logkavg_folds_proc[syst] + ) + else: + # shared across folds: replicate the full logk + logk_folds[:, ibin : ibin + nbinschan, iproc, isyst] = ( + dict_logkavg_proc[syst][None, :] + ) + if self.symmetric_tensor: logk[ibin : ibin + nbinschan, iproc, isyst] = ( dict_logkavg_proc[syst] @@ -2297,6 +2420,12 @@ def create_dataset( ) logk = None + if logk_folds is not None: + nbytes += h5pyutils_write.writeFlatInChunks( + logk_folds, f, "hlogk_folds", maxChunkBytes=self.chunkSize + ) + logk_folds = None + if self.has_beta_variations: nbytes += h5pyutils_write.writeFlatInChunks( beta_variations, f, "hbetavariations", maxChunkBytes=self.chunkSize diff --git a/tests/toy_splitlogk.py b/tests/toy_splitlogk.py new file mode 100644 index 0000000..2997283 --- /dev/null +++ b/tests/toy_splitlogk.py @@ -0,0 +1,64 @@ +"""Small toy with a FOLDED process and a FOLDED systematic, to exercise the +split-logk two-half path (add_systematic(fold_axis=...)). The systematic's +up-variation is built from each fold's OWN events (a per-fold reweight), so its +logk carries MC noise that differs between folds -> split logk is meaningful. + +Writes toy_splitlogk.hdf5 (folded syst) and toy_sharedlogk.hdf5 (same nominal +folds, systematic added WITHOUT fold_axis = shared logk) for comparison.""" +import os +import numpy as np +import hist +from rabbit.tensorwriter import TensorWriter + +NB = 50 +rng = np.random.default_rng(7) +nom = np.linspace(800.0, 1200.0, NB) # smooth nominal shape +data = nom.copy() # Asimov + +# two independent half-samples of the nominal (per-bin binomial split) +full = rng.poisson(nom) +A = rng.binomial(full, 0.5); B = full - A + +# systematic: a tilt reweight applied to each fold's OWN events -> per-fold +# up template = fold * (1 + 0.15*(x-0.5)) with fold-specific Poisson noise +x = np.linspace(0, 1, NB) +w = 1.0 + 0.15 * (x - 0.5) +upA = rng.poisson(np.maximum(A * w, 0.0)) +upB = rng.poisson(np.maximum(B * w, 0.0)) + +ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") +axf = hist.axis.IntCategory([0, 1], name="mcfold") + + +def dhist(v): + h = hist.Hist(ax, storage=hist.storage.Weight()) + h.view()["value"] = v; h.view()["variance"] = v + return h + + +def fhist(f0, f1): + h = hist.Hist(ax, axf, storage=hist.storage.Weight()) + h.view()["value"][:, 0] = f0; h.view()["value"][:, 1] = f1 + h.view()["variance"][:, 0] = f0; h.view()["variance"][:, 1] = f1 + return h + + +def build(outname, split): + tw = TensorWriter(sparse=False) + tw.add_channel([ax], "ch0") + tw.add_data(dhist(data), "ch0") + tw.add_process(fhist(A, B), "sig", "ch0", signal=True, fold_axis="mcfold") + if split: + tw.add_systematic(fhist(upA, upB), "tilt", "sig", "ch0", fold_axis="mcfold") + else: + # shared logk: full up template (folds summed), no fold_axis + tw.add_systematic(dhist(upA + upB), "tilt", "sig", "ch0") + tw.write(outfolder=os.path.dirname(outname) or ".", + outfilename=os.path.basename(outname)) + + +if __name__ == "__main__": + out = os.environ.get("OUT", "/tmp/claude") + build(f"{out}/toy_splitlogk.hdf5", split=True) + build(f"{out}/toy_sharedlogk.hdf5", split=False) + print(f"wrote {out}/toy_splitlogk.hdf5 (split) and toy_sharedlogk.hdf5 (shared)") From 69147e8655ee0bc656fd6f3500c7c2fe6184a245 Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 20:37:57 +0000 Subject: [PATCH 05/15] mcstat: BB-lite (bin-by-bin stat) compatibility for both methods Examined composing the de-biasing with Barlow-Beeston-lite; both diverged with BB-lite ON. Root causes + fixes: - two-half + BB-lite was UNBOUNDED: beta-profiling flattens L_full's POI curvature while -1/2 L_A -1/2 L_B keep full curvature -> A=2H_full,bb-1/2H_A -1/2H_B indefinite. Fix: apply the full-sample profiled beta (per-bin factor nexp/nexp_full_raw) to the half predictions too, so all three terms are consistently profiled and A ~ H_full,bb > 0 (bounded/PD). Composes sensibly: BB-lite inflates the baseline, two-half de-biases on top. - continuous-M + BB-lite: M must use the BB-inflated variance (mu+sumw2) in its denominator (else H-M non-PD, divergence). Fitter warns; add_mc_stat_moment documents it. Works with the correct M. - Both --covMode modes work with BB-lite (fisher ~ observed). tests/verify_bblite.py asserts PD convergence + de-bias for both methods with BB-lite on. No-BB paths unchanged (verify_mcstat / verify_twohalf exact). See RESULTS.md S9g. Co-Authored-By: Claude Opus 4.8 (1M context) --- rabbit/fitter.py | 40 +++++++++++++++--- rabbit/tensorwriter.py | 10 +++++ tests/verify_bblite.py | 94 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 139 insertions(+), 5 deletions(-) create mode 100644 tests/verify_bblite.py diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 87f5ac7..2d47282 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -303,6 +303,15 @@ def __init__( "(RESULTS.md S9a). Use a linear parametrization for the de-biased " "parameters, or the two-half method (--mcStatDebias twoHalf)." ) + if self.mcstat_M is not None and self.bbstat.enabled: + logger.warning( + "--mcStatDebias continuousM with bin-by-bin stat (BB-lite) ON: " + "M must be computed with the BB-lite-inflated variance " + "(mu_b + sum_proc sumw2) in its denominator, NOT the bare data " + "variance. Otherwise M over-subtracts the (already BB-inflated) " + "curvature and H - M becomes non-positive-definite (the fit will " + "diverge / Cholesky will fail). See RABBIT_MCSTAT_DESIGN.md §1b." + ) def init_fit_parms( self, @@ -2196,21 +2205,39 @@ def _compute_nll_components(self, profile=True, full_nll=False): nexp = nexpfullcentral[: self.indata.nbins] - if ( - self.mcStatDebias in ("twoHalf", "kfold") - and self.norm_A is not None - ): + debiased = ( + self.mcStatDebias in ("twoHalf", "kfold") and self.norm_A is not None + ) + if debiased: # Cross-fit jackknife combination L_cf = 2 L_full - 1/2 L_A - 1/2 L_B. # Since mu_bar = 1/2(n_A + n_B) = n_full exactly, this de-biases the # point (gradient = cross-fit score) and curvature (cross-half # Fisher) regardless of nonlinearity. Shared logk; full templates - # carry the BB-lite beta profiling (unchanged). + # carry the BB-lite beta profiling. nexp_A = self._compute_yields_noBBB( full=False, compute_norm=False, templates="A" )[0][: self.indata.nbins] nexp_B = self._compute_yields_noBBB( full=False, compute_norm=False, templates="B" )[0][: self.indata.nbins] + if self.bbstat.enabled: + # Apply the FULL-sample profiled beta (per-bin multiplicative + # factor beta = nexp / nexp_full_raw) to the half predictions too. + # Without this the BB-lite profiling flattens L_full's curvature + # while the -1/2 L_A - 1/2 L_B terms keep full curvature, making + # the jackknife A = 2 H_full,bb - 1/2 H_A - 1/2 H_B indefinite + # (unbounded objective). Sharing beta keeps H_A,bb ~ H_B,bb ~ + # H_full,bb so A ~ H_full,bb > 0. + nexp_full_raw = self._compute_yields_noBBB( + full=False, compute_norm=False, templates="full" + )[0][: self.indata.nbins] + beta_factor = nexp / tf.where( + nexp_full_raw == 0.0, + tf.ones_like(nexp_full_raw), + nexp_full_raw, + ) + nexp_A = nexp_A * beta_factor + nexp_B = nexp_B * beta_factor ln = ( 2.0 * self._compute_ln(nexp, full_nll) - 0.5 * self._compute_ln(nexp_A, full_nll) @@ -2221,6 +2248,9 @@ def _compute_nll_components(self, profile=True, full_nll=False): lc = self._compute_lc(full_nll) + # lbeta (BB-lite MC-stat constraint) and lc (theta priors) are counted + # once. With two-half the same profiled beta is shared across the full + # and half terms (see above), so its constraint enters with weight 1. lbeta = self._compute_lbeta(beta, full_nll) if len(self.regularizers): diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 0bc190b..2dfadd6 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -331,6 +331,16 @@ def add_mc_stat_moment(self, M, param_names, name="mcstat_M"): ``H_data - M``. ``param_names`` label the rows/cols of ``M`` and are resolved against the fit parameter list at fit time. + IMPORTANT — denominator / BB-lite: ``mu_b`` is the per-bin variance used + to weight the score outer products, i.e. the SAME variance the fit's + per-bin term uses. With bin-by-bin stat (BB-lite) ON, that effective + variance is the INFLATED ``mu_b + sum_proc sumw2_b`` (data stat + MC + stat), so ``M`` must use that denominator. Using the bare data variance + while BB-lite is on makes ``M`` roughly a factor ~2 too large, so + ``H_data - M`` goes non-positive-definite and the fit diverges + (RABBIT_MCSTAT_DESIGN.md §1b). With BB-lite OFF use ``mu_b`` = the data + (chisq) or expected (Poisson) variance. + Parameters ---------- M : (n, n) array_like diff --git a/tests/verify_bblite.py b/tests/verify_bblite.py new file mode 100644 index 0000000..c52ce68 --- /dev/null +++ b/tests/verify_bblite.py @@ -0,0 +1,94 @@ +"""Validate MC-stat de-biasing composed with bin-by-bin stat (BB-lite) ON. + +(A) continuous-M + BB-lite: M MUST use the BB-lite-inflated variance + (mu_b + sum_proc sumw2) in its denominator, else H - M is non-PD and the fit + diverges. Builds toy_Mbb with the correct denominator and checks it converges + PD. +(B) two-half + BB-lite: the full-sample profiled beta is shared with the half + predictions (fitter), so the jackknife stays bounded/PD; BB-lite inflates the + baseline and two-half de-biases on top. Uses the well-posed split-logk toy. + +Run with OUT=. Needs toy_splitlogk.hdf5 (tests/toy_splitlogk.py) present. +""" +import os, sys +import numpy as np +import hist +import importlib.util as _ilu + +RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") +sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) +_spec = _ilu.spec_from_file_location( + "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") +) +_rfm = _ilu.module_from_spec(_spec) +_spec.loader.exec_module(_rfm) +from rabbit import inputdata, fitter +from rabbit.param_models import helpers as ph +from rabbit.tensorwriter import TensorWriter + +OUT = os.environ.get("OUT", "/tmp/claude") + + +def build_toy_Mbb(): + NB = 200; A, H1, H2 = 4962.0, 5112.0, 4962.0 + rng = np.random.default_rng(20240614) + n_flat = np.full(NB, A) + n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) + data = (n_flat + n_step).astype(float) + sw_flat = rng.poisson(n_flat).astype(float); sw_step = rng.poisson(n_step).astype(float) + ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") + + def wh(v, var): + h = hist.Hist(ax, storage=hist.storage.Weight()) + h.view()["value"] = v; h.view()["variance"] = var + return h + + tw = TensorWriter(sparse=False) + tw.add_channel([ax], "ch0") + tw.add_data(wh(data, data), "ch0") + tw.add_process(wh(sw_flat, sw_flat), "flat", "ch0", signal=True) + tw.add_process(wh(sw_step, sw_step), "step", "ch0", signal=True) + Vbb = data + (sw_flat + sw_step) # BB-lite inflated variance + M = np.diag([np.sum(sw_flat / Vbb), np.sum(sw_step / Vbb)]) + tw.add_mc_stat_moment(M, ["flat", "step"]) + tw.write(outfolder=OUT, outfilename="toy_Mbb.hdf5") + + +def fit(fn, debias, bbb, sandwich=None): + argv = [fn, "-o", f"{OUT}/vo", "-t", "0", "--chisqFit", + "--allowNegativeParam", "--mcStatDebias", debias] + if not bbb: + argv.append("--noBinByBinStat") + a = _rfm.make_parser().parse_args(argv) + ind = inputdata.FitInputData(fn, None) + f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) + f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + _, g, h = f.loss_val_grad_hess() + _, cov = f.edmval_cov(g, h) + C = np.asarray(cov) + S = None + if sandwich == "continuousM": + S = np.asarray(f.cov_mcstat_sandwich(h)) + elif sandwich == "twoHalf": + S = np.asarray(f.cov_twohalf_sandwich(h, "observed")) + eig = np.linalg.eigvalsh(h.numpy()) + return f.x.numpy(), np.sqrt(np.diag(C)), S, float(np.max(np.abs(g.numpy()))), eig + + +if __name__ == "__main__": + print("(A) continuous-M + BB-lite (toy_Mbb, BB-variance M):") + build_toy_Mbb() + x, sig, S, gmax, eig = fit(f"{OUT}/toy_Mbb.hdf5", "continuousM", True, "continuousM") + pd = bool(np.all(eig > 0)) + print(f" x={np.round(x,4)} |grad|={gmax:.1e} PD={pd} hess_eig={np.round(eig,2)}") + assert pd and gmax < 1e-3, "continuous-M + BB-lite did not converge PD" + + print("(B) two-half + BB-lite (toy_splitlogk):") + fn = f"{OUT}/toy_splitlogk.hdf5" + x0, sig0, _, g0, e0 = fit(fn, "none", True) + x1, sig1, S1, g1, e1 = fit(fn, "twoHalf", True, "twoHalf") + print(f" none x={np.round(x0,4)} tilt sigma={sig0[1]:.4f}") + print(f" twoHalf x={np.round(x1,4)} tilt curv={sig1[1]:.4f} sandwich={np.sqrt(S1[1,1]):.4f}") + assert bool(np.all(e1 > 0)) and g1 < 1e-3, "two-half + BB-lite not converged PD" + assert np.sqrt(S1[1, 1]) > sig0[1], "two-half should inflate the tilt uncertainty" + print("\n ALL BB-lite composition checks passed.") From be872abba3ce8b862a640ab2d300c49696880134 Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 20:59:56 +0000 Subject: [PATCH 06/15] mcstat: propagate de-biased covariance to impacts and outputs (B) - fitter.cov holds the reported de-biased cov (sandwich/curvature), so parameter errors, expected-hist bin errors (--saveHists), and postfit variations already propagate it (verified: sandwich bin errors > curvature). - Traditional impacts now decompose the de-biased CURVATURE A^-1 with the matching covMode-aware bread for total/syst/stat (was: sandwich for syst, raw observed Hessian for stat; --covMode fisher unpropagated). impacts_parms gains cov= and extra_group_vars=. The sandwich's extra coverage term diag(sandwich-curvature) is appended as a new 'mcStatDebias' grouped-impact column (workspace axis label added). Verified columns/labels/values consistent. - global_impacts_parms gains a cov override (decompose the de-biased curvature). - All --mcStatDebias x --covMode x --mcStatDebiasCov combos run --doImpacts / --globalImpacts / --saveHists end-to-end; standard impacts + no-BB/BB-lite validation suites unchanged. See RESULTS.md S9h. Co-Authored-By: Claude Opus 4.8 (1M context) --- bin/rabbit_fit.py | 33 +++++++++++++++++++++------ rabbit/fitter.py | 22 ++++++++++++++---- rabbit/impacts/traditional_impacts.py | 12 ++++++++++ rabbit/workspace.py | 16 ++++++++++++- 4 files changed, 71 insertions(+), 12 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 05119cf..e3241fb 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -548,21 +548,28 @@ def fit(args, fitter, ws, dofit=True): _covmode = getattr(fitter, "covMode", "observed") bread = fitter.fisher_curvature(hess) if _covmode == "fisher" else hess - edmval, cov = fitter.edmval_cov(grad, bread) + edmval, cov_curv = fitter.edmval_cov(grad, bread) + cov = cov_curv logger.info(f"edmval: {edmval}") # Robust (sandwich) covariance for the de-biased fit. Bread = the # (fisher- or observed-) curvature `bread`; continuous-M uses the # analytic meat H = A + M (Sigma = A^-1 + A^-1 M A^-1), two-half uses - # the full-sample meat (both in --covMode). + # the full-sample meat (both in --covMode). cov_curv is kept so the + # impacts can decompose the de-biased CURVATURE consistently and the + # sandwich's extra term is reported as the `mcStatDebias` group. _debias = getattr(fitter, "mcStatDebias", "none") _debiascov = getattr(fitter, "mcStatDebiasCov", "sandwich") - if _debiascov == "sandwich": - if _debias == "continuousM" and getattr(fitter, "mcstat_M", None) is not None: + _is_debiased = ( + (_debias == "continuousM" and getattr(fitter, "mcstat_M", None) is not None) + or (_debias in ("twoHalf", "kfold") and getattr(fitter, "norm_A", None) is not None) + ) + if _debiascov == "sandwich" and _is_debiased: + if _debias == "continuousM": cov = fitter.cov_mcstat_sandwich(bread) logger.info("Reporting continuous-M sandwich covariance " f"(A^-1 + A^-1 M A^-1, covMode={_covmode})") - elif _debias in ("twoHalf", "kfold") and getattr(fitter, "norm_A", None) is not None: + else: cov = fitter.cov_twohalf_sandwich(bread, covMode=_covmode) logger.info("Reporting two-half sandwich covariance " f"(A^-1 H A^-1, covMode={_covmode})") @@ -582,13 +589,25 @@ def fit(args, fitter, ws, dofit=True): logger.info(f"edmvalbeta: {edmvalbeta}") if args.doImpacts: - ws.add_impacts_hists(*fitter.impacts_parms(hess)) + # Decompose the de-biased CURVATURE (cov_curv) with the matching + # curvature `bread`, so total/syst/stat are consistent. When the + # reported cov is the sandwich, append the extra coverage term + # diag(sandwich - curvature) as the `mcStatDebias` impact group. + extra_vars = None + if _is_debiased and _debiascov == "sandwich": + extra_vars = [ + tf.linalg.diag_part(fitter.cov) - tf.linalg.diag_part(cov_curv) + ] + ws.add_impacts_hists( + *fitter.impacts_parms(bread, cov=cov_curv, extra_group_vars=extra_vars) + ) del hess if args.globalImpacts: + # de-biased fit: decompose the de-biased curvature consistently ws.add_impacts_hists( - *fitter.global_impacts_parms(), + *fitter.global_impacts_parms(cov=cov_curv if _is_debiased else None), base_name="global_impacts", global_impacts=True, ) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 2d47282..6664321 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -1235,7 +1235,15 @@ def _cov_stat_floating(self, hess, nstat): ) @tf.function - def impacts_parms(self, hess): + def impacts_parms(self, hess, cov=None, extra_group_vars=None): + # cov: the covariance matrix to DECOMPOSE (per-nuisance + grouped syst + # impacts). Defaults to self.cov. For a de-biased fit pass the de-biased + # CURVATURE covariance (A^-1, in the selected --covMode) and `hess` = the + # matching de-biased curvature, so the whole decomposition (total / syst / + # stat) is internally consistent; the sandwich's extra coverage term is + # then reported as the `mcStatDebias` extra group via extra_group_vars. + if cov is None: + cov = self.cov nstat = ( self.param_model.npoi @@ -1254,7 +1262,7 @@ def impacts_parms(self, hess): param_groups = self._resolved_param_impact_groups() impacts, impacts_grouped = traditional_impacts.impacts_parms( - self.cov, + cov, cov_stat, cov_stat_no_bbb, self.param_model.npoi, @@ -1262,12 +1270,18 @@ def impacts_parms(self, hess): self.indata.systgroupidxs, nmodel_params=self.param_model.npoi + self.param_model.npou, param_groupidxs=[idxs for _, idxs in param_groups], + extra_group_vars=extra_group_vars, ) return impacts, impacts_grouped @tf.function - def global_impacts_parms(self): + def global_impacts_parms(self, cov=None): + # cov: covariance to decompose (default self.cov). For a de-biased fit + # pass the de-biased CURVATURE covariance so the global-impact groups + # (theta0 / nobs / beta0 / syst) decompose it consistently. + if cov is None: + cov = self.cov return global_impacts.global_impacts_parms( self.x, self.bbstat.ubeta, @@ -1282,7 +1296,7 @@ def global_impacts_parms(self): self.bbstat.enabled, self.bbstat.binByBinStatMode, self.globalImpactsFromJVP, - self.cov, + cov, ) @tf.function diff --git a/rabbit/impacts/traditional_impacts.py b/rabbit/impacts/traditional_impacts.py index 04bf116..92aa217 100644 --- a/rabbit/impacts/traditional_impacts.py +++ b/rabbit/impacts/traditional_impacts.py @@ -43,6 +43,7 @@ def impacts_parms( systgroupidxs=[], nmodel_params=None, param_groupidxs=None, + extra_group_vars=None, ): """ Gaussian approximation. @@ -114,4 +115,15 @@ def impacts_parms( ] impacts_grouped = tf.concat([impacts_grouped, *param_cols], axis=1) + if extra_group_vars: + # Extra diagonal-variance impact groups (e.g. the MC-stat de-bias term + # diag(sandwich - curvature)), appended LAST. Each is a full-x variance + # vector; the per-POI/NOI sqrt is gathered like the stat group. + extra_cols = [] + for var in extra_group_vars: + col = tf.sqrt(tf.nn.relu(var)) + col = _gather_poi_noi_vector(col, noiidxs, nsignal_params, nmodel_params) + extra_cols.append(tf.reshape(col, (-1, 1))) + impacts_grouped = tf.concat([impacts_grouped, *extra_cols], axis=1) + return impacts, impacts_grouped diff --git a/rabbit/workspace.py b/rabbit/workspace.py index e069cfa..38f7566 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -70,11 +70,25 @@ def __init__(self, outdir, outname, fitter, postfix=None): param_impact_group_names = [ name for name, _ in fitter._resolved_param_impact_groups() ] + # When a de-biased fit reports the robust (sandwich) covariance, the + # traditional impacts decompose the de-biased CURVATURE and append the + # extra coverage term diag(sandwich - curvature) as an "mcStatDebias" + # group (matching the extra column from traditional_impacts.impacts_parms). + _is_debiased = ( + getattr(fitter, "mcStatDebias", "none") == "continuousM" + and getattr(fitter, "mcstat_M", None) is not None + ) or ( + getattr(fitter, "mcStatDebias", "none") in ("twoHalf", "kfold") + and getattr(fitter, "norm_A", None) is not None + ) + extra_traditional = list(param_impact_group_names) + if _is_debiased and getattr(fitter, "mcStatDebiasCov", "sandwich") == "sandwich": + extra_traditional.append("mcStatDebias") self.grouped_impact_axis = getGroupedImpactsAxes( fitter.indata, bin_by_bin_stat=fitter.bbstat.enabled, per_process=False, - extra_groups=param_impact_group_names, + extra_groups=extra_traditional, ) self.grouped_global_impact_axis = getGroupedImpactsAxes( fitter.indata, From 64fb7dabe2f3097fe5cfb9eac4acacf41032266d Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 21:25:23 +0000 Subject: [PATCH 07/15] mcstat: k-fold averaging via the complete pairwise U-statistic - Generalize the fisher (GN) cross-half curvature to the complete k-fold U-statistic A_k = k/(k-1)[F_full - sum_i F_i_raw] = k/(k-1) sum_{i!=j} cross, via the full-minus-self identity: O(k), streamable, no partition enumeration. Reduces EXACTLY to the 2-half cross-half Fisher at k=2. - _fisher_data_terms returns the raw-per-fold U-statistic terms (residual-free GN); _fisher_core splits H_ln (objective's rescaled halves -> non-data extraction) from F (U-statistic data Fisher). _compute_yields_noBBB gains templates='fold' + fold_index; __init__ stores n_folds and per-fold scaled logk. - Realized in --covMode fisher (auto for k>=2; benefit for k>=3). Observed mode k>=3 uses the single 2-half split; fitter logs an INFO -> use fisher. The point (objective) keeps the rescaled 2-half jackknife (unbiased). --mcStatKfold deprecated (k inferred from folds; U-statistic is parameter-free). - Validated (tests/verify_kfold.py): rabbit == numpy U-statistic (k=4) to 4 d.p.; ensemble median ~ sigma_inf (unbiased), RMS shrinks with k toward the continuous-M floor. No regression. See RESULTS.md S9i. Co-Authored-By: Claude Opus 4.8 (1M context) --- rabbit/fitter.py | 96 +++++++++++++++++++++++++++++++++++++------ rabbit/parsing.py | 6 ++- tests/toy_kfold.py | 51 +++++++++++++++++++++++ tests/verify_kfold.py | 95 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 234 insertions(+), 14 deletions(-) create mode 100644 tests/toy_kfold.py create mode 100644 tests/verify_kfold.py diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 6664321..bfb0050 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -230,6 +230,9 @@ def __init__( # rescaled by k/(k/2)=2 so ==norm_full). Shared logk. self.norm_A = None self.norm_B = None + # n_folds and the raw per-fold templates are kept for the complete + # k-fold U-statistic curvature (k-fold averaging, see _ln_terms_for_fisher). + self.n_folds = None norm_folds = getattr(self.indata, "norm_folds", None) if norm_folds is not None: k = int(norm_folds.shape[0]) @@ -237,6 +240,7 @@ def __init__( raise NotImplementedError( f"two-half de-biasing needs an even number of folds; got k={k}." ) + self.n_folds = k half = k // 2 scale = tf.constant(k / half, dtype=self.indata.dtype) # = 2 self.norm_A = scale * tf.reduce_sum(norm_folds[:half], axis=0) @@ -250,6 +254,9 @@ def __init__( # and the x2 norm rescale leaves the ratio unchanged. self.logk_A = self.logk self.logk_B = self.logk + # Per-fold scaled logk [k, ...] for the U-statistic curvature; None means + # use the shared self.logk for every fold (the common case). + self.logk_folds_scaled = None logk_folds = getattr(self.indata, "logk_folds", None) if logk_folds is not None: if int(logk_folds.shape[0]) != 2: @@ -272,12 +279,27 @@ def __init__( else: self.logk_A = logk_folds[0] self.logk_B = logk_folds[1] + # k=2: per-fold logk == per-half logk, so the U-statistic per-fold + # yields reuse the half logk directly. + self.logk_folds_scaled = tf.stack([self.logk_A, self.logk_B], axis=0) if self.mcStatDebias in ("twoHalf", "kfold") and self.norm_A is None: logger.warning( f"--mcStatDebias {self.mcStatDebias} requested but no fold " "templates (hnorm_folds) found in the input; two-half de-biasing " "is inactive. Add processes with a fold_axis in the TensorWriter." ) + if ( + self.n_folds is not None + and self.n_folds >= 3 + and self.covMode != "fisher" + ): + logger.info( + f"k-fold de-biasing with k={self.n_folds} folds: the k-fold " + "AVERAGING (complete U-statistic, lower curvature-estimate " + "variance) is realized only in --covMode fisher. In observed mode " + "the curvature uses the single 2-half split (no averaging); pass " + "--covMode fisher to use all C(k,k/2)/2 groupings." + ) if self.mcStatDebias == "continuousM": if self.mcstat_M is None: logger.warning( @@ -1009,28 +1031,66 @@ def _ln_terms_for_fisher(self, debiased): return [(nexp_full, 2.0), (nexp_A, -0.5), (nexp_B, -0.5)] return [(nexp_full, 1.0)] + def _fisher_data_terms(self, debiased): + """Terms (nexp, coeff) whose Gauss-Newton sum sum_i coeff_i J_i^T D J_i is + the de-biased data Fisher F_data used as the GN curvature. + + For a fold-based debias this is the COMPLETE k-fold U-statistic (k-fold + AVERAGING): using J_full = sum_i J_i^raw (raw, unrescaled folds), + + F_ch = k/(k-1) [ J_full^T D J_full - sum_i J_i^raw^T D J_i^raw ] + = k/(k-1) sum_{i!=j} J_i^raw^T D J_j^raw + + i.e. the average over ALL fold pairs, computed in O(k) via the + full-minus-self identity (no enumeration of partitions). For k=2 it + equals the single 2-half cross-half Fisher exactly; for k>=3 it averages + the C(k,k/2)/2 half-groupings, reducing the curvature-estimate variance + toward the continuous-M floor as k grows. The RAW folds (~n_full/k) are + used only here (the GN form is residual-independent), never in the + objective ln (which keeps the rescaled halves for a sensible point).""" + nbins = self.indata.nbins + nexp_full = self._compute_yields_with_beta( + profile=True, compute_norm=False, full=len(self.regularizers) + )[0][:nbins] + if ( + debiased + and self.mcStatDebias in ("twoHalf", "kfold") + and self.n_folds is not None + ): + k = self.n_folds + coeff = k / (k - 1.0) + terms = [(nexp_full, coeff)] + for i in range(k): + n_i = self._compute_yields_noBBB( + full=False, compute_norm=False, templates="fold", fold_index=i + )[0][:nbins] + terms.append((n_i, -coeff)) + return terms + return [(nexp_full, 1.0)] + def _fisher_core(self, hess_obj, debiased): """Gauss-Newton (expected-information) curvature corresponding to the observed-Hessian `hess_obj`: F = hess_obj - H_ln_obs + F_data - where H_ln_obs is the observed Hessian of the (active) data term and - F_data = sum_i coeff_i J_i^T D J_i is its Gauss-Newton replacement - (J_i = d nexp_i / d x, D = ell''(nexp_i): 1/V for chisq, data_cov_inv - for covarianceFit, 1/nexp for Poisson expected info). This drops the - residual * d2nexp piece (the only data-term difference) while leaving - the constraint / external / BB-lite curvature in hess_obj untouched, so - the two modes stay consistent (RABBIT_MCSTAT_DESIGN.md §2c/§2d).""" + H_ln_obs is the observed Hessian of the OBJECTIVE's data term (rescaled + halves), so `hess_obj - H_ln_obs` is the non-data curvature (constraints, + external M, BB-lite). F_data is the de-biased data Gauss-Newton Fisher + from _fisher_data_terms (the complete k-fold U-statistic for the bread; + the full J^T D J for the meat). D = ell''(nexp): 1/V chisq, data_cov_inv + covFit, 1/nexp Poisson. Drops the residual * d2nexp piece while keeping + the non-data curvature (RABBIT_MCSTAT_DESIGN.md §2c/§2d).""" with tf.GradientTape(persistent=True) as t2: with tf.GradientTape(persistent=True) as t1: - terms = self._ln_terms_for_fisher(debiased) - ln = tf.add_n([c * self._compute_ln(n) for n, c in terms]) + obj_terms = self._ln_terms_for_fisher(debiased) + fisher_terms = self._fisher_data_terms(debiased) + ln = tf.add_n([c * self._compute_ln(n) for n, c in obj_terms]) g = t1.gradient(ln, self.x) H_ln = t2.jacobian(g, self.x) F = tf.zeros_like(hess_obj) - for n, c in terms: + for n, c in fisher_terms: J = t1.jacobian(n, self.x) # [nbins, nparams] if self.covarianceFit: JT_Cinv = tf.matmul(J, self.data_cov_inv, transpose_a=True) @@ -1815,10 +1875,14 @@ def _init_logk_scaled(self): else: self.logk = self.indata.logk * rnorm_init[..., None, None] - def _compute_yields_noBBB(self, full=True, compute_norm=True, templates="full"): + def _compute_yields_noBBB( + self, full=True, compute_norm=True, templates="full", fold_index=None + ): # templates: "full" uses indata.norm; "A"/"B" use the precomputed # half-sample fold templates norm_A/norm_B (two-half de-biasing, shared - # logk). Only supported in the dense path. + # logk); "fold" uses the RAW per-fold template norm_folds[fold_index] + # (for the complete k-fold U-statistic curvature). Only supported in the + # dense path. # full: compute yields inclduing masked channels # compute_norm: also build the dense [nbins, nproc] normcentral tensor. # In sparse mode this is expensive (forward + backward) and is only @@ -1904,6 +1968,14 @@ def _compute_yields_noBBB(self, full=True, compute_norm=True, templates="full"): elif templates == "B": norm_src = self.norm_B logk_src = self.logk_B + elif templates == "fold": + # raw per-fold template (unrescaled); shared logk unless split-logk + norm_src = self.indata.norm_folds[fold_index] + logk_src = ( + self.logk + if self.logk_folds_scaled is None + else self.logk_folds_scaled[fold_index] + ) else: norm_src = self.indata.norm logk_src = self.logk diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 14f5bcb..14770bf 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -401,8 +401,10 @@ def common_parser(): "--mcStatKfold", default=2, type=int, - help="Number of folds R re-grouped into halves for --mcStatDebias kfold " - "(2 = the two-half case).", + help="DEPRECATED / unused: the number of folds k is inferred from the " + "fold-axis size in the input, and the k-fold averaging uses the COMPLETE " + "pairwise U-statistic (all C(k,k/2)/2 groupings, O(k), parameter-free) in " + "--covMode fisher. No fold/regrouping count needs to be chosen here.", ) parser.add_argument( "--covMode", diff --git a/tests/toy_kfold.py b/tests/toy_kfold.py new file mode 100644 index 0000000..f9b449f --- /dev/null +++ b/tests/toy_kfold.py @@ -0,0 +1,51 @@ +"""Build the 2-process degenerate toy with k=4 MC-stat folds, to exercise the +complete k-fold U-statistic curvature (k-fold averaging). Each process's finite-MC +template is split into 4 independent folds (multinomial 1/4 per bin); the full = +sum of folds is derived by the writer. Writes toy_kfold.hdf5.""" +import os +import numpy as np +import hist +from rabbit.tensorwriter import TensorWriter + +NB = 200 +A, H1, H2 = 4962.0, 5112.0, 4962.0 +K = 4 +rng = np.random.default_rng(20240614) + +n_flat = np.full(NB, A) +n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) +data = (n_flat + n_step).astype(float) + +sw_flat = rng.poisson(n_flat) +sw_step = rng.poisson(n_step) +# split each bin's counts into K folds via a multinomial(count, [1/K]*K) +flat_folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in sw_flat], axis=0).T # [K, NB] +step_folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in sw_step], axis=0).T + +ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") +axf = hist.axis.IntCategory(list(range(K)), name="mcfold") + + +def dhist(v): + h = hist.Hist(ax, storage=hist.storage.Weight()) + h.view()["value"] = v; h.view()["variance"] = v + return h + + +def fhist(folds): + h = hist.Hist(ax, axf, storage=hist.storage.Weight()) + for f in range(K): + h.view()["value"][:, f] = folds[f] + h.view()["variance"][:, f] = folds[f] + return h + + +if __name__ == "__main__": + out = os.environ.get("OUT", "/tmp/claude") + tw = TensorWriter(sparse=False) + tw.add_channel([ax], "ch0") + tw.add_data(dhist(data), "ch0") + tw.add_process(fhist(flat_folds), "flat", "ch0", signal=True, fold_axis="mcfold") + tw.add_process(fhist(step_folds), "step", "ch0", signal=True, fold_axis="mcfold") + tw.write(outfolder=out, outfilename="toy_kfold.hdf5") + print(f"wrote {out}/toy_kfold.hdf5 (k={K} folds)") diff --git a/tests/verify_kfold.py b/tests/verify_kfold.py new file mode 100644 index 0000000..9eb1146 --- /dev/null +++ b/tests/verify_kfold.py @@ -0,0 +1,95 @@ +"""Validate the complete k-fold U-statistic curvature (k-fold averaging). + +(A) Correctness: rabbit's fisher curvature on the k=4 fold toy equals the numpy + complete pairwise U-statistic A_k = k/(k-1)[F_full - sum_i F_i_raw]. +(B) Benefit: over an ensemble the U-statistic curvature is unbiased (median ~ + sigma_inf) and its RMS shrinks as k grows (toward the continuous-M floor). + +Needs toy_kfold.hdf5 (tests/toy_kfold.py). Run with OUT=. +""" +import os, sys +import numpy as np +import importlib.util as _ilu + +RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") +sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) +_spec = _ilu.spec_from_file_location( + "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") +) +_rfm = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_rfm) +from rabbit import inputdata, fitter +from rabbit.param_models import helpers as ph + +NB, A, H1, H2 = 200, 4962.0, 5112.0, 4962.0 +ddif = np.array([1.0, -1.0]) / np.sqrt(2) +n_flat = np.full(NB, A) +n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) +data = (n_flat + n_step).astype(float) + + +def ustat(Ti, V): + Jf = sum(Ti) + Ff = np.einsum("bi,b,bj->ij", Jf, 1 / V, Jf) + Fi = sum(np.einsum("bi,b,bj->ij", T, 1 / V, T) for T in Ti) + K = len(Ti) + return (K / (K - 1.0)) * (Ff - Fi) + + +def numpy_ref_k4(): + rng = np.random.default_rng(20240614) # same seed as toy_kfold.py + swf = rng.poisson(n_flat); sws = rng.poisson(n_step) + ff = np.stack([rng.multinomial(c, [0.25] * 4) for c in swf], 0).T + sf = np.stack([rng.multinomial(c, [0.25] * 4) for c in sws], 0).T + Ti = [np.stack([ff[i], sf[i]], 1).astype(float) for i in range(4)] + return ustat(Ti, data) + + +def rabbit_fisher(fn): + import tensorflow as tf + a = _rfm.make_parser().parse_args( + [fn, "-o", "/tmp/claude/vo", "-t", "0", "--chisqFit", "--noBinByBinStat", + "--allowNegativeParam", "--mcStatDebias", "kfold", "--covMode", "fisher"]) + ind = inputdata.FitInputData(fn, None) + f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) + f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + _, _, h = f.loss_val_grad_hess() + return f.n_folds, np.asarray(f.fisher_curvature(h)) + + +def ensemble_rms(K, ntoy=300): + rng = np.random.default_rng(100 + K) + s = [] + for _ in range(ntoy): + swf = rng.poisson(n_flat); sws = rng.poisson(n_step) + ff = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in swf], 0).T + sf = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in sws], 0).T + Ti = [np.stack([ff[i], sf[i]], 1).astype(float) for i in range(K)] + Ak = ustat(Ti, data) + v = ddif @ np.linalg.inv(Ak) @ ddif + if v > 0: + s.append(np.sqrt(v)) + return np.median(s), np.std(s) + + +if __name__ == "__main__": + out = os.environ.get("OUT", "/tmp/claude") + A4 = numpy_ref_k4() + k, Fb = rabbit_fisher(f"{out}/toy_kfold.hdf5") + print(f"(A) rabbit k={k} fisher curvature == numpy U-statistic A_k: " + f"{np.allclose(Fb, A4, rtol=1e-4)} " + f"(sigma(dif) rabbit={np.sqrt(ddif@np.linalg.inv(Fb)@ddif):.4f}, " + f"numpy={np.sqrt(ddif@np.linalg.inv(A4)@ddif):.4f})") + assert np.allclose(Fb, A4, rtol=1e-4) + + Tt = np.stack([n_flat, n_step], 1) + sinf = np.sqrt(ddif @ np.linalg.inv(np.einsum("bi,b,bj->ij", Tt, 1 / data, Tt)) @ ddif) + print(f"(B) ensemble: sigma_inf={sinf:.4f}") + prev = None + for K in (4, 8, 16): + m, r = ensemble_rms(K) + print(f" k={K:2d} median={m:.4f} RMS={r:.4f}") + assert abs(m - sinf) < 0.02, "median should be ~unbiased" + if prev is not None: + assert r < prev + 1e-9, "RMS should not increase with k" + prev = r + print("\n k-fold averaging checks passed (correct, unbiased, variance-reducing).") From 03e72d3df735bcde382fa3389612b07f05caeb5e Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 21:39:52 +0000 Subject: [PATCH 08/15] mcstat: k>2 split-logk via per-fold logk in the U-statistic curvature - Lift the k=2-only split-logk restriction. Build logk_folds_scaled [k,...] (per- fold scaled logk) for any k; k=2 sets the half logk for the objective (full treatment), k>2 objective falls back to shared logk (point de-biases nominal noise; systematic-template noise de-biased in the curvature) + INFO. - _fisher_data_terms uses n_sumfold = sum_i n_i as the U-statistic 'full' term and the meat, so full-minus-self = sum_{i!=j} cross holds EXACTLY for shared logk (sum_i n_i == nexp_full, common case unchanged) AND split logk. BB-lite beta applied to every per-fold prediction. - toy_splitlogk.py generalized to K folds (env K, default 2). Validated: shared- logk paths unchanged (verify_kfold/twohalf/bblite pass; k=2 tilt sandwich 0.0907 unchanged); k=4 split-logk loads per-fold logk [4,...], converges, de-biases the tilt systematic more than shared (0.130 vs 0.086). See RESULTS.md S9j. Co-Authored-By: Claude Opus 4.8 (1M context) --- rabbit/fitter.py | 81 ++++++++++++++++++++++++++++-------------- tests/toy_splitlogk.py | 38 ++++++++++---------- 2 files changed, 74 insertions(+), 45 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index bfb0050..5d8764e 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -259,11 +259,6 @@ def __init__( self.logk_folds_scaled = None logk_folds = getattr(self.indata, "logk_folds", None) if logk_folds is not None: - if int(logk_folds.shape[0]) != 2: - raise NotImplementedError( - "split-logk (hlogk_folds) is currently supported only for k=2 " - "folds (per-fold logk == per-half logk only when k=2)." - ) if not self.indata.symmetric_tensor: raise NotImplementedError("split-logk requires a symmetric tensor.") # apply the same constant rnorm_init scaling as _init_logk_scaled @@ -274,14 +269,26 @@ def __init__( ), [self.indata.nbinsfull, self.indata.nproc], ) - self.logk_A = logk_folds[0] * rnorm_init[..., None] - self.logk_B = logk_folds[1] * rnorm_init[..., None] + self.logk_folds_scaled = logk_folds * rnorm_init[None, ..., None] + else: + self.logk_folds_scaled = logk_folds + # The CURVATURE (U-statistic) uses the per-fold logk for any k. The + # OBJECTIVE (point) needs the per-HALF logk, which equals the per-fold + # logk only for k=2 (log of a sum != sum of logs). So for k=2 the + # objective uses the split half logk (full treatment); for k>2 the + # objective falls back to the shared logk (the point still de-biases + # the dominant nominal-template noise via norm^A/norm^B; the + # systematic-template noise is de-biased in the curvature). + if int(logk_folds.shape[0]) == 2: + self.logk_A = self.logk_folds_scaled[0] + self.logk_B = self.logk_folds_scaled[1] else: - self.logk_A = logk_folds[0] - self.logk_B = logk_folds[1] - # k=2: per-fold logk == per-half logk, so the U-statistic per-fold - # yields reuse the half logk directly. - self.logk_folds_scaled = tf.stack([self.logk_A, self.logk_B], axis=0) + logger.info( + "split-logk with k>2: the objective uses shared logk (the " + "point de-biases nominal-template noise); the systematic-" + "template noise is de-biased in the k-fold U-statistic " + "curvature (--covMode fisher)." + ) if self.mcStatDebias in ("twoHalf", "kfold") and self.norm_A is None: logger.warning( f"--mcStatDebias {self.mcStatDebias} requested but no fold " @@ -1052,21 +1059,41 @@ def _fisher_data_terms(self, debiased): nexp_full = self._compute_yields_with_beta( profile=True, compute_norm=False, full=len(self.regularizers) )[0][:nbins] - if ( - debiased - and self.mcStatDebias in ("twoHalf", "kfold") - and self.n_folds is not None - ): - k = self.n_folds - coeff = k / (k - 1.0) - terms = [(nexp_full, coeff)] - for i in range(k): - n_i = self._compute_yields_noBBB( - full=False, compute_norm=False, templates="fold", fold_index=i - )[0][:nbins] - terms.append((n_i, -coeff)) - return terms - return [(nexp_full, 1.0)] + fold_active = ( + self.mcStatDebias in ("twoHalf", "kfold") and self.n_folds is not None + ) + if not fold_active: + return [(nexp_full, 1.0)] + + # Per-fold raw predictions n_i (each with its own logk for split-logk), + # times the full-sample BB-lite beta so all terms are consistently + # profiled (beta_factor=1 without BB-lite). The "full" term is the SUM of + # the per-fold predictions, J_sumfold = sum_i J_i, so the full-minus-self + # identity F_sumfold - sum_i F_i = sum_{i!=j} cross holds EXACTLY for both + # shared logk (where sum_i n_i == nexp_full) and split logk (where it does + # not). Using n_sumfold for the meat too keeps bread/meat consistent + # (H - A = M) under split-logk. + if self.bbstat.enabled: + nexp_full_raw = self._compute_yields_noBBB( + full=False, compute_norm=False, templates="full" + )[0][:nbins] + beta_factor = nexp_full / tf.where( + nexp_full_raw == 0.0, tf.ones_like(nexp_full_raw), nexp_full_raw + ) + else: + beta_factor = tf.ones_like(nexp_full) + per_fold = [ + beta_factor + * self._compute_yields_noBBB( + full=False, compute_norm=False, templates="fold", fold_index=i + )[0][:nbins] + for i in range(self.n_folds) + ] + n_sumfold = tf.add_n(per_fold) + if debiased: + coeff = self.n_folds / (self.n_folds - 1.0) + return [(n_sumfold, coeff)] + [(n_i, -coeff) for n_i in per_fold] + return [(n_sumfold, 1.0)] def _fisher_core(self, hess_obj, debiased): """Gauss-Newton (expected-information) curvature corresponding to the diff --git a/tests/toy_splitlogk.py b/tests/toy_splitlogk.py index 2997283..f24ede7 100644 --- a/tests/toy_splitlogk.py +++ b/tests/toy_splitlogk.py @@ -1,33 +1,35 @@ """Small toy with a FOLDED process and a FOLDED systematic, to exercise the -split-logk two-half path (add_systematic(fold_axis=...)). The systematic's -up-variation is built from each fold's OWN events (a per-fold reweight), so its -logk carries MC noise that differs between folds -> split logk is meaningful. +split-logk path (add_systematic(fold_axis=...)). The systematic's up-variation is +built from each fold's OWN events (a per-fold reweight), so its logk carries MC +noise that differs between folds -> split logk is meaningful. Writes toy_splitlogk.hdf5 (folded syst) and toy_sharedlogk.hdf5 (same nominal -folds, systematic added WITHOUT fold_axis = shared logk) for comparison.""" +folds, systematic added WITHOUT fold_axis = shared logk) for comparison. Set the +env var K to the number of folds (default 2; K>=3 exercises k>2 split-logk, whose +per-fold logk feeds the k-fold U-statistic curvature in --covMode fisher).""" import os import numpy as np import hist from rabbit.tensorwriter import TensorWriter NB = 50 +K = int(os.environ.get("K", "2")) rng = np.random.default_rng(7) nom = np.linspace(800.0, 1200.0, NB) # smooth nominal shape data = nom.copy() # Asimov -# two independent half-samples of the nominal (per-bin binomial split) +# K independent folds of the nominal (per-bin multinomial split) full = rng.poisson(nom) -A = rng.binomial(full, 0.5); B = full - A +folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in full], axis=0).T # [K, NB] -# systematic: a tilt reweight applied to each fold's OWN events -> per-fold -# up template = fold * (1 + 0.15*(x-0.5)) with fold-specific Poisson noise +# systematic: a tilt reweight applied to each fold's OWN events -> per-fold up +# template with fold-specific Poisson noise x = np.linspace(0, 1, NB) w = 1.0 + 0.15 * (x - 0.5) -upA = rng.poisson(np.maximum(A * w, 0.0)) -upB = rng.poisson(np.maximum(B * w, 0.0)) +upfolds = np.stack([rng.poisson(np.maximum(folds[f] * w, 0.0)) for f in range(K)], axis=0) ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") -axf = hist.axis.IntCategory([0, 1], name="mcfold") +axf = hist.axis.IntCategory(list(range(K)), name="mcfold") def dhist(v): @@ -36,10 +38,10 @@ def dhist(v): return h -def fhist(f0, f1): +def fhist(fl): h = hist.Hist(ax, axf, storage=hist.storage.Weight()) - h.view()["value"][:, 0] = f0; h.view()["value"][:, 1] = f1 - h.view()["variance"][:, 0] = f0; h.view()["variance"][:, 1] = f1 + for f in range(K): + h.view()["value"][:, f] = fl[f]; h.view()["variance"][:, f] = fl[f] return h @@ -47,12 +49,12 @@ def build(outname, split): tw = TensorWriter(sparse=False) tw.add_channel([ax], "ch0") tw.add_data(dhist(data), "ch0") - tw.add_process(fhist(A, B), "sig", "ch0", signal=True, fold_axis="mcfold") + tw.add_process(fhist(folds), "sig", "ch0", signal=True, fold_axis="mcfold") if split: - tw.add_systematic(fhist(upA, upB), "tilt", "sig", "ch0", fold_axis="mcfold") + tw.add_systematic(fhist(upfolds), "tilt", "sig", "ch0", fold_axis="mcfold") else: # shared logk: full up template (folds summed), no fold_axis - tw.add_systematic(dhist(upA + upB), "tilt", "sig", "ch0") + tw.add_systematic(dhist(upfolds.sum(0)), "tilt", "sig", "ch0") tw.write(outfolder=os.path.dirname(outname) or ".", outfilename=os.path.basename(outname)) @@ -61,4 +63,4 @@ def build(outname, split): out = os.environ.get("OUT", "/tmp/claude") build(f"{out}/toy_splitlogk.hdf5", split=True) build(f"{out}/toy_sharedlogk.hdf5", split=False) - print(f"wrote {out}/toy_splitlogk.hdf5 (split) and toy_sharedlogk.hdf5 (shared)") + print(f"wrote {out}/toy_splitlogk.hdf5 (split) and toy_sharedlogk.hdf5 (shared), K={K}") From 7f5708116da31f1aed51f5f26990248c1fd1bb57 Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 21:53:01 +0000 Subject: [PATCH 09/15] mcstat: asymmetric folded systematics (split-logk on up/down pairs) - add_systematic([up,down], ..., fold_axis=, symmetrize=): _add_systematic_folded handles the up/down PAIR; per fold computes logkavg (+ logkhalfdiff for the asymmetric symmetrize mode), stored in dict_logkavg_folds / dict_logkhalfdiff_ folds. write() assembles the fold logk in the full tensor's shape ([k,..,nsyst] symmetric or [k,..,2,nsyst] asymmetric; folded systs split, others replicated). - Fitter: shape-aware rnorm_init scaling of logk_folds_scaled (sym/asym); removed the symmetric-tensor guard; per-fold asym logk flows through the existing asymmetric interpolation in _compute_yields_noBBB(templates='fold'). - Folded process + asymmetric SHARED syst already worked (verified). 'linear'/ 'quadratic' symmetrize not supported with folds. Validated: k=4 folded asym syst converges PD, logk_folds [4,40,1,2,1]. Symmetric/split paths unchanged. See RESULTS.md S9k. Co-Authored-By: Claude Opus 4.8 (1M context) --- rabbit/fitter.py | 15 +++-- rabbit/tensorwriter.py | 150 ++++++++++++++++++++++++++++------------- 2 files changed, 115 insertions(+), 50 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 5d8764e..ff704fc 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -259,9 +259,11 @@ def __init__( self.logk_folds_scaled = None logk_folds = getattr(self.indata, "logk_folds", None) if logk_folds is not None: - if not self.indata.symmetric_tensor: - raise NotImplementedError("split-logk requires a symmetric tensor.") - # apply the same constant rnorm_init scaling as _init_logk_scaled + # apply the same constant rnorm_init scaling as _init_logk_scaled. + # logk_folds is [k, nbinsfull, nproc, nsyst] (symmetric) or + # [k, nbinsfull, nproc, 2, nsyst] (asymmetric); broadcast rnorm_init + # ([nbinsfull, nproc]) over the leading fold axis and trailing + # syst (and asym) axes. if self.indata.systematic_type == "normal" and self.param_model.nparams > 0: rnorm_init = tf.broadcast_to( self.param_model.compute( @@ -269,7 +271,12 @@ def __init__( ), [self.indata.nbinsfull, self.indata.nproc], ) - self.logk_folds_scaled = logk_folds * rnorm_init[None, ..., None] + ntrail = len(logk_folds.shape) - 3 # 1 (sym) or 2 (asym) + scale = tf.reshape( + rnorm_init, + [1, self.indata.nbinsfull, self.indata.nproc] + [1] * ntrail, + ) + self.logk_folds_scaled = logk_folds * scale else: self.logk_folds_scaled = logk_folds # The CURVATURE (U-statistic) uses the per-fold logk for any k. The diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 2dfadd6..f2f807c 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -59,8 +59,11 @@ def __init__( self.dict_norm_folds = {} # [channel][process] -> [k, nbinschan] self.fold_k = {} # [channel][process] -> k # Split-logk fold templates (only for systematics added with fold_axis): - # dict_logkavg_folds[channel][process][syst] = [k, nbinschan]. + # dict_logkavg_folds[channel][process][syst] = [k, nbinschan]; for + # asymmetric folded systematics dict_logkhalfdiff_folds holds the per-fold + # half-difference (same shape). self.dict_logkavg_folds = {} # [channel][process][syst] -> [k, nbinschan] + self.dict_logkhalfdiff_folds = {} # [channel][process][syst] -> [k, nbinschan] self.dict_logkavg = {} # [channel][proc][syst] self.dict_logkhalfdiff = {} # [channel][proc][syst] self.dict_logkavg_indices = {} @@ -384,6 +387,7 @@ def add_channel(self, axes, name=None, masked=False, flow=False): self.dict_norm_folds[name] = {} self.fold_k[name] = {} self.dict_logkavg_folds[name] = {} + self.dict_logkhalfdiff_folds[name] = {} self.dict_beta_variations[name] = {} # add masked channels last @@ -1402,21 +1406,31 @@ def _add_systematic_folded( self, h, name, process, channel, kfactor, mirror, symmetrize, add_to_data_covariance, as_difference, syst_axes, fold_axis, **kargs ): - """Add a systematic whose variation histogram carries an MC-stat fold + """Add a systematic whose variation histogram(s) carry an MC-stat fold axis, SPLITTING logk per fold (RABBIT_MCSTAT_DESIGN.md §2a). The full systematic is registered by projecting the fold axis out; per-fold - logkavg [k, nbinschan] is stored for the two-half cross-fit so the - systematic's own template noise is de-biased. - - Minimal support: dense histogram, single variation (mirror=True), and the - fold axis the only extra axis beyond the channel. The process must have - been added with the SAME fold_axis (so per-fold nominals exist).""" - if self.sparse and self._issparse(h): + logkavg (and, for an asymmetric symmetrize mode, per-fold logkhalfdiff) + [k, nbinschan] are stored for the cross-fit so the systematic's own + template noise is de-biased. + + `h` is either a single histogram (mirror=True, symmetric) or a list + [up, down] of two histograms (asymmetric); both must carry the named + fold_axis (the only extra axis beyond the channel). The process must have + been added with the SAME fold_axis. Dense only; the 'linear'/'quadratic' + symmetrize modes (which split into two systematics) are not supported with + folds.""" + is_pair = isinstance(h, (list, tuple)) + h0 = h[0] if is_pair else h + if self.sparse and self._issparse(h0): raise NotImplementedError("fold_axis systematics: sparse not supported.") - if isinstance(h, (list, tuple)) or not mirror: + if not is_pair and not mirror: + raise RuntimeError( + "Only one histogram given but mirror=False, can not construct a variation" + ) + if is_pair and symmetrize in ("linear", "quadratic"): raise NotImplementedError( - "fold_axis systematics: only the symmetric single-histogram " - "(mirror=True) case is supported." + "fold_axis systematics do not support the 'linear'/'quadratic' " + "symmetrize modes (they split into two systematics)." ) if process not in self.dict_norm_folds[channel]: raise RuntimeError( @@ -1425,7 +1439,7 @@ def _add_systematic_folded( ) norm_folds = self.dict_norm_folds[channel][process] # [k, nbinschan] k = norm_folds.shape[0] - ax_names = [a.name for a in h.axes] + ax_names = [a.name for a in h0.axes] if isinstance(fold_axis, str): if fold_axis not in ax_names: raise ValueError( @@ -1434,36 +1448,62 @@ def _add_systematic_folded( fold_idx = ax_names.index(fold_axis) else: fold_idx = int(fold_axis) - if h.axes[fold_idx].size != k: + if h0.axes[fold_idx].size != k: raise ValueError( - f"systematic fold axis size {h.axes[fold_idx].size} != process k={k}" + f"systematic fold axis size {h0.axes[fold_idx].size} != process k={k}" ) flow = self.channels[channel]["flow"] systematic_type = "normal" if add_to_data_covariance else self.systematic_type - # per-fold logkavg vs the per-fold nominal - logk_folds = np.zeros_like(norm_folds) + # per-fold logkavg (+ logkhalfdiff for asymmetric) vs the per-fold nominal + logkavg_folds = np.zeros_like(norm_folds) + logkhalfdiff_folds = np.zeros_like(norm_folds) + is_asym = False for f in range(k): - syst_f = self.get_flat_values(h[{fold_idx: f}], flow=flow) norm_f = norm_folds[f] - if as_difference: - syst_f = norm_f + syst_f - logk_folds[f] = self.get_logk( - syst_f, norm_f, kfactor, systematic_type=systematic_type - ) + if is_pair: + up_f = self.get_flat_values(h[0][{fold_idx: f}], flow=flow) + down_f = self.get_flat_values(h[1][{fold_idx: f}], flow=flow) + if as_difference: + up_f = norm_f + up_f + down_f = norm_f + down_f + lu = self.get_logk(up_f, norm_f, kfactor, systematic_type=systematic_type) + ld = -self.get_logk(down_f, norm_f, kfactor, systematic_type=systematic_type) + if symmetrize == "conservative": + logkavg_folds[f] = np.where(np.abs(lu) > np.abs(ld), lu, ld) + elif symmetrize == "average": + logkavg_folds[f] = 0.5 * (lu + ld) + else: # asymmetric: keep the half-difference + is_asym = True + logkavg_folds[f] = 0.5 * (lu + ld) + logkhalfdiff_folds[f] = 0.5 * (lu - ld) + else: + syst_f = self.get_flat_values(h[{fold_idx: f}], flow=flow) + if as_difference: + syst_f = norm_f + syst_f + logkavg_folds[f] = self.get_logk( + syst_f, norm_f, kfactor, systematic_type=systematic_type + ) # register the FULL systematic by projecting the fold axis out - h_full = h[{fold_idx: sum}] + if is_pair: + h_full = [h[0][{fold_idx: sum}], h[1][{fold_idx: sum}]] + else: + h_full = h[{fold_idx: sum}] self.add_systematic( h_full, name, process, channel, kfactor=kfactor, mirror=mirror, symmetrize=symmetrize, add_to_data_covariance=add_to_data_covariance, as_difference=as_difference, syst_axes=syst_axes, **kargs ) - # store the per-fold logk under the resolved (full) systematic name(s) - if process not in self.dict_logkavg_folds[channel]: - self.dict_logkavg_folds[channel][process] = {} - self.dict_logkavg_folds[channel][process][name] = logk_folds.astype(self.dtype) + # store the per-fold logk under the resolved (full) systematic name + self.dict_logkavg_folds[channel].setdefault(process, {}) + self.dict_logkavg_folds[channel][process][name] = logkavg_folds.astype(self.dtype) + if is_asym: + self.dict_logkhalfdiff_folds[channel].setdefault(process, {}) + self.dict_logkhalfdiff_folds[channel][process][name] = ( + logkhalfdiff_folds.astype(self.dtype) + ) def add_beta_variations( self, @@ -2153,11 +2193,11 @@ def write(self, outfolder="./", outfilename="rabbit_input.hdf5", meta_data_dict= if self.has_beta_variations: beta_variations = np.zeros([nbinsfull, nbins, nproc], self.dtype) - # Split-logk fold tensor [k, nbinsfull, nproc, nsyst]: built only when - # at least one systematic was added with a fold_axis. Folded systs get - # their per-fold logk; all other (bin,proc,syst) entries replicate the - # full logk/k-equivalent so logk_A == logk_B for them (shared, no extra - # de-bias). Requires a symmetric tensor. + # Split-logk fold tensor: built only when at least one systematic was + # added with a fold_axis. Same trailing shape as the full `logk` + # ([..., nsyst] symmetric or [..., 2, nsyst] asymmetric) with a leading + # fold axis. Folded systs get their per-fold logk; other (bin,proc,syst) + # entries replicate the full logk so they are identical across folds. any_folded_syst = any( len(self.dict_logkavg_folds[c].get(p, {})) for c in self.channels @@ -2165,13 +2205,14 @@ def write(self, outfolder="./", outfilename="rabbit_input.hdf5", meta_data_dict= ) logk_folds = None if any_folded_syst: - if not self.symmetric_tensor: - raise NotImplementedError( - "fold_axis systematics require a symmetric tensor." - ) kset = {k for c in self.channels for k in self.fold_k[c].values()} kfold = kset.pop() - logk_folds = np.zeros([kfold, nbinsfull, nproc, nsyst], self.dtype) + if self.symmetric_tensor: + logk_folds = np.zeros([kfold, nbinsfull, nproc, nsyst], self.dtype) + else: + logk_folds = np.zeros( + [kfold, nbinsfull, nproc, 2, nsyst], self.dtype + ) for chan in self.channels.keys(): nbinschan = self.nbinschan[chan] @@ -2192,20 +2233,37 @@ def write(self, outfolder="./", outfilename="rabbit_input.hdf5", meta_data_dict= if logk_folds is not None else {} ) + logkhalfdiff_folds_proc = ( + self.dict_logkhalfdiff_folds[chan].get(proc, {}) + if logk_folds is not None + else {} + ) for isyst, syst in enumerate(systs): if syst not in dict_logkavg_proc.keys(): continue if logk_folds is not None: - if syst in logkavg_folds_proc: - logk_folds[:, ibin : ibin + nbinschan, iproc, isyst] = ( - logkavg_folds_proc[syst] - ) + bs = slice(ibin, ibin + nbinschan) + # per-fold avg (folded) or replicated full avg (shared) + avg_f = ( + logkavg_folds_proc[syst] + if syst in logkavg_folds_proc + else dict_logkavg_proc[syst][None, :] + ) + if self.symmetric_tensor: + logk_folds[:, bs, iproc, isyst] = avg_f else: - # shared across folds: replicate the full logk - logk_folds[:, ibin : ibin + nbinschan, iproc, isyst] = ( - dict_logkavg_proc[syst][None, :] - ) + logk_folds[:, bs, iproc, 0, isyst] = avg_f + # per-fold halfdiff (folded asymmetric) or replicated + # full halfdiff (shared asymmetric); 0 if symmetric syst + if syst in logkhalfdiff_folds_proc: + logk_folds[:, bs, iproc, 1, isyst] = ( + logkhalfdiff_folds_proc[syst] + ) + elif syst in dict_logkhalfdiff_proc: + logk_folds[:, bs, iproc, 1, isyst] = ( + dict_logkhalfdiff_proc[syst][None, :] + ) if self.symmetric_tensor: logk[ibin : ibin + nbinschan, iproc, isyst] = ( From ba88a43c4d19f9c36dfa86cab77d37fc0442a0f7 Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 21:58:35 +0000 Subject: [PATCH 10/15] mcstat: sparse-mode fold de-biasing (sparse per-fold yield path) - Fold norm is stored dense [k,nbinsfull,nproc] even in sparse mode (the full template stays sparse CSR); removed the writer's sparse fold_axis raise (folds densified). - Fitter sparse per-fold yield: scatter the shared systematic factor from the sparse logk (exp(logsnorm) log_normal / additive logsnorm normal) into a dense [nbinsfull,nproc] grid and contract with the dense fold/half norm. The fisher U-statistic curvature (dense per-bin yield + dense Jacobian) then works for sparse-input + dense-cov. split-logk in sparse not supported (raises). - Validated (tests/verify_sparsefold.py): sparse k=4 fold de-bias == dense to 1e-6 (point, fisher curvature, sandwich). Sparse non-fold path unchanged; verify_* pass. See RESULTS.md S9l. Co-Authored-By: Claude Opus 4.8 (1M context) --- rabbit/fitter.py | 36 ++++++++++++++++++ rabbit/tensorwriter.py | 10 ++--- tests/verify_sparsefold.py | 75 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 115 insertions(+), 6 deletions(-) create mode 100644 tests/verify_sparsefold.py diff --git a/rabbit/fitter.py b/rabbit/fitter.py index ff704fc..859fa91 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -1959,6 +1959,42 @@ def _compute_yields_noBBB( axis=-1, ) + if templates != "full": + # De-bias fold/half yields in sparse mode: the per-fold/half norm + # templates are stored DENSE, so scatter the (shared) systematic + # factor from the sparse logk into a dense [nbinsfull, nproc] grid + # and contract with the dense norm. Split-logk (per-fold logk) is + # not supported in sparse mode. + if templates == "A": + norm_dense = self.norm_A + elif templates == "B": + norm_dense = self.norm_B + else: + norm_dense = self.indata.norm_folds[fold_index] + if self.logk_folds_scaled is not None: + raise NotImplementedError( + "split-logk (per-fold logk) is not supported in sparse " + "mode; use shared logk (no fold_axis on systematics)." + ) + idx = self.indata.norm.indices # [nnz, 2] = (bin, proc) + shape = [self.indata.nbinsfull, self.indata.nproc] + if self.indata.systematic_type == "log_normal": + factor = tf.tensor_scatter_nd_update( + tf.ones(shape, dtype=norm_dense.dtype), idx, tf.exp(logsnorm) + ) + normcentral = norm_dense * rnorm * factor + else: # "normal": additive variation (norm-independent) + add = tf.tensor_scatter_nd_update( + tf.zeros(shape, dtype=norm_dense.dtype), idx, logsnorm + ) + normcentral = norm_dense * rnorm + add + if not (full or self.indata.nbinsmasked == 0): + normcentral = normcentral[: self.indata.nbins] + nexpcentral = tf.reduce_sum(normcentral, axis=-1) + if not compute_norm: + normcentral = None + return nexpcentral, normcentral + # Build a sparse [nbinsfull, nproc] tensor whose values absorb # the per-entry syst variation and the per-(bin, proc) POI # scaling rnorm. The sparsity pattern is unchanged from diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index f2f807c..32cd750 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -277,14 +277,12 @@ def _add_process_folded(self, h, name, channel, signal, variances, fold_axis): add_process path (so sumw / sumw2 / systematics are unchanged), and * stores the per-fold dense templates `[k, nbinschan]` for the cross-fit (two-half / k-fold) de-biasing in the fitter. - `k` is inferred from the size of `fold_axis`. Sparse storage is not - supported for folded processes. + `k` is inferred from the size of `fold_axis`. In sparse mode the FULL + template is still stored sparse (CSR); the per-fold fold templates are + densified ([k, nbinschan]) — the O(k*nparams*nbins) de-bias cost — and the + fitter uses a dense per-fold yield path. """ self._check_hist_and_channel(h, channel) - if self.sparse and self._issparse(h): - raise NotImplementedError( - "fold_axis is not supported for sparse histograms." - ) # resolve the fold axis (by name or index) and its size k ax_names = [a.name for a in h.axes] if isinstance(fold_axis, str): diff --git a/tests/verify_sparsefold.py b/tests/verify_sparsefold.py new file mode 100644 index 0000000..24d15b5 --- /dev/null +++ b/tests/verify_sparsefold.py @@ -0,0 +1,75 @@ +"""Validate MC-stat fold de-biasing in SPARSE mode: the sparse per-fold yield path +(scatter the shared systematic factor into a dense [nbinsfull,nproc] grid, contract +with the dense fold norm) must give EXACTLY the same de-biased point, fisher +curvature, and sandwich as the equivalent DENSE tensor. Run with OUT=.""" +import os, sys +import numpy as np +import hist +import importlib.util as _ilu + +RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") +sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) +_spec = _ilu.spec_from_file_location( + "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") +) +_rfm = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_rfm) +from rabbit import inputdata, fitter +from rabbit.tensorwriter import TensorWriter +from rabbit.param_models import helpers as ph + +NB, K = 40, 4 +OUT = os.environ.get("OUT", "/tmp/claude") + + +def build(sparse): + rng = np.random.default_rng(13) + nom = np.linspace(900, 1100, NB); data = nom.copy() + full = rng.poisson(nom) + folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in full], 0).T + ax = hist.axis.Regular(NB, 0, 1, name="x") + axf = hist.axis.IntCategory(list(range(K)), name="mcfold") + + def dh(v): + h = hist.Hist(ax, storage=hist.storage.Weight()) + h.view()["value"] = v; h.view()["variance"] = v + return h + + def fh(fl): + h = hist.Hist(ax, axf, storage=hist.storage.Weight()) + for f in range(K): + h.view()["value"][:, f] = fl[f]; h.view()["variance"][:, f] = fl[f] + return h + + x = np.linspace(0, 1, NB) + tw = TensorWriter(sparse=sparse) + tw.add_channel([ax], "ch0"); tw.add_data(dh(data), "ch0") + tw.add_process(fh(folds), "sig", "ch0", signal=True, fold_axis="mcfold") + tw.add_systematic(dh(nom * (1 + 0.1 * (x - 0.5))), "shp", "sig", "ch0") + fn = f"{OUT}/toy_sf_{'sp' if sparse else 'de'}.hdf5" + tw.write(outfolder=os.path.dirname(fn), outfilename=os.path.basename(fn)) + return fn + + +def run(fn): + a = _rfm.make_parser().parse_args( + [fn, "-o", f"{OUT}/vo", "-t", "0", "--chisqFit", "--noBinByBinStat", + "--allowNegativeParam", "--mcStatDebias", "kfold", "--covMode", "fisher"]) + ind = inputdata.FitInputData(fn, None) + f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) + f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + _, _, h = f.loss_val_grad_hess() + bread = f.fisher_curvature(h) + S = np.asarray(f.cov_twohalf_sandwich(bread, "fisher")) + return ind.sparse, f.x.numpy(), np.asarray(bread), S + + +if __name__ == "__main__": + sp, xs, bs, Ss = run(build(True)) + de, xd, bd, Sd = run(build(False)) + print(f"sparse={sp} x={np.round(xs,5)} | dense={not de or de} x={np.round(xd,5)}") + ok_x = np.allclose(xs, xd, rtol=1e-6) + ok_b = np.allclose(bs, bd, rtol=1e-6) + ok_s = np.allclose(Ss, Sd, rtol=1e-6) + print(f" point match: {ok_x} fisher curvature match: {ok_b} sandwich match: {ok_s}") + assert ok_x and ok_b and ok_s, "sparse fold de-bias must equal dense" + print("\n sparse fold de-biasing == dense (point, curvature, sandwich). passed.") From de1f9f829e4dda2a676e40820ae2239c063231a6 Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 22:11:06 +0000 Subject: [PATCH 11/15] mcstat: sparse split-logk via a per-fold logk delta over folded systs - The sparse fold-yield uses the shared systematic factor; split-logk needs per- fold logk. Instead of densifying the full [k,nbins,nproc,nsyst] logk, store a per-fold DELTA over the FOLDED systs only: delta=logk_fold-logk_full, dense [k,nbinsfull,nproc,n_folded], + folded global syst indices. - Writer: _add_systematic_folded computes logk_full densely and stores the delta; write() assembles hlogk_folds_delta + hmcstat_folded_syst_idx (sparse only). inputdata loads them. - Fitter: sparse fold-yield multiplies the shared log_normal factor by exp(delta.theta) for the active fold (curvature path); shared_logk+delta=fold_logk exactly. Objective/halves use shared logk (point de-biases nominal noise). log_normal symmetric only (asymmetric sparse split-logk raises). - Validated (tests/verify_sparse_splitlogk.py): sparse k=4 split-logk == dense to 1e-6 (point, fisher, sandwich). All other fold paths unchanged. See RESULTS.md S9m. Co-Authored-By: Claude Opus 4.8 (1M context) --- rabbit/fitter.py | 34 ++++++++++-- rabbit/inputdata.py | 11 ++++ rabbit/tensorwriter.py | 90 ++++++++++++++++++++++++++++++++ tests/verify_sparse_splitlogk.py | 77 +++++++++++++++++++++++++++ 4 files changed, 207 insertions(+), 5 deletions(-) create mode 100644 tests/verify_sparse_splitlogk.py diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 859fa91..3c92fdd 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -296,6 +296,20 @@ def __init__( "template noise is de-biased in the k-fold U-statistic " "curvature (--covMode fisher)." ) + # Sparse split-logk: per-fold logk delta for folded systs (applied as a + # exp(delta . theta) correction to the shared sparse systematic factor in + # the per-fold curvature yields). log_normal only. + self.logk_folds_delta = getattr(self.indata, "logk_folds_delta", None) + self.mcstat_folded_syst_idx = getattr( + self.indata, "mcstat_folded_syst_idx", None + ) + if ( + self.logk_folds_delta is not None + and self.indata.systematic_type != "log_normal" + ): + raise NotImplementedError( + "sparse split-logk is supported only for log_normal systematics." + ) if self.mcStatDebias in ("twoHalf", "kfold") and self.norm_A is None: logger.warning( f"--mcStatDebias {self.mcStatDebias} requested but no fold " @@ -1971,17 +1985,27 @@ def _compute_yields_noBBB( norm_dense = self.norm_B else: norm_dense = self.indata.norm_folds[fold_index] - if self.logk_folds_scaled is not None: - raise NotImplementedError( - "split-logk (per-fold logk) is not supported in sparse " - "mode; use shared logk (no fold_axis on systematics)." - ) idx = self.indata.norm.indices # [nnz, 2] = (bin, proc) shape = [self.indata.nbinsfull, self.indata.nproc] + # split-logk (sparse): multiply the shared log_normal factor by + # exp(delta . theta) over the folded systs for THIS fold, so the + # systematic-template noise is de-biased. Curvature path only + # (templates='fold'); halves use the shared logk. + split_corr = None + if ( + self.logk_folds_delta is not None + and templates == "fold" + ): + theta_folded = tf.gather(theta, self.mcstat_folded_syst_idx) + split_corr = tf.einsum( + "bpj,j->bp", self.logk_folds_delta[fold_index], theta_folded + ) if self.indata.systematic_type == "log_normal": factor = tf.tensor_scatter_nd_update( tf.ones(shape, dtype=norm_dense.dtype), idx, tf.exp(logsnorm) ) + if split_corr is not None: + factor = factor * tf.exp(split_corr) normcentral = norm_dense * rnorm * factor else: # "normal": additive variation (norm-independent) add = tf.tensor_scatter_nd_update( diff --git a/rabbit/inputdata.py b/rabbit/inputdata.py index 8a2ed48..d3722bf 100644 --- a/rabbit/inputdata.py +++ b/rabbit/inputdata.py @@ -191,6 +191,17 @@ def __init__(self, filename, pseudodata=None): else: self.logk_folds = None + # Sparse split-logk: per-fold logk delta for folded systs (dense, + # reduced syst dim) + the folded global syst indices. + if "hlogk_folds_delta" in f.keys(): + self.logk_folds_delta = maketensor(f["hlogk_folds_delta"]) + self.mcstat_folded_syst_idx = maketensor( + f["hmcstat_folded_syst_idx"] + ) + else: + self.logk_folds_delta = None + self.mcstat_folded_syst_idx = None + # compute indices for channels ibin = 0 for channel, info in self.channel_info.items(): diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 32cd750..13f9047 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -64,6 +64,10 @@ def __init__( # half-difference (same shape). self.dict_logkavg_folds = {} # [channel][process][syst] -> [k, nbinschan] self.dict_logkhalfdiff_folds = {} # [channel][process][syst] -> [k, nbinschan] + # Per-fold logk delta vs the FULL logk (= logkavg_fold - logkavg_full), + # used by the SPARSE split-logk path (multiply the shared systematic factor + # by exp(delta . theta)). [channel][process][syst] -> [k, nbinschan]. + self.dict_logk_folds_delta = {} self.dict_logkavg = {} # [channel][proc][syst] self.dict_logkhalfdiff = {} # [channel][proc][syst] self.dict_logkavg_indices = {} @@ -386,6 +390,7 @@ def add_channel(self, axes, name=None, masked=False, flow=False): self.fold_k[name] = {} self.dict_logkavg_folds[name] = {} self.dict_logkhalfdiff_folds[name] = {} + self.dict_logk_folds_delta[name] = {} self.dict_beta_variations[name] = {} # add masked channels last @@ -1454,6 +1459,32 @@ def _add_systematic_folded( flow = self.channels[channel]["flow"] systematic_type = "normal" if add_to_data_covariance else self.systematic_type + # full logkavg (vs the full nominal) for the sparse split-logk delta + norm_full = self.dict_norm[channel][process] + if self._issparse(norm_full): + _dense = np.zeros([norm_folds.shape[1]], self.dtype) + _dense[norm_full.indices] = norm_full.data + norm_full = _dense + if is_pair: + up_full = self.get_flat_values(h[0][{fold_idx: sum}], flow=flow) + down_full = self.get_flat_values(h[1][{fold_idx: sum}], flow=flow) + if as_difference: + up_full = norm_full + up_full + down_full = norm_full + down_full + lu_full = self.get_logk(up_full, norm_full, kfactor, systematic_type=systematic_type) + ld_full = -self.get_logk(down_full, norm_full, kfactor, systematic_type=systematic_type) + if symmetrize == "conservative": + logkavg_full = np.where(np.abs(lu_full) > np.abs(ld_full), lu_full, ld_full) + else: + logkavg_full = 0.5 * (lu_full + ld_full) + else: + syst_full = self.get_flat_values(h[{fold_idx: sum}], flow=flow) + if as_difference: + syst_full = norm_full + syst_full + logkavg_full = self.get_logk( + syst_full, norm_full, kfactor, systematic_type=systematic_type + ) + # per-fold logkavg (+ logkhalfdiff for asymmetric) vs the per-fold nominal logkavg_folds = np.zeros_like(norm_folds) logkhalfdiff_folds = np.zeros_like(norm_folds) @@ -1494,9 +1525,20 @@ def _add_systematic_folded( symmetrize=symmetrize, add_to_data_covariance=add_to_data_covariance, as_difference=as_difference, syst_axes=syst_axes, **kargs ) + if is_asym and self.sparse: + raise NotImplementedError( + "sparse split-logk does not support asymmetric folded systematics " + "(symmetric symmetrize modes only)." + ) + # store the per-fold logk under the resolved (full) systematic name self.dict_logkavg_folds[channel].setdefault(process, {}) self.dict_logkavg_folds[channel][process][name] = logkavg_folds.astype(self.dtype) + # per-fold delta vs the full logk (for the sparse split-logk path) + self.dict_logk_folds_delta[channel].setdefault(process, {}) + self.dict_logk_folds_delta[channel][process][name] = ( + (logkavg_folds - logkavg_full[None, :]).astype(self.dtype) + ) if is_asym: self.dict_logkhalfdiff_folds[channel].setdefault(process, {}) self.dict_logkhalfdiff_folds[channel][process][name] = ( @@ -1944,6 +1986,43 @@ def write(self, outfolder="./", outfilename="rabbit_input.hdf5", meta_data_dict= systs = self.get_systs() nsyst = len(systs) + # SPARSE split-logk: assemble the per-fold logk DELTA for folded systs + # [k, nbinsfull, nproc, n_folded] + the folded global syst indices, so the + # sparse fold-yield path can multiply the shared factor by exp(delta.theta). + logk_folds_delta = None + mcstat_folded_syst_idx = None + if norm_folds is not None and self.sparse: + folded_systs = sorted( + { + s + for c in self.channels + for p in self.dict_logk_folds_delta[c] + for s in self.dict_logk_folds_delta[c][p] + }, + key=lambda s: systs.index(s), + ) + if folded_systs: + kfold = norm_folds.shape[0] + mcstat_folded_syst_idx = np.array( + [systs.index(s) for s in folded_systs], dtype=np.int64 + ) + logk_folds_delta = np.zeros( + [kfold, nbinsfull, nproc, len(folded_systs)], self.dtype + ) + ibin = 0 + for chan in self.channels.keys(): + nbinschan = self.nbinschan[chan] + for iproc, proc in enumerate(procs): + if proc not in self.dict_norm[chan]: + continue + deltas = self.dict_logk_folds_delta[chan].get(proc, {}) + for j, s in enumerate(folded_systs): + if s in deltas: + logk_folds_delta[ + :, ibin : ibin + nbinschan, iproc, j + ] = deltas[s] + ibin += nbinschan + if self.symmetric_tensor: logger.info("No asymmetric systematics - write fully symmetric tensor") @@ -2455,6 +2534,17 @@ def create_dataset( f.attrs["mcstat_fold_k"] = int(norm_folds.shape[0]) norm_folds = None + # SPARSE split-logk: per-fold logk delta for folded systs + their indices. + if logk_folds_delta is not None: + nbytes += h5pyutils_write.writeFlatInChunks( + logk_folds_delta, f, "hlogk_folds_delta", maxChunkBytes=self.chunkSize + ) + nbytes += h5pyutils_write.writeFlatInChunks( + mcstat_folded_syst_idx, f, "hmcstat_folded_syst_idx", + maxChunkBytes=self.chunkSize, + ) + logk_folds_delta = None + if self.sparse: nbytes += h5pyutils_write.writeSparse( norm_sparse_indices, diff --git a/tests/verify_sparse_splitlogk.py b/tests/verify_sparse_splitlogk.py new file mode 100644 index 0000000..a0a29f4 --- /dev/null +++ b/tests/verify_sparse_splitlogk.py @@ -0,0 +1,77 @@ +"""Validate SPARSE split-logk: a folded (split-logk) systematic in sparse mode must +give EXACTLY the same de-biased point, fisher curvature, and sandwich as the +equivalent DENSE tensor. The sparse path multiplies the shared systematic factor by +exp(delta . theta) over the folded systs (delta = logk_fold - logk_full), curvature +path only. log_normal symmetric folded systematics. Run with OUT=.""" +import os, sys +import numpy as np +import hist +import importlib.util as _ilu + +RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") +sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) +_spec = _ilu.spec_from_file_location( + "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") +) +_rfm = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_rfm) +from rabbit import inputdata, fitter +from rabbit.tensorwriter import TensorWriter +from rabbit.param_models import helpers as ph + +NB, K = 40, 4 +OUT = os.environ.get("OUT", "/tmp/claude") + + +def build(sparse): + rng = np.random.default_rng(21) + nom = np.linspace(900, 1100, NB); data = nom.copy() + full = rng.poisson(nom) + folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in full], 0).T + x = np.linspace(0, 1, NB); w = 1 + 0.15 * (x - 0.5) + upf = np.stack([rng.poisson(np.maximum(folds[f] * w, 0)) for f in range(K)], 0) + ax = hist.axis.Regular(NB, 0, 1, name="x") + axf = hist.axis.IntCategory(list(range(K)), name="mcfold") + + def dh(v): + h = hist.Hist(ax, storage=hist.storage.Weight()) + h.view()["value"] = v; h.view()["variance"] = v + return h + + def fh(fl): + h = hist.Hist(ax, axf, storage=hist.storage.Weight()) + for f in range(K): + h.view()["value"][:, f] = fl[f]; h.view()["variance"][:, f] = fl[f] + return h + + tw = TensorWriter(sparse=sparse) + tw.add_channel([ax], "ch0"); tw.add_data(dh(data), "ch0") + tw.add_process(fh(folds), "sig", "ch0", signal=True, fold_axis="mcfold") + tw.add_systematic(fh(upf), "tilt", "sig", "ch0", fold_axis="mcfold") # split-logk + fn = f"{OUT}/toy_spl_{'sp' if sparse else 'de'}.hdf5" + tw.write(outfolder=os.path.dirname(fn), outfilename=os.path.basename(fn)) + return fn + + +def run(fn): + a = _rfm.make_parser().parse_args( + [fn, "-o", f"{OUT}/vo", "-t", "0", "--chisqFit", "--noBinByBinStat", + "--allowNegativeParam", "--mcStatDebias", "kfold", "--covMode", "fisher"]) + ind = inputdata.FitInputData(fn, None) + f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) + f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + _, _, h = f.loss_val_grad_hess() + bread = f.fisher_curvature(h) + S = np.asarray(f.cov_twohalf_sandwich(bread, "fisher")) + return ind.sparse, f.x.numpy(), np.asarray(bread), S + + +if __name__ == "__main__": + _, xs, bs, Ss = run(build(True)) + _, xd, bd, Sd = run(build(False)) + ok = (np.allclose(xs, xd, rtol=1e-6) and np.allclose(bs, bd, rtol=1e-6) + and np.allclose(Ss, Sd, rtol=1e-6)) + print(f"point/fisher/sandwich match sparse==dense: " + f"{np.allclose(xs,xd,1e-6)}/{np.allclose(bs,bd,1e-6)}/{np.allclose(Ss,Sd,1e-6)}") + print(f" tilt sandwich sparse={np.sqrt(Ss[1,1]):.5f} dense={np.sqrt(Sd[1,1]):.5f}") + assert ok, "sparse split-logk must equal dense" + print("\n sparse split-logk == dense (point, curvature, sandwich). passed.") From 36472a15503b46b6afcf1aaedb603eb84fa3a8ef Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Sun, 14 Jun 2026 23:51:17 +0000 Subject: [PATCH 12/15] mcstat: native template-fluctuating coverage harness + test rabbit's built-in toys fluctuate DATA only; MC-stat coverage needs the TEMPLATES fluctuated. tests/mcstat_coverage.py fluctuates both per pseudo-experiment, builds a rabbit tensor, fits, and measures POI coverage (configurable debias x covMode x debiasCov x BB-lite x chisq/Poisson). Near-degenerate slides toy with NONZERO degenerate-direction truth (rnorm_true=[1.3,0.7]) so the attenuation bias is real. Findings: de-bias REMOVES the central-value bias (naive med_bias ~ -0.29 on a true 0.42 -> ~0) and inflates sigma toward sigma_inf (validates the first two original criteria); the raw curvature/sandwich sigma still underestimates the de-biased point's ENSEMBLE scatter by ~20-25% (MC-noise variance in the score not in the single-toy meat) -> coverage ~0.4-0.55 not 0.683 -> the Bartlett/calibration factor is needed for exact coverage (CONFIRMS the known limitation, now measured natively). tests/verify_coverage.py asserts the robust facts (bias removed, sigma inflated, coverage improved). See RESULTS.md S9n. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/mcstat_coverage.py | 173 +++++++++++++++++++++++++++++++++++++++ tests/verify_coverage.py | 47 +++++++++++ 2 files changed, 220 insertions(+) create mode 100644 tests/mcstat_coverage.py create mode 100644 tests/verify_coverage.py diff --git a/tests/mcstat_coverage.py b/tests/mcstat_coverage.py new file mode 100644 index 0000000..c73680b --- /dev/null +++ b/tests/mcstat_coverage.py @@ -0,0 +1,173 @@ +"""Native template-fluctuating ensemble coverage harness for the MC-stat de-biasing. + +Unlike rabbit's built-in toys (which fluctuate the DATA only), this fluctuates the +TEMPLATES (the MC noise floor) AND the data per pseudo-experiment, builds a rabbit +tensor, fits it, and measures the coverage of a POI combination. This is the single +metric the de-bias is supposed to restore (none undercovers; debias+sandwich ~0.683). + +Two near-degenerate processes (corr tunable via STEP), finite MC (oversample:1), +linear POIs (rnorm_true = 1). Coverage is measured for the degenerate direction +(rnorm0 - rnorm1, truth 0) and the well-constrained sum. + +Usage: + OUT=/tmp/cov python3 tests/mcstat_coverage.py # prints a table + from mcstat_coverage import run_coverage # programmatic +""" +import os, sys +import numpy as np +import hist +import importlib.util as _ilu + +RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") +sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) +_spec = _ilu.spec_from_file_location( + "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") +) +_rfm = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_rfm) +from rabbit import inputdata, fitter +from rabbit.tensorwriter import TensorWriter +from rabbit.param_models import helpers as ph + +OUT = os.environ.get("OUT", "/tmp/claude") +# The near-degenerate 2-process toy from the slides (RESULTS §1): a flat process +# and a "step" process identical to it in the 2nd half and +3% in the 1st half. +# This sits in the regime where the de-biased CURVATURE ~ sigma_inf (M/lambda_min +# ~0.6, so H-M is comfortably positive and the de-biased point is well-behaved) -- +# the naive sigma is ~60% of true so the naive interval undercovers (~0.45), and +# the de-bias restores it. +NB = int(os.environ.get("NB", "200")) +FLAT = 4962.0 +HI = float(os.environ.get("HI", "5112")) # step height; larger = less degenerate +N0 = np.full(NB, FLAT) +N1 = np.concatenate([np.full(NB // 2, HI), np.full(NB // 2, FLAT)]) +# TRUE signal strengths. The degenerate (difference) direction must have a NONZERO +# true value, else the MC-stat attenuation (which biases toward 0) leaves it ~0 and +# there is no central-value bias for the de-bias to correct. R0!=R1 puts a real +# value on the difference so naive attenuates it (biased) and the de-bias restores it. +R0 = float(os.environ.get("R0", "1.3")) +R1 = float(os.environ.get("R1", "0.7")) +RTRUE = np.array([R0, R1]) +MU = R0 * N0 + R1 * N1 # true total (data expectation) +AX = hist.axis.Regular(NB, 0.0, 1.0, name="x") +DDIF = np.array([1.0, -1.0]) / np.sqrt(2.0) +DSUM = np.array([1.0, 1.0]) / np.sqrt(2.0) + + +def _whist(values, variances): + h = hist.Hist(AX, storage=hist.storage.Weight()) + h.view()["value"] = values + h.view()["variance"] = variances + return h + + +def _fhist(folds): + axf = hist.axis.IntCategory(list(range(folds.shape[0])), name="mcfold") + h = hist.Hist(AX, axf, storage=hist.storage.Weight()) + for f in range(folds.shape[0]): + h.view()["value"][:, f] = folds[f] + h.view()["variance"][:, f] = folds[f] + return h + + +def _build(fn, T0, T1, folds0, folds1, data, debias, bbb, poisson): + tw = TensorWriter(sparse=False) + tw.add_channel([AX], "ch0") + tw.add_data(_whist(data, data), "ch0") + if debias in ("twoHalf", "kfold"): + tw.add_process(_fhist(folds0), "p0", "ch0", signal=True, fold_axis="mcfold") + tw.add_process(_fhist(folds1), "p1", "ch0", signal=True, fold_axis="mcfold") + else: + tw.add_process(_whist(T0, T0), "p0", "ch0", signal=True) + tw.add_process(_whist(T1, T1), "p1", "ch0", signal=True) + if debias == "continuousM": + # M_pp = sum_b sumw2_p / var_b ; var = data (chisq) or data+sumw2 (BB-lite) + var = data + (T0 + T1) if bbb else data + M = np.diag([np.sum(T0 / var), np.sum(T1 / var)]) + tw.add_mc_stat_moment(M, ["p0", "p1"]) + tw.write(outfolder=os.path.dirname(fn), outfilename=os.path.basename(fn)) + + +def _fit(fn, debias, covMode, debiasCov, bbb, poisson): + argv = [fn, "-o", f"{OUT}/vo", "-t", "0", "--allowNegativeParam", + "--mcStatDebias", debias, "--covMode", covMode, + "--mcStatDebiasCov", debiasCov] + argv += [] if poisson else ["--chisqFit"] + if not bbb: + argv.append("--noBinByBinStat") + a = _rfm.make_parser().parse_args(argv) + ind = inputdata.FitInputData(fn, None) + f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) + f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + _, grad, hess = f.loss_val_grad_hess() + bread = f.fisher_curvature(hess) if covMode == "fisher" else hess + _, cov_curv = f.edmval_cov(grad, bread) + if debiasCov == "sandwich" and debias == "continuousM": + C = np.asarray(f.cov_mcstat_sandwich(bread)) + elif debiasCov == "sandwich" and debias in ("twoHalf", "kfold"): + C = np.asarray(f.cov_twohalf_sandwich(bread, covMode)) + else: + C = np.asarray(cov_curv) + return f.x.numpy(), C, float(np.max(np.abs(grad.numpy()))) + + +def run_coverage(debias="none", covMode="observed", debiasCov="sandwich", + bbb=False, poisson=False, ntoy=200, k=2, seed=1): + rng = np.random.default_rng(seed) + fn = f"{OUT}/cov_{debias}_{covMode}_{debiasCov}_{int(bbb)}_{int(poisson)}.hdf5" + cov_dif = cov_sum = 0 + bias_dif = [] + res_dif = [] + sig_dif = [] + nbad = 0 + for _ in range(ntoy): + folds0 = np.stack([rng.poisson(N0 / k) for _ in range(k)], 0).astype(float) + folds1 = np.stack([rng.poisson(N1 / k) for _ in range(k)], 0).astype(float) + T0, T1 = folds0.sum(0), folds1.sum(0) + data = rng.poisson(MU).astype(float) + try: + _build(fn, T0, T1, folds0, folds1, data, debias, bbb, poisson) + x, C, gmax = _fit(fn, debias, covMode, debiasCov, bbb, poisson) + if gmax > 1e-2 or not np.all(np.linalg.eigvalsh(C) > 0): + nbad += 1; continue + except Exception: + nbad += 1; continue + r = x - RTRUE # residual vs true signal strengths + s_dif = np.sqrt(DDIF @ C @ DDIF) + s_sum = np.sqrt(DSUM @ C @ DSUM) + cov_dif += abs(DDIF @ r) < s_dif + cov_sum += abs(DSUM @ r) < s_sum + res_dif.append(DDIF @ r); sig_dif.append(s_dif) + n = ntoy - nbad + res = np.array(res_dif) + return dict( + cov_dif=cov_dif / n, cov_sum=cov_sum / n, n=n, nbad=nbad, + # median residual = robust bias (mean is corrupted by the heavy tails the + # de-bias develops near the singular limit); rms_res = the actual point + # scatter (vs med_sig = the reported uncertainty -> coverage = rms/med). + bias_dif=float(np.median(res)), mean_res=float(np.mean(res)), + rms_res=float(np.std(res)), med_sig=float(np.median(sig_dif)), + ) + + +if __name__ == "__main__": + ntoy = int(os.environ.get("NTOY", "300")) + configs = [ + ("none observed sandwich chisq", dict(debias="none")), + ("continuousM observed sandwich chisq", dict(debias="continuousM")), + ("continuousM observed curvature chisq", dict(debias="continuousM", debiasCov="curvature")), + ("continuousM fisher sandwich chisq", dict(debias="continuousM", covMode="fisher")), + ("twoHalf observed sandwich chisq", dict(debias="twoHalf")), + ("twoHalf fisher sandwich chisq", dict(debias="twoHalf", covMode="fisher")), + ("kfold(8) fisher sandwich chisq", dict(debias="kfold", covMode="fisher", k=8)), + ("none observed sandwich chisq BBB", dict(debias="none", bbb=True)), + ("continuousM observed sandwich chisq BBB", dict(debias="continuousM", bbb=True)), + ("twoHalf observed sandwich chisq BBB", dict(debias="twoHalf", bbb=True)), + ("none observed sandwich POISSON", dict(debias="none", poisson=True)), + ("twoHalf observed sandwich POISSON", dict(debias="twoHalf", poisson=True)), + ] + print(f"MC-stat coverage ensemble (ntoy={ntoy}, target cov(dif)=0.683)\n") + print(f"{'config':40s}{'cov(dif)':>9s}{'cov(sum)':>9s}{'bias(dif)':>11s}{'med_sig':>9s}{'nbad':>6s}") + for label, kw in configs: + r = run_coverage(ntoy=ntoy, **kw) + print(f"{label:40s}{r['cov_dif']:9.3f}{r['cov_sum']:9.3f}" + f"{r['bias_dif']:+11.4f}{r['med_sig']:9.4f}{r['nbad']:6d}") diff --git a/tests/verify_coverage.py b/tests/verify_coverage.py new file mode 100644 index 0000000..4650260 --- /dev/null +++ b/tests/verify_coverage.py @@ -0,0 +1,47 @@ +"""Fast assertion test for the MC-stat de-biasing, using the native template- +fluctuating coverage harness (tests/mcstat_coverage.py). Asserts the ROBUST, +reproducible facts on the near-degenerate slides toy with a NONZERO degenerate- +direction truth (rnorm_true=[1.3,0.7]): + + 1. naive (no de-bias) is badly BIASED by the MC-stat attenuation and undercovers; + 2. continuous-M REMOVES the central-value bias; + 3. it INFLATES the uncertainty estimate toward sigma_inf; + 4. and IMPROVES the coverage of the degenerate POI direction. + +Note (RESULTS §9n): the raw sandwich sigma still underestimates the de-biased +point's ensemble scatter by ~20-25% (the MC-noise variance in the score is not +captured by the single-toy meat), so exact 68.3% coverage needs the Bartlett/ +calibration factor -- this test asserts the de-bias DIRECTION (bias removed, sigma +inflated, coverage improved), not exact 0.683. Run with OUT=. +""" +import os +os.environ.setdefault("NB", "200") +os.environ.setdefault("HI", "5112") # near-degenerate (large attenuation) +os.environ.setdefault("R0", "1.3") +os.environ.setdefault("R1", "0.7") + +import tests.mcstat_coverage as mc + +NTOY = int(os.environ.get("NTOY", "50")) + +if __name__ == "__main__": + print(f"coverage harness, slides toy, diff_true={mc.DDIF @ mc.RTRUE:.3f}, ntoy={NTOY}") + none = mc.run_coverage(ntoy=NTOY, seed=11, debias="none") + cmS = mc.run_coverage(ntoy=NTOY, seed=11, debias="continuousM", debiasCov="sandwich") + for lab, r in [("none", none), ("continuousM sand", cmS)]: + print(f" {lab:18s} cov(dif)={r['cov_dif']:.3f} med_bias={r['bias_dif']:+.4f} " + f"med_sig={r['med_sig']:.4f} rms_res={r['rms_res']:.4f}") + + # 1. naive is badly biased by attenuation (degenerate direction pulled to 0) + assert none["bias_dif"] < -0.10, f"naive should be attenuated, got {none['bias_dif']}" + # 2. de-bias removes the central-value bias + assert abs(cmS["bias_dif"]) < 0.4 * abs(none["bias_dif"]), \ + f"de-bias should remove the bias: {none['bias_dif']} -> {cmS['bias_dif']}" + # 3. de-bias inflates the uncertainty estimate toward sigma_inf + assert cmS["med_sig"] > 1.3 * none["med_sig"], \ + f"de-bias should inflate sigma: {none['med_sig']} -> {cmS['med_sig']}" + # 4. de-bias improves coverage of the degenerate POI + assert cmS["cov_dif"] > none["cov_dif"] + 0.10, \ + f"de-bias should improve coverage: {none['cov_dif']} -> {cmS['cov_dif']}" + print("\n coverage de-bias checks passed (bias removed, sigma inflated, " + "coverage improved).") From 15615a86e98e5c6f10cf201d00718f92b99c082a Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Mon, 15 Jun 2026 00:27:07 +0000 Subject: [PATCH 13/15] mcstat: coverage harness - robust direction filter + n=0 handling; reconcile proper sandwich - run_coverage: require positive variance in the MEASURED directions (not full PD), so BB-lite's rank-deficient-other-direction covariance on near-degenerate toys is not spuriously rejected; return nan instead of crashing when n=0. - Reconciliation (RESULTS S9o): S9n measured the BARE sandwich (--noBinByBinStat, meat omits sigma_MC) -> undercovers (cov 0.51), calibration gap. The PROPER sandwich with BB-lite (meat = A+M = H_bb includes sigma_MC) COVERS 0.67 ~ 0.683 with NO calibration (med_sig 0.041 matches the ensemble scatter 0.040) -- matches the numpy result. rabbit's two-half sandwich uses the Hessian meat, not the cross-fold score variance, so it is not yet self-calibrating without BB-lite. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/mcstat_coverage.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/mcstat_coverage.py b/tests/mcstat_coverage.py index c73680b..3bb255b 100644 --- a/tests/mcstat_coverage.py +++ b/tests/mcstat_coverage.py @@ -127,17 +127,26 @@ def run_coverage(debias="none", covMode="observed", debiasCov="sandwich", try: _build(fn, T0, T1, folds0, folds1, data, debias, bbb, poisson) x, C, gmax = _fit(fn, debias, covMode, debiasCov, bbb, poisson) - if gmax > 1e-2 or not np.all(np.linalg.eigvalsh(C) > 0): + v_dif = DDIF @ C @ DDIF + v_sum = DSUM @ C @ DSUM + # require convergence and POSITIVE variance in the measured directions + # (a rank-deficient full matrix is fine as long as the projection is ok; + # BB-lite on a near-degenerate toy can null the other direction). + if gmax > 1e-2 or v_dif <= 0 or v_sum <= 0 or not np.isfinite(v_dif): nbad += 1; continue except Exception: nbad += 1; continue r = x - RTRUE # residual vs true signal strengths - s_dif = np.sqrt(DDIF @ C @ DDIF) - s_sum = np.sqrt(DSUM @ C @ DSUM) + s_dif = np.sqrt(v_dif) + s_sum = np.sqrt(v_sum) cov_dif += abs(DDIF @ r) < s_dif cov_sum += abs(DSUM @ r) < s_sum res_dif.append(DDIF @ r); sig_dif.append(s_dif) n = ntoy - nbad + if n == 0: + return dict(cov_dif=float("nan"), cov_sum=float("nan"), n=0, nbad=nbad, + bias_dif=float("nan"), mean_res=float("nan"), + rms_res=float("nan"), med_sig=float("nan")) res = np.array(res_dif) return dict( cov_dif=cov_dif / n, cov_sum=cov_sum / n, n=n, nbad=nbad, From d6af02bd915e0cdb8ebfa09f04e26d974b71783b Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Mon, 15 Jun 2026 00:53:38 +0000 Subject: [PATCH 14/15] mcstat: add --mcStatDebiasCov dataPropagated (Huber-White Var(score) sandwich) - Fitter.cov_dataprop_sandwich: delta-method sandwich (RESULTS S7d/demo23) propagating the per-bin DATA variance and per-(bin,proc) TEMPLATE variance (sumw2) through the de-biased estimator by implicit diff (Sigma = dx/dnobs diag(Var_nobs) dx/dnobs^T + dx/dnorm diag(sumw2) dx/dnorm^T). Continuous-M, dense (watches indata.norm). - --mcStatDebiasCov gains 'dataPropagated'; rabbit_fit wires it. - Validated: conservative / OVER-covers (slides toy: cov 0.79 vs analytic 0.50, none 0.06) -- matches S7d (~0.81). Corrects S9o: the score-variance meat is NOT self-calibrating, it over-covers (needs a downward k); the calibration-free meat is the analytic/BB-lite one. CLI runs end-to-end. See RESULTS.md S9p. Co-Authored-By: Claude Opus 4.8 (1M context) --- bin/rabbit_fit.py | 7 ++++++- rabbit/fitter.py | 43 ++++++++++++++++++++++++++++++++++++++++ rabbit/parsing.py | 8 ++++++-- tests/mcstat_coverage.py | 4 +++- 4 files changed, 58 insertions(+), 4 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index e3241fb..d000cef 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -564,7 +564,12 @@ def fit(args, fitter, ws, dofit=True): (_debias == "continuousM" and getattr(fitter, "mcstat_M", None) is not None) or (_debias in ("twoHalf", "kfold") and getattr(fitter, "norm_A", None) is not None) ) - if _debiascov == "sandwich" and _is_debiased: + if _debiascov == "dataPropagated" and _is_debiased: + # Huber-White delta-method sandwich (conservative; over-covers). + cov = fitter.cov_dataprop_sandwich(cov_curv) + logger.info("Reporting data-propagated Var(score) sandwich " + "(Huber-White, conservative/over-covers; RESULTS §7d)") + elif _debiascov == "sandwich" and _is_debiased: if _debias == "continuousM": cov = fitter.cov_mcstat_sandwich(bread) logger.info("Reporting continuous-M sandwich covariance " diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 3c92fdd..83574e0 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -1180,6 +1180,49 @@ def cov_twohalf_sandwich(self, A, covMode="observed"): Ainv = tf.linalg.inv(A) return Ainv @ tf.cast(H, Ainv.dtype) @ Ainv + def cov_dataprop_sandwich(self, Ainv): + """Data-propagated Var(score) sandwich (Huber-White / delta-method). + + Propagates the per-bin DATA variance and the per-(bin,proc) TEMPLATE + variance (sumw2 = the MC-stat variance) through the de-biased estimator + theta_hat(nobs, norm) by implicit differentiation: + + Sigma = dx/dnobs . diag(Var[nobs]) . (dx/dnobs)^T + + dx/dnorm . diag(sumw2) . (dx/dnorm)^T + + with dx/dv = -A^-1 d^2L/dx dv (A = de-biased bread; `Ainv` = A^-1). The + FIRST term equals the analytic A^-1 H A^-1 (data Fisher meat); the SECOND + adds the template-noise propagation. This is the conservative route: it + OVER-covers (~0.81 in the numpy tests, RESULTS §7d) and needs a downward + calibration k~0.81, vs the calibration-free analytic/BB-lite meat (~0.67). + Dense only (watches indata.norm); the de-biased POINT (M in the + minimization) is still required for an unbiased centre (§7d).""" + if self.indata.sparse: + raise NotImplementedError( + "data-propagated sandwich is dense-only (it differentiates the " + "loss w.r.t. indata.norm)." + ) + Ainv = tf.cast(Ainv, self.indata.dtype) + norm = self.indata.norm + with tf.GradientTape() as t2: + t2.watch([self.nobs, norm]) + with tf.GradientTape() as t1: + t1.watch([self.nobs, norm]) + val = self._compute_loss() + grad = t1.gradient(val, self.x) + pd2ldxdnobs, pd2ldxdnorm = t2.jacobian( + grad, [self.nobs, norm], unconnected_gradients="zero" + ) + dxdnobs = -Ainv @ pd2ldxdnobs # [npar, nbins] + dxdnorm = -Ainv @ tf.reshape( + pd2ldxdnorm, [tf.shape(pd2ldxdnorm)[0], -1] + ) # [npar, nbinsfull*nproc] + var_nobs = self.varnobs if self.varnobs is not None else self.nobs + sumw2_flat = tf.reshape(self.indata.sumw2, [-1]) + return (dxdnobs * var_nobs[None, :]) @ tf.transpose(dxdnobs) + ( + dxdnorm * sumw2_flat[None, :] + ) @ tf.transpose(dxdnorm) + def cov_mcstat_sandwich(self, A): """Continuous-M robust (sandwich) covariance. diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 14770bf..1c0cf85 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -392,10 +392,14 @@ def common_parser(): parser.add_argument( "--mcStatDebiasCov", default="sandwich", - choices=["curvature", "sandwich"], + choices=["curvature", "sandwich", "dataPropagated"], help="Covariance reported when --mcStatDebias != none. 'curvature' = inverse " "de-biased Hessian A^-1 (undercovers). 'sandwich' = A^-1 H A^-1 = " - "A^-1 + A^-1 M A^-1 (covering robust covariance). Default 'sandwich'.", + "A^-1 + A^-1 M A^-1 (analytic meat; calibration-free ~0.68 with BB-lite, " + "undercovers without). 'dataPropagated' = Huber-White delta-method sandwich " + "propagating the data AND template (sumw2) variances through theta_hat " + "(continuous-M, dense): CONSERVATIVE, over-covers (~0.81, needs a downward " + "calibration; RESULTS §7d). Default 'sandwich'.", ) parser.add_argument( "--mcStatKfold", diff --git a/tests/mcstat_coverage.py b/tests/mcstat_coverage.py index 3bb255b..657b2ea 100644 --- a/tests/mcstat_coverage.py +++ b/tests/mcstat_coverage.py @@ -101,7 +101,9 @@ def _fit(fn, debias, covMode, debiasCov, bbb, poisson): _, grad, hess = f.loss_val_grad_hess() bread = f.fisher_curvature(hess) if covMode == "fisher" else hess _, cov_curv = f.edmval_cov(grad, bread) - if debiasCov == "sandwich" and debias == "continuousM": + if debiasCov == "dataPropagated" and debias != "none": + C = np.asarray(f.cov_dataprop_sandwich(cov_curv)) + elif debiasCov == "sandwich" and debias == "continuousM": C = np.asarray(f.cov_mcstat_sandwich(bread)) elif debiasCov == "sandwich" and debias in ("twoHalf", "kfold"): C = np.asarray(f.cov_twohalf_sandwich(bread, covMode)) From 78105f75d58ba9c12a5e1d15aa7bf46e7e55a76e Mon Sep 17 00:00:00 2001 From: Josh Bendavid Date: Mon, 15 Jun 2026 01:16:13 +0000 Subject: [PATCH 15/15] mcstat: apply isort/black formatting and remove unused imports/vars (CI lint) - isort (--profile black --line-length 88) + black across the mcstat source and test files; remove an unused 'import tensorflow as tf' (F401) and a dead 'bias_dif' list. No functional change (verify_twohalf still matches numpy). Co-Authored-By: Claude Opus 4.8 (1M context) --- bin/rabbit_fit.py | 33 +++++--- rabbit/fitter.py | 23 ++---- rabbit/inputdata.py | 9 +-- rabbit/tensorwriter.py | 85 +++++++++++++++----- rabbit/workspace.py | 5 +- tests/mcstat_coverage.py | 128 +++++++++++++++++++++++-------- tests/toy_kfold.py | 12 ++- tests/toy_mcstat.py | 23 ++++-- tests/toy_splitlogk.py | 28 ++++--- tests/toy_twohalf.py | 23 ++++-- tests/verify_bblite.py | 51 ++++++++---- tests/verify_coverage.py | 43 +++++++---- tests/verify_kfold.py | 58 ++++++++++---- tests/verify_mcstat.py | 60 ++++++++++----- tests/verify_sparse_splitlogk.py | 69 ++++++++++++----- tests/verify_sparsefold.py | 57 ++++++++++---- tests/verify_twohalf.py | 80 +++++++++++++------ 17 files changed, 551 insertions(+), 236 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index d000cef..42a13c1 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -561,23 +561,32 @@ def fit(args, fitter, ws, dofit=True): _debias = getattr(fitter, "mcStatDebias", "none") _debiascov = getattr(fitter, "mcStatDebiasCov", "sandwich") _is_debiased = ( - (_debias == "continuousM" and getattr(fitter, "mcstat_M", None) is not None) - or (_debias in ("twoHalf", "kfold") and getattr(fitter, "norm_A", None) is not None) + _debias == "continuousM" + and getattr(fitter, "mcstat_M", None) is not None + ) or ( + _debias in ("twoHalf", "kfold") + and getattr(fitter, "norm_A", None) is not None ) if _debiascov == "dataPropagated" and _is_debiased: # Huber-White delta-method sandwich (conservative; over-covers). cov = fitter.cov_dataprop_sandwich(cov_curv) - logger.info("Reporting data-propagated Var(score) sandwich " - "(Huber-White, conservative/over-covers; RESULTS §7d)") + logger.info( + "Reporting data-propagated Var(score) sandwich " + "(Huber-White, conservative/over-covers; RESULTS §7d)" + ) elif _debiascov == "sandwich" and _is_debiased: if _debias == "continuousM": cov = fitter.cov_mcstat_sandwich(bread) - logger.info("Reporting continuous-M sandwich covariance " - f"(A^-1 + A^-1 M A^-1, covMode={_covmode})") + logger.info( + "Reporting continuous-M sandwich covariance " + f"(A^-1 + A^-1 M A^-1, covMode={_covmode})" + ) else: cov = fitter.cov_twohalf_sandwich(bread, covMode=_covmode) - logger.info("Reporting two-half sandwich covariance " - f"(A^-1 H A^-1, covMode={_covmode})") + logger.info( + "Reporting two-half sandwich covariance " + f"(A^-1 H A^-1, covMode={_covmode})" + ) ws.add_cov_hist(cov) @@ -604,7 +613,9 @@ def fit(args, fitter, ws, dofit=True): tf.linalg.diag_part(fitter.cov) - tf.linalg.diag_part(cov_curv) ] ws.add_impacts_hists( - *fitter.impacts_parms(bread, cov=cov_curv, extra_group_vars=extra_vars) + *fitter.impacts_parms( + bread, cov=cov_curv, extra_group_vars=extra_vars + ) ) del hess @@ -612,7 +623,9 @@ def fit(args, fitter, ws, dofit=True): if args.globalImpacts: # de-biased fit: decompose the de-biased curvature consistently ws.add_impacts_hists( - *fitter.global_impacts_parms(cov=cov_curv if _is_debiased else None), + *fitter.global_impacts_parms( + cov=cov_curv if _is_debiased else None + ), base_name="global_impacts", global_impacts=True, ) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 83574e0..d3548dd 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -266,9 +266,7 @@ def __init__( # syst (and asym) axes. if self.indata.systematic_type == "normal" and self.param_model.nparams > 0: rnorm_init = tf.broadcast_to( - self.param_model.compute( - self.param_model.xparamdefault, full=True - ), + self.param_model.compute(self.param_model.xparamdefault, full=True), [self.indata.nbinsfull, self.indata.nproc], ) ntrail = len(logk_folds.shape) - 3 # 1 (sym) or 2 (asym) @@ -316,11 +314,7 @@ def __init__( "templates (hnorm_folds) found in the input; two-half de-biasing " "is inactive. Add processes with a fold_axis in the TensorWriter." ) - if ( - self.n_folds is not None - and self.n_folds >= 3 - and self.covMode != "fisher" - ): + if self.n_folds is not None and self.n_folds >= 3 and self.covMode != "fisher": logger.info( f"k-fold de-biasing with k={self.n_folds} folds: the k-fold " "AVERAGING (complete U-statistic, lower curvature-estimate " @@ -1213,10 +1207,10 @@ def cov_dataprop_sandwich(self, Ainv): pd2ldxdnobs, pd2ldxdnorm = t2.jacobian( grad, [self.nobs, norm], unconnected_gradients="zero" ) - dxdnobs = -Ainv @ pd2ldxdnobs # [npar, nbins] + dxdnobs = -Ainv @ pd2ldxdnobs # [npar, nbins] dxdnorm = -Ainv @ tf.reshape( pd2ldxdnorm, [tf.shape(pd2ldxdnorm)[0], -1] - ) # [npar, nbinsfull*nproc] + ) # [npar, nbinsfull*nproc] var_nobs = self.varnobs if self.varnobs is not None else self.nobs sumw2_flat = tf.reshape(self.indata.sumw2, [-1]) return (dxdnobs * var_nobs[None, :]) @ tf.transpose(dxdnobs) + ( @@ -2035,10 +2029,7 @@ def _compute_yields_noBBB( # systematic-template noise is de-biased. Curvature path only # (templates='fold'); halves use the shared logk. split_corr = None - if ( - self.logk_folds_delta is not None - and templates == "fold" - ): + if self.logk_folds_delta is not None and templates == "fold": theta_folded = tf.gather(theta, self.mcstat_folded_syst_idx) split_corr = tf.einsum( "bpj,j->bp", self.logk_folds_delta[fold_index], theta_folded @@ -2428,9 +2419,7 @@ def _compute_nll_components(self, profile=True, full_nll=False): nexp = nexpfullcentral[: self.indata.nbins] - debiased = ( - self.mcStatDebias in ("twoHalf", "kfold") and self.norm_A is not None - ) + debiased = self.mcStatDebias in ("twoHalf", "kfold") and self.norm_A is not None if debiased: # Cross-fit jackknife combination L_cf = 2 L_full - 1/2 L_A - 1/2 L_B. # Since mu_bar = 1/2(n_A + n_B) = n_full exactly, this de-biases the diff --git a/rabbit/inputdata.py b/rabbit/inputdata.py index d3722bf..e25a4de 100644 --- a/rabbit/inputdata.py +++ b/rabbit/inputdata.py @@ -177,8 +177,9 @@ def __init__(self, filename, pseudodata=None): # the two-half / k-fold de-biasing in the fitter. if "hnorm_folds" in f.keys(): self.norm_folds = maketensor(f["hnorm_folds"]) - self.mcstat_fold_k = int(f.attrs.get("mcstat_fold_k", - self.norm_folds.shape[0])) + self.mcstat_fold_k = int( + f.attrs.get("mcstat_fold_k", self.norm_folds.shape[0]) + ) else: self.norm_folds = None self.mcstat_fold_k = None @@ -195,9 +196,7 @@ def __init__(self, filename, pseudodata=None): # reduced syst dim) + the folded global syst indices. if "hlogk_folds_delta" in f.keys(): self.logk_folds_delta = maketensor(f["hlogk_folds_delta"]) - self.mcstat_folded_syst_idx = maketensor( - f["hmcstat_folded_syst_idx"] - ) + self.mcstat_folded_syst_idx = maketensor(f["hmcstat_folded_syst_idx"]) else: self.logk_folds_delta = None self.mcstat_folded_syst_idx = None diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 13f9047..f071b57 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -369,7 +369,9 @@ def add_mc_stat_moment(self, M, param_names, name="mcstat_M"): ax0 = _hist.axis.StrCategory(names, name="mcstat_params0") ax1 = _hist.axis.StrCategory(names, name="mcstat_params1") h = _hist.Hist(ax0, ax1, storage=_hist.storage.Double()) - h.view(flow=False)[...] = -M # hess = -M -> +1/2 x^T H x = -1/2 theta^T M theta + h.view(flow=False)[...] = ( + -M + ) # hess = -M -> +1/2 x^T H x = -1/2 theta^T M theta self.add_external_likelihood_term(hess=h, name=name) # tag for the fitter's sandwich covariance self.mcstat_moment_term = name @@ -1279,8 +1281,18 @@ def add_systematic( if fold_axis is not None: return self._add_systematic_folded( - h, name, process, channel, kfactor, mirror, symmetrize, - add_to_data_covariance, as_difference, syst_axes, fold_axis, **kargs + h, + name, + process, + channel, + kfactor, + mirror, + symmetrize, + add_to_data_covariance, + as_difference, + syst_axes, + fold_axis, + **kargs, ) # Fast batched path for SparseHist multi-systematic input. Conditions: @@ -1406,8 +1418,19 @@ def add_systematic( ) def _add_systematic_folded( - self, h, name, process, channel, kfactor, mirror, symmetrize, - add_to_data_covariance, as_difference, syst_axes, fold_axis, **kargs + self, + h, + name, + process, + channel, + kfactor, + mirror, + symmetrize, + add_to_data_covariance, + as_difference, + syst_axes, + fold_axis, + **kargs, ): """Add a systematic whose variation histogram(s) carry an MC-stat fold axis, SPLITTING logk per fold (RABBIT_MCSTAT_DESIGN.md §2a). The full @@ -1440,7 +1463,7 @@ def _add_systematic_folded( f"add_systematic(fold_axis=...): process '{process}' has no fold " f"templates; add it with the same fold_axis first." ) - norm_folds = self.dict_norm_folds[channel][process] # [k, nbinschan] + norm_folds = self.dict_norm_folds[channel][process] # [k, nbinschan] k = norm_folds.shape[0] ax_names = [a.name for a in h0.axes] if isinstance(fold_axis, str): @@ -1471,10 +1494,16 @@ def _add_systematic_folded( if as_difference: up_full = norm_full + up_full down_full = norm_full + down_full - lu_full = self.get_logk(up_full, norm_full, kfactor, systematic_type=systematic_type) - ld_full = -self.get_logk(down_full, norm_full, kfactor, systematic_type=systematic_type) + lu_full = self.get_logk( + up_full, norm_full, kfactor, systematic_type=systematic_type + ) + ld_full = -self.get_logk( + down_full, norm_full, kfactor, systematic_type=systematic_type + ) if symmetrize == "conservative": - logkavg_full = np.where(np.abs(lu_full) > np.abs(ld_full), lu_full, ld_full) + logkavg_full = np.where( + np.abs(lu_full) > np.abs(ld_full), lu_full, ld_full + ) else: logkavg_full = 0.5 * (lu_full + ld_full) else: @@ -1497,8 +1526,12 @@ def _add_systematic_folded( if as_difference: up_f = norm_f + up_f down_f = norm_f + down_f - lu = self.get_logk(up_f, norm_f, kfactor, systematic_type=systematic_type) - ld = -self.get_logk(down_f, norm_f, kfactor, systematic_type=systematic_type) + lu = self.get_logk( + up_f, norm_f, kfactor, systematic_type=systematic_type + ) + ld = -self.get_logk( + down_f, norm_f, kfactor, systematic_type=systematic_type + ) if symmetrize == "conservative": logkavg_folds[f] = np.where(np.abs(lu) > np.abs(ld), lu, ld) elif symmetrize == "average": @@ -1521,9 +1554,17 @@ def _add_systematic_folded( else: h_full = h[{fold_idx: sum}] self.add_systematic( - h_full, name, process, channel, kfactor=kfactor, mirror=mirror, - symmetrize=symmetrize, add_to_data_covariance=add_to_data_covariance, - as_difference=as_difference, syst_axes=syst_axes, **kargs + h_full, + name, + process, + channel, + kfactor=kfactor, + mirror=mirror, + symmetrize=symmetrize, + add_to_data_covariance=add_to_data_covariance, + as_difference=as_difference, + syst_axes=syst_axes, + **kargs, ) if is_asym and self.sparse: raise NotImplementedError( @@ -1533,12 +1574,14 @@ def _add_systematic_folded( # store the per-fold logk under the resolved (full) systematic name self.dict_logkavg_folds[channel].setdefault(process, {}) - self.dict_logkavg_folds[channel][process][name] = logkavg_folds.astype(self.dtype) + self.dict_logkavg_folds[channel][process][name] = logkavg_folds.astype( + self.dtype + ) # per-fold delta vs the full logk (for the sparse split-logk path) self.dict_logk_folds_delta[channel].setdefault(process, {}) self.dict_logk_folds_delta[channel][process][name] = ( - (logkavg_folds - logkavg_full[None, :]).astype(self.dtype) - ) + logkavg_folds - logkavg_full[None, :] + ).astype(self.dtype) if is_asym: self.dict_logkhalfdiff_folds[channel].setdefault(process, {}) self.dict_logkhalfdiff_folds[channel][process][name] = ( @@ -1951,9 +1994,7 @@ def write(self, outfolder="./", outfilename="rabbit_input.hdf5", meta_data_dict= any_folded = any(len(self.fold_k[c]) for c in self.channels) norm_folds = None if any_folded: - kset = { - k for c in self.channels for k in self.fold_k[c].values() - } + kset = {k for c in self.channels for k in self.fold_k[c].values()} if len(kset) != 1: raise NotImplementedError( f"All folded processes must share a common number of folds k; " @@ -2540,7 +2581,9 @@ def create_dataset( logk_folds_delta, f, "hlogk_folds_delta", maxChunkBytes=self.chunkSize ) nbytes += h5pyutils_write.writeFlatInChunks( - mcstat_folded_syst_idx, f, "hmcstat_folded_syst_idx", + mcstat_folded_syst_idx, + f, + "hmcstat_folded_syst_idx", maxChunkBytes=self.chunkSize, ) logk_folds_delta = None diff --git a/rabbit/workspace.py b/rabbit/workspace.py index 38f7566..e08132b 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -82,7 +82,10 @@ def __init__(self, outdir, outname, fitter, postfix=None): and getattr(fitter, "norm_A", None) is not None ) extra_traditional = list(param_impact_group_names) - if _is_debiased and getattr(fitter, "mcStatDebiasCov", "sandwich") == "sandwich": + if ( + _is_debiased + and getattr(fitter, "mcStatDebiasCov", "sandwich") == "sandwich" + ): extra_traditional.append("mcStatDebias") self.grouped_impact_axis = getGroupedImpactsAxes( fitter.indata, diff --git a/tests/mcstat_coverage.py b/tests/mcstat_coverage.py index 657b2ea..49480e8 100644 --- a/tests/mcstat_coverage.py +++ b/tests/mcstat_coverage.py @@ -13,20 +13,24 @@ OUT=/tmp/cov python3 tests/mcstat_coverage.py # prints a table from mcstat_coverage import run_coverage # programmatic """ -import os, sys -import numpy as np -import hist + import importlib.util as _ilu +import os +import sys + +import hist +import numpy as np RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) _spec = _ilu.spec_from_file_location( "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") ) -_rfm = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_rfm) -from rabbit import inputdata, fitter -from rabbit.tensorwriter import TensorWriter +_rfm = _ilu.module_from_spec(_spec) +_spec.loader.exec_module(_rfm) +from rabbit import fitter, inputdata from rabbit.param_models import helpers as ph +from rabbit.tensorwriter import TensorWriter OUT = os.environ.get("OUT", "/tmp/claude") # The near-degenerate 2-process toy from the slides (RESULTS §1): a flat process @@ -37,7 +41,7 @@ # the de-bias restores it. NB = int(os.environ.get("NB", "200")) FLAT = 4962.0 -HI = float(os.environ.get("HI", "5112")) # step height; larger = less degenerate +HI = float(os.environ.get("HI", "5112")) # step height; larger = less degenerate N0 = np.full(NB, FLAT) N1 = np.concatenate([np.full(NB // 2, HI), np.full(NB // 2, FLAT)]) # TRUE signal strengths. The degenerate (difference) direction must have a NONZERO @@ -47,7 +51,7 @@ R0 = float(os.environ.get("R0", "1.3")) R1 = float(os.environ.get("R1", "0.7")) RTRUE = np.array([R0, R1]) -MU = R0 * N0 + R1 * N1 # true total (data expectation) +MU = R0 * N0 + R1 * N1 # true total (data expectation) AX = hist.axis.Regular(NB, 0.0, 1.0, name="x") DDIF = np.array([1.0, -1.0]) / np.sqrt(2.0) DSUM = np.array([1.0, 1.0]) / np.sqrt(2.0) @@ -88,16 +92,31 @@ def _build(fn, T0, T1, folds0, folds1, data, debias, bbb, poisson): def _fit(fn, debias, covMode, debiasCov, bbb, poisson): - argv = [fn, "-o", f"{OUT}/vo", "-t", "0", "--allowNegativeParam", - "--mcStatDebias", debias, "--covMode", covMode, - "--mcStatDebiasCov", debiasCov] + argv = [ + fn, + "-o", + f"{OUT}/vo", + "-t", + "0", + "--allowNegativeParam", + "--mcStatDebias", + debias, + "--covMode", + covMode, + "--mcStatDebiasCov", + debiasCov, + ] argv += [] if poisson else ["--chisqFit"] if not bbb: argv.append("--noBinByBinStat") a = _rfm.make_parser().parse_args(argv) ind = inputdata.FitInputData(fn, None) - f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) - f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + f = fitter.Fitter( + ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False + ) + f.defaultassign() + f.set_nobs(ind.data_obs) + f.minimize() _, grad, hess = f.loss_val_grad_hess() bread = f.fisher_curvature(hess) if covMode == "fisher" else hess _, cov_curv = f.edmval_cov(grad, bread) @@ -112,12 +131,19 @@ def _fit(fn, debias, covMode, debiasCov, bbb, poisson): return f.x.numpy(), C, float(np.max(np.abs(grad.numpy()))) -def run_coverage(debias="none", covMode="observed", debiasCov="sandwich", - bbb=False, poisson=False, ntoy=200, k=2, seed=1): +def run_coverage( + debias="none", + covMode="observed", + debiasCov="sandwich", + bbb=False, + poisson=False, + ntoy=200, + k=2, + seed=1, +): rng = np.random.default_rng(seed) fn = f"{OUT}/cov_{debias}_{covMode}_{debiasCov}_{int(bbb)}_{int(poisson)}.hdf5" cov_dif = cov_sum = 0 - bias_dif = [] res_dif = [] sig_dif = [] nbad = 0 @@ -135,28 +161,43 @@ def run_coverage(debias="none", covMode="observed", debiasCov="sandwich", # (a rank-deficient full matrix is fine as long as the projection is ok; # BB-lite on a near-degenerate toy can null the other direction). if gmax > 1e-2 or v_dif <= 0 or v_sum <= 0 or not np.isfinite(v_dif): - nbad += 1; continue + nbad += 1 + continue except Exception: - nbad += 1; continue - r = x - RTRUE # residual vs true signal strengths + nbad += 1 + continue + r = x - RTRUE # residual vs true signal strengths s_dif = np.sqrt(v_dif) s_sum = np.sqrt(v_sum) cov_dif += abs(DDIF @ r) < s_dif cov_sum += abs(DSUM @ r) < s_sum - res_dif.append(DDIF @ r); sig_dif.append(s_dif) + res_dif.append(DDIF @ r) + sig_dif.append(s_dif) n = ntoy - nbad if n == 0: - return dict(cov_dif=float("nan"), cov_sum=float("nan"), n=0, nbad=nbad, - bias_dif=float("nan"), mean_res=float("nan"), - rms_res=float("nan"), med_sig=float("nan")) + return dict( + cov_dif=float("nan"), + cov_sum=float("nan"), + n=0, + nbad=nbad, + bias_dif=float("nan"), + mean_res=float("nan"), + rms_res=float("nan"), + med_sig=float("nan"), + ) res = np.array(res_dif) return dict( - cov_dif=cov_dif / n, cov_sum=cov_sum / n, n=n, nbad=nbad, + cov_dif=cov_dif / n, + cov_sum=cov_sum / n, + n=n, + nbad=nbad, # median residual = robust bias (mean is corrupted by the heavy tails the # de-bias develops near the singular limit); rms_res = the actual point # scatter (vs med_sig = the reported uncertainty -> coverage = rms/med). - bias_dif=float(np.median(res)), mean_res=float(np.mean(res)), - rms_res=float(np.std(res)), med_sig=float(np.median(sig_dif)), + bias_dif=float(np.median(res)), + mean_res=float(np.mean(res)), + rms_res=float(np.std(res)), + med_sig=float(np.median(sig_dif)), ) @@ -165,20 +206,39 @@ def run_coverage(debias="none", covMode="observed", debiasCov="sandwich", configs = [ ("none observed sandwich chisq", dict(debias="none")), ("continuousM observed sandwich chisq", dict(debias="continuousM")), - ("continuousM observed curvature chisq", dict(debias="continuousM", debiasCov="curvature")), - ("continuousM fisher sandwich chisq", dict(debias="continuousM", covMode="fisher")), + ( + "continuousM observed curvature chisq", + dict(debias="continuousM", debiasCov="curvature"), + ), + ( + "continuousM fisher sandwich chisq", + dict(debias="continuousM", covMode="fisher"), + ), ("twoHalf observed sandwich chisq", dict(debias="twoHalf")), - ("twoHalf fisher sandwich chisq", dict(debias="twoHalf", covMode="fisher")), - ("kfold(8) fisher sandwich chisq", dict(debias="kfold", covMode="fisher", k=8)), + ( + "twoHalf fisher sandwich chisq", + dict(debias="twoHalf", covMode="fisher"), + ), + ( + "kfold(8) fisher sandwich chisq", + dict(debias="kfold", covMode="fisher", k=8), + ), ("none observed sandwich chisq BBB", dict(debias="none", bbb=True)), - ("continuousM observed sandwich chisq BBB", dict(debias="continuousM", bbb=True)), + ( + "continuousM observed sandwich chisq BBB", + dict(debias="continuousM", bbb=True), + ), ("twoHalf observed sandwich chisq BBB", dict(debias="twoHalf", bbb=True)), ("none observed sandwich POISSON", dict(debias="none", poisson=True)), ("twoHalf observed sandwich POISSON", dict(debias="twoHalf", poisson=True)), ] print(f"MC-stat coverage ensemble (ntoy={ntoy}, target cov(dif)=0.683)\n") - print(f"{'config':40s}{'cov(dif)':>9s}{'cov(sum)':>9s}{'bias(dif)':>11s}{'med_sig':>9s}{'nbad':>6s}") + print( + f"{'config':40s}{'cov(dif)':>9s}{'cov(sum)':>9s}{'bias(dif)':>11s}{'med_sig':>9s}{'nbad':>6s}" + ) for label, kw in configs: r = run_coverage(ntoy=ntoy, **kw) - print(f"{label:40s}{r['cov_dif']:9.3f}{r['cov_sum']:9.3f}" - f"{r['bias_dif']:+11.4f}{r['med_sig']:9.4f}{r['nbad']:6d}") + print( + f"{label:40s}{r['cov_dif']:9.3f}{r['cov_sum']:9.3f}" + f"{r['bias_dif']:+11.4f}{r['med_sig']:9.4f}{r['nbad']:6d}" + ) diff --git a/tests/toy_kfold.py b/tests/toy_kfold.py index f9b449f..0f9b741 100644 --- a/tests/toy_kfold.py +++ b/tests/toy_kfold.py @@ -2,9 +2,12 @@ complete k-fold U-statistic curvature (k-fold averaging). Each process's finite-MC template is split into 4 independent folds (multinomial 1/4 per bin); the full = sum of folds is derived by the writer. Writes toy_kfold.hdf5.""" + import os -import numpy as np + import hist +import numpy as np + from rabbit.tensorwriter import TensorWriter NB = 200 @@ -19,7 +22,9 @@ sw_flat = rng.poisson(n_flat) sw_step = rng.poisson(n_step) # split each bin's counts into K folds via a multinomial(count, [1/K]*K) -flat_folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in sw_flat], axis=0).T # [K, NB] +flat_folds = np.stack( + [rng.multinomial(c, [1.0 / K] * K) for c in sw_flat], axis=0 +).T # [K, NB] step_folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in sw_step], axis=0).T ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") @@ -28,7 +33,8 @@ def dhist(v): h = hist.Hist(ax, storage=hist.storage.Weight()) - h.view()["value"] = v; h.view()["variance"] = v + h.view()["value"] = v + h.view()["variance"] = v return h diff --git a/tests/toy_mcstat.py b/tests/toy_mcstat.py index 11b3e16..bdb42d2 100644 --- a/tests/toy_mcstat.py +++ b/tests/toy_mcstat.py @@ -4,9 +4,12 @@ toy_M.hdf5 : same + frozen noise-floor matrix M over the 2 POIs A chisqFit on each should give a *larger*, de-biased POI uncertainty with M (curvature H - M; RABBIT_MCSTAT_DESIGN §1).""" + import os -import numpy as np + import hist +import numpy as np + from rabbit.tensorwriter import TensorWriter NB = 200 @@ -15,21 +18,23 @@ n_flat = np.full(NB, A) n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) -mu_true = n_flat + n_step # true total per bin +mu_true = n_flat + n_step # true total per bin # finite-MC (1:1) templates: one Poisson realization, sumw2 = counts (unit weights) sw_flat = rng.poisson(n_flat).astype(float) sw_step = rng.poisson(n_step).astype(float) -data = mu_true.copy() # Asimov data (clean curvature comparison) +data = mu_true.copy() # Asimov data (clean curvature comparison) ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") + def whist(values, variances): h = hist.Hist(ax, storage=hist.storage.Weight()) h.view()["value"] = values h.view()["variance"] = variances return h + def build(outname, with_M): tw = TensorWriter(sparse=False) tw.add_channel([ax], "ch0") @@ -44,10 +49,14 @@ def build(outname, with_M): M = np.diag([Mflat, Mstep]) tw.add_mc_stat_moment(M, ["flat", "step"]) print(f" M = diag({Mflat:.2f}, {Mstep:.2f})") - tw.write(outfolder=os.path.dirname(outname) or ".", - outfilename=os.path.basename(outname)) + tw.write( + outfolder=os.path.dirname(outname) or ".", outfilename=os.path.basename(outname) + ) + if __name__ == "__main__": out = os.environ.get("OUT", "/tmp/claude") - print("building toy_noM.hdf5"); build(f"{out}/toy_noM.hdf5", with_M=False) - print("building toy_M.hdf5"); build(f"{out}/toy_M.hdf5", with_M=True) + print("building toy_noM.hdf5") + build(f"{out}/toy_noM.hdf5", with_M=False) + print("building toy_M.hdf5") + build(f"{out}/toy_M.hdf5", with_M=True) diff --git a/tests/toy_splitlogk.py b/tests/toy_splitlogk.py index f24ede7..f73b5c0 100644 --- a/tests/toy_splitlogk.py +++ b/tests/toy_splitlogk.py @@ -7,16 +7,19 @@ folds, systematic added WITHOUT fold_axis = shared logk) for comparison. Set the env var K to the number of folds (default 2; K>=3 exercises k>2 split-logk, whose per-fold logk feeds the k-fold U-statistic curvature in --covMode fisher).""" + import os -import numpy as np + import hist +import numpy as np + from rabbit.tensorwriter import TensorWriter NB = 50 K = int(os.environ.get("K", "2")) rng = np.random.default_rng(7) -nom = np.linspace(800.0, 1200.0, NB) # smooth nominal shape -data = nom.copy() # Asimov +nom = np.linspace(800.0, 1200.0, NB) # smooth nominal shape +data = nom.copy() # Asimov # K independent folds of the nominal (per-bin multinomial split) full = rng.poisson(nom) @@ -26,7 +29,9 @@ # template with fold-specific Poisson noise x = np.linspace(0, 1, NB) w = 1.0 + 0.15 * (x - 0.5) -upfolds = np.stack([rng.poisson(np.maximum(folds[f] * w, 0.0)) for f in range(K)], axis=0) +upfolds = np.stack( + [rng.poisson(np.maximum(folds[f] * w, 0.0)) for f in range(K)], axis=0 +) ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") axf = hist.axis.IntCategory(list(range(K)), name="mcfold") @@ -34,14 +39,16 @@ def dhist(v): h = hist.Hist(ax, storage=hist.storage.Weight()) - h.view()["value"] = v; h.view()["variance"] = v + h.view()["value"] = v + h.view()["variance"] = v return h def fhist(fl): h = hist.Hist(ax, axf, storage=hist.storage.Weight()) for f in range(K): - h.view()["value"][:, f] = fl[f]; h.view()["variance"][:, f] = fl[f] + h.view()["value"][:, f] = fl[f] + h.view()["variance"][:, f] = fl[f] return h @@ -55,12 +62,15 @@ def build(outname, split): else: # shared logk: full up template (folds summed), no fold_axis tw.add_systematic(dhist(upfolds.sum(0)), "tilt", "sig", "ch0") - tw.write(outfolder=os.path.dirname(outname) or ".", - outfilename=os.path.basename(outname)) + tw.write( + outfolder=os.path.dirname(outname) or ".", outfilename=os.path.basename(outname) + ) if __name__ == "__main__": out = os.environ.get("OUT", "/tmp/claude") build(f"{out}/toy_splitlogk.hdf5", split=True) build(f"{out}/toy_sharedlogk.hdf5", split=False) - print(f"wrote {out}/toy_splitlogk.hdf5 (split) and toy_sharedlogk.hdf5 (shared), K={K}") + print( + f"wrote {out}/toy_splitlogk.hdf5 (split) and toy_sharedlogk.hdf5 (shared), K={K}" + ) diff --git a/tests/toy_twohalf.py b/tests/toy_twohalf.py index e53edba..b381b78 100644 --- a/tests/toy_twohalf.py +++ b/tests/toy_twohalf.py @@ -4,9 +4,12 @@ Each process's finite-MC template is split into two independent halves (fold 0/1) via a per-bin binomial(count, 0.5); the full template = fold0 + fold1 is derived by the writer. Writes toy_folds.hdf5.""" + import os -import numpy as np + import hist +import numpy as np + from rabbit.tensorwriter import TensorWriter NB = 200 @@ -16,13 +19,15 @@ n_flat = np.full(NB, A) n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) mu_true = n_flat + n_step -data = mu_true.copy() # Asimov data +data = mu_true.copy() # Asimov data # finite-MC full templates, then split into two folds per bin sw_flat = rng.poisson(n_flat) sw_step = rng.poisson(n_step) -flat_A = rng.binomial(sw_flat, 0.5); flat_B = sw_flat - flat_A -step_A = rng.binomial(sw_step, 0.5); step_B = sw_step - step_A +flat_A = rng.binomial(sw_flat, 0.5) +flat_B = sw_flat - flat_A +step_A = rng.binomial(sw_step, 0.5) +step_B = sw_step - step_A ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") axf = hist.axis.IntCategory([0, 1], name="mcfold") @@ -49,9 +54,11 @@ def folded_hist(fold0, fold1): tw = TensorWriter(sparse=False) tw.add_channel([ax], "ch0") tw.add_data(data_hist(data), "ch0") - tw.add_process(folded_hist(flat_A, flat_B), "flat", "ch0", - signal=True, fold_axis="mcfold") - tw.add_process(folded_hist(step_A, step_B), "step", "ch0", - signal=True, fold_axis="mcfold") + tw.add_process( + folded_hist(flat_A, flat_B), "flat", "ch0", signal=True, fold_axis="mcfold" + ) + tw.add_process( + folded_hist(step_A, step_B), "step", "ch0", signal=True, fold_axis="mcfold" + ) tw.write(outfolder=out, outfilename="toy_folds.hdf5") print(f"wrote {out}/toy_folds.hdf5 (k=2 folds for flat, step)") diff --git a/tests/verify_bblite.py b/tests/verify_bblite.py index c52ce68..f416af0 100644 --- a/tests/verify_bblite.py +++ b/tests/verify_bblite.py @@ -10,10 +10,13 @@ Run with OUT=. Needs toy_splitlogk.hdf5 (tests/toy_splitlogk.py) present. """ -import os, sys -import numpy as np -import hist + import importlib.util as _ilu +import os +import sys + +import hist +import numpy as np RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) @@ -22,7 +25,7 @@ ) _rfm = _ilu.module_from_spec(_spec) _spec.loader.exec_module(_rfm) -from rabbit import inputdata, fitter +from rabbit import fitter, inputdata from rabbit.param_models import helpers as ph from rabbit.tensorwriter import TensorWriter @@ -30,17 +33,20 @@ def build_toy_Mbb(): - NB = 200; A, H1, H2 = 4962.0, 5112.0, 4962.0 + NB = 200 + A, H1, H2 = 4962.0, 5112.0, 4962.0 rng = np.random.default_rng(20240614) n_flat = np.full(NB, A) n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) data = (n_flat + n_step).astype(float) - sw_flat = rng.poisson(n_flat).astype(float); sw_step = rng.poisson(n_step).astype(float) + sw_flat = rng.poisson(n_flat).astype(float) + sw_step = rng.poisson(n_step).astype(float) ax = hist.axis.Regular(NB, 0.0, 1.0, name="x") def wh(v, var): h = hist.Hist(ax, storage=hist.storage.Weight()) - h.view()["value"] = v; h.view()["variance"] = var + h.view()["value"] = v + h.view()["variance"] = var return h tw = TensorWriter(sparse=False) @@ -48,21 +54,34 @@ def wh(v, var): tw.add_data(wh(data, data), "ch0") tw.add_process(wh(sw_flat, sw_flat), "flat", "ch0", signal=True) tw.add_process(wh(sw_step, sw_step), "step", "ch0", signal=True) - Vbb = data + (sw_flat + sw_step) # BB-lite inflated variance + Vbb = data + (sw_flat + sw_step) # BB-lite inflated variance M = np.diag([np.sum(sw_flat / Vbb), np.sum(sw_step / Vbb)]) tw.add_mc_stat_moment(M, ["flat", "step"]) tw.write(outfolder=OUT, outfilename="toy_Mbb.hdf5") def fit(fn, debias, bbb, sandwich=None): - argv = [fn, "-o", f"{OUT}/vo", "-t", "0", "--chisqFit", - "--allowNegativeParam", "--mcStatDebias", debias] + argv = [ + fn, + "-o", + f"{OUT}/vo", + "-t", + "0", + "--chisqFit", + "--allowNegativeParam", + "--mcStatDebias", + debias, + ] if not bbb: argv.append("--noBinByBinStat") a = _rfm.make_parser().parse_args(argv) ind = inputdata.FitInputData(fn, None) - f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) - f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + f = fitter.Fitter( + ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False + ) + f.defaultassign() + f.set_nobs(ind.data_obs) + f.minimize() _, g, h = f.loss_val_grad_hess() _, cov = f.edmval_cov(g, h) C = np.asarray(cov) @@ -78,7 +97,9 @@ def fit(fn, debias, bbb, sandwich=None): if __name__ == "__main__": print("(A) continuous-M + BB-lite (toy_Mbb, BB-variance M):") build_toy_Mbb() - x, sig, S, gmax, eig = fit(f"{OUT}/toy_Mbb.hdf5", "continuousM", True, "continuousM") + x, sig, S, gmax, eig = fit( + f"{OUT}/toy_Mbb.hdf5", "continuousM", True, "continuousM" + ) pd = bool(np.all(eig > 0)) print(f" x={np.round(x,4)} |grad|={gmax:.1e} PD={pd} hess_eig={np.round(eig,2)}") assert pd and gmax < 1e-3, "continuous-M + BB-lite did not converge PD" @@ -88,7 +109,9 @@ def fit(fn, debias, bbb, sandwich=None): x0, sig0, _, g0, e0 = fit(fn, "none", True) x1, sig1, S1, g1, e1 = fit(fn, "twoHalf", True, "twoHalf") print(f" none x={np.round(x0,4)} tilt sigma={sig0[1]:.4f}") - print(f" twoHalf x={np.round(x1,4)} tilt curv={sig1[1]:.4f} sandwich={np.sqrt(S1[1,1]):.4f}") + print( + f" twoHalf x={np.round(x1,4)} tilt curv={sig1[1]:.4f} sandwich={np.sqrt(S1[1,1]):.4f}" + ) assert bool(np.all(e1 > 0)) and g1 < 1e-3, "two-half + BB-lite not converged PD" assert np.sqrt(S1[1, 1]) > sig0[1], "two-half should inflate the tilt uncertainty" print("\n ALL BB-lite composition checks passed.") diff --git a/tests/verify_coverage.py b/tests/verify_coverage.py index 4650260..8de8822 100644 --- a/tests/verify_coverage.py +++ b/tests/verify_coverage.py @@ -14,9 +14,11 @@ calibration factor -- this test asserts the de-bias DIRECTION (bias removed, sigma inflated, coverage improved), not exact 0.683. Run with OUT=. """ + import os + os.environ.setdefault("NB", "200") -os.environ.setdefault("HI", "5112") # near-degenerate (large attenuation) +os.environ.setdefault("HI", "5112") # near-degenerate (large attenuation) os.environ.setdefault("R0", "1.3") os.environ.setdefault("R1", "0.7") @@ -25,23 +27,36 @@ NTOY = int(os.environ.get("NTOY", "50")) if __name__ == "__main__": - print(f"coverage harness, slides toy, diff_true={mc.DDIF @ mc.RTRUE:.3f}, ntoy={NTOY}") + print( + f"coverage harness, slides toy, diff_true={mc.DDIF @ mc.RTRUE:.3f}, ntoy={NTOY}" + ) none = mc.run_coverage(ntoy=NTOY, seed=11, debias="none") - cmS = mc.run_coverage(ntoy=NTOY, seed=11, debias="continuousM", debiasCov="sandwich") + cmS = mc.run_coverage( + ntoy=NTOY, seed=11, debias="continuousM", debiasCov="sandwich" + ) for lab, r in [("none", none), ("continuousM sand", cmS)]: - print(f" {lab:18s} cov(dif)={r['cov_dif']:.3f} med_bias={r['bias_dif']:+.4f} " - f"med_sig={r['med_sig']:.4f} rms_res={r['rms_res']:.4f}") + print( + f" {lab:18s} cov(dif)={r['cov_dif']:.3f} med_bias={r['bias_dif']:+.4f} " + f"med_sig={r['med_sig']:.4f} rms_res={r['rms_res']:.4f}" + ) # 1. naive is badly biased by attenuation (degenerate direction pulled to 0) - assert none["bias_dif"] < -0.10, f"naive should be attenuated, got {none['bias_dif']}" + assert ( + none["bias_dif"] < -0.10 + ), f"naive should be attenuated, got {none['bias_dif']}" # 2. de-bias removes the central-value bias - assert abs(cmS["bias_dif"]) < 0.4 * abs(none["bias_dif"]), \ - f"de-bias should remove the bias: {none['bias_dif']} -> {cmS['bias_dif']}" + assert abs(cmS["bias_dif"]) < 0.4 * abs( + none["bias_dif"] + ), f"de-bias should remove the bias: {none['bias_dif']} -> {cmS['bias_dif']}" # 3. de-bias inflates the uncertainty estimate toward sigma_inf - assert cmS["med_sig"] > 1.3 * none["med_sig"], \ - f"de-bias should inflate sigma: {none['med_sig']} -> {cmS['med_sig']}" + assert ( + cmS["med_sig"] > 1.3 * none["med_sig"] + ), f"de-bias should inflate sigma: {none['med_sig']} -> {cmS['med_sig']}" # 4. de-bias improves coverage of the degenerate POI - assert cmS["cov_dif"] > none["cov_dif"] + 0.10, \ - f"de-bias should improve coverage: {none['cov_dif']} -> {cmS['cov_dif']}" - print("\n coverage de-bias checks passed (bias removed, sigma inflated, " - "coverage improved).") + assert ( + cmS["cov_dif"] > none["cov_dif"] + 0.10 + ), f"de-bias should improve coverage: {none['cov_dif']} -> {cmS['cov_dif']}" + print( + "\n coverage de-bias checks passed (bias removed, sigma inflated, " + "coverage improved)." + ) diff --git a/tests/verify_kfold.py b/tests/verify_kfold.py index 9eb1146..e8c165c 100644 --- a/tests/verify_kfold.py +++ b/tests/verify_kfold.py @@ -7,17 +7,21 @@ Needs toy_kfold.hdf5 (tests/toy_kfold.py). Run with OUT=. """ -import os, sys -import numpy as np + import importlib.util as _ilu +import os +import sys + +import numpy as np RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) _spec = _ilu.spec_from_file_location( "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") ) -_rfm = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_rfm) -from rabbit import inputdata, fitter +_rfm = _ilu.module_from_spec(_spec) +_spec.loader.exec_module(_rfm) +from rabbit import fitter, inputdata from rabbit.param_models import helpers as ph NB, A, H1, H2 = 200, 4962.0, 5112.0, 4962.0 @@ -37,7 +41,8 @@ def ustat(Ti, V): def numpy_ref_k4(): rng = np.random.default_rng(20240614) # same seed as toy_kfold.py - swf = rng.poisson(n_flat); sws = rng.poisson(n_step) + swf = rng.poisson(n_flat) + sws = rng.poisson(n_step) ff = np.stack([rng.multinomial(c, [0.25] * 4) for c in swf], 0).T sf = np.stack([rng.multinomial(c, [0.25] * 4) for c in sws], 0).T Ti = [np.stack([ff[i], sf[i]], 1).astype(float) for i in range(4)] @@ -45,13 +50,29 @@ def numpy_ref_k4(): def rabbit_fisher(fn): - import tensorflow as tf a = _rfm.make_parser().parse_args( - [fn, "-o", "/tmp/claude/vo", "-t", "0", "--chisqFit", "--noBinByBinStat", - "--allowNegativeParam", "--mcStatDebias", "kfold", "--covMode", "fisher"]) + [ + fn, + "-o", + "/tmp/claude/vo", + "-t", + "0", + "--chisqFit", + "--noBinByBinStat", + "--allowNegativeParam", + "--mcStatDebias", + "kfold", + "--covMode", + "fisher", + ] + ) ind = inputdata.FitInputData(fn, None) - f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) - f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + f = fitter.Fitter( + ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False + ) + f.defaultassign() + f.set_nobs(ind.data_obs) + f.minimize() _, _, h = f.loss_val_grad_hess() return f.n_folds, np.asarray(f.fisher_curvature(h)) @@ -60,7 +81,8 @@ def ensemble_rms(K, ntoy=300): rng = np.random.default_rng(100 + K) s = [] for _ in range(ntoy): - swf = rng.poisson(n_flat); sws = rng.poisson(n_step) + swf = rng.poisson(n_flat) + sws = rng.poisson(n_step) ff = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in swf], 0).T sf = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in sws], 0).T Ti = [np.stack([ff[i], sf[i]], 1).astype(float) for i in range(K)] @@ -75,14 +97,18 @@ def ensemble_rms(K, ntoy=300): out = os.environ.get("OUT", "/tmp/claude") A4 = numpy_ref_k4() k, Fb = rabbit_fisher(f"{out}/toy_kfold.hdf5") - print(f"(A) rabbit k={k} fisher curvature == numpy U-statistic A_k: " - f"{np.allclose(Fb, A4, rtol=1e-4)} " - f"(sigma(dif) rabbit={np.sqrt(ddif@np.linalg.inv(Fb)@ddif):.4f}, " - f"numpy={np.sqrt(ddif@np.linalg.inv(A4)@ddif):.4f})") + print( + f"(A) rabbit k={k} fisher curvature == numpy U-statistic A_k: " + f"{np.allclose(Fb, A4, rtol=1e-4)} " + f"(sigma(dif) rabbit={np.sqrt(ddif@np.linalg.inv(Fb)@ddif):.4f}, " + f"numpy={np.sqrt(ddif@np.linalg.inv(A4)@ddif):.4f})" + ) assert np.allclose(Fb, A4, rtol=1e-4) Tt = np.stack([n_flat, n_step], 1) - sinf = np.sqrt(ddif @ np.linalg.inv(np.einsum("bi,b,bj->ij", Tt, 1 / data, Tt)) @ ddif) + sinf = np.sqrt( + ddif @ np.linalg.inv(np.einsum("bi,b,bj->ij", Tt, 1 / data, Tt)) @ ddif + ) print(f"(B) ensemble: sigma_inf={sinf:.4f}") prev = None for K in (4, 8, 16): diff --git a/tests/verify_mcstat.py b/tests/verify_mcstat.py index 8809867..851bd7d 100644 --- a/tests/verify_mcstat.py +++ b/tests/verify_mcstat.py @@ -8,9 +8,12 @@ cancels the correction; pass --allowNegativeParam (LINEAR rnorm=x) to see the de-bias. See RESULTS.md S9a. Run with OUT= pointing at toy_{noM,M}.hdf5. """ -import os, sys -import numpy as np + import importlib.util as _ilu +import os +import sys + +import numpy as np RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) @@ -20,13 +23,20 @@ _rfm = _ilu.module_from_spec(_spec) _spec.loader.exec_module(_rfm) -from rabbit import inputdata, fitter +from rabbit import fitter, inputdata from rabbit.param_models import helpers as ph def run(filename, linear): - argv = [filename, "-o", "/tmp/claude/verify_out", "-t", "0", - "--noBinByBinStat", "--chisqFit"] + argv = [ + filename, + "-o", + "/tmp/claude/verify_out", + "-t", + "0", + "--noBinByBinStat", + "--chisqFit", + ] if linear: argv.append("--allowNegativeParam") args = _rfm.make_parser().parse_args(argv) @@ -38,14 +48,14 @@ def run(filename, linear): f.minimize() _, grad, hess = f.loss_val_grad_hess() _, cov = f.edmval_cov(grad, hess) - C = np.asarray(cov) # curvature cov A^-1 + C = np.asarray(cov) # curvature cov A^-1 Csand = None if f.mcstat_M is not None: - Csand = np.asarray(f.cov_mcstat_sandwich(hess)) # A^-1 + A^-1 M A^-1 + Csand = np.asarray(f.cov_mcstat_sandwich(hess)) # A^-1 + A^-1 M A^-1 x = f.x.numpy() - rnorm = x if linear else x ** 2 # physical signal strength - J = np.diag(np.ones_like(x) if linear else 2 * x) # d rnorm / d x - Cr = J @ C @ J # curvature cov in rnorm space + rnorm = x if linear else x**2 # physical signal strength + J = np.diag(np.ones_like(x) if linear else 2 * x) # d rnorm / d x + Cr = J @ C @ J # curvature cov in rnorm space Csr = None if Csand is None else J @ Csand @ J return f.parms.astype(str), rnorm, Cr, Csr @@ -54,23 +64,35 @@ def run(filename, linear): out = os.environ.get("OUT", "/tmp/claude") ddif = np.array([1.0, -1.0]) / np.sqrt(2) for linear in (False, True): - print(f"\n=== POI {'LINEAR (--allowNegativeParam)' if linear else 'SQUARED (rabbit default)'} ===") + print( + f"\n=== POI {'LINEAR (--allowNegativeParam)' if linear else 'SQUARED (rabbit default)'} ===" + ) res = {} for tag in ("noM", "M"): parms, rnorm, Cr, Csr = run(f"{out}/toy_{tag}.hdf5", linear) res[tag] = (rnorm, Cr, Csr) extra = "" if Csr is not None: - extra = (f" sandwich(flat)={np.sqrt(Csr[0,0]):.4f} " - f"sandwich(dif)={np.sqrt(ddif@Csr@ddif):.4f}") - print(f" {tag:3s} rnorm={np.round(rnorm, 4)} " - f"curv_rnorm(flat)={np.sqrt(Cr[0,0]):.4f} " - f"curv_rnorm(dif)={np.sqrt(ddif@Cr@ddif):.4f}{extra}") + extra = ( + f" sandwich(flat)={np.sqrt(Csr[0,0]):.4f} " + f"sandwich(dif)={np.sqrt(ddif@Csr@ddif):.4f}" + ) + print( + f" {tag:3s} rnorm={np.round(rnorm, 4)} " + f"curv_rnorm(flat)={np.sqrt(Cr[0,0]):.4f} " + f"curv_rnorm(dif)={np.sqrt(ddif@Cr@ddif):.4f}{extra}" + ) s_no = np.sqrt(ddif @ res["noM"][1] @ ddif) s_M = np.sqrt(ddif @ res["M"][1] @ ddif) - verdict = "DE-BIASED (sigma inflated)" if s_M > 1.2 * s_no else "no de-bias (cancelled)" + verdict = ( + "DE-BIASED (sigma inflated)" + if s_M > 1.2 * s_no + else "no de-bias (cancelled)" + ) print(f" -> curvature sigma(dif): {s_no:.4f} -> {s_M:.4f} {verdict}") if res["M"][2] is not None: s_sand = np.sqrt(ddif @ res["M"][2] @ ddif) - print(f" -> sandwich sigma(dif): {s_sand:.4f} " - f"(>= curvature {s_M:.4f}: robust over-coverage margin)") + print( + f" -> sandwich sigma(dif): {s_sand:.4f} " + f"(>= curvature {s_M:.4f}: robust over-coverage margin)" + ) diff --git a/tests/verify_sparse_splitlogk.py b/tests/verify_sparse_splitlogk.py index a0a29f4..779302b 100644 --- a/tests/verify_sparse_splitlogk.py +++ b/tests/verify_sparse_splitlogk.py @@ -3,20 +3,24 @@ equivalent DENSE tensor. The sparse path multiplies the shared systematic factor by exp(delta . theta) over the folded systs (delta = logk_fold - logk_full), curvature path only. log_normal symmetric folded systematics. Run with OUT=.""" -import os, sys -import numpy as np -import hist + import importlib.util as _ilu +import os +import sys + +import hist +import numpy as np RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) _spec = _ilu.spec_from_file_location( "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") ) -_rfm = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_rfm) -from rabbit import inputdata, fitter -from rabbit.tensorwriter import TensorWriter +_rfm = _ilu.module_from_spec(_spec) +_spec.loader.exec_module(_rfm) +from rabbit import fitter, inputdata from rabbit.param_models import helpers as ph +from rabbit.tensorwriter import TensorWriter NB, K = 40, 4 OUT = os.environ.get("OUT", "/tmp/claude") @@ -24,27 +28,32 @@ def build(sparse): rng = np.random.default_rng(21) - nom = np.linspace(900, 1100, NB); data = nom.copy() + nom = np.linspace(900, 1100, NB) + data = nom.copy() full = rng.poisson(nom) folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in full], 0).T - x = np.linspace(0, 1, NB); w = 1 + 0.15 * (x - 0.5) + x = np.linspace(0, 1, NB) + w = 1 + 0.15 * (x - 0.5) upf = np.stack([rng.poisson(np.maximum(folds[f] * w, 0)) for f in range(K)], 0) ax = hist.axis.Regular(NB, 0, 1, name="x") axf = hist.axis.IntCategory(list(range(K)), name="mcfold") def dh(v): h = hist.Hist(ax, storage=hist.storage.Weight()) - h.view()["value"] = v; h.view()["variance"] = v + h.view()["value"] = v + h.view()["variance"] = v return h def fh(fl): h = hist.Hist(ax, axf, storage=hist.storage.Weight()) for f in range(K): - h.view()["value"][:, f] = fl[f]; h.view()["variance"][:, f] = fl[f] + h.view()["value"][:, f] = fl[f] + h.view()["variance"][:, f] = fl[f] return h tw = TensorWriter(sparse=sparse) - tw.add_channel([ax], "ch0"); tw.add_data(dh(data), "ch0") + tw.add_channel([ax], "ch0") + tw.add_data(dh(data), "ch0") tw.add_process(fh(folds), "sig", "ch0", signal=True, fold_axis="mcfold") tw.add_systematic(fh(upf), "tilt", "sig", "ch0", fold_axis="mcfold") # split-logk fn = f"{OUT}/toy_spl_{'sp' if sparse else 'de'}.hdf5" @@ -54,11 +63,28 @@ def fh(fl): def run(fn): a = _rfm.make_parser().parse_args( - [fn, "-o", f"{OUT}/vo", "-t", "0", "--chisqFit", "--noBinByBinStat", - "--allowNegativeParam", "--mcStatDebias", "kfold", "--covMode", "fisher"]) + [ + fn, + "-o", + f"{OUT}/vo", + "-t", + "0", + "--chisqFit", + "--noBinByBinStat", + "--allowNegativeParam", + "--mcStatDebias", + "kfold", + "--covMode", + "fisher", + ] + ) ind = inputdata.FitInputData(fn, None) - f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) - f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + f = fitter.Fitter( + ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False + ) + f.defaultassign() + f.set_nobs(ind.data_obs) + f.minimize() _, _, h = f.loss_val_grad_hess() bread = f.fisher_curvature(h) S = np.asarray(f.cov_twohalf_sandwich(bread, "fisher")) @@ -68,10 +94,15 @@ def run(fn): if __name__ == "__main__": _, xs, bs, Ss = run(build(True)) _, xd, bd, Sd = run(build(False)) - ok = (np.allclose(xs, xd, rtol=1e-6) and np.allclose(bs, bd, rtol=1e-6) - and np.allclose(Ss, Sd, rtol=1e-6)) - print(f"point/fisher/sandwich match sparse==dense: " - f"{np.allclose(xs,xd,1e-6)}/{np.allclose(bs,bd,1e-6)}/{np.allclose(Ss,Sd,1e-6)}") + ok = ( + np.allclose(xs, xd, rtol=1e-6) + and np.allclose(bs, bd, rtol=1e-6) + and np.allclose(Ss, Sd, rtol=1e-6) + ) + print( + f"point/fisher/sandwich match sparse==dense: " + f"{np.allclose(xs,xd,1e-6)}/{np.allclose(bs,bd,1e-6)}/{np.allclose(Ss,Sd,1e-6)}" + ) print(f" tilt sandwich sparse={np.sqrt(Ss[1,1]):.5f} dense={np.sqrt(Sd[1,1]):.5f}") assert ok, "sparse split-logk must equal dense" print("\n sparse split-logk == dense (point, curvature, sandwich). passed.") diff --git a/tests/verify_sparsefold.py b/tests/verify_sparsefold.py index 24d15b5..f9abdfb 100644 --- a/tests/verify_sparsefold.py +++ b/tests/verify_sparsefold.py @@ -2,20 +2,24 @@ (scatter the shared systematic factor into a dense [nbinsfull,nproc] grid, contract with the dense fold norm) must give EXACTLY the same de-biased point, fisher curvature, and sandwich as the equivalent DENSE tensor. Run with OUT=.""" -import os, sys -import numpy as np -import hist + import importlib.util as _ilu +import os +import sys + +import hist +import numpy as np RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) _spec = _ilu.spec_from_file_location( "rabbit_fit_main", os.path.join(RABBIT_BASE, "bin", "rabbit_fit.py") ) -_rfm = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_rfm) -from rabbit import inputdata, fitter -from rabbit.tensorwriter import TensorWriter +_rfm = _ilu.module_from_spec(_spec) +_spec.loader.exec_module(_rfm) +from rabbit import fitter, inputdata from rabbit.param_models import helpers as ph +from rabbit.tensorwriter import TensorWriter NB, K = 40, 4 OUT = os.environ.get("OUT", "/tmp/claude") @@ -23,7 +27,8 @@ def build(sparse): rng = np.random.default_rng(13) - nom = np.linspace(900, 1100, NB); data = nom.copy() + nom = np.linspace(900, 1100, NB) + data = nom.copy() full = rng.poisson(nom) folds = np.stack([rng.multinomial(c, [1.0 / K] * K) for c in full], 0).T ax = hist.axis.Regular(NB, 0, 1, name="x") @@ -31,18 +36,21 @@ def build(sparse): def dh(v): h = hist.Hist(ax, storage=hist.storage.Weight()) - h.view()["value"] = v; h.view()["variance"] = v + h.view()["value"] = v + h.view()["variance"] = v return h def fh(fl): h = hist.Hist(ax, axf, storage=hist.storage.Weight()) for f in range(K): - h.view()["value"][:, f] = fl[f]; h.view()["variance"][:, f] = fl[f] + h.view()["value"][:, f] = fl[f] + h.view()["variance"][:, f] = fl[f] return h x = np.linspace(0, 1, NB) tw = TensorWriter(sparse=sparse) - tw.add_channel([ax], "ch0"); tw.add_data(dh(data), "ch0") + tw.add_channel([ax], "ch0") + tw.add_data(dh(data), "ch0") tw.add_process(fh(folds), "sig", "ch0", signal=True, fold_axis="mcfold") tw.add_systematic(dh(nom * (1 + 0.1 * (x - 0.5))), "shp", "sig", "ch0") fn = f"{OUT}/toy_sf_{'sp' if sparse else 'de'}.hdf5" @@ -52,11 +60,28 @@ def fh(fl): def run(fn): a = _rfm.make_parser().parse_args( - [fn, "-o", f"{OUT}/vo", "-t", "0", "--chisqFit", "--noBinByBinStat", - "--allowNegativeParam", "--mcStatDebias", "kfold", "--covMode", "fisher"]) + [ + fn, + "-o", + f"{OUT}/vo", + "-t", + "0", + "--chisqFit", + "--noBinByBinStat", + "--allowNegativeParam", + "--mcStatDebias", + "kfold", + "--covMode", + "fisher", + ] + ) ind = inputdata.FitInputData(fn, None) - f = fitter.Fitter(ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False) - f.defaultassign(); f.set_nobs(ind.data_obs); f.minimize() + f = fitter.Fitter( + ind, ph.load_models([["Mu"]], ind, **vars(a)), a, do_blinding=False + ) + f.defaultassign() + f.set_nobs(ind.data_obs) + f.minimize() _, _, h = f.loss_val_grad_hess() bread = f.fisher_curvature(h) S = np.asarray(f.cov_twohalf_sandwich(bread, "fisher")) @@ -70,6 +95,8 @@ def run(fn): ok_x = np.allclose(xs, xd, rtol=1e-6) ok_b = np.allclose(bs, bd, rtol=1e-6) ok_s = np.allclose(Ss, Sd, rtol=1e-6) - print(f" point match: {ok_x} fisher curvature match: {ok_b} sandwich match: {ok_s}") + print( + f" point match: {ok_x} fisher curvature match: {ok_b} sandwich match: {ok_s}" + ) assert ok_x and ok_b and ok_s, "sparse fold de-bias must equal dense" print("\n sparse fold de-biasing == dense (point, curvature, sandwich). passed.") diff --git a/tests/verify_twohalf.py b/tests/verify_twohalf.py index cc75b61..05f3b69 100644 --- a/tests/verify_twohalf.py +++ b/tests/verify_twohalf.py @@ -4,9 +4,12 @@ Compares: standard fit (no debias), two-half curvature A^-1, two-half sandwich A^-1 H A^-1. Uses a LINEAR POI (--allowNegativeParam) so the closed-form numpy reference applies. Run with OUT= pointing at toy_folds.hdf5.""" -import os, sys -import numpy as np + import importlib.util as _ilu +import os +import sys + +import numpy as np RABBIT_BASE = os.environ.get("RABBIT_BASE", ".") sys.path.insert(0, os.path.join(RABBIT_BASE, "bin")) @@ -15,21 +18,32 @@ ) _rfm = _ilu.module_from_spec(_spec) _spec.loader.exec_module(_rfm) -from rabbit import inputdata, fitter +from rabbit import fitter, inputdata from rabbit.param_models import helpers as ph ddif = np.array([1.0, -1.0]) / np.sqrt(2) def rabbit_fit(filename, debias): - argv = [filename, "-o", "/tmp/claude/verify_out", "-t", "0", - "--noBinByBinStat", "--chisqFit", "--allowNegativeParam", - "--mcStatDebias", debias] + argv = [ + filename, + "-o", + "/tmp/claude/verify_out", + "-t", + "0", + "--noBinByBinStat", + "--chisqFit", + "--allowNegativeParam", + "--mcStatDebias", + debias, + ] args = _rfm.make_parser().parse_args(argv) indata = inputdata.FitInputData(filename, None) pm = ph.load_models([["Mu"]], indata, **vars(args)) f = fitter.Fitter(indata, pm, args, do_blinding=False) - f.defaultassign(); f.set_nobs(indata.data_obs); f.minimize() + f.defaultassign() + f.set_nobs(indata.data_obs) + f.minimize() _, grad, hess = f.loss_val_grad_hess() _, cov = f.edmval_cov(grad, hess) curv = np.asarray(cov) @@ -41,23 +55,31 @@ def rabbit_fit(filename, debias): def numpy_ref(): # regenerate the SAME folds as tests/toy_twohalf.py - NB = 200; A, H1, H2 = 4962.0, 5112.0, 4962.0 + NB = 200 + A, H1, H2 = 4962.0, 5112.0, 4962.0 rng = np.random.default_rng(20240614) n_flat = np.full(NB, A) n_step = np.concatenate([np.full(NB // 2, H1), np.full(NB // 2, H2)]) data = (n_flat + n_step).astype(float) - sw_flat = rng.poisson(n_flat); sw_step = rng.poisson(n_step) - flat_A = rng.binomial(sw_flat, 0.5); flat_B = sw_flat - flat_A - step_A = rng.binomial(sw_step, 0.5); step_B = sw_step - step_A - Tf = np.stack([sw_flat, sw_step], 1).astype(float) # full - TA = 2 * np.stack([flat_A, step_A], 1).astype(float) # half A (x2) - TB = 2 * np.stack([flat_B, step_B], 1).astype(float) # half B (x2) + sw_flat = rng.poisson(n_flat) + sw_step = rng.poisson(n_step) + flat_A = rng.binomial(sw_flat, 0.5) + flat_B = sw_flat - flat_A + step_A = rng.binomial(sw_step, 0.5) + step_B = sw_step - step_A + Tf = np.stack([sw_flat, sw_step], 1).astype(float) # full + TA = 2 * np.stack([flat_A, step_A], 1).astype(float) # half A (x2) + TB = 2 * np.stack([flat_B, step_B], 1).astype(float) # half B (x2) V = data - def H(T): return np.einsum('bi,bj,b->ij', T, T, 1 / V) - def b(T): return np.einsum('bi,b,b->i', T, 1 / V, data) + def H(T): + return np.einsum("bi,bj,b->ij", T, T, 1 / V) + + def b(T): + return np.einsum("bi,b,b->i", T, 1 / V, data) + Hf = H(Tf) - A_cf = 2 * Hf - 0.5 * H(TA) - 0.5 * H(TB) # jackknife Hessian + A_cf = 2 * Hf - 0.5 * H(TA) - 0.5 * H(TB) # jackknife Hessian b_cf = 2 * b(Tf) - 0.5 * b(TA) - 0.5 * b(TB) r_std = np.linalg.solve(Hf, b(Tf)) r_cf = np.linalg.solve(A_cf, b_cf) @@ -74,17 +96,27 @@ def b(T): return np.einsum('bi,b,b->i', T, 1 / V, data) r_std, Cstd, r_cf, Ccurv, Csand = numpy_ref() print("NUMPY REFERENCE:") print(f" std r={np.round(r_std,4)} sigma(dif)={np.sqrt(ddif@Cstd@ddif):.4f}") - print(f" cf r={np.round(r_cf,4)} curv(dif)={np.sqrt(ddif@Ccurv@ddif):.4f} " - f"sandwich(dif)={np.sqrt(ddif@Csand@ddif):.4f}") + print( + f" cf r={np.round(r_cf,4)} curv(dif)={np.sqrt(ddif@Ccurv@ddif):.4f} " + f"sandwich(dif)={np.sqrt(ddif@Csand@ddif):.4f}" + ) xr_std, Cr_std, _ = rabbit_fit(fn, "none") xr_cf, Cr_curv, Cr_sand = rabbit_fit(fn, "twoHalf") print("\nRABBIT:") print(f" std r={np.round(xr_std,4)} sigma(dif)={np.sqrt(ddif@Cr_std@ddif):.4f}") - print(f" cf r={np.round(xr_cf,4)} curv(dif)={np.sqrt(ddif@Cr_curv@ddif):.4f} " - f"sandwich(dif)={np.sqrt(ddif@Cr_sand@ddif):.4f}") + print( + f" cf r={np.round(xr_cf,4)} curv(dif)={np.sqrt(ddif@Cr_curv@ddif):.4f} " + f"sandwich(dif)={np.sqrt(ddif@Cr_sand@ddif):.4f}" + ) - ok = (np.allclose(xr_cf, r_cf, atol=1e-3) - and np.isclose(np.sqrt(ddif@Cr_curv@ddif), np.sqrt(ddif@Ccurv@ddif), atol=1e-3) - and np.isclose(np.sqrt(ddif@Cr_sand@ddif), np.sqrt(ddif@Csand@ddif), atol=1e-3)) + ok = ( + np.allclose(xr_cf, r_cf, atol=1e-3) + and np.isclose( + np.sqrt(ddif @ Cr_curv @ ddif), np.sqrt(ddif @ Ccurv @ ddif), atol=1e-3 + ) + and np.isclose( + np.sqrt(ddif @ Cr_sand @ ddif), np.sqrt(ddif @ Csand @ ddif), atol=1e-3 + ) + ) print(f"\n rabbit == numpy reference: {ok}")