import torch
dist = torch.distributions
import matplotlib.pyplot as plt
import pandas as pd
%matplotlib inline

Simple rejection sampling

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 0x14871f760>]

mix.mean, mix.stddev

(tensor(2.4200), tensor(2.9356))

q = dist.Normal(loc = 0., scale=1.)
plt.plot(theta_range, q.log_prob(theta_range).exp())

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

plt.plot(theta_range, mix.log_prob(theta_range).exp(), label='p_tilde')


plt.plot(theta_range, q.log_prob(theta_range).exp(), label='q')
plt.legend()

<matplotlib.legend.Legend at 0x1487f25e0>

Choosing proposal q with high variance

q = dist.Normal(loc = 0., scale=5.)


plt.plot(theta_range, mix.log_prob(theta_range).exp(), label='p_tilde')


plt.plot(theta_range, q.log_prob(theta_range).exp(), label='q')
plt.legend()

<matplotlib.legend.Legend at 0x147ce8a00>

Scaling q so that cq(x) >= p(x) for all x

theta_more = theta_range[mix.log_prob(theta_range).exp() >  q.log_prob(theta_range).exp()]
c = (mix.log_prob(theta_more).exp()/q.log_prob(theta_more).exp()).max().item()

q = dist.Normal(loc = 0., scale=5.)


plt.plot(theta_range, mix.log_prob(theta_range).exp(), label='p_tilde')


plt.plot(theta_range, c*q.log_prob(theta_range).exp(), label='c*q')
plt.legend()

<matplotlib.legend.Legend at 0x148b6d610>

plt.gca().get_xlim()

(0.0, 1.0)

def plot_rejection(ax=None):
    if ax is None:
        fig, ax = plt.subplots()
    s = q.sample()


    ax.plot(theta_range, mix.log_prob(theta_range).exp(), label='p_tilde')


    ax.plot(theta_range, c*q.log_prob(theta_range).exp(), label='c*q')

    ax.scatter([s], 0, label='s (sample from q)', s=100,color='k')
    cqs = c*q.log_prob(s).exp().item()
    ax.axhline(cqs, label='c*q(s)', color='k')
    ax.plot([s, s], [0, cqs], color='k', linestyle='-.')
    u = dist.Uniform(0, cqs).sample()
    if u > mix.log_prob(s).exp():
        ax.scatter([s], [u], label='u', s=100,color='r')
    else:
        ax.scatter([s], [u], label='u', s=100,color='g')


    ax.legend(bbox_to_anchor=(1.04,1), borderaxespad=0)

plot_rejection()

