Radial Velocities

In this tutorial we will learn how to use jaxoplanet to compute the radial velocities (RVs) of a star hosting a single exoplanet, and how to fit this dataset using numpyro.

Note

This tutorial requires some extra packages that are not included in the jaxoplanet dependencies.

Setup

We first setup the number of CPUs to use and enable the use of double-precision numbers with jax.

import jax
import numpyro

numpyro.set_host_device_count(2)
jax.config.update("jax_enable_x64", True)

Model and dataset

Let’s first generate our dataset, consisting in the radial velocities of a star orbited by a unique exoplanet. We start by defining the system (see for more details)

Important

Note that the radial velocity is returned in Rsun/day. Checkout the units convention to learn more about units in jaxoplanet and keep reading to see how you can easily convert velocities to m/s.

import numpy as np
import matplotlib.pyplot as plt
from jaxoplanet.orbits import keplerian
from jaxoplanet import constants

# Convert 1.5 Mjup to Msun for jaxoplanet
truth = dict(mass=1.5 * constants.M_jup, period=3.0)
star = keplerian.Central(mass=1.0, radius=1.0)
system = keplerian.System(star).add_body(**truth)

As we left many parameters as default, let’s check the parameters of the system

system

System(
  central=Central(mass=1.0, radius=1.0, density=0.238732414637843),
  _body_stack=ObjectStack(...)
)

We can now compute the radial velocities of the star and add some noise to simulate our dataset

Note that jaxoplanet follows a consistent units convention and returns RVs in Rsun/day. Since it is much more common to work with RVs in m/s for exoplanet research, we will convert all velocities in m/s for this tutorial. The jaxoplanet.constants module has a helper m_per_s constant to help us do that.

np.random.seed(10)
over_time = np.linspace(0, 10, 1000)
over_rvs = system.radial_velocity(over_time)[0] / constants.m_per_s
time = np.sort(np.random.uniform(0, 10, 40))
rv_obs = system.radial_velocity(time)[0] / constants.m_per_s
rv_err = 20.0  # m/s
rv_obs += rv_err * np.random.normal(size=len(time))


def plot_data():
    plt.errorbar(time, rv_obs, yerr=rv_err, fmt="+k")
    plt.xlabel("time (days)")
    plt.ylabel(r"radial velocity (m/s)")


plt.plot(over_time, over_rvs)
plot_data()

Inference

We will infer the value and associated uncertainty of the system orbital parameters using numpyro. In order to do that we first define a callable model function

Since it is more common to fit the planet mass in Jupiter masses, we convert the planet mass internally in the model via the M_jup constant so that the parameters can be expressed in Jupiter masses instead of solar masses, which jaxoplanet expects.

from numpyro import distributions as dist, infer


def rv_model(time, params):
    system = keplerian.System(star).add_body(
        mass=params["mass"] * constants.M_jup, period=params["period"]
    )
    return system.radial_velocity(time)[0] / constants.m_per_s


def model(time, y=None):
    mass = numpyro.sample("mass", dist.Uniform(0.1, 10.0))
    period = numpyro.sample("period", dist.Uniform(1.0, 10.0))
    error = numpyro.sample("error", dist.Uniform(0.1, 100.0))
    rv = rv_model(time, {"mass": mass, "period": period})

    # the likelihood function
    numpyro.sample("y", dist.Normal(rv, error), obs=y)

Once the model defined, we will sample the posterior likelihood of the model parameters given the observed radial velocities. As we will need to provide some initial values for these parameters, it is a good idea to check that these values provide a good starting point for the model.

init_values = {"mass": 1.0, "period": 3.01, "error": rv_err}
init_model = rv_model(over_time, init_values)

plt.plot(over_time, init_model, "C0")
plot_data()

Starting from this initial guess, we use numpyro’s MCMC No-U-Turn Sampler (NUTS) to sample the posterior likelihood of the parameters.

sampler = infer.MCMC(
    infer.NUTS(model, init_strategy=infer.init_to_value(values=init_values)),
    num_warmup=2000,
    num_samples=10000,
    progress_bar=True,
    num_chains=2,
)

sampler.run(jax.random.PRNGKey(6), time, y=rv_obs)

We can plot the inferred model with

samples = sampler.get_samples()
posterior_rvs = infer.Predictive(model, samples)(jax.random.PRNGKey(0), over_time)["y"]

plot_data()
plt.plot(over_time, posterior_rvs.mean(0), "C0")
_ = plt.fill_between(
    over_time,
    *np.percentile(posterior_rvs, [16, 84], axis=0),
    alpha=0.3,
    color="C0",
)

and check the inferred model parameters in a corner plot

import arviz as az
import corner
import arviz as az

idata = az.from_numpyro(sampler)

idata = az.from_numpyro(sampler)

import arviz as az

idata = az.from_numpyro(sampler)

_ = corner.corner(
    idata,
    var_names=["mass", "period"],
    truths=[truth["mass"] / constants.M_jup, truth["period"]],
    show_titles=True,
    quantiles=[0.16, 0.5, 0.84],
    title_fmt=".4f",
)