Simulation based calibration and related checks

Data set download

---

%load_ext autoreload
%autoreload 2
%load_ext blackcellmagic

import numpy as np
import pandas as pd

import cmdstanpy
import arviz as az

import bokeh_catplot

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

---

Important notes: You should update your bebi103 package prior to running this lesson using pip install --upgrade bebi103.

You should also set the value of the variable cores to be the number of cores you have on your machine. I will be using 16 cores in this notebook, which is probably more cores that your machines has.

cores = 16

---

In the previous lecture, we laid out the a principled pipeline for constructing and testing a generative model and associated inference procedures. Be sure you are familiar with that lecture before working through this lesson.

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 from last term here.

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('../data/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(bokeh_catplot.ecdf(n, x_axis_label='mRNA count'))

The generative model

In Lesson 4, 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} &\alpha \sim \text{LogNorm}(0,2), \\[1em] &b \sim \text{LogNorm}(2, 3), \\[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 alpha = lognormal_rng(0.0, 2.0);
  real b = lognormal_rng(2.0, 3.0);
  real beta_ = 1.0 / 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> alpha;
  real<lower=0> b;
}


transformed parameters {
  real beta_ = 1.0 / b;
}


model {
  // Priors
  alpha ~ lognormal(0.0, 2.0);
  b ~ lognormal(2.0, 3.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:stan to c++ (/home/ec2-user/bebi103_course/2020/b/content/lessons/lesson_09/prior_pred.hpp)
INFO:cmdstanpy:compiling c++
INFO:cmdstanpy:compiled model file: /home/ec2-user/bebi103_course/2020/b/content/lessons/lesson_09/prior_pred
INFO:cmdstanpy:stan to c++ (/home/ec2-user/bebi103_course/2020/b/content/lessons/lesson_09/model.hpp)
INFO:cmdstanpy:compiling c++
INFO:cmdstanpy:compiled model file: /home/ec2-user/bebi103_course/2020/b/content/lessons/lesson_09/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, sampling_iters=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 = bokeh_catplot.ecdf(
        n, marker_kwargs=dict(fill_alpha=0.2, line_alpha=0.2), p=p, x_axis_type="log"
    )

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

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"],
    parameters=["alpha", "b"],
    measured_data_dtypes=dict(n=int),
    cores=cores,
    N=1000,
    progress_bar=True,
)

100%|██████████| 1000/1000 [04:32<00:00,  3.67it/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.155109	1186	0.162454	0.013691	1.000000	0.536481	1.001381	1962.759105	0.490690	2034.311326	0.508578	0	0	0	0	4000	0	alpha
1	2.295170	883	2.478863	0.229915	0.999937	0.798960	1.003249	754.591060	0.188648	1105.485187	0.276371	0	0	0	0	4000	1	alpha
2	1.014250	422	1.124240	0.088548	0.999991	1.242141	1.003101	1037.403485	0.259351	1094.752458	0.273688	0	0	0	0	4000	2	alpha
3	0.012971	1942	0.028399	0.088738	0.999991	0.173859	1.003397	932.538163	0.233135	1057.835287	0.264459	0	0	0	0	4000	3	alpha
4	16.508400	3334	15.317094	1.226472	0.998217	-0.971327	1.004364	718.427544	0.179607	1028.065017	0.257016	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')]
).groupby(
    'parameter'
).opts(
    size=2
).overlay(
)

points

There are several data sets that led to the posterior actually being broader than the prior. Let’s display the plot again, this time zooming in on the region where most of the samples are.

points.opts(xlim=(0, 1.05))

There are some problems here. First, we get z-scores below -5, which suggests we are having problems capturing the ground truth. Second, we do have several data sets that show poor shrinkage, so we are not learning from the data.

Let’s think about the z-score problem first. We will pull out all data sets with z-scores with magnitude greater than 5 and generate data sets from them by drawing out of a Negative Binomial distribution with the ground truth parameters.

# Extract trials that had large magnitude z-scores
trials = df_sbc.loc[np.abs(df_sbc["z_score"]) > 5, "trial"]

# Generate and plot simulated data sets that were problematic
rg = np.random.default_rng()


def gen_and_plot_data(df_sbc, trials, logx=False):
    dfs = []
    for trial, g in df_sbc.loc[df_sbc["trial"].isin(trials), :].groupby("trial"):
        alpha = g.loc[g["parameter"] == "alpha", "ground_truth"].values[0]
        b = g.loc[g["parameter"] == "b", "ground_truth"].values[0]
        beta = 1 / b
        n = rg.negative_binomial(alpha, beta / (1 + beta), size=len(df))
        dfs.append(pd.DataFrame(dict(trial=trial, n=n)))

    df_check = pd.concat(dfs, ignore_index=True)

    if logx:
        x_axis_type = "log"
        x_axis_label = "n + 1"
        df_check["n"] += 1
    else:
        x_axis_type = "linear"
        x_axis_label = "n"

    return bokeh_catplot.ecdf(
        df_check,
        cats="trial",
        val="n",
        x_axis_type=x_axis_type,
        x_axis_label=x_axis_label,
    )


bokeh.io.show(gen_and_plot_data(df_sbc, trials))

In all of the cases with poor z-score, we got mostly zero counts. Let’s take a look at what the ground truth parameter values were.

df_sbc.loc[
    df_sbc["trial"].isin(trials), ["parameter", "trial", "ground_truth", "shrinkage"]
].sort_values(by=["trial", 'parameter'])

	parameter	trial	ground_truth	shrinkage
201	alpha	201	0.007923	0.995891
1201	b	201	2.302770	1.000000
323	alpha	323	0.099423	0.999996
1323	b	323	5.796090	1.000000
716	alpha	716	16.444700	0.999047
1716	b	716	0.000331	1.000000
986	alpha	986	0.011020	0.999999
1986	b	986	25.218200	1.000000

We have problems when either the burst size or burst frequency is tiny. Either case results in a large number of zero copy numbers. So, we have discovered that if we have very low mRNA counts in our data, we will possibly not be able to capture the ground truth. We will keep this warning in mind going forward.

Now, we should check to see for what ground truth parameter values we go poor shrinkage.

trials = df_sbc.loc[df_sbc["shrinkage"] < 0, "trial"]
df_sbc.loc[
    df_sbc["trial"].isin(trials), ["parameter", "trial", "ground_truth", "shrinkage"]
].sort_values(by=["trial", 'parameter'])

	parameter	trial	ground_truth	shrinkage
66	alpha	66	524.907000	-4.801162e+01
1066	b	66	0.278872	1.000000e+00
97	alpha	97	113.396000	-2.257682e-01
1097	b	97	0.532460	1.000000e+00
182	alpha	182	721.125000	-1.187942e+02
1182	b	182	0.435883	1.000000e+00
208	alpha	208	56.418900	-8.252752e+00
1208	b	208	0.092603	1.000000e+00
288	alpha	288	22.981300	-2.853881e+16
1288	b	288	0.529094	1.000000e+00
384	alpha	384	21.041300	-6.973032e-01
1384	b	384	0.483136	1.000000e+00
424	alpha	424	40.958600	-2.824073e+00
1424	b	424	0.097227	1.000000e+00
432	alpha	432	47.571300	-4.888162e+00
1432	b	432	0.159237	1.000000e+00
454	alpha	454	172.453000	-1.561828e+01
1454	b	454	0.315042	1.000000e+00
530	alpha	530	0.920082	9.999941e-01
1530	b	530	342510.000000	-1.437984e+00
575	alpha	575	105.199000	-1.833792e+01
1575	b	575	0.082234	1.000000e+00
654	alpha	654	912.355000	-5.194055e+00
1654	b	654	25.637000	1.000000e+00
668	alpha	668	4.344020	-2.119141e+00
1668	b	668	0.159192	1.000000e+00
680	alpha	680	26.192700	-1.656187e-02
1680	b	680	0.055961	1.000000e+00
704	alpha	704	22.489200	-1.938570e+00
1704	b	704	0.138149	1.000000e+00
797	alpha	797	6.160340	-2.932768e-01
1797	b	797	0.338923	1.000000e+00
903	alpha	903	199.062000	-7.437519e-01
1903	b	903	1.518970	1.000000e+00

Again, in many cases, we see that small burst size is problematic. But we also see that when we on occasion get enormous burst size, this is also a problem. Many of these burst sizes are unphysical, too. Making tens of thousands of transcripts in a single burst is unlikely.

To make sure the data sets giving poor shrinkage are not somehow absurd, let’s generate data sets by drawing out of the likelihood for the respective ground truths. We’ll then plot the ECDFs to see if they are somehow pathological.

bokeh.io.show(gen_and_plot_data(df_sbc, trials, logx=True))

The ECDFs do not seem particularly pathological.

From this analysis, we see that the z-score is poor (meaning that we are failing to capture the ground truth) when the number of transcripts is very small, which happens with the burst size is tiny. The shrinkage is poor when the burst size very small or very large and the burst frequency \(\alpha\) is high.

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}(2, 0.25). \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 several thousand transcripts per burst.

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

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}(2, 0.25), \\[1em] &b \sim \text{Gamma}(2, 0.05), \\[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(2.0, 0.25);
  real b = gamma_rng(2.0, 0.05);
  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, sampling_iters=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 = bokeh_catplot.ecdf(
        n, marker_kwargs=dict(fill_alpha=0.2, line_alpha=0.2), p=p, x_axis_type="log"
    )

bokeh.io.show(p)

INFO:cmdstanpy:stan to c++ (/home/ec2-user/bebi103_course/2020/b/content/lessons/lesson_09/prior_pred_2.hpp)
INFO:cmdstanpy:compiling c++
INFO:cmdstanpy:compiled model file: /home/ec2-user/bebi103_course/2020/b/content/lessons/lesson_09/prior_pred_2
INFO:cmdstanpy:start chain 1
INFO:cmdstanpy:finish chain 1

Most of the data sets have reasonable ECDFs. These data sets again seem to make our intuition. Let’s check the mean and variance of trancsript 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)
)

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:stan to c++ (/home/ec2-user/bebi103_course/2020/b/content/lessons/lesson_09/model_2.hpp)
INFO:cmdstanpy:compiling c++
INFO:cmdstanpy:compiled model file: /home/ec2-user/bebi103_course/2020/b/content/lessons/lesson_09/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"],
    parameters=["alpha", "b"],
    measured_data_dtypes=dict(n=int),
    posterior_predictive_var_names=["n_ppc"],
    cores=cores,
    N=1000,
    progress_bar=True,
)

