Figures for my talk on MCMC for data analysis

README.md

You'll need numpy, matplotlib, emcee, triangle.py and acor to run this example.

Here's the model:

line.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import (division, print_function, absolute_import,
                        unicode_literals)

import numpy as np
import matplotlib.pyplot as pl
from matplotlib.ticker import MaxNLocator
import emcee
import triangle

np.random.seed(678)


def lnprobfn(p, x, y, yerr):
    m, b, d2 = p
    if not 0 <= d2 <= 1:
        return -np.inf
    ivar = 1.0 / (yerr ** 2 + d2)
    chi2 = np.sum((y - m * x - b) ** 2 * ivar)
    return -0.5 * (chi2 - np.sum(np.log(ivar)))


# Generate some fake linear data with underestimated error bars.
m0, b0, d0 = 0.56, 5.2, 0.25
x = np.sort(10 * np.random.rand(20))
yerr = 0.2 + 0.5 * np.random.rand(len(x))
y = m0 * x + b0 + np.sqrt(yerr * yerr + d0 * d0) * np.random.randn(len(x))

# Plot the initial dataset.
pl.errorbar(x, y, yerr=yerr, fmt=".k")
pl.plot(x, m0 * x + b0, ":k")
pl.xlabel("$x$")
pl.ylabel("$y$")
pl.savefig("data.png", transparent=True, dpi=300)

# Initialize the walkers.
nwalkers = 100
p0 = [[np.random.rand(), 10 * np.random.rand(), np.random.rand()]
      for k in range(nwalkers)]

# Set up the sampler.
sampler = emcee.EnsembleSampler(nwalkers, 3, lnprobfn, args=(x, y, yerr))

# Inherit the current random state so that we always get the same results.
sampler.random_state = np.random.get_state()

# Run MCMC.
sampler.run_mcmc(p0, 1100)

# Plot the results projected into the data space.
xs = np.array([0, 10])
samples = sampler.chain[:, 100:, :].reshape([-1, 3])
for m, b, d2 in samples[np.random.randint(len(samples), size=24)]:
    ys = m * xs + b
    pl.fill_between(xs, ys + np.sqrt(d2), ys - np.sqrt(d2), color="k",
                    alpha=0.1)
pl.savefig("results.png", transparent=True, dpi=300)

# Plot the corner plot.
labels = ["$m$", "$b$", "$\delta^2$"]
triangle.corner(samples, labels=labels,
                truths=[m0, b0, d0 * d0], quantiles=[0.16, 0.5, 0.84])
pl.savefig("corner.png", transparent=True, dpi=300)

# Plot the time series.
fig1 = pl.figure(figsize=[8, 8])
fig2 = pl.figure(figsize=[8, 8])
for i in range(samples.shape[-1]):
    ax = fig1.add_subplot(3, 1, i + 1)
    ax.plot(sampler.chain[:, :, i].T, "k", alpha=0.3)
    ax.set_xlim(0, sampler.chain.shape[1])
    if i < 2:
        ax.set_xticklabels([])
    else:
        ax.set_xlabel("Step Number")
    ax.set_ylabel(labels[i])
    ax.yaxis.set_major_locator(MaxNLocator(4))

    ax = fig2.add_subplot(3, 1, i + 1)
    ax.plot(sampler.chain[:, :, i].T, "k", alpha=0.6)
    ax.set_xlim(0, 200)
    if i < 2:
        ax.set_xticklabels([])
    else:
        ax.set_xlabel("Step Number")
    ax.set_ylabel(labels[i])
    ax.yaxis.set_major_locator(MaxNLocator(4))

fig1.savefig("time.png", transparent=True, dpi=300)
fig2.savefig("burnin.png", transparent=True, dpi=300)

# Compute and plot the autocorrelation function.
import acor
f = [acor.function(np.mean(sampler.chain[:, 100:, i], axis=0))
     for i in range(3)]
print([acor.acor(np.mean(sampler.chain[:, 100:, i], axis=0))[0]
       for i in range(3)])
ax = pl.figure().add_subplot(111)
for i in range(3):
    ax.plot(f[i][:200], "k", lw=2)
ax.axhline(0, color="k")
ax.set_xlim(0, 200)
ax.set_xlabel(r"$\tau$")
ax.set_ylabel(r"Autocorrelation")
ax.figure.savefig("acor.png", transparent=True, dpi=300)

model.png