asymmetric traditional impacts; unblinding groups

bin/rabbit_fit.py

@@ -226,6 +226,107 @@ def make_parser():
         action="store_true",
         help="compute impacts of frozen (non-profiled) systematics",
     )
+    parser.add_argument(
+        "--asymImpacts",
+        default=False,
+        action="store_true",
+        help="Compute traditional asymmetric impacts on POIs by running a "
+        "Delta(2NLL)=1 contour scan per nuisance. Skips structurally "
+        "symmetric and unconstrained nuisances by default.",
+    )
+    parser.add_argument(
+        "--asymImpactsInclude",
+        default=None,
+        nargs="+",
+        help="Regex(es) restricting which nuisances are scanned for --asymImpacts.",
+    )
+    parser.add_argument(
+        "--asymImpactsExclude",
+        default=None,
+        nargs="+",
+        help="Regex(es) excluding nuisances from --asymImpacts.",
+    )
+    parser.add_argument(
+        "--asymImpactsAll",
+        default=False,
+        action="store_true",
+        help="Disable the structural-symmetric skip in --asymImpacts and scan "
+        "every constrained nuisance regardless of template content.",
+    )
+    parser.add_argument(
+        "--asymImpactsHess",
+        default="exact",
+        choices=["exact", "hvp", "frozen", "bfgs", "sr1"],
+        help="Constraint-Hessian mode for the contour-scan in --asymImpacts. "
+        "'exact' rebuilds the full N x N NLL Hessian each iteration (slow, "
+        "reference). 'hvp' uses Hessian-vector products via a LinearOperator "
+        "(exact, typically fastest for large N). 'frozen' uses cov^-1 from "
+        "the postfit as a constant Hessian (cheapest but produces silent "
+        "failures on non-Gaussian profiles -- speed reference only). "
+        "'bfgs'/'sr1' use a quasi-Newton estimate built up from the gradient "
+        "sequence (no extra TF calls per iteration; may need more iterations).",
+    )
+    parser.add_argument(
+        "--asymImpactsMaxiter",
+        default=200,
+        type=int,
+        help="trust-constr maxiter for the --asymImpacts contour-scan. Lower "
+        "values cap the cost of nuisances that don't converge (useful for "
+        "sanity benchmarks); higher values give genuinely-non-Gaussian "
+        "nuisances more chances to converge.",
+    )
+    parser.add_argument(
+        "--asymImpactsTol",
+        default=1e-6,
+        type=float,
+        help="trust-constr xtol/gtol for the --asymImpacts contour-scan. "
+        "Looser values (e.g. 1e-3) terminate before the iterate is on the "
+        "Delta(2NLL)=q contour, producing silent constraint violations. "
+        "Tighter values (1e-5, 1e-6) are slower with no benefit unless your "
+        "fit has nuisances whose profile is far from quadratic.",
+    )
+    parser.add_argument(
+        "--globalAsymImpacts",
+        default=False,
+        action="store_true",
+        help="Compute fully likelihood-based asymmetric global impacts on POIs "
+        "by shifting each constrained nuisance's theta0 by +/- 1 prefit sigma "
+        "and re-running the fit. In the Gaussian limit this reproduces "
+        "--gaussianGlobalImpacts; deviations measure non-Gaussianity of the "
+        "joint profile. Cost is N_selected x 2 full minimizations -- gate with "
+        "--globalAsymImpactsInclude in practice.",
+    )
+    parser.add_argument(
+        "--globalAsymImpactsInclude",
+        default=None,
+        nargs="+",
+        help="Regex(es) restricting which nuisances are scanned for "
+        "--globalAsymImpacts.",
+    )
+    parser.add_argument(
+        "--globalAsymImpactsExclude",
+        default=None,
+        nargs="+",
+        help="Regex(es) excluding nuisances from --globalAsymImpacts.",
+    )
+    parser.add_argument(
+        "--globalAsymImpactsSigma",
+        default=1.0,
+        type=float,
+        help="theta0 shift magnitude for --globalAsymImpacts, in units of the "
+        "prefit constraint width (1.0 = 1 prefit sigma).",
+    )
+    parser.add_argument(
+        "--globalAsymImpactsLinearWarmstart",
+        default=False,
+        action="store_true",
+        help="EXPERIMENTAL: warm-start each --globalAsymImpacts refit at the "
+        "Gaussian-approximation new minimum x_nom + dxdtheta0[:, i] * shift "
+        "(same Jacobian as --gaussianGlobalImpacts). On near-Gaussian "
+        "nuisances this should reduce per-nuisance refit cost by 10-50x. "
+        "Adds one --gaussianGlobalImpacts-equivalent precompute up front. "
+        "Off by default until validated on real tensors.",
+    )
     parser.add_argument(
         "--lCurveScan",
         default=False,
@@ -358,6 +459,37 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False):
                 logger.info(f"    2*deltaNLL: {round(chi2val, 2)}")
                 logger.info(rf"    p-value: {round(p_val * 100, 2)}%")
 
