posterior.py

#%%
# Compute posterior of theta from coin tosses

#%%
%config InlineBackend.figure_format = 'retina'

#%%
from matplotlib import pyplot as plt
import numpy as np
import seaborn as sns
from scipy import stats
sns.set_style('white')

# Theta values
theta = np.linspace(0, 1, 1000)

#%%
hidden_theta = .6
generator = stats.bernoulli.rvs

#%%
# Discrete uniform
prior = [1/len(theta)] * len(theta)

def bernoulli_likelihood(theta_p: float, y_p: int) -> float:
    return (theta_p**y_p) * ((1-theta_p)**(1-y_p))

#%%
# Prior for n theta
plt.plot(theta, prior, 'ko', ms=8)
plt.vlines(theta, 0, prior, colors='k', lw=5, alpha=0.5)
plt.title('Prior')

#%%
ks = 10
fig, axs = plt.subplots(ks, 1, figsize=(5, ks*1), sharex=True)
for k in range(ks):
    y = generator(hidden_theta)
    posterior = []
    for ix, t in enumerate(theta):
        posterior.append(prior[ix] * bernoulli_likelihood(t, y))
    color = 'r' if y else 'k'
    axs[k].plot(theta, posterior, 'ko', ms=5, color=color)
    axs[k].vlines(theta, 0, posterior, colors='k', lw=5, alpha=0.5)
    prior = posterior
    prior /= sum(prior)