def test_predict(seed=42): np.random.seed(seed) x = np.linspace(1, 59, 300) t = np.sort(np.random.uniform(10, 50, 100)) yerr = np.random.uniform(0.1, 0.5, len(t)) y = np.sin(t) kernel = terms.RealTerm(0.1, 0.5) for term in [(0.6, 0.7, 1.0), (0.1, 0.05, 0.5, -0.1)]: kernel += terms.ComplexTerm(*term) gp = GP(kernel) gp.compute(t, yerr) K = gp.get_matrix(include_diagonal=True) Ks = gp.get_matrix(x, t) true_mu = np.dot(Ks, np.linalg.solve(K, y)) true_cov = gp.get_matrix(x, x) - np.dot(Ks, np.linalg.solve(K, Ks.T)) mu, cov = gp.predict(y, x) _, var = gp.predict(y, x, return_var=True) assert np.allclose(mu, true_mu) assert np.allclose(cov, true_cov) assert np.allclose(var, np.diag(true_cov)) mu0, cov0 = gp.predict(y, t) mu, cov = gp.predict(y) assert np.allclose(mu0, mu) assert np.allclose(cov0, cov)
def __call__(self, rho: Optional[float] = 2., niter: int = 5): logger.info("Running Celerite Matern 3/2 detrending") time = self.ts.time.copy() flux = self.ts.flux.copy() mask = ones(time.size, bool) for i in range(niter): time_learn = time[mask] flux_learn = flux[mask] wn = mad_std(diff(flux_learn)) / sqrt(2) log_sigma = log(mad_std(flux_learn)) log_rho = log(rho) kernel = Matern32Term(log_sigma, log_rho) gp = GP(kernel, mean=1.0) gp.compute(time_learn, yerr=wn) self.prediction = gp.predict(flux_learn, time, return_cov=False) mask &= ~sigma_clip(flux - self.prediction, sigma=3).mask residuals = flux - self.prediction self.mask = m = ~(sigma_clip(residuals, sigma_lower=inf, sigma_upper=4).mask) self.ts._data.update('celerite_m32', time[m], (flux - self.prediction + 1)[m], self.ts.ferr[m])
def test_predict(method, seed=42): np.random.seed(seed) x = np.sort(np.random.rand(10)) yerr = np.random.uniform(0.1, 0.5, len(x)) y = np.sin(x) kernel = terms.RealTerm(0.1, 0.5) for term in [(0.6, 0.7, 1.0)]: kernel += terms.ComplexTerm(*term) gp = GP(kernel, method=method) gp.compute(x, yerr) mu0, cov0 = gp.predict(y, x) mu, cov = gp.predict(y) assert np.allclose(mu0, mu) assert np.allclose(cov0, cov)
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 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 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 __call__(self, period: Optional[float] = None): logger.info("Running Celerite") assert hasattr( self.ts, 'ls'), "Celerite step requires a prior Lomb-Scargle step" self.period = period if period is not None else self.ts.ls.period gp = GP(SHOTerm(log(self.ts.flux.var()), log(10), log(2 * pi / self.period)), mean=1.) gp.freeze_parameter('kernel:log_omega0') gp.compute(self.ts.time, yerr=self.ts.ferr) def minfun(pv): gp.set_parameter_vector(pv) gp.compute(self.ts.time, yerr=self.ts.ferr) return -gp.log_likelihood(self.ts.flux) res = minimize(minfun, gp.get_parameter_vector(), jac=False, method='powell') self.result = res self.parameters = res.x self.prediction = gp.predict(self.ts.flux, return_cov=False)
class LineFitter(object): def __init__(self, x, flux, ivar=None, absorp_emiss=None, n_bg_coef=2, target_x=None): """ Parameters ---------- x : array_like Dispersion parameter. Usually either pixel or wavelength. flux : array_like Flux array. ivar : array_like Inverse-variance array. absorp_emiss : numeric [-1, 1] (optional) If -1, absorption line. If +1, emission line. If not specified, will try to guess. n_bg_coef : int (optional) The order of the polynomial used to fit the background. 1 = constant, 2 = linear, 3 = quadratic, etc. target_x : numeric (optional) Where we think the line of interest is. """ self.x = np.array(x) self.flux = np.array(flux) if ivar is not None: self.ivar = np.array(ivar) self._err = 1 / np.sqrt(self.ivar) else: self.ivar = np.ones_like(flux) self._err = None if absorp_emiss is None: raise NotImplementedError( "You must supply an absorp_emiss for now") self.absorp_emiss = float(absorp_emiss) self.target_x = target_x self.n_bg_coef = int(n_bg_coef) # result of fitting self.gp = None self.success = None @classmethod def transform_pars(cls, p): new_p = OrderedDict() for k, v in p.items(): if k in cls._log_pars: k = 'ln_{0}'.format(k) v = np.log(v) new_p[k] = v return new_p def get_init_generic(self, amp=None, x0=None, bg=None, bg_buffer=None): """ Get initial guesses for the generic line parameters. Parameters ---------- amp : numeric Line amplitude. This should be strictly positive; the ``absorp_emiss`` parameter controls whether it is absorption or emission. x0 : numeric Line centroid. bg : iterable Background polynomial coefficients. bg_buffer : int The number of values to use on either end of the flux array to estimate the background parameters. """ target_x = self.target_x if x0 is None: # estimate the initial guess for the centroid relmins = argrelmin(-self.absorp_emiss * self.flux)[0] if len(relmins) > 1 and target_x is None: logger.log( 0, "no target_x specified - taking largest line " "in spec region") target_x = self.x[np.argmin(-self.absorp_emiss * self.flux)] if len(relmins) == 1: x0 = self.x[relmins[0]] else: x0_idx = relmins[np.abs(self.x[relmins] - target_x).argmin()] x0 = self.x[x0_idx] # shift x array so that line is approximately at 0 _x = np.array(self.x, copy=True) x = np.array(_x) - x0 # background polynomial parameters if bg_buffer is None: bg_buffer = len(x) // 4 bg_buffer = max(1, bg_buffer) if bg is None: if self.n_bg_coef < 2: bg = np.array([0.]) bg[0] = np.median(self.flux) else: # estimate linear background model bg = np.array([0.] * self.n_bg_coef) i1 = argmedian(self.flux[:bg_buffer]) i2 = argmedian(self.flux[-bg_buffer:]) f1 = self.flux[:bg_buffer][i1] f2 = self.flux[-bg_buffer:][i2] x1 = x[:bg_buffer][i1] x2 = x[-bg_buffer:][i2] bg[1] = (f2 - f1) / (x2 - x1) # slope bg[0] = f2 - bg[1] * x2 # estimate constant term else: if len(bg) != self.n_bg_coef: raise ValueError('Number of bg polynomial coefficients does ' 'not match n_bg_coef specified when fitter ' 'was created.') if amp is None: # then estimate the initial guess for amplitude _i = np.argmin(np.abs(x)) if len(bg) > 1: _bg = bg[0] + bg[1] * x[_i] else: _bg = bg[0] # amp = np.sqrt(2*np.pi) * (flux[_i] - bg) amp = self.flux[_i] - _bg p0 = OrderedDict() p0['amp'] = np.abs(amp) p0['x0'] = x0 p0['bg_coef'] = bg return p0 def get_init_gp(self, log_sigma=None, log_rho=None, x0=None): """ Set up the GP kernel parameters. """ if log_sigma is None: if x0 is not None: # better initial guess for GP parameters mask = np.abs(self.x - x0) > 2. y_MAD = np.median( np.abs(self.flux[mask] - np.median(self.flux[mask]))) else: y_MAD = np.median(np.abs(self.flux - np.median(self.flux))) log_sigma = np.log(3 * y_MAD) if log_rho is None: log_rho = np.log(5.) p0 = OrderedDict() p0['log_sigma'] = log_sigma p0['log_rho'] = log_rho return p0 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 get_gp_mean_pars(self): """ Return a parameter dictionary for the mean model parameters only. """ fit_pars = OrderedDict() for k, v in self.gp.get_parameter_dict().items(): if 'mean' not in k: continue k = k[5:] # remove 'mean:' if k.startswith('ln'): if 'amp' in k: fit_pars[k[3:]] = self.absorp_emiss * np.exp(v) else: fit_pars[k[3:]] = np.exp(v) elif k.startswith('bg'): if 'bg_coef' not in fit_pars: fit_pars['bg_coef'] = [] fit_pars['bg_coef'].append(v) else: fit_pars[k] = v return fit_pars def _neg_log_like(self, params): self.gp.set_parameter_vector(params) try: ll = self.gp.log_likelihood(self.flux) except (RuntimeError, LinAlgError): return np.inf if np.isnan(ll): return np.inf return -ll def fit(self, bounds=None, **kwargs): """ """ init_params = self.init_gp() if bounds is None: bounds = OrderedDict() # default bounds for mean model parameters i = self.gp.models['mean'].get_parameter_names().index('x0') mean_bounds = self.gp.models['mean'].get_parameter_bounds() if mean_bounds[i][0] is None and mean_bounds[i][1] is None: mean_bounds[i] = (min(self.x), max(self.x)) for i, k in enumerate(self.gp.models['mean'].parameter_names): mean_bounds[i] = bounds.get(k, mean_bounds[i]) self.gp.models['mean'].parameter_bounds = mean_bounds # HACK: default bounds for kernel parameters self.gp.models['kernel'].parameter_bounds = [(None, None), (np.log(0.5), None)] soln = minimize(self._neg_log_like, init_params, method="L-BFGS-B", bounds=self.gp.get_parameter_bounds()) self.success = soln.success self.gp.set_parameter_vector(soln.x) if self.success: logger.debug("Success: {0}, Final log-likelihood: {1}".format( soln.success, -soln.fun)) else: logger.warning("Fit failed! Final log-likelihood: {0}, " "Final parameters: {1}".format(-soln.fun, soln.x)) return self def plot_fit(self, axes=None, fit_alpha=0.5): unbinned_color = '#3182bd' binned_color = '#2ca25f' gp_color = '#ff7f0e' noise_color = '#de2d26' if self.gp is None: raise ValueError("You must run .fit() first!") # ---------------------------------------------------------------------- # Plot the maximum likelihood model if axes is None: fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharex=True) # data for ax in axes[:2]: ax.plot(self.x, self.flux, drawstyle='steps-mid', marker='', color='#777777', zorder=2) if self._err is not None: ax.errorbar(self.x, self.flux, self._err, marker='', ls='none', ecolor='#666666', zorder=1) # mean model wave_grid = np.linspace(self.x.min(), self.x.max(), 256) mu, var = self.gp.predict(self.flux, self.x, return_var=True) std = np.sqrt(var) axes[0].plot(wave_grid, self.gp.mean.get_unbinned_value(wave_grid), marker='', alpha=fit_alpha, zorder=10, color=unbinned_color) axes[0].plot(self.x, self.gp.mean.get_value(self.x), marker='', alpha=fit_alpha, zorder=10, drawstyle='steps-mid', color=binned_color) # full GP model axes[1].plot(self.x, mu, color=gp_color, drawstyle='steps-mid', marker='', alpha=fit_alpha, zorder=10) axes[1].fill_between(self.x, mu + std, mu - std, color=gp_color, alpha=fit_alpha / 10., edgecolor="none", step='mid') # just GP noise component mean_model = self.gp.mean.get_value(self.x) axes[2].plot(self.x, mu - mean_model, marker='', alpha=fit_alpha, zorder=10, drawstyle='steps-mid', color=noise_color) axes[2].plot(self.x, self.flux - mean_model, drawstyle='steps-mid', marker='', color='#777777', zorder=2) if self._err is not None: axes[2].errorbar(self.x, self.flux - mean_model, self._err, marker='', ls='none', ecolor='#666666', zorder=1) axes[0].set_ylabel('flux') axes[0].set_xlim(self.x.min(), self.x.max()) axes[0].set_title('Line model') axes[1].set_title('Line + noise model') axes[2].set_title('Noise model') fig = axes[0].figure fig.tight_layout() return fig
class CeleriteLogLikelihood: def __init__(self, lpf, name: str = 'gp', noise_ids=None, fixed_hps=None): if not with_celerite: raise ImportError("CeleriteLogLikelihood requires celerite.") self.name = name self.lpf = lpf if fixed_hps is None: self.free = True else: self.hps = asarray(fixed_hps) self.free = False if lpf.lcids is None: raise ValueError( 'The LPF data needs to be initialised before initialising CeleriteLogLikelihood.' ) self.noise_ids = noise_ids if noise_ids is not None else unique( lpf.noise_ids) self.mask = m = zeros(lpf.lcids.size, bool) for lcid, nid in enumerate(lpf.noise_ids): if nid in self.noise_ids: m[lcid == lpf.lcids] = 1 if m.sum() == lpf.lcids.size: self.times = lpf.timea self.fluxes = lpf.ofluxa else: self.times = lpf.timea[m] self.fluxes = lpf.ofluxa[m] self.gp = GP(Matern32Term(0, 0)) if self.free: self.init_parameters() else: self.compute_gp(None, force=True, hps=self.hps) def init_parameters(self): name = self.name wns = log10(mad_std(diff(self.fluxes)) / sqrt(2)) pgp = [ LParameter(f'{name}_ln_out', f'{name} ln output scale', '', NP(-6, 1.5), bounds=(-inf, inf)), LParameter(f'{name}_ln_in', f'{name} ln input scale', '', UP(-8, 8), bounds=(-inf, inf)), LParameter(f'{name}_log10_wn', f'{name} log10 white noise sigma', '', NP(wns, 0.025), bounds=(-inf, inf)) ] self.lpf.ps.thaw() self.lpf.ps.add_global_block(self.name, pgp) self.lpf.ps.freeze() self.pv_slice = self.lpf.ps.blocks[-1].slice self.pv_start = self.lpf.ps.blocks[-1].start setattr(self.lpf, f"_sl_{name}", self.pv_slice) setattr(self.lpf, f"_start_{name}", self.pv_start) def compute_gp(self, pv, force: bool = False, hps=None): if self.free or force: parameters = pv[self.pv_slice] if hps is None else hps self.gp.set_parameter_vector(parameters[:-1]) self.gp.compute(self.times, yerr=10**parameters[-1]) def compute_gp_lnlikelihood(self, pv, model): self.compute_gp(pv) return self.gp.log_likelihood(self.fluxes - model[self.mask]) def predict_baseline(self, pv): self.compute_gp(pv) residuals = self.fluxes - squeeze( self.lpf.transit_model(pv))[self.mask] bl = zeros_like(self.lpf.timea) bl[self.mask] = self.gp.predict(residuals, self.times, return_cov=False) return 1. + bl def __call__(self, pvp, model): if pvp.ndim == 1: lnlike = self.compute_gp_lnlikelihood(pvp, model) else: lnlike = zeros(pvp.shape[0]) for ipv, pv in enumerate(pvp): if all(isfinite(model[ipv])): lnlike[ipv] = self.compute_gp_lnlikelihood(pv, model[ipv]) else: lnlike[ipv] = -inf return lnlike
class GP_celerite(): """ A simple Gaussian process class for regression problems, based on celerite """ def __init__(self, kernel, t, yerr=None, white_noise=1e-10): """ t : the independent variable where the GP is "trained" (opt) yerr : array of observational uncertainties (opt) white_noise : value/array added to the diagonal of the kernel matrix """ self.kernel = kernel self.t = t self.yerr = yerr self.white_noise = white_noise self._GP = GPcel(self.kernel) self._GP.compute(self.t, self.yerr) def _cov(self, x1, x2=None, add2diag=True): raise NotImplementedError def log_likelihood(self, y): """ The log marginal likelihood of observations y under the GP model """ raise NotImplementedError def sample(self, t=None, size=1): """ Draw samples from the GP prior distribution; Output shape: (t.size, size) """ raise NotImplementedError def predict(self, y, t=None, return_std=False, return_cov=False): """ Conditional predictive distribution of the GP model given observations y, evaluated at coordinates t. """ return self._GP.predict(y, t, return_cov=return_cov, return_var=return_std) def predict_with_hyperpars(self, results, sample, t=None, add_parts=True): """ Given the parameters in `sample`, return the GP predictive mean. If `t` is None, the prediction is done at the observed times and will contain other model components (systemic velocity, trend, instrument offsets) if `add_parts` is True. Otherwise, when `t` is given, the prediction is made at times `t` and *does not* contain these extra model components. """ raise NotImplementedError def sample_conditional(self, y, t, size=1): """ Sample from the conditional predictive distribution of the GP model given observations y, at coordinates t. """ raise NotImplementedError def sample_from_posterior(self, results, size=1): """ Given the posterior sample in `results`, take one sample of the GP hyperparameters and the white noise, and return a sample from the GP prior given those parameters. """ raise NotImplementedError def sample_with_hyperpars(self, results, sample, size=1): """ Given the value of the hyperparameters and the white noise in `sample`, return a sample from the GP prior. """ raise NotImplementedError def sample_conditional_from_posterior(self, results, t, size=1): """ Given the posterior sample in `results`, take one sample of the GP hyperparameters and the white noise, and return a sample from the GP predictive given those parameters. """ raise NotImplementedError def sample_conditional_with_hyperpars(self, results, sample, t, size=1): """ Given the value of the hyperparameters and the white noise in `sample`, return a sample from the GP predictive. """ raise NotImplementedError
# 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 print 'Simulating spectrum %d' % (i + 1) xpred = x_sim[i,:].flatten() mu, _ = gp.predict(y, xpred, return_var = True) flux_sim[i,:] = mu pl.figure(figsize = (8,6)) ax1 = pl.subplot(411) pl.plot(wav_rest_sim[:3,:].T, flux_sim[:3,:].T, '.', ms = 8, alpha=0.5,lw=0.5) pl.ylabel('rest') ax2 = pl.subplot(412, sharex = ax1, sharey = ax1) pl.plot(wav_earth_sim[:3,:].T, flux_sim[:3,:].T, '-', alpha=0.5,lw=0.5) pl.ylabel('true') # Add noise for j,snr in enumerate(SNR): print j, snr sigma = (flux_sim/snr) print np.median(sigma) flux_noisy = np.copy(flux_sim) for i in np.arange(NSMAX):
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 ]