def plot_sky(trace, m):
# plot sky position for a full orbit
xs_phase = np.linspace(0, 1, num=1000)
with m.model:
ts_full = xs_phase * m.P + m.t_periastron
predict_full = m.orbit.get_relative_angles(ts_full, m.parallax)
fig, ax = plt.subplots(nrows=1, figsize=(4, 4))
for sample in xo.get_samples_from_trace(trace, size=20):
# we'll want to cache these functions when we evaluate many samples
rho_full, theta_full = xo.eval_in_model(
predict_full, point=sample, model=m.model
)
x_full = rho_full * np.cos(theta_full) # X North
y_full = rho_full * np.sin(theta_full)
ax.plot(y_full, x_full, color="C0", lw=0.8, alpha=0.7)
xs = d.wds[1] * np.cos(d.wds[3]) # X is north
ys = d.wds[1] * np.sin(d.wds[3]) # Y is east
ax.plot(ys, xs, "ko")
ax.set_ylabel(r"$\Delta \delta$ ['']")
ax.set_xlabel(r"$\Delta \alpha \cos \delta$ ['']")
ax.invert_xaxis()
ax.plot(0, 0, "k*")
ax.set_aspect("equal", "datalim")
fig.subplots_adjust(left=0.18, right=0.82)
return fig
def evaluate_prediction_quantiles(event, pm_model, trace, t_grid, alert_time):
"""
Evaluates median model prediciton in data space.
"""
n_samples = 500
samples = xo.get_samples_from_trace(trace, size=n_samples)
with pm_model:
# Compute the trajectory of the lens
trajectory = ca.trajectory.Trajectory(
event, alert_time + T.exp(pm_model.ln_delta_t0), pm_model.u0,
pm_model.tE)
u_dense = trajectory.compute_trajectory(t_grid)
# Compute the magnification
mag_dense = (u_dense**2 + 2) / (u_dense * T.sqrt(u_dense**2 + 4))
F_base = 10**(-(pm_model.m_b - 22.0) / 2.5)
# Compute the mean model
prediction = pm_model.f * F_base * mag_dense + (1 -
pm_model.f) * F_base
# Evaluate model for each sample on a fine grid
n_pts_dense = T.shape(t_grid)[0].eval()
n_bands = len(event.light_curves)
prediction_eval = np.zeros((n_samples, n_bands, n_pts_dense))
# Evaluate predictions in model context
with pm_model:
for i, sample in enumerate(samples):
prediction_eval[i] = xo.eval_in_model(prediction, sample)
for i in range(n_bands):
q = np.percentile(prediction_eval[:, i, :], [16, 50, 84], axis=0)
return q
def plot_sep_pa(trace, m):
# plot separation and PA across the observed dates
ts_obs = np.linspace(
np.min(d.wds[0]) - 300, np.max(d.wds[0]) + 300, num=1000
) # days
with m.model:
predict_fine = m.orbit.get_relative_angles(ts_obs, m.parallax)
# we can plot the maximum posterior solution to see
# pkw = {'marker':".", "color":"k", 'ls':""}
ekw = {"color": "C1", "marker": "o", "ls": ""}
fig, ax = plt.subplots(nrows=2, sharex=True, figsize=(6, 4))
ax[0].set_ylabel(r'$\rho\,$ ["]')
ax[1].set_ylabel(r"P.A. [radians]")
ax[1].set_xlabel("JD [days]")
for sample in xo.get_samples_from_trace(trace, size=20):
# we'll want to cache these functions when we evaluate many samples
rho_fine, theta_fine = xo.eval_in_model(
predict_fine, point=sample, model=m.model
)
ax[0].plot(ts_obs, rho_fine, "C0")
ax[1].plot(ts_obs, theta_fine, "C0")
# get map sol for tot_rho_err
tot_rho_err = np.sqrt(d.wds[2] ** 2 + np.exp(2 * np.median(trace["logRhoS"])))
tot_theta_err = np.sqrt(d.wds[4] ** 2 + np.exp(2 * np.median(trace["logThetaS"])))
# ax[0].plot(d.wds[0], d.wds[1], **pkw)
ax[0].errorbar(d.wds[0], d.wds[1], yerr=tot_rho_err, **ekw)
# ax[1].plot(jds, theta_data, **pkw)
ax[1].errorbar(d.wds[0], d.wds[3], yerr=tot_theta_err, **ekw)
return fig
# truths=truth,
labels=["S1", "S2", "w1", "w2", "log_Q"],
)
fig.savefig('../results/test_results/gp_model/test_gp_corner.png')
plt.close('all')
#
# Generate 50 realizations of the prediction sampling randomly from the chain
#
N_pred = 50
pred_mu = np.empty((N_pred, len(true_t)))
pred_var = np.empty((N_pred, len(true_t)))
with model:
pred = gp.predict(true_t, return_var=True, predict_mean=True)
for i, sample in enumerate(
xo.get_samples_from_trace(trace, size=N_pred)):
pred_mu[i], pred_var[i] = xo.eval_in_model(pred, sample)
# Plot the predictions
for i in range(len(pred_mu)):
mu = pred_mu[i]
sd = np.sqrt(pred_var[i])
label = None if i else "prediction"
art = plt.fill_between(true_t, mu + sd, mu - sd, color="C1", alpha=0.1)
art.set_edgecolor("none")
plt.plot(true_t, mu, color="C1", label=label, alpha=0.1)
plt.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0, label="data")
plt.plot(true_t, true_y, "k", lw=1.5, alpha=0.3, label="truth")
plt.legend(fontsize=12, loc=2)
plt.xlabel("t")
)
_ = corner.corner(samples)
# %% [markdown]
# The "posterior predictive" plot that I like to make isn't the same as a "posterior predictive check" (which can be a good thing to do too).
# Instead, I like to look at the predictions of the model in the space of the data.
# We could have saved these predictions using a `pymc3.Deterministic` distribution, but that adds some overhead to each evaluation of the model so instead, we can use :func:`exoplanet.utils.get_samples_from_trace` to loop over a few random samples from the chain and then the :func:`exoplanet.eval_in_model` function to evaluate the prediction just for those samples.
# %%
# Generate 50 realizations of the prediction sampling randomly from the chain
N_pred = 50
pred_mu = np.empty((N_pred, len(true_t)))
pred_var = np.empty((N_pred, len(true_t)))
with model:
pred = gp.predict(true_t, return_var=True, predict_mean=True)
for i, sample in enumerate(xo.get_samples_from_trace(trace, size=N_pred)):
pred_mu[i], pred_var[i] = xo.eval_in_model(pred, sample)
# Plot the predictions
for i in range(len(pred_mu)):
mu = pred_mu[i]
sd = np.sqrt(pred_var[i])
label = None if i else "prediction"
art = plt.fill_between(true_t, mu + sd, mu - sd, color="C1", alpha=0.1)
art.set_edgecolor("none")
plt.plot(true_t, mu, color="C1", label=label, alpha=0.1)
plt.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0, label="data")
plt.plot(true_t, true_y, "k", lw=1.5, alpha=0.3, label="truth")
plt.legend(fontsize=12, loc=2)
plt.xlabel("t")
"""
phase = get_phase(data[0])
# evaluate the model with the current parameter settings
offset = offset_dict[label]
d = data[1] - offset
resid = d - model
err = np.sqrt(data[2]**2 + np.exp(2 * err_dict[label]))
color = color_dict[label]
a.errorbar(phase, d, yerr=err, label=label, **ekw, color=color)
a_r.errorbar(phase, resid, yerr=err, **ekw, color=color)
for sample in xo.get_samples_from_trace(trace, size=1):
# we'll want to cache these functions when we evaluate many samples
rv1_m = xo.eval_in_model(rv1, point=sample, model=m.model)
rv2_m = xo.eval_in_model(rv2, point=sample, model=m.model)
ax1.plot(xs_phase, rv1_m, **pkw)
ax2.plot(xs_phase, rv2_m, **pkw)
err_dict = {
"CfA": sample["logjittercfa"],
"Keck": sample["logjitterkeck"],
"FEROS": sample["logjitterferos"],
"du Pont": sample["logjitterdupont"],
}
offset_dict = {
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 plot_sep_pa(trace, m):
# Plot the astrometric fits.
# 1 square panel on left, then 3 rows of panels on right (rho, theta, v_B)
# set up the figure
# plot styles
pkw = {"marker": "o", "ms": 3, "color": "k", "ls": "", "zorder": 20}
lkw = {"ls": "-", "lw": 0.5}
ekw = {
"marker": "o",
"ms": 3,
"ls": "",
"elinewidth": 0.8,
"zorder": 20,
"color": "k",
}
kekw = {**ekw, "color": "maroon"}
xx = 7.1 # [in] textwidth
hmargin = 0.0
tmargin = 0.45
bmargin = 0.5
lmargin = 0.5
rmargin = 0.5
mmargin = 0.75
ax_width = (xx - lmargin - rmargin - mmargin) / 2
cax_frac = 0.7
cax_margin = 0.05
cax_height = 0.08
ax_height = ax_width
sax_height = (ax_height - hmargin) / 3
yy = ax_height + tmargin + bmargin
fig = plt.figure(figsize=(xx, yy))
ax_sky = fig.add_axes(
[lmargin / xx, bmargin / yy, ax_width / xx, ax_height / yy])
ax_sky.set_xlabel(r"$\Delta \alpha \cos \delta\; [{}^{\prime\prime}]$")
ax_sky.set_ylabel(r"$\Delta \delta\; [{}^{\prime\prime}]$")
# ax_sky.invert_xaxis()
cax = fig.add_axes([
(lmargin + (1 - cax_frac) * ax_width / 2) / xx,
(bmargin + ax_height + cax_margin) / yy,
cax_frac * ax_width / xx,
cax_height / yy,
])
# load and plot the gas image here; routine in gas.py
# frame width also set there
plot_gas(ax_sky, cax)
# set colorbar label
cax.tick_params(
axis="both",
labelsize="small",
labeltop=True,
labelbottom=False,
which="both",
direction="in",
bottom=False,
top=True,
pad=2,
)
cax.xaxis.set_label_position("top")
cax.set_xlabel(r"$v_\mathrm{BARY} \quad {\rm[km\,s^{-1}]}$", labelpad=2)
# plot the star
ax_sky.plot(0, 0, "*", ms=5, color="k", mew=0.1, zorder=99)
# plot the sky positions
X_ABs = d.wds[1] * np.cos(d.wds[3]) # north
Y_ABs = d.wds[1] * np.sin(d.wds[3]) # east
ax_sky.plot(Y_ABs, X_ABs, **pkw)
plot_errorbar(ax_sky,
d.wds[3],
d.wds[1],
d.wds[4],
d.wds[2],
color="C0",
lw=0.8)
# make the right-column plots
yr_lim = (1990, 2020)
ax_sep = fig.add_axes([
(lmargin + ax_width + mmargin) / xx,
(2 * (sax_height + hmargin) + bmargin) / yy,
ax_width / xx,
sax_height / yy,
])
ax_pa = fig.add_axes([
(lmargin + ax_width + mmargin) / xx,
(sax_height + hmargin + bmargin) / yy,
ax_width / xx,
sax_height / yy,
])
ax_V = fig.add_axes([
(lmargin + ax_width + mmargin) / xx,
bmargin / yy,
ax_width / xx,
sax_height / yy,
])
# get the mean value of all samples to use for errorbars and offset
sep_err = np.sqrt(d.wds[2]**2 + np.exp(2 * np.mean(trace["logRhoS"])))
ax_sep.errorbar(jd_to_year(d.wds[0]), d.wds[1], yerr=sep_err, **ekw)
ax_sep.set_ylabel(r"$\rho\;[{}^{\prime\prime}]$")
ax_sep.set_xlim(*yr_lim)
ax_sep.xaxis.set_ticklabels([])
pa_err = np.sqrt(d.wds[4]**2 + np.exp(2 * np.mean(trace["logThetaS"])))
ax_pa.errorbar(jd_to_year(d.wds[0]),
d.wds[3] / deg + 360,
yerr=pa_err / deg,
**ekw)
ax_pa.set_ylabel(r"$\theta\ [{}^\circ]$")
ax_pa.set_xlim(*yr_lim)
ax_pa.xaxis.set_ticklabels([])
VB_err = np.sqrt(d.keck3[2]**2 +
np.exp(2 * np.mean(trace["logjitterkeck"])))
ax_V.errorbar(
jd_to_year(d.keck3[0]),
d.keck3[1] - np.mean(trace["offsetKeck"]),
yerr=VB_err,
**kekw,
label="Keck",
)
ax_V.set_xlim(*yr_lim)
ax_V.set_ylim(6, 11)
ax_V.set_ylabel(r"$v_\mathrm{B} \quad [\mathrm{km\;s}^{-1}$]")
ax_V.set_xlabel("Epoch [yr]")
ax_V.legend(
loc="upper left",
fontsize="xx-small",
labelspacing=0.5,
handletextpad=0.2,
borderpad=0.4,
borderaxespad=1.0,
)
######
# Plot the orbits on the sep_pa.pdf figure
######
# define new Theano variables to get the
# 1) absolute X,Y,Z positions evaluated over the full orbital period
# 2) the sep, PA values evaluated over the observational window
# 3) the vB evaluated over the observational window
# phases for the full orbital period
phases = np.linspace(0, 1.0, num=500)
# times for the observational window (unchanging)
t_yr = Time(np.linspace(*yr_lim, num=500), format="byear").byear # [yrs]
t_obs = Time(np.linspace(*yr_lim, num=500),
format="byear").jd - jd0 # [days]
with m.model:
# the reason why we can refer to m.orbit_outer here (rather than, say)
# m.model.orbit_outer, is because orbit_outer is actually in the scope of the
# model module. It is in the *context* of the pymc3 model.model.
# so we need both, for it to make sense.
# convert phases to times stretching the full orbital period [days]
t_period = phases * m.P_outer # [days]
# 1) absolute X,Y,Z positions evaluated over the full orbital period
pos_outer = m.orbit_outer.get_relative_position(t_period, m.parallax)
# 2) the sep, PA values evaluated over the observational window
angles_outer = m.orbit_outer.get_relative_angles(t_obs, m.parallax)
# 3) the vB evaluated over the observational window
rv3 = (
conv * m.orbit_outer.get_planet_velocity(t_obs)[2] + m.gamma_outer
# do not include the offset, because we are plotting in the CfA frame
# + m.offset_keck
)
# iterate among several samples to plot the orbits on the sep_pa.pdf figure.
for sample in xo.get_samples_from_trace(trace, size=30):
vBs = xo.eval_in_model(rv3, point=sample, model=m.model)
# we can select orbits which have a higher vB[-1] than vB[0]
# and colorize them
# increasing = vBs[-1] > vBs[0]
if sample["increasing"]:
lkw["color"] = "C0"
ax_V.plot(t_yr, vBs, **lkw, zorder=1)
else:
lkw["color"] = "C1"
ax_V.plot(t_yr, vBs, **lkw, zorder=0)
rho, theta = xo.eval_in_model(angles_outer,
point=sample,
model=m.model)
ax_sep.plot(t_yr, rho, **lkw)
ax_pa.plot(t_yr, theta / deg + 360, **lkw)
X, Y, Z = xo.eval_in_model(pos_outer, point=sample, model=m.model)
ax_sky.plot(Y, X, **lkw)
return fig
def plot_interior_RV(trace, m):
"""
plot the interior RV
Args:
trace : pymc3 trace
m : the reference to the model module
"""
# choose the plot styles
color_dict = {"CfA": "C5", "Keck": "C2", "FEROS": "C3", "du Pont": "C4"}
hkw = {"lw": 1.0, "color": "0.4", "ls": ":"}
pkw = {"ls": "-", "lw": 0.5, "color": "C0"}
ekw = {"marker": ".", "ms": 2.5, "ls": "", "elinewidth": 0.6, "zorder": 20}
# set the figure dimensions
lmargin = 0.47
rmargin = lmargin
bmargin = 0.4
tmargin = 0.05
mmargin = 0.07
mmmargin = 0.15
ax_height = 1.0
ax_r_height = 0.4
# \textwidth=7.1in
# \columnsep=0.3125in
# column width = (7.1 - 0.3125)/2 = 3.393
xx = 3.393
ax_width = xx - lmargin - rmargin
yy = bmargin + tmargin + 2 * mmargin + mmmargin + 2 * ax_height + 2 * ax_r_height
# fill out the axes
fig = plt.figure(figsize=(xx, yy))
ax1 = fig.add_axes([
lmargin / xx, 1 - (tmargin + ax_height) / yy, ax_width / xx,
ax_height / yy
])
ax1_r = fig.add_axes([
lmargin / xx,
1 - (tmargin + mmargin + ax_height + ax_r_height) / yy,
ax_width / xx,
ax_r_height / yy,
])
ax1_r.axhline(0.0, **hkw)
ax2 = fig.add_axes([
lmargin / xx,
(bmargin + mmargin + ax_r_height) / yy,
ax_width / xx,
ax_height / yy,
])
ax2_r = fig.add_axes(
[lmargin / xx, bmargin / yy, ax_width / xx, ax_r_height / yy])
ax2_r.axhline(0.0, **hkw)
xs_phase = np.linspace(0, 1, num=500)
# define new Theano variables to get the continuous model RVs and the
# discrete models RVs to plot
with m.model:
# the reason why we can refer to m.P_inner here (rather than, say)
# m.model.P_inner, is because P_inner is actually in the scope of the
# model module. It is in the *context* of the pymc3 model.model.
# so we need both, for it to make sense.
ts_phases = xs_phase * m.P_inner + m.t_periastron_inner # new theano var
rv1, rv2 = m.get_RVs(ts_phases, ts_phases, 0.0)
# We want to plot the model and data phased to a single reference orbit
# this means we need to take out any additional motion to gamma_A relative
# to the reference orbit from both the model and data
gamma_A_epoch = m.get_gamma_A(ts_phases[0])
err_dict = {
"CfA": m.logjit_cfa,
"Keck": m.logjit_keck,
"FEROS": m.logjit_feros,
"du Pont": m.logjit_dupont,
}
offset_dict = {
"CfA": 0.0,
"Keck": m.offset_keck,
"FEROS": m.offset_feros,
"du Pont": m.offset_dupont,
}
def get_centered_rvs(label, data1, data2):
"""
Return the data and model discrete velocities centered on the reference
epoch bary velocity
"""
t1 = tt.as_tensor_variable(data1[0])
t2 = tt.as_tensor_variable(data2[0])
offset = offset_dict[label]
d1_corrected = data1[1] - offset - m.get_gamma_A(
t1) + gamma_A_epoch
d2_corrected = data2[1] - offset - m.get_gamma_A(
t2) + gamma_A_epoch
rv1_base, rv2_base = m.get_RVs(t1, t2, 0.0)
rv1_corrected = rv1_base - m.get_gamma_A(t1) + gamma_A_epoch
rv2_corrected = rv2_base - m.get_gamma_A(t2) + gamma_A_epoch
e1 = np.sqrt(data1[2]**2 + np.exp(2 * err_dict[label]))
e2 = np.sqrt(data2[2]**2 + np.exp(2 * err_dict[label]))
return (
t1,
d1_corrected,
e1,
rv1_corrected,
t2,
d2_corrected,
e2,
rv2_corrected,
)
# repack to yield a set of corrected data, errors, and corrected model velocities
rv_cfa0 = get_centered_rvs("CfA", d.cfa1, d.cfa2)
rv_keck0 = get_centered_rvs("Keck", d.keck1, d.keck2)
rv_feros0 = get_centered_rvs("FEROS", d.feros1, d.feros2)
rv_dupont0 = get_centered_rvs("du Pont", d.dupont1, d.dupont2)
# data, then phase them, then plot them.
def phase_and_plot(rv_package, label):
"""
Calculate inner orbit radial velocities and residuals for a given dataset and star.
Args:
data: tuple containing (date, rv, err)
model: model rvs evaluated at date (with no offset)
a: primary matplotlib axes
a_r: residual matplotlib axes
label: instrument that acquired data
Returns:
None
"""
(
t1,
d1_corrected,
e1,
rv1_corrected,
t2,
d2_corrected,
e2,
rv2_corrected,
) = rv_package
phase1 = get_phase(t1)
phase2 = get_phase(t2)
# evaluate the model with the current parameter settings
color = color_dict[label]
ax1.errorbar(phase1,
d1_corrected,
yerr=e1,
label=label,
**ekw,
color=color)
ax1_r.errorbar(phase1,
d1_corrected - rv1_corrected,
yerr=e1,
**ekw,
color=color)
ax2.errorbar(phase2,
d2_corrected,
yerr=e2,
label=label,
**ekw,
color=color)
ax2_r.errorbar(phase2,
d2_corrected - rv2_corrected,
yerr=e2,
**ekw,
color=color)
# plot many samples for the RV trace
for sample in xo.get_samples_from_trace(trace, size=10):
rv1_m = xo.eval_in_model(rv1, point=sample, model=m.model)
rv2_m = xo.eval_in_model(rv2, point=sample, model=m.model)
ax1.plot(xs_phase, rv1_m, **pkw)
ax2.plot(xs_phase, rv2_m, **pkw)
# get just a single sample to plot the data and model
for sample in xo.get_samples_from_trace(trace, size=1):
# create a phasing function local to these sampled values of
# t_periastron and period
def get_phase(dates):
return ((dates - sample["tPeriastronInner"]) %
sample["PInner"]) / sample["PInner"]
# Plot the data and residuals for this sample
phase_and_plot(xo.eval_in_model(rv_cfa0, point=sample, model=m.model),
"CfA")
phase_and_plot(xo.eval_in_model(rv_keck0, point=sample, model=m.model),
"Keck")
phase_and_plot(
xo.eval_in_model(rv_feros0, point=sample, model=m.model), "FEROS")
phase_and_plot(
xo.eval_in_model(rv_dupont0, point=sample, model=m.model),
"du Pont")
ax1.legend(
loc="upper left",
fontsize="xx-small",
labelspacing=0.5,
handletextpad=0.2,
borderpad=0.4,
borderaxespad=1.0,
)
ax1.set_ylabel(r"$v_\mathrm{Aa}$ [$\mathrm{km s}^{-1}$]", labelpad=0)
ax1_r.set_ylabel(r"$O-C$", labelpad=-1)
ax2.set_ylabel(r"$v_\mathrm{Ab}$ [$\mathrm{km s}^{-1}$]", labelpad=-5)
ax2_r.set_ylabel(r"$O-C$", labelpad=-2)
ax2_r.set_xlabel("phase")
ax = [ax1, ax1_r, ax2, ax2_r]
for a in ax:
a.set_xlim(0, 1)
ax1.xaxis.set_ticklabels([])
ax1_r.xaxis.set_ticklabels([])
ax2.xaxis.set_ticklabels([])
return fig
