This page was generated from docs/examples/getting_started.ipynb. Interactive online version:

Getting started

caustic is a code for probabilistic modeling of single lens gravitational microlensing events. It is designed primarily to serve as a testbed for new approaches to modeling microlensing events rather than something that can be used in existing pipelines. caustic is built on top of the probabilistic programming language PyMC3 which is used in many fields for the purporse of building and testing Bayesian models. The main feature of PyMC3 is that it enables the use of Hamiltonian Monte Carlo (HMC) via the No U-Turn Sampler (NUTS). HMC enables efficient sampling of high dimensional probabilistic models by utilizing the gradients of the log probability with respect the model parameters. These gradients are obtained through automatic differentiation, a method widely used in machine learning to obtain exact gradientss of computer code by repeatedly apply the chain rule from calculus.

Automatic differention (AD) is different from symbolic differentiation (like for example in Wolfram Mathematica) and the finite difference method used in numerical solvers for differential equations. Automatic differentiation is not trivial to do so one usually relies existing libraries which implement it. PyMC3 uses Theano under the hood. You can think of Theano as a sort of compiler that takes code written in Python, compiles it to C, builds a graph of all mathematical computations in the code and propagates the gradients through every node in the graph using reverse mode automatic differentiation (backpropagation). Theano is best known today as a predecessor to ultra popular machine learning libraries Tensorflow and PyTorch whose main purpose is also automatic differentiation. Because of the requirement that all models written in PyMC3 have to be compliant with Theano, building physical models is somewhat more complicated than if we were to use numpy, although the API is very similar.

Check out the PyMC3 docs for an introduction to PyMC3 and probabilistic programming, docs on Theano for an overview of Theano, and the very recent exoplanet code for building differentiable astrophysical models on top of PyMC3 in the context of explanet transits and radial velocity measurements. caustic is very similar to exoplanet and uses it as a dependancy.

In this notebook, I’ll walk you through loading a microlensing dataset and fitting a simple model with caustic.

Loading microlensing datasets

All of the data handling is done using the caustic.data module. The purpose of the module is to take raw lightcurves provided by an observatory or collaboration and transform it into a common format which can then used within a context of a PyMC3 model for inference. Here’s an example where we load an OGLE light curve using the caustic.data.OGLEData class.

import numpy as np
from matplotlib import pyplot as plt

import caustic as ca

np.random.seed(42)

event = ca.data.OGLEData("../../data/OGLE-2017-BLG-0324")

The data is stored as a list of astropy.table.Table tables, each table corresponds to one time series observation in a certain band.

event.light_curves

[<Table length=626>
      HJD        mag   mag_err mask
    float64    float64 float64 bool
 ------------- ------- ------- ----
 2457069.87628   18.64   0.035 True
 2457074.88664  18.604   0.038 True
 2457075.86271  18.639   0.036 True
 2457078.86521  18.638   0.031 True
 2457082.86705   18.58   0.026 True
 2457086.89543   18.58   0.032 True
 2457090.82554  18.652   0.039 True
  2457097.8947   18.65   0.023 True
 2457100.87989  18.614   0.028 True
 2457110.80722  18.677   0.027 True
           ...     ...     ...  ...
 2458014.55412  18.645   0.034 True
 2458015.55245  18.586   0.041 True
 2458016.56098  18.607   0.036 True
  2458018.4914  18.641   0.039 True
 2458019.52922  18.593   0.027 True
 2458020.56318  18.613   0.034 True
 2458026.52622  18.645   0.052 True
  2458029.5112  18.641   0.049 True
 2458030.51879  18.696   0.039 True
 2458040.49979  18.614   0.055 True]

Each table contains metadata information about the observatory and filter used

for table in event.light_curves:
    print(table.meta)

{'filter': 'I', 'observatory': 'OGLE'}

Plotting the data is straightforward

fig, ax = plt.subplots(figsize=(8, 4))
event.plot(ax)

