26  Parameter estimation with Markov chain Monte Carlo

| Download notebook

Data set download

---

# 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 numpy as np
import scipy.stats as st
import polars as pl

import cmdstanpy
import arviz as az

import iqplot
import bebi103

import bokeh.io
import bokeh.plotting
bokeh.io.output_notebook()

Loading BokehJS ...

---

In this lesson, we will learn how to use Markov chain Monte Carlo to do parameter estimation. To get the basic idea behind MCMC, imagine for a moment that we can draw samples out of the posterior distribution. This means that the probability of choosing given values of a set of parameters is proportional to the posterior probability of that set of values. If we drew many many such samples, we could reconstruct the posterior from the samples, e.g., by making histograms. That’s a big thing to imagine: that we can draw properly weighted samples. But, it turns out that we can! That is what MCMC allows us to do.

We discussed some theory behind this seemingly miraculous capability in lecture. For this lesson, we will just use the fact that we can do the sampling to learn about posterior distributions in the context of parameter estimation.

26.1 The data set

The data come from the Elowitz lab, published in Singer et al., Dynamic Heterogeneity and DNA Methylation in Embryonic Stem Cells, Molec. Cell, 55, 319-331, 2014, available here. In the following paragraphs, I repeat the description of the data set and EDA from last term:

In this paper, the authors investigated cell populations of embryonic stem cells using RNA single molecule fluorescence in situ hybridization (smFISH), a technique that enables them to count the number of mRNA transcripts in a cell for a given gene. They were able to measure four different genes in the same cells. So, for one experiment, they get the counts of four different genes in a collection of cells.

The authors focused on genes that code for pluripotency-associated regulators to study cell differentiation. Indeed, differing gene expression levels are a hallmark of differentiated cells. The authors do not just look at counts in a given cell at a given time. The temporal nature of gene expression is also important. While the authors do not directly look at temporal data using smFISH (since the technique requires fixing the cells), they did look at time lapse fluorescence movies of other regulators. We will not focus on these experiments here, but will discuss how the distribution of mRNA counts acquired via smFISH can serve to provide some insight about the dynamics of gene expression.

The data set we are analyzing now comes from an experiment where smFISH was performed in 279 cells for the genes rex1, rest, nanog, and prdm14. The data set may be downloaded at https://s3.amazonaws.com/bebi103.caltech.edu/data/singer_transcript_counts.csv.

26.2 ECDFs of mRNA counts

We will do a quick EDA to get a feel for the data set by generating ECDFs for the mRNA counts for each of the four genes.

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

genes = ["Nanog", "Prdm14", "Rest", "Rex1"]

plots = [
    iqplot.ecdf(
        data=df[gene],
        q=gene,
        x_axis_label="mRNA count",
        title=gene,
        frame_height=150,
        frame_width=200,
    )
    for gene in genes
]

bokeh.io.show(
    bokeh.layouts.column(bokeh.layouts.row(*plots[:2]), bokeh.layouts.row(*plots[2:]))
)

Note the difference in the \(x\)-axis scales. Clearly, prdm14 has far fewer mRNA copies than the other genes. The presence of two inflection points in the Rex1 EDCF implies bimodality.

26.3 Building a generative model

We can model the transcript counts, which result from bursty gene expression, as being Negative Binomially distributed. (The details behind this model are a bit nuanced, and you can read about them here.) For a given gene, the likelihood for the counts is

\[ \begin{align} n_i \mid \alpha, b \sim \text{NegBinom}(\alpha, 1/b) \;\forall i, \end{align} \tag{26.1}\]

where \(\alpha\) is the burst frequency (higher \(\alpha\) means gene expression comes on more frequently) and \(b\) is the burst size, i.e., the typical number of transcripts made per burst. We have therefore identified the two parameters we need to estimate, \(\alpha\) and \(b\).

Because the Negative Binomial distribution is often parametrized in terms of \(\alpha\) and \(\beta= 1/b\), we can alternatively state our likelihood as

\[ \begin{align} &\beta = 1/b,\\[1em] &n_i \mid \alpha, \beta \sim \text{NegBinom}(\alpha, \beta)\;\; \forall i. \end{align} \tag{26.2}\]

