From 92b6f2bbdad29561d81d2e66824d4b85261ad790 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 16 May 2026 20:46:03 -0400 Subject: [PATCH] Implement observed asymptotic CLs likelihood limit The observed likelihood-based upper limit was never produced: the call in rabbit_limit.py was commented out ("TODO: make it work"), the channel/mapping case was a `pass`, and the old compute_likelihood_limit had a Jacobian bug (it differentiated sqrt(2*NLL) instead of sqrt(2*(NLL-min)), mis-scaling the constraint Jacobian fed to the joint trust-constr solve). Rewrite compute_likelihood_limit to reuse the same profiled-NLL machinery as the (working) expected limit: evaluate the profile-likelihood test statistic with the data fit (q_mu) and the background-only Asimov fit (q_{mu,A}) and root-find mu over CLs(mu) = cl_s, using the same algebra as the gaussian path (Cowan et al. arXiv:1007.1727). Add the modified statistic q-tilde_mu for a negative best-fit signal strength, consistent with the gaussian xobs<0 branch. - fitter.py: add Fitter.profiled_nll_at_poi (freeze/minimize/reduced_nll, restoring state), the literal profile-likelihood test-statistic definition. - asymptotic_limits.py: replace the buggy bespoke two-fitter solver with the CLs root-find; handle the channel/mapping case by reporting the mapped observable at the limit point. - rabbit_limit.py: wire the observed likelihood limit for both the parameter and channel cases; enable the likelihood_asymptoticLimits_observed output. - main.yml: assert the observed likelihood limit is finite and within the expected band in the bsm CI jobs (it was silently NaN before). Co-Authored-By: Claude Opus 4.7 --- .github/workflows/main.yml | 59 +++++++- bin/rabbit_limit.py | 45 +++--- rabbit/asymptotic_limits.py | 282 ++++++++++++++---------------------- rabbit/fitter.py | 27 ++++ 4 files changed, 224 insertions(+), 189 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 7543221c..dddd2819 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -429,11 +429,68 @@ jobs: rabbit_limit.py $RABBIT_OUTDIR/test_tensor_bsm.hdf5 -o $RABBIT_OUTDIR/ --postfix bsm -t 0 --unblind --expectSignal bsm1 0 --expectSignal bsm2 1 --allowNegativeParam --asymptoticLimits bsm1 --levels 0.05 - - name: bsm limits on channel + - name: check bsm observed likelihood limit + run: | + python - <<'PY' + import os + import numpy as np + from rabbit import io_tools + + fr = io_tools.get_fitresult( + os.path.join(os.environ["RABBIT_OUTDIR"], "limits_bsm.hdf5") + ) + keys = list(fr.keys()) + assert "likelihood_asymptoticLimits_observed" in keys, keys + obs = np.array( + fr["likelihood_asymptoticLimits_observed"].get().values() + ).flatten() + assert np.all(np.isfinite(obs)) and np.all(obs > 0), obs + # observed must sit inside the expected +/-2 sigma band (clb axis + # ordered 0.025 .. 0.975), with padding for the gaussian/likelihood + # method difference; this catches NaN / sign / scale regressions + exp = np.array(fr["likelihood_asymptoticLimits_expected"].get().values()) + lo = exp[..., 0].flatten() + hi = exp[..., -1].flatten() + assert np.all(obs >= 0.5 * lo) and np.all(obs <= 2.0 * hi), (obs, lo, hi) + gobs = np.array( + fr["gaussian_asymptoticLimits_observed"].get().values() + ).flatten() + assert np.all(obs > 0.3 * gobs) and np.all(obs < 3.0 * gobs), (obs, gobs) + print("likelihood observed limit OK:", obs) + PY + + - name: bsm limits on channel run: >- rabbit_limit.py $RABBIT_OUTDIR/test_tensor_bsm.hdf5 -o $RABBIT_OUTDIR/ --postfix bsm -t 0 --unblind --paramModel Mixture bsm1 sig --expectSignal bsm1_sig_mixing 0 --asymptoticLimits bsm1_sig_mixing --levels 0.05 --allowNegativeParam + - name: check bsm observed likelihood limit on channel + run: | + python - <<'PY' + import os + import numpy as np + from rabbit import io_tools + + fr = io_tools.get_fitresult( + os.path.join(os.environ["RABBIT_OUTDIR"], "limits_bsm.hdf5") + ) + keys = list(fr.keys()) + assert "likelihood_asymptoticLimits_observed" in keys, keys + obs = np.array( + fr["likelihood_asymptoticLimits_observed"].get().values() + ).flatten() + assert np.all(np.isfinite(obs)) and np.all(obs > 0), obs + exp = np.array(fr["likelihood_asymptoticLimits_expected"].get().values()) + lo = exp[..., 0].flatten() + hi = exp[..., -1].flatten() + assert np.all(obs >= 0.5 * lo) and np.all(obs <= 2.0 * hi), (obs, lo, hi) + gobs = np.array( + fr["gaussian_asymptoticLimits_observed"].get().values() + ).flatten() + assert np.all(obs > 0.3 * gobs) and np.all(obs < 3.0 * gobs), (obs, gobs) + print("likelihood observed limit on channel OK:", obs) + PY + plotting: runs-on: [self-hosted, linux, x64] needs: [setenv, fitting] diff --git a/bin/rabbit_limit.py b/bin/rabbit_limit.py index 26383ccf..cf07d708 100755 --- a/bin/rabbit_limit.py +++ b/bin/rabbit_limit.py @@ -166,6 +166,9 @@ def do_asymptotic_limits(args, fitter, ws, data_values, mapping=None, fit_data=F # this is the denominator of q for likelihood based limits nllvalreduced_asimov = fitter_asimov.reduced_nll().numpy() + if fit_data: + # global minimum on data (unconditional fit, done above) + nllvalreduced = fitter.reduced_nll().numpy() for j, cl_s in enumerate(args.levels): logger.info(f" -- AsymptoticLimits ( CLs={round(cl_s*100,1):4.1f}% ) -- ") @@ -200,16 +203,24 @@ def do_asymptotic_limits(args, fitter, ws, data_values, mapping=None, fit_data=F key, xobs, xobs_err, xerr, cl_s ) - if "likelihood" in args.modes: - # Likelihood approximation - if mapping is not None: - # TODO: implement for channels - pass - else: - # TODO: make it work - # nllvalreduced = fitter.reduced_nll().numpy() - # limits_nll_obs[i, j] = asymptotic_limits.compute_likelihood_limit(fitter, fitter_asimov, nllvalreduced, nllvalreduced_asimov, key, cl_s) - pass + if "likelihood" in args.modes and fit_data: + # Likelihood approximation (covers both the parameter case + # fun=None and the channel/mapping case fun=mapping.compute_flat) + limits_nll_obs[i, j] = asymptotic_limits.compute_likelihood_limit( + fitter, + fitter_asimov, + nllvalreduced, + nllvalreduced_asimov, + key, + cl_s, + fun=fun, + r_init=( + limits_obs[i, j] + if (mapping is None and "gaussian" in args.modes) + else None + ), + allow_negative=args.allowNegativeParam, + ) if "gaussian" in args.modes: ws.add_limits_hist( @@ -228,7 +239,6 @@ def do_asymptotic_limits(args, fitter, ws, data_values, mapping=None, fit_data=F ) if "likelihood" in args.modes: - # TODO: implement for channels ws.add_limits_hist( limits_nll, args.asymptoticLimits, @@ -237,13 +247,12 @@ def do_asymptotic_limits(args, fitter, ws, data_values, mapping=None, fit_data=F base_name="likelihood_asymptoticLimits_expected", ) - # TODO: make it work - # ws.add_limits_hist( - # limits_nll_obs, - # args.asymptoticLimits, - # args.levels, - # base_name="likelihood_asymptoticLimits_observed", - # ) + ws.add_limits_hist( + limits_nll_obs, + args.asymptoticLimits, + args.levels, + base_name="likelihood_asymptoticLimits_observed", + ) def main(): diff --git a/rabbit/asymptotic_limits.py b/rabbit/asymptotic_limits.py index dee1e244..5e528e6c 100644 --- a/rabbit/asymptotic_limits.py +++ b/rabbit/asymptotic_limits.py @@ -1,189 +1,131 @@ import numpy as np -import scipy.optimize -import tensorflow as tf from scipy.optimize import root from wums import logging -from rabbit import tfhelpers as tfh +import tensorflow as tf # isort: skip -logger = logging.child_logger(__name__) - - -@tf.function -def limit_constraint(f): - val = f() - return val +from rabbit import tfhelpers as tfh # isort: skip - -@tf.function -def limit_constraint_val_grad(fitter, fitter_asimov, f): - with tf.GradientTape() as t: - t.watch([fitter.x, fitter_asimov.x]) - val = f() - grad = t.gradient(val, [fitter.x, fitter_asimov.x]) - grad = tf.concat(grad, 0) - - return val, grad +logger = logging.child_logger(__name__) def compute_likelihood_limit( - fitter, fitter_asimov, nllvalreduced, nllvalreduced_asimov, param, cl + fitter, + fitter_asimov, + nllvalreduced, + nllvalreduced_asimov, + param, + cl_s, + muhat=None, + fun=None, + r_init=None, + allow_negative=True, ): - - # We have 2 fitters, one to fit the asimov dataset and one to fit the real data - # The two are allowed to have different parameters; but the POI has to be the same - - # initial values + """Observed asymptotic CLs upper limit from the profile likelihood. + + Reuses the same profiled-NLL machinery as the (working) expected limit: + the profile-likelihood test statistic is evaluated with `fitter` fitted to + data (q_mu) and with `fitter_asimov` fitted to the background-only Asimov + dataset (q_{mu,A}). The limit is the value of `param` for which + CLs(mu) = cl_s, following the asymptotic formulae of Cowan, Cranmer, + Gross, Vitells, arXiv:1007.1727. + + For a negative best-fit signal strength the modified test statistic + q-tilde_mu (Eq. 16) is used, taking the mu=0 conditional fit as reference + instead of the global minimum, consistent with the `xobs < 0` branch of + `compute_gaussian_limit`. + + If `fun` is given (channel/mapping case) the limit is reported on the + mapped observable at the limit point, matching the convention of the + expected channel limit produced via `contour_scan(fun=...)`. + """ idx = np.where(fitter.parms.astype(str) == param)[0][0] - xval = tf.identity(fitter.x).numpy() - - xval_asimov = np.zeros_like(xval) - xval_asimov[idx] = xval[idx] - fitter_asimov.x.assign(xval_asimov) - - # perform an asumov fit at xobs to get good starting values - fitter_asimov.freeze_params(param) - fitter_asimov.minimize() - fitter_asimov.defreeze_params(param) - - xval_asimov = tf.identity(fitter_asimov.x).numpy() - xval_asimov = np.delete(xval_asimov, idx) - xval_init = np.concatenate([xval, xval_asimov]) - - # to find the upper limit we set the mu equal to mu_obs also for the asimov dataset - # xval_init[len(xval_init)//2:][idx] = xval_init[idx] - - # TODO: modified statistics - - def _compute_limit_constraint(): - qmu = tf.sqrt(2 * (fitter._compute_nll() - nllvalreduced)) - qA = tf.sqrt(2 * (fitter_asimov._compute_nll() - nllvalreduced_asimov)) - constraint = cl * tfh.normal_cdf(qA - qmu) + tfh.normal_cdf(qmu) - 1 - - logger.debug(f"q_mu = {qmu}") - logger.debug(f"q_A = {qA}") - logger.debug(f"constraint = {constraint}") - - logger.debug(f"fitter_asimov.x = {fitter_asimov.x}") - logger.debug(f"fitter.x = {fitter.x}") - - return constraint - - def _compute_limit_constraint_grad(): - qmu = tf.sqrt(2 * (fitter._compute_nll() - nllvalreduced)) - qA = tf.sqrt(2 * (fitter_asimov._compute_nll() - nllvalreduced_asimov)) - - phi_qmu = tfh.normal_pdf(qmu) - phi_qA_qmu = tfh.normal_pdf(qA - qmu) - - with tf.GradientTape() as t: - val = tf.sqrt(2 * fitter._compute_nll()) - dqmu_dx = t.gradient(val, fitter.x) - - with tf.GradientTape() as t: - val = tf.sqrt(2 * fitter_asimov._compute_nll()) - dqA_dx = t.gradient(val, fitter_asimov.x) - - # from first fitter - dconstraint_dx1 = (phi_qmu - cl * phi_qA_qmu) * dqmu_dx + if muhat is None: + # best-fit POI in the internal (self.x) space, consistent with + # profiled_nll_at_poi (also correct for the channel/mapping case + # where `param` is the POI driving the mapped observable) + muhat = float(fitter.x.numpy()[idx]) - # from second fitter - dconstraint_dx2 = cl * phi_qA_qmu * dqA_dx - - dconstraint_dx1 = dconstraint_dx1.numpy() - dconstraint_dx2 = dconstraint_dx2.numpy() - - # common x[idx] - dconstraint_dx1[idx] = dconstraint_dx1[idx] + dconstraint_dx2[idx] - - # drop x[idx] from second fitter - dconstraint_dx2 = np.delete(dconstraint_dx2, idx) - - dconstraint_dx = np.concatenate([dconstraint_dx1, dconstraint_dx2]) - - logger.debug(f"q_mu = {qmu}") - logger.debug(f"q_A = {qA}") - logger.debug(f"dqmu_dx = {dqmu_dx}") - logger.debug(f"dqA_dx = {dqA_dx}") - logger.debug(f"dconstraint_dx1 = {dconstraint_dx1}") - logger.debug(f"dconstraint_dx2 = {dconstraint_dx2}") - - logger.debug(f"dconstraint/dx = {dconstraint_dx}") - - return dconstraint_dx - - def scipy_constraint(xval): - logger.info(f"scipy_constraint: xval = {xval}") - - nx = (len(xval) + 1) // 2 - x1 = xval[:nx] - x2 = xval[nx:] - x2 = np.insert(x2, idx, x1[idx]) - - fitter.x.assign(x1) - fitter_asimov.x.assign(x2) - val = _compute_limit_constraint() - # val = limit_constraint(_compute_limit_constraint) - return val.numpy() - - def scipy_constraint_grad(xval): - logger.info(f"scipy_constraint_grad: xval = {xval}") - - nx = (len(xval) + 1) // 2 - x1 = xval[:nx] - x2 = xval[nx:] - x2 = np.insert(x2, idx, x1[idx]) - - fitter.x.assign(x1) - fitter_asimov.x.assign(x2) - - # val, grad = limit_constraint_val_grad(fitter, fitter_asimov, _compute_limit_constraint) - grad = _compute_limit_constraint_grad() - return grad - - nlc = scipy.optimize.NonlinearConstraint( - fun=scipy_constraint, - lb=-0.5, # -0.5 at qmu = 0 - ub=0, - jac=scipy_constraint_grad, - hess=scipy.optimize.SR1(), # TODO: use exact hessian or hessian vector product - ) - - # Objective function and its derivatives - def objective(params): - logger.debug(f"param[{idx}] = {params[idx]}") - return -params[idx] - - def objective_jac(params): - jac = np.zeros_like(params) - jac[idx] = -0.001 - return jac - - def objective_hessp(params, v): - return np.zeros_like(v) - - res = scipy.optimize.minimize( - objective, - xval_init, - method="trust-constr", - jac=objective_jac, - hessp=objective_hessp, - constraints=[nlc], - options={ - "maxiter": 500, - "xtol": 1e-10, - "gtol": 1e-10, - # "verbose": 3 - }, + logger.info( + f"Compute observed (Likelihood) limit for {param}, " + f"mu_hat = {muhat}, CLs = {cl_s}" ) - if not res["success"]: - logger.warning("No success") - - limit = res["x"][idx] + # Reference NLL for the data test statistic q_mu (denominator of the + # profile likelihood ratio). Modified statistics (mu_hat < 0): use the + # mu = 0 conditional fit as reference instead of the global minimum. + if allow_negative and muhat < 0: + logger.debug("Use modified statistics (mu_hat < 0)") + ref = fitter.profiled_nll_at_poi(param, 0.0) + mu_lo = 0.0 + else: + ref = nllvalreduced + mu_lo = muhat + + def sqrt_qmu(mu): + # one-sided test statistic for an upper limit + if mu <= mu_lo: + return 0.0 + dnll = fitter.profiled_nll_at_poi(param, mu) - ref + return np.sqrt(max(0.0, 2.0 * dnll)) + + def sqrt_qA(mu): + # background-only Asimov: best fit at mu = 0, reference is its minimum + if mu <= 0.0: + return 0.0 + dnll = fitter_asimov.profiled_nll_at_poi(param, mu) - nllvalreduced_asimov + return np.sqrt(max(0.0, 2.0 * dnll)) + + def cls_minus_alpha(mu): + mu = float(np.asarray(mu).reshape(-1)[0]) + # np.float64 (not python float) so tfh.normal_cdf can read .dtype + qmu = np.float64(sqrt_qmu(mu)) + qA = np.float64(sqrt_qA(mu)) + # CLs+b - cl_s * CLb ; same algebra as compute_gaussian_limit + val = tfh.normal_cdf(-qmu).numpy() - cl_s * tfh.normal_cdf(qA - qmu).numpy() + logger.debug( + f"mu = {mu}, sqrt(q_mu) = {qmu}, sqrt(q_A) = {qA}, " + f"CLs+b - cl_s*CLb = {val}" + ) + return val + + if r_init is None or not np.isfinite(r_init) or r_init <= mu_lo: + # scale-aware default: ~2 sigma above the best fit + sigma = float(fitter.cov[idx, idx].numpy()) ** 0.5 + if not np.isfinite(sigma) or sigma <= 0: + sigma = 1.0 + r_init = max(mu_lo, muhat) + 2.0 * sigma + + res = root(cls_minus_alpha, r_init) + if not res.success: + logger.warning(f"Root finding for observed limit on {param} did not converge") + + mu_limit = float(np.asarray(res.x).reshape(-1)[0]) + + if fun is not None: + # profile at the limit POI and report the mapped observable, same + # convention as the expected channel limit (contour_scan(fun=...)) + xsave = tf.identity(fitter.x) + xval = xsave.numpy() + xval[idx] = mu_limit + fitter.x.assign(xval) + fitter.freeze_params(param) + fitter.minimize() + fitter.defreeze_params(param) + exp, *_ = fitter.expected_with_variance( + fun, + profile=True, + compute_cov=False, + compute_global_impacts=False, + need_observables=True, + inclusive=True, + ) + fitter.x.assign(xsave) + limit = float(exp.numpy()[0]) + else: + limit = mu_limit logger.info(f"Observed (Likelihood): {param} < {limit}") - return limit diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 000692fe..54bdf3e1 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -2147,6 +2147,33 @@ def objective_hessp(x, pval): return intervals, params_values + def profiled_nll_at_poi(self, param, value): + """Reduced NLL profiled over all other parameters with `param` fixed. + + `value` is interpreted in the internal parameter (`self.x`) space, the + same convention used by `contour_scan` (for `allowNegativeParam` POIs + this coincides with the physical signal strength). The unconditional + best-fit state of `self.x` is restored before returning so the fitter + is left untouched. + """ + xsave = tf.identity(self.x) + idx = np.where(self.parms.astype(str) == param)[0][0] + + xval = xsave.numpy() + xval[idx] = value + self.x.assign(xval) + + self.freeze_params(param) + self.minimize() + self.defreeze_params(param) + + nll = self.reduced_nll().numpy() + + # restore the unconditional best fit + self.x.assign(xsave) + + return nll + def contour_scan2D(self, param_tuple, nll_min, cl=1, n_points=16): # Not yet working def scipy_loss(xval):