# + [markdown] colab_type="text" id="GF3kfGD4eFWH"
# *Author*: <NAME>
#
# [Adapted from <NAME>'s book in the chapter on MCMC (chapter 24.2)](https://github.com/probml/pmtk3/blob/master/pmtk/examples/isingImageDenoiseDemo.m). 

# + [markdown] colab_type="text" id="5q83umLwvRfU"
# # Gibbs sampling in an Ising model
#
# Suppose we have a noisy image that we want to denoise. We assume that underlying non-noisy image (the hidden image) is binary, and that the image statistics are defined by pairwise potentials; the Ising model.
#
# We write the image as a vector $\mathbf{x}$. We assume that each element of $\mathbf{x}$ can take the values +1 and -1. We use an Ising model for the underlying field:
#
# $$\log p(\mathbf{x}) = k + \sum_{s \in nbr(t)} \log \phi(x_s, x_t)$$
#
# Where k is an unknown normalization constant, $nbr(\mathbf{x_s})$ is the 4 neighborhood of a point, and $\log \phi(x_s, x_t)$ is a potential:
#
# $$\log \phi(x_s, x_t) = J x_s x_t$$
#
# For simplicity we'll assume $J = 1$. We then assume an iid Gaussian observation model:
#
# $$\mathbf{y} \sim \text{Normal}(\mathbf{x}, \sigma^2)$$
#
# To estimate the underlying model parameters $\mathbf{x}$, proceed by Gibbs sampling. We sample from the conditional distribution of the $t$th pixel given every other parameter and observation. We exploit the conditional independence properties of the model as follows:
#
# The first important thing to note is that this distribution is independent of every observation but the one for the pixel in question:
#
# $$p(x_t|\mathbf{x_{\setminus t}}, \mathbf{y}, J, \sigma) \equiv p(x_t|\mathbf{x_{\setminus t}}, y_t, J, \sigma)$$
#

# + [markdown] colab_type="text" id="-mkDTAJhq_jB"
#
# That means that for the purpose of conditional sampling, all the other observations don't matter; only the underlying hidden variables matter. The second thing to note is that not all hidden variables matter:
#
# $$p(x_t|\mathbf{x_{\setminus t}}, y_t, J, \sigma) \equiv \prod_{s\in nbr(t)} p(x_t|x_{s}, y_t, J, \sigma)$$
#
# So we only need to look at neighbors of the pixel we're sampling from. What is this probability then? By Bayes' theorem:
#
# $$\prod_{s\in nbr(t)} p(x_t|x_{s}, y_t, J, \sigma) \propto p(y_t|x_t, \sigma) \prod_{s\in nbr(t)} p(x_t|x_{s}, J)$$
#
# Let's normalize this quantity to obtain the conditional probability, which we'll call $q(x_t = +1)$ for short:
#
# $$q(x_t = +1) = \frac{p(y_t|x_t = +1, \sigma) \prod_{s\in nbr(t)} p(x_t = +1|x_{s}, J)}{\sum_{v\in(-1,+1)} p(y_t|x_t = v, \sigma) \prod_{s\in nbr(t)} p(x_n = v|x_{s}, J)}$$
#
# Let's call $\eta_t$ the sum of the values of the neighbors of $n$. Then we have:
#
# $$q(x_t = +1) = \frac{\exp \left(-\frac{1}{2\sigma^2} (y_t - 1) ^ 2 + J\eta_t \right)}{\sum_{v\in(-1,+1)} \exp \left(-\frac{1}{2\sigma^2} (y_t - v) ^ 2 + vJ\eta_t \right)}$$
#

# + [markdown] colab_type="text" id="kqEh1CIvt7Da"
#
# Finally, if we divide by the numerator, we obtain a familiar sigmoid function:
#
# $$q(x_t = +1) = \text{sigmoid} \left(  2J\eta_t + \frac{2 y_t}{\sigma^2} \right)$$
#
# Sanity checks: 
#
#   * If the observation $y_t$ is positive, it's more likely that the hidden pixel $x_t$ is positive. 
#   * If most of its neighbours are positive, it's more likely that the hidden pixel $x_t$ is positive.
#   * As the coupling strength $J$ becomes larger, $x$ is more likely to follow its neighbors
#   * If the noise is large, we're less likely to follow the neighbours.
#   
# ## Gibbs sampling.
#
# Gibbs sampling simply consists of:
#
#   * For N outer iterations:
#     * For k inner iterations corresponding to pixel t
#       *  sample from the conditional distribution $q(x_t)$
#       
# If the chain mixes well, the samples we obtian will be samples from the model posterior, from which we can calculate, e.g. the mean to obtain our best estimate of the image.
#
# Let's start by loading the data and then apply Gibbs sampling to denoise it.

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}} colab_type="code" id="I2EzAs_8vMH8"
import numpy as np
import StringIO
import matplotlib.pyplot as plt
from PIL import Image

path = 'https://raw.githubusercontent.com/probml/pmtk3/master/data/letterA/letterA.bmp'
sigma = 2

with urllib.urlopen(path) as f:
  im = np.asarray(Image.open(StringIO.StringIO(f.read())))

# Create the image and a noisy version of it.
im = ((im > 128).astype(np.float32) - 0.5) * 2.0
im_noisy = im + sigma*np.random.randn(im.shape[0], im.shape[1])

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 7059, "status": "ok", "timestamp": 1474048628348, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="8O2_4oMkzviJ" outputId="89c60e04-7ec5-4ce0-e40a-55c73a90037b"
plt.subplot(121)
plt.imshow(im, vmin = -1, vmax=1)
plt.title('Original image')
plt.subplot(122)
plt.imshow(im_noisy, vmin = -1, vmax=1)
plt.title('Noisy image')


# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 1659, "status": "ok", "timestamp": 1474048636542, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="gHD03Vlni4Up" outputId="14be4061-6baa-43be-e962-18223d68112e"
def gibbs_sampler(im, J, sigma, niter):
  # Start somewhere reasonable. 
  # Random would work but would be slower.
  x = ((im > 0).astype(np.float32)) * 2 - 1
  samples = [x.copy()]
  for k in range(niter):
    for i in range(x.shape[0]):
      for j in range(x.shape[1]):
        # Here (i, j) is isomorphic to t
        # Select the neighbors.
        eta = 0
        if i > 0:
          eta += x[i - 1, j]
        if i < x.shape[0] - 1:
          eta += x[i + 1, j]
        if j > 0:
          eta += x[i, j - 1]
        if j < x.shape[1] - 1:
          eta += x[i, j + 1]
        # Compute the energy
        gamma = 2 * J * eta + 2 * im[i, j] / sigma ** 2
        q = 1 / (1 + np.exp(- gamma))
        # And sample from the conditional distribution
        x[i, j] = (np.random.rand() < q) * 2 - 1
    samples.append(x.copy())
  return samples

im_samples = gibbs_sampler(im_noisy, 1, sigma, 2)
plt.subplot(131)
plt.imshow(im_samples[0], vmin=-1, vmax=1)
plt.title('sample 0')
plt.subplot(132)
plt.imshow(im_samples[1], vmin=-1, vmax=1)
plt.title('sample 1')
plt.subplot(133)
plt.imshow(im_samples[2], vmin=-1, vmax=1)
plt.title('sample 2')

# + [markdown] colab_type="text" id="52RZDZlfxWpj"
# We see that very rapidly, the model is moving towards a pretty good guess for the underlying image. Let's see how many iterations is takes to start to settle around the mode of the posterior.

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}} colab_type="code" id="LZMmQ3kUxqav"
im_samples = gibbs_sampler(im_noisy, 1, sigma, 100)
mse = ((im_samples - im.reshape((1, im.shape[0], im.shape[1])))**2).sum(2).sum(1)

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 18645, "status": "ok", "timestamp": 1474048666517, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="hrXnQcJmyfOR" outputId="670510f6-ba35-4dcf-c82c-51b1f4be4102"
plt.plot(mse)
plt.title('MSE of posterior sample vs. time')
plt.xlabel('# iterations')
plt.ylim((0, 3000))

