22: Simulation based calibration and related checks in practice

Data set download

---

import numpy as np
import pandas as pd
import scipy.stats as st

import cmdstanpy
import arviz as az

import iqplot

import bebi103

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

/home/ec2-user/miniconda/envs/bebi103/lib/python3.11/site-packages/dask/dataframe/_pyarrow_compat.py:23: UserWarning: You are using pyarrow version 11.0.0 which is known to be insecure. See https://www.cve.org/CVERecord?id=CVE-2023-47248 for further details. Please upgrade to pyarrow>=14.0.1 or install pyarrow-hotfix to patch your current version.
  warnings.warn(

Loading BokehJS ...

---

You should set the value of the variable cores to be the number of cores you have available on your machine. I will be using 95 cores in this notebook, since I am using a big machine on AWS to do my SBC.

cores = 95

---

In the previous lesson, we laid out the a principled pipeline for constructing and testing a generative model and associated inference procedures. In this lesson, we work through the implementation of the principled pipeline on a familiar data set. We will again look at the RNA FISH data set from this paper from the Elowitz lab. You can download the data set here. If you want to refresh yourself about this data set, you can read its description in a previous lesson.

ECDFs of mRNA counts

Let’s go ahead and load the data. In our analysis here, we will use the Rest gene.

# Load DataFrame
df = pd.read_csv(os.path.join(data_path, 'singer_transcript_counts.csv'), comment='#')

# Pull out data for Stan
n = df['Rest'].values
data = dict(N=len(n), n=n)

# Take a look
bokeh.io.show(iqplot.ecdf(n, x_axis_label='mRNA count'))

The generative model

When we first used MCMC with this data set, we used a Negative Binomial likelihood (which has both a theoretical and empirical justification), parametrized by the burst size \(b\) and the burst frequency \(\alpha\). We had the following generative 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}

We can code up prior predictive checks and the model in Stan. First, the prior predictive checks.

data {
  int N;
}


generated quantities {
  array[N] int n;

  real log10_alpha = normal_rng(0.0, 1.0);
  real log10_b = normal_rng(2.0, 1.0);
  real alpha = 10^log10_alpha;
  real b = 10^log10_b;
  real beta_ = 1 / b;

  for (i in 1:N) {
    n[i] = neg_binomial_rng(alpha, beta_);
  }
}

And also the model.

data {
  int N;
  array[N] int n;
}


parameters {
  real<lower=0> log10_alpha;
  real<lower=0> 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.0, 1.0);
  log10_b ~ normal(2.0, 1.0);

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

For now, we are not going to bother with posterior predictive checks or computing the log likelihood.

Let’s compile the models.

with bebi103.stan.disable_logging():
    sm_prior_pred = cmdstanpy.CmdStanModel(stan_file='prior_pred.stan')
    sm = cmdstanpy.CmdStanModel(stan_file='model.stan')

We can now perform prior predictive checks. We will plot the resulting checks as ECDFs so we can see how the mRNA counts are distributed. For the plot, to avoid choking the browser, we will only plot 100 ECDFS.

with bebi103.stan.disable_logging():
    samples_prior_pred = sm_prior_pred.sample(
        data=data, fixed_param=True, chains=1, iter_sampling=1000
    )

samples_prior_pred = az.from_cmdstanpy(
    posterior=samples_prior_pred, prior=samples_prior_pred, prior_predictive="n"
)

p = None
for n in samples_prior_pred.prior_predictive.n.squeeze()[::10]:
    p = iqplot.ecdf(
        n, marker_kwargs=dict(fill_alpha=0.2, line_alpha=0.2), p=p, x_axis_type="log"
    )

p.x_range = bokeh.models.Range1d(0.3, 3e5)

bokeh.io.show(p)

We can also plot the mean and variance of all of the generated data sets as to further characterize the prior predictive distribution.

means = samples_prior_pred.prior_predictive.n.squeeze().mean(axis=1).values
variances = samples_prior_pred.prior_predictive.n.squeeze().var(axis=1).values

