21 Display of MCMC samples

Code

# Colab setup ------------------
import os, shutil, sys, subprocess, urllib.request
if "google.colab" in sys.modules:
    cmd = "pip install --upgrade polars iqplot colorcet datashader bebi103 arviz cmdstanpy watermark"
    process = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate()
    from cmdstanpy.install_cmdstan import latest_version
    cmdstan_version = latest_version()
    cmdstan_url = f"https://github.com/stan-dev/cmdstan/releases/download/v{cmdstan_version}/"
    fname = f"colab-cmdstan-{cmdstan_version}.tgz"
    urllib.request.urlretrieve(cmdstan_url + fname, fname)
    shutil.unpack_archive(fname)
    os.environ["CMDSTAN"] = f"./cmdstan-{cmdstan_version}"
    data_path = "https://s3.amazonaws.com/bebi103.caltech.edu/data/"
else:
    data_path = "../data/"
# ------------------------------

import warnings

import numpy as np
import polars as pl

import cmdstanpy
import arviz as az

import iqplot
import bebi103

import bokeh.io
bokeh.io.output_notebook()

Loading BokehJS ...

In the previous lesson, we draw samples out of a posterior distribution from a single molecule FISH experiment. We made a corner plot to investigate the results of the sampling. Here, we explore other commonly used methods of displaying MCMC samples. Ultimately, we will conclude that for most applications, a corner plot is an excellent way to display samples and is usually unrivaled.

As a reminder, the model we used is

\[ \begin{align} &\log_{10} \alpha \sim \text{Norm}(0, 1),\\[1em] &\log_{10} b \sim \text{Norm}(2, 1),\\[1em] &\beta = 1/b,\\[1em] &n_i \sim \text{NegBinom}(\alpha, \beta) \;\forall i. \end{align} \]

The corresponding Stan model is

data {
  int<lower=0> N;
  array[N] int<lower=0> n;
}


parameters {
  real log10_alpha;
  real log10_b;
}


transformed parameters {
  real alpha = 10^log10_alpha;
  real b = 10^log10_b;
  real beta_ = 1.0 / b;
}


model {
  // Priors
  log10_alpha ~ normal(0, 1);
  log10_b ~ normal(2, 1);

  // Likelihood
  n ~ neg_binomial(alpha, beta_);
}

Let’s acquire samples for this model applied to the mRNA counts of the rest gene, just as we did in the previous lesson.

# Load data set
df = pl.read_csv(os.path.join(data_path, "singer_transcript_counts.csv"), comment_prefix="#")

# Construct data dict, making sure data are ints
data = dict(N=len(df), n=df["Rest"].to_numpy())

# Sample using Stan
with bebi103.stan.disable_logging():
    sm = cmdstanpy.CmdStanModel(stan_file='smfish.stan')
    samples = sm.sample(    
        data=data,
        chains=4,
        iter_sampling=1000,
        show_progress=False,
    )

# Convert to ArviZ InferenceData instance
samples = az.from_cmdstanpy(posterior=samples)

In our visualizations, we will only investigate the \(\alpha\) and \(b\) parameters, as the others are all transformations of these.

samples.posterior = samples.posterior[['alpha', 'b']]

21.1 Visualization with ArviZ

ArviZ has lots of handy visualization of MCMC results, and it is a very promising package to use to quickly get plots. It currently still has some drawbacks. First, some plots are not rendered properly with Bokeh. This problem will likely be fixed very soon with future releases. Second, some of the aesthetic choices in the plots would not be my first choice. There are other specific drawbacks to the ArviZ visualizations that I will point out as we go along.

For each of the visualizations below, I show how to construct the visualization using ArviZ, but also using the bebi103 package.

21.2 Examining traces

The first type of visualization we will explore is useful for diagnosing potential problems with the sampler.

21.2.1 Trace plots

Trace plots are a common way of visualizing the trajectory a sampler takes through parameter space. We plot the value of the parameter versus step number. You can make these plots using bebi103.viz.trace_plot(). By default, the traces are colored according to chain number. In our sampling we just did, we used Stan’s default of four chains.

