Imports


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

Import the RV dataset generated by wobble


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

Define our priors


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()

Define our model and likelihood


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

Sample from the initial conditions a few times

... to verify we're off to a good start


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)

Let's do MCMC


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 [ ]: