def sample_mod(self):
'''
'''
# Sampling model
np.random.seed(42)
sampler = xo.PyMC3Sampler(finish=300, chains=4)
with self.model:
#burnin = sampler.tune(tune=800, start=self.soln, step_kwargs=dict(target_accept=0.9))
burnin = sampler.tune(tune=4500, start=self.soln, step_kwargs=dict(target_accept=0.9))
self.trace = sampler.sample(draws=2000)
def mcmc(self,
draws=500,
tune=500,
target_accept=0.9,
progressbar=False,
cores=4):
sampler = xo.PyMC3Sampler(finish=200)
with self.model:
#sampler.tune(tune=tune, start=self.map_soln, step=xo.get_dense_nuts_step(), step_kwargs=dict(target_accept=target_accept), progressbar=progressbar, cores=cores)
sampler.tune(tune=tune,
start=self.map_soln,
step_kwargs=dict(target_accept=target_accept),
progressbar=progressbar,
cores=cores)
trace = sampler.sample(draws=draws,
progressbar=progressbar,
cores=cores)
return trace
def best_model(model):
sampler = xo.PyMC3Sampler(finish=200)
with model:
sampler.tune(tune=2000, step_kwargs=dict(target_accept=0.9))
trace = sampler.sample(draws=2000)
return trace
def gp_rotation(self,
init_period=None,
tune=2000,
draws=2000,
prediction=True,
cores=None):
"""
Calculate a rotation period using a Gaussian process method.
Args:
init_period (Optional[float]): Your initial guess for the rotation
period. The default is the Lomb-Scargle period.
tune (Optional[int]): The number of tuning samples. Default is
2000.
draws (Optional[int]): The number of samples. Default is 2000.
prediction (Optional[Bool]): If true, a prediction will be
calculated for each sample. This is useful for plotting the
prediction but will slow down the whole calculation.
cores (Optional[int]): The number of cores to use. Default is
None (for running one process).
Returns:
gp_period (float): The GP rotation period in days.
errp (float): The upper uncertainty on the rotation period.
errm (float): The lower uncertainty on the rotation period.
logQ (float): The Q factor.
Qerrp (float): The upper uncertainty on the Q factor.
Qerrm (float): The lower uncertainty on the Q factor.
"""
self.prediction = prediction
x = np.array(self.time, dtype=float)
# Median of data must be zero
y = np.array(self.flux, dtype=float) - np.median(self.flux)
yerr = np.array(self.flux_err, dtype=float)
if init_period is None:
# Calculate ls period
init_period = self.ls_rotation()
with pm.Model() as model:
# The mean flux of the time series
mean = pm.Normal("mean", mu=0.0, sd=10.0)
# A jitter term describing excess white noise
logs2 = pm.Normal("logs2", mu=2 * np.log(np.min(yerr)), sd=5.0)
# The parameters of the RotationTerm kernel
logamp = pm.Normal("logamp", mu=np.log(np.var(y)), sd=5.0)
logperiod = pm.Normal("logperiod", mu=np.log(init_period), sd=5.0)
logQ0 = pm.Normal("logQ0", mu=1.0, sd=10.0)
logdeltaQ = pm.Normal("logdeltaQ", mu=2.0, sd=10.0)
mix = pm.Uniform("mix", lower=0, upper=1.0)
# Track the period as a deterministic
period = pm.Deterministic("period", tt.exp(logperiod))
# Set up the Gaussian Process model
kernel = xo.gp.terms.RotationTerm(log_amp=logamp,
period=period,
log_Q0=logQ0,
log_deltaQ=logdeltaQ,
mix=mix)
gp = xo.gp.GP(kernel, x, yerr**2 + tt.exp(logs2), J=4)
# Compute the Gaussian Process likelihood and add it into the
# the PyMC3 model as a "potential"
pm.Potential("loglike", gp.log_likelihood(y - mean))
# Compute the mean model prediction for plotting purposes
if prediction:
pm.Deterministic("pred", gp.predict())
# Optimize to find the maximum a posteriori parameters
self.map_soln = xo.optimize(start=model.test_point)
# print(self.map_soln)
# print(xo.utils.eval_in_model(model.logpt, self.map_soln))
# assert 0
# Sample from the posterior
np.random.seed(42)
sampler = xo.PyMC3Sampler()
with model:
print("sampling...")
sampler.tune(tune=tune,
start=self.map_soln,
step_kwargs=dict(target_accept=0.9),
cores=cores)
trace = sampler.sample(draws=draws, cores=cores)
# Save samples
samples = pm.trace_to_dataframe(trace)
self.samples = samples
self.period_samples = trace["period"]
self.gp_period = np.median(self.period_samples)
lower = np.percentile(self.period_samples, 16)
upper = np.percentile(self.period_samples, 84)
self.errm = self.gp_period - lower
self.errp = upper - self.gp_period
self.logQ = np.median(trace["logQ0"])
upperQ = np.percentile(trace["logQ0"], 84)
lowerQ = np.percentile(trace["logQ0"], 16)
self.Qerrp = upperQ - self.logQ
self.Qerrm = self.logQ - lowerQ
self.trace = trace
return self.gp_period, self.errp, self.errm, self.logQ, self.Qerrp, \
self.Qerrm
rho_save_sky = pm.Deterministic("rhoSaveSky", rho)
theta_save_sky = pm.Deterministic("thetaSaveSky", theta)
rho, theta = orbit.get_relative_angles(t_data, parallax)
rho_save_data = pm.Deterministic("rhoSaveData", rho)
theta_save_data = pm.Deterministic("thetaSaveData", theta)
# save RV plots
t_dense = pm.Deterministic("tDense", xs_phase * P + jd0)
rv1_dense = pm.Deterministic(
"RV1Dense",
conv * orbit.get_star_velocity(t_dense - jd0)[2] + gamma_keck)
rv2_dense = pm.Deterministic(
"RV2Dense",
conv * orbit.get_planet_velocity(t_dense - jd0)[2] + gamma_keck)
with model:
map_sol0 = xo.optimize(vars=[a_ang, phi])
map_sol1 = xo.optimize(map_sol0, vars=[a_ang, phi, omega, Omega])
map_sol2 = xo.optimize(map_sol1,
vars=[a_ang, logP, phi, omega, Omega, incl, e])
map_sol3 = xo.optimize(map_sol2)
# now let's actually explore the posterior for real
sampler = xo.PyMC3Sampler(finish=500, chains=4)
with model:
burnin = sampler.tune(tune=2000, step_kwargs=dict(target_accept=0.9))
trace = sampler.sample(draws=3000)
pm.backends.ndarray.save_trace(trace, directory="current", overwrite=True)
def PLD(tpf,
planet_mask=None,
aperture=None,
return_soln=False,
return_quick_corrected=False,
sigma=5,
trim=0,
ndraws=1000):
''' Use exoplanet, pymc3 and theano to perform PLD correction
Parameters
----------
tpf : lk.TargetPixelFile
Target Pixel File to Correct
planet_mask : np.ndarray
Boolean array. Cadences where planet_mask is False will be excluded from the
PLD correction. Use this to mask out planet transits.
'''
if planet_mask is None:
planet_mask = np.ones(len(tpf.time), dtype=bool)
if aperture is None:
aperture = tpf.pipeline_mask
time = np.asarray(tpf.time, np.float64)
if trim > 0:
flux = np.asarray(tpf.flux[:, trim:-trim, trim:-trim], np.float64)
flux_err = np.asarray(tpf.flux_err[:, trim:-trim, trim:-trim],
np.float64)
aper = np.asarray(aperture, bool)[trim:-trim, trim:-trim]
else:
flux = np.asarray(tpf.flux, np.float64)
flux_err = np.asarray(tpf.flux_err, np.float64)
aper = np.asarray(aperture, bool)
raw_flux = np.asarray(np.nansum(flux[:, aper], axis=(1)), np.float64)
raw_flux_err = np.asarray(
np.nansum(flux_err[:, aper]**2, axis=(1))**0.5, np.float64)
raw_flux_err /= np.median(raw_flux)
raw_flux /= np.median(raw_flux)
raw_flux -= 1
# Setting to Parts Per Thousand keeps us from hitting machine precision errors...
raw_flux *= 1e3
raw_flux_err *= 1e3
# Build the first order PLD basis
# X_pld = np.reshape(flux[:, aper], (len(flux), -1))
saturation = (np.nanpercentile(flux, 100, axis=0) > 175000)
X_pld = np.reshape(flux[:, aper & ~saturation], (len(tpf.flux), -1))
extra_pld = np.zeros((len(time), np.any(saturation, axis=0).sum()))
idx = 0
for column in saturation.T:
if column.any():
extra_pld[:, idx] = np.sum(flux[:, column, :], axis=(1, 2))
idx += 1
X_pld = np.hstack([X_pld, extra_pld])
# Remove NaN pixels
X_pld = X_pld[:, ~((~np.isfinite(X_pld)).all(axis=0))]
X_pld = X_pld / np.sum(flux[:, aper], axis=-1)[:, None]
# Build the second order PLD basis and run PCA to reduce the number of dimensions
X2_pld = np.reshape(X_pld[:, None, :] * X_pld[:, :, None], (len(flux), -1))
# Remove NaN pixels
X2_pld = X2_pld[:, ~((~np.isfinite(X2_pld)).all(axis=0))]
U, _, _ = np.linalg.svd(X2_pld, full_matrices=False)
X2_pld = U[:, :X_pld.shape[1]]
# Construct the design matrix and fit for the PLD model
X_pld = np.concatenate((np.ones((len(flux), 1)), X_pld, X2_pld), axis=-1)
# Create a matrix with small numbers along diagonal to ensure that
# X.T * sigma^-1 * X is not singular, and prevent it from being non-invertable
diag = np.diag((1e-8 * np.ones(X_pld.shape[1])))
if (~np.isfinite(X_pld)).any():
raise ValueError('NaNs in components.')
if (np.any(raw_flux_err == 0)):
raise ValueError('Zeros in raw_flux_err.')
def build_model(mask=None, start=None):
''' Build a PYMC3 model
Parameters
----------
mask : np.ndarray
Boolean array to mask cadences. Cadences that are False will be excluded
from the model fit
start : dict
MAP Solution from exoplanet
Returns
-------
model : pymc3.model.Model
A pymc3 model
map_soln : dict
Best fit solution
'''
if mask is None:
mask = np.ones(len(time), dtype=bool)
with pm.Model() as model:
# GP
# --------
logs2 = pm.Normal("logs2", mu=np.log(np.var(raw_flux[mask])), sd=4)
# pm.Potential("logs2_prior1", tt.switch(logs2 < -3, -np.inf, 0.0))
logsigma = pm.Normal("logsigma",
mu=np.log(np.std(raw_flux[mask])),
sd=4)
logrho = pm.Normal("logrho", mu=np.log(150), sd=4)
kernel = xo.gp.terms.Matern32Term(log_rho=logrho,
log_sigma=logsigma)
gp = xo.gp.GP(kernel, time[mask],
tt.exp(logs2) + raw_flux_err[mask]**2)
# Motion model
#------------------
A = tt.dot(X_pld[mask].T, gp.apply_inverse(X_pld[mask]))
A = A + diag # Add small numbers to diagonal to prevent it from being singular
B = tt.dot(X_pld[mask].T, gp.apply_inverse(raw_flux[mask, None]))
C = tt.slinalg.solve(A, B)
motion_model = pm.Deterministic("motion_model",
tt.dot(X_pld[mask], C)[:, 0])
# Likelihood
#------------------
pm.Potential("obs",
gp.log_likelihood(raw_flux[mask] - motion_model))
# gp predicted flux
gp_pred = gp.predict()
pm.Deterministic("gp_pred", gp_pred)
pm.Deterministic("weights", C)
# Optimize
#------------------
if start is None:
start = model.test_point
map_soln = xo.optimize(start=start, vars=[logsigma])
map_soln = xo.optimize(start=start, vars=[logrho, logsigma])
map_soln = xo.optimize(start=start, vars=[logsigma])
map_soln = xo.optimize(start=start, vars=[logrho, logsigma])
map_soln = xo.optimize(start=map_soln, vars=[logs2])
map_soln = xo.optimize(start=map_soln,
vars=[logrho, logsigma, logs2])
return model, map_soln, gp
# First rough correction
log.info('Optimizing roughly')
with silence():
model0, map_soln0, gp = build_model(mask=planet_mask)
# Remove outliers, make sure to remove a few nearby points incase of flares.
with model0:
motion = np.dot(X_pld, map_soln0['weights']).reshape(-1)
stellar = xo.eval_in_model(gp.predict(time), map_soln0)
corrected = raw_flux - motion - stellar
mask = ~sigma_clip(corrected, sigma=sigma).mask
mask = ~(convolve(mask, Box1DKernel(3), fill_value=1) != 1)
mask &= planet_mask
# Optimize PLD
log.info('Optimizing without outliers')
with silence():
model, map_soln, gp = build_model(mask, map_soln0)
lc_fig = _plot_light_curve(map_soln, model, mask, time, raw_flux,
raw_flux_err, X_pld, gp)
if return_soln:
motion = np.dot(X_pld, map_soln['weights']).reshape(-1)
with model:
stellar = xo.eval_in_model(gp.predict(time), map_soln)
return model, map_soln, motion, stellar
if return_quick_corrected:
raw_lc = tpf.to_lightcurve()
clc = lk.KeplerLightCurve(
time=time,
flux=(raw_flux - stellar - motion) * 1e-3 + 1,
flux_err=(raw_flux_err) * 1e-3,
time_format=raw_lc.time_format,
centroid_col=tpf.estimate_centroids()[0],
centroid_row=tpf.estimate_centroids()[0],
quality=raw_lc.quality,
channel=raw_lc.channel,
campaign=raw_lc.campaign,
quarter=raw_lc.quarter,
mission=raw_lc.mission,
cadenceno=raw_lc.cadenceno,
targetid=raw_lc.targetid,
ra=raw_lc.ra,
dec=raw_lc.dec,
label='{} PLD Corrected'.format(raw_lc.targetid))
return clc
# Burn in
sampler = xo.PyMC3Sampler()
with model:
burnin = sampler.tune(tune=np.max([int(ndraws * 0.3), 150]),
start=map_soln,
step_kwargs=dict(target_accept=0.9),
chains=4)
# Sample
with model:
trace = sampler.sample(draws=ndraws, chains=4)
varnames = ["logrho", "logsigma", "logs2"]
pm.traceplot(trace, varnames=varnames)
samples = pm.trace_to_dataframe(trace, varnames=varnames)
corner.corner(samples)
# Generate 50 realizations of the prediction sampling randomly from the chain
N_pred = 50
pred_mu = np.empty((N_pred, len(time)))
pred_motion = np.empty((N_pred, len(time)))
with model:
pred = gp.predict(time)
for i, sample in enumerate(
tqdm(xo.get_samples_from_trace(trace, size=N_pred),
total=N_pred)):
pred_mu[i] = xo.eval_in_model(pred, sample)
pred_motion[i, :] = np.dot(X_pld, sample['weights']).reshape(-1)
star_model = np.mean(pred_mu + pred_motion, axis=0)
star_model_err = np.std(pred_mu + pred_motion, axis=0)
raw_lc = tpf.to_lightcurve()
meta = {
'samples': samples,
'trace': trace,
'pred_mu': pred_mu,
'pred_motion': pred_motion
}
clc = lk.KeplerLightCurve(
time=time,
flux=(raw_flux - star_model) * 1e-3 + 1,
flux_err=((raw_flux_err**2 + star_model_err**2)**0.5) * 1e-3,
time_format=raw_lc.time_format,
centroid_col=tpf.estimate_centroids()[0],
centroid_row=tpf.estimate_centroids()[0],
quality=raw_lc.quality,
channel=raw_lc.channel,
campaign=raw_lc.campaign,
quarter=raw_lc.quarter,
mission=raw_lc.mission,
cadenceno=raw_lc.cadenceno,
targetid=raw_lc.targetid,
ra=raw_lc.ra,
dec=raw_lc.dec,
label='{} PLD Corrected'.format(raw_lc.targetid),
meta=meta)
return clc
def fit_planets(lc,
period_value,
t0_value,
depth_value,
R_star,
M_star,
T_star,
texp=0.0204335,
ndraws=1000):
'''Fit planet parameters using exoplanet'''
lc = lc.copy()
shape = len(period_value)
# if (np.asarray([len(period_value), len(t0_value), len(depth_value)]) == shape).all() == False
# raise ValueError('All planet parameters must have the same shape!')
x, y, yerr = np.asarray(lc.time, float), np.asarray(lc.flux,
float), np.asarray(
lc.flux_err, float)
yerr /= np.median(y)
y /= np.median(y)
y -= 1
y *= 1e3
yerr *= 1e3
def build_model(mask=None, start=None):
if mask is None:
mask = np.ones(len(x), dtype=bool)
with pm.Model() as model:
# Parameters for the stellar properties
mean = pm.Normal("mean", mu=0.0, sd=10.0)
u_star = xo.distributions.QuadLimbDark("u_star")
m_star = pm.Normal("m_star", mu=M_star[0], sd=M_star[1])
r_star = pm.Normal("r_star", mu=R_star[0], sd=R_star[1])
t_star = pm.Normal("t_star", mu=T_star[0], sd=T_star[1])
# Prior to require physical parameters
pm.Potential("m_star_prior", tt.switch(m_star > 0, 0, -np.inf))
pm.Potential("r_star_prior", tt.switch(r_star > 0, 0, -np.inf))
# Orbital parameters for the planets
logP = pm.Normal("logP",
mu=np.log(period_value),
sd=0.01,
shape=shape)
t0 = pm.Normal("t0", mu=t0_value, sd=0.01, shape=shape)
b = pm.Uniform("b", lower=0, upper=1, testval=0.5, shape=shape)
logr = pm.Normal("logr",
sd=1.0,
mu=0.5 * np.log(np.array(depth_value)) +
np.log(R_star[0]),
shape=shape)
r_pl = pm.Deterministic("r_pl", tt.exp(logr))
ror = pm.Deterministic("ror", r_pl / r_star)
# Tracking planet parameters
period = pm.Deterministic("period", tt.exp(logP))
# Orbit model
orbit = xo.orbits.KeplerianOrbit(r_star=r_star,
m_star=m_star,
period=period,
t0=t0,
b=b)
incl = pm.Deterministic('incl', orbit.incl)
a = pm.Deterministic('a', orbit.a)
teff = pm.Deterministic('teff', t_star * tt.sqrt(0.5 * (1 / a)))
# Compute the model light curve using starry
light_curves = xo.StarryLightCurve(u_star).get_light_curve(
orbit=orbit, r=r_pl, t=x[mask], texp=texp) * 1e3
light_curve = pm.math.sum(light_curves, axis=-1) + mean
pm.Deterministic("light_curves", light_curves)
pm.Normal('obs', mu=light_curve, sd=yerr[mask], observed=y[mask])
# Optimize
#------------------
if start is None:
start = model.test_point
map_soln = xo.optimize(start=start, vars=[logr])
map_soln = xo.optimize(start=map_soln, vars=[b])
map_soln = xo.optimize(start=map_soln, vars=[logP, t0])
map_soln = xo.optimize(start=map_soln, vars=[u_star])
map_soln = xo.optimize(start=map_soln, vars=[logr])
map_soln = xo.optimize(start=map_soln, vars=[b])
map_soln = xo.optimize(start=map_soln, vars=[mean])
map_soln = xo.optimize(start=map_soln)
return model, map_soln
with silence():
model0, map_soln0 = build_model()
corrected = y.copy()
for planet in map_soln0['light_curves'].T:
corrected -= planet
mask = ~sigma_clip(corrected, sigma_upper=4, sigma_lower=10).mask
mask = ~(convolve(mask, Box1DKernel(2), fill_value=1) != 1)
with silence():
model, map_soln = build_model(mask=mask, start=map_soln0)
# Burn in
sampler = xo.PyMC3Sampler()
with model:
burnin = sampler.tune(tune=np.max([int(ndraws * 0.3), 150]),
start=map_soln,
step_kwargs=dict(target_accept=0.9),
chains=4)
# Sample
with model:
trace = sampler.sample(draws=ndraws, chains=4)
return trace, mask
# Optimize
map_soln = pm.find_MAP(start=model.test_point, vars=[trend])
map_soln = pm.find_MAP(start=map_soln)
def check_convergence(samples):
tau = emcee.autocorr.integrated_time(samples, tol=0)
num = samples.shape[0] * samples.shape[1]
converged = np.all(tau * target_n_eff < num)
converged &= np.all(len(samples) > 50 * tau)
return converged, num / tau
# Run the PyMC3 sampler
chains = 2
sampler = xo.PyMC3Sampler(start=500, finish=500, window=500)
with model:
burnin = sampler.tune(tune=100000,
start=map_soln,
chains=chains,
cores=1,
progressbar=False)
tottime = 0
trace = None
with model:
while True:
strt = time.time()
trace = sampler.sample(draws=2000,
trace=trace,
chains=chains,