# + [markdown] colab_type="text" id="d-RmejCCyqJ0"
# If I was to eyeball it, I would say that the sampler starts to converge around 20 iterations. It's conventional with MCMC to forget the initial transient samples to focus on the steady state samples after a *burn-in* period. To get an estimate of the posterior of the model, nothing is simpler - just average the samples after the burnin period.

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 2314, "status": "ok", "timestamp": 1474048793501, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPT<KEY>AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="NU1L5TkLzN0O" outputId="6aa5e1bb-bda0-44c4-c4fd-ac269aa8d904"
plt.figure(figsize=(12, 4))
burnin = 50
im_denoised = np.array(im_samples[burnin:]).mean(0)
plt.subplot(141)
plt.imshow(im, vmin = -1, vmax=1)
plt.title('Original image')
plt.subplot(142)
plt.imshow(im_noisy, vmin = -1, vmax=1)
plt.title('Noisy image')
plt.subplot(143)
plt.imshow(im_samples[burnin], vmin = -1, vmax=1)
plt.title('One steady-state sample')
plt.subplot(144)
plt.imshow(im_denoised, vmin = -1, vmax=1)
plt.title('Denoised image')

# + [markdown] colab_type="text" id="hBA6Va92I0Gl"
# Let's look at the delta between the guessed image and the true image.

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 879, "status": "ok", "timestamp": 1474048803645, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="IF4iOZZlEOGN" outputId="6a7fdcd3-be74-4760-b813-43e755b7f05c"
plt.imshow(np.array(im_samples)[burnin:,:,:].mean(0) - im)


# + [markdown] colab_type="text" id="Sb0N9WJE1EcO"
# ![magic](http://i.imgur.com/YsbKHg1.gif)

# + [markdown] colab_type="text" id="VFI1X4o21Nxw"
# # À trous blocked Gibbs sampler
#
# That's all well and good, but the sampler is super slow. We could push everything to C, which would definitely help - [Gibbs samplers are notoriously poorly adapted to high-level languages](https://darrenjw.wordpress.com/2011/07/16/gibbs-sampler-in-various-languages-revisited/). 
#
# However, we can take advantage of the conditional independence properties of the graph to speed up the sampling tremendously while keeping in good old Python. The neighborhood potential implies that if two (hidden) pixels lie in each other's neighborhoods, sampling from each other can't be done independently. However, the converse is also true. Recall this equation:
#
# $$p(x_t|\mathbf{x_{\setminus t}}, y_t, J, \sigma) \equiv \prod_{s\in nbr(t)} p(x_t|x_{s}, y_t, J, \sigma)$$
#
# If you squint you will notice that you can sample from the conditional distributions of two pixels *simultaneously* as long as they lie in different neighbordhoods, because these conditional distributions are independent of each other.
#
# In fact, you can sample from 1/4 of the pixels simultaneously, as long as you leave 2 pixels between each hidden variable. You sample from pixel (0, 0), pixel (2, 0), pixel (0, 2), pixel (2, 2), etc. Then you can repeat this 4 times with different offsets. I call this an *à trous* sampler, because of its similarity to an *à trous* or strided convolution.
#
# ![a trous](http://i.stack.imgur.com/f2RiP.gif)

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 1304, "status": "ok", "timestamp": 1474049531176, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="CFthtnV95qrM" outputId="257584d1-2371-4a72-9deb-3fcf671fbf99"
def gibbs_sampler_a_trous(im, J, sigma, niter, stride = 2):
  # Pad the image to avoid edge problems
  im = np.pad(im, (1, 1), 'constant')
  x = np.sign(im)
  samples = [x.copy()]
  for k in range(niter):
    for i in range(stride):
      for j in range(stride):
        # Here (i, j) is isomorphic to t
        # Select the neighbors.
        islice_l = slice(i,     im.shape[0] - 1 + i - stride, 2)
        islice_0 = slice(i + 1, im.shape[0]     + i - stride, 2)
        islice_r = slice(i + 2, im.shape[0] + 1 + i - stride, 2)
        
        jslice_l = slice(j,     im.shape[1] - 1 + j - stride, 2)
        jslice_0 = slice(j + 1, im.shape[1]     + j - stride, 2)
        jslice_r = slice(j + 2, im.shape[1] + 1 + j - stride, 2)
        
        etas = (x[islice_l, jslice_0] + 
                x[islice_r, jslice_0] + 
                x[islice_0, jslice_l] + 
                x[islice_0, jslice_r])
        
        # Compute the energy
        gamma = 2 * J * etas + 2 * im[islice_0, jslice_0] / sigma ** 2
        q = 1 / (1 + np.exp(- gamma))
        
        # And sample from the conditional distribution
        x[islice_0, jslice_0] = np.sign(q - np.random.rand(q.shape[0], q.shape[1]))
        assert np.all(x[:, 0] == 0)
    samples.append(x.copy())
  return np.array(samples)[:,1:-1,1:-1]

