Bayesian Inference for Beta-Binomial Model

import math

def beta_binomial_posterior(k: int, n: int,
                            alpha_prior: float = 1.0,
                            beta_prior: float = 1.0) -> dict:
    # Posterior parameters (conjugate update)
    alpha_post = alpha_prior + k
    beta_post = beta_prior + (n - k)

    # Posterior mean
    mean = alpha_post / (alpha_post + beta_post)

    # Posterior mode (MAP estimate)
    if alpha_post > 1 and beta_post > 1:
        mode = (alpha_post - 1) / (alpha_post + beta_post - 2)
    else:
        mode = mean  # fallback

    # Posterior variance
    total = alpha_post + beta_post
    variance = (alpha_post * beta_post) / (total ** 2 * (total + 1))

    # 95% credible interval using beta quantile approximation
    lower = _beta_quantile(0.025, alpha_post, beta_post)
    upper = _beta_quantile(0.975, alpha_post, beta_post)

    return {
        "alpha_posterior": alpha_post,
        "beta_posterior": beta_post,
        "posterior_mean": round(mean, 6),
        "posterior_mode": round(mode, 6),
        "posterior_variance": round(variance, 6),
        "credible_interval_95": (round(lower, 4), round(upper, 4)),
    }

def _beta_quantile(p, a, b, tol=1e-8, max_iter=100):
    """Find x such that I_x(a, b) = p using bisection."""
    lo, hi = 0.0, 1.0
    for _ in range(max_iter):
        mid = (lo + hi) / 2
        if _beta_cdf(mid, a, b) < p:
            lo = mid
        else:
            hi = mid
        if hi - lo < tol:
            break
    return (lo + hi) / 2

def _beta_cdf(x, a, b):
    """Regularized incomplete beta function via continued fraction."""
    if x <= 0:
        return 0.0
    if x >= 1:
        return 1.0
    if x > (a + 1) / (a + b + 2):
        return 1.0 - _beta_cdf(1 - x, b, a)

    ln_pre = a * math.log(x) + b * math.log(1 - x) - _log_beta(a, b)
    # Continued fraction
    f, c, d = 1.0, 1.0, 0.0
    for m in range(200):
        if m == 0:
            num = 1.0
        elif m % 2 == 0:
            k = m // 2
            num = k * (b - k) * x / ((a + 2 * k - 1) * (a + 2 * k))
        else:
            k = (m + 1) // 2
            num = -(a + k - 1 + (a + b) * 0 + k) * 0
            k2 = m // 2 + 1
            num = -(a + k2 - 1) * (a + b + k2 - 1) * x / (
                (a + 2 * k2 - 2) * (a + 2 * k2 - 1))
        d = 1.0 + num * d
        if abs(d) < 1e-30:
            d = 1e-30
        c = 1.0 + num / c
        if abs(c) < 1e-30:
            c = 1e-30
        d = 1.0 / d
        f *= c * d
    return math.exp(ln_pre) * (f - 1.0) / a

def _beta_cdf(x, a, b):
    """Regularized incomplete beta function - simple numerical integration."""
    if x <= 0:
        return 0.0
    if x >= 1:
        return 1.0
    n_steps = 1000
    dx = x / n_steps
    total = 0.0
    ln_beta = _log_beta(a, b)
    for i in range(n_steps):
        t = (i + 0.5) * dx
        total += math.exp((a - 1) * math.log(t) + (b - 1) * math.log(1 - t) - ln_beta) * dx
    return total

def _log_beta(a, b):
    return math.lgamma(a) + math.lgamma(b) - math.lgamma(a + b)