2 parameter gibbs sampler with pretty plots

gistfile1.py

import random
import math
import matplotlib.pyplot as plt
import numpy as np

def acf(x, length=35):
  return np.array([1] + [np.corrcoef(x[:-i], x[i:])[0,1] for i in range(1,length)])

# The number of samples to observe
n = 100

# The gibbs parameters
M = 1000 # sample size
burn_in = 100
acf_window = 35
acf_xvals = [i for i in range(1,acf_window+1)]

# Choose initial guesses for parameters
theta = 0.
phi = 0.

# Gibbs sampler
draws = np.zeros((M,2)) # results matrix
for i in range(M):
	# Sample theta
	theta = random.gammavariate(4, 1./(3.+phi)) # Note: b is 1/b in python
	# Sample phi
	phi = random.gammavariate(3, 1./(2.+theta)) # Note: b is 1/b in python
	# Record samples
	draws[i] = np.array([theta,phi])

# Markov chains + marginal posterior
names = ['theta', 'phi']
x_vals = np.array([i for i in range(M)])
for i in range(len(names)):
	plt.subplot(len(names),3,i*3+1)
	plt.subplots_adjust(hspace=0.5)
	plt.plot(x_vals, draws[:,i], color='black')
	plt.axvline(burn_in, color='blue')
	plt.axhline(np.mean(draws[:,i]), color='red')
	plt.title('{0} Samples'.format(names[i]))
	plt.setp(plt.xticks()[1], rotation=90)
	plt.subplot(len(names),3,i*3+2)
	plt.subplots_adjust(hspace=0.5)
	plt.title('{0} Autocorrelation'.format(names[i]))
	plt.xlabel('Lag')
	plt.setp(plt.xticks()[1], rotation=90)
	plt.plot(acf_xvals, acf(draws[:,i], acf_window))
	plt.subplot(len(names),3,i*3+3)
	plt.subplots_adjust(hspace=0.5)
	plt.title('f({0})'.format(names[i]))
	plt.hist(draws[:,i])
	plt.setp(plt.xticks()[1], rotation=90)

plt.savefig('hw3_results.pdf')
plt.clf()