Dependencies¶

The only required dependencies for exoplanet are NumPy, PyMC3 and AstroPy. These can be installed using conda or pip:

conda install numpy pymc3 astropy

or

pip install numpy pymc3 astropy

Using pip¶

exoplanet can also be installed using pip:

pip install exoplanet

From Source¶

The source code for exoplanet can be downloaded and installed from GitHub by running

git clone --recursive https://github.com/dfm/exoplanet.git
cd exoplanet
python setup.py install

Testing¶

To run the unit tests, install the following dependencies using pip or conda (you’ll need to use the conda-forge channel to get starry):

conda install -c conda-forge numpy scipy astropy pymc3 pytest starry pip
pip install batman-package parameterized nose

and then execute:

py.test -v

All of the tests should (of course) pass. If any of the tests don’t pass and if you can’t sort out why, open an issue on GitHub.

Orbits¶

class exoplanet.orbits.KeplerianOrbit(period=None, a=None, t0=None, t_periastron=None, incl=None, b=None, duration=None, ecc=None, omega=None, Omega=None, m_planet=0.0, m_star=None, r_star=None, rho_star=None, m_planet_units=None, rho_star_units=None, model=None, contact_points_kwargs=None, **kwargs)¶

A system of bodies on Keplerian orbits around a common central

Given the input parameters, the values of all other parameters will be computed so a KeplerianOrbit instance will always have attributes for each argument. Note that the units of the computed attributes will all be in the standard units of this class (R_sun, M_sun, and days) except for rho_star which will be in g / cm^3.

There are only specific combinations of input parameters that can be used:

1. First, either period or a must be given. If values are given for both parameters, then neither m_star or rho_star can be defined because the stellar density implied by each planet will be computed in rho_star.
2. Only one of incl and b can be given.
3. If a value is given for ecc then omega must also be given.
4. If no stellar parameters are given, the central body is assumed to be the sun. If only rho_star is defined, the radius of the central is assumed to be 1 * R_sun. Otherwise, at most two of m_star, r_star, and rho_star can be defined.
5. Either t0 (reference transit) or t_periastron must be given, but not both.

Parameters:

period – The orbital periods of the bodies in days. a – The semimajor axes of the orbits in R_sun. t0 – The time of a reference transit for each orbits in days. t_periastron – The epoch of a reference periastron passage in days. incl – The inclinations of the orbits in radians. b – The impact parameters of the orbits. ecc – The eccentricities of the orbits. Must be 0 <= ecc < 1. omega – The arguments of periastron for the orbits in radians. Omega – The position angles of the ascending nodes in radians. m_planet – The masses of the planets in units of m_planet_units. m_star – The mass of the star in M_sun. r_star – The radius of the star in R_sun. rho_star – The density of the star in units of rho_star_units. m_planet_units – An astropy.units compatible unit object giving the units of the planet masses. If not given, the default is M_sun. rho_star_units – An astropy.units compatible unit object giving the units of the stellar density. If not given, the default is g / cm^3.

get_planet_position(t, parallax=None)¶

The planets’ positions in the barycentric frame

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_planet_velocity(t)¶

Get the planets’ velocity vectors

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of M_sun/day.

get_radial_velocity(t, K=None, output_units=None)¶

Get the radial velocity of the star

Note

The convention in exoplanet is that positive z points towards the observer. However, for consistency with radial velocity literature this method returns values where positive radial velocity corresponds to a redshift as expected.

Parameters:	t – The times where the radial velocity should be evaluated. K (Optional) – The semi-amplitudes of the orbits. If provided, the m_planet and incl parameters will be ignored and this amplitude will be used instead. output_units (Optional) – An AstroPy velocity unit. If not given, the output will be evaluated in m/s. This is ignored if a value is given for K.
Returns:	The reflex radial velocity evaluated at t in units of output_units. For multiple planets, this will have one row for each planet.

get_relative_angles(t, parallax=None)¶

The planets’ relative position to the star in the sky plane, in separation, position angle coordinates.

Note

This treats each planet independently and does not take the other planets into account when computing the position of the star. This is fine as long as the planet masses are small.

Parameters:	t – The times where the position should be evaluated.
Returns:	The separation (arcseconds) and position angle (radians, measured east of north) of the planet relative to the star.

get_relative_position(t, parallax=None)¶

The planets’ positions relative to the star in the X,Y,Z frame.

Note

This treats each planet independently and does not take the other planets into account when computing the position of the star. This is fine as long as the planet masses are small. In other words, the reflex motion of the star caused by the other planets is neglected when computing the relative coordinates of one of the planets.

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_relative_velocity(t)¶

The planets’ velocity relative to the star

Note

This treats each planet independently and does not take the other planets into account when computing the position of the star. This is fine as long as the planet masses are small.

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of R_sun/day.

get_star_position(t, parallax=None)¶

The star’s position in the barycentric frame

Note

If there are multiple planets in the system, this will return one column per planet with each planet’s contribution to the motion. The star’s full position can be computed by summing over the last axis.

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_star_velocity(t)¶

Get the star’s velocity vector

Note

For a system with multiple planets, this will return one column per planet with the contributions from each planet. The total velocity can be found by summing along the last axis.

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of M_sun/day.

in_transit(t, r=0.0, texp=None)¶

Get a list of timestamps that are in transit

Parameters:	t (vector) – A vector of timestamps to be evaluated. r (Optional) – The radii of the planets. texp (Optional[float]) – The exposure time.
Returns:	The indices of the timestamps that are in transit.

class exoplanet.orbits.TTVOrbit(*args, **kwargs)¶

A generalization of a Keplerian orbit with transit timing variations

Only one of the arguments ttvs or transit_times can be given and the other will be computed from the one that was provided.

In practice the way this works is that the time axis is shifted to account for the TTVs before solving for a standard Keplerian orbit. To identify which times corrorspond to which transits, this will find the closest labelled transit to the timestamp and adjust the timestamp accordingly. This means that things will go very badly if the time between neighboring transits is larger than 2*period.

Parameters:

ttvs – A list (with on entry for each planet) of “O-C” vectors for each transit of each planet in units of days. “O-C” means the difference between the observed transit time and the transit time expected for a regular periodic orbit. transit_times – A list (with on entry for each planet) of transit times for each transit of each planet in units of days. These times will be used to compute the implied (least squares) ttv_period and t0. It is possible to supply a separate period parameter that will set the shape of the transits, but care should be taken to make sure that period and ttv_period don’t diverge because things will break if the time between neighboring transits is larger than 2*period. transit_inds – A list of integer value tensors giving the transit number for each transit in transit_times'' or ``ttvs. This is useful when not all transits are observed. This should always be zero indexed. delta_log_period – If using the transit_times argument, this parameter specifies the difference (in natural log) between the leqast squares period and the effective period of the transit.

