23: Simulation based calibration and related checks in practice

Data set download

---

# Colab setup ------------------
import os, sys, subprocess
if "google.colab" in sys.modules:
    cmd = "pip install --upgrade colorcet bebi103 arviz cmdstanpy watermark"
    process = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate()
    import cmdstanpy; cmdstanpy.install_cmdstan()
    data_path = "https://s3.amazonaws.com/bebi103.caltech.edu/data/"
else:
    data_path = "../data/"
# ------------------------------

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

import cmdstanpy
import arviz as az

import iqplot

import bebi103

import holoviews as hv
hv.extension('bokeh')
bebi103.hv.set_defaults()

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

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 3 cores in this notebook, since I have four cores on my desktop computer, reserving one core other work.

cores = 3

---

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 {
  int n[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;
  int n[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.

sm_prior_pred = cmdstanpy.CmdStanModel(stan_file='prior_pred.stan')
sm = cmdstanpy.CmdStanModel(stan_file='model.stan')

INFO:cmdstanpy:found newer exe file, not recompiling
INFO:cmdstanpy:compiled model file: /Users/bois/Dropbox/git/bebi103_course/2021/b/content/lessons/23/prior_pred
INFO:cmdstanpy:found newer exe file, not recompiling
INFO:cmdstanpy:compiled model file: /Users/bois/Dropbox/git/bebi103_course/2021/b/content/lessons/23/model

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.

samples_prior_pred = sm_prior_pred.sample(
    data=data, fixed_param=True, 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)

INFO:cmdstanpy:start chain 1
INFO:cmdstanpy:finish chain 1

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

hv.Points(
    data=(means, variances),
    kdims=['mean of counts', 'variance of counts'],
).opts(
    logx=True,
    logy=True,
    size=2,
    xlim=(1, None),
    ylim=(1, None)
)

Data type cannot be displayed:

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 [12:59<00:00,  1.28it/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.925864	0	1.042250	0.038034	0.999999	3.060028	1.003168	1121.500225	0.280375	1329.352480	0.332338	9	0	0	4	4000	0	alpha
1	0.126949	0	1.001268	0.001303	1.000000	670.978560	1.000649	1709.691520	0.427423	1174.324897	0.293581	0	0	0	0	4000	1	alpha
2	1.676310	4000	1.208541	0.094007	0.999996	-4.975887	1.002458	983.970311	0.245993	667.008918	0.166752	0	0	0	0	4000	2	alpha
3	0.075720	0	1.001272	0.001344	1.000000	688.661030	1.000850	1408.754982	0.352189	1018.430103	0.254608	0	0	0	0	4000	3	alpha
4	0.111288	0	1.001286	0.001340	1.000000	664.134392	1.001206	1631.910772	0.407978	1418.852365	0.354713	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!

points = hv.Points(
    data=df_sbc,
    kdims=['shrinkage', ('z_score', 'z-score')],
    vdims=['parameter', 'trial']
).groupby(
    'parameter'
).opts(
    tools=['hover'],
    size=2
).overlay(
)

points

Data type cannot be displayed:

Oof! We are severely overfitting the model, as evidenced by very 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',
    jitter=True
)

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. It is therefore also not surprising that we also miss the ground truth when \(b\) is small, again giving us many zero measurements.

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)
hv.Curve(
    (alpha, g),
    kdims='α',
    vdims='g(α)'
)

Data type cannot be displayed:

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)
hv.Curve(
    (b, g),
    kdims='b',
    vdims='g(b)'
)

Data type cannot be displayed:

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 {
  int n[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.

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, 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)

INFO:cmdstanpy:found newer exe file, not recompiling
INFO:cmdstanpy:compiled model file: /Users/bois/Dropbox/git/bebi103_course/2021/b/content/lessons/23/prior_pred_2
INFO:cmdstanpy:start chain 1
INFO:cmdstanpy:finish chain 1

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 30% or so of the counts. These data sets again seem to make 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

hv.Points(
    data=(means, variances),
    kdims=['mean of counts', 'variance of counts'],
).opts(
    logx=True,
    logy=True,
    size=2,
    xlim=(1, None),
    ylim=(1, None)
)

Data type cannot be displayed:

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;
  int n[N];
}


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


transformed parameters {
  real beta_ = 1.0 / b;
}


model {
  // Priors
  alpha ~ gamma(2.0, 0.25);
  b ~ gamma(2.0, 0.05);

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


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

Let’s compile!

sm_2 = cmdstanpy.CmdStanModel(stan_file='model_2.stan')

INFO:cmdstanpy:found newer exe file, not recompiling
INFO:cmdstanpy:compiled model file: /Users/bois/Dropbox/git/bebi103_course/2021/b/content/lessons/23/model_2

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 [20:29<00:00,  1.23s/it]

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    888.0
1      4.0
2     96.0
3     12.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 [38:34<00:00,  2.31s/it]

Let’s again check the diagnostics.

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

warning_code
0    997.0
2      3.0
dtype: float64

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

hv.Points(
    data=df_sbc,
    kdims=['shrinkage', ('z_score', 'z-score')]
).groupby(
    'parameter'
).opts(
    size=2
).overlay(
)

Data type cannot be displayed:

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 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.

samples = sm_2.sample(data=data)
samples = az.from_cmdstanpy(posterior=samples, posterior_predictive='n_ppc')

bebi103.stan.check_all_diagnostics(samples)

INFO:cmdstanpy:start chain 1
INFO:cmdstanpy:start chain 2
INFO:cmdstanpy:start chain 3
INFO:cmdstanpy:start chain 4
INFO:cmdstanpy:finish chain 2
INFO:cmdstanpy:finish chain 1
INFO:cmdstanpy:finish chain 4
INFO:cmdstanpy:finish chain 3

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 9, 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,holoviews,iqplot,bebi103,jupyterlab
print("cmdstan   :", bebi103.stan.cmdstan_version())

Python implementation: CPython
Python version       : 3.8.5
IPython version      : 7.19.0

numpy     : 1.19.2
pandas    : 1.2.1
cmdstanpy : 0.9.67
arviz     : 0.11.0
bokeh     : 2.2.3
holoviews : 1.14.0
iqplot    : 0.2.0
bebi103   : 0.1.2
jupyterlab: 2.2.6

cmdstan   : 2.26.0