Given that we have a Negative Binomial likelihood, we are left to specify priors the burst size \(b\) and the burst frequency \(\alpha\).

26.3.1 Priors for burst size and inter-burst time

We will apply the bet-the-farm technique to get our priors for the burst size and inter-burst times. I would expect the time between bursts to be longer than a second, since it takes time for the transcriptional machinery to assemble. I would expect it to be shorter than a few hours, since an organism would need to adapt its gene expression based on environmental changes on that time scale or faster. The time between bursts needs to be in units of RNA lifetimes, and bacterial RNA lifetimes are of order minutes. So, the range of values of \(\alpha\) is \(10^{-2}\) to \(10^2\), leading to a prior of

\[\begin{align} \log_{10} \alpha \sim \text{Norm}(0, 1). \end{align} \]

I would expect the burst size to depend on promoter strength and/or strength of transcriptional activators. I could imagine anywhere from a few to a few thousand transcripts per burst, giving a range of \(10^0\) to \(10^4\), and a prior of

\[\begin{align} \log_{10} b \sim \text{Norm}(2, 1). \end{align} \]

We then have the following model.

\[ \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} \tag{26.3}\]

26.4 Sampling the posterior

To draw samples out of the posterior, we need to use some new Stan syntax. Here is the Stan code we will use with some notes about Stan syntax.

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_);
}

• Note that the raise-to-power operator is ^, not ** as in Python.
• The data block contains the counts \(n\) of the mRNA transcripts. There are \(N\) cells that are measured. Most data blocks look like this. There is an integer parameter that specifies the size of the data set, and then the data set is given as an array. The declaration that n is an array of length N is array [N], followed by the type of data, is integer. We specified a lower bound on the data (as we will do on the parameters) using the <lower=0> syntax.
• The parameters block tells us what the parameters of the posterior are. In this case, we wish to sample out of the posterior \(g(\alpha, b \mid \mathbf{n})\), where \(\mathbf{n}\) is the set of transcript counts for the gene. So, the two parameters are \(\alpha\) and \(b\). However, since defining the prior was more easily defined in terms of logarithms, we specify \(\log_{10} \alpha\) and \(\log_{10} b\) as the parameters.
• The transformed parameters block allows you to do any transformation of the parameters you are sampling for convenience. In this case, Stan’s Negative Binomial distribution is parametrized by \(\beta = 1/b\), so we make the transformation of the b to beta_. Notice that I have called this variable beta_ and not beta. I did this because beta is one of Stan’s distributions, and you should avoid naming a variable after a word that is already in the Stan language. The other transformations we need to make involve converting the logarithms to the actual parameter values.
• Finally, the model block is where the model is specified. The syntax of the model block is almost identical to that of the hand-written model.

Now that we have specified our model, we can compile it.

sm = cmdstanpy.CmdStanModel(stan_file='smfish.stan')

With our compiled model, we just need to specify the data and let Stan’s sampler do the work! When using CmdStanPy, the data has to be passed in as a dictionary with keys corresponding to the variable names declared in the data block of the Stan program and values as Numpy arrays with the appropriate data type. For this calculation, we will use the data set for the rest gene.

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

# Sample using Stan
samples = sm.sample(    
    data=data,
    chains=4,
    iter_sampling=1000,
)

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

12:15:48 - cmdstanpy - INFO - CmdStan start processing

12:15:48 - cmdstanpy - INFO - CmdStan done processing.
12:15:48 - cmdstanpy - WARNING - Non-fatal error during sampling:
Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Inverse scale parameter is 0, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Inverse scale parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Inverse scale parameter is 0, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
    Exception: neg_binomial_lpmf: Shape parameter is inf, but must be positive finite! (in 'smfish.stan', line 26, column 2 to column 33)
Consider re-running with show_console=True if the above output is unclear!

We got lots of warnings! In particular, we get warnings that some of the parameters fed into the Negative Binomial distribution are invalid, being either zero or infinite. These warnings are arising during Stan’s warm-up phase as it is assessing optimal settings for sampling, and should not be of concern. It is generally a bad idea to silence warnings, but if you are sure that the warnings that the sampler will throw are of no concern, you can silence logging using the bebi103.stan.disable_logging context. In most notebooks in these notes, to avoid clutter for pedagogical purposes, we will disable the warnings.

