LinearModel.py

File contents

#!/usr/bin/env python

import emcee
import pymc
import numpy as np
from emceeBridge import emceeBridge
from matplotlib import pyplot as plt

# Choose the "true" parameters.
m_true = -0.9594
b_true = 4.294
f_true = 0.534

# Generate some synthetic data from the model.
N = 50
x = np.sort(10*np.random.rand(N))
yerr = 0.1+0.5*np.random.rand(N)
y = m_true*x+b_true
y += np.abs(f_true*y) * np.random.randn(N)
y += yerr * np.random.randn(N)

fig1 = plt.figure()
ax1 = fig1.add_subplot(111)
ax1.errorbar(x, y, yerr=yerr, fmt='o', capsize=0, color='k')
ax1.set_xlabel('x', fontsize=16)
ax1.set_ylabel('y', fontsize=16)

# Generate the pymc model
with pymc.Model() as model:
   m = pymc.Uniform('m', -5.0, 0.5)
   b = pymc.Uniform('b', 0.0, 10.0)
   lnf = pymc.Uniform('lnf', -10.0, 1.0)
   mod = pymc.Deterministic('mod',  m*x + b)

   yobs = pymc.Normal('yobs', 
                      mu=mod,
                      sd=pymc.sqrt(yerr**2 + mod**2*pymc.exp(2*lnf)),
                      observed=y)

# Find the MAP as a starting point
with model:
   start = pymc.find_MAP()

# Create the emcee bridge:
br = emceeBridge(model, start)

# Now, we convert from the dict representation of pymc to array needed for
# emcee and initialize the walkers:
start = br.toarray(start)
ndim,nwalkers = br.dim,100
pos = [start + 1e-4*np.random.randn(ndim) for i in range(nwalkers)]

# set up the sampler, and pass the bridge as the lnprob function:
sampler = emcee.EnsembleSampler(nwalkers, ndim, br)

sampler.run_mcmc(pos, 500)

# Have a look at the samples
fig2 = plt.figure()
true_vals = {'m':m_true, 'b':b_true, 'lnf':np.log(f_true)}
for i in range(3):
   ax = fig2.add_subplot(3,1,i)
   ax.plot(sampler.chain[:,:,i].T, color='k', alpha=0.1)
   ax.set_ylabel(br.varnames[i])
   ax.axhline(true_vals[br.varnames[i]], color='grey')
fig2.savefig('ensemble_plot.png')

# Throw out the first 200 as burn-in
samples = sampler.chain[:,100:,:].reshape((-1,ndim))

# Plot out 100 sample realizations from the chain
xx = np.array([x.min(), x.max()])
for i in np.random.randint(len(samples), size=100):
   # extract the array solution and convert to dict rep
   sol = br.todict(samples[i])
   ax1.plot(xx, sol['b'] + sol['m']*xx, '-', color='k', alpha=0.1)
# Plot the true solution
ax1.plot(xx, b_true + m_true*xx, '-', color='red')
fig1.savefig('solution.png')

try:
   import triangle
   labels = [{'m':'$m$', 'b':'$b$', 'lnf':'$\ln\,f$'}[var] \
         for var in br.varnames]
   truths = [true_vals[var] for var in br.varnames]
   fig = triangle.corner(samples, labels=labels, truths=truths)
   fig.savefig("triangle.png")
except:
   print "Sorry, to get the triangle plot, you need the triangle.py module"