p = bokeh.plotting.figure(
    frame_height=250,
    frame_width=250,
    x_axis_label='mean of counts',
    y_axis_label='variance of counts',
    x_axis_type='log',
    y_axis_type='log',
    x_range=[1, np.nanmax(means)],
    y_range=[1, np.nanmax(variances)],
)

p.circle(means, variances, size=2)

bokeh.io.show(p)

This also makes sense. We get Poissonian behavior (mean = variance) for some samples, and then a range of dispersion beyond that.

The prior predictive checks show a wide range of mRNA counts, and all seem reasonable. We do get some large number of counts, upwards of 10,000, considering that the typical total mRNA count in a mammalian cell is about 100,000. But this is not dominant, and we get good coverage over what we might expect, so this seems like a pretty good prior.

Performing SBC

Performing SBC really only requires a few ingredients. First, we need the requisite data to be used for prior predictive checks. In this case, it is just the number of measurements we are making, \(N\). Second, we need a Stan model to generate the prior predictive data sets. Finally, we need a Stan model to sample out of the posterior. The bebi103.stan.sbc() function will then perform SBC and give the results back in a data frame. That is, it will draw a prior predictive data set, use that data set in a posterior sampling by MCMC calculation, and then compute the useful diagnostics and statistics (z-score, shrinkage, and rank statistic) from those samples. It does this N times (not to be confused with \(N\), the number of measurements in the experiment). Let’s now put it to use to perform SBC.

df_sbc = bebi103.stan.sbc(
    prior_predictive_model=sm_prior_pred,
    posterior_model=sm,
    prior_predictive_model_data=data,
    posterior_model_data=data,
    measured_data=["n"],
    var_names=["alpha", "b"],
    measured_data_dtypes=dict(n=int),
    cores=cores,
    N=1000,
    progress_bar=True,
)

100%|███████████████████████████████████████| 1000/1000 [00:42<00:00, 23.30it/s]

The bebi103.stan.sbc() function gives a data frame with the SBC analysis results. Let’s take a look at the data frame to see what it has.

df_sbc.head()

	ground_truth	rank_statistic	mean	sd	shrinkage	z_score	Rhat	ESS	ESS_per_iter	tail_ESS	tail_ESS_per_iter	n_divergences	n_bad_ebfmi	n_max_treedepth	warning_code	L	trial	parameter
0	0.418733	0	1.004741	0.004720	1.000000	124.159304	1.002847	2063.132436	0.515783	1795.518948	0.448880	2	0	0	4	4000	0	alpha
1	0.193699	0	1.002572	0.002566	1.000000	315.236969	1.000743	1938.450809	0.484613	1478.583786	0.369646	0	0	0	0	4000	1	alpha
2	22.689700	3446	20.879789	1.709294	0.998972	-1.058865	1.007017	579.443127	0.144861	378.743290	0.094686	0	0	0	1	4000	2	alpha
3	1.026900	2032	1.034750	0.029606	1.000000	0.265146	1.002357	1211.397615	0.302849	1097.282499	0.274321	0	0	0	0	4000	3	alpha
4	0.298719	0	1.003109	0.003085	1.000000	228.290468	1.001857	1673.597065	0.418399	1116.781554	0.279195	0	0	0	0	4000	4	alpha

For each trial, for each parameter, we get diagnostic results, z-scores, shrinkage, rank statistic, posterior mean and standard deviations for each ground truth, as well as the ground truth used in the posterior sampling. The warning_code column gives a succinct summary of the diagnostic warnings. You can parse a warning code using the bebi103.stan.parse_warning_code() function. As an example, I’ll test it on warning code 14.

bebi103.stan.parse_warning_code(14)

Rhat warning
divergence warning
treedepth warning

To visualize the results of SBC, we can first make a plot of the z-scores and of shrinkage. Ideally, the shrinkage should all be close to one, and the magnitude of the z-scores should all be less than five. Let’s take a look!

p = bokeh.plotting.figure(
    frame_height=250,
    frame_width=300,
    x_axis_label="shrinkage",
    y_axis_label="z-score",
    tooltips=[("trial", "@trial")],
)

