ci-math.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from scipy.stats import beta as beta_dist

def credible_interval(t, f, prob=.5, prior='uniform'):
  """Calculates posterior credible interval for a beta-binomial model with prior `Beta(a, b)` after `t` successes in `t+f` trials.

  @param t: Number of successes.
  @param f: Number of failures.
  @param prob: Probability mass inside the posterior interval.
  @param prior: Either a tuple (a, b) parametrizing the Beta prior or one of "uniform" (a=b=1) or "jeffreys" (a=b=1/2).
  @returns Credible interval bounds as a tuple `(lower, upper)`.
  """
  if prior == 'uniform':
    a, b = 1, 1
  elif prior == 'jeffreys':
    a, b = .5, .5
  else:
    a, b = prior

  # NB: Beta(a, b) and Binom(k, n) yield the same CDF for the continuous parameter between (0, 1) (the probability) when a=k+1 and b=n-k+1
  ppf = beta_dist(a+t, b+f).ppf

  lo = (1 - prob) / 2
  hi = lo + prob

  return ppf(lo), ppf(hi)