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
]