+                # DEBUG: decompose chi2_sat = 2*(NLL_nom - NLL_sat)
+                #   = 2*(ln_nom-ln_sat)    [data-vs-prediction misfit]
+                #   + 2*(lc_nom-lc_sat)    [constrained-nuisance pull cost]
+                #   + 2*(lbeta_nom-lbeta_sat) [BBB pull cost]
+                ln_sat_t, lc_sat_t, lbeta_sat_t, _, _ = (
+                    fitter_saturated._compute_nll_components(profile=True)
+                )
+                ln_sat = float(ln_sat_t.numpy())
+                lc_sat = float(lc_sat_t.numpy())
+                lbeta_sat = (
+                    float(lbeta_sat_t.numpy()) if lbeta_sat_t is not None else 0.0
+                )
+                d_ln = 2.0 * (ws.results["ln_nom"] - ln_sat)
+                d_lc = 2.0 * (ws.results["lc_nom"] - lc_sat)
+                d_lbeta = 2.0 * (ws.results["lbeta_nom"] - lbeta_sat)
+                d_sum = d_ln + d_lc + d_lbeta
+                logger.info(f"Saturated chi2 decomposition for {mapping.key}:")
+                logger.info(f"    2*Delta ln    (data-vs-pred) : {d_ln:.4f}")
+                logger.info(f"    2*Delta lc    (nuis. pulls)  : {d_lc:.4f}")
+                logger.info(f"    2*Delta lbeta (BBB pulls)    : {d_lbeta:.4f}")
+                logger.info(
+                    f"    sum                          : {d_sum:.4f}"
+                    f"  (closure vs chi2val={chi2val:.4f}: {d_sum - chi2val:+.4f})"
+                )
+                if chi2val > 0:
+                    logger.info(
+                        f"    fractions: data {d_ln / chi2val * 100:5.1f}%  "
+                        f"nuis {d_lc / chi2val * 100:5.1f}%  "
+                        f"BBB  {d_lbeta / chi2val * 100:5.1f}%"
+                    )
+
                 ws.add_chi2(
                     chi2val, ndf, prefit, mapping, saturated=True, edmval=edmval
                 )
@@ -434,24 +566,41 @@ def fit(args, fitter, ws, dofit=True):
         ws.add_1D_integer_hist([max_curvature], "best", "lcurve")
 
     if dofit:
+        _t_min = time.perf_counter()
         cb = fitter.minimize()
+        logger.debug(
+            f"[timing] fitter.minimize(): {time.perf_counter() - _t_min:.1f} s"
+        )
 
         # force profiling of beta with final parameter values
         # TODO avoid the extra calculation and jitting if possible since the relevant calculation
         # usually would have been done during the minimization
         if fitter.bbstat.enabled and not args.noPostfitProfileBB:
+            _t_bb = time.perf_counter()
             logger.info(f"profile beta")
             fitter._profile_beta()
+            logger.debug(
+                f"[timing] _profile_beta(): {time.perf_counter() - _t_bb:.1f} s"
+            )
 
         if cb is not None:
             ws.add_1D_integer_hist(cb.loss_history, "epoch", "loss")
             ws.add_1D_integer_hist(cb.time_history, "epoch", "time")
+            try:
+                n_iter = len(cb.loss_history)
+                logger.debug(f"[timing] L-BFGS: {n_iter} epochs recorded")
+            except Exception:
+                pass
 
     # prefit variances as the default fallback for add_parms_hist below
     parms_variances = None
 
-    # take covariance from externalPostfit in case the fit was skipped
-    if dofit or args.externalPostfit is None:
+    # take covariance from externalPostfit in case the fit was skipped —
+    # unless a Hessian is explicitly wanted (no --noHessian): the two-pass
+    # straight-through covariance recipe (--externalPostfit ... --noFit,
+    # WITHOUT --noHessian) relies on recomputing the Hessian at the loaded
+    # postfit point.
+    if dofit or args.externalPostfit is None or not args.noHessian:
         if not args.noEDM and not args.noHessian:
             # compute the covariance matrix and estimated distance to minimum
             _, grad, hess = fitter.loss_val_grad_hess()
