def grid(theta, x, y, yerr, periods, lhf, fname):
L = np.zeros_like(periods)
results = np.zeros((len(theta), len(L)))
for i, p in enumerate(periods):
print 'minimising'
result = fmin(lhf, theta, args=(x, y, yerr, p))
L[i] = -lhf(result, x, y, yerr, p)
results[:, i] = result
l = L==max(L)
new_theta = results[:, l].T[0]
period = periods[l][0]
print 'max likelihood = ', L[l][0], 'period = ', period
print 'best params: ', new_theta
plt.clf()
plt.subplot(2,1,1)
plt.plot(periods, L, color=ocols[1])
plt.xlabel('$\mathrm{Period~(days)}$')
plt.ylabel('$\mathrm{Log~likelihood}$')
plt.savefig('%s_likelihood' % fname)
plt.subplot(2,1,2)
plt.errorbar(x, y, yerr=yerr, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x), max(x), 100)
mu, cov = predict(new_theta, xs, x, y, yerr, period)
plt.plot(xs, mu, color=ocols[0], label="$\mathrm{Period}=%s$" % period)
plt.xlabel('$\mathrm{Time~(days)}$')
plt.ylabel('$\mathrm{RV~(ms}_{-1}\mathrm{)}$')
plt.subplots_adjust(hspace=.2)
plt.legend()
plt.savefig('%s_result' % fname)
return L, results, L[l][0], new_theta, period
def grid(theta, x, y, yerr, periods, lhf, fname):
L = np.zeros_like(periods)
results = np.zeros((len(theta), len(L)))
for i, p in enumerate(periods):
print 'minimising'
result = fmin(lhf, theta, args=(x, y, yerr, p))
L[i] = -lhf(result, x, y, yerr, p)
results[:, i] = result
l = L == max(L)
new_theta = results[:, l].T[0]
period = periods[l][0]
print 'max likelihood = ', L[l][0], 'period = ', period
print 'best params: ', new_theta
plt.clf()
plt.subplot(2, 1, 1)
plt.plot(periods, L, color=ocols[1])
plt.xlabel('$\mathrm{Period~(days)}$')
plt.ylabel('$\mathrm{Log~likelihood}$')
plt.savefig('%s_likelihood' % fname)
plt.subplot(2, 1, 2)
plt.errorbar(x, y, yerr=yerr, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x), max(x), 100)
mu, cov = predict(new_theta, xs, x, y, yerr, period)
plt.plot(xs, mu, color=ocols[0], label="$\mathrm{Period}=%s$" % period)
plt.xlabel('$\mathrm{Time~(days)}$')
plt.ylabel('$\mathrm{RV~(ms}_{-1}\mathrm{)}$')
plt.subplots_adjust(hspace=.2)
plt.legend()
plt.savefig('%s_result' % fname)
return L, results, L[l][0], new_theta, period
def MCMC(theta, x1, x2, x3, x4, x5, y1, y2, y3, y4, y5, yerr1, yerr2, yerr3, \
yerr4, yerr5, fname, burn_in, nsteps, nruns):
# setup sampler
nwalkers, ndim = 32, len(theta)
p0 = [theta+1e-4*np.random.rand(ndim) for i in range(nwalkers)]
args = [x1, x2, x3, x4, x5, y1, y2, y3, y4, y5, yerr1, yerr2, yerr3, yerr4, yerr5]
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=args)
print("Burning in...")
p0, lp, state = sampler.run_mcmc(p0, burn_in)
sampler.reset()
for i in range(nruns):
print 'Running... ', i
p0, lp, state = sampler.run_mcmc(p0, nsteps)
# results
samples = sampler.chain[:, 50:, :].reshape((-1, ndim))
mcmc_result = map(lambda v: (v[1], v[2]-v[1], v[1]-v[0]),
zip(*np.percentile(samples, [16, 50, 84], axis=0)))
mres = np.array(mcmc_result)[:, 0]
print 'mcmc_result = ', np.exp(mres)
np.savetxt("parameters_%s.txt" % fname, np.array(mcmc_result))
print "saving samples"
f = h5py.File("samples%s" % fname, "w")
data = f.create_dataset("samples", np.shape(sampler.chain))
data[:,:] = np.array(sampler.chain)
f.close()
# make triangle plot
fig_labels = ["$A$", "$l1$", "$l2$", "$wn1$", "$wn2$", "$wn3$", "$wn4$", "$wn5$", "$P$"]
fig = triangle.corner(samples, truths=mres, labels=fig_labels)
fig.savefig("triangle_%s.png" % fname)
# plot results
new_theta, period = np.exp(mres), np.exp(mres[-1])
plt.clf()
plt.subplot(4,1,1)
plt.errorbar(x1, y1, yerr=yerr1, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x1), max(x1), 100)
pars = np.array([new_theta[0], new_theta[1], new_theta[2], new_theta[3], period])
mu, cov = predict(pars, xs, x1, y1, yerr1, period)
plt.plot(xs, mu, color=ocols.lightblue)
plt.subplot(4,1,2)
plt.errorbar(x2, y2, yerr=yerr2, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x2), max(x2), 100)
pars = np.array([new_theta[0], new_theta[1], new_theta[2], new_theta[4], period])
mu, cov = predict(pars, xs, x2, y2, yerr2, period)
plt.plot(xs, mu, color=ocols.orange)
plt.subplot(4,1,3)
plt.errorbar(x3, y3, yerr=yerr3, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x3), max(x3), 100)
pars = np.array([new_theta[0], new_theta[1], new_theta[2], new_theta[5], period])
mu, cov = predict(pars, xs, x3, y3, yerr3, period)
plt.plot(xs, mu, color=ocols.pink)
plt.subplot(4,1,4)
plt.errorbar(x4, y4, yerr=yerr4, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x4), max(x4), 100)
pars = np.array([new_theta[0], new_theta[1], new_theta[2], new_theta[6], period])
mu, cov = predict(pars, xs, x4, y4, yerr4, period)
plt.plot(xs, mu, color=ocols.green)
plt.savefig('all_data_results_emcee_wasp')
def MCMC(theta, x, y, yerr, fname, burn_in, nsteps, nruns):
# setup sampler
nwalkers, ndim = 32, len(theta)
p0 = [theta + 1e-4 * np.random.rand(ndim) for i in range(nwalkers)]
args = [x, y, yerr]
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=args)
print("Burning in...")
p0, lp, state = sampler.run_mcmc(p0, burn_in)
sampler.reset()
for i in range(nruns):
print 'Running... ', i
p0, lp, state = sampler.run_mcmc(p0, nsteps)
# results
samples = sampler.chain[:, 50:, :].reshape((-1, ndim))
mcmc_result = map(lambda v: (v[1], v[2] - v[1], v[1] - v[0]),
zip(*np.percentile(samples, [16, 50, 84], axis=0)))
mres = np.array(mcmc_result)[:, 0]
print 'mcmc_result = ', np.exp(mres)
np.savetxt("parameters_%s.txt" % fname, np.array(mcmc_result))
print "saving samples"
f = h5py.File("samples%s" % fname, "w")
data = f.create_dataset("samples", np.shape(sampler.chain))
data[:, :] = np.array(sampler.chain)
f.close()
# make triangle plot
fig_labels = [
"$A$", "$l1$", "$l2$", "$wn1$", "$wn2$", "$wn3$", "$wn4$", "$P$"
]
fig = triangle.corner(samples, truths=mres, labels=fig_labels)
fig.savefig("triangle_%s.png" % fname)
# plot results
new_theta, period = np.exp(mres), np.exp(mres[-1])
plt.clf()
plt.subplot(4, 1, 1)
plt.errorbar(x1, y1, yerr=yerr1, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x1), max(x1), 100)
pars = np.array(
[new_theta[0], new_theta[1], new_theta[2], new_theta[3], period])
mu, cov = predict(pars, xs, x1, y1, yerr1, period)
plt.plot(xs, mu, color=ocols.lightblue)
plt.subplot(4, 1, 2)
plt.errorbar(x2, y2, yerr=yerr2, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x2), max(x2), 100)
pars = np.array(
[new_theta[0], new_theta[1], new_theta[2], new_theta[4], period])
mu, cov = predict(pars, xs, x2, y2, yerr2, period)
plt.plot(xs, mu, color=ocols.orange)
plt.subplot(4, 1, 3)
plt.errorbar(x3, y3, yerr=yerr3, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x3), max(x3), 100)
pars = np.array(
[new_theta[0], new_theta[1], new_theta[2], new_theta[5], period])
mu, cov = predict(pars, xs, x3, y3, yerr3, period)
plt.plot(xs, mu, color=ocols.pink)
plt.subplot(4, 1, 4)
plt.errorbar(x4, y4, yerr=yerr4, fmt='k.', capsize=0, ecolor='.8')
xs = np.linspace(min(x4), max(x4), 100)
pars = np.array(
[new_theta[0], new_theta[1], new_theta[2], new_theta[6], period])
mu, cov = predict(pars, xs, x4, y4, yerr4, period)
plt.plot(xs, mu, color=ocols.green)
plt.savefig('all_data_results_emcee')