with bebi103.stan.disable_logging():
    samples = sm.sample(    
        data=data,
        chains=4,
        iter_sampling=1000,
    )

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

Now, let’s take a quick look at the samples.

samples.posterior

<xarray.Dataset> Size: 168kB
Dimensions:      (chain: 4, draw: 1000)
Coordinates:
  * chain        (chain) int64 32B 0 1 2 3
  * draw         (draw) int64 8kB 0 1 2 3 4 5 6 ... 993 994 995 996 997 998 999
Data variables:
    log10_alpha  (chain, draw) float64 32kB 0.5804 0.6351 ... 0.646 0.6322
    log10_b      (chain, draw) float64 32kB 1.32 1.252 1.21 ... 1.237 1.241
    alpha        (chain, draw) float64 32kB 3.805 4.316 4.407 ... 4.425 4.288
    b            (chain, draw) float64 32kB 20.91 17.89 16.2 ... 17.27 17.42
    beta_        (chain, draw) float64 32kB 0.04783 0.05591 ... 0.05789 0.05742
Attributes:
    created_at:                 2025-08-20T09:15:49.005458+00:00
    arviz_version:              0.22.0
    inference_library:          cmdstanpy
    inference_library_version:  1.2.5

xarray.Dataset

• Dimensions:
  • chain: 4
  • draw: 1000
• Coordinates: (2)
  • chain
    (chain)
    int64
    0 1 2 3

    array([0, 1, 2, 3])

  • draw
    (draw)
    int64
    0 1 2 3 4 5 ... 995 996 997 998 999

    array([  0,   1,   2, ..., 997, 998, 999], shape=(1000,))

• Data variables: (5)
  • log10_alpha
    (chain, draw)
    float64
    0.5804 0.6351 ... 0.646 0.6322

    array([[0.580387, 0.635072, 0.644183, ..., 0.633148, 0.655338, 0.703629],
           [0.607035, 0.604947, 0.660854, ..., 0.720163, 0.71799 , 0.593938],
           [0.713801, 0.715316, 0.703968, ..., 0.630683, 0.628306, 0.706312],
           [0.610528, 0.600729, 0.641074, ..., 0.598044, 0.645953, 0.632221]],
          shape=(4, 1000))

  • log10_b
    (chain, draw)
    float64
    1.32 1.252 1.21 ... 1.237 1.241

    array([[1.32032, 1.25249, 1.20962, ..., 1.22404, 1.23303, 1.16735],
           [1.2667 , 1.25687, 1.23603, ..., 1.15223, 1.15658, 1.29306],
           [1.15525, 1.16491, 1.15836, ..., 1.24806, 1.24704, 1.16517],
           [1.26757, 1.26331, 1.22422, ..., 1.26765, 1.23738, 1.24093]],
          shape=(4, 1000))

  • alpha
    (chain, draw)
    float64
    3.805 4.316 4.407 ... 4.425 4.288

    array([[3.80528, 4.31591, 4.40741, ..., 4.29682, 4.52207, 5.05393],
           [4.04608, 4.02668, 4.57988, ..., 5.25005, 5.22384, 3.92589],
           [5.1737 , 5.19178, 5.05787, ..., 4.27251, 4.24919, 5.08524],
           [4.07876, 3.98776, 4.37596, ..., 3.96318, 4.42541, 4.28767]],
          shape=(4, 1000))

  • b
    (chain, draw)
    float64
    20.91 17.89 16.2 ... 17.27 17.42

    array([[20.9083, 17.885 , 16.2038, ..., 16.751 , 17.1012, 14.7012],
           [18.4798, 18.0665, 17.2197, ..., 14.1982, 14.3409, 19.6364],
           [14.2971, 14.6188, 14.3999, ..., 17.7034, 17.6622, 14.6274],
           [18.5168, 18.3363, 16.7579, ..., 18.5204, 17.2736, 17.4153]],
          shape=(4, 1000))

  • beta_
    (chain, draw)
    float64
    0.04783 0.05591 ... 0.05789 0.05742

    array([[0.0478279, 0.0559126, 0.0617138, ..., 0.059698 , 0.0584756,
            0.0680215],
           [0.0541131, 0.055351 , 0.058073 , ..., 0.0704317, 0.0697308,
            0.0509258],
           [0.0699441, 0.0684051, 0.0694452, ..., 0.0564862, 0.0566181,
            0.068365 ],
           [0.054005 , 0.0545367, 0.0596733, ..., 0.0539944, 0.0578917,
            0.0574209]], shape=(4, 1000))

• Indexes: (2)
  • chain
    PandasIndex

    PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

  • draw
    PandasIndex

    PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
           ...
           990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
          dtype='int64', name='draw', length=1000))

• Attributes: (4)

  created_at :

  2025-08-20T09:15:49.005458+00:00

  arviz_version :

  0.22.0

  inference_library :

  cmdstanpy

  inference_library_version :

  1.2.5

As we have already seen, the samples are indexed by chain and draw. Parameters represented in the parameters and transformed parameters blocks are reported.

26.5 Plots of the samples

There are many ways of looking at the samples. In this case, since we have two parameters of interest, the pulse frequency and pulse size, we can plot the samples as a scatter plot to get the approximate density.

def plot_scatter(samples, title=None):
    p = bokeh.plotting.figure(
        frame_width=200,
        frame_height=200,
        x_axis_label='α',
        y_axis_label='b',
        title=title,
    )
    
    p.scatter(
        samples.posterior['alpha'].values.ravel(), 
        samples.posterior['b'].values.ravel(),
        size=2,
        alpha=0.2,
    )

    return p

bokeh.io.show(plot_scatter(samples))

We see very strong correlation between \(\alpha\) and \(b\). This does not necessarily mean that they depend on each other. Rather, it means that our degree of belief about their values depends on both in a correlated way. The measurements we made cannot effectively separate the effects of \(\alpha\) and \(b\) on the transcript counts.

26.6 Marginalizing the posterior

We can also plot the marginalized posterior distributions. Remember that the marginalized distributions properly take into account the effects of the other variable, including the strong correlation I just mentioned. To obtain the marginalized distribution, we simply ignore the samples of the parameters we are marginalizing out. It is convenient to look at the marginalized distributions as ECDFs.

plots = [
    iqplot.ecdf(
        samples.posterior[param].values.ravel(),
        q=param,
        frame_height=200,
        frame_width=250,
    )
    for param in ["alpha", "b"]
]

bokeh.io.show(bokeh.layouts.row(*plots))

Alternatively, we can visualize the marginalized posterior PDFs as histograms. Because we have such a large number of samples, binning bias from histograms is less of a concern.

plots = [
    iqplot.histogram(
        samples.posterior[param].values.ravel(),
        q=param,
        rug=False,
        frame_height=200,
        frame_width=250,
    )
    for param in ["alpha", "b"]
]

bokeh.io.show(bokeh.layouts.row(*plots))

26.7 Corner plots

We now have a two-dimensional posterior distribution, with our two parameters being \(\alpha\) and \(b\). We can combine a plot of the samples from the full posterior with the histograms (or ECDFs) of samples from the marginal distributions in a corner plot, available via bebi103.viz.corner().

bokeh.io.show(
    bebi103.viz.corner(
        samples, parameters=[('alpha', 'α'), 'b']
    )
)

Corner plots generalize to dimensions beyond two. The off-diagonal plots are of samples from marginal distributions where two parameters remain and the diagonals are plots from univariate marginal distributions.

bebi103.stan.clean_cmdstan()

26.8 Computing environment

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

Python implementation: CPython
Python version       : 3.13.5
IPython version      : 9.4.0

numpy     : 2.2.6
scipy     : 1.16.0
polars    : 1.31.0
cmdstanpy : 1.2.5
arviz     : 0.22.0
bokeh     : 3.7.3
iqplot    : 0.3.7
bebi103   : 0.1.28
jupyterlab: 4.4.5

cmdstan   : 2.36.0