MCMC

The maximum likelihood estimation (MLE) only gives you a point estimate, what if you want to get a distribution over the parameter space? That’s when Markov Chain Monte Carlo (MCMC) comes into play. EzTao provides a simple function, eztao.ts.carma_mcmc.mcmc, to quickly run MCMC using \(\mathit{emcee}\) (the MCMC hammer). However, since it is just a wrapper, it has very limited options. A more advanced/flexible solution is to use EzTao for likelihood computation and hook it with a MCMC sampler of your choice.

In this notebook, we will show how to run MCMC using EzTao’s built-in MCMC function as well as your own MCMC sampler.

Built-in MCMC

# general packages
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
%matplotlib inline

# eztao, emcee, corner
import eztao
from eztao.carma import CARMA_term
from eztao.ts import gpSimRand
from eztao.ts import mcmc
import emcee
import corner

mpl.rc_file(os.path.join(eztao.__path__[0], "viz/eztao.rc"))

Simulate a CARMA(3,0) process and plot the time series

# simulate a CARMA(3,0) process
p = 3
q = 0
SNR = 100
duration = 365*10.0
npts = 2000
carma30_kernel = CARMA_term(np.log([3, 3.189, 0.05]), np.log([0.5]))
t, y, yerr = gpSimRand(carma30_kernel, SNR, duration, npts, log_flux=False)

# plot the simulated process
fig, ax = plt.subplots(1,1, dpi=120, figsize=(8,3))
ax.errorbar(t, y, yerr, fmt='.', markersize=4)
ax.set_xlabel('Time (day)')
ax.set_ylabel('Flux (arb. unit)')
ax.set_title('Simulated CARMA(3,0) process', fontsize=14)

Text(0.5, 1.0, 'Simulated CARMA(3,0) process')

Run MCMC using the built-in function & Generate a corner plot

# use built-in function to run MCMC
sampler, carma_flatchain, carma_chain = mcmc(t, y, yerr, p, q)

Searching for best-fit CARMA parameters...
Running burn-in...
Running production...

# remove points with low prob for the sake of making good corner plot
prob_threshold = np.percentile(sampler.flatlnprobability, 5)
clean_chain = carma_flatchain[sampler.flatlnprobability > prob_threshold, :]

# make corner plot
labels = [name for name in carma30_kernel.get_parameter_names()]
corner.corner(clean_chain, truths=carma30_kernel.get_parameter_vector(),
              quantiles=[0.16, 0.5, 0.84], labels=labels, show_titles=True,
              title_kwargs={"fontsize": 12});

Note

The mcmc function returns three variables, the emcee sampler object, the flatchain and the chain in regular CARMA space. For p > 2, given the reasons mentioned in the Fiitting notebook (CARMA space vs. polynomial space), the MCMC sampler runs in the polynomial space (which regular user does not need to worry about), mcmc returns the chain and flatchain separately from the emcee sampler. If the p order is 2 or smaller, the last two variables will just be empty arrays.

Use your own MCMC sampler

Here, I will use emcee for demonstration. To run MCMC with emcee, you need to define your own probability function. You can either write a simple wrapper around the EzTao likelihood function, neg_param_ll or neg_fcoeff_ll (depending on the p order of the CARMA model), to return the POSITIVE log likelihood, or write your own probability function.

p <= 2

from eztao.ts import neg_param_ll, drw_fit
from eztao.carma import DRW_term
from celerite import GP

Simulate a DRW process and plot the time series

# simulate a DRW process
amp = 0.2
tau = 120
drw_kernel = DRW_term(np.log(amp), np.log(tau))

SNR = 10
duration = 365*10.0
npts = 100
t2, y2, yerr2 = gpSimRand(drw_kernel, SNR, duration, npts, log_flux=False)

# plot the simulated process
fig, ax = plt.subplots(1,1, dpi=120, figsize=(8,3))
ax.errorbar(t2, y2, yerr2, fmt='.', markersize=4)
ax.set_xlabel('Time (day)')
ax.set_ylabel('Flux (arb. unit)')
ax.set_title('Simulated DRW process')

Text(0.5, 1.0, 'Simulated DRW process')

Obtain best-fit parameters and use that to initialize a GP model

# obtain best-fit
best_drw = drw_fit(t2, y2, yerr2)
print(f'Best-fit DRW: {best_drw}')

# define celerite GP model
drw_gp = GP(DRW_term(*np.log(best_drw)), mean=np.median(y2))
drw_gp.compute(t2, yerr2)

Best-fit DRW: [  0.20333737 134.83935452]

Run MCMC & Generate a corner plot

# define log prob function
def param_ll(*args):
    return -neg_param_ll(*args)

# initialize the walker, specify number of walkers, prob function, args and etc.
initial = np.array(np.log(best_drw))
ndim, nwalkers = len(initial), 32
sampler_drw = emcee.EnsembleSampler(nwalkers, ndim, param_ll, args=[y2, drw_gp])