100%|██████████| 1000/1000 [08:21<00:00,  1.99it/s]

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

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

warning_code
0    1812
1      12
2     174
3       2
dtype: int64

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"],
    parameters=["alpha", "b"],
    measured_data_dtypes=dict(n=int),
    posterior_predictive_var_names=['n_ppc'],
    sampling_kwargs=dict(warmup_iters=2000, sampling_iters=2000),
    cores=cores,
    N=1000,
    progress_bar=True,
)

100%|██████████| 1000/1000 [16:05<00:00,  1.04it/s]

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

warning_code
0    1998
2       2
dtype: int64

Our diagnostics are much better. We can check any problems quickly.

trials = df_sbc.loc[df_sbc['warning_code'] > 0, "trial"]
df_sbc.loc[
    df_sbc["trial"].isin(trials), ["parameter", "trial", "ground_truth", "shrinkage", "Rhat"]
].sort_values(by=["trial", 'parameter'])

	parameter	trial	ground_truth	shrinkage	Rhat
665	alpha	665	7.31457	0.989161	1.011710
1665	b	665	36.81310	0.985080	1.011994

The R-hats are still very close to 1.01, so we will proceed. 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(
)

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 a few samples for either parameter escape the 99% envelope. This is not unexpected for the 1000 samples we took, so the check of the rank statistics looks fine.

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 3
INFO:cmdstanpy:finish chain 4

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, pars=['alpha', 'b']))

This result looks very much like what we achieved in Lesson 4, 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 we 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=True,
    )
)

The model mostly captures the distribution, with just a few data points, about as many as would be expected, falling outside the 80th percentile of the posterior predictive samples.

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,bebi103,jupyterlab

CPython 3.7.6
IPython 7.12.0

numpy 1.18.1
pandas 0.24.2
cmdstanpy 0.8.0
arviz 0.6.1
bokeh 1.4.0
holoviews 1.12.7
bebi103 0.0.53
jupyterlab 1.2.6