"""Beta-binomial distribution -- an overdispersed binomial (Beta-binomial compound).
If the success probability of a binomial is itself Beta(a, b) distributed and integrated out, the
count ``k`` in ``n`` trials follows the beta-binomial:
P(k; n, a, b) = C(n, k) B(k + a, n - k + b) / B(a, b), k = 0, ..., n,
with mean ``n a/(a+b)`` and a variance inflated over the binomial by the intra-class correlation
``rho = 1/(a+b+1)``. It is the standard model for overdispersed bounded counts (clustered trials,
batch defect rates). The number of trials ``n`` is a fixed, known parameter; ``a`` and ``b`` are
estimated by moments -- ``rho`` from the dispersion and the mean fixing ``a/(a+b)``.
"""
import math
from collections.abc import Sequence
from typing import Any
import numpy as np
from numpy.random import RandomState
from scipy.special import betaln, gammaln
from mixle.stats.compute.pdist import (
DataSequenceEncoder,
DistributionEnumerator,
DistributionSampler,
ParameterEstimator,
SequenceEncodableProbabilityDistribution,
SequenceEncodableStatisticAccumulator,
StatisticAccumulatorFactory,
)
from mixle.stats.univariate.discrete._count_contracts import exact_integer_observations, nonnegative_weights
[docs]
class BetaBinomialDistribution(SequenceEncodableProbabilityDistribution):
"""Beta-binomial distribution over ``{0, ..., n}`` with shape parameters ``a, b > 0``."""
def __init__(self, n: int, a: float, b: float, name: str | None = None, keys: str | None = None) -> None:
# n must land exactly on a non-negative integer -- checked on the float value before any
# int() cast, since int(10.5) truncates to 10 and int(nan)/int(inf) raise the wrong
# exception type (ValueError/OverflowError instead of this constructor's own ValueError).
n_val = float(n)
if (
not np.isfinite(n_val)
or n_val != math.floor(n_val)
or n_val < 0
or a <= 0.0
or b <= 0.0
or not (np.isfinite(a) and np.isfinite(b))
):
raise ValueError("BetaBinomialDistribution requires n to be a non-negative integer and a, b > 0.")
self.n = int(n_val)
self.a = float(a)
self.b = float(b)
self.name = name
self.keys = keys
self._log_beta_ab = betaln(self.a, self.b)
def __setattr__(self, name: str, value: Any) -> None:
"""Keep ``_log_beta_ab`` tied to the parameter(s) ``a``/``b`` they derive from.
Computed once in ``__init__`` and read by ``log_density``, so a later assignment used to
leave them stale and the scorer kept reporting the *previous* parameters' density with no
error at all (MXR-080-1192).
Recompute rather than validate: callers legitimately install out-of-domain or non-finite
parameters -- deserialized legacy states and NaN-propagation checks both do -- so a value
outside the domain yields a NaN constant that propagates honestly instead of rejecting a
state the library is expected to be able to hold.
"""
object.__setattr__(self, name, value)
if name not in ("a", "b"):
return
try:
object.__setattr__(self, "_log_beta_ab", float(betaln(self.a, self.b)))
except (ValueError, TypeError, OverflowError, ZeroDivisionError, AttributeError, FloatingPointError):
# AttributeError covers __init__, where the first parameter is assigned before the rest.
object.__setattr__(self, "_log_beta_ab", float("nan"))
def __str__(self) -> str:
return "BetaBinomialDistribution(%s, %s, %s, name=%s, keys=%s)" % (
repr(self.n),
repr(self.a),
repr(self.b),
repr(self.name),
repr(self.keys),
)
[docs]
def density(self, x: int) -> float:
"""Return the probability mass at a single count ``x``."""
return math.exp(self.log_density(x))
[docs]
def log_density(self, x: int) -> float:
"""Return the log probability mass at ``x`` (``-inf`` outside ``{0, ..., n}`` or non-integer --
a fractional count is not a valid beta-binomial outcome)."""
try:
xx = float(x)
except Exception: # noqa: BLE001
return -np.inf
if not np.isfinite(xx) or np.floor(xx) != xx or xx < 0 or xx > self.n:
return -np.inf
k = int(xx)
log_choose = gammaln(self.n + 1) - gammaln(k + 1) - gammaln(self.n - k + 1)
return float(log_choose + betaln(k + self.a, self.n - k + self.b) - self._log_beta_ab)
[docs]
def seq_log_density(self, x: np.ndarray) -> np.ndarray:
"""Return vectorized log-mass for an array of counts (``-inf`` outside ``{0, ..., n}`` or
non-integer -- a fractional count is not a valid beta-binomial outcome)."""
k = np.asarray(x, dtype=np.float64)
good = np.isfinite(k) & (np.floor(k) == k) & (k >= 0) & (k <= self.n)
log_choose = gammaln(self.n + 1) - gammaln(k + 1) - gammaln(self.n - k + 1)
rv = log_choose + betaln(k + self.a, self.n - k + self.b) - self._log_beta_ab
return np.where(good, rv, -np.inf)
[docs]
def support_size(self) -> int:
"""``n + 1`` outcomes ``{0, ..., n}``."""
return int(self.n) + 1
def _log_cdf_table(self) -> np.ndarray:
"""Log cumulative masses ``log P(X <= k)`` for k = 0..n by stable log-space accumulation."""
return np.logaddexp.accumulate(self.seq_log_density(np.arange(self.n + 1, dtype=np.float64)))
[docs]
def cdf(self, x: float) -> float:
"""Cumulative distribution function P(X <= x): a log-space partial sum of the pmf over {0..n}."""
k = math.floor(float(x))
if k < 0:
return 0.0
if k >= self.n:
return 1.0
return float(min(np.exp(self._log_cdf_table()[k]), 1.0))
[docs]
def quantile(self, q: float) -> float:
"""Inverse CDF F^{-1}(q) over {0..n}: the smallest count whose cdf reaches ``q`` (search on the cdf)."""
qq = float(q)
if qq <= 0.0:
return 0.0
cum = np.exp(self._log_cdf_table())
return float(min(int(np.searchsorted(cum, qq, side="left")), self.n))
[docs]
def enumerator(self) -> "BetaBinomialEnumerator":
"""Returns BetaBinomialEnumerator iterating the support in descending probability order."""
return BetaBinomialEnumerator(self)
[docs]
def sampler(self, seed: int | None = None) -> "BetaBinomialSampler":
"""Return a sampler for drawing counts from this distribution."""
return BetaBinomialSampler(self, seed)
[docs]
def estimator(self, pseudo_count: float | None = None) -> "BetaBinomialEstimator":
"""Return a method-of-moments estimator for ``a, b`` at the fixed number of trials ``n``."""
if pseudo_count is None:
return BetaBinomialEstimator(self.n, name=self.name, keys=self.keys)
# Convert this distribution's own (n, a, b) into the raw first two moments -- the space
# estimate() accumulates in -- so pseudo_count can blend a prior pseudo-sample toward them
# (mirrors GumbelEstimator / WeibullEstimator's suff_stat pattern). mean = n*pi,
# var = binom_var*(1 + rho*(n-1)) with pi = a/(a+b), rho = 1/(a+b+1) -- the same
# mean-inflation relation estimate() inverts (rho = (var/binom_var - 1)/(n-1)).
pi0 = self.a / (self.a + self.b)
rho0 = 1.0 / (self.a + self.b + 1.0)
mean0 = self.n * pi0
binom_var0 = self.n * pi0 * (1.0 - pi0)
var0 = binom_var0 * (1.0 + rho0 * (self.n - 1.0))
second0 = var0 + mean0 * mean0
return BetaBinomialEstimator(
self.n, pseudo_count=pseudo_count, suff_stat=(mean0, second0), name=self.name, keys=self.keys
)
[docs]
def dist_to_encoder(self) -> "BetaBinomialDataEncoder":
"""Return the data encoder used by this distribution for vectorized methods."""
return BetaBinomialDataEncoder(self.n)
[docs]
class BetaBinomialEnumerator(DistributionEnumerator):
"""Enumerate the bounded beta-binomial support {0, ..., n} in descending probability order.
Unlike the binomial, the beta-binomial pmf is not always unimodal (it is U-shaped for
``a < 1, b < 1``), so a two-pointer walk from a mode is not safe here; the finite support is
materialized and sorted instead (the CategoricalEnumerator pattern), with value-order tie-breaks.
"""
def __init__(self, dist: BetaBinomialDistribution) -> None:
super().__init__(dist)
lp = dist.seq_log_density(np.arange(dist.n + 1, dtype=np.float64))
order = np.argsort(-lp, kind="stable")
self._entries = [(int(k), float(lp[k])) for k in order if np.isfinite(lp[k])]
self._pos = 0
def __next__(self) -> tuple[int, float]:
if self._pos >= len(self._entries):
raise StopIteration
rv = self._entries[self._pos]
self._pos += 1
return rv
[docs]
class BetaBinomialSampler(DistributionSampler):
"""Draw counts as ``p ~ Beta(a, b)`` then ``k ~ Binomial(n, p)``."""
def __init__(self, dist: BetaBinomialDistribution, seed: int | None = None) -> None:
self.rng = RandomState(seed)
self.dist = dist
[docs]
def sample(self, size: int | None = None, *, batched: bool = True) -> int | np.ndarray:
"""Draw one count or an array of iid counts."""
d = self.dist
p = self.rng.beta(d.a, d.b, size=size)
return self.rng.binomial(d.n, p)
[docs]
class BetaBinomialAccumulator(SequenceEncodableStatisticAccumulator):
"""Accumulate weighted first and second moments for beta-binomial estimation."""
def __init__(self, n: int, name: str | None = None, keys: str | None = None) -> None:
self.n = n
self.sum = 0.0
self.sum2 = 0.0
self.count = 0.0
self.name = name
self.keys = keys
[docs]
def update(self, x: int, weight: float, estimate: BetaBinomialDistribution | None) -> None:
"""Accumulate weighted first and second moments for one count."""
weights = nonnegative_weights([weight], shape=(1,))
if weights[0] == 0.0:
return
value = int(
exact_integer_observations(
[x],
label="Beta-binomial observations",
minimum=0,
maximum=self.n,
)[0]
)
self.sum += weight * value
self.sum2 += weight * value * value
self.count += weight
[docs]
def initialize(self, x: int, weight: float, rng: RandomState | None) -> None:
"""Initialize statistics from one count."""
self.update(x, weight, None)
[docs]
def seq_update(self, x: np.ndarray, weights: np.ndarray, estimate: BetaBinomialDistribution | None) -> None:
"""Accumulate weighted moments from encoded counts."""
xx = exact_integer_observations(
x,
label="Beta-binomial observations",
minimum=0,
maximum=self.n,
).astype(np.float64)
ww = nonnegative_weights(weights, shape=xx.shape)
self.sum += float(np.dot(xx, ww))
self.sum2 += float(np.dot(xx * xx, ww))
self.count += float(np.sum(ww))
[docs]
def seq_initialize(self, x: np.ndarray, weights: np.ndarray, rng: RandomState | None) -> None:
"""Initialize statistics from encoded counts."""
self.seq_update(x, weights, None)
[docs]
def combine(self, suff_stat: tuple[float, float, float]) -> "BetaBinomialAccumulator":
"""Merge another beta-binomial sufficient-statistic tuple."""
self.sum += suff_stat[0]
self.sum2 += suff_stat[1]
self.count += suff_stat[2]
return self
[docs]
def value(self) -> tuple[float, float, float]:
"""Return accumulated sum, second moment sum, and count."""
return self.sum, self.sum2, self.count
[docs]
def from_value(self, x: tuple[float, float, float]) -> "BetaBinomialAccumulator":
"""Replace accumulator contents from a sufficient-statistic tuple."""
self.sum, self.sum2, self.count = float(x[0]), float(x[1]), float(x[2])
return self
[docs]
def acc_to_encoder(self) -> "BetaBinomialDataEncoder":
"""Return the encoder used by this accumulator."""
return BetaBinomialDataEncoder(self.n)
[docs]
class BetaBinomialAccumulatorFactory(StatisticAccumulatorFactory):
"""Factory for BetaBinomialAccumulator."""
def __init__(self, n: int, name: str | None = None, keys: str | None = None) -> None:
self.n = n
self.name = name
self.keys = keys
[docs]
def make(self) -> BetaBinomialAccumulator:
"""Create a fresh beta-binomial accumulator."""
return BetaBinomialAccumulator(self.n, name=self.name, keys=self.keys)
[docs]
class BetaBinomialEstimator(ParameterEstimator):
"""Method-of-moments estimator for the beta-binomial shape parameters."""
def __init__(
self,
n: int,
pseudo_count: float | None = None,
suff_stat: tuple[float, float] | None = None,
min_conc: float = 1.0e-6,
max_conc: float = 1.0e8,
name: str | None = None,
keys: str | None = None,
) -> None:
# Mirrors BetaBinomialDistribution's own n validation: previously this cast bare int(n),
# which silently truncates a fractional n (10.5 -> 10) instead of rejecting it, and the
# bad n would only surface later -- if at all -- once estimate() built a distribution from it.
n_val = float(n)
if not np.isfinite(n_val) or n_val != math.floor(n_val) or n_val < 0:
raise ValueError("BetaBinomialEstimator requires n to be a non-negative integer.")
self.n = int(n_val)
self.pseudo_count = pseudo_count
self.suff_stat = suff_stat
self.min_conc = min_conc
self.max_conc = max_conc
self.name = name
self.keys = keys
[docs]
def accumulator_factory(self) -> BetaBinomialAccumulatorFactory:
"""Return an accumulator factory for beta-binomial moments."""
return BetaBinomialAccumulatorFactory(self.n, name=self.name, keys=self.keys)
[docs]
def estimate(self, nobs: float | None, suff_stat: tuple[float, float, float]) -> BetaBinomialDistribution:
"""Estimate beta-binomial shape parameters from weighted moments."""
sum_x, sum_x2, count = suff_stat
if self.pseudo_count is not None and self.suff_stat is not None:
mean0, second0 = self.suff_stat
sum_x += self.pseudo_count * mean0
sum_x2 += self.pseudo_count * second0
count += self.pseudo_count
n = self.n
if count <= 0.0 or n == 0:
return BetaBinomialDistribution(n, 1.0, 1.0, name=self.name, keys=self.keys)
mean = sum_x / count
var = sum_x2 / count - mean * mean
pi = min(max(mean / n, 1.0e-12), 1.0 - 1.0e-12) # success probability a/(a+b)
binom_var = n * pi * (1.0 - pi)
if n == 1 or binom_var <= 0.0:
s = self.max_conc # no overdispersion information -> binomial limit
else:
rho = (var / binom_var - 1.0) / (n - 1.0) # intra-class correlation
rho = min(max(rho, 1.0 / (self.max_conc + 1.0)), 1.0 - 1.0e-9)
s = 1.0 / rho - 1.0 # a + b
s = min(max(s, self.min_conc), self.max_conc)
return BetaBinomialDistribution(n, pi * s, (1.0 - pi) * s, name=self.name, keys=self.keys)
[docs]
class BetaBinomialDataEncoder(DataSequenceEncoder):
"""Encode beta-binomial counts as a float array."""
def __init__(self, n: int) -> None:
self.n = n
def __str__(self) -> str:
return f"BetaBinomialDataEncoder(n={self.n!r})"
def __eq__(self, other: object) -> bool:
return isinstance(other, BetaBinomialDataEncoder) and self.n == other.n
[docs]
def seq_encode(self, x: Sequence[int]) -> np.ndarray:
"""Encode exact counts in the fixed support ``{0, ..., n}``."""
return exact_integer_observations(
x,
label="Beta-binomial observations",
minimum=0,
maximum=self.n,
)