In [ ]:
import numpy as np
import emcee
from tqdm import tqdm
import corner
In [ ]:
import starry
from ylm_rot import get_ylm_coeffs
print("Using `starry` version %s." % starry.__version__)
In [ ]:
import matplotlib
import matplotlib.pyplot as plt
%matplotlib inline
%config InlineBackend.figure_format = "retina"
from matplotlib import rcParams
rcParams["savefig.dpi"] = 100
rcParams["figure.dpi"] = 100
matplotlib.rcParams['axes.formatter.useoffset'] = False
In [ ]:
time, rv, err = np.loadtxt("HD189733_rvs.txt", unpack=True, skiprows=1, delimiter=',')
idx = np.argsort(time)
time = time[idx]
rv = rv[idx]
err = err[idx]
tref = 2454279.0
In [ ]:
class Normal(object):
"""A normal distribution."""
def __init__(self, mean, sigma):
self.mean = mean
self.sigma = sigma
def sample(self):
return self.mean + self.sigma * np.random.randn()
def evaluate(self, x):
return -0.5 * (x - self.mean) ** 2 / self.sigma ** 2
class Uniform(object):
"""A uniform distribution."""
def __init__(self, low, high):
self.low = low
self.high = high
def sample(self):
return np.random.uniform(self.low, self.high)
def evaluate(self, x):
if (x < self.low) or (x > self.high):
return -np.inf
else:
return 0
class Sine(object):
"""
A sine distribution.
This is an uninformative distribution for the
inclination of the stellar rotation axis.
"""
def __init__(self):
pass
def sample(self):
x = np.random.random()
y = np.arccos(1 - x) * 180 / np.pi
z = np.random.random()
if z < 0.5:
return 180 - y
else:
return y
def evaluate(self, x):
if x < 0 or x > 180:
return -np.inf
else:
return np.log10(np.sin(x * np.pi / 180))
In [ ]:
class Prior:
"""A class containing all the information on our priors."""
def __init__(self):
self.params = ['veq', 'obl', 'inc', 'alpha', 'q1', 'q2',
'baseline', 'K', 'b_t0', 'b_per',
'b_inc', 'b_r', 'b_a', 'ln_err']
self.names = [r'$v_\mathrm{eq}$',
r'$\lambda$',
r'$i$',
r'$\alpha$',
r'$q_1$',
r'$q_2$',
r'$v_\mathrm{0}$',
r'$K$',
r'$t_\mathrm{0,p}$',
r'$P_\mathrm{p}$',
r'$i_\mathrm{p}$',
r'$r_\mathrm{p}$',
r'$a$',
r'$\ln\sigma$']
# The actual priors we use.
# Some of these taken from Table 1 in Cegla et al. (2016)
# Others are uniform / uninformative priors
self.veq = Uniform(0, 1e4)
self.obl = Uniform(-90, 90)
self.inc = Sine()
self.alpha = Uniform(-1, 1)
self.q1 = Uniform(0, 1)
self.q2 = Uniform(0, 1)
self.baseline = Uniform(-10000, -9900)
self.K = Uniform(180, 220)
self.b_t0 = Normal(2454279.436714, 0.000015)
self.b_per = Normal(2.21857567, 0.00000015)
self.b_inc = Normal(85.710, 0.024)
self.b_r = Normal(0.15667, 0.00012)
self.b_a = Normal(8.863, 0.020)
self.ln_err = Uniform(-3, 3)
# Distributions for the initial MCMC step
self.veq_guess = Normal(4500, 100)
self.obl_guess = Normal(-0.4, 0.3)
self.inc_guess = Normal(100, 2.0)
self.alpha_guess = Uniform(-0.5, 0.5)
self.q1_guess = Uniform(0, 1)
self.q2_guess = Uniform(0, 1)
self.baseline_guess = Normal(-9950.0, 3.0)
self.K_guess = Normal(200.56, 0.88)
self.b_t0_guess = Normal(2454279.436714, 0.000015)
self.b_per_guess = Normal(2.21857567, 0.00000015)
self.b_inc_guess = Normal(85.710, 0.024)
self.b_r_guess = Normal(0.15667, 0.00012)
self.b_a_guess = Normal(8.863, 0.020)
self.ln_err_guess = Normal(0, 0.1)
def evaluate(self, x):
"""Evaluate the log prior."""
return np.sum([getattr(self, p).evaluate(x[i]) for i, p in enumerate(self.params)])
def sample(self):
"""Sample from the prior distribution."""
return [getattr(self, p).sample() for i, p in enumerate(self.params)]
def guess(self):
"""Sample from the `guess' distribution."""
return [getattr(self, p + "_guess").sample() for i, p in enumerate(self.params)]
In [ ]:
# Instantiate the prior
prior = Prior()
In [ ]:
# Instantiate our `starry` system
star = starry.kepler.Primary(5)
planet = starry.kepler.Secondary(0)
system = starry.kepler.System(star, planet)
map_unif = starry.Map(2)
In [ ]:
def compute(x):
"""Compute the RV model given a parameter vector `x`."""
# Get our params
veq, obl, inc, alpha, q1, q2, baseline, K, b_t0, b_per, b_inc, b_r, b_a, ln_err = x
# Planet params
planet.tref = b_t0
planet.porb = b_per
planet.inc = b_inc
planet.r = b_r
planet.a = b_a
# Stellar brightness-weighted velocity profile
star.reset()
star[:3, :] = get_ylm_coeffs(veq=veq, inc=inc, obl=obl, alpha=np.abs(alpha))
star[0, 0] = 1
sqrtq1 = np.sqrt(q1)
u1 = 2 * sqrtq1 * q2
u2 = sqrtq1 * (1 - 2 * q2)
star[1] = u1
star[2] = u2
# Compute the integral of the brightness-weighted velocity field.
# As we explain in `DifferentialRotationWithSphericalHarmonics.ipynb`,
# what we're actually computing here is the integral of (Iv + I)
system.compute(time)
intIv_plus_I = star.lightcurve
# Compute the integral of the brightness, I; this is the
# RM effect normalization.
map_unif[0, 0] = 1
map_unif[1] = u1
map_unif[2] = u2
intI = map_unif.flux(xo=planet.X, yo=planet.Y, ro=planet.r)
# The RM effect is the integral of Iv divided by the integral of I
# Note that we must subtract out the I term from the numerator
model = (intIv_plus_I - intI) / intI
# Add a baseline
model += baseline - K * np.sin(2 * np.pi / b_per * (time - b_t0))
return model
In [ ]:
def lnlike(x):
"""Return the prior times the log-likelihood for a parameter vector `x`."""
# Evaluate the prior
lp = prior.evaluate(x)
if np.isinf(lp):
return lp
# Compute the model
model = compute(x)
# Compute the likelihood from chi-squared
ln_err = x[-1]
N = len(time)
return lp - 0.5 * np.sum(((model - rv) / np.exp(ln_err)) ** 2) - N * ln_err
In [ ]:
fig = plt.figure(figsize=(14, 6))
plt.plot(time - tref, rv, '.')
plt.xlabel("Time [BJD - %d]" % tref, fontsize=16)
plt.ylabel("RV [m / s]", fontsize=16);
for i in range(100):
model = compute(prior.guess())
plt.plot(time - tref, model, 'C1', alpha=0.1)
In [ ]:
# Chain params
nsteps = 150000
nwalk = 100
ndim = 14
thin = 10
nburn = int(nsteps / 5)
# Run the chain, or load from disk?
RUN_CHAIN = False
In [ ]:
# Instantiate the sampler
sampler = emcee.EnsembleSampler(nwalk, ndim, lnlike, threads=8)
if RUN_CHAIN:
# Initial guesses
p0 = [prior.guess() for k in range(nwalk)]
# Run our MCMC chain
for i in tqdm(sampler.sample(p0, iterations=nsteps * thin, thin=thin), total=nsteps * thin):
pass
np.savez("hd189_rm.npz", chain=sampler.chain,
lnprobability=sampler.lnprobability, iterations=sampler.iterations,
naccepted=sampler.naccepted)
else:
# Load the data
data = np.load("hd189_rm.npz")
sampler._chain = data['chain']
sampler._lnprob = data['lnprobability']
sampler.iterations = int(data['iterations'])
sampler.naccepted = data['naccepted']
In [ ]:
# Transform some of the params
ndim = sampler.dim
labels = list(prior.names)
chain = np.array(sampler.chain)
lnprobability = np.array(sampler.lnprobability)
# Take the absolute value of alpha
chain[:, :, 3] = np.abs(chain[:, :, 3])
# Subtract the reference time from t0
chain[:, :, 8] -= tref
prior.b_t0.mean -= tref
# Compute the total 3D obliquity as in Cegla et al.
inc = chain[:, :, 2]
obl = chain[:, :, 1]
b_inc = chain[:, :, 10]
obl3d = np.arccos(np.sin(inc * np.pi / 180) *
np.cos(obl * np.pi / 180) *
np.sin(b_inc * np.pi / 180) +
np.cos(inc * np.pi / 180) *
np.cos(b_inc * np.pi / 180)).reshape(nwalk, nsteps, 1)
obl3d *= 180 / np.pi
# Compute u1 and u2
q1 = chain[:, :, 4]
q2 = chain[:, :, 5]
sqrtq1 = np.sqrt(q1)
u1 = (2 * sqrtq1 * q2).reshape(nwalk, nsteps, 1)
u2 = (sqrtq1 * (1 - 2 * q2)).reshape(nwalk, nsteps, 1)
chain = np.concatenate((chain, obl3d, u1, u2), axis=-1)
ndim += 3
labels += [r"$\psi$", r"$u_1$", r"$u_2$"]
In [ ]:
# Plot the chains
fig = plt.figure(figsize=(12, 8))
fig.subplots_adjust(bottom=0.05, top=0.95, hspace=0.1)
axc = [plt.subplot2grid((9, 23), (n, 0), colspan=8, rowspan=1) for n in range(9)]
axc += [plt.subplot2grid((9, 23), (n, 13), colspan=8, rowspan=1) for n in range(9)]
axh = [plt.subplot2grid((9, 23), (n, 8), colspan=2, rowspan=1, sharey=axc[n]) for n in range(9)]
axh += [plt.subplot2grid((9, 23), (n, 21), colspan=2, rowspan=1, sharey=axc[9 + n]) for n in range(9)]
alpha = 0.3
# Thin them a little
plot_thin = 100
burn_idx = np.arange(0, nburn, plot_thin)
prod_idx = np.arange(nburn, nsteps, plot_thin)
for i, label in enumerate(labels):
# Plot the production part of the chains and fix the axis limits
for k in range(nwalk):
axc[i].plot(prod_idx, chain[k, nburn:, i][::plot_thin], alpha=alpha, lw=1)
axc[i].set_ylabel(label, fontsize=18)
axc[i].set_ylim(*axc[i].get_ylim())
axc[i].set_xticklabels([])
# Go back and plot the burn-in part
axc[i].set_prop_cycle(None)
for k in range(nwalk):
axc[i].plot(burn_idx, chain[k, :nburn, i][::plot_thin], alpha=alpha, lw=1)
axc[i].margins(0, None)
axc[i].axvline(nburn, color="r", lw=1)
# Plot the histogram for the production part
axh[i].hist(chain[:, nburn:, i].flatten(), bins=30,
orientation="horizontal", histtype='step',
fill=False, color='k', lw=1)
plt.setp(axh[i].get_yticklabels(), visible=False)
plt.setp(axh[i].get_xticklabels(), visible=False)
# Plot the likelihood
for k in range(nwalk):
axc[-1].plot(prod_idx, lnprobability[k, nburn:][::plot_thin], alpha=alpha, lw=1)
axc[-1].set_ylabel(r"$\ln P$", fontsize=18)
axc[-1].set_ylim(*axc[-1].get_ylim())
axc[-1].set_prop_cycle(None)
for k in range(nwalk):
axc[-1].plot(burn_idx, lnprobability[k, :nburn][::plot_thin], alpha=alpha, lw=1)
axc[-1].margins(0, None)
axc[-1].axvline(nburn, color="r", lw=1)
axc[-1].set_xticklabels([])
axh[-1].axis('off')
plt.setp(axh[-1].get_yticklabels(), visible=False)
plt.setp(axh[-1].get_xticklabels(), visible=False);
for ax in axc:
for tick in ax.yaxis.get_major_ticks():
tick.label.set_fontsize(6)
fig.align_ylabels(axc[:9]);
fig.align_ylabels(axc[9:]);
# Save!
fig.savefig("hd189_chains.pdf", bbox_inches="tight")
In [ ]:
trefprime = tref + 62
rv_offset = np.median(rv)
COLOR3 = '#ff7f0e'
# Plot the data + model
fig, ax = plt.subplots(2, 1, figsize=(12, 8), sharex=True, gridspec_kw = {'height_ratios':[2, 1]})
ax[0].errorbar(time - trefprime, rv - rv_offset, err, ms=5, fmt='o', color='k', zorder=10)
ax[0].set_ylabel(r"RV (m s$^{-1}$)", fontsize=16)
# Plot 500 random samples
for i in range(500):
idx = np.random.randint(0, nwalk * nsteps)
model = compute(sampler.flatchain[idx])
ax[0].plot(time - trefprime, model - rv_offset, color=COLOR3, alpha=0.05, lw=1)
# Plot the max like model
idx = np.argmax(sampler.flatlnprobability)
model = compute(sampler.flatchain[idx])
ax[0].plot(time - trefprime, model - rv_offset, color='k', lw=1, label="MAP");
# Hack a legend
'''
ax[0].set_xlim(*ax[0].get_xlim())
ax[0].set_ylim(*ax[0].get_ylim())
ax[0].plot([0, 0], [0, 0], color=COLOR3, alpha=1, lw=1, label="500 random samples")
ax[0].legend(fontsize=12)
'''
# Plot the residuals
#ax[1].plot(time - trefprime, rv - model, 'k.');
#ax[1].plot(time - trefprime, rv - model, 'k-', color='k', lw=1);
ax[1].errorbar(time - trefprime, rv - model, err, fmt='o', ms=5, color='k')
ax[1].axhline(0, color='k', ls='--', alpha=0.5);
ax[1].set_xlabel("Time (BJD - %d)" % trefprime, fontsize=16, labelpad=20)
ax[1].set_ylabel(r"Residuals (m s$^{-1}$)", fontsize=16);
fig.align_ylabels()
for axis in ax:
for tick in axis.xaxis.get_major_ticks() + axis.yaxis.get_major_ticks():
tick.label.set_fontsize(14)
fig.tight_layout()
fig.subplots_adjust(hspace=.05)
# Save!
fig.savefig("hd189_samples.pdf")
fig.savefig("../paper/figures/hd189_rm.pdf");
In [ ]:
# Plot the corner plot for all params
samples = chain[:, nburn:, :].reshape(nwalk * (nsteps - nburn), ndim)
# Thin them a tiny bit
corner_thin = 10
samples = samples[::corner_thin]
# Plot
fig = corner.corner(samples, labels=labels, bins=50, plot_datapoints=False);
# Plot the priors
for ind in range(ndim):
if ind < len(prior.params):
x = np.array([getattr(prior, prior.params[ind]).sample()
for i in range(len(samples))])
ax = fig.axes[(ndim + 1) * ind]
n, _, _ = ax.hist(x, bins=50, weights=None,
range=np.sort([x.min(), x.max()]),
histtype="step", color="C0")
for ax in fig.axes:
ax.xaxis.label.set_fontsize(20)
ax.yaxis.label.set_fontsize(20)
# Save!
fig.savefig("hd189_corner.pdf", bbox_inches="tight")
In [ ]:
# Reproduce Figure 4 in Cegla et al.
inds = [3, 0, 2, 1, 14]
samples_cegla = np.array([samples[:, inds[0]],
samples[:, inds[1]],
samples[:, inds[2]],
samples[:, inds[3]],
samples[:, inds[4]]]).T
labels_cegla = [labels[i] for i in inds]
# Corner plot
fig = corner.corner(samples_cegla, labels=labels_cegla, bins=50, plot_datapoints=False);
for ax in fig.axes:
ax.xaxis.label.set_fontsize(20)
ax.yaxis.label.set_fontsize(20)
# Save!
fig.savefig("hd189_corner_cegla.pdf", bbox_inches="tight")
In [ ]: