Note

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

PyMC3 has support for Gaussian Processes
(GPs), but this implementation is too
slow for many applications in time series astrophysics. So *exoplanet*
comes with an implementation of scalable GPs powered by
celerite. More information about
the algorithm can be found in the celerite
docs and in the papers (Paper
1 and Paper
2), but this tutorial will give a
hands on demo of how to use celerite in PyMC3.

Note

For the best results, we generally recommend the use of the `exoplanet.gp.terms.SHOTerm`

, `exoplanet.gp.terms.Matern32Term`

, and `exoplanet.gp.terms.RotationTerm`

“terms” because the other terms tend to have unphysical behavior at high frequency.

Let’s start with the quickstart demo from the celerite
docs.
We’ll fit the following simulated dataset using the sum of two
`exoplanet.gp.terms.SHOTerm`

objects.

First, generate the simulated data:

```
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(42)
t = np.sort(
np.append(np.random.uniform(0, 3.8, 57), np.random.uniform(5.5, 10, 68))
) # The input coordinates must be sorted
yerr = np.random.uniform(0.08, 0.22, len(t))
y = (
0.2 * (t - 5)
+ np.sin(3 * t + 0.1 * (t - 5) ** 2)
+ yerr * np.random.randn(len(t))
)
true_t = np.linspace(0, 10, 5000)
true_y = 0.2 * (true_t - 5) + np.sin(3 * true_t + 0.1 * (true_t - 5) ** 2)
plt.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0, label="data")
plt.plot(true_t, true_y, "k", lw=1.5, alpha=0.3, label="truth")
plt.legend(fontsize=12)
plt.xlabel("t")
plt.ylabel("y")
plt.xlim(0, 10)
_ = plt.ylim(-2.5, 2.5)
```

This plot shows the simulated data as black points with error bars and the true function is shown as a gray line.

Now let’s build the PyMC3 model that we’ll use to fit the data. We can see that there’s some roughly periodic signal in the data as well as a longer term trend. To capture these two features, we will model this as a mixture of two stochastically driven simple harmonic oscillators (SHO) with the power spectrum:

\[S(\omega) = \sqrt{\frac{2}{\pi}}\frac{S_1\,{\omega_1}^4}{(\omega^2 - {\omega_1}^2)^2 + 2\,{\omega_1}^2\,\omega^2}
+ \sqrt{\frac{2}{\pi}}\frac{S_2\,{\omega_2}^4}{(\omega^2 - {\omega_2}^2)^2 + {\omega_2}^2\,\omega^2/Q^2}\]

The first term is `exoplanet.gp.terms.SHOterm`

with
\(Q=1/\sqrt{2}\) and the second is regular
`exoplanet.gp.terms.SHOterm`

. This model has g free parameters:
\(S_1\), \(\omega_1\), \(S_2\), \(\omega_2\), \(Q\),
and a constant mean value. Most of the parameters will have weakly
informative inverse Gamma priors (see this blog
post
for a discussion of these priors) where the parameters are chosen to
have reasonable tail probabilities. Using *exoplanet*, this is how you
would build this model:

```
import pymc3 as pm
import theano.tensor as tt
from exoplanet.gp import terms, GP
from exoplanet import estimate_inverse_gamma_parameters
with pm.Model() as model:
mean = pm.Normal("mean", mu=0.0, sigma=1.0)
S1 = pm.InverseGamma(
"S1", **estimate_inverse_gamma_parameters(0.5 ** 2, 10.0 ** 2)
)
S2 = pm.InverseGamma(
"S2", **estimate_inverse_gamma_parameters(0.25 ** 2, 1.0 ** 2)
)
w1 = pm.InverseGamma(
"w1", **estimate_inverse_gamma_parameters(2 * np.pi / 10.0, np.pi)
)
w2 = pm.InverseGamma(
"w2", **estimate_inverse_gamma_parameters(0.5 * np.pi, 2 * np.pi)
)
log_Q = pm.Uniform("log_Q", lower=np.log(2), upper=np.log(10))
# Set up the kernel an GP
kernel = terms.SHOTerm(S_tot=S1, w0=w1, Q=1.0 / np.sqrt(2))
kernel += terms.SHOTerm(S_tot=S2, w0=w2, log_Q=log_Q)
gp = GP(kernel, t, yerr ** 2, mean=mean)
# Condition the GP on the observations and add the marginal likelihood
# to the model
gp.marginal("gp", observed=y)
```

A few comments here:

The

`term`

interface in*exoplanet*only accepts keyword arguments with names given by the`parameter_names`

property of the term. But it will also interpret keyword arguments with the name prefaced by`log_`

to be the log of the parameter. For example, in this case, we used`log_Q`

as a parameter, but`Q=tt.exp(log_Q)`

would have been equivalent. This is useful because many of the parameters are required to be positive so fitting the log of those parameters is often best.The third argument to the

`exoplanet.gp.GP`

constructor should be the*variance*to add along the diagonal, not the standard deviation as in the original celerite implementation.

To start, let’s fit for the maximum a posteriori (MAP) parameters and look the the predictions that those make.

```
import exoplanet as xo
with model:
map_soln = xo.optimize(start=model.test_point)
```

```
optimizing logp for variables: [log_Q, w2, w1, S2, S1, mean]
18it [00:00, 32.45it/s, logp=1.008735e+01]
message: Optimization terminated successfully.
logp: -4.1063696496869095 -> 10.087352559326797
```

We’ll use the `exoplanet.eval_in_model()`

