Note

This tutorial was generated from an IPython notebook that can be downloaded here.

theano version: 1.0.4
pymc3 version: 3.6
exoplanet version: 0.1.5

Radial velocity fitting

In this tutorial, we will demonstrate how to fit radial velocity observations of an exoplanetary system using exoplanet. We will follow the getting started tutorial from the exellent RadVel package where they fit for the parameters of the two planets in the K2-24 system.

First, let’s download the data from RadVel:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

url = "https://raw.githubusercontent.com/California-Planet-Search/radvel/master/example_data/epic203771098.csv"
data = pd.read_csv(url, index_col=0)

x = np.array(data.t)
y = np.array(data.vel)
yerr = np.array(data.errvel)

plt.errorbar(x, y, yerr=yerr, fmt=".k")
plt.xlabel("time [days]")
plt.ylabel("radial velocity [m/s]");
../../_images/rv_3_0.png

Now, we know the periods and transit times for the planets from the K2 light curve, so let’s start by using the exoplanet.estimate_semi_amplitude() function to estimate the expected RV semi-amplitudes for the planets.

import exoplanet as xo

periods = [20.8851, 42.3633]
period_errs = [0.0003, 0.0006]
t0s = [2072.7948, 2082.6251]
t0_errs = [0.0007, 0.0004]
Ks = xo.estimate_semi_amplitude(periods, x, y, yerr, t0s=t0s)
print(Ks, "m/s")

[5.05069163 5.50983542] m/s

The radial velocity model in PyMC3

Now that we have the data and an estimate of the initial values for the parameters, let’s start defining the probabilistic model in PyMC3 (take a look at intro-to-pymc3 if you’re new to PyMC3). First, we’ll define our priors on the parameters:

import pymc3 as pm
import theano.tensor as tt

with pm.Model() as model:

    # Gaussian priors based on transit data (from Petigura et al.)
    t0 = pm.Normal("t0", mu=np.array(t0s), sd=np.array(t0_errs), shape=2)
    P = pm.Normal("P", mu=np.array(periods), sd=np.array(period_errs), shape=2)

    # Wide log-normal prior for semi-amplitude
    logK = pm.Normal("logK", mu=np.log(Ks), sd=10.0, shape=2)

    # This is a sanity check that restricts the semiamplitude to reasonable
    # values because things can get ugly as K -> 0
    pm.Potential("logK_bound", tt.switch(logK < 0, -np.inf, 0.0))

    # We also want to keep period physical but this probably won't be hit
    pm.Potential("P_bound", tt.switch(P <= 0, -np.inf, 0.0))

    # Eccentricity & argument of periasteron
    ecc = pm.Uniform("ecc", lower=0, upper=0.99, shape=2,
                     testval=np.array([0.1, 0.1]))
    omega = xo.distributions.Angle("omega", shape=2, testval=np.zeros(2))

    # Jitter & a quadratic RV trend
    logs = pm.Normal("logs", mu=np.log(np.median(yerr)), sd=5.0)
    trend = pm.Normal("trend", mu=0, sd=10.0**-np.arange(3)[::-1], shape=3)

Now we’ll define the orbit model:

with model:

    # Set up the orbit
    orbit = xo.orbits.KeplerianOrbit(
        period=P, t0=t0,
        ecc=ecc, omega=omega)

    # Set up the RV model and save it as a deterministic
    # for plotting purposes later
    vrad = orbit.get_radial_velocity(x, K=tt.exp(logK))
    pm.Deterministic("vrad", vrad)

    # Define the background model
    A = np.vander(x - 0.5*(x.min() + x.max()), 3)
    bkg = pm.Deterministic("bkg", tt.dot(A, trend))

    # Sum over planets and add the background to get the full model
    rv_model = pm.Deterministic("rv_model", tt.sum(vrad, axis=-1) + bkg)

For plotting purposes, it can be useful to also save the model on a fine grid in time.

t = np.linspace(x.min()-5, x.max()+5, 1000)

with model:
    vrad_pred = orbit.get_radial_velocity(t, K=tt.exp(logK))
    pm.Deterministic("vrad_pred", vrad_pred)
    A_pred = np.vander(t - 0.5*(x.min() + x.max()), 3)
    bkg_pred = pm.Deterministic("bkg_pred", tt.dot(A_pred, trend))
    rv_model_pred = pm.Deterministic("rv_model_pred",
                                     tt.sum(vrad_pred, axis=-1) + bkg_pred)