for color, (parameter, sub_df) in zip(
    ["#1f77b4", "orange"], df_sbc.groupby("parameter")
):
    p.circle(
        source=sub_df[["shrinkage", "z_score", "trial"]],
        x="shrinkage",
        y="z_score",
        size=2,
        color=color,
    )

bokeh.io.show(p)

Oof! We are severely overfitting the model, as evidenced by z-scores of very large magnitude. We are missing the ground truth.

To diagnose why, let’s look at which samples have reasonable z-scores. We’ll make a strip plot categorizing the results of SBC by parameter and whether or not the z-score is good.

df_sbc['good_z'] = np.abs(df_sbc['z_score']) < 5
p = iqplot.strip(
    df_sbc,
    cats=['parameter', 'good_z'],
    color_column='good_z',
    order=(('alpha', True), ('alpha', False), ('b', True), ('b', False)),
    q='ground_truth',
    x_axis_type='log',
    spread='jitter',
)

bokeh.io.show(p)

Most strikingly, the z-score is poor for \(\alpha < 1\). Recall that for a Negative Binomial distribution, the mean is \(\alpha b\). So, when \(\alpha\) is small, the mean can be less than one, meaning that most of the counts generated by the model are zero. It makes sense, then, that we will miss the ground truth, since the data are almost all zero; there is nothing to properly inform the posterior.

This immediately identifies a possible problem with our inference pipeline. If a data set comes through with mostly zero measurements, we will not be able to make reliable inferences. SBC has thus identified a problem area look out for when doing our inference.

Having a typical burst size less than one is actually unphysical, since no transcripts are created. To be “on,” we would need to make at least one transcript. So, the SBC has exposed a problem in our modeling that we didn’t see before. Not only can the data fail to inform the prior for these parameter values, we have also discovered that our model can give unphysical parameter values. We will abort continued analysis of our SBC results and instead adapt our model.

An adjusted prior

I would expect the time between bursts to be of order minutes, since that is a typical response time to signaling of a cell. This is of the same order of magnitude of an RNA lifetime, so I might then expect \(\alpha\) to be of order unity.

\begin{align} \alpha \sim \text{Gamma}(1.25, 0.1). \end{align}

We can make a quick plot.

alpha = np.linspace(0, 50, 200)
g = st.gamma.pdf(alpha, 1.25, loc=0, scale=1/0.1)

p = bokeh.plotting.figure(
    frame_width=300,
    frame_height=200,
    x_axis_label='α',
    y_axis_label='g(α)',
)

p.line(alpha, g, line_width=2)

bokeh.io.show(p)

This is still pretty broad and pushes some of the prior probability mass away from zero.

Turning now to the burst size, I would expect \(b\) to depend on promoter strength and/or strength of transcriptional activators. I could imagine anywhere from a few to several thousand transcripts per burst.

\begin{align} b \sim \text{Gamma}(2, 0.002). \end{align}

Again, with a plot.

b = np.linspace(0, 5000, 200)
g = st.gamma.pdf(b, 2, loc=0, scale=1/0.002)

p = bokeh.plotting.figure(
    frame_width=300,
    frame_height=200,
    x_axis_label='b',
    y_axis_label='g(b)',
)

p.line(b, g, line_width=2)

bokeh.io.show(p)

This prior moves \(b\) off of zero, which we saw was problematic in our previous prior. The Gamma prior also decays faster than our original Log-Normal prior, which ended up getting us very large burst sizes. We then have the following model.

\begin{align} &\alpha \sim \text{Gamma}(1.25, 0.1), \\[1em] &b \sim \text{Gamma}(2, 0.002), \\[1em] &\beta = 1/b,\\[1em] &n_i \sim \text{NegBinom}(\alpha, \beta) \;\forall i. \end{align}

We can code this model up and check the prior predictive checks. The Stan code is as follows.

data {
  int N;
}


generated quantities {
  array[N] int n;

  real alpha = gamma_rng(1.25, 0.1);
  real b = gamma_rng(2.0, 0.002);
  real beta_ = 1.0 / b;

  for (i in 1:N) {
    n[i] = neg_binomial_rng(alpha, beta_);
  }
}

