"""
Binned Likelihood and Feldman-Cousins Toy Generation Module
This file contains the core mathematical and statistical routines required for
performing binned frequentist analyses, tailored for the Feldman-Cousins construction.
It implements expected rate calculations, standard Poisson and finite-Monte Carlo
(Poisson-Gamma mixture) likelihood evaluations, grid-based profile likelihood
optimizations, and parametric bootstrap pseudo-experiment (toy) generation.
Created: v0.1.0 (July 24, 2026)
Author: Mauricio Bustamante (mbustamante@gmail.com)
This file was released as part of the PyFC code, stored at
https://github.com/mbustama/FeldmanCousins, which exists under a GNU GPL v3 License.
"""
import math
import numpy as np
# --- 1. Dynamic Dependency Injection for Binned Math ---
try:
from numba import njit, prange, set_num_threads
NUMBA_AVAILABLE = True
print("Numba detected: JIT compilation enabled for maximum speed.")
except ImportError:
NUMBA_AVAILABLE = False
print("Numba not found: Falling back to pure Python (will be significantly slower).")
def njit(*args, **kwargs):
def wrapper(func):
return func
if len(args) == 1 and callable(args[0]):
return args[0]
return wrapper
prange = range
def set_num_threads(n):
pass
# --- 2. Core Math (Binned) ---
[docs]
@njit(fastmath=True, nogil=True)
def calc_nll(params, N_obs, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func):
"""
[BINNED] Computes the negative log-likelihood for arbitrary N parameters.
Statistical Theory & Assumptions:
This function computes the negative log-likelihood (NLL) for comparing observed
data counts against expected model counts. It handles two distinct statistical regimes:
1. Standard Poisson Likelihood (use_finite_mc = False):
Assumes template predictions have zero statistical uncertainty (infinite MC).
It uses the Baker-Cousins chi-square equivalent form (based on the likelihood
ratio with a saturated model).
Expression: NLL = 2 * SUM [ mu_i - n_i + n_i * ln(n_i / mu_i) ]
2. Finite Monte Carlo Likelihood (use_finite_mc = True):
Accounts for statistical uncertainty in the simulation templates. It marginalizes
over the unknown true Poisson rate using a Gamma conjugate prior, resulting in a
Negative Binomial (Poisson-Gamma mixture) likelihood.
Expression:
alpha = (mu^2 / sigma^2) + 1
beta = mu / sigma^2
ln(L) = alpha*ln(beta) + ln(Gamma(n+alpha)) - (n+alpha)*ln(1+beta) - ln(Gamma(alpha))
NLL = -2 * ln(L)
This is exactly L_Eff, Eq. (3.16) of Arguelles, Schneider & Yuan,
"A binned likelihood for stochastic models", JHEP 06 (2019) 030
[arXiv:1901.04645] (their main, recommended result; see also
docs/source/refs.bib's `Arguelles2019izp` entry and
docs/source/methodology.rst). alpha/beta here are their Eq. (3.12)
with the uniform-prior choice a=1, b=0 in their Eq. (3.17) -- the
"+1" in alpha is a deliberate, load-bearing part of that choice, not
a typo: it is what distinguishes L_Eff from the paper's alternative
L_Mean parameterization (a=b=0, alpha = mu^2/sigma^2 with no "+1"),
which matches mean-and-variance moment-matching but has worse
coverage properties in the paper's own toy-experiment tests (their
Sec. 4.2, Fig. 5). The paper's own ln(k!) term is omitted here, as
it is a data-only constant that cancels in any NLL difference
(the same convention documented in `unbinned.calc_nll_unbinned`'s
own dropped ln(N_obs!) term).
Numerically stable evaluation (not the naive formula above): for
well-simulated templates (small sigma^2 relative to mu^2), alpha
grows large, and evaluating `ln(Gamma(n+alpha)) - ln(Gamma(alpha))`
as a literal difference of two `lgamma` calls loses catastrophic
precision -- both terms individually scale like alpha*ln(alpha)
while their difference is only O(n*ln(alpha)), so double precision's
~16 significant digits get eaten by the leading digits both terms
share. Verified directly against a 50-digit `mpmath` reference: at
mu=1000, sigma^2=1e-9 (a realistic well-simulated-template regime),
the naive difference is off by 7.5 in absolute NLL -- large enough
to distort t = NLL_cond - NLL_uncond and bias interval boundaries.
Fixed via two exact reformulations, not an approximation:
(a) `ln(Gamma(n+alpha)) - ln(Gamma(alpha))` is computed as the exact
integer-recurrence sum `SUM_{k=0}^{n-1} ln(alpha + k)` (valid since
n_obs is always a non-negative integer count), which adds n
well-conditioned terms instead of subtracting two enormous ones;
(b) `alpha*ln(beta) - (n+alpha)*ln(1+beta)` is rewritten as
`-alpha*log1p(1/beta) - n*log1p(beta)`, avoiding the same kind of
cancellation between `ln(beta)` and `ln(1+beta)` when beta is large.
Both reformulations are algebraically exact (zero added
approximation error), not truncated series -- see
tests/test_finite_mc_likelihood.py's high-alpha regression tests,
which fail against the naive formula and pass against this one.
Note on the `sigma2 <= 1e-10` Poisson fallback below: this is not a
numerical workaround for the instability above -- it is the
mathematically correct limit. As sigma^2 -> 0, alpha -> infinity
and this Negative Binomial collapses onto Poisson(mu) (zero
simulation variance is exactly the "infinite MC" assumption the
plain-Poisson branch above already encodes), so evaluating the
finite-MC formula at sigma^2 = 0 exactly would only divide by zero,
not disagree with the limit. The two branches use different
additive normalizations (this branch drops only ln(n!); the
use_finite_mc=False branch is the fully saturated Baker-Cousins
form), so their raw NLL values differ by a data-only constant even
in this limit -- only t = NLL_cond - NLL_uncond is meaningful to
compare, and that constant cancels there regardless of which branch
is taken, as long as both the conditional and unconditional fit for
a given call use the same use_finite_mc setting (they always do).
Parameters:
-----------
params : array_like, 1D
The physical model parameters to evaluate.
N_obs : array_like, any shape
The observed data counts (or toy data) in each bin. May be an
arbitrary N-dimensional histogram (e.g. shape (E_bins, cos_theta_bins));
it is flattened internally, so the NLL is a sum over all bins
regardless of how they are laid out spatially. Non-contiguous views
(e.g. a transposed histogram) are handled by copying to a contiguous
layout before flattening, since numba's nopython `.reshape(-1)` only
accepts C-contiguous input.
S_sumw2 : array_like, same shape as N_obs
Sum of squared MC weights (sum(w_i^2)) per bin for the signal
template -- the standard per-bin variance estimator for a weighted
MC sample (unweighted MC is the w_i=1 special case).
B_sumw2 : array_like, same shape as N_obs
Sum of squared MC weights per bin for the background template.
use_finite_mc : bool
If True, applies the Poisson-Gamma mixture likelihood. If False, uses standard Poisson.
compute_rates_func : callable
User-provided mapping function `(params, S_sumw2, B_sumw2) -> (mu, sigma2)`
that calculates bin expectations (mu) and variances (sigma2), of the same
shape as N_obs. Any fixed template arrays it needs (signal/background
shapes) should be referenced via closure or module-global rather than
passed in as arguments.
Returns:
--------
nll : float
The calculated negative log-likelihood value. For bins where the model
yields an unphysical or near-zero expected count mu_i (<= 1e-12) while
data was observed there (n_obs > 0), a smooth quadratic barrier is
added to that bin's contribution instead of returning a flat constant
(see note below).
Note on unphysical (mu_i <= mu_floor = 1e-12) bins:
----------------------------------------------------
Earlier versions of this function returned a flat penalty (1e10) the instant
any bin's expected count went non-positive, discarding whatever NLL had
already been accumulated from other bins in this call. Because
`scipy.optimize.minimize(..., method='L-BFGS-B')` estimates gradients via
finite differences, a locally-constant NLL surface reads as "gradient ~ 0,
already converged," which can permanently trap the optimizer if its starting
guess lands in a mu <= 0 region (e.g. the bounds midpoint, for a user model
that is unphysical there). Instead, the contribution from an unphysical bin
is now a continuous extension of the real Poisson NLL term (evaluated at a
tiny positive floor) plus a quadratic barrier that grows with how far mu_i
actually is below that floor, so the gradient w.r.t. mu_i stays informative
and points back toward mu_i > 0. This bin-local contribution is added to the
running total rather than short-circuiting the whole function. The trigger
is `mu_i <= mu_floor` rather than `mu_i <= 0` specifically so there is no
razor-thin discontinuity in the tiny sliver `(0, mu_floor)`: the real
Poisson term diverges as mu_i -> 0+, while the floor-based term is constant
(evaluated at mu_floor, not at the actual mu_i), so triggering only at
mu_i <= 0 would make the NLL drop right at that boundary instead of
continuing to climb. For all mu_i > mu_floor (i.e. all physically
meaningful values -- mu_floor is astronomically small), behavior is
practically identical to before this change (matches to within 1e-12
relative/absolute tolerance; see tests/test_smoothing.py).
"""
nll = 0.0
mu, sigma2_arr = compute_rates_func(params, S_sumw2, B_sumw2)
mu_flat = np.ascontiguousarray(mu).reshape(-1)
N_obs_flat = np.ascontiguousarray(N_obs).reshape(-1)
sigma2_flat = np.ascontiguousarray(sigma2_arr).reshape(-1)
if mu_flat.shape[0] != N_obs_flat.shape[0]:
raise ValueError("compute_rates_func's returned `mu` does not have the same number of bins as `N_obs`.")
for i in range(len(N_obs_flat)):
mu_i = mu_flat[i]
n_obs = float(N_obs_flat[i])
mu_floor = 1e-12
# NaN mu_i is not "unphysical" in the sense the barrier below
# handles -- `mu_i <= mu_floor` is False for NaN, so without this
# check NaN would silently fall through into math.log(n_obs / mu_i)
# below, producing a NaN nll with no barrier, no warning, and no
# gradient direction for the optimizer to recover from (a quadratic
# barrier needs a finite distance-from-floor to be meaningful; NaN
# has none). A NaN mu_i almost always means a bug in the user's
# compute_rates_func (e.g. division by zero, sqrt/log of a negative
# number), so fail loudly here instead of corrupting the fit silently.
#
# NaN detection here deliberately avoids math.isnan(mu_i)/np.isnan(mu_i)/
# (mu_i != mu_i): this function (and every caller in its nopython
# call chain, e.g. unconditional_fit_grid) is @njit(fastmath=True),
# and numba's fastmath compiles under LLVM's "assume no NaN/Inf"
# flags, which makes ALL THREE of those standard NaN checks silently
# return False for a genuine NaN under fastmath -- verified directly
# before writing this. Reinterpreting the float64's raw bits as an
# int64 and checking the IEEE-754 NaN bit pattern (exponent all 1s,
# mantissa nonzero) uses only integer bitwise ops, which fastmath's
# floating-point assumptions do not touch.
mu_i_bits = np.array([mu_i], dtype=np.float64).view(np.int64)[0]
mu_i_is_nan = ((mu_i_bits >> 52) & 0x7FF) == 0x7FF and (mu_i_bits & 0xFFFFFFFFFFFFF) != 0
if mu_i_is_nan:
raise ValueError("compute_rates_func returned a NaN expected count (mu). Check compute_rates_func for division by zero, sqrt/log of a negative number, or similar.")
# Triggered for mu_i <= mu_floor, not just mu_i <= 0: the physical
# branch's own Poisson term below diverges to +inf as mu_i -> 0+, while
# this branch's base term is constant (evaluated at the fixed
# mu_floor, not at the actual mu_i). If the trigger were mu_i <= 0,
# the tiny physical sliver (0, mu_floor) would still take the
# diverging physical branch, so crossing from mu_i = +epsilon to
# mu_i = -epsilon would make the NLL drop rather than keep climbing.
# Extending the trigger to include that sliver removes the
# discontinuity: at mu_i == mu_floor exactly, both formulations
# already agree (barrier is 0 there by construction).
if mu_i <= mu_floor:
if n_obs > 0:
# Continuous extension of the Poisson NLL term below mu=0, using a
# tiny positive floor, plus a quadratic barrier scaled by how far
# mu_i actually is below that floor (keeps the gradient informative
# instead of flat, unlike the old constant 1e10 penalty).
barrier = 1e6 * (mu_floor - mu_i)**2
nll += 2.0 * (mu_floor - n_obs + n_obs * math.log(n_obs / mu_floor)) + barrier
# else: n_obs == 0 and mu_i <= mu_floor has no real statistical
# content here; behave like a no-op contribution (matches the
# prior `continue`). Note: for n_obs == 0 and mu_i in the sliver
# (0, mu_floor), the physical branch below would instead have
# contributed 2.0 * mu_i (at most ~2e-12) -- utterly negligible,
# and the right trade for keeping the n_obs > 0 case continuous.
continue
if use_finite_mc:
sigma2 = sigma2_flat[i]
if sigma2 > 1e-10:
alpha = (mu_i**2) / sigma2 + 1.0
beta = max(mu_i / sigma2, 1e-300)
# Numerically stable, algebraically exact reformulation --
# see calc_nll's own docstring "Numerically stable
# evaluation" note for the derivation and why the naive
# lgamma-difference form above loses catastrophic precision
# once alpha (~ mu^2/sigma^2) gets large.
n_int = int(round(n_obs))
lgamma_diff = 0.0
for k in range(n_int):
lgamma_diff += math.log(alpha + k)
lnL = -alpha * math.log1p(1.0 / beta) - n_obs * math.log1p(beta) + lgamma_diff
nll += -2.0 * lnL
else:
lnL_poisson = n_obs * math.log(mu_i) - mu_i
nll += -2.0 * lnL_poisson
else:
if n_obs > 0:
nll += 2.0 * (mu_i - n_obs + n_obs * math.log(n_obs / mu_i))
else:
nll += 2.0 * mu_i
return nll
# --- 3. Grid Search Optimizers (Binned) ---
[docs]
@njit(fastmath=True, nogil=True)
def unconditional_fit_grid(N_obs, full_grid_points, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func):
"""
Performs an unconditional maximum likelihood fit via exhaustive grid search.
Statistical Theory:
To construct the denominator of the Profile Likelihood Ratio (PLR) test statistic,
we must find the global maximum of the likelihood function (minimum of the NLL)
allowing all parameters in the model to vary freely across their physical bounds.
Expression: NLL_best = min( NLL(theta) ) for all theta in parameter space.
Parameters:
-----------
N_obs : array_like, any shape
The observed data or toy counts.
full_grid_points : array_like, 2D
A pre-computed matrix where each row represents a complete N-dimensional
parameter combination to evaluate.
S_sumw2, B_sumw2 : array_like, same shape as N_obs
Sum of squared MC weights (sumw2) per bin, for signal and background respectively.
use_finite_mc : bool
Flag to use finite MC likelihood formulation.
compute_rates_func : callable
User-provided physics mapping function.
Returns:
--------
min_nll : float
The absolute minimum negative log-likelihood found in the grid.
best_params : array_like, 1D
The parameter combination that produced min_nll.
"""
min_nll = 1e10
best_params = full_grid_points[0].copy()
for row in range(len(full_grid_points)):
p = full_grid_points[row]
nll = calc_nll(p, N_obs, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func)
if nll < min_nll:
min_nll = nll
best_params = p.copy()
return min_nll, best_params
[docs]
@njit(fastmath=True, nogil=True)
def conditional_fit_grid_1d(test_val, fix_idx, n_params, N_obs, cond_grid_points, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func):
"""
Performs a conditional maximum likelihood fit with one parameter fixed (Profiling).
Statistical Theory:
To construct the numerator of the Profile Likelihood Ratio for a 1D confidence
interval, the likelihood is maximized (NLL minimized) while keeping the parameter
of interest fixed to a specific test value. The remaining (nuisance) parameters
are allowed to vary (profiled out).
Expression: NLL_cond = min( NLL(theta_test, nu) ) over all nuisance parameters nu.
Parameters:
-----------
test_val : float
The fixed value for the parameter of interest.
fix_idx : int
The index of the parameter to fix.
n_params : int
Total number of parameters in the model.
N_obs, S_sumw2, B_sumw2 : array_like, any shape (all matching)
Data counts and variances.
cond_grid_points : array_like, 2D
Pre-computed grid of the (N-1) free nuisance parameters to scan over.
use_finite_mc : bool
Flag to use finite MC likelihood formulation.
compute_rates_func : callable
User-provided physics mapping function.
Returns:
--------
min_nll : float
The minimum negative log-likelihood given the fixed parameter.
best_params : array_like, 1D
The full parameter combination (including the fixed one) that produced min_nll.
"""
min_nll = 1e10
best_params = np.zeros(n_params)
p = np.zeros(n_params)
p[fix_idx] = test_val
for row in range(len(cond_grid_points)):
free_p = cond_grid_points[row]
free_i = 0
for i in range(n_params):
if i != fix_idx:
p[i] = free_p[free_i]
free_i += 1
nll = calc_nll(p, N_obs, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func)
if nll < min_nll:
min_nll = nll
best_params = p.copy()
return min_nll, best_params
[docs]
@njit(fastmath=True, nogil=True)
def conditional_fit_grid_2d(test_vA, test_vB, fix_A, fix_B, n_params, N_obs, cond_grid_points, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func):
"""
Performs a conditional maximum likelihood fit with two parameters fixed.
Statistical Theory:
Analogous to the 1D case, this computes the profiled likelihood for a 2D joint
confidence region. Two parameters of interest are fixed to a grid point, and
the NLL is minimized over all remaining (N-2) nuisance parameters.
Parameters:
-----------
test_vA, test_vB : float
The fixed values for the two parameters of interest.
fix_A, fix_B : int
The indices of the parameters being fixed.
n_params : int
Total number of parameters in the model.
N_obs, S_sumw2, B_sumw2 : array_like, any shape (all matching)
Data counts and variances.
cond_grid_points : array_like, 2D
Pre-computed grid of the (N-2) free nuisance parameters to scan over.
use_finite_mc : bool
Flag to use finite MC likelihood formulation.
compute_rates_func : callable
User-provided physics mapping function.
Returns:
--------
min_nll : float
The minimum negative log-likelihood given the two fixed parameters.
best_params : array_like, 1D
The full parameter combination that produced min_nll.
"""
min_nll = 1e10
best_params = np.zeros(n_params)
p = np.zeros(n_params)
p[fix_A] = test_vA
p[fix_B] = test_vB
for row in range(len(cond_grid_points)):
free_p = cond_grid_points[row]
free_i = 0
for i in range(n_params):
if i != fix_A and i != fix_B:
p[i] = free_p[free_i]
free_i += 1
nll = calc_nll(p, N_obs, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func)
if nll < min_nll:
min_nll = nll
best_params = p.copy()
return min_nll, best_params
# --- 4. Toy Generators (Binned) ---
[docs]
@njit(fastmath=True, parallel=True, nogil=True)
def generate_and_fit_toys_grid_1d(test_val, fix_idx, true_params, n_params,
full_grid_points, cond_grid_points, n_toys, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func):
"""
Generates Poisson toys and computes the 1D test statistic distribution (Feldman-Cousins).
Statistical Theory & Assumptions:
To ensure proper frequentist coverage (as per the Feldman-Cousins unified approach),
we cannot assume the Profile Likelihood Ratio strictly follows an asymptotic chi-square
distribution (Wilks' Theorem), especially near physical boundaries.
Instead, we build the empirical distribution of the test statistic (t) by running
parametric bootstraps (toys).
For each toy:
1. A synthetic dataset (toy_N) is sampled from Poisson(mu_true).
2. The unconditional global minimum NLL is found for the toy.
3. The conditional minimum NLL (fixing the parameter of interest) is found.
4. The test statistic is calculated: t = NLL_cond - NLL_uncond.
(Due to numerical imprecision in grid searches, t is bounded at a minimum of 0.0).
Parameters:
-----------
test_val : float
The fixed value for the parameter of interest.
fix_idx : int
The index of the fixed parameter.
true_params : array_like, 1D
The parameter combination used to generate the true expected counts for the toys.
n_params : int
Total number of parameters.
S_sumw2, B_sumw2 : array_like, any shape
Variances (matching whatever shape compute_rates_func's `mu` produces).
full_grid_points : array_like, 2D
Grid for the unconditional fit.
cond_grid_points : array_like, 2D
Grid for the conditional (nuisance parameter) fit.
n_toys : int
The number of pseudo-experiments to generate.
use_finite_mc : bool
Flag to use finite MC likelihood during toy fitting.
compute_rates_func : callable
User-provided physics mapping function.
Returns:
--------
t_statistics : array_like, 1D
Array of length `n_toys` containing the calculated test statistic for each toy.
"""
t_statistics = np.zeros(n_toys)
mu_true, _ = compute_rates_func(true_params, S_sumw2, B_sumw2)
mu_true_flat = np.ascontiguousarray(mu_true).reshape(-1)
n_bins = mu_true_flat.shape[0]
for t in prange(n_toys):
toy_N = np.zeros(n_bins)
for i in range(n_bins):
toy_N[i] = np.random.poisson(mu_true_flat[i])
uncond_nll, _ = unconditional_fit_grid(toy_N, full_grid_points, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func)
cond_nll, _ = conditional_fit_grid_1d(test_val, fix_idx, n_params, toy_N, cond_grid_points, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func)
t_statistics[t] = max(0.0, cond_nll - uncond_nll)
return t_statistics
[docs]
@njit(fastmath=True, parallel=True, nogil=True)
def generate_and_fit_toys_grid_2d(test_vA, test_vB, fix_A, fix_B, true_params, n_params,
full_grid_points, cond_grid_points, n_toys, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func):
"""
Generates Poisson toys and computes the 2D test statistic distribution.
Statistical Theory:
Similar to the 1D case, this function calculates empirical distributions of the
Profile Likelihood Ratio test statistic for 2D parameter joint combinations.
t = NLL_cond_2D - NLL_uncond.
Parameters:
-----------
test_vA, test_vB : float
The fixed values for the two parameters of interest.
fix_A, fix_B : int
The indices of the parameters being fixed.
true_params : array_like, 1D
The parameter combination used to generate expected counts for the toys.
n_params : int
Total number of parameters.
S_sumw2, B_sumw2 : array_like, any shape
Variances (matching whatever shape compute_rates_func's `mu` produces).
full_grid_points : array_like, 2D
Grid for the unconditional fit.
cond_grid_points : array_like, 2D
Grid for the conditional (nuisance parameter) fit.
n_toys : int
The number of pseudo-experiments to generate.
use_finite_mc : bool
Flag to use finite MC likelihood during toy fitting.
compute_rates_func : callable
User-provided physics mapping function.
Returns:
--------
t_statistics : array_like, 1D
Array of length `n_toys` containing the calculated test statistic for each toy.
"""
t_statistics = np.zeros(n_toys)
mu_true, _ = compute_rates_func(true_params, S_sumw2, B_sumw2)
mu_true_flat = np.ascontiguousarray(mu_true).reshape(-1)
n_bins = mu_true_flat.shape[0]
for t in prange(n_toys):
toy_N = np.zeros(n_bins)
for i in range(n_bins):
toy_N[i] = np.random.poisson(mu_true_flat[i])
uncond_nll, _ = unconditional_fit_grid(toy_N, full_grid_points, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func)
cond_nll, _ = conditional_fit_grid_2d(test_vA, test_vB, fix_A, fix_B, n_params, toy_N, cond_grid_points, S_sumw2, B_sumw2, use_finite_mc, compute_rates_func)
t_statistics[t] = max(0.0, cond_nll - uncond_nll)
return t_statistics