# run a burn-in surrounding the best-fit parameters obtained above
print("Running burn-in...")
p0 = initial + 1e-8 * np.random.randn(nwalkers, ndim)
p0, lp, _ = sampler_drw.run_mcmc(p0, 500)

# clear up the stored chain from burn-in, rerun the MCMC
print("Running production...")
sampler_drw.reset()
sampler_drw.run_mcmc(p0, 2000);

Running burn-in...
Running production...

# remove points with low prob for the sake of making good corner plot
prob_threshold_drw = np.percentile(sampler_drw.flatlnprobability, 3)
clean_chain_drw = sampler_drw.flatchain[sampler_drw.flatlnprobability > prob_threshold_drw, :]

# make corner plot
labels = [name for name in drw_gp.kernel.get_parameter_names()]
corner.corner(clean_chain_drw, truths=drw_kernel.get_parameter_vector(),
              quantiles=[0.16, 0.5, 0.84], labels=labels, show_titles=True,
              title_kwargs={"fontsize": 12});

Note

It is always a good practice to start your MCMC walker at a position that is close to the truth (assuming the best-fit is a good estimation) as MCMC is a sampler not an optimizer.

p > 2

In this section, we will reuse the simulated CARMA(3,0) process in Built-in MCMC. Since p > 2, we need to start and run MCMC in the factored polynomial space of CARMA, and we will also need to transform the chain from the polynomial space to regular CARMA space once it is done. EzTao provides the functions (see here) to facilitate this process as we will see below.

Obtain best-fit CARMA parameters & Convert them into the polynomial space

from eztao.ts import carma_fit

# obtain best-fit as intial position for MCMC
best_cm30 = carma_fit(t, y, yerr, p, q, n_opt=50)
print(f'Best-fit CARMA(3,0) parameters:{best_cm30}')

# get the representation of the best-fit in the polynomial space
best_fcoeffs_log = CARMA_term.carma2fcoeffs_log(np.log(best_cm30[:p]),
                                            np.log(best_cm30[p:]))
best_fcoeffs = np.exp(best_fcoeffs_log)
print(f'Best-fit in polynomial space: {best_fcoeffs}')

Best-fit CARMA(3,0) parameters:[3.18331428 3.44275143 0.05733812 0.52477284]
Best-fit in polynomial space: [3.1663963  3.3891824  0.01691798 0.52477284]

Initialize a celerite GP model

from eztao.ts import neg_fcoeff_ll

# init GP
cm30_kernel_mcmc = CARMA_term(np.log(best_cm30)[:p], np.log(best_cm30)[p:])
cm30_gp = GP(cm30_kernel_mcmc, mean=np.median(y))
cm30_gp.compute(t, yerr)

Run MCMC

# define log prob function
def fcoeff_ll(*args):
    return -neg_fcoeff_ll(*args)

# run MCMC
initial_cm30 = np.log(best_fcoeffs)
ndim, nwalkers = len(initial_cm30), 32
sampler_cm30 = emcee.EnsembleSampler(nwalkers, ndim, fcoeff_ll, args=[y, cm30_gp])

print("Running burn-in...")
p0 = initial_cm30 + 1e-8 * np.random.randn(nwalkers, ndim)
p0, lp, _ = sampler_cm30.run_mcmc(p0, 500)

print("Running production...")
sampler_cm30.reset()
sampler_cm30.run_mcmc(p0, 2000);

Running burn-in...
Running production...

Convert the chain from the polynomial space to the CARMA space

# from the polynomial space to CARMA space
vec_fcoeff2carma_log = np.vectorize(CARMA_term.fcoeffs2carma_log, excluded=[1,],
                                signature="(n)->(m),(k)")
logAR, logMA = vec_fcoeff2carma_log(sampler_cm30.flatchain, p)
cm30_flatchain = np.hstack((logAR, logMA))

Make a corner plot

# remove points with low prob for the sake of making good corner plot
prob_threshold_cm30 = np.percentile(sampler_cm30.flatlnprobability, 5)
clean_chain_cm30 = cm30_flatchain[sampler_cm30.flatlnprobability > prob_threshold_cm30, :]

# make corner plot
labels = [name for name in carma30_kernel.get_parameter_names()]
corner.corner(clean_chain_cm30, truths=carma30_kernel.get_parameter_vector(),
              quantiles=[0.16, 0.5, 0.84], labels=labels, show_titles=True,
              title_kwargs={"fontsize": 12});

Note:

Extra cautions should be taken when running MCMC with CARMA models of p > 2:

1. Use neg_fcoeff_ll as the base log likelihood function.

2. Remember to transform best-fit parameters to the polynomial space when initializing the walker and the chain back to the regular CARMA space once the MCMC is finished.