get_planet_position(t, parallax=None)¶

The planets’ positions in the barycentric frame

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_planet_velocity(t)¶

Get the planets’ velocity vectors

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of M_sun/day.

get_radial_velocity(t, K=None, output_units=None)¶

Get the radial velocity of the star

Note

The convention in exoplanet is that positive z points towards the observer. However, for consistency with radial velocity literature this method returns values where positive radial velocity corresponds to a redshift as expected.

Parameters:	t – The times where the radial velocity should be evaluated. K (Optional) – The semi-amplitudes of the orbits. If provided, the m_planet and incl parameters will be ignored and this amplitude will be used instead. output_units (Optional) – An AstroPy velocity unit. If not given, the output will be evaluated in m/s. This is ignored if a value is given for K.
Returns:	The reflex radial velocity evaluated at t in units of output_units. For multiple planets, this will have one row for each planet.

get_relative_angles(t, parallax=None)¶

The planets’ relative position to the star in the sky plane, in separation, position angle coordinates.

Note

This treats each planet independently and does not take the other planets into account when computing the position of the star. This is fine as long as the planet masses are small.

Parameters:	t – The times where the position should be evaluated.
Returns:	The separation (arcseconds) and position angle (radians, measured east of north) of the planet relative to the star.

get_relative_position(t, parallax=None)¶

The planets’ positions relative to the star in the X,Y,Z frame.

Note

This treats each planet independently and does not take the other planets into account when computing the position of the star. This is fine as long as the planet masses are small. In other words, the reflex motion of the star caused by the other planets is neglected when computing the relative coordinates of one of the planets.

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_relative_velocity(t)¶

The planets’ velocity relative to the star

Note

This treats each planet independently and does not take the other planets into account when computing the position of the star. This is fine as long as the planet masses are small.

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of R_sun/day.

get_star_position(t, parallax=None)¶

The star’s position in the barycentric frame

Note

If there are multiple planets in the system, this will return one column per planet with each planet’s contribution to the motion. The star’s full position can be computed by summing over the last axis.

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_star_velocity(t)¶

Get the star’s velocity vector

Note

For a system with multiple planets, this will return one column per planet with the contributions from each planet. The total velocity can be found by summing along the last axis.

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of M_sun/day.

in_transit(t, r=0.0, texp=None)¶

Get a list of timestamps that are in transit

Parameters:	t (vector) – A vector of timestamps to be evaluated. r (Optional) – The radii of the planets. texp (Optional[float]) – The exposure time.
Returns:	The indices of the timestamps that are in transit.

class exoplanet.orbits.ReboundOrbit(*args, **kwargs)¶

An N-body system powered by the rebound integrator

This takes all the same arguments as the KeplerianOrbit, but these arguments define the orbital elements at some reference time (given by the rebound_t parameter). The positions and velocities of the bodies are then computed by numerically integrating the gravitational N-body system.

rebound-specific parameters can be provided as keyword arguments prefixed by rebound_. These will then be applied to the rebound.Simulation object as properties. Therefore, if you want to change the integrator, you could use: rebound_integrator = "whfast", for example. All of these parameters are passed directly through to rebound except rebound_t (the reference time) which is converted from days to years over two pi (the default time units in rebound).

Note

exoplanet and rebound use different base units, but all of the unit conversions are handled behind the scenes in this object so that means that you should mostly use the exoplanet units when interacting with this class and you should be very cautious about setting the rebound_G argument. One example of a case where you’ll need to use the rebound units is when you want to set the integrator step size using the rebound_dt parameter.

get_planet_position(t)¶

The planets’ positions in the barycentric frame

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_planet_velocity(t)¶

Get the planets’ velocity vectors

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of M_sun/day.

get_radial_velocity(t, output_units=None)¶

Get the radial velocity of the star

Note

The convention in exoplanet is that positive z points towards the observer. However, for consistency with radial velocity literature this method returns values where positive radial velocity corresponds to a redshift as expected.

Note

Unlike the KeplerianOrbit implementation, the semi-amplitude K cannot be used with the ReboundOrbit. Also, the contributions of each planet are not returned separately; this will always return a single time series.

Parameters:	t – The times where the radial velocity should be evaluated. output_units (Optional) – An AstroPy velocity unit. If not given, the output will be evaluated in m/s. This is ignored if a value is given for K.
Returns:	The reflex radial velocity evaluated at t in units of output_units.

get_relative_angles(t, parallax=None)¶

The planets’ relative position to the star in the sky plane, in separation, position angle coordinates.

Note

This treats each planet independently and does not take the other planets into account when computing the position of the star. This is fine as long as the planet masses are small.

Parameters:	t – The times where the position should be evaluated.
Returns:	The separation (arcseconds) and position angle (radians, measured east of north) of the planet relative to the star.

get_relative_position(t)¶

The planets’ positions relative to the star in the X,Y,Z frame.

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_relative_velocity(t)¶

The planets’ velocity relative to the star

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of R_sun/day.

get_star_position(t)¶

The star’s position in the barycentric frame

Note

Unlike the KeplerianOrbit, this will not return the contributions from each planet separately.

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

get_star_velocity(t)¶

Get the star’s velocity vector

Note

Unlike the KeplerianOrbit, this will not return the contributions from each planet separately.

Parameters:	t – The times where the velocity should be evaluated.
Returns:	The components of the velocity vector at t in units of M_sun/day.

in_transit(t, r=0.0, texp=None)¶

This is a no-op and all points are assumed to be in transit

class exoplanet.orbits.SimpleTransitOrbit(period=None, t0=0.0, b=0.0, duration=None, r_star=1.0)¶

An orbit representing a set of planets transiting a common central

This orbit is parameterized by the observables of a transiting system, period, phase, duration, and impact parameter.

Parameters:	period – The orbital period of the planets in days. t0 – The midpoint time of a reference transit for each planet in days. b – The impact parameters of the orbits. duration – The durations of the transits in days. r_star – The radius of the star in R_sun.

get_relative_position(t)¶

The planets’ positions relative to the star

Parameters:	t – The times where the position should be evaluated.
Returns:	The components of the position vector at t in units of R_sun.

in_transit(t, r=None, texp=None)¶

Get a list of timestamps that are in transit

Parameters:	t (vector) – A vector of timestamps to be evaluated. r (Optional) – The radii of the planets. texp (Optional[float]) – The exposure time.
Returns:	The indices of the timestamps that are in transit.

Light curve models¶

class exoplanet.LimbDarkLightCurve(u, model=None)¶

A limb darkened light curve computed using starry

Parameters:	u (vector) – A vector of limb darkening coefficients.

get_light_curve(orbit=None, r=None, t=None, texp=None, oversample=7, order=0, use_in_transit=True)¶

Get the light curve for an orbit at a set of times

