Exercise - Estimating Mean and Standard Deviation of Normal Distribution with Pyro

import numpy as np

import scipy.stats
from scipy.stats import norm

from matplotlib import pyplot as plt
from IPython.core.pylabtools import figsize

%matplotlib inline

import torch
from torch.distributions import constraints

import pyro
import pyro.infer
import pyro.optim as optim
import pyro.distributions as dist

dtype=torch.float32

torch.manual_seed(101)
pyro.set_rng_seed(101)
np.random.seed(12);

# generate observed data
N = 10
mu_ = 10.
sigma_= 2.
X = np.random.normal(mu_, sigma_, N)
X = np.array(X, dtype=np.float32)

X

x = np.arange(3,18,0.01)
p_x = scipy.stats.norm.pdf(x, loc=mu_, scale=sigma_)
plt.plot(x, p_x, label="true distribution")
plt.plot(X, np.zeros_like(X), "ro", label="observed data")
plt.title("")
plt.xlabel("x")
plt.ylabel("p(x)")
plt.legend();

def model(X):
    # Prior
    mu = pyro.sample("mu", dist.Uniform(torch.tensor(-25.), torch.tensor(+25.))) 
    tau = torch.tensor(1/4)


    # pyro plate mark the samples as conditional independet
    with pyro.plate("observed_data", size=len(X)):
        sample = pyro.sample("gaussian_data", dist.Normal(mu, 1/torch.sqrt(tau)), obs=X)
    
    return sample

### same (function) arguments for guide and model !
def guide(X):
    mean_loc = torch.randn((1)) 
    # note that we initialize the scale to be pretty narrow
    mean_scale = torch.tensor(0.001)
    
    # mu_loc and mu_scale are the (variational) parameters
    # they will be learned by SVI from the data
    
    mu_loc = pyro.param("guide_mu_mean", mean_loc)
    
    # the scale must be positive.
    mu_scale = pyro.param("guide_mu_scale", mean_scale, constraint=constraints.positive)

    # note the same name "mu" here as in our model
    mu = pyro.sample("mu", dist.Normal(mu_loc, mu_scale)) 

pyro.clear_param_store()

adam_params = {"lr": 0.003, "betas": (0.95, 0.999)}
optimizer = optim.Adam(adam_params)

svi = pyro.infer.SVI(model=model,
                     guide=guide,
                     optim=optimizer,
                     loss=pyro.infer.Trace_ELBO())

### to keep track of our loss history
losses = []

### convert observed data to a torch tensor object
X_ = torch.tensor(X, dtype=dtype)

### training / inference
for t in range(10000):
    ### svi.step takes same parameters as inpust as our defined model(X) and guide(X) function
    loss = svi.step(X_)
    losses.append(loss)
    ### for monitoring
    if t%100==0:
        print (t, "\t", loss)

### Let us plot the costs / iteration curve

plt.xlabel("# iteration")
plt.ylabel("MC-Estimate of ELBO")
plt.plot(range(len(losses)), losses)

# Adjust the strings according to your names for
# the parameters "mu_mean", etc...
mu_mean_param = pyro.param("guide_mu_mean")
mu_scale_param = pyro.param("guide_mu_scale")
mu_mean_param, mu_scale_param

plt.figure(figsize=(12,4))

mu_mean = mu_mean_param.detach().numpy()
mu_scale = mu_scale_param.detach().numpy()

x = np.arange(5,15,0.01)
p_mu = scipy.stats.norm.pdf(x, loc=mu_mean, scale=np.sqrt(mu_scale))
ax = plt.subplot(121)
ax.plot(x, p_mu)
ax.set_xlabel("$\mu$")
ax.set_ylabel("q($\mu$)")
ax.set_title("Mean: q($\\mu$)")
print("true mu: ", mu_)

concentration = 1.5
rate = 1.

x = np.arange(0,4,0.01)
p_tau = scipy.stats.gamma.pdf(x, a=concentration, scale=1/rate)
plt.xlabel("x")
plt.ylabel("p(x)")
plt.plot(x, p_tau)

# Note that dist.Gamma has a different parameter signature: __init__(self, concentration, rate, validate_args=None)
# see:
#help(dist.Gamma)

def model_with_tau(X):
    
    ######################
    ### Your Code here ###
    ######################

    
    return

def guide_with_tau(X):
    
    ######################
    ### Your Code here ###
    ######################

    
    return

### Initilize pyro.infer.SVI object

######################
### Your Code here ###
######################

### Training

######################
### Your Code here ###
######################

### Let us plot the costs / iteration curve

plt.xlabel("# iteration")
plt.ylabel("MC-Estimate of ELBO")
plt.plot(range(len(losses)), losses)

# Adjust the strings according to your names for
# the parameters "mu_mean", etc...
mu_mean_param = pyro.param("guide_mu_mean")
mu_scale_param = pyro.param("guide_mu_scale")
mu_mean_param, mu_scale_param

# Adjust the strings according to your names for
# the parameters "mu_mean", etc...
tau_concentration_param = pyro.param("guide_tau_concentration")
tau_rate_param = pyro.param("guide_tau_rate")
tau_concentration_param, tau_rate_param

plt.figure(figsize=(12,4))

mu_mean = mu_mean_param.detach().numpy()
mu_scale = mu_scale_param.detach().numpy()

x = np.arange(5,15,0.01)
p_mu = scipy.stats.norm.pdf(x, loc=mu_mean, scale=np.sqrt(mu_scale))
ax = plt.subplot(121)
ax.plot(x, p_mu)
ax.set_xlabel("$\mu$")
ax.set_ylabel("q($\mu$)")
ax.set_title("Mean: q($\\mu$)")
print("true mu: ", mu_)

tau_concentration =tau_concentration_param.detach().numpy()
tau_rate = tau_rate_param.detach().numpy()

x = np.arange(0,1,0.01)
p_tau = scipy.stats.gamma.pdf(x, a=tau_concentration, scale=1/tau_rate)
ax = plt.subplot(122)
ax.plot(x, p_tau)
ax.set_xlabel("$\\tau$")
ax.set_ylabel("q($\\tau$)")
ax.set_title("Precision: q($\\tau$)")
print("true tau: ", 1/sigma_**2)