21.2.1.1 Trace plots with ArviZ

To make plots with ArviZ, the basic syntax is az.plot_***(samples, backend='bokeh'), where *** is the kind of plot you want. The backend="bokeh" kwarg indicates that you want the plot rendered with Bokeh. Let’s make a trace plot.

# Have to catch Bokeh deprecation warnings because ArviZ still uses circle
# and not scatter
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    az.plot_trace(samples, backend="bokeh");

ArviZ generates two plots for each variable. The plots to the right show the trace of the sampler. In this plot, the x-axis is the iteration number of the steps of the walker and the y-axis is the values of the parameter for that step.

ArviZ also plots a picture of the marginalized posterior distribution for each trace to the left. The issue I have these plots because they use kernel density estimation for the posterior plot. KDE involves selecting a bandwidth, which leaves an adjustable parameter. I prefer simply to plot the ECDFs of the samples to visualize the marginal distributions, as I show below.

21.2.1.2 Trace plots with bebi103

To make trace plots with the bebi103 module, you can do the following.

bokeh.io.show(bebi103.viz.trace(samples))

21.2.1.3 Interpetation of trace plots

These trace plots look pretty good; the sampler is bounding around a central value. Trace plots are very commonly used, which is why I present them here, but are not particularly useful. This is not just my opinion. Here is what Dan Simpson has to say about trace plots.

21.2.2 Parallel coordinate plots

As Dan Simpson pointed out in his talk I linked to above, plots that visualize diagnostics effectively are better. We will talk more about diagnostics of MCMC in later lessons, but for now, I will display one of the diagnostic plots Dan showed in his talk. Here, we put the parameter names on the x-axis, and we plot the values of the parameters on the y-axis. Each line in the plot represents a single sample of the set of parameters.

21.2.2.1 Parallel coordinate plots with ArviZ

To make sure we compare things of the same magnitude so we can better see the details of how the parameters vary with respect to each other, we scale the samples by the minimum and maximum values.

bokeh.io.show(
    az.plot_parallel(
        samples, var_names=["alpha", "b"], backend="bokeh", norm_method="minmax"
    )
)

21.2.2.2 Parallel coordinate plots with bebi103

The bebi103 package offers parallel coordinate plots as well.

bokeh.io.show(
    bebi103.viz.parcoord(
        samples,
        transformation="minmax",
    )
)

21.2.2.3 Intepretation of parallel coordinate plots

Normally, samples with problems are plotted in another color to help diagnose the problems. In this particular set of samples, there were no problems, so everything looks fine.

Interestingly, the “neck” between alpha and b is indicative of anticorrelation. When the burst rate \(\alpha\) is high in the posterior, the burst size \(b\) is low, and vice-versa.

21.3 Plots of marginalized distributions

While we cannot in general plot a multidimensional distribution, we can plot marginalized distributions, either of one parameter or of two.

21.3.1 Plotting marginalized distributions of one parameter

There are three main options for plotting marginalized distributions of a single parameter.

A kernel density estimate of the marginalized posterior.
A histogram of the marginalized posterior.
An ECDF of the samples out of the posterior for a particular parameter. This approximates the CDF of the marginalized distribution.

I prefer (3), though (2) is a good option as well. For continuous parameters, ArviZ offers only a KDE plot. (It will offer an ECDF plot similar to what iqplot offers in the upcoming 1.0 release.)

21.3.2 Plotting marginalized distributions with ArviZ

We already saw that ArviZ can give these plots along with trace plots. We can also directly plot them using the az.plot_density(). When using this function, the distributions are truncated at the bound of the highest probability density region, or HPD. If we’re considering a 95% credible interval, the HPD interval is the shortest interval that contains 95% of the probability of the posterior. By default, az.plot_density() truncates the plot of the KDE of the PDF for a 94% HPD.

az.plot_density(
    samples, backend="bokeh", backend_kwargs=dict(frame_width=200, frame_height=150)
);