Parameters:

orbit – An object with a get_relative_position method that takes a tensor of times and returns a list of Cartesian coordinates of a set of bodies relative to the central source. This method should return three tensors (one for each coordinate dimension) and each tensor should have the shape append(t.shape, r.shape) or append(t.shape, oversample, r.shape) when texp is given. The first two coordinate dimensions are treated as being in the plane of the sky and the third coordinate is the line of sight with positive values pointing away from the observer. For an example, take a look at orbits.KeplerianOrbit. r (tensor) – The radius of the transiting body in the same units as r_star. This should have a shape that is consistent with the coordinates returned by orbit. In general, this means that it should probably be a scalar or a vector with one entry for each body in orbit. t (tensor) – The times where the light curve should be evaluated. texp (Optional[tensor]) – The exposure time of each observation. This can be a scalar or a tensor with the same shape as t. If texp is provided, t is assumed to indicate the timestamp at the middle of an exposure of length texp. oversample (Optional[int]) – The number of function evaluations to use when numerically integrating the exposure time. order (Optional[int]) – The order of the numerical integration scheme. This must be one of the following: 0 for a centered Riemann sum (equivalent to the “resampling” procedure suggested by Kipping 2010), 1 for the trapezoid rule, or 2 for Simpson’s rule. use_in_transit (Optional[bool]) – If True, the model will only be evaluated for the data points expected to be in transit as computed using the in_transit method on orbit.

Scalable Gaussian processes¶

class exoplanet.gp.GP(kernel, x, diag, J=-1, model=None)¶

The interface for computing Gaussian Process models with celerite

This class implements the method described in Foreman-Mackey et al. (2017) and Foreman-Mackey (2018) for scalable evaluation of Gaussian Process (GP) models in 1D.

Note

The input coordinates x must be sorted in ascending order, but this is not checked in the code. If the values are not sorted, the behavior of the algorithm is undefined.

Parameters:

kernel – A exoplanet.gp.terms.Term object the specifies the GP kernel. x – The input coordinates. This should be a 1D array and the elements must be sorted. Otherwise the results are undefined. diag – The extra diagonal to add to the covariance matrix. This should have the same length as x and correspond to the excess variance for each data point. Note: this is different from the usage in the celerite package where the standard deviation (instead of variance) is provided. J (Optional) – The width of the system. This is the J parameter from Foreman-Mackey (2018) (not the original paper) so a real term contributes J += 1 and a complex term contributes J += 2. If you know this value in advance, you can provide it. Otherwise, the code will try to work it out.

class exoplanet.gp.terms.Term(**kwargs)¶

The abstract base “term” that is the superclass of all other terms

Subclasses should overload the terms.Term.get_real_coefficients() and terms.Term.get_complex_coefficients() methods.

class exoplanet.gp.terms.RealTerm(**kwargs)¶

The simplest celerite term

This term has the form

\[k(\tau) = a_j\,e^{-c_j\,\tau}\]

with the parameters a and c.

Strictly speaking, for a sum of terms, the parameter a could be allowed to go negative but since it is somewhat subtle to ensure positive definiteness, we recommend keeping both parameters strictly positive. Advanced users can build a custom term that has negative coefficients but care should be taken to ensure positivity.

Parameters:	a or log_a (tensor) – The amplitude of the term. c or log_c (tensor) – The exponent of the term.

class exoplanet.gp.terms.ComplexTerm(**kwargs)¶

A general celerite term

This term has the form

\[k(\tau) = \frac{1}{2}\,\left[(a_j + b_j)\,e^{-(c_j+d_j)\,\tau} + (a_j - b_j)\,e^{-(c_j-d_j)\,\tau}\right]\]

with the parameters a, b, c, and d.

This term will only correspond to a positive definite kernel (on its own) if \(a_j\,c_j \ge b_j\,d_j\).

Parameters:	a or log_a (tensor) – The real part of amplitude. b or log_b (tensor) – The imaginary part of amplitude. c or log_c (tensor) – The real part of the exponent. d or log_d (tensor) – The imaginary part of exponent.

class exoplanet.gp.terms.SHOTerm(*args, **kwargs)¶

A term representing a stochastically-driven, damped harmonic oscillator

The PSD of this term is

\[S(\omega) = \sqrt{\frac{2}{\pi}} \frac{S_0\,\omega_0^4} {(\omega^2-{\omega_0}^2)^2 + {\omega_0}^2\,\omega^2/Q^2}\]

with the parameters S0, Q, and w0.

Parameters:

S0 or log_S0 (tensor) – The parameter \(S_0\). Q or log_Q (tensor) – The parameter \(Q\). w0 or log_w0 (tensor) – The parameter \(\omega_0\). Sw4 or log_Sw4 (tensor) – It can sometimes be more efficient to parameterize the amplitude of a SHO kernel using \(S_0\,{\omega_0}^4\) instead of \(S_0\) directly since \(S_0\) and \(\omega_0\) are strongly correlated. If provided, S0 will be computed from Sw4 and w0.

class exoplanet.gp.terms.Matern32Term(**kwargs)¶

A term that approximates a Matern-3/2 function

This term is defined as

\[k(\tau) = \sigma^2\,\left[ \left(1+1/\epsilon\right)\,e^{-(1-\epsilon)\sqrt{3}\,\tau/\rho} \left(1-1/\epsilon\right)\,e^{-(1+\epsilon)\sqrt{3}\,\tau/\rho} \right]\]

with the parameters sigma and rho. The parameter eps controls the quality of the approximation since, in the limit \(\epsilon \to 0\) this becomes the Matern-3/2 function

\[\lim_{\epsilon \to 0} k(\tau) = \sigma^2\,\left(1+ \frac{\sqrt{3}\,\tau}{\rho}\right)\, \exp\left(-\frac{\sqrt{3}\,\tau}{\rho}\right)\]

Parameters:	sigma or log_sigma (tensor) – The parameter \(\sigma\). rho or log_rho (tensor) – The parameter \(\rho\). eps (Optional[float]) – The value of the parameter \(\epsilon\). (default: 0.01)

class exoplanet.gp.terms.RotationTerm(**kwargs)¶

A mixture of two SHO terms that can be used to model stellar rotation

This term has two modes in Fourier space: one at period and one at 0.5 * period. This can be a good descriptive model for a wide range of stochastic variability in stellar time series from rotation to pulsations.

Parameters:

amp or log_amp (tensor) – The amplitude of the variability. period or log_period (tensor) – The primary period of variability. Q0 or log_Q0 (tensor) – The quality factor (or really the quality factor minus one half) for the secondary oscillation. deltaQ or log_deltaQ (tensor) – The difference between the quality factors of the first and the second modes. This parameterization (if deltaQ > 0) ensures that the primary mode alway has higher quality. mix – The fractional amplitude of the secondary mode compared to the primary. This should probably always be 0 < mix < 1.

