<jupyter_start><jupyter_text>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](../python/ts.rstmodule-eztao.ts.carma_mcmc), to quickly run MCMC using $\mathit{emcee}$ ([the MCMC hammer](https://emcee.readthedocs.io/en/stable/)). 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<jupyter_code># 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"))<jupyter_output><empty_output><jupyter_text>Simulate a CARMA(3,0) process and plot the time series<jupyter_code># 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)<jupyter_output><empty_output><jupyter_text>Run MCMC using the built-in function & Generate a corner plot<jupyter_code># use built-in function to run MCMC
sampler, carma_flatchain, carma_chain = mcmc(t, y, yerr, p, q)
# 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});<jupyter_output><empty_output><jupyter_text>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](03_Fit.ipynbFit-CARMA-models-that-are-more-complex-than-CARMA(2,1)) 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<jupyter_code>from eztao.ts import neg_param_ll, drw_fit
from eztao.carma import DRW_term
from celerite import GP<jupyter_output><empty_output><jupyter_text>Simulate a DRW process and plot the time series<jupyter_code># 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')<jupyter_output><empty_output><jupyter_text>Obtain best-fit parameters and use that to initialize a GP model<jupyter_code># 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)<jupyter_output>Best-fit DRW: [  0.20333737 134.83935452]<jupyter_text>Run MCMC & Generate a corner plot<jupyter_code># 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);
# 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});<jupyter_output><empty_output><jupyter_text>NoteIt 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 > 2In this section, we will reuse the simulated CARMA(3,0) process in [Built-in MCMC](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](../python/carma.rsteztao.carma.CARMATerm.CARMA_term.carma2fcoeffs)) to facilitate this process as we will see below.  Obtain best-fit CARMA parameters & Convert them into the polynomial space<jupyter_code>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}')<jupyter_output>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]<jupyter_text>Initialize a celerite GP model<jupyter_code>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)<jupyter_output><empty_output><jupyter_text>Run MCMC<jupyter_code># 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);<jupyter_output>Running burn-in...
Running production...<jupyter_text>Convert the chain from the polynomial space to the CARMA space<jupyter_code># 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))<jupyter_output><empty_output><jupyter_text>Make a corner plot<jupyter_code># 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});<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10747663551401869, 'fraction_numerical': 0.03571428571428571, 'mean_word_length': 5.379701916252661, 'pattern_counts': {'":': 3, '<': 47, '<?xml version=': 0, '>': 53, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '397080', 'license': 'MIT', 'n_tokens_mistral': 3468, 'n_tokens_neox': 3234, 'n_words': 906, 'path': 'docs/notebooks/04_MCMC.ipynb', 'repo_name': 'ywx649999311/EzTao', 'size': '8988'}