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