fig, ax = plt.subplots(nrows=4, ncols=4, figsize=(20, 20))
for i in range(16):
    plot_rejection(ax[i//4, i%4])
fig.tight_layout()

N = 5000
s = q.sample([N])

fig, ax = plt.subplots()
ax.plot(theta_range, mix.log_prob(theta_range).exp(), label='p_tilde')


ax.plot(theta_range, c*q.log_prob(theta_range).exp(), label='c*q')
ax.scatter(s, torch.zeros_like(s), label='s (sample from q)', s=5,color='k', marker='|')
cqs = c*q.log_prob(s).exp()
u = dist.Uniform(0, cqs).sample()
gt_index = u > mix.log_prob(s).exp()
gt = u[gt_index]
s_gt = s[gt_index]

s_lt = s[~gt_index]
lt = u[~gt_index]
ax.scatter(s_gt, gt, s=1,color='r', label='Rejected')
ax.scatter(s_lt, lt, label='Accepted', s=1,color='g')
ax.legend()

<matplotlib.legend.Legend at 0x14c6efb20>

sns.kdeplot(s_lt)
sns.kdeplot(mix.sample([1000]))

<AxesSubplot:ylabel='Density'>

Linear Regression

y = 4*x + torch.randn(5)

plt.scatter(x, y)

<matplotlib.collections.PathCollection at 0x146cd6f40>

prior = dist.Normal(loc = 0., scale = 1.)

likelihood = lambda t: dist.Normal(loc = x*t, scale=1.).log_prob(y).sum(axis=0).exp().item()

log_likelihood = lambda t: dist.Normal(loc = x*t, scale=1.).log_prob(y).sum(axis=0)

lls = {}
ls = {}
theta_range = torch.linspace(-1, 7, 100)
for t in theta_range:
    lls[t.item()] = log_likelihood(t).item()
    ls[t.item()] = log_likelihood(t).exp().item()

pd.Series(ls).plot()

<AxesSubplot:>

pd.Series(lls).plot()

<AxesSubplot:>

plt.plot(theta_range, prior.log_prob(theta_range).exp())

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

unnorm_post = lambda t:log_likelihood(t).item() + prior.log_prob(t).item()

up = {}
for t in theta_range:
    up[t.item()] = unnorm_post(t)

fig, ax = plt.subplots(nrows=3, sharex=True, figsize= (8, 6))
pd.Series(up).plot(ax=ax[0], title="Unnorm Posterior")
pd.Series(lls).plot(ax=ax[1], title='Loglik')
ax[2].plot(theta_range, prior.log_prob(theta_range))

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

torch.exp(torch.tensor(pd.Series(up).values))

tensor([7.6395e-249, 1.9184e-241, 3.6856e-234, 5.4180e-227, 6.0933e-220,
        5.2437e-213, 3.4526e-206, 1.7393e-199, 6.7039e-193, 1.9769e-186,
        4.4607e-180, 7.7007e-174, 1.0172e-167, 1.0280e-161, 7.9493e-156,
        4.7028e-150, 2.1288e-144, 7.3728e-139, 1.9538e-133, 3.9613e-128,
        6.1447e-123, 7.2932e-118, 6.6227e-113, 4.6016e-108, 2.4463e-103,
         9.9501e-99,  3.0965e-94,  7.3730e-90,  1.3432e-85,  1.8722e-81,
         1.9967e-77,  1.6293e-73,  1.0172e-69,  4.8588e-66,  1.7758e-62,
         4.9658e-59,  1.0625e-55,  1.7392e-52,  2.1783e-49,  2.0875e-46,
         1.5306e-43,  8.5864e-41,  3.6854e-38,  1.2103e-35,  3.0411e-33,
         5.8465e-31,  8.5998e-29,  9.6786e-27,  8.3341e-25,  5.4907e-23,
         2.7677e-21,  1.0675e-19,  3.1500e-18,  7.1120e-17,  1.2286e-15,
         1.6238e-14,  1.6421e-13,  1.2706e-12,  7.5216e-12,  3.4069e-11,
         1.1807e-10,  3.1306e-10,  6.3512e-10,  9.8585e-10,  1.1708e-09,
         1.0639e-09,  7.3967e-10,  3.9346e-10,  1.6014e-10,  4.9867e-11,
         1.1881e-11,  2.1659e-12,  3.0209e-13,  3.2237e-14,  2.6322e-15,
         1.6444e-16,  7.8598e-18,  2.8744e-19,  8.0428e-21,  1.7219e-22,
         2.8205e-24,  3.5348e-26,  3.3896e-28,  2.4869e-30,  1.3960e-32,
         5.9956e-35,  1.9702e-37,  4.9536e-40,  9.5292e-43,  1.4026e-45,
         1.5795e-48,  1.3609e-51,  8.9718e-55,  4.5254e-58,  1.7464e-61,
         5.1568e-65,  1.1650e-68,  2.0139e-72,  2.6634e-76,  2.6951e-80],
       dtype=torch.float64)

q = dist.Normal(loc = 3., scale=.4)
c = 1e-1

q.log_prob(theta_range) + torch.log(torch.tensor(c))

tensor([-52.3052, -50.3054, -48.3465, -46.4283, -44.5509, -42.7144, -40.9186,
        -39.1637, -37.4496, -35.7763, -34.1438, -32.5521, -31.0013, -29.4912,
        -28.0220, -26.5936, -25.2060, -23.8592, -22.5532, -21.2880, -20.0636,
        -18.8801, -17.7373, -16.6354, -15.5743, -14.5540, -13.5745, -12.6358,
        -11.7379, -10.8809, -10.0646,  -9.2892,  -8.5546,  -7.8608,  -7.2078,
         -6.5956,  -6.0242,  -5.4937,  -5.0039,  -4.5550,  -4.1469,  -3.7796,
         -3.4531,  -3.1674,  -2.9225,  -2.7185,  -2.5552,  -2.4328,  -2.3511,
         -2.3103,  -2.3103,  -2.3511,  -2.4328,  -2.5552,  -2.7185,  -2.9225,
         -3.1674,  -3.4531,  -3.7796,  -4.1469,  -4.5550,  -5.0039,  -5.4937,
         -6.0242,  -6.5956,  -7.2078,  -7.8608,  -8.5546,  -9.2892, -10.0646,
        -10.8809, -11.7379, -12.6358, -13.5745, -14.5540, -15.5743, -16.6354,
        -17.7373, -18.8801, -20.0636, -21.2880, -22.5532, -23.8592, -25.2059,
        -26.5936, -28.0220, -29.4912, -31.0013, -32.5521, -34.1438, -35.7763,
        -37.4496, -39.1637, -40.9186, -42.7144, -44.5509, -46.4283, -48.3464,
        -50.3054, -52.3052])

plt.plot(theta_range, q.log_prob(theta_range)+torch.log(torch.tensor(c)), label='q(theta)')
pd.Series(up).plot(ax=plt.gca(), label='Unnorm Posterior')
plt.legend()

<matplotlib.legend.Legend at 0x1471fb7c0>

plt.plot(theta_range, q.log_prob(theta_range)+torch.log(torch.tensor(c)), label='q(theta)')

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

s = q.sample()
s

tensor(3.0192)

q_s = q.log_prob(s).exp()
q_s

tensor(0.9962)

u = dist.Uniform(low=0., high=q_s).sample()

u

tensor(0.8261)

log_q_s = q.log_prob(s).item()

log_p_bar_s = unnorm_post(s)

log_u = torch.log(u).item()

log_q_s, log_p_bar_s, log_u

(-0.0037949681282043457, -48.340909004211426, -0.1910155713558197)

if log_u > log_p_bar_s:
    pass

considered = []
for i in range(10000):
    s = q.sample()
    q_s = q.log_prob(s).exp()

    u = c*dist.Uniform(low=0., high=q_s).sample()
    log_q_s = q.log_prob(s).item()

    log_p_bar_s = unnorm_post(s)
    log_u = torch.log(u).item()
    if log_u > log_p_bar_s:
        pass
    else:
        considered.append(u.item())
        
    

import seaborn as sns
sns.kdeplot(considered)

<AxesSubplot:>