!pip install pyro-ppl

import matplotlib.pyplot as plt
import numpy as np
import torch

import pyro
import pyro.infer
import pyro.optim
import pyro.distributions as dist
from torch.distributions import constraints
from pyro.infer import MCMC, NUTS, Predictive, HMC
from pyro.infer import SVI, Trace_ELBO
from pyro.infer import EmpiricalMarginal

from pyro.distributions import Beta, Binomial, HalfCauchy, Normal, Pareto, Uniform
from pyro.distributions.util import scalar_like
from pyro.infer.mcmc.util import initialize_model, summary
from pyro.util import ignore_experimental_warning

pyro.set_rng_seed(101)

"""
Explanation: <a href="https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/pyro_intro.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Pyro is a probabilistic programming system built on top of PyTorch. It supports posterior inference based on MCMC and stochastic variational inference; discrete latent variables can be marginalized out exactly using dynamic programmming.
End of explanation
"""


def model(hparams, data=None):
    prior_mean, prior_sd, obs_sd = hparams
    theta = pyro.sample("theta", dist.Normal(prior_mean, prior_sd))
    y = pyro.sample("y", dist.Normal(theta, obs_sd), obs=data)
    return y

"""
Explanation: Example: inferring mean of 1d Gaussian .
We use the simple example from the Pyro intro. The goal is to infer the weight $\theta$ of an object, given noisy measurements $y$. We assume the following model:
$$
\begin{align}
\theta &\sim N(\mu=8.5, \tau^2=1.0)\ 
y \sim &N(\theta, \sigma^2=0.75^2)
\end{align}
$$
Where $\mu=8.5$ is the initial guess.
End of explanation
"""


mu = 8.5
tau = 1.0
sigma = 0.75
hparams = (mu, tau, sigma)

y = 9.5
m = (sigma**2 * mu + tau**2 * y) / (sigma**2 + tau**2)  # posterior mean
s2 = (sigma**2 * tau**2) / (sigma**2 + tau**2)  # posterior variance
s = np.sqrt(s2)
print(m)
print(s)

"""
Explanation: Exact inference
By Bayes rule for Gaussians, we know that the exact posterior,
given a single observation $y=9.5$, is given by
$$
\begin{align}
\theta|y &\sim N(m, s^s) \
m &=\frac{\sigma^2 \mu + \tau^2 y}{\sigma^2 + \tau^2} 
  = \frac{0.75^2 \times 8.5 + 1 \times 9.5}{0.75^2 + 1^2}
  = 9.14 \
s^2 &= \frac{\sigma^2 \tau^2}{\sigma^2  + \tau^2} 
= \frac{0.75^2 \times 1^2}{0.75^2 + 1^2}= 0.6^2
\end{align}
$$
End of explanation
"""


def model2(hparams, data=None):
    prior_mean, prior_sd, obs_sd = hparams
    theta = pyro.sample("theta", dist.Normal(prior_mean, prior_sd))
    y = pyro.sample("y", dist.Normal(theta, obs_sd), obs=data)
    return theta, y


for i in range(5):
    theta, y = model2(hparams)
    print([theta, y])

"""
Explanation: Ancestral sampling
End of explanation
"""


nuts_kernel = NUTS(model)
obs = torch.tensor(y)
mcmc = MCMC(nuts_kernel, num_samples=1000, warmup_steps=50)
mcmc.run(hparams, obs)
print(type(mcmc))

samples = mcmc.get_samples()
print(type(samples))
print(samples.keys())
print(samples["theta"].shape)

mcmc.diagnostics()

thetas = samples["theta"].numpy()
print(np.mean(thetas))
print(np.std(thetas))

"""
Explanation: MCMC
See the documentation
End of explanation
"""


# the guide must have the same signature as the model
def guide(hparams, data):
    y = data
    prior_mean, prior_sd, obs_sd = hparams
    m = pyro.param("m", torch.tensor(y))  # location
    s = pyro.param("s", torch.tensor(prior_sd), constraint=constraints.positive)  # scale
    return pyro.sample("theta", dist.Normal(m, s))


# initialize variational parameters
pyro.clear_param_store()

# set up the optimizer
# optimizer = pyro.optim.Adam({"lr": 0.001, "betas": (0.90, 0.999)})
optimizer = pyro.optim.SGD({"lr": 0.001, "momentum": 0.1})

# setup the inference algorithm
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

n_steps = 2000
# do gradient steps
obs = torch.tensor(y)
loss_history, m_history, s_history = [], [], []
for t in range(num_steps):
    loss_history.append(svi.step(hparams, obs))
    m_history.append(pyro.param("m").item())
    s_history.append(pyro.param("s").item())