Estimators¶

exoplanet.estimate_semi_amplitude(periods, x, y, yerr=None, t0s=None)¶

Estimate the RV semi-amplitudes for planets in an RV series

Parameters:	periods – The periods of the planets. Assumed to be in days if not an AstroPy Quantity. x – The observation times. Assumed to be in days if not an AstroPy Quantity. y – The radial velocities. Assumed to be in m/s if not an AstroPy Quantity. yerr (Optional) – The uncertainty on y. t0s (Optional) – The time of a reference transit for each planet, if known.
Returns:	An estimate of the semi-amplitude of each planet in units of m/s.

exoplanet.estimate_minimum_mass(periods, x, y, yerr=None, t0s=None, m_star=1)¶

Estimate the minimum mass(es) for planets in an RV series

Parameters:

periods – The periods of the planets. Assumed to be in days if not an AstroPy Quantity. x – The observation times. Assumed to be in days if not an AstroPy Quantity. y – The radial velocities. Assumed to be in m/s if not an AstroPy Quantity. yerr (Optional) – The uncertainty on y. t0s (Optional) – The time of a reference transit for each planet, if known. m_star (Optional) – The mass of the star. Assumed to be in M_sun if not an AstroPy Quantity.

Returns:

An estimate of the minimum mass of each planet as an AstroPy Quantity with units of M_jupiter.

exoplanet.lomb_scargle_estimator(x, y, yerr=None, min_period=None, max_period=None, filter_period=None, max_peaks=2, **kwargs)¶

Estimate period of a time series using the periodogram

Parameters:

x (ndarray[N]) – The times of the observations y (ndarray[N]) – The observations at times x yerr (Optional[ndarray[N]]) – The uncertainties on y min_period (Optional[float]) – The minimum period to consider max_period (Optional[float]) – The maximum period to consider filter_period (Optional[float]) – If given, use a high-pass filter to down-weight period longer than this max_peaks (Optional[int]) – The maximum number of peaks to return (default: 2)

Returns:

A dictionary with the computed periodogram and the parameters for up to max_peaks peaks in the periodogram.

exoplanet.autocorr_estimator(x, y, yerr=None, min_period=None, max_period=None, oversample=2.0, smooth=2.0, max_peaks=10)¶

Estimate the period of a time series using the autocorrelation function

Note

The signal is interpolated onto a uniform grid in time so that the autocorrelation function can be computed.

Parameters:

x (ndarray[N]) – The times of the observations y (ndarray[N]) – The observations at times x yerr (Optional[ndarray[N]]) – The uncertainties on y min_period (Optional[float]) – The minimum period to consider max_period (Optional[float]) – The maximum period to consider oversample (Optional[float]) – When interpolating, oversample the times by this factor (default: 2.0) smooth (Optional[float]) – Smooth the autocorrelation function by this factor times the minimum period (default: 2.0) max_peaks (Optional[int]) – The maximum number of peaks to identify in the autocorrelation function (default: 10)

Returns:

A dictionary with the computed autocorrelation function and the estimated period. For compatibility with the lomb_scargle_estimator(), the period is returned as a list with the key peaks.

Distributions¶

Utilities¶

exoplanet.optimize(start=None, vars=None, model=None, return_info=False, verbose=True, **kwargs)¶

Maximize the log prob of a PyMC3 model using scipy

All extra arguments are passed directly to the scipy.optimize.minimize function.

Parameters:	start – The PyMC3 coordinate dictionary of the starting position vars – The variables to optimize model – The PyMC3 model return_info – Return both the coordinate dictionary and the result of scipy.optimize.minimize verbose – Print the success flag and log probability to the screen

exoplanet.eval_in_model(var, point=None, return_func=False, model=None, **kwargs)¶

Evaluate a Theano tensor or PyMC3 variable in a PyMC3 model

This method builds a Theano function for evaluating a node in the graph given the required parameters. This will also cache the compiled Theano function in the current pymc3.Model to reduce the overhead of calling this function many times.

Parameters:

var – The variable or tensor to evaluate. point (Optional) – A dict of input parameter values. This can be model.test_point (default), the result of pymc3.find_MAP, a point in a pymc3.MultiTrace or any other representation of the input parameters. return_func (Optional[bool]) – If False (default), return the evaluated variable. If True, return the result, the Theano function and the list of arguments for that function.

Returns:

Depending on return_func, either the value of var at point, or this value, the Theano function, and the input arguments.

exoplanet.get_samples_from_trace(trace, size=1)¶

Generate random samples from a PyMC3 MultiTrace

Parameters:	trace – The MultiTrace. size – The number of samples to generate.

exoplanet.get_dense_nuts_step(start=None, adaptation_window=101, doubling=True, model=None, **kwargs)¶

Get a NUTS step function with a dense mass matrix

The entries in the mass matrix will be tuned based on the sample covariances during tuning. All extra arguments are passed directly to pymc3.NUTS.

Parameters:	start (dict, optional) – A starting point in parameter space. If not provided, the model’s test_point is used. adaptation_window (int, optional) – The (initial) size of the window used for sample covariance estimation. doubling (bool, optional) – If True (default) the adaptation window is doubled each time the matrix is updated.

exoplanet.orbits.ttv.compute_expected_transit_times(min_time, max_time, period, t0)¶

Compute the expected transit times within a dataset

Parameters:	min_time (float) – The start time of the dataset max_time (float) – The end time of the dataset period (array) – The periods of the planets t0 (array) – The reference transit times for the planets
Returns:	A list of arrays of expected transit times for each planet

Units¶

exoplanet.units.with_unit(obj, unit)¶

Decorate a Theano tensor with Astropy units

Parameters:	obj – The Theano tensor unit (astropy.Unit) – The units for this object
Raises:	TypeError – If the tensor already has units

exoplanet.units.has_unit(obj)¶

Does an object have units as defined by exoplanet?

exoplanet.units.to_unit(obj, target)¶

Convert a Theano tensor with units to a target set of units

Parameters:	obj – The Theano tensor target (astropy.Unit) – The target units
Returns:	A Theano tensor in the right units

Citations¶

exoplanet.citations.get_citations_for_model(model=None, width=79)¶

Get the citations for the components used an exoplanet PyMC3

Returns: The acknowledgement text for exoplanet and its dependencies and a string containing the BibTeX entries for the citations in the acknowledgement.

Hello world (AKA fitting a line to data)¶

My standard intro to a new modeling language or inference framework is to fit a line to data. So. Let’s do that with PyMC3.

To start, we’ll generate some fake data using a linear model. Feel free to change the random number seed to try out a different dataset.

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

true_m = 0.5
true_b = -1.3
true_logs = np.log(0.3)