@@ -499,13 +648,26 @@ def fit(args, fitter, ws, dofit=True):
             # Hessian-vector products. The CG solves touch O(npar) memory
             # per call instead of O(npar^2), so this works on problems
             # where the full covariance would be infeasible.
+            _t_lg = time.perf_counter()
             _, grad = fitter.loss_val_grad()
+            logger.debug(
+                f"[timing] loss_val_grad() (postfit): {time.perf_counter() - _t_lg:.1f} s"
+            )
+
             npoi = int(fitter.param_model.npoi)
             noi_idx_in_x = np.asarray(fitter.indata.noiidxs, dtype=np.int64) + npoi
             poi_noi_idx = np.concatenate(
                 [np.arange(npoi, dtype=np.int64), noi_idx_in_x]
             )
+            logger.debug(
+                f"[timing] EDM CG: solving for {len(poi_noi_idx)} POI+NOI rows"
+            )
+
+            _t_edm = time.perf_counter()
             edmval, cov_rows = fitter.edmval_cov_rows_hessfree(grad, poi_noi_idx)
+            logger.debug(
+                f"[timing] edmval_cov_rows_hessfree(): {time.perf_counter() - _t_edm:.1f} s"
+            )
             logger.info(f"edmval: {edmval}")
 
             # Build a full-length variance vector with the POI+NOI entries
@@ -520,6 +682,22 @@ def fit(args, fitter, ws, dofit=True):
 
     nllvalreduced = fitter.reduced_nll().numpy()
 