plt.plot(loss_history)
plt.title("ELBO")
plt.xlabel("step")
plt.ylabel("loss")

post_mean = pyro.param("m").item()
post_std = pyro.param("s").item()
print([post_mean, post_std])

"""
Explanation: Variational Inference
See the documentation
For the guide (approximate posterior), we use a pytorch.distributions.normal.
End of explanation
"""


alpha0 = 10.0
beta0 = 10.0


def model(data):
    alpha0_tt = torch.tensor(alpha0)
    beta0_tt = torch.tensor(beta0)
    f = pyro.sample("theta", dist.Beta(alpha0_tt, beta0_tt))
    # loop over the observed data
    for i in range(len(data)):
        pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])


def model_binom(data):
    alpha0_tt = torch.tensor(alpha0)
    beta0_tt = torch.tensor(beta0)
    theta = pyro.sample("theta", dist.Beta(alpha0_tt, beta0_tt))
    data_np = [x.item() for x in data]
    N = len(data_np)
    N1 = np.sum(data_np)
    N0 = N - N1
    pyro.sample("obs", dist.Binomial(N, theta))

# create some data with 6 observed heads and 4 observed tails
data = []
for _ in range(6):
    data.append(torch.tensor(1.0))
for _ in range(4):
    data.append(torch.tensor(0.0))

data_np = [x.item() for x in data]
print(data)
print(data_np)

N = len(data_np)
N1 = np.sum(data_np)
N0 = N - N1
print([N1, N0])

"""
Explanation: Example: beta-bernoulli model
Example is from SVI tutorial
The model is
$$
\begin{align}
\theta &\sim \text{Beta}(\alpha, \beta) \
x_i &\sim \text{Ber}(\theta)
\end{align}
$$
where $\alpha=\beta=10$.
End of explanation
"""


alpha1 = alpha0 + N1
beta1 = beta0 + N0
print("exact posterior: alpha={:0.3f}, beta={:0.3f}".format(alpha1, beta1))
post_mean = alpha1 / (alpha1 + beta1)
post_var = (post_mean * beta1) / ((alpha1 + beta1) * (alpha1 + beta1 + 1))
post_std = np.sqrt(post_var)
print([post_mean, post_std])

"""
Explanation: Exact inference
The posterior is given by
$$
\begin{align}
\theta &\sim \text{Ber}(\alpha + N_1, \beta  + N_0) \
N_1 &= \sum_{i=1}^N [x_i=1] \
N_0 &= \sum_{i=1}^N [x_i=0]
\end{align}
$$
End of explanation
"""


nuts_kernel = NUTS(model)
mcmc = MCMC(nuts_kernel, num_samples=1000, warmup_steps=50)
mcmc.run(data)
print(mcmc.diagnostics())
samples = mcmc.get_samples()
print(samples["theta"].shape)

thetas = samples["theta"].numpy()
print(np.mean(thetas))
print(np.std(thetas))

nuts_kernel = NUTS(model_binom)
mcmc = MCMC(nuts_kernel, num_samples=1000, warmup_steps=50)
mcmc.run(data)
print(mcmc.diagnostics())
samples = mcmc.get_samples()
print(samples["theta"].shape)

thetas = samples["theta"].numpy()
print(np.mean(thetas))
print(np.std(thetas))

"""
Explanation: MCMC
End of explanation
"""


def guide(data):
    alpha_q = pyro.param("alpha_q", torch.tensor(15.0), constraint=constraints.positive)
    beta_q = pyro.param("beta_q", torch.tensor(15.0), constraint=constraints.positive)
    pyro.sample("theta", dist.Beta(alpha_q, beta_q))

# optimizer = pyro.optim.Adam({"lr": 0.0005, "betas": (0.90, 0.999)})
optimizer = pyro.optim.SGD({"lr": 0.001, "momentum": 0.1})

svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

n_steps = 2000
loss_history = []
for step in range(n_steps):
    loss_history.append(svi.step(data))

plt.plot(loss_history)
plt.title("ELBO")
plt.xlabel("step")
plt.ylabel("loss");

# grab the learned variational parameters
alpha_q = pyro.param("alpha_q").item()
beta_q = pyro.param("beta_q").item()
print("variational posterior: alpha={:0.3f}, beta={:0.3f}".format(alpha_q, beta_q))

post_mean = alpha_q / (alpha_q + beta_q)
post_var = (post_mean * beta_q) / ((alpha_q + beta_q) * (alpha_q + beta_q + 1))
post_std = np.sqrt(post_var)
print([post_mean, post_std])

"""
Explanation: Variational inference
End of explanation
"""