im_samples = gibbs_sampler_a_trous(im_noisy, 1, sigma, 2)
plt.subplot(131)
plt.imshow(im_samples[0], vmin=-1, vmax=1)
plt.title('sample 0')
plt.subplot(132)
plt.imshow(im_samples[1], vmin=-1, vmax=1)
plt.title('sample 1')
plt.subplot(133)
plt.imshow(im_samples[2], vmin=-1, vmax=1)
plt.title('sample 2')

# + [markdown] colab_type="text" id="lv2DlwfP5kx9"
# The old Gibbs sampler took ~ 250 ms per iteration.

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}]} colab_type="code" executionInfo={"elapsed": 2765, "status": "ok", "timestamp": 1474049325375, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="2GvECOHK5dJz" outputId="bdde5c70-1e8c-45a6-fdfe-b75fb7a8c559"
# %time im_samples = gibbs_sampler(im_noisy, 1, sigma, 10)

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}]} colab_type="code" executionInfo={"elapsed": 3172, "status": "ok", "timestamp": 1474049540484, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="DAUQx4qQ-ZRi" outputId="4c1148e4-9200-4035-fac9-d9fb5a9f6d35"
# %time im_samples = gibbs_sampler_a_trous(im_noisy, 1, sigma, 10)

# + [markdown] colab_type="text" id="tetJrqbg-hJ6"
# The block sampler is 100X faster than the old one. 

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}} colab_type="code" id="7wd0tTYd59Es"
im_samples = gibbs_sampler_a_trous(im_noisy, 1, sigma, 100)
mse = ((im_samples - im.reshape((1, im.shape[0], im.shape[1])))**2).sum(2).sum(1)

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 747, "status": "ok", "timestamp": 1474049358176, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="tdKQy2LX-6rE" outputId="10bc90f7-38fe-4ce7-8f8e-e4ec5044914e"
plt.plot(mse)
plt.title('MSE of posterior sample vs. time')
plt.xlabel('# iterations')
plt.ylim((0, 3000))

# + [markdown] colab_type="text" id="kLI_aTHfJFXX"
# It converges quite acceptably compared to the pixel-by-pixel solution - in Monto Carlo parlance, the chain is mixing pretty well.

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 6482, "status": "ok", "timestamp": 1474049442555, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="byXZdcV1_2Sx" outputId="fe537075-6dfd-4a33-f231-f400c4d97cf2"
plt.figure(figsize=(12, 4))
burnin = 50
im_denoised = np.array(im_samples[burnin:]).mean(0)
plt.subplot(141)
plt.imshow(im, vmin = -1, vmax=1)
plt.title('Original image')
plt.subplot(142)
plt.imshow(im_noisy, vmin = -1, vmax=1)
plt.title('Noisy image')
plt.subplot(143)
plt.imshow(im_samples[burnin], vmin = -1, vmax=1)
plt.title('One steady-state sample')
plt.subplot(144)
plt.imshow(im_denoised, vmin = -1, vmax=1)
plt.title('Denoised image')

# + cellView colab={"autoexec": {"startup": false, "wait_interval": 0}, "output_extras": [{"item_id": 1}, {"item_id": 2}]} colab_type="code" executionInfo={"elapsed": 1122, "status": "ok", "timestamp": 1474049476191, "user": {"displayName": "<NAME>", "photoUrl": "//lh3.googleusercontent.com/-CdPTZM1TPbU/AAAAAAAAAAI/AAAAAAAAABQ/MOmbIGKtlQA/s50-c-k-no/photo.jpg", "userId": "103444865736961301998"}, "user_tz": 420} id="bNAZdtwpMfxg" outputId="62b81b10-3e20-4c34-f932-77eeb3b969f4"
plt.imshow(im_samples[burnin:,:,:].mean(0) - im)

# + [markdown] colab_type="text" id="ofJoJTpRMhRQ"
# Beautiful!

StarcoderdataJupyterScriptsDedupFiltered