Code from Python lunch at MPIA

data.dat

   0.29012    4.69497    0.58980
   0.39687    4.53831    0.61233
   0.58085    4.58511    0.72283
   0.79724    3.88045    0.58649
   0.96492    3.28931    0.59474
   1.44055    2.63219    0.52092
   2.00743    1.76821    0.57815
   2.71773    0.46981    0.64042
   2.73360    0.32956    0.50058
   3.04199   -0.22867    0.63828
   3.40496   -0.54015    0.64521
   3.78026   -0.40934    0.67248
   3.83754   -1.41567    0.71468
   4.12843   -1.92919    0.76586
   4.17218   -2.47955    0.52952
   4.23041   -2.05766    0.56216
   4.23079   -1.98897    0.58624
   4.29481   -1.80686    0.61830
   4.56428   -1.82784    0.65145
   4.67689   -3.04755    0.64531

livedemo.py

import matplotlib.pyplot as pl
import emcee
import numpy as np


def chi2(p, x, y, yerr):
    m, b = p
    return np.sum(((y - (m * x + b)) / yerr) ** 2)


data = np.array([line.split() for line in open("data.dat")], dtype=float)
m_true, b_true = np.array(open("truth.dat").read().split(), dtype=float)

x, y, yerr = data[:, 0], data[:, 1], data[:, 2]
xgrid = np.array([0, 5])

pl.errorbar(x, y, yerr=yerr, fmt="ok")
pl.plot(xgrid, m_true * xgrid + b_true, "--r")

import scipy.optimize as op

p0 = [1, 1]

m, b = op.fmin(chi2, p0, args=[x, y, yerr])

pl.errorbar(x, y, yerr=yerr, fmt="ok")
pl.plot(xgrid, m_true * xgrid + b_true, "--r")
pl.plot(xgrid, m * xgrid + b, "g")


def lnlike(p, x, y, yerr):
    return -0.5 * chi2(p, x, y, yerr) \
            - 0.5 * np.sum(np.log(2 * np.pi * yerr ** 2))


nwalkers, ndim = 100, 2
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnlike, args=[x, y, yerr])

initial = [np.random.rand(ndim) for k in range(nwalkers)]

pos, lnprob, state = sampler.run_mcmc(initial, 100)

sampler.reset()

ret = sampler.run_mcmc(pos, 500)

chain = sampler.flatchain
m, b = chain[:, 0], chain[:, 1]

pl.plot(xgrid, m[None, :] * xgrid[:, None] + b[None, :], color="#444444",
        alpha=0.03)
pl.errorbar(x, y, yerr=yerr, fmt="ok")
pl.plot(xgrid, m_true * xgrid + b_true, "--r")

truth.dat

-1.6 5.2