Limited-MC-statistics de-biasing for binned template fits (continuous-M + cross-fit)

bin/rabbit_fit.py

@@ -540,9 +540,54 @@ 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_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). 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")
+            _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 == "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 "
+                        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})"
+                    )
+
             ws.add_cov_hist(cov)
 
             fitter.cov.assign(cov)
@@ -558,13 +603,29 @@ 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,
                 )

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

rabbit/inputdata.py

@@ -172,6 +172,35 @@ 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
+
+            # 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
+
+            # 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():

rabbit/parsing.py

@@ -377,6 +377,49 @@ 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", "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 (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",
+        default=2,
+        type=int,
+        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",
+        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,