def Fit0_Jitter(x, y, yerr = None, verbose = True, doPlot = False, \
xpred = None):
k = terms.Matern32Term(log_sigma=0.0, log_rho=0.0)
if yerr is None:
wn = np.median(abs(np.diff(y)))
else:
wn = np.median(yerr)
k += terms.JitterTerm(log_sigma=np.log(wn))
gp = GP(k, mean=1.0)
gp.compute(x)
HP_init = gp.get_parameter_vector()
soln = minimize(NLL0, gp.get_parameter_vector(), jac=True, args=(gp, y))
gp.set_parameter_vector(soln.x)
if verbose:
print 'Initial pars:', HP_init
print 'Fitted pars:', soln.x
if xpred is None:
return soln.x
mu, var = gp.predict(y, xpred, return_var=True)
std = np.sqrt(var)
if doPlot:
plt.errorbar(x, y, yerr=yerr, fmt=".k", capsize=0)
plt.plot(xpred, mu, 'C0')
plt.fill_between(xpred,
mu + std,
mu - std,
color='C0',
alpha=0.4,
lw=0)
return soln.x, mu, std
def Pred1_2D(par, x2d, y2d, y2derr, doPlot=True, x2dpred=None):
K = x2d.shape[0]
k = terms.Matern32Term(log_sigma=par[-2], log_rho=par[-1])
gp = GP(k, mean=1.0)
shifts = np.append(0, par[:K - 1])
x1d = (x2d + shifts[:, None]).flatten()
inds = np.argsort(x1d)
y1d = y2d.flatten()
y1derr = y2derr.flatten()
gp.compute(x1d[inds], yerr=y1derr[inds])
if x2dpred is None:
x2dpred = np.copy(x2d)
x2dpreds = (x2dpred + shifts[:, None])
y2dpred = np.zeros_like(x2dpred)
y2dprederr = np.zeros_like(x2dpred)
for k in range(K):
print 'Prediction for spectrum %d' % (k + 1)
x1dpred = x2dpreds[k, :].flatten()
indspred = np.argsort(x1dpred)
mu, var = gp.predict(y1d[inds], x1dpred[indspred], return_var=True)
std = np.sqrt(var)
y2dpred[k, indspred] = mu
y2dprederr[k, indspred] = std
if doPlot:
for i in range(K):
plt.errorbar(x2d[i, :],
y2d[i, :] - i,
yerr=y2derr[i, :],
fmt=".k",
capsize=0,
alpha=0.5)
plt.plot(x2dpred[i, :], y2dpred[i, :] - i, 'C0')
plt.fill_between(x2dpred[i,:], y2dpred[i,:] + y2dprederr[i,:] - i, \
y2dpred[i,:] - y2dprederr[i,:] - i, color = 'C0', alpha = 0.4, lw = 0)
return x2dpred, y2dpred, y2dprederr
def Fit0(x, y, yerr, verbose = True, doPlot = False, \
xpred = None, HP_init = None):
if HP_init is None:
HP_init = np.zeros(2)
k = terms.Matern32Term(log_sigma=HP_init[0], log_rho=HP_init[1])
gp = GP(k, mean=1.0)
gp.compute(x, yerr=yerr)
soln = minimize(NLL0, HP_init, jac=True, args=(gp, y))
gp.set_parameter_vector(soln.x)
if verbose:
print 'Initial pars:', HP_init
print 'Fitted pars:', soln.x
if xpred is None:
return soln.x
mu, var = gp.predict(y, xpred, return_var=True)
std = np.sqrt(var)
if doPlot:
plt.errorbar(x, y, yerr=yerr, fmt=".k", capsize=0)
plt.plot(xpred, mu, 'C0')
plt.fill_between(xpred,
mu + std,
mu - std,
color='C0',
alpha=0.4,
lw=0)
return soln.x, mu, std
def baseline_hybrid_GP(*args):
x, y, yerr_w, xx, params, inst, key = args
kernel = terms.Matern32Term(log_sigma=1., log_rho=1.)
gp = celerite.GP(kernel, mean=np.nanmean(y))
gp.compute(x, yerr=yerr_w) #constrain on x/y/yerr
def neg_log_like(gp_params, y, gp):
gp.set_parameter_vector(gp_params)
return -gp.log_likelihood(y)
def grad_neg_log_like(gp_params, y, gp):
gp.set_parameter_vector(gp_params)
return -gp.grad_log_likelihood(y)[1]
initial_params = gp.get_parameter_vector()
bounds = gp.get_parameter_bounds()
soln = minimize(neg_log_like,
initial_params,
jac=grad_neg_log_like,
method="L-BFGS-B",
bounds=bounds,
args=(y, gp))
gp.set_parameter_vector(soln.x)
baseline = gp.predict(y, xx)[0] #constrain on x/y/yerr, evaluate on xx (!)
return baseline
def __init__(self, x, y, n_bg_coef, wave_err=0.05, # MAGIC NUMBER: wavelength error hack
log_sigma0=0., log_rho0=np.log(10.), # initial params for GP
x_shift=None):
self.x = np.array(x)
self.y = np.array(y)
if n_bg_coef >= 2:
a_kw = dict([('a{}'.format(i),0.) for i in range(2,n_bg_coef)])
# estimate background
a_kw['a1'] = (y[-1]-y[0])/(x[-1]-x[0]) # slope
a_kw['a0'] = y[-1] - a_kw['a1']*x[-1] # estimate constant term
else:
a_kw = dict(a0=np.mean([y[0], y[-1]]))
# initialize model
self.mean_model = MeanModel(n_bg_coef=n_bg_coef, **a_kw)
self.kernel = terms.Matern32Term(log_sigma=log_sigma0, log_rho=log_rho0)
# set up the gp
self.gp = GP(self.kernel, mean=self.mean_model, fit_mean=True)
self.gp.compute(x, yerr=wave_err)
logger.debug("Initial log-likelihood: {0}"
.format(self.gp.log_likelihood(y)))
if x_shift is None:
self.x_shift = 0.
else:
self.x_shift = x_shift
def init_gp(self,
log_sigma=None,
log_rho=None,
amp=None,
x0=None,
bg=None,
bg_buffer=None,
**kwargs):
"""
**kwargs:
"""
# Call the different get_init methods with the correct kwargs passed
sig = inspect.signature(self.get_init)
kw = OrderedDict()
for k in list(sig.parameters.keys()):
kw[k] = kwargs.pop(k, None)
p0 = self.get_init(**kw)
# Call the generic init method - all line models must have these params
p0_generic = self.get_init_generic(amp=amp,
x0=x0,
bg=bg,
bg_buffer=bg_buffer)
for k, v in p0_generic.items():
p0[k] = v
# expand bg parameters
bgs = p0.pop('bg_coef')
for i in range(self.n_bg_coef):
p0['bg{}'.format(i)] = bgs[i]
# transform
p0 = self.transform_pars(p0)
# initialize model
mean_model = self.MeanModel(n_bg_coef=self.n_bg_coef,
absorp_emiss=self.absorp_emiss,
**p0)
p0_gp = self.get_init_gp(log_sigma=log_sigma, log_rho=log_rho)
kernel = terms.Matern32Term(log_sigma=p0_gp['log_sigma'],
log_rho=p0_gp['log_rho'])
# set up the gp model
self.gp = GP(kernel, mean=mean_model, fit_mean=True)
if self._err is not None:
self.gp.compute(self.x, self._err)
else:
self.gp.compute(self.x)
init_params = self.gp.get_parameter_vector()
init_ll = self.gp.log_likelihood(self.flux)
logger.log(0, "Initial log-likelihood: {0}".format(init_ll))
return init_params
def Fit1(x2d, y2d, y2derr, par_in=None, verbose=True):
K = x2d.shape[0]
if par_in is None:
par_in = np.zeros(K + 1)
k = terms.Matern32Term(log_sigma=par_in[-2], log_rho=par_in[-1])
gp = GP(k, mean=1.0)
soln = minimize(NLL1, par_in, args=(gp, x2d, y2d, y2derr))
if verbose:
print 'Initial pars:', par_in
print 'Fitted pars:', soln.x
return soln.x
def call_gp(params):
log_sigma, log_rho, log_error_scale = params
if GP_CODE=='celerite':
kernel = terms.Matern32Term(log_sigma=log_sigma, log_rho=log_rho)
gp = celerite.GP(kernel, mean=MEAN, fit_mean=False) #log_white_noise=np.log(yerr),
gp.compute(xx, yerr=yyerr/err_norm*np.exp(log_error_scale))
return gp
elif GP_CODE=='george':
kernel = np.exp(log_sigma) * kernels.Matern32Kernel(log_rho)
gp = george.GP(kernel, mean=MEAN, fit_mean=False) #log_white_noise=np.log(yerr),
gp.compute(xx, yerr=yyerr/err_norm*np.exp(log_error_scale))
return gp
else:
raise ValueError('A bad thing happened.')
def call_gp(params):
log_sigma, log_rho, log_error_scale, contrast, rv, fwhm = params
if GP_CODE=='celerite':
mean_model = Model_celerite(**dict(**dict(Contrast=contrast, RV=rv, FWHM=fwhm)))
kernel = terms.Matern32Term(log_sigma=log_sigma, log_rho=log_rho)
gp = celerite.GP(kernel, mean=mean_model)
gp.compute(xx, yerr=yyerr/err_norm*np.exp(log_error_scale))
return gp
elif GP_CODE=='george':
mean_model = Model_george(**dict(**dict(Contrast=contrast, RV=rv, FWHM=fwhm)))
kernel = np.exp(log_sigma) * kernels.Matern32Kernel(log_rho)
gp = george.GP(kernel, mean=mean_model)
gp.compute(xx, yerr=yyerr/err_norm*np.exp(log_error_scale))
return gp
else:
raise ValueError('gp_code must be "celerite" or "george".')
def calculate_lnlike(params, inst, key):
#::: if no GP baseline sampling, then calculate lnlike per hand
if config.BASEMENT.settings['baseline_' + key + '_' + inst] != 'sample_GP':
model = calculate_model(params, inst, key)
yerr_w = calculate_yerr_w(params, inst, key)
baseline = calculate_baseline(params,
inst,
key,
model=model,
yerr_w=yerr_w)
residuals = config.BASEMENT.data[inst][key] - model - baseline
inv_sigma2_w = 1. / yerr_w**2
return -0.5 * (np.nansum((residuals)**2 * inv_sigma2_w -
np.log(inv_sigma2_w)))
#::: if GP baseline sampling, use the GP lnlike instead
else:
model = calculate_model(params, inst, key)
x = config.BASEMENT.data[inst]['time']
y = config.BASEMENT.data[inst][key] - model
yerr_w = calculate_yerr_w(params, inst, key)
kernel = terms.Matern32Term(
log_sigma=params['baseline_gp1_' + key + '_' + inst],
log_rho=params['baseline_gp2_' + key + '_' + inst])
gp = celerite.GP(kernel)
gp.compute(x, yerr=yerr_w)
lnlike = gp.log_likelihood(y)
# baseline2 = gp.predict(y, x)[0]
# plt.figure()
# plt.plot(x,y,'k.', color='grey')
# plt.plot(xx,baseline,'r-', lw=2)
# plt.plot(x,baseline2,'ro', lw=2)
# plt.title(inst+' '+key+' '+str(gp.get_parameter_vector()))
# plt.show()
# raw_input('press enter to continue')
return lnlike
def lnlike(params_in, y_obs):
if params_in[N_pixel+1]>0.0 and 0.02<params_in[N_pixel+2]<0.04 and 0.0<params_in[N_pixel+3]<time_array[-1] and -10.0<params_in[N_pixel+4]<11.4 and params_in[N_pixel+5]<0.0 and 85.0<params_in[N_pixel+6]<95.0 and 20.0<params_in[N_pixel+7]<30.0 and 0.0<params_in[N_pixel+8]<100.0:
print(params_in)
# Set up the GP model
params_mean_model = np.resize(params_in, int(N_pixel+4))
params_mean_model = np.append(params_mean_model, params_in[N_pixel+6])
params_mean_model = np.append(params_mean_model, params_in[N_pixel+7])
mean_model = MeanModel(*params_mean_model)
mean_model_out = mean_model.get_value()
kernel = terms.Matern32Term(log_sigma=params_in[N_pixel+5], log_rho=params_in[N_pixel+4], eps = 0.000001)
gp = celerite.GP(kernel, mean=0.0, fit_mean=False)
gp.compute(time_array, err_frac*params_in[N_pixel+8], check_sorted=True)
logprob = -np.sum(((y_obs-mean_model_out)/err_frac)**2.0/2.0) - log(err_frac)*len(time_array) - 0.5*len(time_array)*log(2.0*pi) + gp.log_likelihood(y_obs-mean_model_out)
#reduced_chisq = np.sum(((y_obs-mu)/err_frac)**2.0)/(len(y_obs)-len(params_in))
print(logprob)
return logprob
print("-inf")
return -np.inf
def baseline_sample_GP(*args):
x, y, yerr_w, xx, params, inst, key = args
kernel = terms.Matern32Term(
log_sigma=params['baseline_gp1_' + key + '_' + inst],
log_rho=params['baseline_gp2_' + key + '_' + inst])
gp = celerite.GP(kernel)
gp.compute(x, yerr=yerr_w)
baseline = gp.predict(y, xx)[0]
# baseline2 = gp.predict(y, x)[0]
# plt.figure()
# plt.plot(x,y,'k.', color='grey')
# plt.plot(xx,baseline,'r-', lw=2)
# plt.plot(x,baseline2,'ro', lw=2)
# plt.title(inst+' '+key+' '+str(gp.get_parameter_vector()))
# plt.show()
# raw_input('press enter to continue')
return baseline
def get_initial_guess_gp(datadir):
import celerite
from celerite import terms
from scipy.optimize import minimize
from .computer import update_params
config.init(datadir)
base = config.BASEMENT
params = update_params(base.theta_0)
for inst in base.settings['inst_phot']:
key = 'flux'
model = calculate_model(params, inst, key, xx=None)
x = base.data[inst]['time']
y = base.data[inst][key] - model
# yerr_weights = config.BASEMENT.data[inst]['err_scales_'+key]
yerr = np.nanstd(y) #overwrite yerr; works better for removing smooth global trends
kernel = terms.Matern32Term(log_sigma=1., log_rho=1.)
gp = celerite.GP(kernel, mean=np.nanmean(y))
gp.compute(x, yerr=yerr) #constrain on x/y/yerr
def neg_log_like(gp_params, y, gp):
gp.set_parameter_vector(gp_params)
return -gp.log_likelihood(y)
def grad_neg_log_like(gp_params, y, gp):
gp.set_parameter_vector(gp_params)
return -gp.grad_log_likelihood(y)[1]
initial_params = gp.get_parameter_vector()
bounds = gp.get_parameter_bounds()
soln = minimize(neg_log_like, initial_params, jac=grad_neg_log_like,
method="L-BFGS-B", bounds=bounds, args=(y, gp))
# gp.set_parameter_vector(soln.x)
inv_sigma2_w = calculate_inv_sigma2_w(params, inst, key)
print('baseline_gp1_'+key+'_'+inst + ':', soln.x[0])
print('baseline_gp2_'+key+'_'+inst + ':', soln.x[1])
print('inv_sigma2_'+key+'_'+inst + ':', np.nanmean(inv_sigma2_w))
def init_gp(self, log_amp=-5, log_tau=-3, log_sigma=-5):
bounds = [(-np.inf, np.inf), (0, np.inf), (0, 1), (0, np.inf), (0, 1),
(0, 1), (0, 1)]
t0, p, k, r, b, q1, q2 = [
self.init_params.get(i) for i in 't0,p,k,r,b,q1,q2'.split(',')
]
mean_model = TransitMeanModel(t0=t0,
p=p,
k=k,
r=r,
b=b,
q1=q1,
q2=q2,
bounds=bounds)
kernel = terms.Matern32Term(log_sigma=log_amp,
log_rho=log_tau,
bounds=[(-15, 5), (-15, 5)])
kernel += terms.JitterTerm(log_sigma=log_sigma, bounds=[(-10, 0)])
self.gp = celerite.GP(kernel, mean=mean_model, fit_mean=True)
self.gp.compute(self.t)
def fit_lc(time, mag,mag_err, ellc_params, ellc_bounds, ellc_priors, plot_model = True, minimize_lc=True, emcee_lc=False, emcee_draws = 100,emcee_chain_file = 'emcee_chain.fits', GP=True, free_params = ['t_zero', 'k', 'radius_1','b'],not_included_params=[], return_handler=False, return_model=False, gp_kernel = terms.Matern32Term(log_sigma=-2, log_rho = 2), gp_kernel_bounds=dict(), gp_kernel_freeze=[]):
if GP:
##########################################################################
# Ok so this is a bit tricky. We will do the following to please celerite
#
# 1. Create the model and
# 1. Define a subset of ellc_params called ellc_GP_params which are just
# floats and integeres (exclude None and string entries).
##########################################################################
ellc_GP_params,not_included_params,ellc_gp_bounds = get_ellc_GP_params(ellc_params, ellc_bounds, free_params)
lc_model = ellc_GP_wrap(**ellc_GP_params, bounds = ellc_gp_bounds)# not_included_params, ellc_gp_bounds)
lc_model.ellc_GP_params, lc_model.not_included_params, lc_model.ellc_gp_bounds, lc_model.ellc_gp_priors = ellc_GP_params,not_included_params,ellc_gp_bounds,ellc_priors
gp = celerite.GP(gp_kernel, mean=lc_model, fit_mean=True)
gp.compute(time, yerr = mag_err)
lglike = gp.log_likelihood(mag)
########################################
# Freeze parameter we do ot wish to fit
########################################
all_parameters = np.array(tuple(ellc_GP_params.keys()))
params_to_freeze = np.setdiff1d(all_parameters, free_params)
for i in params_to_freeze:
gp.freeze_parameter('mean:{}'.format(i))
for i in gp_kernel_freeze:
gp.freeze_parameter('kernel:{}'.format(i))
if return_handler:
gp.get_param = types.MethodType(get_param, gp)
gp.calculate_log_like_prior = types.MethodType(calculate_log_like_prior, gp)
return gp
if plot_model:
plt.scatter(time, mag, c='k', s= 10, alpha=0.5)
plt.plot(time, lc_model.get_lc(time), 'r', alpha=0.8, label='Mean model')
# Plot the variance with upsampled time series
mu, var = gp.predict(mag, np.linspace(min(time), max(time), \
len(time)*10), return_var=True)
std = np.sqrt(abs(var))
color = "#ff7f0e"
plt.fill_between( np.linspace(min(time), max(time), len(time)*10), \
mu+std, \
mu-std, color=color, alpha=0.8, edgecolor="none", label = 'GP model')
plt.legend()
plt.gca().invert_yaxis()
plt.ylabel('Mag')
plt.xlabel('Time')
print('~'*80)
print('Initial fit')
print('~'*80)
print('Loglike: {:.2f}'.format(lglike))
print('Chi**2: {:.2f}'.format(-2*lglike))
print('Reduced Chi**2: {:.2f}'.format(-2*lglike / (len(time) - len(free_params) ) ))
print('BIC: {:.2f}'.format(len(free_params)*np.log(len(time)) - 2*np.log(lglike)))
print('~'*80)
plt.show()
return np.linspace(min(time), max(time), len(time)*10), mu, var
if minimize_lc:
plt.scatter(time, mag, c='k', s= 10, alpha=0.5)
plt.plot(time, lc_model.get_lc(time), 'g', alpha=0.8, label='Old')
#plt.gca().invert_yaxis()
plt.ylabel('Mag')
plt.xlabel('Time')
print('Minimisation of lightcurve with GP')
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print("{:>20}: {:.3f}".format('Initial log-like',gp.log_likelihood(mag)))
print('\n\n')
print('~'*80)
print("Minization using Scipy's L-BFGS-R algorithm with GP")
print('~'*80)
print('\nRunning L-BFGS-s algorithm . . .', end='')
##################
# Now optimise
##################
initial_params = gp.get_parameter_vector()
bounds_ = gp.get_parameter_bounds()
soln = minimize(neg_log_like, initial_params, method="L-BFGS-B", bounds=bounds_, args=(mag, gp))
print(soln)
###################
# Print the output
###################
print("{:>20}: {:.3f} ({} iterations)".format('Final log-like', -soln.fun, soln.nit))
print('Chi**2: {:.2f}'.format(-2*-soln.fun))
print('Reduced Chi**2: {:.2f}'.format(-2*-soln.fun / (len(time) - len(free_params) ) ))
print('BIC: {:.2f}'.format(len(free_params)*np.log(len(time)) - 2*np.log(-soln.fun)))
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
#for i in range(len(gp.get_parameter_names())):
# print('{:.3f} {}'.format(soln.x[i], gp.get_parameter_names()[i]))
############################################################################
# Now get the best model as a fictionary so we can put it back in and plot
############################################################################
best_params = ellc_GP_params.copy()
best_params_passback = ellc_params.copy()
for i in range(len(gp.parameter_names)):
if gp.parameter_names[i][:5]=='mean:':
best_params[gp.parameter_names[i][5:]] = gp.parameter_vector[i]
best_params_passback[gp.parameter_names[i][5:]] = gp.parameter_vector[i]
#print('{:>12}: {:.3f}'.format(gp.parameter_names[i][5:], gp.parameter_vector[i]))
#########################
# Now get the best model
#########################
plt.cla()
# plot data
plt.scatter(time,mag,s=2, c='k', alpha = 0.3)
# Plot the best model (just to be sure)
mean_model_better = ellc_GP_wrap(**best_params, bounds = ellc_gp_bounds)
mean_model_better.ellc_GP_params, mean_model_better.not_included_params, mean_model_better.ellc_gp_bounds, mean_model_better.ellc_gp_priors = ellc_GP_params,not_included_params,ellc_gp_bounds,ellc_priors
plt.plot(time,mean_model_better.get_value(time), 'b--', linewidth=2, \
label='Model selected')
# Plot the variance with upsampled time series
mu, var = gp.predict(mag, np.linspace(min(time), max(time), \
len(time)*10), return_var=True)
std = np.sqrt(abs(var))
color = "#ff7f0e"
plt.fill_between( np.linspace(min(time), max(time), len(time)*10), \
mu+std, \
mu-std, color=color, alpha=0.8, edgecolor="none", label = 'Best fit')
plt.xlabel('BJD')
plt.ylabel('Mag')
plt.gca().invert_yaxis()
plt.legend()
plt.show()
return best_params_passback, gp
if emcee_lc==True:
plt.scatter(time, mag, c='k', s= 10, alpha=0.5)
plt.plot(time, lc_model.get_lc(time), 'g', alpha=0.8, label='Old')
#plt.gca().invert_yaxis()
plt.ylabel('Mag')
plt.xlabel('Time')
print('\n\n')
print('~'*80)
print("Posterior inference of lightcurve using emcee with GP")
print('~'*80)
print('\nRunning production chain')
initial = gp.get_parameter_vector()
ndim, nwalkers = len(initial), 32
sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability, args = (gp, mag, ellc_priors), threads = 8)
print("Running burn-in...")
p0 = initial + 1e-5 * np.random.randn(nwalkers, ndim)
width=30
start_time = time_.time()
i, result=[], []
for i, result in enumerate(sampler.sample(p0, iterations=emcee_draws)):
n = int((width+1) * float(i) / emcee_draws)
delta_t = time_.time()-start_time# time to do float(i) / n_steps % of caluculations
time_incr = delta_t/(float(i+1) / emcee_draws) # seconds per increment
time_left = time_incr*(1- float(i) / emcee_draws)
m, s = divmod(time_left, 60)
h, m = divmod(m, 60)
sys.stdout.write('\r[{0}{1}] {2}% - {3}h:{4}m:{5:.2f}s'.format('#' * n, ' ' * (width - n), 100*float(i) / emcee_draws ,h, m, s))
gp.compute(time, yerr = mag_err)
mu, var = gp.predict(mag, np.linspace(min(time), max(time), \
len(time)*10), return_var=True)
std = np.sqrt(abs(var))
color = "#ff7f0e"
plt.fill_between( np.linspace(min(time), max(time), len(time)*10), \
mu+std, \
mu-std, color=color, alpha=0.8, edgecolor="none", label = 'Best fit')
plt.xlabel('BJD')
plt.ylabel('Mag')
plt.gca().invert_yaxis()
plt.legend()
#plt.show()
#names = gp.get_parameter_names()
samples = sampler.chain[:, int(np.floor(emcee_draws*0.75)):, :].reshape((-1, ndim))
corner.corner(samples, labels = gp.get_parameter_names())
########################
# Now save to fits file
########################
t = Table(sampler.flatchain, names=gp.get_parameter_names())
t.add_column(Column(sampler.flatlnprobability,name='loglike'))
indices = np.mgrid[0:nwalkers,0:emcee_draws]
step = indices[1].flatten()
walker = indices[0].flatten()
t.add_column(Column(step,name='step'))
t.add_column(Column(walker,name='walker'))
try:
t.write(emcee_chain_file)
except:
os.remove(emcee_chain_file)
t.write(emcee_chain_file)
lglike = t['loglike'][np.argmax(np.array(t['loglike']))]
print('~'*80)
print('Initial fit')
print('~'*80)
print('Loglike: {:.2f}'.format(lglike))
print('Chi**2: {:.2f}'.format(-2*lglike))
print('Reduced Chi**2: {:.2f}'.format(-2*lglike / (len(time) - len(gp.get_parameter_names()) ) ))
print('BIC: {:.2f}'.format(len(gp.get_parameter_names())*np.log(len(time)) - 2*np.log(lglike)))
print('~'*80)
plt.show()
return
else:
# Initiate the model and get loglike
lc_model = ellc_wrap()
lc_model.set_param(ellc_params=ellc_params)
if return_handler:
return lc_model
lglike = lc_model.get_log_like_lc(time, mag, mag_err, ellc_params, ellc_bounds, ellc_priors)
if plot_model:
if return_model:
return lc_model.get_lc(time)
else:
# Make the plot
plt.scatter(time, mag, c='k', s= 10, alpha=0.5)
plt.plot(time, lc_model.get_lc(time), 'r', alpha=0.8)
plt.gca().invert_yaxis()
plt.ylabel('Mag')
plt.xlabel('Time')
# Return parameters where needed
print('~'*80)
print('Initial fit')
print('~'*80)
print('Loglike: {:.2f}'.format(lglike))
print('Chi**2: {:.2f}'.format(-2*lglike))
print('Reduced Chi**2: {:.2f}'.format(-2*lglike / (len(time) - len(free_params) ) ))
print('BIC: {:.2f}'.format(len(free_params)*np.log(len(time)) - 2*np.log(-lglike)))
print('~'*80)
plt.show()
return
if minimize_lc:
# Output
print('\n\n')
print('~'*80)
print("Minization using Scipy's L-BFGS-R algorithm")
print('~'*80)
print('\nRunning L-BFGS-s algorithm . . .', end='')
# Plot the initial data
plt.scatter(time, mag, c='k', s= 10, alpha=0.5)
plt.plot(time, lc_model.get_lc(time), 'g', alpha=0.8, label='Old')
plt.gca().invert_yaxis()
plt.ylabel('Mag')
plt.xlabel('Time')
# get param vector and bounds from free_params
param_vector = [ellc_params[i] for i in free_params]
bounds = [ellc_bounds[i] for i in free_params]
# Cals scipy minimize with the L-BFGS_B algorithm
soln = minimize(lc_model.scipy_minimize, param_vector, method="L-BFGS-B", bounds = bounds, args=(time, mag, mag_err,free_params,ellc_bounds, ellc_priors),options={'disp': True})
# Display output
if soln.success:
print(' Success in {} iterations'.format(soln.nit))
else:
print(' Failure in {} iterations'.format(soln.nit))
print('Old loglike: {:.2f}'.format(lglike))
print('new loglike: {:.2f}'.format(-soln.fun))
print('Chi**2: {:.2f}'.format(-2*-soln.fun))
print('Reduced Chi**2: {:.2f}'.format(-2*-soln.fun / (len(time) - len(free_params) ) ))
print('BIC: {:.2f}'.format(len(free_params)*np.log(len(time)) - 2*np.log(2*soln.fun)))
print('~'*70)
print('{:>12} {:>10} {:>10}'.format('Parameter', 'Old', 'New'))
for i in range(len(free_params)):
print('{:>12} {:>12.5f} {:>12.5f}'.format(free_params[i], param_vector[i], soln.x[i] ))
print('~'*70)
# Create the dictionary with best params for pass back
bett_params = ellc_params.copy()
for i in range(len(free_params)):
bett_params[free_params[i]] = soln.x[i]
# Now set and plot the best model
lc_model.set_param(ellc_params=bett_params)
bett_params = lc_model.get_param()
plt.plot(time, lc_model.get_lc(time), 'r', alpha=0.8, label='L-BFGS-R')
plt.legend()
plt.show()
return bett_params
if emcee_lc:
# output
print('\n\n')
print('~'*80)
print("Posterior inference of lightcurve using emcee")
print('~'*80)
print('\nRunning production chain')
# Plot data and initial model
plt.scatter(time, mag, c='k', s= 10, alpha=0.5)
plt.plot(time, lc_model.get_lc(time), 'g', alpha=0.8, label='Old')
plt.gca().invert_yaxis()
plt.ylabel('Mag')
plt.xlabel('Time')
# get param vector and bounds from free_params
param_vector = [ellc_params[i] for i in free_params]
ndim, nwalkers = len(param_vector), 4*len(param_vector)
# initialse start position with finite loglike value
p0 = []
while len(p0) != nwalkers:
p0_trial = np.random.normal(param_vector, 1e-6)
lnlike_trial = lc_model.emcee_sampler(p0_trial, time, mag, mag_err,free_params, ellc_bounds, ellc_priors)
if lnlike_trial!=-np.inf:
p0.append(p0_trial)
# Initiate the sampler
sampler = emcee.EnsembleSampler(nwalkers, ndim, lc_model.emcee_sampler, args = (time, mag, mag_err,free_params, ellc_bounds, ellc_priors), threads = 8)
# run the sampler with custom progress bar
width=30
start_time = time_.time()
i, result=[], []
for i, result in enumerate(sampler.sample(p0, iterations=emcee_draws)):
n = int((width+1) * float(i) / emcee_draws)
delta_t = time_.time()-start_time# time to do float(i) / n_steps % of caluculations
time_incr = delta_t/(float(i+1) / emcee_draws) # seconds per increment
time_left = time_incr*(1- float(i) / emcee_draws)
m, s = divmod(time_left, 60)
h, m = divmod(m, 60)
sys.stdout.write('\r[{0}{1}] {2:>5}% - {3}h:{4}m:{5:.2f}s'.format('#' * n, ' ' * (width - n), 100*float(i) / emcee_draws ,h, m, s))
# Make the table
print('Making the table')
t = Table(sampler.flatchain, names=free_params)
# Valculate quadratic limb darkening values if needed
if 'quad_ldc_1_1_' in free_params:
quad_ldc_1_1, quad_ldc_1_2 = (np.array(t['quad_ldc_1_1_'])+ np.array(t['quad_ldc_1_2_']))**2, 0.5*np.array(t['quad_ldc_1_1_'])*(np.array(t['quad_ldc_1_1_']) + np.array(t['quad_ldc_1_2_']))**(-1)
t.add_column(Column(quad_ldc_1_1,name='quad_ldc_1_1'))
t.add_column(Column(quad_ldc_1_2,name='quad_ldc_1_2'))
if 'quad_ldc_2_1_' in free_params:
quad_ldc_2_1, quad_ldc_2_2 = (np.array(t['quad_ldc_2_1_'])+ np.array(t['quad_ldc_2_2_']))**2, 0.5*np.array(t['quad_ldc_2_1_'])*(np.array(t['quad_ldc_2_1_']) + np.array(t['quad_ldc_2_2_']))**(-1)
t.add_column(Column(quad_ldc_2_1,name='quad_ldc_2_1'))
t.add_column(Column(quad_ldc_2_2,name='quad_ldc_2_2'))
# add others
t.add_column(Column(sampler.flatlnprobability,name='loglike'))
indices = np.mgrid[0:nwalkers,0:emcee_draws]
step = indices[1].flatten()
walker = indices[0].flatten()
t.add_column(Column(step,name='step'))
t.add_column(Column(walker,name='walker'))
try:
t.write(emcee_chain_file)
except:
os.remove(emcee_chain_file)
t.write(emcee_chain_file)
# Find and plot the best solution
print('Finding and plotting the best model')
best_index = np.argmax(t['loglike'])
bett_params = ellc_params.copy()
best_params = np.array(list(t[best_index]))[:-3]
for i in range(len(free_params)):
bett_params[free_params[i]] = best_params[i]
lc_model.set_param(ellc_params=bett_params)
bett_params = lc_model.get_param()
plt.plot(time, lc_model.get_lc(time), 'r', alpha=0.8, label='emcee')
plt.legend()
# Create a corner plot from the last 25% of sampler value
samples = sampler.chain[:, int(np.floor(emcee_draws*0.75)):, :].reshape((-1, ndim))
std_vals = np.std(samples, axis=0)
means_vals = np.mean(samples, axis=0)
corner_fig = corner.corner(samples, labels = free_params)
# Output the best values
print('~'*70)
print('{:>12} {:>10} {:>10}'.format('Parameter', 'Old', 'New'))
print('~'*70)
for i in range(len(free_params)):
print('{:>12} {:>12.5f} {:>12.5f} +- {:>12.5f}'.format(free_params[i], param_vector[i], means_vals[i],std_vals[i] ))
print('~'*70)
plt.show()
return bett_params
binned_flux = sum_points/num_points
plt.scatter((bin_edge[:-1]+bin_edge[1:])/2.0/24./60/60+SpitzerStartTime, binned_flux, s=5, zorder=1, c='blue')
plt.plot(time_array/24./60/60+SpitzerStartTime, transit_flux, color = 'r', linewidth=1.5)
#ylim([0.995, 1.005])
xlabel(r'time(BJD)', fontsize=30)
ylabel(r'fraction', fontsize=30)
plt.savefig(time+"_starTransit.pdf", bbox_inches='tight')
plt.show()
plt.close(fig2)
fig = plt.figure(figsize = (9,5), dpi=120)
plt.scatter(time_array, tot_flux, s=3, zorder=0)
# calculate 50 posterior light curves######################################
#samples2 = sampler.chain[:, 2:, :].reshape(-1, ndim)
for s in samples[np.random.randint(len(samples), size=30)]:
kernel = terms.Matern32Term(log_sigma=s[N_pixel+5], log_rho=s[N_pixel+4], eps = 0.000001)
# print(s)
s_in = np.resize(s, N_pixel+4)
s_in = np.append(s_in, s[N_pixel+6])
s_in = np.append(s_in, s[N_pixel+7])
mean_model = MeanModel(*s_in)
mean_model_out = mean_model.get_value()
gp = celerite.GP(kernel, mean=0.0, fit_mean=False)
gp.compute(time_array, err_frac*s[N_pixel+8], check_sorted=True)
# print(mean_model_out)
mu = gp.predict(tot_flux-mean_model_out, time_array, return_cov=False) + mean_model_out
plt.plot(time_array, mu, color="pink", alpha=0.3, linewidth = 1.0)
################################################################
#plt.plot(time_array, out_flux_final, color = 'r')
#ylim([0.98, 1.02])
plt.savefig(time+"_star.pdf", bbox_inches='tight')
def neo_update_kernel(theta, params):
gp = george.GP(mean=0.0, fit_mean=False, white_noise=jitt)
pass
from celerite import terms as cterms
# 2 or sp.log(10.) ?
T = {
'Constant': 1.**2,
'RealTerm': cterms.RealTerm(log_a=2., log_c=2.),
'ComplexTerm': cterms.ComplexTerm(log_a=2., log_b=2., log_c=2., log_d=2.),
'SHOTerm': cterms.SHOTerm(log_S0=2., log_Q=2., log_omega0=2.),
'Matern32Term': cterms.Matern32Term(log_sigma=2., log_rho=2.0),
'JitterTerm': cterms.JitterTerm(log_sigma=2.0)
}
def neo_term(terms):
t_out = T[terms[0][0]]
for f in range(len(terms[0])):
if f == 0:
pass
else:
t_out *= T[terms[0][f]]
for i in range(len(terms)):
if i == 0:
pass
log_c=np.log(1.102),
log_d=np.log(1.05),
),
cterms.SHOTerm(log_S0=np.log(1.1),
log_Q=np.log(0.1),
log_omega0=np.log(1.2)),
cterms.SHOTerm(log_S0=np.log(1.1),
log_Q=np.log(2.5),
log_omega0=np.log(1.2)),
cterms.SHOTerm(
log_S0=np.log(1.1), log_Q=np.log(2.5), log_omega0=np.log(1.2)) +
cterms.RealTerm(log_a=np.log(1.345), log_c=np.log(2.4)),
cterms.SHOTerm(
log_S0=np.log(1.1), log_Q=np.log(2.5), log_omega0=np.log(1.2)) *
cterms.RealTerm(log_a=np.log(1.345), log_c=np.log(2.4)),
cterms.Matern32Term(log_sigma=0.1, log_rho=0.4),
]
@pytest.mark.parametrize("oterm", test_terms)
@pytest.mark.parametrize("mean", [0.0, 10.5])
def test_consistency(oterm, mean, data):
x, diag, y, t = data
# Setup the original GP
original_gp = original_celerite.GP(oterm, mean=mean)
original_gp.compute(x, np.sqrt(diag))
# Setup the new GP
term = terms.OriginalCeleriteTerm(oterm)
gp = celerite2.GaussianProcess(term, mean=mean)
def fit_gaussian_process(lc,
objid,
passbands,
plot,
extrapolate,
bad_loglike_thresh=-2000):
print(f"Fitting GP to {objid}")
gp_lc = {}
if plot:
plt.figure()
kernel = terms.Matern32Term(log_sigma=5., log_rho=3.)
times, fluxes, fluxerrs = {}, {}, {}
for pbidx, pb in enumerate(passbands):
pbmask = lc['passband'] == pb
sortedidx = np.argsort(lc[pbmask]['time'].data)
times[pb] = lc[pbmask]['time'].data[sortedidx]
fluxes[pb] = lc[pbmask]['flux'].data[sortedidx]
fluxerrs[pb] = lc[pbmask]['fluxErr'].data[sortedidx]
try:
gp_lc[pb] = celerite.GP(kernel)
gp_lc[pb].compute(times[pb], fluxerrs[pb])
except Exception as e:
print("Failed object", objid, e)
return
# print("Initial log likelihood: {0}".format(gp_lc[pb].log_likelihood(fluxes[pb])))
initial_params = gp_lc[pb].get_parameter_vector(
) # This should be the same across passbands
bounds = gp_lc[pb].get_parameter_bounds()
# Optimise parameters
try:
r = minimize(combined_neg_log_like,
initial_params,
method="L-BFGS-B",
bounds=bounds,
args=(fluxes, gp_lc, passbands))
# print(r)
except Exception as e:
print("Failed object", objid, e)
return
for pbidx, pb in enumerate(passbands):
gp_lc[pb].set_parameter_vector(r.x)
time = times[pb]
flux = fluxes[pb]
fluxerr = fluxerrs[pb]
# print("Final log likelihood: {0}".format(gp_lc[pb].log_likelihood(flux)))
# Remove objects with bad fits
if extrapolate is True:
x = np.linspace(min(min(time), -70), max(max(time), 80), 5000)
else:
x = np.linspace(min(time), max(time), 5000)
pred_mean, pred_var = gp_lc[pb].predict(flux, x, return_var=True)
if np.any(~np.isfinite(pred_mean)
) or gp_lc[pb].log_likelihood(flux) < bad_loglike_thresh:
print("Bad fit for object", objid)
return
# Plot GP fit
if plot:
# Predict with GP
if extrapolate:
x = np.linspace(min(min(time), -70), max(max(time), 80), 5000)
else:
x = np.linspace(min(time), max(time), 5000)
pred_mean, pred_var = gp_lc[pb].predict(flux, x, return_var=True)
pred_std = np.sqrt(pred_var)
color = {
'g': 'tab:green',
'r': "tab:red",
'i': "tab:purple",
'z': "tab:brown"
}
# plt.plot(time, flux, "k", lw=1.5, alpha=0.3)
plt.errorbar(time,
flux,
yerr=fluxerr,
fmt=".",
capsize=0,
color=color[pb])
plt.plot(x, pred_mean, color=color[pb], label=pb)
plt.fill_between(x,
pred_mean + pred_std,
pred_mean - pred_std,
color=color[pb],
alpha=0.3,
edgecolor="none")
if plot:
plt.xlabel("Days since trigger", fontsize=15)
plt.ylabel("Flux", fontsize=15)
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)
plt.legend(fontsize=15)
plt.show()
# if extrapolate:
# plt.savefig(f'/Users/danmuth/PycharmProjects/transomaly/plots/gp_fits/extrapolated/gp_{objid}.pdf')
# else:
# plt.savefig(f'/Users/danmuth/PycharmProjects/transomaly/plots/gp_fits/gp_{objid}.pdf')
plt.close()
return gp_lc, objid
def fit_SB1_gp(time, rv,rv_err, initial_params, bounds, priors, flux_weighted=False, plot_guess=True,
freeze_parameters=['r1', 'k', 'sbratio', 'b', 'q', \
'Pdot', 'T_ld', 'M_ld',\
'Logg_ld', 'dV0'], emcee_fit=False, draws =1000):
print('\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print('RV fitting procedure with Gaussian Processes')
print('S. Gill ([email protected])')
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n')
mean_model = _SB1_Model(**initial_params, bounds=bounds)
time_upsampled = np.linspace(min(time), max(time), len(time) * 1000)
if plot_guess:
print('Plotting initial Guess. . .')
plt.plot(time_upsampled,mean_model.get_value(time_upsampled), 'r--', linewidth=2, \
label='Starting guess', alpha = 0.7)
plt.errorbar(time, rv, yerr=rv_err, fmt='ko')
plt.legend()
plt.show()
return
######################
# Initiate the kernel
######################
#kernel = (terms.RealTerm(log_a=-9.77, log_c=1.91) )
kernel = (terms.Matern32Term(log_sigma=-2, log_rho=2))
########################
# Initiate the GP model
########################
gp = celerite.GP(kernel, mean=mean_model, fit_mean=True)
########################################
# Freeze parameter we do ot wish to fit
########################################
for i in freeze_parameters:
gp.freeze_parameter('mean:{}'.format(i))
#############################################
# Compute and calculate the initial log like
#############################################
gp.compute(time, rv_err)
print('Minimisation of lightcurve with GP')
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print("{:>20}: {:.3f}".format('Initial log-like', gp.log_likelihood(rv)))
##################
# Now optimise
##################
initial_params = gp.get_parameter_vector()
bounds_ = gp.get_parameter_bounds()
soln = minimize(neg_log_like,
initial_params,
method="L-BFGS-B",
bounds=bounds_,
args=(rv, gp))
###################
# Print the output
###################
print("{:>20}: {:.3f} ({} iterations)".format('Final log-like', -soln.fun,
soln.nit))
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
#for i in range(len(gp.get_parameter_names())):
# print('{:.3f} {}'.format(soln.x[i], gp.get_parameter_names()[i]))
############################################################################
# Now get the best model as a fictionary so we can put it back in and plot
############################################################################
best_params = {}
for i in range(len(gp.parameter_names)):
if gp.parameter_names[i][:5] == 'mean:':
best_params[gp.parameter_names[i][5:]] = gp.parameter_vector[i]
print('{:>12}: {:.3f}'.format(gp.parameter_names[i][5:],
gp.parameter_vector[i]))
#########################
# Now get the best model
#########################
plt.cla()
plt.errorbar(time, rv, yerr=rv_err, fmt='ko', alpha=0.3)
# Plot the best model (just to be sure)
mean_model_better = _SB1_Model(**best_params, bounds=bounds)
plt.plot(time,mean_model_better.get_value(time), 'b--', linewidth=2, \
label='Model selected')
# Plot the variance with upsampled time series
mu, var = gp.predict(rv, time_upsampled, return_var=True)
std = np.sqrt(abs(var))
color = "#ff7f0e"
plt.fill_between( time_upsampled, \
mu+std, \
mu-std, color=color, alpha=0.8, edgecolor="none", label = 'Best fit')
plt.xlabel('Time')
plt.ylabel('RV [km/s]')
plt.legend()
plt.show()
if emcee_fit:
initial = np.array(soln.x)
ndim, nwalkers = len(initial), 32
sampler = emcee.EnsembleSampler(nwalkers,
ndim,
log_probability,
args=(gp, rv, priors, flux_weighted),
threads=8)
print("Running burn-in...")
p0 = initial + 1e-5 * np.random.randn(nwalkers, ndim)
burn_in = draws
width = 30
start_time = time_.time()
i, result = [], []
for i, result in enumerate(sampler.sample(p0, iterations=burn_in)):
n = int((width + 1) * float(i) / burn_in)
delta_t = time_.time(
) - start_time # time to do float(i) / n_steps % of caluculations
time_incr = delta_t / (float(i + 1) / burn_in
) # seconds per increment
time_left = time_incr * (1 - float(i) / burn_in)
m, s = divmod(time_left, 60)
h, m = divmod(m, 60)
sys.stdout.write('\r[{0}{1}] {2}% - {3}h:{4}m:{5:.2f}s'.format(
'#' * n, ' ' * (width - n), 100 * float(i) / burn_in, h, m, s))
names = gp.get_parameter_names()
cols = mean_model.get_parameter_names()
inds = np.array([names.index("mean:" + k) for k in cols])
samples = sampler.chain[:, int(np.floor(draws * 0.75)):, :].reshape(
(-1, ndim))
print(samples.shape, cols)
corner.corner(samples, labels=gp.get_parameter_names())
plt.show()
########################
# Now save to fits file
########################
from astropy.table import Table, Column
chain_file = 'rv_fit.fits'
try:
os.remove(chain_file)
except:
pass
t = Table(sampler.flatchain, names=gp.get_parameter_names())
t.add_column(Column(sampler.flatlnprobability, name='loglike'))
indices = np.mgrid[0:nwalkers, 0:burn_in]
step = indices[1].flatten()
walker = indices[0].flatten()
t.add_column(Column(step, name='step'))
t.add_column(Column(walker, name='walker'))
t.write(chain_file)
def GPSpec_2Comp(wav,
flux,
flux_err,
shifts_in=None,
nsteps=2000,
nrange=3,
prefix='RR2'):
# NB: input wavelengths should be in nm, flux continuum should be about 1
K, N = wav.shape
# Create 2-D array of scaled log wavelengths for fitting
lwav = np.log(wav * 1e-9) # in m
lw0, lw1 = lwav.min(), lwav.max()
x = (lwav - lw0) / (lw1 - lw0)
# First do GP fit to individual spectra to get estimate of GP HPs
print 'GP fit to individual spectra'
HPs = np.zeros((K, 3))
for i in range(K):
xx = x[i, :].flatten()
yy = flux[i, :].flatten()
ee = flux_err[i, :].flatten()
HPs[i, :] = Fit0_Jitter(xx, yy, ee, verbose=False)
HPs = np.median(HPs, axis=0)
print 'GP HPs:', HPs
k = terms.Matern32Term(log_sigma=HPs[0], log_rho=HPs[1])
k += terms.JitterTerm(log_sigma=HPs[2])
gp1 = GP(k, mean=1.0)
gp2 = GP(k, mean=1.0)
# Initial (ML) estimate of parameters
print "Starting ML fit"
if shifts_in is None:
shifts_in = np.zeros(2 * (K - 1))
par_in = shifts_in / SPEED_OF_LIGHT / (lw1 - lw0)
ML_par = np.array(Fit2(x, flux, gp1, gp2, verbose=False, par_in=par_in))
par_ML = np.copy(ML_par)
par_ML *= (lw1 - lw0) * SPEED_OF_LIGHT * 1e-3
print "ML fit done"
# MCMC
print "Starting MCMC"
ndim = len(ML_par)
nwalkers = ndim * 4
p0 = ML_par + 1e-4 * np.random.randn(nwalkers, ndim)
sampler = emcee.EnsembleSampler(nwalkers,
ndim,
LP2,
args=[gp1, gp2, x, flux])
for i, result in enumerate(sampler.sample(p0, iterations=nsteps)):
n = int((30 + 1) * float(i) / nsteps)
print i
sys.stdout.write("\r[{0}{1}]".format('#' * n, ' ' * (30 - n)))
sys.stdout.write("\n")
print("MCMC done")
# find MAP parameters
iMAP = np.argmax(sampler.flatlnprobability)
MAP_par = sampler.flatchain[iMAP, :].flatten()
# extract MCMC chains
samples = sampler.chain
Lprob = sampler.lnprobability
# convert chains back to physical units: shifts in km/s
samples_tpl = np.copy(samples)
samples_tpl *= (lw1 - lw0) * SPEED_OF_LIGHT * 1e-3
par_MAP = np.copy(MAP_par)
par_MAP *= (lw1 - lw0) * SPEED_OF_LIGHT * 1e-3
# parameter names for plots
labels = []
for i in range(K - 1):
labels.append(r'$\delta v^1_{%d}$ (km/s)' % (i + 1))
for i in range(K - 1):
labels.append(r'$\delta v^2_{%d}$ (km/s)' % (i + 1))
labels = np.array(labels)
names = []
for i in range(K - 1):
names.append('dv1_%d (km/s)' % (i + 1))
for i in range(K - 1):
names.append('dv2_%d (km/s)' % (i + 1))
names = np.array(names)
# Plot the chains
fig1 = plt.figure(figsize=(12, 2 * (K - 1) + 2))
gs1 = gridspec.GridSpec(ndim + 1, 1)
gs1.update(left=0.1, right=0.98, bottom=0.07, top=0.98, hspace=0)
ax1 = plt.subplot(gs1[0, 0])
plt.setp(ax1.get_xticklabels(), visible=False)
plt.plot(Lprob.T, 'k-', alpha=0.2)
plt.ylabel(r'$\ln P$')
for i in range(ndim):
print i, ndim, len(labels)
axc = plt.subplot(gs1[i + 1, 0], sharex=ax1)
if i < (ndim - 1):
plt.setp(axc.get_xticklabels(), visible=False)
plt.plot(samples_tpl[:, :, i].T, 'k-', alpha=0.2)
plt.ylabel(labels[i])
plt.xlim(0, nsteps)
plt.xlabel('iteration number')
# Discard burnout
nburn = int(raw_input('Enter no. steps to discard as burnout: '))
plt.axvline(nburn)
# Evaluate and print the parameter ranges
print '\n{:20s}: {:10s} {:10s} {:10s} - {:7s} + {:7s}'.format('Parameter', 'ML', 'MAP', \
'Median','Error','Error')
par50 = np.zeros(ndim)
par84 = np.zeros(ndim)
par16 = np.zeros(ndim)
for i in range(ndim):
sam = samples_tpl[:, :, i].flatten()
b, m, f = np.percentile(sam, [16, 50, 84])
par50[i] = m
par16[i] = b
par84[i] = f
print '{:20s}: {:10.5f} {:10.5f} {:10.5f} - {:7.5f} + {:7.5f}'.format(names[i], \
par_ML[i], \
par_MAP[i], \
m, m-b, f-m)
if prefix is None:
return par_MAP, par50, par50 - par16, par84 - par50
plt.savefig('%s_chains.png' % prefix)
samples_flat = samples[:, nburn:, :].reshape(-1, ndim)
samples_tpl_flat = samples_tpl[:, nburn:, :].reshape(-1, ndim)
# Plot the parameter distributions
fig2 = corner.corner(samples_tpl_flat, truths = par_MAP, labels = labels, show_titles = True, \
quantiles = [0.16, 0.84])
plt.savefig('%s_corner.png' % prefix)
# Plot the individual spectra with MAP fit
xpred = np.copy(x)
fpred, fpred_err, f1pred, f2pred = Pred2_2D(MAP_par,
gp1,
gp2,
x,
flux,
flux_err,
xpred=xpred)
lwpred = (lw1 - lw0) * xpred + lw0
wpred = np.exp(lwpred) * 1e9
fig3 = plt.figure(figsize=(12, K + 1))
gs3 = gridspec.GridSpec(K, 1)
gs3.update(left=0.1, right=0.98, bottom=0.07, top=0.98, hspace=0)
for i in range(K):
if i == 0:
ax1 = plt.subplot(gs3[0, 0])
else:
axc = plt.subplot(gs3[i, 0], sharex=ax1, sharey=ax1)
if i < (K - 1):
plt.setp(ax1.get_xticklabels(), visible=False)
plt.plot(wpred[i, :], f1pred[i, :], 'C1')
plt.plot(wpred[i, :], f2pred[i, :], 'C2')
plt.fill_between(wpred[i,:], fpred[i,:] + 2 * fpred_err[i,:], \
fpred[i,:] - fpred_err[i,:], color = 'C0', alpha = 0.4, lw = 0)
plt.plot(wpred[i, :], fpred[i, :], 'C0')
plt.ylabel('spec. %d' % (i + 1))
plt.errorbar(wav[i,:], flux[i,:], yerr = flux_err[i,:], \
fmt = ".k", ms = 3, mec = 'none', capsize = 0, alpha = 0.5, lw=0.5)
plt.xlim(wav.min(), wav.max())
plt.xlabel('wavelength (nm)')
plt.savefig('%s_spectra.png' % prefix)
# Plot the combined spectra with samples from MCMC chain
s1 = np.append(0.0, MAP_par[:K - 1])
x11d = (x + s1[:, None]).flatten()
lw11d = (lw1 - lw0) * x11d + lw0
w11d = np.exp(lw11d) * 1e9
K1 = gp1.get_matrix(x11d)
s2 = np.append(0.0, MAP_par[K - 1:])
x21d = (x + s2[:, None]).flatten()
lw21d = (lw1 - lw0) * x21d + lw0
w21d = np.exp(lw21d) * 1e9
K2 = gp2.get_matrix(x21d)
y1derr = flux_err.flatten()
Ktot = K1 + K2 + np.diag(y1derr**2)
y1d = flux.flatten() - 1.0
y11d = (flux - f2pred).flatten() + 1
y21d = (flux - f1pred).flatten() + 1
offset = 1.5 * (y11d.min() - 1)
L = sla.cho_factor(Ktot)
b = sla.cho_solve(L, y1d)
fig4 = plt.figure(figsize=(12, 2 * nrange + 1))
gs4 = gridspec.GridSpec(nrange, 1)
gs4.update(left=0.1, right=0.98, bottom=0.07, top=0.98, hspace=0.15)
ws = min(w11d.min(), w21d.min())
wr = (max(w11d.max(), w21d.max()) - ws) / float(nrange)
for i in range(nrange):
if i == 0:
ax1 = plt.subplot(gs4[0, 0])
else:
axc = plt.subplot(gs4[i, 0], sharey=ax1)
wmin = ws + (i - 0.05) * wr
wmax = ws + (i + 1.05) * wr
l = (w11d >= wmin) * (w11d <= wmax)
plt.errorbar(w11d[l], y11d[l], yerr = y1derr[l], fmt = ".k", capsize = 0, \
alpha = 0.5, ms = 2, mec='none')
l = (w21d >= wmin) * (w21d <= wmax)
plt.errorbar(w21d[l], y21d[l] + offset, yerr = y1derr[l], fmt = ".k", capsize = 0, \
alpha = 0.5, ms = 2, mec='none')
wpred = np.linspace(wmin, wmax, 1000)
lwpred = np.log(wpred * 1e-9)
xpred = (lwpred - lw0) / (lw1 - lw0)
isamp = np.random.randint(nsteps - nburn, size=10)
for j in isamp:
samp_params = samples_flat[j, :].flatten()
s1 = samp_params[:K - 1]
x1pred = (xpred + s1[:, None]).flatten()
lw1pred = (lw1 - lw0) * x1pred + lw0
w1pred = np.exp(lw1pred) * 1e9
K1s = gp1.get_matrix(x1pred, x11d)
s2 = samp_params[K - 1:]
x2pred = (xpred + s2[:, None]).flatten()
lw2pred = (lw1 - lw0) * x2pred + lw0
w2pred = np.exp(lw2pred) * 1e9
K2s = gp2.get_matrix(x2pred, x21d)
Ks = K1s + K2s
K1ss = gp1.get_matrix(x1pred)
K2ss = gp2.get_matrix(x2pred)
Kss = K1ss + K2ss
mu1 = np.dot(K1s, b).reshape(x1pred.shape) + 1
mu2 = np.dot(K2s, b).reshape(x2pred.shape) + 1
inds1 = np.argsort(w1pred)
plt.plot(w1pred[inds1], mu1[inds1], 'C0-', lw=0.5, alpha=0.5)
inds2 = np.argsort(w2pred)
plt.plot(w2pred[inds2],
mu2[inds2] + offset,
'C1-',
lw=0.5,
alpha=0.5)
plt.xlim(wmin, wmax)
plt.ylabel('flux')
plt.xlabel('wavelength (nm)')
plt.savefig('%s_combined.png' % prefix)
return par_MAP, par50, par50 - par16, par84 - par50, [
fig1, fig2, fig3, fig4
]
def fitspec(wav, flux, flux_err, nsteps = 2000, nrange = 3, prefix = 'RR1'):
K, N = wav.shape
lwav = np.log(wav * 1e-9) # in m
lw0, lw1 = lwav.min(), lwav.max()
x = (lwav - lw0) / (lw1 - lw0)
x = np.copy(lwav)
# First do GP fit to one spectrum to get estimate of GP HPs
i = np.random.randint(K)
xx = x[i,:].flatten()
yy = flux[i,:].flatten()
ee = flux_err[i,:].flatten()
xp = np.linspace(xx.min(), xx.max(), 100)
print xx
HPs, _, _ = u.Fit0(xx, yy, ee, verbose = False, xpred = xp, HP_init=np.array([-1.4,-9.03]))
print 'Initial GP HPs:', HPs
# Initial (ML) estimate of parameters
print "Starting ML fit"
par_in = np.zeros(K+1)
par_in[-2:] = HPs
ML_par = np.array(u.Fit1(x, flux, flux_err, verbose = False, par_in = par_in))
print ML_par
par_ML = np.copy(ML_par)
par_ML[:K-1] *= (lw1 - lw0) * SPEED_OF_LIGHT
print "ML fit done"
return par_ML
# MCMC
print "Starting MCMC"
ndim = K+1
nwalkers = ndim * 4
p0 = ML_par + 1e-6 * np.random.randn(nwalkers, ndim)
k = terms.Matern32Term(log_sigma = ML_par[-2], log_rho = ML_par[-1])
gp = GP(k, mean = 1.0)
sampler = emcee.EnsembleSampler(nwalkers, ndim, u.LP1,
args = [gp, x, flux, flux_err])
for i, result in enumerate(sampler.sample(p0, iterations=nsteps)):
n = int((30+1) * float(i) / nsteps)
sys.stdout.write("\r[{0}{1}]".format('#' * n, ' ' * (30 - n)))
sys.stdout.write("\n")
print("MCMC done")
# find MAP parameters
iMAP = np.argmax(sampler.flatlnprobability)
MAP_par = sampler.flatchain[iMAP,:].flatten()
# extract MCMC chains
samples = sampler.chain
Lprob = sampler.lnprobability
# convert chains back to physical units: shifts in km/s
samples_tpl = np.copy(samples)
samples_tpl[:,:,:K-1] *= (lw1 - lw0) * SPEED_OF_LIGHT
par_MAP = np.copy(MAP_par)
par_MAP[:K-1] *= (lw1 - lw0) * SPEED_OF_LIGHT
# parameter names for plots
labels = []
for i in range(K-1):
labels.append(r'$\delta v_{%d}$ (m/s)' % (i+1))
labels.append(r'$\ln \sigma$')
labels.append(r'$\ln \rho$')
labels = np.array(labels)
names = []
for i in range(K-1):
names.append('dv_%d (m/s)' % (i+1))
names.append('ln(sig)')
names.append('ln(rho)')
names = np.array(names)
# Plot the chains
fig1 = pl.figure(figsize = (12,K+3))
gs1 = gridspec.GridSpec(ndim+1,1)
gs1.update(left=0.1, right=0.98, bottom = 0.07, top = 0.98, hspace=0)
ax1 = pl.subplot(gs1[0,0])
axs = [ax1]
pl.setp(ax1.get_xticklabels(), visible=False)
pl.plot(Lprob.T, 'k-', alpha = 0.2)
pl.ylabel(r'$\ln P$')
for i in range(ndim):
axc = pl.subplot(gs1[i+1,0], sharex = ax1)
axs.append(axc)
if i < (ndim-1):
pl.setp(axc.get_xticklabels(), visible=False)
pl.plot(samples_tpl[:,:,i].T, 'k-', alpha = 0.2)
pl.ylabel(labels[i])
pl.xlim(0,nsteps)
pl.xlabel('iteration number')
# Discard burnout
nburn = int(0.25*nsteps)
for ax in axs:
ax.axvline(nburn)
pl.savefig('../plots/%s_chains.png' % prefix)
# Evaluate and print the parameter ranges
print '\n{:20s}: {:10s} {:10s} {:10s} - {:7s} + {:7s}'.format('Parameter', 'ML', 'MAP', \
'Median','Error','Error')
par50 = np.zeros(ndim)
par84 = np.zeros(ndim)
par16 = np.zeros(ndim)
for i in range(ndim):
sam = samples_tpl[:,:,i].flatten()
b, m, f = np.percentile(sam, [16,50,84])
par50[i] = m
par16[i] = b
par84[i] = f
print '{:20s}: {:10.5f} {:10.5f} {:10.5f} - {:7.5f} + {:7.5f}'.format(names[i], \
par_ML[i], \
par_MAP[i], \
m, m-b, f-m)
samples_flat = samples[:,nburn:,:].reshape(-1, ndim)
samples_tpl_flat = samples_tpl[:,nburn:,:].reshape(-1, ndim)
# Plot the parameter distributions
fig2 = corner.corner(samples_tpl_flat, truths = par_MAP, labels = labels, show_titles = True, \
quantiles = [0.16, 0.84])
pl.savefig('../plots/%s_corner.png' % prefix)
return par_MAP, par50, par50-par16, par84-par50
def GPSpec_1Comp(wav, flux, flux_err, nsteps=2000, nrange=3, prefix='RR1'):
# NB: input wavelengths should be in nm, flux continuum should be about 1
K, N = wav.shape
# Create 2-D array of scaled log wavelengths for fitting
lwav = np.log(wav * 1e-9) # in m
lw0, lw1 = lwav.min(), lwav.max()
x = (lwav - lw0) / (lw1 - lw0)
# First do GP fit to individual spectra to get estimate of GP HPs
print 'GP fit to individual spectra'
HPs = np.zeros((K, 2))
for i in range(K):
xx = x[i, :].flatten()
yy = flux[i, :].flatten()
ee = flux_err[i, :].flatten()
HPs[i, :] = Fit0(xx, yy, ee, verbose=False, xpred=None)
HPs = np.median(HPs, axis=0)
print 'Initial GP HPs:', HPs
# Initial (ML) estimate of parameters
print "Starting ML fit"
par_in = np.zeros(K + 1)
par_in[-2:] = HPs
ML_par = np.array(Fit1(x, flux, flux_err, verbose=False, par_in=par_in))
par_ML = np.copy(ML_par)
par_ML[:K - 1] *= (lw1 - lw0) * SPEED_OF_LIGHT * 1e-3
par_ML[-1] *= (lw1 - lw0)
k = terms.Matern32Term(log_sigma=ML_par[-2], log_rho=ML_par[-1])
gp = GP(k, mean=1.0)
print "ML fit done"
# MCMC
print "Starting MCMC"
ndim = K + 1
nwalkers = ndim * 4
p0 = ML_par + 1e-4 * np.random.randn(nwalkers, ndim)
sampler = emcee.EnsembleSampler(nwalkers,
ndim,
LP1,
args=[gp, x, flux, flux_err])
for i, result in enumerate(sampler.sample(p0, iterations=nsteps)):
n = int((30 + 1) * float(i) / nsteps)
sys.stdout.write("\r[{0}{1}]".format('#' * n, ' ' * (30 - n)))
sys.stdout.write("\n")
print("MCMC done")
# find MAP parameters
iMAP = np.argmax(sampler.flatlnprobability)
MAP_par = sampler.flatchain[iMAP, :].flatten()
# extract MCMC chains
samples = sampler.chain
Lprob = sampler.lnprobability
# convert chains back to physical units: shifts in km/s
samples_tpl = np.copy(samples)
samples_tpl[:, :, :K - 1] *= (lw1 - lw0) * SPEED_OF_LIGHT * 1e-3
samples_tpl[:, :, -1] *= (lw1 - lw0)
par_MAP = np.copy(MAP_par)
par_MAP[:K - 1] *= (lw1 - lw0) * SPEED_OF_LIGHT * 1e-3
par_MAP[-1] *= (lw1 - lw0)
# parameter names for plots
labels = []
for i in range(K - 1):
labels.append(r'$\delta v_{%d}$ (km/s)' % (i + 1))
labels.append(r'$\ln \sigma$')
labels.append(r'$\ln \rho$')
labels = np.array(labels)
names = []
for i in range(K - 1):
names.append('dv_%d (km/s)' % (i + 1))
names.append('ln(sig)')
names.append('ln(rho)')
names = np.array(names)
# Plot the chains
fig1 = plt.figure(figsize=(12, K + 3))
gs1 = gridspec.GridSpec(ndim + 1, 1)
gs1.update(left=0.1, right=0.98, bottom=0.07, top=0.98, hspace=0)
ax1 = plt.subplot(gs1[0, 0])
plt.setp(ax1.get_xticklabels(), visible=False)
plt.plot(Lprob.T, 'k-', alpha=0.2)
plt.ylabel(r'$\ln P$')
for i in range(ndim):
axc = plt.subplot(gs1[i + 1, 0], sharex=ax1)
if i < (ndim - 1):
plt.setp(axc.get_xticklabels(), visible=False)
plt.plot(samples_tpl[:, :, i].T, 'k-', alpha=0.2)
plt.ylabel(labels[i])
plt.xlim(0, nsteps)
plt.xlabel('iteration number')
# Discard burnout
nburn = int(raw_input('Enter no. steps to discard as burnout: '))
plt.axvline(nburn)
# Evaluate and print the parameter ranges
print '\n{:20s}: {:10s} {:10s} {:10s} - {:7s} + {:7s}'.format('Parameter', 'ML', 'MAP', \
'Median','Error','Error')
par50 = np.zeros(ndim)
par84 = np.zeros(ndim)
par16 = np.zeros(ndim)
for i in range(ndim):
sam = samples_tpl[:, nburn:, i].flatten()
b, m, f = np.percentile(sam, [16, 50, 84])
par50[i] = m
par16[i] = b
par84[i] = f
print '{:20s}: {:10.5f} {:10.5f} {:10.5f} - {:7.5f} + {:7.5f}'.format(names[i], \
par_ML[i], \
par_MAP[i], \
m, m-b, f-m)
if prefix is None:
return par_MAP, par50, par50 - par16, par84 - par50
plt.savefig('%s_chains.png' % prefix)
samples_flat = samples[:, nburn:, :].reshape(-1, ndim)
samples_tpl_flat = samples_tpl[:, nburn:, :].reshape(-1, ndim)
# Plot the parameter distributions
fig2 = corner.corner(samples_tpl_flat, truths = par_MAP, labels = labels, show_titles = True, \
quantiles = [0.16, 0.84])
plt.savefig('%s_corner.png' % prefix)
# Plot the individual spectra with MAP fit
xpred, fpred, fpred_err = Pred1_2D(MAP_par,
x,
flux,
flux_err,
doPlot=False)
lwpred = (lw1 - lw0) * xpred + lw0
wpred = np.exp(lwpred) * 1e9
fig3 = plt.figure(figsize=(12, K + 1))
gs3 = gridspec.GridSpec(K, 1)
gs3.update(left=0.1, right=0.98, bottom=0.07, top=0.98, hspace=0)
for i in range(K):
if i == 0:
ax1 = plt.subplot(gs3[0, 0])
else:
axc = plt.subplot(gs3[i, 0], sharex=ax1, sharey=ax1)
if i < (K - 1):
plt.setp(ax1.get_xticklabels(), visible=False)
plt.errorbar(wav[i,:], flux[i,:], yerr = flux_err[i,:], \
fmt = ".k", ms = 2, mec = 'none', capsize = 0, alpha = 0.5)
plt.plot(wpred[i, :], fpred[i, :], 'C0')
plt.fill_between(wpred[i,:], fpred[i,:] + 2 * fpred_err[i,:], \
fpred[i,:] - fpred_err[i,:], color = 'C0', alpha = 0.4, lw = 0)
plt.ylabel('spec. %d' % (i + 1))
plt.xlim(wav.min(), wav.max())
plt.xlabel('wavelength (nm)')
plt.savefig('%s_spectra.png' % prefix)
# Plot the combined spectra with samples from MCMC chain
shifts = np.append(0, MAP_par[:K - 1])
x1d = (x + shifts[:, None]).flatten()
lw1d = (lw1 - lw0) * x1d + lw0
w1d = np.exp(lw1d) * 1e9
y1d = flux.flatten()
y1derr = flux_err.flatten()
inds = np.argsort(x1d)
gp.set_parameter_vector(MAP_par[-2:])
gp.compute(x1d[inds], yerr=y1derr[inds])
fig4 = plt.figure(figsize=(12, nrange + 1))
gs4 = gridspec.GridSpec(nrange, 1)
gs4.update(left=0.1, right=0.98, bottom=0.07, top=0.98, hspace=0.05)
ws = w1d.min()
wr = (w1d.max() - ws) / float(nrange)
for i in range(nrange):
if i == 0:
ax1 = plt.subplot(gs4[0, 0])
else:
axc = plt.subplot(gs4[i, 0], sharey=ax1)
if i < (nrange - 1):
plt.setp(ax1.get_xticklabels(), visible=False)
wmin = ws + (i - 0.05) * wr
wmax = ws + (i + 1.05) * wr
l = (w1d >= wmin) * (w1d <= wmax)
plt.errorbar(w1d[l], y1d[l], yerr = y1derr[l], fmt = ".k", capsize = 0, \
alpha = 0.5, ms = 2, mec='none')
wpred = np.linspace(wmin, wmax, 1000)
lwpred = np.log(wpred * 1e-9)
xpred = (lwpred - lw0) / (lw1 - lw0)
isamp = np.random.randint(nsteps - nburn, size=10)
for j in isamp:
samp_params = samples_flat[j, :].flatten()
samp_shifts = np.append(0, samp_params[:K - 1])
x1_samp = (x + samp_shifts[:, None]).flatten()
inds_samp = np.argsort(x1_samp)
k_samp = terms.Matern32Term(log_sigma=samp_params[-2],
log_rho=samp_params[-1])
gp_samp = GP(k_samp, mean=1.)
gp_samp.compute(x1_samp[inds_samp], yerr=y1derr[inds_samp])
mu, _ = gp.predict(y1d[inds_samp], xpred, return_var=True)
plt.plot(wpred, mu, 'C0-', lw=0.5, alpha=0.5)
plt.xlim(wmin, wmax)
plt.ylabel('flux')
plt.xlabel('wavelength (nm)')
plt.savefig('%s_combined.png' % prefix)
return par_MAP, par50, par50 - par16, par84 - par50, [
fig1, fig2, fig3, fig4
]
george_solvers = {
"basic": george.BasicSolver,
"HODLR": george.HODLRSolver,
}
celerite_terms = {
"N": terms.Term(),
"B": terms.RealTerm(log_a=-6., log_c=-np.inf,
bounds={"log_a": [-30, 30],
"log_c": [-np.inf, np.inf]}),
"W": terms.JitterTerm(log_sigma=-25,
bounds={"log_sigma": [-30, 30]}),
"Mat32": terms.Matern32Term(
log_sigma=1.,
log_rho=1.,
bounds={"log_sigma": [-30, 30],
# The `celerite` version of the Matern-3/2
# kernel has problems with very large `log_rho`
# values. -7.4 is empirical.
"log_rho": [-7.4, 16]}),
"SHO0": terms.SHOTerm(log_S0=-6, log_Q=1.0 / np.sqrt(2.), log_omega0=0.,
bounds={"log_S0": [-30, 30],
"log_Q": [-30, 30],
"log_omega0": [-30, 30]}),
"SHO1": terms.SHOTerm(log_S0=-6, log_Q=-2., log_omega0=0.,
bounds={"log_S0": [-10, 10],
"log_omega0": [-10, 10]}),
"SHO2": terms.SHOTerm(log_S0=-6, log_Q=0.5, log_omega0=0.,
bounds={"log_S0": [-10, 10],
"log_Q": [-10, 10],
"log_omega0": [-10, 10]}),
# see Foreman-Mackey et al. 2017, AJ 154, 6, pp. 220
def calculate_lnlike(params, inst, key):
#if fitting flares, force them to be in time order
# if config.BASEMENT.settings['N_flares'] > 0:
# flare_times = [ params['flare_tpeak_'+str(i)] for i in range(1,config.BASEMENT.settings['N_flares']+1) ]
# if sorted(flare_times) != flare_times:
# return -np.inf
#::: calculate the model. if there are any NaN, return -np.inf
model = calculate_model(params, inst, key)
if any(np.isnan(model)) or any(np.isinf(model)):
return -np.inf
#::: if no GP baseline sampling, then calculate lnlike per hand
if config.BASEMENT.settings['baseline_'+key+'_'+inst] != 'sample_GP':
yerr_w = calculate_yerr_w(params, inst, key)
baseline = calculate_baseline(params, inst, key, model=model, yerr_w=yerr_w)
residuals = config.BASEMENT.data[inst][key] - model - baseline
inv_sigma2_w = 1./yerr_w**2
# print('###############################################################################')
# print('model',model)
# print('baseline',baseline)
## try:
# fig = plt.figure()
# plt.plot(config.BASEMENT.data[inst]['time'][0:200], config.BASEMENT.data[inst]['flux'][0:200], 'b.')
# plt.plot(config.BASEMENT.data[inst]['time'][0:200], model[0:200]+baseline, 'r-')
# plt.title( 'lnlike ' + str(-0.5*(np.nansum((residuals)**2 * inv_sigma2_w - np.log(inv_sigma2_w)))) )
# plt.savefig( os.path.join(config.BASEMENT.outdir,'fig_'+str(params['b_period'])+'.jpg') )
# plt.close(fig)
## except:
## pass
return -0.5*(np.sum((residuals)**2 * inv_sigma2_w - np.log(inv_sigma2_w/2./np.pi))) #use np.sum to catch any nan and then set lnlike to nan
#::: if GP baseline sampling, use the GP lnlike instead
#::: this is MUCH MUCH MUUUUUCH FASTER than gp.predict
else:
x = config.BASEMENT.data[inst]['time']
y = config.BASEMENT.data[inst][key] - model
yerr_w = calculate_yerr_w(params, inst, key)
# print(params['baseline_gp1_'+key+'_'+inst])
# print(params['baseline_gp2_'+key+'_'+inst])
kernel = terms.Matern32Term(log_sigma=params['baseline_gp1_'+key+'_'+inst],
log_rho=params['baseline_gp2_'+key+'_'+inst])
gp = celerite.GP(kernel)
try:
gp.compute(x, yerr=yerr_w)
lnlike = gp.log_likelihood(y)
# print('runs')
# print(lnlike)
# try:
# #::: debug
# baseline2 = gp.predict(y, x)[0]
# plt.figure()
# plt.plot(x,y,'k.', color='grey')
# plt.plot(xx,baseline,'r-', lw=2)
# plt.plot(x,baseline2,'ro', lw=2)
# plt.title(inst+' '+key+' '+str(gp.get_parameter_vector()))
# plt.show()
# raw_input('press enter to continue')
except:
lnlike = -np.inf
# print('fails')
return lnlike
def __init__(self,
host=None,
planets=None,
use_gps=False,
log_sigma=6,
log_rho=5,
log_sigma_error=(-0.5, 0.5),
log_rho_error=(-0.5, 0.5)):
self.use_gps = use_gps
self.fit_params = deepcopy(fit_params)
if not self.use_gps:
self.fit_params['GP'] = []
# Validate everything
if host is not None:
host._validate()
if planets is not None:
for planet in planets:
planet._validate()
self.host = host
if isinstance(planets, list):
self.planets = planets
elif planets is None:
self.planets = []
else:
self.planets = [planets]
self.nplanets = len(self.planets)
for planet in self.planets:
planet._validate()
self.host._validate()
l = []
for idx in range(self.nplanets):
for f in self.fit_params['planet']:
l.append(
u.Quantity(np.copy(getattr(self.planets[idx], f))).value)
self._best_guess = l
self.initial_guess = l
self.system = None
if self.use_gps:
log.info(
'Using Gaussian Process to fit long term trends. This will make fitting slower, but more accurate.'
)
# Store GP model Parameters
self.log_sigma = log_sigma
self.log_rho = log_rho
self.log_sigma_error = log_sigma_error
self.log_rho_error = log_rho_error
# Set up the GP model
kernel = terms.Matern32Term(log_sigma=self.log_sigma,
log_rho=self.log_rho,
bounds=self.bounds[-2:])
self.gp = celerite.GP(kernel)
for f in self.fit_params['GP']:
l.append(getattr(self, f))
self._best_guess = l
l = []
for jdx in range(self.nplanets):
for idx, f in enumerate(self.fit_params['planet']):
l.append(
tuple(
np.asarray(
np.copy(getattr(self.planets[jdx], f +
'_error'))) +
self._best_guess[
(jdx * len(self.fit_params['planet'])) + idx]))
for f in self.fit_params['GP']:
l.append(
tuple([
getattr(self, f) + getattr(self, f + '_error')[0],
getattr(self, f) + getattr(self, f + '_error')[1]
]))
self.initial_bounds = l
self._is_eccen = np.where(
[l.split('.')[-1] == 'eccentricity' for l in self._fit_labels])[0]
self._is_inc = np.where(
[l.split('.')[-1] == 'inclination' for l in self._fit_labels])[0]
self.nwalkers = 100
self.burnin = 200
self.nsteps = 1000
# Work around for the starry pickle bug.
global _wellfit_toy_model
_wellfit_toy_model = None
self._initialize_system()
xpred = np.linspace(0.29,0.34,1000)
HPs, mu, std = u.Fit0(x, y, yerr, verbose = True, doPlot = False, \
xpred = xpred, HP_init = HP_in)
pl.clf()
pl.errorbar(x,y,yerr=yerr,fmt='k.',ms=8,capsize=0,lw=0.5,alpha=0.5)
pl.fill_between(xpred, mu + 2 * std, mu - 2 * std, alpha=0.2, color='C0', lw=0)
pl.fill_between(xpred, mu + std, mu - std, alpha=0.2, color='C0', lw=0)
pl.plot(xpred, mu, 'C0-')
pl.xlim(xpred.min(), xpred.max())
pl.ylim((mu - 5 * std).min(), (mu + 5 * std).max())
# ----------------
# Simulate spectra
# ----------------
# First set up GP object
k = terms.Matern32Term(log_sigma = HPs[0], log_rho = HPs[1])
gp = GP(k, mean = 1.0)
gp.compute(x, yerr = yerr)
# Barycentric shifts
baryvel = np.random.uniform(low = -BVMAX, high = BVMAX, size = NSMAX)
dlw = baryvel / SPEED_OF_LIGHT
lwav_sim = np.tile(lwav, (NSMAX, 1)) + dlw[:,None]
x_sim = (lwav_sim - lw0) / (lw1 - lw0)
wav_rest_sim = np.exp(lwav_sim) * 1e9
wav_earth_sim = np.tile(wav_rest, (NSMAX, 1))
# Evaluate each spectrum using predictive mean of GP conditioned on
# observed spectrum, and wavelength shifts caused by barycentric
# velocity changes
flux_sim = np.zeros((NSMAX, N))
for i in range(NSMAX):
print baryvel[i], dlw[i], np.median(x_sim[i,:]-x)*N
