def GP(kernel, kernel_params, white=False):
'''
'''
if kernel == 'Basic':
w, a, t = kernel_params
if white:
if OLDGEORGE:
return george.GP(
WhiteKernel(w**2) + a**2 * Matern32Kernel(t**2))
else:
return george.GP(a**2 * Matern32Kernel(t**2),
white_noise=np.log(w**2),
fit_white_noise=True)
else:
return george.GP(a**2 * Matern32Kernel(t**2))
elif kernel == 'QuasiPeriodic':
w, a, g, p = kernel_params
if white:
if OLDGEORGE:
return george.GP(
WhiteKernel(w**2) + a**2 * ExpSine2Kernel(g, p))
else:
return george.GP(a**2 * ExpSine2Kernel(g, p),
white_noise=np.log(w**2),
fit_white_noise=True)
else:
return george.GP(a**2 * ExpSine2Kernel(g, p))
else:
raise ValueError('Invalid value for `kernel`.')
def GP(kernel, kernel_params, white = False):
if kernel == 'Basic':
w, a, t = kernel_params
if white:
return george.GP(WhiteKernel(w ** 2) + a ** 2 * Matern32Kernel(t ** 2))
else:
return george.GP(a ** 2 * Matern32Kernel(t ** 2))
elif kernel == 'QuasiPeriodic':
w, a, g, p = kernel_params
if white:
return george.GP(WhiteKernel(w ** 2) + a ** 2 * ExpSine2Kernel(g, p))
else:
return george.GP(a ** 2 * ExpSine2Kernel(g, p))
else:
raise ValueError('Invalid value for `kernel`.')
def make_plots(id, RESULTS_DIR="/Users/ruthangus/projects/GProtation/code/" \
"results_acfprior_03_10"):
"""
Make a plot of the fit to the light curve and the posteriors.
"""
""" load lc """
x, y = load_suzanne_lcs(id)
yerr = np.ones(len(y)) * 1e-5
m = x < 100
x, y, yerr = x[m], y[m], yerr[m]
""" load posteriors """
fn = os.path.join(RESULTS_DIR, "{}.h5".format(id))
df = pd.read_hdf(fn, key="samples")
""" find medians """
theta = [np.median(df.iloc[:, i]) for i in range(5)]
""" fit GP """
print(np.exp(theta[-1]), "period")
k = theta[0] * ExpSquaredKernel(theta[1]) \
* ExpSine2Kernel(theta[2], theta[4]) + WhiteKernel(theta[3])
gp = george.GP(k, solver=george.HODLRSolver)
gp.compute(x - x[0], yerr)
xs = np.linspace((x - x[0])[0], (x - x[0])[-1], 1000)
mu, cov = gp.predict(y, xs)
""" plot fit """
plt.clf()
plt.plot(x, y, "k.")
plt.plot(xs, mu)
plt.xlim(0, 100)
# v = np.std(y)
# plt.ylim(-10*v, 10*v)
plt.savefig("{}_fit".format(id))
def gp_kernel(self, theta):
A = np.exp(theta[0])
l = np.exp(theta[1])
G = np.exp(theta[2])
sigma = np.exp(theta[3])
P = np.exp(theta[4])
return A * ExpSquaredKernel(l) * ExpSine2Kernel(G, P) + WhiteKernel(sigma)
def MCMC(theta, x, y, yerr, fname, burn_in, nsteps, nruns):
# calculate initial likelihood and plot initial hparams
xs = np.linspace(min(x), max(x), 1000)
k = theta[0] * ExpSquaredKernel(theta[1]) * ExpSine2Kernel(theta[2], theta[4])
k += WhiteKernel(theta[3])
gp = george.GP(k)
print 'initial lnlike = ', lnlike(theta, x, y, yerr)
mu, cov = predict(theta, xs, x, y, yerr)
plt.clf()
plt.errorbar(x, y, yerr=yerr, fmt='k.', capsize=0)
plt.plot(xs, mu, 'r')
std = np.sqrt(np.diag(cov))
# plt.fill_between(mu-std, mu+std, color='r', alpha='.5')
plt.savefig('%s_init' % fname)
# 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$", "$wn$", "$P$"]
fig = triangle.corner(samples, truths=mres, labels=fig_labels)
fig.savefig("triangle_%s.png" % fname)
# plot result
mu, cov = predict(mres, xs, x, y, yerr)
plt.clf()
plt.errorbar(x, y, yerr=yerr, fmt='k.', capsize=0)
plt.plot(xs, mu, 'r')
plt.savefig('%s_final' % fname)
def lnlike(theta, x, y, yerr):
theta = np.exp(theta)
k = theta[0] * ExpSquaredKernel(theta[1]) \
* ExpSine2Kernel(theta[2], theta[4]) + WhiteKernel(theta[3])
gp = george.GP(k, solver=george.HODLRSolver)
try:
gp.compute(x, np.sqrt(theta[3]+yerr**2))
except (ValueError, np.linalg.LinAlgError):
return 10e25
return gp.lnlikelihood(y, quiet=True)
def plot_lc(koi):
"""
Make demo plot of a light curve.
"""
# Load the data
print(LC_DIR)
x, y, yerr = kd.load_kepler_data(LC_DIR)
x -= x[0]
m = x < 500
x, y, yerr = x[m], y[m], yerr[m]
# Load the posterior samples.
df = pd.read_hdf(os.path.join(DATA_DIR, "KOI-{}.h5".format(int(koi))),
key="samples")
a = np.exp(MAP(df.ln_A.values))
l = np.exp(MAP(df.ln_l.values))
g = np.exp(MAP(df.ln_G.values))
s = np.exp(MAP(df.ln_sigma.values))
p = np.exp(MAP(df.ln_period.values))
print("ln(a) = ", np.log(a), "ln(l) = ", np.log(l), "ln(G) = ", np.log(g),
"ln(s) = ", np.log(s), "ln(p) = ", np.log(p), "p = ", p)
xs = np.linspace(min(x), max(x), 500)
k = a * ExpSquaredKernel(l) \
* ExpSine2Kernel(g, p) + WhiteKernel(s)
gp = george.GP(k)
gp.compute(x, yerr)
mu, cov = gp.predict(y, xs)
plt.clf()
plt.plot(x, y, "k.")
plt.plot(xs, mu, color="CornFlowerBlue")
plt.xlabel("$\mathrm{Time~(days)}$")
plt.ylabel("$\mathrm{Normalised~flux}$")
plt.subplots_adjust(left=.18)
plt.savefig(os.path.join(FIG_DIR, "koi_lc_demo.pdf"))
def gp_kernel(self, theta):
A = np.exp(theta[0])
l = np.exp(theta[1])
sigma = np.exp(theta[2])
P = np.exp(theta[3])
return A * ExpSquaredKernel(l) * CosineKernel(P) + WhiteKernel(sigma)
def make_plot(sampler, xb, yb, yerrb, ID, RESULTS_DIR, trths, traces=False,
tri=False, prediction=True):
nwalkers, nsteps, ndims = np.shape(sampler)
flat = np.reshape(sampler, (nwalkers * nsteps, ndims))
mcmc_res = map(lambda v: (v[1], v[2]-v[1], v[1]-v[0]),
zip(*np.percentile(flat, [16, 50, 84], axis=0)))
med = np.concatenate([np.array(mcmc_res[i]) for i in
range(len(mcmc_res))])
print("median values = ", med[::3])
logprob = flat[:, -1]
ml = logprob == max(logprob)
maxlike = flat[np.where(ml)[0][0], :][:-1]
print("max like = ", maxlike)
print("\n", np.exp(np.array(maxlike[-1])), "period (days)", "\n")
r = np.concatenate((maxlike, med))
# calculate autocorrelation times
acorr_t = emcee.autocorr.integrated_time(flat)
# save data
df = pd.DataFrame({"N": [ID], "A_max": [r[0]], "l_max": [r[1]],
"gamma_max": [r[2]], "period_max": [r[3]],
"sigma_max": [r[4]], "A": [r[5]], "A_errp": [r[6]],
"A_errm": [r[7]], "l": [r[8]], "l_errp": [r[9]],
"l_errm": [r[10]], "gamma": [r[11]],
"gamma_errp": [r[12]], "gamma_errm": [r[13]],
"sigma": [r[14]], "sigma_errp": [r[15]],
"sigma_errm": [r[16]], "period": [r[17]],
"period_errp": [r[18]], "period_errm": [r[19]],
"acorr_A": acorr_t[0], "acorr_l": acorr_t[1],
"acorr_gamma": acorr_t[2], "acorr_sigma": acorr_t[3],
"acorr_period": acorr_t[4]})
df.to_csv(os.path.join(RESULTS_DIR, "{0}_mcmc_results.txt".format(ID)))
fig_labels = ["ln(A)", "ln(l)", "ln(G)", "ln(s)", "ln(P)", "lnprob"]
if traces:
print("Plotting traces")
for i in range(ndims):
plt.clf()
plt.plot(sampler[:, :, i].T, 'k-', alpha=0.3)
plt.ylabel(fig_labels[i])
plt.savefig(os.path.join(RESULTS_DIR, "{0}_{1}.png".format(ID,
fig_labels[i])))
if tri:
DIR = "../code/simulations/kepler_diffrot_full/par/"
truths = pd.read_csv(os.path.join(DIR, "final_table.txt"),
delimiter=" ")
true_p = np.log(truths.P_MIN.values[truths.N.values ==
int(filter(str.isdigit, ID))][0])
trths[-1] = true_p
print("Making triangle plot")
fig = corner.corner(flat[:, :-1], labels=fig_labels,
quantiles=[.16, .5, .84], show_titles=True,
truths=trths)
fig.savefig(os.path.join(RESULTS_DIR, "{0}_triangle".format(ID)))
print(os.path.join("{0}_triangle.png".format(ID)))
if prediction:
if len(xb) > 1: # if the data is a list of lists.
try:
x = [i for j in xb for i in j]
y = [i for j in yb for i in j]
yerr = [i for j in yerrb for i in j]
except: # if the data are just a single list.
TypeError
x, y, yerr = xb, yb, yerrb
else: # if the data is a list of a single list.
x, y, yerr = xb[0], yb[0], yerrb[0]
print("plotting prediction")
theta = np.exp(np.array(maxlike))
k = theta[0] * ExpSquaredKernel(theta[1]) \
* ExpSine2Kernel(theta[2], theta[4]) + WhiteKernel(theta[3])
gp = george.GP(k, solver=george.HODLRSolver)
gp.compute(x-x[0], yerr)
xs = np.linspace((x-x[0])[0], (x-x[0])[-1], 1000)
mu, cov = gp.predict(y, xs)
plt.clf()
plt.errorbar(x-x[0], y, yerr=yerr, fmt="k.", capsize=0)
plt.xlabel("Time (days)")
plt.ylabel("Normalised Flux")
plt.plot(xs, mu, color='#0066CC')
plt.xlim(min(x-x[0]), max(x-x[0]))
plt.savefig(os.path.join(RESULTS_DIR, "{0}_prediction".format(ID)))
print(os.path.join(RESULTS_DIR, "{0}_prediction.png".format(ID)))
return r