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