Let’s get some samples and look at the ECDFs of the copy numbers again.

with bebi103.stan.disable_logging():
    sm_prior_pred_2 = cmdstanpy.CmdStanModel(stan_file='prior_pred_2.stan')
    samples_prior_pred = sm_prior_pred_2.sample(
        data=data, fixed_param=True, chains=1, iter_sampling=1000
    )

samples_prior_pred = az.from_cmdstanpy(
    posterior=samples_prior_pred, prior=samples_prior_pred, prior_predictive="n"
)

p = None
for n in samples_prior_pred.prior_predictive.n.squeeze()[::10]:
    p = iqplot.ecdf(
        n, marker_kwargs=dict(fill_alpha=0.2, line_alpha=0.2), p=p, x_axis_type="log",
        x_range=[0.3, 1e6]
    )

bokeh.io.show(p)

Most of the data sets have reasonable ECDFs. Importantly, we see that the most number of zeros we get in any one data set is about 60% or so of the counts. These data sets again seem to match our intuition. Let’s check the mean and variance of transcript counts.

means = samples_prior_pred.prior_predictive.n.squeeze().mean(axis=1).values
variances = samples_prior_pred.prior_predictive.n.squeeze().var(axis=1).values

p = bokeh.plotting.figure(
    frame_height=250,
    frame_width=250,
    x_axis_label='mean of counts',
    y_axis_label='variance of counts',
    x_axis_type='log',
    y_axis_type='log',
    x_range=[1, np.nanmax(means)],
    y_range=[1, np.nanmax(variances)],
)

p.circle(means, variances, size=2)

bokeh.io.show(p)

This looks good. We can now code up the Stan model and run SBC on this, hopefully improved, model. We will now include posterior predictive checks because we will ultimately use this model. The Stan code is as follows.

data {
  int N;
  array[N] int n;
}


parameters {
  real<lower=0> alpha;
  real<lower=0> b;
}


transformed parameters {
  real beta_ = 1.0 / b;
}


