"""Likelihood-statistic kernels for spectral fitting.
Implements the Gaussian (``Gstat``), Poisson (``Pstat``), Poisson-source
plus Poisson-background (``PPstat``/``cstat``), Poisson-source plus
Gaussian-background (``PGstat``) statistics, and their upper-limit
variants. ``StatisticNB`` exposes the numba-accelerated fast path (fused
stat + signed residual per bin); ``Statistic`` is a pure-numpy fallback
with the same interface.
Every statistic takes the same keyword arguments (``S``, ``B``, ``m``,
``ts``, ``tb``, ``sigma_S``, ``sigma_B``) and returns ``(stat,
residual)`` where ``residual`` is the signed square-root of the per-bin
statistic contribution.
"""
import numba as nb
import numpy as np
from ..util.significance import pgsig, ppsig
@nb.njit(cache=True, fastmath=True)
def _gstat_core(S, B, m, ts, tb, sigma_S, sigma_B):
"""Numba kernel for ``Gstat``; returns ``(total_stat, signed_residual)``."""
n = S.shape[0]
residual = np.empty(n, dtype=np.float64)
stat = 0.0
ratio = 0.0
if tb != 0.0:
ratio = ts / tb
for i in range(n):
bi = B[i]
sigma_bi = sigma_B[i]
if tb != 0.0:
bi = bi * ratio
sigma_bi = sigma_bi * ratio
sigma = np.sqrt(sigma_S[i] * sigma_S[i] + sigma_bi * sigma_bi)
di = S[i] - bi
mi = m[i] * ts
z = (di - mi) / sigma
logli = -0.5 * z * z
stati = -2.0 * logli
stat += stati
sign = 0.0
delta = di - mi
if delta > 0.0:
sign = 1.0
elif delta < 0.0:
sign = -1.0
residual[i] = sign * np.sqrt(stati)
return stat, residual
@nb.njit(cache=True, fastmath=True)
def _pstat_core(S, m, ts):
"""Numba kernel for the pure-Poisson ``Pstat``; background is absent."""
n = S.shape[0]
residual = np.empty(n, dtype=np.float64)
stat = 0.0
for i in range(n):
si = S[i]
mu = m[i] * ts
klogmu = 0.0
if si != 0.0:
klogmu = si * np.log(mu)
klogk = 0.0
if si != 0.0:
klogk = si * np.log(si)
logli = klogmu - mu - klogk + si
stati = -2.0 * logli
stat += stati
delta = si - mu
sign = 0.0
if delta > 0.0:
sign = 1.0
elif delta < 0.0:
sign = -1.0
residual[i] = sign * np.sqrt(stati)
return stat, residual
@nb.njit(cache=True, fastmath=True)
def _ppstat_core(S, B, m, ts, tb):
"""Numba kernel for ``PPstat``/``cstat``: Poisson source + Poisson background.
Profiles out the true background rate ``b`` per bin via the standard
``cstat`` closed-form, then sums the profiled Poisson log-likelihoods.
"""
n = S.shape[0]
residual = np.empty(n, dtype=np.float64)
stat = 0.0
aa = ts + tb
for i in range(n):
si = S[i]
bi = B[i]
mi = m[i]
bb = aa * mi - si - bi
cc = -bi * mi
dd = np.sqrt(bb * bb - 4.0 * aa * cc)
b = -2.0 * cc / (bb + dd) if bb >= 0.0 else -(bb - dd) / (2.0 * aa)
mu_s = ts * (b + mi)
mu_b = tb * b
s_klogmu = 0.0
if si != 0.0:
s_klogmu = si * np.log(mu_s)
s_klogk = 0.0
if si != 0.0:
s_klogk = si * np.log(si)
b_klogmu = 0.0
if bi != 0.0:
b_klogmu = bi * np.log(mu_b)
b_klogk = 0.0
if bi != 0.0:
b_klogk = bi * np.log(bi)
logli = (s_klogmu - mu_s - s_klogk + si) + (b_klogmu - mu_b - b_klogk + bi)
stati = -2.0 * logli
stat += stati
delta = si / ts - bi / tb - mi
sign = 0.0
if delta > 0.0:
sign = 1.0
elif delta < 0.0:
sign = -1.0
residual[i] = sign * np.sqrt(stati)
return stat, residual
@nb.njit(cache=True, fastmath=True)
def _pgstat_core(S, B, m, ts, tb, sigma_B):
"""Numba kernel for ``PGstat``: Poisson source + Gaussian background.
Profiles the background rate via the quadratic closed-form used in
XSPEC's ``pgstat`` and combines the profiled Poisson source and
Gaussian background log-likelihoods.
"""
n = S.shape[0]
residual = np.empty(n, dtype=np.float64)
stat = 0.0
aa = tb * tb
for i in range(n):
si = S[i]
bi = B[i]
mi = m[i]
sigma = sigma_B[i]
bb = ts * sigma * sigma - tb * bi + tb * tb * mi
cc = ts * sigma * sigma * mi - si * sigma * sigma - tb * bi * mi
dd = np.sqrt(bb * bb - 4.0 * aa * cc)
sgn = 1.0
if bb < 0.0:
sgn = -1.0
qq = -0.5 * (bb + sgn * dd)
b1 = qq / aa
b2 = cc / qq
b = b1 if b1 > 0.0 else b2
mu_s = ts * (b + mi)
s_klogmu = 0.0
if si != 0.0:
s_klogmu = si * np.log(mu_s)
s_klogk = 0.0
if si != 0.0:
s_klogk = si * np.log(si)
pois_logli = s_klogmu - mu_s - s_klogk + si
z = (bi - tb * b) / sigma
gauss_logli = -0.5 * z * z
logli = pois_logli + gauss_logli
stati = -2.0 * logli
stat += stati
delta = si / ts - bi / tb - mi
sign = 0.0
if delta > 0.0:
sign = 1.0
elif delta < 0.0:
sign = -1.0
residual[i] = sign * np.sqrt(stati)
return stat, residual
[docs]
class StatisticNB:
"""Numba-accelerated statistic dispatch table.
Every method takes the standard keyword bundle (``S``, ``B``, ``m``,
``ts``, ``tb``, ``sigma_S``, ``sigma_B``) and returns ``(stat,
residual)``. Upper-limit variants compare a significance against a
target (default 3σ) so that a minimiser can search for the
model normalization that produces the requested sigma.
"""
[docs]
@staticmethod
def Gstat(**kwargs):
"""Gaussian statistic; delegates to :func:`_gstat_core`."""
return _gstat_core(
kwargs['S'],
kwargs['B'],
kwargs['m'],
kwargs['ts'],
kwargs['tb'],
kwargs['sigma_S'],
kwargs['sigma_B'],
)
[docs]
@staticmethod
def Pstat(**kwargs):
"""Pure-Poisson statistic (no background); delegates to :func:`_pstat_core`."""
return _pstat_core(kwargs['S'], kwargs['m'], kwargs['ts'])
[docs]
@staticmethod
def PPstat(**kwargs):
"""Profile Poisson-Poisson statistic (``cstat``); uses :func:`_ppstat_core`."""
return _ppstat_core(kwargs['S'], kwargs['B'], kwargs['m'], kwargs['ts'], kwargs['tb'])
[docs]
@staticmethod
def PGstat(**kwargs):
"""Profile Poisson-Gaussian statistic; uses :func:`_pgstat_core`."""
return _pgstat_core(
kwargs['S'], kwargs['B'], kwargs['m'], kwargs['ts'], kwargs['tb'], kwargs['sigma_B']
)
[docs]
@staticmethod
def PPstat_UL(**kwargs):
"""Upper-limit driver for ``PPstat``: minimize ``(sigma - 3)^2``."""
B = kwargs['B']
m = kwargs['m']
ts = kwargs['ts']
tb = kwargs['tb']
alpha = ts / tb
bkg_cts = np.sum(B)
mo_cts = np.sum(m * ts)
ul_sigma = 3.0
sigma = ppsig(mo_cts + bkg_cts * alpha, bkg_cts, alpha)
stat = (sigma - ul_sigma) ** 2
residual = np.array([sigma - ul_sigma], dtype=np.float64)
return stat, residual
[docs]
@staticmethod
def PGstat_UL(**kwargs):
"""Upper-limit driver for ``PGstat``: minimize ``(sigma - 3)^2``."""
B = kwargs['B']
m = kwargs['m']
ts = kwargs['ts']
tb = kwargs['tb']
alpha = ts / tb
sigma_B = kwargs['sigma_B']
bkg_cts = np.sum(B)
mo_cts = np.sum(m * ts)
bkg_err = np.sqrt(np.sum(sigma_B * sigma_B))
ul_sigma = 3.0
sigma = pgsig(mo_cts + bkg_cts * alpha, bkg_cts * alpha, bkg_err * alpha)
stat = (sigma - ul_sigma) ** 2
residual = np.array([sigma - ul_sigma], dtype=np.float64)
return stat, residual
[docs]
class Statistic:
"""Pure-numpy fallback mirror of :class:`StatisticNB`.
Same ``(stat, residual)`` contract as the numba-accelerated class,
intended for debugging and for environments where numba cannot be
used. Every method returns the total statistic plus the signed
per-bin residual.
"""
[docs]
@staticmethod
def xlogy(x, y):
"""Return ``x * log(y)`` element-wise, treating ``0 * log(y)`` as 0 for any ``y``."""
res = np.zeros_like(x, dtype=np.float64)
zero = x == 0
res[~zero] = x[~zero] * np.log(y[~zero])
return res
[docs]
@staticmethod
def xdivy(x, y):
"""Return ``x / y`` element-wise, treating ``0 / 0`` as 0."""
res = np.zeros_like(x, dtype=np.float64)
zero = (x == 0) & (y == 0)
res[~zero] = x[~zero] / y[~zero]
return res
[docs]
@staticmethod
def poisson_logpmf(k, mu):
"""Poisson log-PMF with Stirling's approximation for ``log(k!)``.
Drops the ``log(2πk) / 2`` term, so the result is accurate up to
an additive constant that cancels in likelihood ratios.
"""
return Statistic.xlogy(k, mu) - mu - Statistic.xlogy(k, k) + k
[docs]
@staticmethod
def gaussian_logpdf(x, loc, scale):
"""Gaussian log-PDF dropping the ``log(2π σ²)`` normalization constant."""
return -0.5 * (Statistic.xdivy(x - loc, scale) ** 2)
[docs]
@staticmethod
def Gstat(**kwargs):
"""Gaussian source-background statistic (``Gstat``/``chi2``)."""
S = kwargs['S']
B = kwargs['B']
m = kwargs['m']
ts = kwargs['ts']
tb = kwargs['tb']
sigma_S = kwargs['sigma_S']
sigma_B = kwargs['sigma_B']
if tb != 0:
B = B / tb * ts
sigma_B = sigma_B / tb * ts
sigma = np.sqrt(sigma_S**2 + sigma_B**2)
sign = np.sign(S - B - m * ts)
loglike = Statistic.gaussian_logpdf(S - B, m * ts, sigma)
stat = (-2 * loglike).sum()
residual = sign * np.sqrt(-2 * loglike)
return stat, residual
[docs]
@staticmethod
def Pstat(**kwargs):
"""Pure-Poisson statistic when the background is ignored."""
S = kwargs['S']
m = kwargs['m']
ts = kwargs['ts']
sign = np.sign(S - m * ts)
loglike = Statistic.poisson_logpmf(S, m * ts)
stat = (-2 * loglike).sum()
residual = sign * np.sqrt(-2 * loglike)
return stat, residual
[docs]
@staticmethod
def PPstat(**kwargs):
"""Profile Poisson-Poisson statistic (``cstat`` equivalent)."""
S = kwargs['S']
B = kwargs['B']
m = kwargs['m']
ts = kwargs['ts']
tb = kwargs['tb']
aa = ts + tb
bb = (ts + tb) * m - S - B
cc = -B * m
dd = np.sqrt(bb * bb - 4 * aa * cc)
po = bb >= 0
b = np.empty_like(B, dtype=np.float64)
b[po] = -2 * cc[po] / (bb[po] + dd[po])
b[~po] = -(bb[~po] - dd[~po]) / (2 * aa)
sign = np.sign(S / ts - B / tb - m)
loglike = Statistic.poisson_logpmf(S, ts * (b + m)) + Statistic.poisson_logpmf(B, tb * b)
stat = (-2 * loglike).sum()
residual = sign * np.sqrt(-2 * loglike)
return stat, residual
[docs]
def PGstat(**kwargs):
"""Profile Poisson-Gaussian statistic (XSPEC-style ``pgstat``)."""
S = kwargs['S']
B = kwargs['B']
m = kwargs['m']
ts = kwargs['ts']
tb = kwargs['tb']
sigma = kwargs['sigma_B']
aa = tb**2
bb = ts * sigma**2 - tb * B + tb**2 * m
cc = ts * sigma**2 * m - S * sigma**2 - tb * B * m
dd = np.sqrt(bb**2 - 4 * aa * cc)
sign = np.where(bb >= 0, 1, -1)
qq = -0.5 * (bb + sign * dd)
b1 = qq / aa
b2 = cc / qq
b = np.where(b1 > 0, b1, b2)
sign = np.sign(S / ts - B / tb - m)
loglike = Statistic.poisson_logpmf(S, ts * (b + m)) + Statistic.gaussian_logpdf(
B, tb * b, sigma
)
stat = (-2 * loglike).sum()
residual = sign * np.sqrt(-2 * loglike)
return stat, residual
[docs]
@staticmethod
def PPstat_UL(**kwargs):
"""Upper-limit driver for ``PPstat``: minimize ``(sigma - 3)^2``."""
B = kwargs['B']
m = kwargs['m']
ts = kwargs['ts']
tb = kwargs['tb']
alpha = ts / tb
bkg_cts = np.sum(B)
mo_cts = np.sum(m * ts)
ul_sigma = 3.0
sigma = ppsig(mo_cts + bkg_cts * alpha, bkg_cts, alpha)
stat = (sigma - ul_sigma) ** 2
residual = np.array([sigma - ul_sigma], dtype=np.float64)
return stat, residual
[docs]
@staticmethod
def PGstat_UL(**kwargs):
"""Upper-limit driver for ``PGstat``: minimize ``(sigma - 3)^2``."""
B = kwargs['B']
m = kwargs['m']
ts = kwargs['ts']
tb = kwargs['tb']
alpha = ts / tb
sigma_B = kwargs['sigma_B']
bkg_cts = np.sum(B)
mo_cts = np.sum(m * ts)
bkg_err = np.sqrt(np.sum(sigma_B * sigma_B))
ul_sigma = 3.0
sigma = pgsig(mo_cts + bkg_cts * alpha, bkg_cts * alpha, bkg_err * alpha)
stat = (sigma - ul_sigma) ** 2
residual = np.array([sigma - ul_sigma], dtype=np.float64)
return stat, residual