Now, we can plot the initial model:

plt.errorbar(x, y, yerr=yerr, fmt=".k")

with model:
    plt.plot(t, xo.eval_in_model(vrad_pred), "--k", alpha=0.5)
    plt.plot(t, xo.eval_in_model(bkg_pred), ":k", alpha=0.5)
    plt.plot(t, xo.eval_in_model(rv_model_pred), label="model")

plt.legend(fontsize=10)
plt.xlim(t.min(), t.max())
plt.xlabel("time [days]")
plt.ylabel("radial velocity [m/s]")
plt.title("initial model");
../../_images/rv_13_0.png

In this plot, the background is the dotted line, the individual planets are the dashed lines, and the full model is the blue line.

It doesn’t look amazing so let’s add in the likelihood and fit for the maximum a posterior parameters.

with model:

    err = tt.sqrt(yerr**2 + tt.exp(2*logs))
    pm.Normal("obs", mu=rv_model, sd=err, observed=y)

    map_soln = xo.optimize(start=model.test_point, vars=[trend])
    map_soln = xo.optimize(start=map_soln)

optimizing logp for variables: ['trend']
message: Optimization terminated successfully.
logp: -87.51672438985798 -> -72.63226415698486
optimizing logp for variables: ['trend', 'logs', 'omega_angle__', 'ecc_interval__', 'logK', 'P', 't0']
message: Desired error not necessarily achieved due to precision loss.
logp: -72.63226415698486 -> -22.056179633452878

plt.errorbar(x, y, yerr=yerr, fmt=".k")
plt.plot(t, map_soln["vrad_pred"], "--k", alpha=0.5)
plt.plot(t, map_soln["bkg_pred"], ":k", alpha=0.5)
plt.plot(t, map_soln["rv_model_pred"], label="model")

plt.legend(fontsize=10)
plt.xlim(t.min(), t.max())
plt.xlabel("time [days]")
plt.ylabel("radial velocity [m/s]")
plt.title("MAP model");
../../_images/rv_16_0.png

That looks better.

Sampling

Now that we have our model set up and a good estimate of the initial parameters, let’s start sampling. There are substantial covariances between some of the parameters so we’ll use a exoplanet.PyMC3Sampler to tune the sampler (see the pymc3-extras tutorial for more information).

np.random.seed(42)
sampler = xo.PyMC3Sampler(start=200, window=100, finish=300)
with model:
    burnin = sampler.tune(tune=3000, start=model.test_point,
                          step_kwargs=dict(target_accept=0.9))
    trace = sampler.sample(draws=3000)

Sampling 4 chains: 100%|██████████| 808/808 [00:10<00:00, 74.18draws/s]
Sampling 4 chains: 100%|██████████| 408/408 [00:05<00:00, 80.47draws/s]
Sampling 4 chains: 100%|██████████| 808/808 [00:04<00:00, 167.87draws/s]
Sampling 4 chains: 100%|██████████| 1608/1608 [00:03<00:00, 413.90draws/s]
Sampling 4 chains: 100%|██████████| 8408/8408 [00:15<00:00, 534.00draws/s]
Sampling 4 chains: 100%|██████████| 1208/1208 [00:02<00:00, 443.56draws/s]
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [trend, logs, omega, ecc, logK, P, t0]
Sampling 4 chains: 100%|██████████| 12000/12000 [00:18<00:00, 645.68draws/s]
There were 108 divergences after tuning. Increase target_accept or reparameterize.
The acceptance probability does not match the target. It is 0.788312912990659, but should be close to 0.9. Try to increase the number of tuning steps.
There were 10 divergences after tuning. Increase target_accept or reparameterize.
The acceptance probability does not match the target. It is 0.8402226010187613, but should be close to 0.9. Try to increase the number of tuning steps.
There were 27 divergences after tuning. Increase target_accept or reparameterize.
The acceptance probability does not match the target. It is 0.8258238424536315, but should be close to 0.9. Try to increase the number of tuning steps.
There were 28 divergences after tuning. Increase target_accept or reparameterize.
The acceptance probability does not match the target. It is 0.8206821919006299, but should be close to 0.9. Try to increase the number of tuning steps.
The number of effective samples is smaller than 10% for some parameters.

After sampling, it’s always a good idea to do some convergence checks. First, let’s check the number of effective samples and the Gelman-Rubin statistic for our parameters of interest:

