Medium example of Bayes' theorem

water_proportion.py

# Number of elements in the grid
N_grid = 100

# Define the grid
p_grid = np.linspace(0, 1, N_grid)

# Define the prior
prior = np.concatenate([np.repeat(0, int(N_grid/2)), np.repeat(2, int(N_grid/2))])

# Define the likelihood
likelihood = binom.pmf(6, 10, p_grid)

posterior = prior*likelihood
posterior /= np.sum(posterior)


# Plot results
_, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5))


titles = ['Prior', "Likelihood", "Posterior"]
lim_factor = 1.1
y_lim = [[0, np.max(prior)*lim_factor], [0, np.max(likelihood)*lim_factor], [0, np.max(posterior)*lim_factor]]
x_ticks = [0, 0.5, 1]

for i, ax in enumerate([ax1, ax2, ax3]):    
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.set_title(titles[i], fontsize=18)
    ax.set_xticks(x_ticks)
    ax.set_xticklabels(x_ticks, fontsize=14)
    ax.set_yticks([])
    ax.set_ylim(y_lim[i])
    ax.set_xlabel("$\it{s}$", fontsize=16)

ax1.plot(p_grid, prior)
ax1.fill_between(p_grid, 0, prior, alpha=0.5)
ax2.plot(p_grid, likelihood)
ax2.fill_between(p_grid, 0, likelihood, alpha=0.5)
ax3.plot(p_grid, posterior)
ax3.fill_between(p_grid, 0, posterior, alpha=0.5)
ax2.set_ylabel("X", fontsize=24, rotation=0)
ax3.set_ylabel("$\propto$", fontsize=48, rotation=0)

plt.show()