findfont: Font family ['sans-serif'] not found. Falling back to DejaVu Sans.
findfont: Font family ['sans-serif'] not found. Falling back to DejaVu Sans.
findfont: Font family ['sans-serif'] not found. Falling back to DejaVu Sans.

The times of observation are specified in Helocentric Julian Days (HJD), the source brightness is either units of magnitudes or flux (defined such that a magnitude of 22 corresponds to unit flux).

event.units = "fluxes"

fig, ax = plt.subplots(figsize=(8, 4))
event.plot(ax)

The coordinates of an event are stored as a astropy.coordinates.SkyCoord object

event.event_coordinates

<SkyCoord (ICRS): (ra, dec) in deg
    (264.61391667, -29.90372222)>

It’s possible to add different datasets together as event_combined = event1 + event2.

Fitting a simple model

All models we build with caustic are just PyMC3 models, all caustic does is that it provides some extra functionality specific to microlensing. The user initializes a base model caustic.models.SingleLensModel inherits from pymc3.model.Model class and provides basic functionality such as transforming the data into theano tensors and methods for computing the likelihood. One can then build on that model by initializing PyMC3 random variables and calling methods for computing the trajectory of the lens and the resulting magnification.

Let’s first initialize a SingleLensModel object which takes a caustic.data object as an input, transforms the data to a standardized format where all of the light curves are independently rescaled to zero median and unit standard deviation and converted to theano tensors. If you want to fit the data without this rescaling step, pass standardize=False to the class constructor.

# Initialize a SingleLensModel object
model = ca.models.SingleLensModel(event)

The data is stored as a list of theano.tensor objects for time observations \(\{\mathbf{t}_j\}\), fluxes \(\{\mathbf{F}_j\}\) and reported uncertainties \(\{\boldsymbol{\sigma}_j\}\) where \(j\) goes over all bands.

import theano.tensor as T

# Time
T.shape(model.t[0]).eval()  # .eval() evaluates the theano computation graph

array([626])

model.t

[TensorConstant{[7069.8762..040.49979]}]

# Flux
model.F

[TensorConstant{[22.091522...6442127 ]}]

# Error bars
model.sigF

[TensorConstant{[0.7123316...14782127]}]

The model we are going to fit is the simplest in microlensing. It models the magnification of a background source star by a single point lens assuming uniform colinear motion between the observer, the lens and the source. The magnification is then given by the following simple expression

\[A(u)=\frac{u^2+2}{u\sqrt{u^2+4}}\]

where the trajectory \(u(t)\) is

\[u(t)=\sqrt{u_0^2+\left( \frac{t-t_0}{t_E}\right)^2}\]

The predicted flux \(\mathbf{f}_j\) in the \(j\)-th band is then

\[\mathbf{f}_j=\boldsymbol{\Phi}\,\mathbf w_j\]

where the matrix \(\boldsymbol{\Phi}\) is defined by

\[\begin{split}\boldsymbol{\Phi}=\left(\begin{array}{cc} {\tilde{A}_{1}\left(t_{1}\right)} & {1} \\ {\tilde{A}_{2}\left(t_{2}\right)} & {1} \\ {\vdots} & {\vdots} \\ {\tilde{A}_{N}\left(t_{N}\right)} & {1} \end{array}\right)\end{split}\]

with \(\tilde{A}(t)\equiv (A(t)-1)/(A(t_0)-1)\) and \(\mathbf{w}_j\) is a vector of linear parameters

\[\begin{split}\mathbf{w}_j=\left(\begin{array}{c}{\Delta F}_j \\ {F_{\mathrm{base}, j}}\end{array}\right)\end{split}\]

where the parameter \(\Delta F_j\) represents the difference between the flux at peak magnification and the baseline flux \(F_\mathrm{base, j}\) in the \(j\)-th band.

We assume that the likelihood function for the \(j\)-th band is a multivariate gaussian with mean \(\mathbf{f}_j\) and a covariance matrix \(\mathbf C\). The likelihood for the time series data in the \(j\)-th band is then

