diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 7c2fb30..42a13c1 100755 --- a/bin/rabbit_fit.py +++ b/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, ) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index d9258ae..d3548dd 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -214,6 +214,149 @@ 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() + # 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 + # 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]) + if k % 2 != 0: + 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) + 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 + # 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: + # 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(self.param_model.xparamdefault, full=True), + [self.indata.nbinsfull, self.indata.nproc], + ) + 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 + # 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: + 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)." + ) + # 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 " + "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( + "--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)." + ) + 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, param_model, @@ -789,6 +932,322 @@ 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'. + + 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 _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 _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_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] + 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 + observed-Hessian `hess_obj`: + + F = hess_obj - H_ln_obs + F_data + + 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: + 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 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) + 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) 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_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 + + 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. + + 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 @@ -921,7 +1380,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 @@ -940,7 +1407,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, @@ -948,12 +1415,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, @@ -968,7 +1441,7 @@ def global_impacts_parms(self): self.bbstat.enabled, self.bbstat.binByBinStatMode, self.globalImpactsFromJVP, - self.cov, + cov, ) @tf.function @@ -1487,7 +1960,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): + 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); "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 @@ -1530,6 +2010,49 @@ def _compute_yields_noBBB(self, full=True, compute_norm=True): 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] + 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( + 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 @@ -1567,14 +2090,32 @@ 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 + logk_src = self.logk_A + 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 + if full or self.indata.nbinsmasked == 0: nbins = self.indata.nbinsfull - logk = self.logk - norm = self.indata.norm + logk = logk_src + norm = norm_src else: nbins = self.indata.nbins - logk = self.logk[:nbins] - norm = self.indata.norm[:nbins] + logk = logk_src[:nbins] + norm = norm_src[:nbins] if self.indata.symmetric_tensor: mlogk = tf.reshape( @@ -1878,10 +2419,50 @@ def _compute_nll_components(self, profile=True, full_nll=False): nexp = nexpfullcentral[: self.indata.nbins] - ln = self._compute_ln(nexp, full_nll) + 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. + 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) + - 0.5 * self._compute_ln(nexp_B, full_nll) + ) + else: + ln = self._compute_ln(nexp, full_nll) 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/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/inputdata.py b/rabbit/inputdata.py index 250a43f..e25a4de 100644 --- a/rabbit/inputdata.py +++ b/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(): diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 085ba1e..1c0cf85 100644 --- a/rabbit/parsing.py +++ b/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, diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 1ba08e3..f071b57 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -52,6 +52,22 @@ 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 + # Split-logk fold templates (only for systematics added with fold_axis): + # 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] + # 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 = {} @@ -187,7 +203,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 +268,114 @@ 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`. 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) + # 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. + + ``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. + + 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 + 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( @@ -258,6 +388,11 @@ 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_logkavg_folds[name] = {} + self.dict_logkhalfdiff_folds[name] = {} + self.dict_logk_folds_delta[name] = {} self.dict_beta_variations[name] = {} # add masked channels last @@ -1122,11 +1257,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 @@ -1134,6 +1279,22 @@ 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) @@ -1256,6 +1417,177 @@ 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(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 (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 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 do not support the 'linear'/'quadratic' " + "symmetrize modes (they split into two systematics)." + ) + 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 h0.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 h0.axes[fold_idx].size != k: + raise ValueError( + 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 + + # 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) + is_asym = False + for f in range(k): + norm_f = norm_folds[f] + 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 + 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, + ) + 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] = ( + logkhalfdiff_folds.astype(self.dtype) + ) + def add_beta_variations( self, h, @@ -1653,9 +1985,85 @@ 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) + # 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") @@ -1903,6 +2311,27 @@ 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: 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 + for p in self.dict_logkavg_folds[c] + ) + logk_folds = None + if any_folded_syst: + kset = {k for c in self.channels for k in self.fold_k[c].values()} + kfold = kset.pop() + 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] dict_norm_chan = self.dict_norm[chan] @@ -1917,10 +2346,43 @@ 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 {} + ) + 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: + 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: + 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] = ( dict_logkavg_proc[syst] @@ -2104,6 +2566,28 @@ 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 + + # 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, @@ -2135,6 +2619,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/rabbit/workspace.py b/rabbit/workspace.py index e069cfa..e08132b 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -70,11 +70,28 @@ 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, diff --git a/tests/mcstat_coverage.py b/tests/mcstat_coverage.py new file mode 100644 index 0000000..49480e8 --- /dev/null +++ b/tests/mcstat_coverage.py @@ -0,0 +1,244 @@ +"""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 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 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 +# 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 == "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)) + 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 + 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) + 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(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, + # 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/toy_kfold.py b/tests/toy_kfold.py new file mode 100644 index 0000000..0f9b741 --- /dev/null +++ b/tests/toy_kfold.py @@ -0,0 +1,57 @@ +"""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 hist +import numpy as np + +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/toy_mcstat.py b/tests/toy_mcstat.py new file mode 100644 index 0000000..bdb42d2 --- /dev/null +++ b/tests/toy_mcstat.py @@ -0,0 +1,62 @@ +"""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 hist +import numpy as np + +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/toy_splitlogk.py b/tests/toy_splitlogk.py new file mode 100644 index 0000000..f73b5c0 --- /dev/null +++ b/tests/toy_splitlogk.py @@ -0,0 +1,76 @@ +"""Small toy with a FOLDED process and a FOLDED systematic, to exercise the +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. 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 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 + +# K independent folds of the nominal (per-bin multinomial split) +full = rng.poisson(nom) +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 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 +) + +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(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 + + +def build(outname, split): + tw = TensorWriter(sparse=False) + tw.add_channel([ax], "ch0") + tw.add_data(dhist(data), "ch0") + tw.add_process(fhist(folds), "sig", "ch0", signal=True, fold_axis="mcfold") + if split: + 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(upfolds.sum(0)), "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), K={K}" + ) diff --git a/tests/toy_twohalf.py b/tests/toy_twohalf.py new file mode 100644 index 0000000..b381b78 --- /dev/null +++ b/tests/toy_twohalf.py @@ -0,0 +1,64 @@ +"""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 hist +import numpy as np + +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_bblite.py b/tests/verify_bblite.py new file mode 100644 index 0000000..f416af0 --- /dev/null +++ b/tests/verify_bblite.py @@ -0,0 +1,117 @@ +"""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 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 fitter, inputdata +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.") diff --git a/tests/verify_coverage.py b/tests/verify_coverage.py new file mode 100644 index 0000000..8de8822 --- /dev/null +++ b/tests/verify_coverage.py @@ -0,0 +1,62 @@ +"""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)." + ) diff --git a/tests/verify_kfold.py b/tests/verify_kfold.py new file mode 100644 index 0000000..e8c165c --- /dev/null +++ b/tests/verify_kfold.py @@ -0,0 +1,121 @@ +"""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 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 fitter, inputdata +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): + 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).") diff --git a/tests/verify_mcstat.py b/tests/verify_mcstat.py new file mode 100644 index 0000000..851bd7d --- /dev/null +++ b/tests/verify_mcstat.py @@ -0,0 +1,98 @@ +"""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 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 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", + ] + 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)" + ) diff --git a/tests/verify_sparse_splitlogk.py b/tests/verify_sparse_splitlogk.py new file mode 100644 index 0000000..779302b --- /dev/null +++ b/tests/verify_sparse_splitlogk.py @@ -0,0 +1,108 @@ +"""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 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 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") + + +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.") diff --git a/tests/verify_sparsefold.py b/tests/verify_sparsefold.py new file mode 100644 index 0000000..f9abdfb --- /dev/null +++ b/tests/verify_sparsefold.py @@ -0,0 +1,102 @@ +"""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 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 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") + + +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.") diff --git a/tests/verify_twohalf.py b/tests/verify_twohalf.py new file mode 100644 index 0000000..05f3b69 --- /dev/null +++ b/tests/verify_twohalf.py @@ -0,0 +1,122 @@ +"""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 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 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, + ] + 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}")