x = np.sort(np.random.uniform(0, 5, 50))
y = true_b + true_m * x + np.exp(true_logs) * np.random.randn(len(x))

plt.plot(x, y, ".k")
plt.ylim(-2, 2)
plt.xlabel("x")
plt.ylabel("y");

To fit a model to these data, our model will have 3 parameters: the slope \(m\), the intercept \(b\), and the log of the uncertainty \(\log(\sigma)\). To start, let’s choose broad uniform priors on these parameters:

Then, the log-likelihood function will be

[Note: the second normalization term is needed in this model because we are fitting for \(\sigma\) and the second term is not a constant.]

Another way of writing this model that might not be familiar is the following:

This is the way that a model like this is often defined in statistics and it will be useful when we implement out model in PyMC3 so take a moment to make sure that you understand the notation.

Now, let’s implement this model in PyMC3. The documentation for the distributions available in PyMC3’s modeling language can be found here and these will come in handy as you go on to write your own models.

import pymc3 as pm

with pm.Model() as model:

    # Define the priors on each parameter:
    m = pm.Uniform("m", lower=-5, upper=5)
    b = pm.Uniform("b", lower=-5, upper=5)
    logs = pm.Uniform("logs", lower=-5, upper=5)

    # Define the likelihood. A few comments:
    #  1. For mathematical operations like "exp", you can't use
    #     numpy. Instead, use the mathematical operations defined
    #     in "pm.math".
    #  2. To condition on data, you use the "observed" keyword
    #     argument to any distribution. In this case, we want to
    #     use the "Normal" distribution (look up the docs for
    #     this).
    pm.Normal("obs", mu=m * x + b, sd=pm.math.exp(logs), observed=y)

    # This is how you will sample the model. Take a look at the
    # docs to see that other parameters that are available.
    trace = pm.sample(draws=1000, tune=1000, chains=2)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 4 jobs)
NUTS: [logs, b, m]
Sampling 2 chains: 100%|██████████| 4000/4000 [00:01<00:00, 2048.29draws/s]

Now since we now have samples, let’s make some diagnostic plots. The first plot to look at is the “traceplot” implemented in PyMC3. In this plot, you’ll see the marginalized distribution for each parameter on the left and the trace plot (parameter value as a function of step number) on the right. In each panel, you should see two lines with different colors. These are the results of different independent chains and if the results are substantially different in the different chains then there is probably something going wrong.

pm.traceplot(trace, varnames=["m", "b", "logs"]);

/mnt/home/dforeman/miniconda3/envs/autoexoplanet/lib/python3.7/site-packages/pymc3/plots/__init__.py:40: UserWarning: Keyword argument varnames renamed to var_names, and will be removed in pymc3 3.8
  warnings.warn('Keyword argument {old} renamed to {new}, and will be removed in pymc3 3.8'.format(old=old, new=new))

It’s also good to quantify that “looking substantially different” argument. This is implemented in PyMC3 as the “summary” function. In this table, some of the key columns to look at are n_eff and Rhat. * n_eff shows an estimate of the number of effective (or independent) samples for that parameter. In this case, n_eff should probably be around 500 per chain (there should have been 2 chains run). * Rhat shows the Gelman–Rubin statistic and it should be close to 1.

pm.summary(trace, varnames=["m", "b", "logs"])

The last diagnostic plot that we’ll make here is the corner plot made using corner.py. The easiest way to do this using PyMC3 is to first convert the trace to a Pandas DataFrame and then pass that to corner.py.

import corner  # https://corner.readthedocs.io

samples = pm.trace_to_dataframe(trace, varnames=["m", "b", "logs"])
corner.corner(samples, truths=[true_m, true_b, true_logs]);

Extra credit: Here are a few suggestions for things to try out while getting more familiar with PyMC3:

1. Try initializing the parameters using the testval argument to the distributions. Does this improve performance in this case? It will substantially improve performance in more complicated examples.
2. Try changing the priors on the parameters. For example, try the “uninformative” prior recommended by Jake VanderPlas on his blog.
3. What happens as you substantially increase or decrease the simulated noise? Does the performance change significantly? Why?

A more realistic example: radial velocity exoplanets¶

While the above example was cute, it doesn’t really fully exploit the power of PyMC3 and it doesn’t really show some of the real issues that you will face when you use PyMC3 as an astronomer. To get a better sense of how you might use PyMC3 in Real Life™, let’s take a look at a more realistic example: fitting a Keplerian orbit to radial velocity observations.

One of the key aspects of this problem that I want to highlight is the fact that PyMC3 (and the underlying model building framework Theano) don’t have out-of-the-box support for the root-finding that is required to solve Kepler’s equation. As part of the process of computing a Keplerian RV model, we must solve the equation:

for the eccentric anomaly \(E\) given some mean anomaly \(M\) and eccentricity \(e\). There are commonly accepted methods of solving this equation using Newton’s method, but if we want to expose that to PyMC3, we have to define a custom Theano operation with a custom gradient. I won’t go into the details of the math (because I blogged about it) and I won’t go into the details of the implementation (because you can take a look at it on GitHub). So, for this tutorial, we’ll use the custom Kepler solver that is implemented as part of exoplanet and fit the publicly available radial velocity observations of the famous exoplanetary system 51 Peg using PyMC3.

First, we need to download the data from the exoplanet archive:

import requests
import pandas as pd
import matplotlib.pyplot as plt

# Download the dataset from the Exoplanet Archive:
url = "https://exoplanetarchive.ipac.caltech.edu/data/ExoData/0113/0113357/data/UID_0113357_RVC_001.tbl"
r = requests.get(url)
if r.status_code != requests.codes.ok:
    r.raise_for_status()
data = np.array(
    [
        l.split()
        for l in r.text.splitlines()
        if not l.startswith("\\") and not l.startswith("|")
    ],
    dtype=float,
)
t, rv, rv_err = data.T
t -= np.mean(t)

# Plot the observations "folded" on the published period:
# Butler et al. (2006) https://arxiv.org/abs/astro-ph/0607493
lit_period = 4.230785
plt.errorbar((t % lit_period) / lit_period, rv, yerr=rv_err, fmt=".k", capsize=0)
plt.xlim(0, 1)
plt.ylim(-110, 110)
plt.annotate(
    "period = {0:.6f} days".format(lit_period),
    xy=(1, 0),
    xycoords="axes fraction",
    xytext=(-5, 5),
    textcoords="offset points",
    ha="right",
    va="bottom",
    fontsize=12,
)
plt.ylabel("radial velocity [m/s]")
plt.xlabel("phase");

Now, here’s the implementation of a radial velocity model in PyMC3. Some of this will look familiar after the Hello World example, but things are a bit more complicated now. Take a minute to take a look through this and see if you can follow it. There’s a lot going on, so I want to point out a few things to pay attention to:

1. All of the mathematical operations (for example exp and sqrt) are being performed using Theano instead of NumPy.
2. All of the parameters have initial guesses provided. This is an example where this makes a big difference because some of the parameters (like period) are very tightly constrained.
3. Some of the lines are wrapped in Deterministic distributions. This can be useful because it allows us to track values as the chain progresses even if they’re not parameters. For example, after sampling, we will have a sample for bkg (the background RV trend) for each step in the chain. This can be especially useful for making plots of the results.
4. Similarly, at the end of the model definition, we compute the RV curve for a single orbit on a fine grid. This can be very useful for diagnosing fits gone wrong.
5. For parameters that specify angles (like \(\omega\), called w in the model below), it can be inefficient to sample in the angle directly because of the fact that the value wraps around at \(2\pi\). Instead, it can be better to sample the unit vector specified by the angle. In practice, this can be achieved by sampling a 2-vector from an isotropic Gaussian and normalizing the components by the norm. This is implemented as part of exoplanet in the exoplanet.distributions.Angle class.

import theano.tensor as tt

from exoplanet.orbits import get_true_anomaly
from exoplanet.distributions import Angle

with pm.Model() as model:

    # Parameters
    logK = pm.Uniform(
        "logK",
        lower=0,
        upper=np.log(200),
        testval=np.log(0.5 * (np.max(rv) - np.min(rv))),
    )
    logP = pm.Uniform("logP", lower=0, upper=np.log(10), testval=np.log(lit_period))
    phi = pm.Uniform("phi", lower=0, upper=2 * np.pi, testval=0.1)
    e = pm.Uniform("e", lower=0, upper=1, testval=0.1)
    w = Angle("w")
    logjitter = pm.Uniform(
        "logjitter", lower=-10, upper=5, testval=np.log(np.mean(rv_err))
    )
    rv0 = pm.Normal("rv0", mu=0.0, sd=10.0, testval=0.0)
    rvtrend = pm.Normal("rvtrend", mu=0.0, sd=10.0, testval=0.0)

    # Deterministic transformations
    n = 2 * np.pi * tt.exp(-logP)
    P = pm.Deterministic("P", tt.exp(logP))
    K = pm.Deterministic("K", tt.exp(logK))
    cosw = tt.cos(w)
    sinw = tt.sin(w)
    s2 = tt.exp(2 * logjitter)
    t0 = (phi + w) / n

    # The RV model
    bkg = pm.Deterministic("bkg", rv0 + rvtrend * t / 365.25)
    M = n * t - (phi + w)

    # This is the line that uses the custom Kepler solver
    f = get_true_anomaly(M, e + tt.zeros_like(M))
    rvmodel = pm.Deterministic(
        "rvmodel", bkg + K * (cosw * (tt.cos(f) + e) - sinw * tt.sin(f))
    )

    # Condition on the observations
    err = tt.sqrt(rv_err ** 2 + tt.exp(2 * logjitter))
    pm.Normal("obs", mu=rvmodel, sd=err, observed=rv)

    # Compute the phased RV signal
    phase = np.linspace(0, 1, 500)
    M_pred = 2 * np.pi * phase - (phi + w)
    f_pred = get_true_anomaly(M_pred, e + tt.zeros_like(M_pred))
    rvphase = pm.Deterministic(
        "rvphase", K * (cosw * (tt.cos(f_pred) + e) - sinw * tt.sin(f_pred))
    )

In this case, I’ve found that it is useful to first optimize the parameters to find the “maximum a posteriori” (MAP) parameters and then start the sampler from there. This is useful here because MCMC is not designed to find the maximum of the posterior; it’s just meant to sample the shape of the posterior. The performance of all MCMC methods can be really bad when the initialization isn’t good (especially when some parameters are very well constrained). To find the maximum a posteriori parameters using PyMC3, you can use the exoplanet.optimize() function:

from exoplanet import optimize

with model:
    map_params = optimize()

optimizing logp for variables: [rvtrend, rv0, logjitter, w, e, phi, logP, logK]
70it [00:05, 12.53it/s, logp=-8.351913e+02]
message: Desired error not necessarily achieved due to precision loss.
logp: -1325.8433398976674 -> -835.1912653583977

Let’s make a plot to check that this initialization looks reasonable. In the top plot, we’re looking at the RV observations as a function of time with the initial guess for the long-term trend overplotted in blue. In the lower panel, we plot the “folded” curve where we have wrapped the observations onto the best-fit period and the prediction for a single overplotted in orange. If this doesn’t look good, try adjusting the initial guesses for the parameters and see if you can get a better fit.

Exercise: Try changing the initial guesses for the parameters (as specified by the testval argument) and see how sensitive the results are to these values. Are there some parameters that are less important? Why is this?

fig, axes = plt.subplots(2, 1, figsize=(8, 8))

period = map_params["P"]

ax = axes[0]
ax.errorbar(t, rv, yerr=rv_err, fmt=".k")
ax.plot(t, map_params["bkg"], color="C0", lw=1)
ax.set_ylim(-110, 110)
ax.set_ylabel("radial velocity [m/s]")
ax.set_xlabel("time [days]")

ax = axes[1]
ax.errorbar(t % period, rv - map_params["bkg"], yerr=rv_err, fmt=".k")
ax.plot(phase * period, map_params["rvphase"], color="C1", lw=1)
ax.set_ylim(-110, 110)
ax.set_ylabel("radial velocity [m/s]")
ax.set_xlabel("phase [days]")

plt.tight_layout()

Now let’s sample the posterior starting from our MAP estimate.

with model:
    trace = pm.sample(draws=2000, tune=1000, start=map_params, chains=2)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 4 jobs)
NUTS: [rvtrend, rv0, logjitter, w, e, phi, logP, logK]
Sampling 2 chains: 100%|██████████| 6000/6000 [00:28<00:00, 156.19draws/s]
There were 2 divergences after tuning. Increase target_accept or reparameterize.
There were 56 divergences after tuning. Increase target_accept or reparameterize.

As above, it’s always a good idea to take a look at the summary statistics for the chain. If everything went as planned, there should be more than 1000 effective samples per chain and the Rhat values should be close to 1. (Not too bad for less than 30 seconds of run time!)

pm.summary(
    trace, varnames=["logK", "logP", "phi", "e", "w", "logjitter", "rv0", "rvtrend"]
)

Similarly, we can make the corner plot again for this model.

samples = pm.trace_to_dataframe(trace, varnames=["K", "P", "e", "w"])
corner.corner(samples);

Finally, the last plot that we’ll make here is of the posterior predictive density. In this case, this means that we want to look at the distribution of predicted models that are consistent with the data. As above, the top plot shows the raw observations as black error bars and the RV trend model is overplotted in blue. But, this time, the blue line is actually composed of 25 lines that are samples from the posterior over trends that are consistent with the data. In the bottom panel, the orange lines indicate the same 25 posterior samples for the RV curve of one orbit.