+    # DEBUG: decompose nominal-fit NLL into (Poisson data, nuisance constraints, BBB)
+    # so the per-projection saturated chi2 below can be split into:
+    #   2*Delta ln  -> data-vs-prediction misfit
+    #   2*Delta lc  -> constrained-nuisance pull cost
+    #   2*Delta lb  -> BBB pull cost
+    ln_nom_t, lc_nom_t, lbeta_nom_t, _, _ = fitter._compute_nll_components(profile=True)
+    ln_nom = float(ln_nom_t.numpy())
+    lc_nom = float(lc_nom_t.numpy())
+    lbeta_nom = float(lbeta_nom_t.numpy()) if lbeta_nom_t is not None else 0.0
+    logger.info("NLL decomposition (nominal postfit):")
+    logger.info(f"    ln (Poisson data) : {ln_nom:.4f}")
+    logger.info(f"    lc (nuis. constr.): {lc_nom:.4f}")
+    logger.info(f"    lbeta (BBB)       : {lbeta_nom:.4f}")
+    logger.info(f"    sum               : {ln_nom + lc_nom + lbeta_nom:.4f}")
+    logger.info(f"    nllvalreduced     : {nllvalreduced:.4f}")
+
     ndfsat = (
         tf.size(fitter.nobs)
         - fitter.param_model.nparams
@@ -537,6 +715,9 @@ def fit(args, fitter, ws, dofit=True):
     ws.results.update(
         {
             "nllvalreduced": nllvalreduced,
+            "ln_nom": ln_nom,
+            "lc_nom": lc_nom,
+            "lbeta_nom": lbeta_nom,
             "ndfsat": ndfsat,
             "edmval": edmval,
             "postfit_profile": not args.noPostfitProfileBB,
@@ -560,6 +741,32 @@ def fit(args, fitter, ws, dofit=True):
             *fitter.nonprofiled_impacts_parms(), base_name="nonprofiled_impacts_asym"
         )
 
+    if args.asymImpacts:
+        ws.add_impacts_asym_hist(
+            *fitter.asym_impacts_parms(
+                nll_min=fitter.reduced_nll().numpy(),
+                include=args.asymImpactsInclude,
+                exclude=args.asymImpactsExclude,
+                skip_symmetric=not args.asymImpactsAll,
+                hess_mode=args.asymImpactsHess,
+                contour_xtol=args.asymImpactsTol,
+                contour_gtol=args.asymImpactsTol,
+                contour_maxiter=args.asymImpactsMaxiter,
+            ),
+            base_name="impacts_asym",
+        )
+
+    if args.globalAsymImpacts:
+        ws.add_impacts_asym_hist(
+            *fitter.global_asym_impacts_parms(
+                include=args.globalAsymImpactsInclude,
+                exclude=args.globalAsymImpactsExclude,
+                sigma=args.globalAsymImpactsSigma,
+                linear_warmstart=args.globalAsymImpactsLinearWarmstart,
+            ),
+            base_name="global_impacts_asym",
+        )
+
     # Likelihood scans
     if args.scan is not None:
         parms = np.array(fitter.parms).astype(str) if len(args.scan) == 0 else args.scan
@@ -714,6 +921,17 @@ def main():
         "nois": ifitter.parms[ifitter.param_model.nparams :][indata.noiidxs],
     }
 
+    # ParamModel Gaussian priors (if --paramModelPriors was used and the
+    # model declared sigmas). Stored as a small dict so downstream tooling
+    # can see what was applied without parsing the rabbit log.
+    if getattr(ifitter, "param_prior_active", False):
+        meta["param_priors"] = {
+            "params": ifitter.param_model.params,  # all nparams names
+            "mask": ifitter.param_prior_mask_np,  # bool array
+            "sigmas": ifitter.param_prior_sigmas_np,  # NaN where mask False
+            "means": ifitter.param_prior_means_np,  # NaN where mask False
+        }
+
     with workspace.Workspace(
         args.outpath,
         args.outname,
@@ -748,6 +966,19 @@ def main():
             for j, (data_values, data_variances) in enumerate(datasets):
 
                 ifitter.defaultassign()
+                # If --externalPostfit was supplied AND this is an
+                # asimov toy (-t -1), load the postfit values BEFORE we
+                # compute the expected yield. Otherwise expected_yield()
+                # below is evaluated at the prefit point (because
+                # defaultassign() just reset x), so we'd get a
+                # prefit-Asimov regardless of --externalPostfit. The
+                # in-fit() load (around line 551) then re-runs harmlessly.
+                if ifit == -1 and args.externalPostfit is not None:
+                    ifitter.load_fitresult(
+                        args.externalPostfit,
+                        args.externalPostfitResult,
+                        profile=not args.noPostfitProfileBB,
+                    )
                 if ifit == -1:
                     group.append("asimov")
                     ifitter.set_nobs(ifitter.expected_yield())

bin/rabbit_print_pulls_and_constraints.py

@@ -50,6 +50,12 @@ def make_parser():
         default=None,
         help="Exclude nuisances by regular expression",
     )
+    parser.add_argument(
+        "--noPrefit",
+        default=False,
+        action="store_true",
+        help="Suppress the prefit pull and constraint columns from the printout",
+    )
     return parser
 
 
@@ -109,15 +115,22 @@ def main():
             constraints_asym = constraints_asym[order]
 
     nround = 5
+    prefit_header = (
+        "" if args.noPrefit else f" ({'pull prefit':>11} +/- {'constraint prefit':>17})"
+    )
     if args.asym:
-        print(
-            f"   {'Parameter':<30} {'pull':>6} +/- {'constraint':>10} + {'up':>10} - {'down':>10} ({'pull prefit':>11} +/- {'constraint prefit':>17})"
-        )
-        print("   " + "-" * 100)
+        header = f"   {'Parameter':<30} {'pull':>6} +/- {'constraint':>10} + {'up':>10} - {'down':>10}{prefit_header}"
+        print(header)
+        print("   " + "-" * (len(header) - 3))
         print(
             "\n".join(
                 [
-                    f"   {l:<30} {round(p, nround):>6} +/- {round(c, nround):>10} + {round(c_asym[0], nround):>10} - {round(c_asym[1], nround):>10} ({round(pp, nround):>11} +/- {round(pc, nround):>17})"
+                    f"   {l:<30} {round(p, nround):>6} +/- {round(c, nround):>10} + {round(c_asym[0], nround):>10} - {round(c_asym[1], nround):>10}"
+                    + (
+                        ""
+                        if args.noPrefit
+                        else f" ({round(pp, nround):>11} +/- {round(pc, nround):>17})"
+                    )
                     for l, p, c, c_asym, pp, pc in zip(
                         labels,
                         pulls,
@@ -130,14 +143,20 @@ def main():
             )
         )
     else:
-        print(
-            f"   {'Parameter':<30} {'pull':>6} +/- {'constraint':>10} ({'pull prefit':>11} +/- {'constraint prefit':>17})"
+        header = (
+            f"   {'Parameter':<30} {'pull':>6} +/- {'constraint':>10}{prefit_header}"
         )
-        print("   " + "-" * 100)
+        print(header)
+        print("   " + "-" * (len(header) - 3))
         print(
             "\n".join(
                 [
-                    f"   {l:<30} {round(p, nround):>6} +/- {round(c, nround):>10} ({round(pp, nround):>11} +/- {round(pc, nround):>17})"
+                    f"   {l:<30} {round(p, nround):>6} +/- {round(c, nround):>10}"
+                    + (
+                        ""
+                        if args.noPrefit
+                        else f" ({round(pp, nround):>11} +/- {round(pc, nround):>17})"
+                    )
                     for l, p, c, pp, pc in zip(
                         labels, pulls, constraints, pulls_prefit, constraints_prefit
                     )