MCMC Parameter Trace Plots

Looking at how each parameter chain evolves across the chains can show if any particular chains are diverging from the rest, if there are multiple modes being explored, or how wide the distribution of parameters are relative to each other.

Each line corresponds to one chain in the ensemble.

import matplotlib.pyplot as plt
import numpy as np
from espei.analysis import truncate_arrays

# ESPEI writes trace.npy and lnprob.npy; the documentation stores this run as a
# compressed .npz archive to keep the repository small.
mcmc_results = np.load('../_data/mcmc-results.npz')
trace = mcmc_results['trace']
lnprob = mcmc_results['lnprob']
trace, lnprob = truncate_arrays(trace, lnprob)

num_chains = trace.shape[0]
num_parameters = trace.shape[2]
iterations = np.arange(1,trace.shape[1]+1)
fig, axs = plt.subplots(nrows=num_parameters, gridspec_kw=dict(hspace=0.4), figsize=(4, 3*num_parameters))
for parameter in range(num_parameters):
    ax = axs[parameter]
    ax.set_xlabel('Iterations')
    ax.set_ylabel('Parameter value')
    ax.plot(iterations, trace[..., parameter].T)
fig.show()

and if we plot again after 500 iterations of burn-in:

burn_in_iterations = 500
fig, axs = plt.subplots(nrows=num_parameters, gridspec_kw=dict(hspace=0.4), figsize=(4, 3*num_parameters))
for parameter in range(num_parameters):
    ax = axs[parameter]
    ax.set_xlabel('Iterations')
    ax.set_ylabel('Parameter value')
    ax.plot(iterations[burn_in_iterations:], trace[:, burn_in_iterations:, parameter].T)
fig.show()