Source code for bayspec.infer.statistic

"""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