\[p(\mathbf{F}_j|\boldsymbol{\theta})\sim\mathcal{N}(\mathbf{F}_j; \mathbf{f}_j,\mathbf{C}_j)\]

where the symbol \(\mathcal{N}\) denotes a multivariate normal distribution. In this example, we take the covariance matrix to be diagonal

\[\mathbf{C}_j=\mathrm{diag}\left(\sigma_{1,j}^2, \sigma_{2,j}^2,\dots,\sigma_{N,j}^2\right)\]

where the elements on the diagonal are the squares of the uncertainties for each point. In the the gaussian_processes notebook we relax this assumption by modelling off diagonal terms with a GP. Since the reported uncertainties in microlensing are generally underestimated we’ll introduce a constant rescaling by a parameter \(c\)

\[\tilde{\sigma}_{i,j}=c\,\sigma_{i,j}\]

The total likelihood for data observed in multiple bands is then the product of the individual likelihoods for each band because the the noise properties are conditionally independent on the parameters between the different bands.

\[p(\{\mathbf{F}\}^{J}_{j=1}|\boldsymbol{\theta})=\prod_{j=1}^J p(\mathbf{F}_j|\boldsymbol{\theta})\]

First, we initialize the linear flux parameters as PyMC3 random variables. The shape of these parameters has to be \((N_\mathrm{bands}, 1)\). It’s good practice to make use of the testval= keyword for every RV in PyMC3 to initialize the RV because MCMC is extremely sensitive to initial values of the model parameters.

import pymc3 as pm

with model:
    n_bands = len(event.light_curves)

    # Initialize linear parameters
    testval_ln_DeltaF = T.log(
        ca.utils.estimate_peak_flux(event) - ca.utils.estimate_baseline_flux(event)
    )  # helper functions
    ln_DeltaF = pm.Normal("ln_DeltaF", mu=4.0, sd=4, testval=testval_ln_DeltaF[0])
    DeltaF = T.exp(ln_DeltaF)

    ln_Fbase = pm.Normal(
        "ln_Fbase", mu=2, sd=4, testval=T.log(ca.utils.estimate_baseline_flux(event)[0])
    )
    Fbase = T.exp(ln_Fbase)

To print all free parameters in the model, we do the following

model.vars

[ln_DeltaF, ln_Fbase]

Let’s specify the rest of the model

with model:
    # Initialize nonlinear parameters
    t0 = pm.Uniform(
        "t0", model.t_min, 2 * model.t_max, testval=ca.utils.estimate_t0(event)
    )
    ln_tE = pm.Normal("ln_tE", mu=3.0, sd=6, testval=2.0)

    tE = pm.Deterministic("tE", T.exp(ln_tE))

    u0 = pm.Exponential("u0", 0.5, testval=0.1)

    # Compute the source magnitude and blending fraction
    m_source, g = ca.utils.compute_source_mag_and_blend_fraction(
        event, DeltaF, Fbase, u0
    )
    pm.Deterministic("m_source", m_source)
    pm.Deterministic("g", g)

    # Compute the trajectory of the les
    trajectory = ca.trajectory.Trajectory(event, t0, u0, tE)
    u = trajectory.compute_trajectory(model.t)

    # Compute the magnification
    mag = model.compute_magnification(u, u0)

    # Compute the mean model
    mean = DeltaF * mag + Fbase

    # We allow for rescaling of the error bars by a constant factor
    ln_c = pm.Exponential("ln_c", 0.5, testval=0.8 * T.ones(n_bands))

    # Diagonal terms of the covariance matrix
    var_F = (T.exp(ln_c) * model.sigF) ** 2

    # Compute the Gaussian log_likelihood, add it as a potential term to the model
    ll = model.compute_log_likelihood(model.F - mean, var_F)
    pm.Potential("log_likelihood", ll)