model {
  // Priors
  alpha ~ gamma(1.25, 0.1);
  b ~ gamma(2.0, 0.002);

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


generated quantities {
  array[N] int n_ppc;
  for (i in 1:N) {
    n_ppc[i] = neg_binomial_rng(alpha, beta_);
  }
}

Let’s compile!

with bebi103.stan.disable_logging():
    sm_2 = cmdstanpy.CmdStanModel(stan_file='model_2.stan')

And now we can conduct SBC with this updated model. Because we have posterior predictive checks, we need to make sure to tell bebi103.stan.sbc() which variables are posterior predictive (or log likelihood, though we do not have that in this model).

df_sbc = bebi103.stan.sbc(
    prior_predictive_model=sm_prior_pred_2,
    posterior_model=sm_2,
    prior_predictive_model_data=data,
    posterior_model_data=data,
    measured_data=["n"],
    var_names=["alpha", "b"],
    measured_data_dtypes=dict(n=int),
    posterior_predictive_var_names=["n_ppc"],
    cores=cores,
    N=1000,
    progress_bar=True,
)

100%|███████████████████████████████████████| 1000/1000 [01:09<00:00, 14.34it/s]

This time, let’s check the diagnostics first. We can get the count of each warning type.

# Divide by two because diagnostics are listed for each parameter
df_sbc.groupby('warning_code').size() / 2

warning_code
0    922.0
1      5.0
2     70.0
3      3.0
dtype: float64

We have two warning types, type 1 (ESS warning) and type 2 (Rhat warning). (A type 3 warning is both Rhat and ESS.) To deal with these, we can increase the number of iterations we take. Note that this is an important feature of performing these SBC calculations; we can see what kinds of difficulties we might encounter in our sampling.

df_sbc = bebi103.stan.sbc(
    prior_predictive_model=sm_prior_pred_2,
    posterior_model=sm_2,
    prior_predictive_model_data=data,
    posterior_model_data=data,
    measured_data=["n"],
    var_names=["alpha", "b"],
    measured_data_dtypes=dict(n=int),
    posterior_predictive_var_names=['n_ppc'],
    sampling_kwargs=dict(iter_warmup=2000, iter_sampling=2000),
    cores=cores,
    N=1000,
    progress_bar=True,
)

100%|███████████████████████████████████████| 1000/1000 [02:14<00:00,  7.41it/s]

Let’s again check the diagnostics.

df_sbc.groupby('warning_code').size() / 2

warning_code
0    1000.0
dtype: float64

Our diagnostics are much better! Now, let’s make a plot of the z-score versus shrinkage.

p = bokeh.plotting.figure(
    frame_height=250,
    frame_width=300,
    x_axis_label="shrinkage",
    y_axis_label="z-score",
    tooltips=[("trial", "@trial")],
)

for color, (parameter, sub_df) in zip(
    ["#1f77b4", "orange"], df_sbc.groupby("parameter")
):
    p.circle(
        source=sub_df[["shrinkage", "z_score", "trial"]],
        x="shrinkage",
        y="z_score",
        size=2,
        color=color,
    )

bokeh.io.show(p)

We have good z-scores for all trials, and decent shrinkage. This all looks good. Let’s now do the self-consistency check with the rank statistic. Recall that the rank statistics should be Uniformly distributed. Therefore, the ECDFs of the rank statistics should fall on a diagonal line. When we plot the ECDF, we can also plot an envelope which encompasses the 99% confidence interval for the ECDF of a Uniformly distributed random variable.

bokeh.io.show(bebi103.viz.sbc_rank_ecdf(df_sbc, diff=False))

It looks like the rank statistic is Uniformly distributed. We can see this more clearly if we instead plot the difference of the ECDF to the theoretical ECDF of a Uniformly distributed random variable.

bokeh.io.show(bebi103.viz.sbc_rank_ecdf(df_sbc))

In this clearer view, we see that most of the rank statistics all live within the 99% envelope, so we are in good shape.

With everything checking out, we can perform our sampling with real data!

Sampling with our new model

We’ll now use our model with updated priors to perform parameter estimation using our real data set, checking all diagnostics after the fact, of course.

with bebi103.stan.disable_logging():
    samples = sm_2.sample(data=data)
    samples = az.from_cmdstanpy(posterior=samples, posterior_predictive='n_ppc')

bebi103.stan.check_all_diagnostics(samples)

Effective sample size looks reasonable for all parameters.

Rhat looks reasonable for all parameters.

0 of 4000 (0.0%) iterations ended with a divergence.

0 of 4000 (0.0%) iterations saturated the maximum tree depth of 10.

E-BFMI indicated no pathological behavior.

0

Let’s take a look at the corner plot.

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

This result looks very much like what we achieved in Lesson 8, so the small adjustment in prior did not affect our results. Nonetheless, making that adjustment to our model improved it, since we caught a problem in the prior (it gave burst sizes that were too small). In my experience, taking a principled approach to model building often uncovers issues in your model, even in simple ones like this one, that you were not aware of before performing checks.

Finally, let’s perform a posterior predictive check to make sure the model adequately captures our data.

n_ppc = samples.posterior_predictive.n_ppc.stack(
    {"sample": ("chain", "draw")}
).transpose("sample", "n_ppc_dim_0")

bokeh.io.show(
    bebi103.viz.predictive_ecdf(
        n_ppc,
        data=np.array(data["n"]),
        x_axis_label="mRNA transcript count",
        diff='ecdf',
    )
)

The model completely captures the data set; excellent!

Conclusions

The simulation-based calibration procedure (and the associated sensitivity analysis) is effective at identifying problem areas in Bayesian modeling. After passing the checks in this procedure, you can have more confidence in your modeling and the inferences you draw.

bebi103.stan.clean_cmdstan()

Computing environment

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

Python implementation: CPython
Python version       : 3.11.7
IPython version      : 8.20.0

numpy     : 1.26.3
pandas    : 2.1.4
cmdstanpy : 1.2.0
arviz     : 0.17.0
bokeh     : 3.3.4
iqplot    : 0.3.5
bebi103   : 0.1.20
jupyterlab: 4.0.12

cmdstan   : 2.34.1