estimate a univariate Gaussian with Gibbs sampling

gibbs.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm, invgamma

# estimate a univariate gaussian by Gibbs sampling
true_mu = 10
true_sigma = 2
true_v = true_sigma**0.5
N = 50
R = 5000
D = norm.rvs(true_mu, true_sigma, size=N)
# D = np.concatenate((norm.rvs(true_mu, true_sigma, size=N*9//10), norm.rvs(0, 1, size=N//10)))
Dbar = D.mean()
# priors
mu_0 = 0
v_mu = 1000
alpha = 1
beta = 1
# init state
mu_state = 0
v_state = 1
mus = [mu_state]
vs = [v_state]
xs = np.linspace(-5, 20, 1000)
p_x = 0
for _ in range(R):
    # update mu
    posterior_mu = (v_mu*Dbar+v_state/N*mu_0)/(v_mu+v_state/N)
    posterior_sigma = np.sqrt((v_mu*v_state/N)/(v_mu+v_state/N))
    mu_state = norm.rvs(posterior_mu, posterior_sigma)
    # update v
    posterior_alpha = N/2 + alpha
    posterior_beta = ((D - mu_state)**2).sum()/2 + beta
    v_state = invgamma.rvs(posterior_alpha)*posterior_beta
    
    mus.append(mu_state)
    vs.append(v_state)
    p_x += norm.pdf(xs, mu_state, v_state**0.5)
    
p_x/=R
true_p_x = norm.pdf(xs, true_mu, true_sigma)
# true_p_x = norm.pdf(xs, true_mu, true_sigma)*0.9 + norm.pdf(xs, 0, 1)*0.1
ml_p_x = norm.pdf(xs, D.mean(), D.std())
plt.plot(xs, true_p_x)
plt.plot(xs, p_x)
plt.plot(xs, ml_p_x)
plt.bar(D, np.ones_like(D)/D.shape, width=0.005)
print(xs[np.argmax(p_x)])
print(D.mean())
# black: data
# blue: ground truth
# green: gibbs
# red: maximum likelihood (sample mean & variance)