function to evaluate the
MAP GP model.

```
with model:
mu, var = xo.eval_in_model(
gp.predict(true_t, return_var=True, predict_mean=True), map_soln
)
```

```
plt.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0, label="data")
plt.plot(true_t, true_y, "k", lw=1.5, alpha=0.3, label="truth")
# Plot the prediction and the 1-sigma uncertainty
sd = np.sqrt(var)
art = plt.fill_between(true_t, mu + sd, mu - sd, color="C1", alpha=0.3)
art.set_edgecolor("none")
plt.plot(true_t, mu, color="C1", label="prediction")
plt.legend(fontsize=12)
plt.xlabel("t")
plt.ylabel("y")
plt.xlim(0, 10)
_ = plt.ylim(-2.5, 2.5)
```

Now we can sample this model using PyMC3. There are strong covariances
between the parameters so we’ll use the custom
`exoplanet.get_dense_nuts_step()`

to fit for these covariances
during burn-in.

```
with model:
trace = pm.sample(
tune=2000,
draws=2000,
start=map_soln,
cores=2,
chains=2,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
```

```
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [log_Q, w2, w1, S2, S1, mean]
Sampling 2 chains, 0 divergences: 100%|██████████| 8000/8000 [00:25<00:00, 316.14draws/s]
```

Now we can compute the standard PyMC3 convergence statistics (using
`pymc3.summary`

) and make a trace plot (using `pymc3.traceplot`

).

```
pm.traceplot(trace)
pm.summary(trace)
```

mean | sd | hpd_3% | hpd_97% | mcse_mean | mcse_sd | ess_mean | ess_sd | ess_bulk | ess_tail | r_hat | |
---|---|---|---|---|---|---|---|---|---|---|---|

mean | 0.019 | 0.444 | -0.778 | 0.855 | 0.007 | 0.007 | 4411.0 | 1981.0 | 4477.0 | 2975.0 | 1.0 |

S1 | 0.986 | 0.626 | 0.241 | 1.974 | 0.011 | 0.008 | 3423.0 | 3038.0 | 4644.0 | 2904.0 | 1.0 |

S2 | 0.271 | 0.118 | 0.098 | 0.495 | 0.002 | 0.001 | 4380.0 | 3626.0 | 4875.0 | 3041.0 | 1.0 |

w1 | 0.971 | 0.239 | 0.543 | 1.403 | 0.004 | 0.003 | 4426.0 | 4039.0 | 4536.0 | 3002.0 | 1.0 |

w2 | 3.198 | 0.207 | 2.805 | 3.592 | 0.003 | 0.002 | 4359.0 | 4359.0 | 4613.0 | 2545.0 | 1.0 |

log_Q | 1.807 | 0.351 | 1.182 | 2.302 | 0.005 | 0.004 | 4667.0 | 4667.0 | 4385.0 | 2399.0 | 1.0 |

That all looks pretty good, but I like to make two other results plots: (1) a corner plot and (2) a posterior predictive plot.

The corner plot is easy using `pymc3.trace_to_dataframe`

and I find it
useful for understanding the covariances between parameters when
debugging.

```
import corner
samples = pm.trace_to_dataframe(
trace, varnames=["S1", "S2", "w1", "w2", "log_Q"]
)
_ = corner.corner(samples)
```

The “posterior predictive” plot that I like to make isn’t the same as a
“posterior predictive check” (which can be a good thing to do too).
Instead, I like to look at the predictions of the model in the space of
the data. We could have saved these predictions using a
`pymc3.Deterministic`

distribution, but that adds some overhead to
each evaluation of the model so instead, we can use
`exoplanet.utils.get_samples_from_trace()`

to loop over a few
random samples from the chain and then the
`exoplanet.eval_in_model()`

function to evaluate the prediction
just for those samples.

```
# Generate 50 realizations of the prediction sampling randomly from the chain
N_pred = 50
pred_mu = np.empty((N_pred, len(true_t)))
pred_var = np.empty((N_pred, len(true_t)))
with model:
pred = gp.predict(true_t, return_var=True, predict_mean=True)
for i, sample in enumerate(xo.get_samples_from_trace(trace, size=N_pred)):
pred_mu[i], pred_var[i] = xo.eval_in_model(pred, sample)
# Plot the predictions
for i in range(len(pred_mu)):
mu = pred_mu[i]
sd = np.sqrt(pred_var[i])
label = None if i else "prediction"
art = plt.fill_between(true_t, mu + sd, mu - sd, color="C1", alpha=0.1)
art.set_edgecolor("none")
plt.plot(true_t, mu, color="C1", label=label, alpha=0.1)
plt.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0, label="data")
plt.plot(true_t, true_y, "k", lw=1.5, alpha=0.3, label="truth")
plt.legend(fontsize=12, loc=2)
plt.xlabel("t")
plt.ylabel("y")
plt.xlim(0, 10)
_ = plt.ylim(-2.5, 2.5)
```

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:exoplanet, exoplanet:foremanmackey17, exoplanet:foremanmackey18, exoplanet:pymc3, exoplanet:theano}.

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

```
@misc{exoplanet:exoplanet,
author = {Daniel Foreman-Mackey and Rodrigo Luger and Ian Czekala and
Eric Agol and Adrian Price-Whelan and Tom Barclay},
title = {exoplanet-dev/exoplanet v0.3.2},
month = may,
year = 2020,
doi = {10.5281/zenodo.1998447},
url = {https://doi.org/10.5281/zenodo.1998447}
}
...
```