fig, axes = plt.subplots(2, 1, figsize=(8, 8))

period = map_params["P"]

ax = axes[0]
ax.errorbar(t, rv, yerr=rv_err, fmt=".k")
ax.set_ylabel("radial velocity [m/s]")
ax.set_xlabel("time [days]")

ax = axes[1]
ax.errorbar(t % period, rv - map_params["bkg"], yerr=rv_err, fmt=".k")
ax.set_ylabel("radial velocity [m/s]")
ax.set_xlabel("phase [days]")

for i in np.random.randint(len(trace) * trace.nchains, size=25):
    axes[0].plot(t, trace["bkg"][i], color="C0", lw=1, alpha=0.3)
    axes[1].plot(phase * period, trace["rvphase"][i], color="C1", lw=1, alpha=0.3)

axes[0].set_ylim(-110, 110)
axes[1].set_ylim(-110, 110)

plt.tight_layout()

Dense mass matrices¶

The main extra is the exoplanet.get_dense_nuts_step() function that extends the PyMC3 sampling procedure to include support for learning off-diagonal elements of the mass matrix. This is very important for any problems where there are covariances between the parameters (this is true for pretty much all exoplanet models). A thorough discussion of this can be found elsewhere online, but here is a simple demo where we sample a covariant Gaussian using exoplanet.get_dense_nuts_step().

First, we generate a random positive definite covariance matrix for the Gaussian:

import numpy as np

ndim = 5
np.random.seed(42)
L = np.random.randn(ndim, ndim)
L[np.diag_indices_from(L)] = 0.1 * np.exp(L[np.diag_indices_from(L)])
L[np.triu_indices_from(L, 1)] = 0.0
cov = np.dot(L, L.T)

And then we can sample this using PyMC3 and exoplanet.get_dense_nuts_step():

import pymc3 as pm
import exoplanet as xo

with pm.Model() as model:
    pm.MvNormal("x", mu=np.zeros(ndim), chol=L, shape=(ndim,))
    trace = pm.sample(tune=2000, draws=2000, chains=4, step=xo.get_dense_nuts_step())

Multiprocess sampling (4 chains in 4 jobs)
NUTS: [x]
Sampling 4 chains: 100%|██████████| 16000/16000 [00:07<00:00, 2115.59draws/s]

This is a little more verbose than the standard use of PyMC3, but the performance is several orders of magnitude better than you would get without the mass matrix tuning. As you can see from the pymc3.summary, the autocorrelation time of this chain is about 1 as we would expect for a simple problem like this.

pm.summary(trace)

Evaluating model components for specific samples¶

I find that when I’m debugging a PyMC3 model, I often want to inspect the value of some part of the model for a given set of parameters. As far as I can tell, there isn’t a simple way to do this in PyMC3, so exoplanet comes with a hack for doing this: exoplanet.eval_in_model(). This function handles the mapping between named PyMC3 variables and the input required by the Theano function that can evaluate the requested variable or tensor.

As a demo, let’s say that we’re fitting a parabola to some data:

np.random.seed(42)
x = np.sort(np.random.uniform(-1, 1, 50))
with pm.Model() as model:
    logs = pm.Normal("logs", mu=-3.0, sd=1.0)
    a0 = pm.Normal("a0")
    a1 = pm.Normal("a1")
    a2 = pm.Normal("a2")
    mod = a0 + a1 * x + a2 * x ** 2

    # Sample from the prior
    prior_sample = pm.sample_prior_predictive(samples=1)
    y = xo.eval_in_model(mod, prior_sample)
    y += np.exp(prior_sample["logs"]) * np.random.randn(len(y))

    # Add the likelihood
    pm.Normal("obs", mu=mod, sd=pm.math.exp(logs), observed=y)

    # Fit the data
    map_soln = xo.optimize()
    trace = pm.sample()

optimizing logp for variables: [a2, a1, a0, logs]
26it [00:00, 292.51it/s, logp=4.272312e+01]
message: Optimization terminated successfully.
logp: -180.51036900636572 -> 42.72311721910288
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [a2, a1, a0, logs]
Sampling 4 chains: 100%|██████████| 4000/4000 [00:01<00:00, 3206.57draws/s]
The acceptance probability does not match the target. It is 0.8788002200800235, but should be close to 0.8. Try to increase the number of tuning steps.
The acceptance probability does not match the target. It is 0.8834353310005734, but should be close to 0.8. Try to increase the number of tuning steps.
The acceptance probability does not match the target. It is 0.8827839735098203, but should be close to 0.8. Try to increase the number of tuning steps.

After running the fit, it might be interesting to look at the predictions of the model. We could have added a pymc3.Deterministic node for eveything, but that can end up taking up a lot of memory and sometimes its useful to be able to experiement with different outputs. Using exoplanet.utils.eval_in_model() we can, for example, evaluate the maximum a posteriori (MAP) model prediction on a fine grid:

import matplotlib.pyplot as plt

x_grid = np.linspace(-1.1, 1.1, 5000)
with model:
    pred = xo.eval_in_model(a0 + a1 * x_grid + a2 * x_grid ** 2, map_soln)

plt.plot(x, y, ".k", label="data")
plt.plot(x_grid, pred, label="map")
plt.legend(fontsize=12)
plt.xlabel("x")
plt.ylabel("y")
plt.xlim(-1.1, 1.1);

We can also combine this with exoplanet.get_samples_from_trace() to plot this prediction for a set of samples in the trace.

samples = np.empty((50, len(x_grid)))
with model:
    y_grid = a0 + a1 * x_grid + a2 * x_grid ** 2
    for i, sample in enumerate(xo.get_samples_from_trace(trace, size=50)):
        samples[i] = xo.eval_in_model(y_grid, sample)

plt.plot(x, y, ".k", label="data")
plt.plot(x_grid, pred, label="map")
plt.plot(x_grid, samples[0], color="C1", alpha=0.1, label="posterior")
plt.plot(x_grid, samples[1:].T, color="C1", alpha=0.1)
plt.legend(fontsize=12)
plt.xlabel("x")
plt.ylabel("y")
plt.xlim(-1.1, 1.1);

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 A quick intro to PyMC3 for exoplaneteers 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.Bound(pm.Normal, lower=0)(
        "P",
        mu=np.array(periods),
        sd=np.array(period_errs),
        shape=2,
        testval=np.array(periods),
    )

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

    # Eccentricity & argument of periasteron
    ecc = xo.distributions.UnitUniform("ecc", shape=2, testval=np.array([0.1, 0.1]))
    omega = xo.distributions.Angle("omega", shape=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");

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]
13it [00:04,  3.01it/s, logp=-6.484820e+01]
message: Optimization terminated successfully.
logp: -79.73266285618838 -> -64.8482026233154
optimizing logp for variables: [trend, logs, omega, ecc, logK, P, t0]
94it [00:00, 180.82it/s, logp=-1.427676e+01]
message: Optimization terminated successfully.
logp: -64.8482026233154 -> -14.276760262380932

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");

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.get_dense_nuts_step() to tune the sampler (see the PyMC3 extras tutorial for more information).