The _lowerbound__ and _interval__ words mean that the original bounded parameters (for example restricted to be positive) have been transformed using a deterministic transform such that they end up being defined on the entire domain of real numbers. PyMC3 does this by default to improve sampling efficiency. In the final trace containing the samples one can access both the transformed and the original parameters.

model.vars

[ln_DeltaF, ln_Fbase, t0_interval__, ln_tE, u0_log__, ln_c_log__]

We’re now ready to sample the model with NUTS. We use Dan Foreman-Mackey’s modification to the default PyMC3 sampler exoplanet.get_dense_nuts_step() which includes an improved tuning schedule. By default, the NUTS implementation in PyMC3 assumes a diagonal mass matrix. The modified sampling procedure in exoplanet uses a dense matrix which is tuned in 4 short runs to estimate the inverse of the covariance matrix of the model parameters. Effectively this means that the parameter space is rescaled such that the posterior is as close as possible to a multivariate Gaussian with an identity matrix as the covariance matrix. For more details of the imporoved tuning schedule, see here.

import exoplanet as xo

with model:
    # Print initial logps
    print("Model test point:\n", model.check_test_point())

    # Run sampling
    trace = pm.sample(tune=500, draws=1000, cores=4, step=xo.get_dense_nuts_step())

Model test point:
 ln_DeltaF       -2.32
ln_Fbase        -2.35
t0_interval__   -2.52
ln_tE           -2.72
u0_log__        -3.05
ln_c_log__      -1.32
Name: Log-probability of test_point, dtype: float64

Multiprocess sampling (4 chains in 4 jobs)
NUTS: [ln_c, u0, ln_tE, t0, ln_Fbase, ln_DeltaF]
Sampling 4 chains: 100%|██████████| 6000/6000 [00:27<00:00, 219.10draws/s]
The acceptance probability does not match the target. It is 0.949121747900346, 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.9364621813178776, 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.9331330374691215, 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.9288506137586271, but should be close to 0.8. Try to increase the number of tuning steps.

To check how the sampling turned out, we can use built in methods from PyMC3

pm.summary(trace)

	mean	sd	mc_error	hpd_2.5	hpd_97.5	n_eff	Rhat
ln_DeltaF	3.347099	0.007774	0.000151	3.330426	3.361087	2583.825845	1.000448
ln_Fbase	3.108515	0.002098	0.000042	3.104439	3.112666	2320.372424	1.000301
ln_tE	3.888819	0.065415	0.001426	3.765297	4.018028	1900.108533	1.000432
t0	7864.489414	0.142709	0.002214	7864.207866	7864.767268	3605.545327	0.999997
tE	48.957640	3.199823	0.069948	43.121147	55.522479	1914.344009	1.000416
u0	0.356397	0.034552	0.000762	0.289729	0.420631	1855.049830	1.000673
m_source	19.089149	0.139616	0.003062	18.829419	19.365581	1875.257228	1.000633
g	0.545908	0.195856	0.004315	0.185496	0.943943	1904.287496	1.000616
ln_c__0	0.189425	0.028463	0.000661	0.131577	0.242081	1885.092701	0.999722

Not bad, we have thousands of effective samples for less than a minute of sampling time.

We can also plot a traceplot of the chain

pm.traceplot(trace, figsize=(8, 12));

findfont: Font family ['sans-serif'] not found. Falling back to DejaVu Sans.

The traceplot also includes the deterministic variables we’ve specified in the model. In this case we’ve added a variable m_source, the posteior over the source star brightness in magnitudes and the blending ratio g which is the ratio of the background flux due to other unresolved stars in the vicinity of the source star and the source star flux.

Let’s check the autocorrelation times of the chains

pm.autocorrplot(trace, figsize=(6, 15));

and plot the posterior means and highest posterior density intervals

pm.plot_posterior(trace, figsize=(8, 8));

findfont: Font family ['sans-serif'] not found. Falling back to DejaVu Sans.

and a corner plot

pm.pairplot();