pm.summary(trace, varnames=["trend", "logs", "omega", "ecc", "t0", "logK", "P"])
mean sd mc_error hpd_2.5 hpd_97.5 n_eff Rhat
trend__0 0.000949 0.000772 0.000012 -0.000563 0.002457 4878.820745 0.999918
trend__1 -0.038582 0.022634 0.000271 -0.082582 0.006290 8292.115487 1.000142
trend__2 -1.935402 0.815588 0.014953 -3.493321 -0.326454 2710.218672 1.000945
logs 1.041404 0.223987 0.003006 0.621729 1.493258 5310.698420 0.999961
omega__0 -0.288851 0.807204 0.018146 -1.500855 1.516015 1997.827922 1.000371
omega__1 -0.673141 2.101381 0.037115 -3.141370 2.917788 3310.355618 1.000380
ecc__0 0.240976 0.115811 0.003410 0.003181 0.451056 871.427379 1.003720
ecc__1 0.210433 0.165469 0.007705 0.000053 0.544293 283.094505 1.005294
t0__0 2072.794804 0.000697 0.000010 2072.793477 2072.796187 8464.441338 1.000201
t0__1 2082.625109 0.000400 0.000007 2082.624334 2082.625892 3738.995835 1.001570
logK__0 1.547680 0.260473 0.006368 1.023810 2.004923 1569.419195 1.002243
logK__1 1.583210 0.246858 0.008785 1.033171 2.034524 495.671764 1.002659
P__0 20.885105 0.000303 0.000006 20.884541 20.885718 2217.966801 1.002084
P__1 42.363296 0.000600 0.000007 42.362160 42.364478 8481.018450 1.000213

It looks like everything is pretty much converged here. Not bad for 14 parameters and about a minute of runtime…

Then we can make a corner plot of any combination of the parameters. For example, let’s look at period, semi-amplitude, and eccentricity:

import corner
samples = pm.trace_to_dataframe(trace, varnames=["P", "logK", "ecc"])
corner.corner(samples);
../../_images/rv_22_0.png

Finally, let’s plot the plosterior constraints on the RV model and compare those to the data:

plt.errorbar(x, y, yerr=yerr, fmt=".k")

# Compute the posterior predictions for the RV model
pred = np.percentile(trace["rv_model_pred"], [16, 50, 84], axis=0)
plt.plot(t, pred[1], color="C0", label="model")
art = plt.fill_between(t, pred[0], pred[2], color="C0", alpha=0.3)
art.set_edgecolor("none")

plt.legend(fontsize=10)
plt.xlim(t.min(), t.max())
plt.xlabel("time [days]")
plt.ylabel("radial velocity [m/s]")
plt.title("posterior constraints");
../../_images/rv_24_0.png

Phase plots

It might be also be interesting to look at the phased plots for this system. Here we’ll fold the dataset on the median of posterior period and then overplot the posterior constraint on the folded model orbits.

for n, letter in enumerate("bc"):
    plt.figure()

    # Get the posterior median orbital parameters
    p = np.median(trace["P"][:, n])
    t0 = np.median(trace["t0"][:, n])

    # Compute the median of posterior estimate of the background RV
    # and the contribution from the other planet. Then we can remove
    # this from the data to plot just the planet we care about.
    other = np.median(trace["vrad"][:, :, (n + 1) % 2], axis=0)
    other += np.median(trace["bkg"], axis=0)

    # Plot the folded data
    x_fold = (x - t0 + 0.5*p) % p - 0.5*p
    plt.errorbar(x_fold, y - other, yerr=yerr, fmt=".k")

    # Compute the posterior prediction for the folded RV model for this
    # planet
    t_fold = (t - t0 + 0.5*p) % p - 0.5*p
    inds = np.argsort(t_fold)
    pred = np.percentile(trace["vrad_pred"][:, inds, n], [16, 50, 84], axis=0)
    plt.plot(t_fold[inds], pred[1], color="C0", label="model")
    art = plt.fill_between(t_fold[inds], pred[0], pred[2], color="C0", alpha=0.3)
    art.set_edgecolor("none")

    plt.legend(fontsize=10)
    plt.xlim(-0.5*p, 0.5*p)
    plt.xlabel("phase [days]")
    plt.ylabel("radial velocity [m/s]")
    plt.title("K2-24{0}".format(letter));
../../_images/rv_26_0.png ../../_images/rv_26_1.png