Metropolis Hastings

import torch
dist = torch.distributions
import matplotlib.pyplot as plt
import pandas as pd
%matplotlib inline
import warnings
warnings.filterwarnings('ignore')

mix = dist.MixtureSameFamily(
    mixture_distribution=dist.Categorical(torch.tensor([0.2, 0.8])),
    component_distribution=dist.Normal(torch.tensor([0.1, 3]), torch.tensor([0.6, 3])),
)

x = torch.linspace(-4., 4, 5)

theta_range = torch.linspace(-8, 12, 100)
plt.plot(theta_range, mix.log_prob(theta_range).exp())

[<matplotlib.lines.Line2D at 0x165fa8c10>]

mix.mean, mix.stddev

(tensor(2.4200), tensor(2.9356))

next_sample = lambda cur_sample: dist.Normal(loc = cur_sample, scale=1).sample().item()

next_sample(1)

2.0207135677337646

p = lambda x: mix.log_prob(torch.tensor(x)).exp().item()
lp = lambda x: mix.log_prob(torch.tensor(x)).item()

p(2)

0.10151920467615128

num_iter = 10
xs = [None]*num_iter
xs[0] = 0.
plt.plot(theta_range, mix.log_prob(theta_range).exp())
plt.scatter(xs[0], p(xs[0]))

xs[1] = next_sample(xs[0])
plt.scatter(xs[1], p(xs[1]))

<matplotlib.collections.PathCollection at 0x166017fa0>

num_iter = 10
xs = [None]*num_iter
xs[0] = 0.2
plt.plot(theta_range, mix.log_prob(theta_range).exp())
plt.scatter(xs[0], p(xs[0]))

xs[1] = next_sample(xs[0])
plt.scatter(xs[1], p(xs[1]))

a = p(xs[1])/p(xs[0])
a

0.27602488498335154

u = dist.Uniform(0, 1).sample().item()
print(u)

0.049220144748687744

if u < a:
    xs[1] = xs[0]

x_start = 0.
num_iter = 20000
xs = torch.empty(num_iter)
xs[0] = x_start
lu = torch.log(dist.Uniform(0, 1).sample([num_iter]))
for i in range(1, num_iter):
    xs[i] = next_sample(xs[i-1])
   
    la = lp(xs[i]) - lp(xs[i-1])

    if lu[i] > la:
        xs[i] = xs[i-1]
    
    

plt.plot(xs)

[<matplotlib.lines.Line2D at 0x1660a9d90>]

import seaborn as sns

plt.plot(theta_range, mix.log_prob(theta_range).exp(), label='True')
sns.kdeplot(torch.tensor(xs[:10000]), label='Samples obtained from MH')
plt.legend()

<matplotlib.legend.Legend at 0x16613f790>

xs[:1000].mean(), mix.mean,  xs[:1000].std(), mix.stddev

(tensor(3.5472), tensor(2.4200), tensor(3.0889), tensor(2.9356))

import numpy as np

g = emcee.moves.GaussianMove(cov = 1.)

import emcee

nwalkers = 2
ndim = 1
log_prob = lp
p0 = np.random.rand(nwalkers, ndim)

sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob,moves=g)
state = sampler.run_mcmc(p0, 100)
sampler.reset()
sampler.run_mcmc(state, 10000, progress=True);

100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:03<00:00, 3105.40it/s]

samples = sampler.get_chain(flat=True)

plt.plot(samples)

[<matplotlib.lines.Line2D at 0x16619fd60>]

sns.kdeplot(samples.flatten(), label='EmCEE')

plt.plot(theta_range, mix.log_prob(theta_range).exp(), label='True')
sns.kdeplot(torch.tensor(xs[:10000]), label='Samples obtained from MH')
plt.legend()

<matplotlib.legend.Legend at 0x1661b3490>

Creating a function

def mh(log_p, next_sample, num_iter, x_start):
    lu = torch.log(dist.Uniform(0, 1).sample([num_iter]))
    try: 
        l = len(x_start)
        xs = torch.empty((num_iter, len(x_start)))
    except:
        xs = torch.empty(num_iter)
    xs[0] = x_start
    
    for i in range(1, num_iter):
        xs[i] = next_sample(xs[i-1])

        la = log_p(xs[i]) - log_p(xs[i-1])

        if lu[i] > la:
            xs[i] = xs[i-1]
    return xs

xs = mh(lp, next_sample, 10000, 0.)

plt.plot(xs)

[<matplotlib.lines.Line2D at 0x166342280>]

2d

dist_2d = dist.MultivariateNormal(loc = torch.zeros(2), covariance_matrix=torch.tensor([[1., 0.5], [0.5, 2.]]))
dist_2d

MultivariateNormal(loc: torch.Size([2]), covariance_matrix: torch.Size([2, 2]))

log_p_2d = lambda x: dist_2d.log_prob(torch.tensor(x)).item()
next_sample_2d = lambda cur_sample: dist.MultivariateNormal(
    loc=cur_sample, covariance_matrix=torch.eye(len(cur_sample))
).sample()

x_2d = mh(log_p_2d, next_sample_2d, 5000, torch.tensor([0., 0.]))

sns.kdeplot(x_2d[:, 0], x_2d[:, 1], bw_adjust=4)

<AxesSubplot:>

import corner

corner.corner(x_2d.numpy(), smooth=2, show_titles=True, labels=[r"$x_1$", "$x_2$"], smooth1d=1, );

Linear Regression with 1 parameter

y = 4*x + 0.5*torch.randn(5)
prior = dist.Normal(loc = 0., scale = 1.)
log_likelihood = lambda t: dist.Normal(loc = x*t, scale=1.).log_prob(y).sum(axis=0)
unnorm_post = lambda t:log_likelihood(torch.tensor(t)).item() - prior.log_prob(torch.tensor(t)).item()

xs = mh(unnorm_post, next_sample, 10000, 0.)

sns.kdeplot(xs)

<AxesSubplot:ylabel='Density'>

import seaborn as sns

sns.kdeplot(xs.numpy(), label='Samples obtained from MH')
plt.axvline(0., label='Prior mean', color='k', linestyle='--')
plt.axvline(4, label='True value', color='g', linestyle='-.')
plt.legend()

<matplotlib.legend.Legend at 0x176e6b880>

plt.plot(xs)

[<matplotlib.lines.Line2D at 0x176ec3880>]

for i in range(100):
    plt.plot(x, xs[i]*x, alpha=0.1, color='k')
plt.scatter(x, y, zorder=10)

<matplotlib.collections.PathCollection at 0x1770017f0>

xs_mean = xs.mean()
xs_std = xs.std()

plt.plot(x, xs_mean*x, color='k')
plt.scatter(x, y, zorder=10)
plt.fill_between(x, (xs_mean-2*xs_std)*x, (xs_mean+2*xs_std)*x, color='k', alpha=0.1)

<matplotlib.collections.PolyCollection at 0x177104eb0>

Linear Regression with 2 parameter

y = 4*x + 0.5*torch.randn(5)
prior = dist.MultivariateNormal(loc = torch.zeros([2]), covariance_matrix=torch.eye(2))
log_likelihood = lambda t: dist.Normal(loc = x*t[1] + t[0], scale=1.).log_prob(y).sum(axis=0)
unnorm_post = lambda t:log_likelihood(torch.tensor(t)).item() - prior.log_prob(torch.tensor(t)).item()

References

1. https://www.youtube.com/watch?v=yCv2N7wGDCw