When interpreting this plot, we have to remember that the actual sampling happens in the sapace of the transformed parameters while the pm.pairplot() function plots the samples in the original space.

Finally, we can plot posterior realizations of the model in data space. This is more complicated than it sounds because we have to evaluate the model prediction which is a theano.tensor for multiple draws from the posterior. To do this, we recompute the model prediction tensor for a denser grid in time, and pass that tensor to the function plot_model_and_residuals defined in caustic.utils. Internally, the prediction tensor is then evaluated for different parameter dictionaries corresponding to different samples from the posterior.

with model:
    # Create dense grid
    t_dense = np.tile(np.linspace(model.t_min, model.t_max, 2000), (n_bands, 1))
    t_dense_tensor = T.as_tensor_variable(t_dense)

    # Compute the trajectory of the lens
    u_dense = trajectory.compute_trajectory(t_dense_tensor)

    # Compute the magnification
    mag_dense = model.compute_magnification(u_dense, u0)

    # Compute the mean model
    mean_dense = DeltaF * mag_dense + Fbase

def plot_model_and_residuals(
    ax, data, samples, t_grid, prediction, gp_list=None, model=None, **kwargs
):
    model = pm.modelcontext(model)

    n_samples = len(samples)

    # Load data
    if model.is_standardized is True:
        tables = data.get_standardized_data()
    else:
        tables = data.get_standardized_data(rescale=False)

    # Evaluate model for each sample on a fine grid
    n_pts_dense = T.shape(t_grid)[1].eval()
    n_bands = len(data.light_curves)

    prediction_eval = np.zeros((n_samples, n_bands, n_pts_dense))

    # Evaluate predictions in model context
    if gp_list is None:
        for i, sample in enumerate(samples):
            prediction_eval[i] = xo.eval_in_model(prediction, sample)

    else:
        prediction_tensors = [gp_list[n].predict(t_grid[n]) for n in range(n_bands)]
        for i, sample in enumerate(samples):
            for n in range(n_bands):
                prediction_eval[i, n] = xo.eval_in_model(prediction_tensors[n], sample)

            # Add mean model to GP prediction
            prediction_eval[i] += xo.eval_in_model(prediction, sample)

    # Plot model predictions for each different samples from posterior on dense
    # grid
    for n in range(n_bands):  # iterate over bands
        for i in range(n_samples):
            ax[0].plot(
                t_grid[n].eval(),
                prediction_eval[i, n, :],
                color="C" + str(n),
                alpha=0.2,
                **kwargs,
            )

    # Compute median of the predictions
    median_predictions = np.zeros((n_bands, n_pts_dense))
    for n in range(n_bands):
        median_predictions[n] = np.percentile(prediction_eval[n], [16, 50, 84], axis=0)[
            1
        ]

    # Plot data
    data.plot_standardized_data(ax[0], rescale=model.is_standardized)
    ax[0].set_xlabel(None)
    ax[1].set_xlabel("HJD - 2450000")
    ax[1].set_ylabel("Residuals")
    ax[0].set_xlim(T.min(t_grid).eval(), T.max(t_grid).eval())

    # Compute residuals with respect to median model
    for n in range(n_bands):
        # Interpolate median predictions onto a grid of observed times
        median_pred_interp = np.interp(
            tables[n]["HJD"], t_grid[n].eval(), median_predictions[n]
        )

        residuals = tables[n]["flux"] - median_pred_interp

        ax[1].errorbar(
            tables[n]["HJD"],
            residuals,
            tables[n]["flux_err"],
            fmt="o",
            color="C" + str(n),
            alpha=0.5,
            **kwargs,
        )
        ax[1].grid(True)

# Plot model
fig, ax = plt.subplots(
    2, 1, gridspec_kw={"height_ratios": [3, 1]}, figsize=(8, 4), sharex=True
)


# Get random samples from trace
samples = list(xo.get_samples_from_trace(trace, size=50))

with model:
    plot_model_and_residuals(ax, event, samples, t_dense_tensor, mean_dense)