np.random.seed(42)
with model:
    trace = pm.sample(
        tune=4000, draws=4000, step=xo.get_dense_nuts_step(target_accept=0.95)
    )

Multiprocess sampling (4 chains in 4 jobs)
NUTS: [trend, logs, omega, ecc, logK, P, t0]
Sampling 4 chains: 100%|██████████| 32000/32000 [01:38<00:00, 324.88draws/s]
There was 1 divergence after tuning. Increase target_accept or reparameterize.
There were 5 divergences after tuning. Increase target_accept or reparameterize.
There were 4 divergences after tuning. Increase target_accept or reparameterize.
There were 6 divergences after tuning. Increase target_accept or reparameterize.
The number of effective samples is smaller than 25% 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"])

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", "omega"])
corner.corner(samples);

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");

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

Citations¶

As described in the Citing exoplanet & its dependencies tutorial, we can use exoplanet.citations.get_citations_for_model() to construct an acknowledgement and BibTeX listing that includes the relevant citations for this model.

with model:
    txt, bib = xo.citations.get_citations_for_model()
print(txt)

This research made use of textsf{exoplanet} citep{exoplanet} and its
dependencies citep{exoplanet:astropy13, exoplanet:astropy18,
exoplanet:exoplanet, exoplanet:pymc3, exoplanet:theano}.

print("\n".join(bib.splitlines()[:10]) + "\n...")

@misc{exoplanet:exoplanet,
  author = {Daniel Foreman-Mackey and Ian Czekala and Rodrigo Luger and
            Eric Agol and Geert Barentsen and Tom Barclay},
   title = {dfm/exoplanet: exoplanet v0.2.1},
   month = sep,
    year = 2019,
     doi = {10.5281/zenodo.3462740},
     url = {https://doi.org/10.5281/zenodo.3462740}
}
...

The transit model in PyMC3¶

In this section, we will construct a simple transit fit model using PyMC3 and then we will fit a two planet model to simulated data. To start, let’s randomly sample some periods and phases and then define the time sampling:

np.random.seed(123)
periods = np.random.uniform(5, 20, 2)
t0s = periods * np.random.rand(2)
t = np.arange(0, 80, 0.02)
yerr = 5e-4

Then, define the parameters. In this simple model, we’ll just fit for the limb darkening parameters of the star, and the period, phase, impact parameter, and radius ratio of the planets (note: this is already 10 parameters and running MCMC to convergence using emcee would probably take at least an hour). For the limb darkening, we’ll use a quadratic law as parameterized by Kipping (2013). This reparameterizations is implemented in exoplanet as custom PyMC3 distribution exoplanet.distributions.QuadLimbDark.

import pymc3 as pm

with pm.Model() as model:

    # The baseline flux
    mean = pm.Normal("mean", mu=0.0, sd=1.0)

    # The time of a reference transit for each planet
    t0 = pm.Normal("t0", mu=t0s, sd=1.0, shape=2)

    # The log period; also tracking the period itself
    logP = pm.Normal("logP", mu=np.log(periods), sd=0.1, shape=2)
    period = pm.Deterministic("period", pm.math.exp(logP))

    # The Kipping (2013) parameterization for quadratic limb darkening paramters
    u = xo.distributions.QuadLimbDark("u", testval=np.array([0.3, 0.2]))

    r = pm.Uniform("r", lower=0.01, upper=0.1, shape=2, testval=np.array([0.04, 0.06]))
    b = xo.distributions.ImpactParameter("b", ror=r, shape=2, testval=np.random.rand(2))

    # Set up a Keplerian orbit for the planets
    orbit = xo.orbits.KeplerianOrbit(period=period, t0=t0, b=b)

    # Compute the model light curve using starry
    light_curves = xo.LimbDarkLightCurve(u).get_light_curve(orbit=orbit, r=r, t=t)
    light_curve = pm.math.sum(light_curves, axis=-1) + mean

    # Here we track the value of the model light curve for plotting
    # purposes
    pm.Deterministic("light_curves", light_curves)

    # In this line, we simulate the dataset that we will fit
    y = xo.eval_in_model(light_curve)
    y += yerr * np.random.randn(len(y))

    # The likelihood function assuming known Gaussian uncertainty
    pm.Normal("obs", mu=light_curve, sd=yerr, observed=y)

    # Fit for the maximum a posteriori parameters given the simuated
    # dataset
    map_soln = xo.optimize(start=model.test_point)

optimizing logp for variables: [b, r, u, logP, t0, mean]
99it [00:05, 18.44it/s, logp=2.479354e+04]
message: Desired error not necessarily achieved due to precision loss.
logp: 24787.977771807487 -> 24793.5394256112

Now we can plot the simulated data and the maximum a posteriori model to make sure that our initialization looks ok.

plt.plot(t, y, ".k", ms=4, label="data")
for i, l in enumerate("bc"):
    plt.plot(t, map_soln["light_curves"][:, i], lw=1, label="planet {0}".format(l))
plt.xlim(t.min(), t.max())
plt.ylabel("relative flux")
plt.xlabel("time [days]")
plt.legend(fontsize=10)
plt.title("map model");

Sampling¶

Now, let’s sample from the posterior defined by this model. As usual, there are strong covariances between some of the parameters so we’ll use exoplanet.get_dense_nuts_step().

np.random.seed(42)
with model:
    trace = pm.sample(
        tune=3000,
        draws=3000,
        start=map_soln,
        chains=4,
        step=xo.get_dense_nuts_step(target_accept=0.9),
    )

Multiprocess sampling (4 chains in 4 jobs)
NUTS: [b, r, u, logP, t0, mean]
Sampling 4 chains: 100%|██████████| 24000/24000 [01:11<00:00, 336.21draws/s]

After sampling, it’s important that we assess convergence. We can do that using the pymc3.summary function:

pm.summary(trace, varnames=["period", "t0", "r", "b", "u", "mean"])

That looks pretty good! Fitting this without exoplanet would have taken a lot more patience.

Now we can also look at the corner plot of some of that parameters of interest:

import corner

samples = pm.trace_to_dataframe(trace, varnames=["period", "r"])
truth = np.concatenate(xo.eval_in_model([period, r], model.test_point, model=model))
corner.corner(
    samples, truths=truth, labels=["period 1", "period 2", "radius 1", "radius 2"]
);