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')