21.3.2.1 Plotting marginalized distributions with iqplot

To get samples out of the marginalized posterior for a single parameter, we simply ignore the values of the parameters that are not the one of interest. We can then use iqplot to make plots of histograms or ECDFs.

To do so, we need to convert the posterior samples to a tidy data frame.

# Convert to data frame
df_mcmc = bebi103.stan.arviz_to_dataframe(samples)

# Take a look
df_mcmc.head()

shape: (5, 5)

alpha	b	chain__	draw__	diverging__
f64	f64	i64	i64	bool
3.87773	18.999069	0	0	false
3.9700863	19.687562	0	1	false
3.8875494	19.524372	0	2	false
4.2078469	17.78808	0	3	false
4.3169623	18.041961	0	4	false

Now, we can use iqplot to make histograms.

hists = [
    iqplot.histogram(
        df_mcmc, q=param, density=True, rug=False, frame_height=150
    )
    for param in ["alpha", "b"]
]

bokeh.io.show(bokeh.layouts.column(*hists))

A better option, in my opinion, is to make ECDFs of the samples.

ecdfs = [
    iqplot.ecdf(df_mcmc, q=param, frame_height=150)
    for param in ["alpha", "b"]
]

bokeh.io.show(bokeh.layouts.column(*ecdfs))

21.3.3 Marginal posteriors of two parameters and corner plots

We can also plot the two-dimensional distribution (of most interest here are the parameters \(\phi\) and \(\gamma\)). The simplest way to plot these is simply to plot each point, possibly with some transparency.

21.3.3.1 Pair plots with ArviZ

We might like to do this for all pairs of plots. ArviZ enables this to be done conveniently. (It uses a previous version of Bokeh, so the glyphs do not show properly.)

az.plot_pair(
    samples,
    var_names=['alpha', 'b'],
    scatter_kwargs=dict(fill_alpha=0.1, size=2, radius_units='screen'),
    backend="bokeh",
    backend_kwargs=dict(frame_width=200, frame_height=200),
);

By adding the marginals=True kwarg, the pair plots now also include representation of the univariate marginal distributions.

az.plot_pair(
    samples,
    var_names=['alpha', 'b'],
    marginals=True,
    backend="bokeh",
    backend_kwargs=dict(frame_width=200, frame_height=200),
);

A plot like this is often called a corner plot. It shows all bi- and uni-variate marginal distributions.

ArviZ also offers two-dimensional binning in the form of a hex plot.

az.plot_pair(
    samples,
    var_names=["alpha", "b"],
    kind="hexbin",
    marginals=True,
    backend='bokeh',
);

21.3.4 Corner plots with bebi103

Corner plots are also implemented in the bebi103 package.

bokeh.io.show(
    bebi103.viz.corner(
        samples,
        parameters=[("alpha", "α"), "b"],
        xtick_label_orientation=np.pi / 4,
    )
)

This is a nice way to summarize the posterior and is useful for visualizing how various parameters covary. We can also do a corner plot with the one-parameter marginalized posteriors represented as CDFs.

bokeh.io.show(
    bebi103.viz.corner(
        samples,
        parameters=[("alpha", "α"), "b"],
        xtick_label_orientation=np.pi / 4,
        plot_ecdf=True,
    )
)

Corner plots are my preferred method of displaying results. All possible display-able marginal posteriors are plotted and laid out in a logical way.

bebi103.stan.clean_cmdstan()

21.4 Computing environment

%load_ext watermark
%watermark -v -p numpy,polars,cmdstanpy,arviz,bokeh,iqplot,bebi103,jupyterlab
print("cmdstan   :", bebi103.stan.cmdstan_version())

Python implementation: CPython
Python version       : 3.13.11
IPython version      : 9.8.0

numpy     : 2.2.6
polars    : 1.36.1
cmdstanpy : 1.3.0
arviz     : 0.23.0
bokeh     : 3.8.1
iqplot    : 0.3.7
bebi103   : 0.1.28
jupyterlab: 4.5.1

cmdstan   : 2.37.0