Hamiltonian Monte Carlo utilities

[1]:

import re

import numpy as np
import polars as pl

import cmdstanpy
import arviz as az

import bebi103

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

bebi103.hv.set_defaults()

import bokeh.io
bokeh.io.output_notebook()

ⓘ

Loading BokehJS ...

The bebi103 package contains a set of convenience functions for use with HMC, primarily with Stan. Much of the functionality is included in ArviZ, but this package has customized visualization that its developer prefers. Furthermore, the HMC and expectand diagnostics in this package are ports of those from Michael Betancourt’s ``mcmc_diagnostics` package <https://github.com/betanalpha/mcmc_diagnostics/>`__, which can differ quantitatively from those in ArviZ.

Additionally, the bebi103 package has functions that may be included in functions block for Stan programs.

Model and sample data

To demonstrate the usage of the Stan utilities of bebi103, we will consider the following hierarchical generative model.

\begin{align} &\theta_i \sim \text{Norm}(5, 5),\;\;i\in\{1, 2\},\\[1em] &\tau_i \sim \text{HalfNorm}(0, 10),\;\;i\in\{1, 2\},\\[1em] &\sigma_i \sim \text{HalfNorm}(0, 10),\;\;i\in\{1, 2\},\\[1em] &\rho \sim \text{Uniform}(-1, 1),\\[1em] &\mathsf{T} = \begin{pmatrix}\tau_1^2 & 0 \\ 0 & \tau_2^2\end{pmatrix}, \\[1em] &\mathsf{\Sigma} = \begin{pmatrix}\sigma_1^2 & \rho\sigma_1\sigma_2 \\ \rho\sigma_1\sigma_2 & \sigma_2^2\end{pmatrix}, \\[1em] &\theta_{1, i} \sim \text{Norm}(\theta, \mathsf{T}) \;\forall\,i\in\{1, 2, 3\},\\[1em] & \begin{pmatrix} x_{i, j} \\ y_{i,j} \end{pmatrix} \sim \text{Norm}(\theta_{1,i}, \mathsf{\Sigma})\;\forall\,i \in \{1, 2, 3\}, \; j \in \{1, \ldots, n_i \}. \end{align}

This is a hierarchical model in which the three vector-valued parameters \(\theta_{1,1}\), \(\theta_{1,2}\), and \(\theta_{1,3}\) are conditioned on the vector-valued hyperparameter \(\theta\). We will use a fabricated data set generated from this model with parameters described in the Using the user guide section. It is useful to understand the data set, so let’s look at the data frame.

[2]:

df = pl.read_csv('sample_data.csv')
df.head()

[2]:

shape: (5, 3)

x	y	trial
f64	f64	i64
1.762886	11.955616	1
4.364957	11.136633	1
3.457626	12.301267	1
-0.839319	10.401899	1
4.694602	11.925334	1

Each trial (there are three, numbered 1, 2, and 3) has an x measurment and a y measurement.

To begin our analysis, we will use the following Stan model, which directly follows from the above generative model.

data {
  // Total number of data points
  int N;

  // Number of entries in each level of the hierarchy
  int J_1;

  //Index array to keep track of hierarchical structure
  array[N] int index_1;

  // The measurements
  array[N] real x;
  array[N] real y;
}


transformed data {
  // Data are two-dimensional, so store in a vector
  array[N] vector[2] xy;

  for (i in 1:N) {
    xy[i, 1] = x[i];
    xy[i, 2] = y[i];
  }
}


parameters {
  // Hyperparameters level 0
  vector[2] theta;

  // How hyperparameters vary
  vector<lower=0>[2] tau;

  // Parameters
  array[J_1] vector[2] theta_1;
  vector<lower=0>[2] sigma;
  real<lower=-1, upper=1> rho;
}


transformed parameters {
  // Covariance matrix for hyperparameters
  matrix[2, 2] Tau = [
    [tau[1]^2,  0       ],
    [0,         tau[2]^2]
  ];

  // Covariance matrix for likelihood
  matrix[2, 2] Sigma = [
    [sigma[1]^2,                 rho * sigma[1] * sigma[2]],
    [rho * sigma[1] * sigma[2],  sigma[2]^2               ]
  ];
}


model {
  // Hyperpriors
  theta ~ normal(5, 5);
  tau ~ normal(0, 10);

  // Priors
  theta_1 ~ multi_normal(theta, Tau);
  sigma ~ normal(0, 10);
  rho ~ uniform(-1, 1);

  // Likelihood
  for (i in 1:N) {
    xy[i] ~ multi_normal(theta_1[index_1[i]], Sigma);
  }
}


generated quantities {
  // Posterior retrodictive checks
  array[N] real x_prc;
  array[N] real y_prc;

  // Log likelihood
  array[N] real log_lik;

  {
    vector[2] xy_prc;

    for (i in 1:N) {
      xy_prc = multi_normal_rng(theta_1[index_1[i]], Sigma);
      log_lik[i] = multi_normal_lpdf(xy_prc | theta_1[index_1[i]], Sigma);
      x_prc[i] = xy_prc[1];
      y_prc[i] = xy_prc[2];
    }
  }
}

A few important notes about the Stan file.

The array index_1 keeps track of the dependence of the data on each parameter, which are themselves conditioned on hyperparameters.
Posterior predictive checks and the log-likelihood function are calculated in the generated quantities block.

A first pass through a simple workflow

We will now proceed through a MCMC workflow, doing the following steps:

Compile the Stan model.
Convert the data frame into input data for the Stan model.
Check the sampling diagnostics.
Create visualizations of the samples.

Compiling the Stan model

Starting with the first step making a compiled Stan model that we can use for sampling, we will use CmdStanPy. CmdStanPy is preferred over PyStan because of the simplicity of its interface and shorter model compilation time. The bebi103 package does not support PyStan.

If we wish to suppress messages to the screen, we can use bebi103.hmc.disable_logging() with context management.

[3]:

with bebi103.hmc.disable_logging():
    sm = cmdstanpy.CmdStanModel(stan_file="sample_model.stan")

Preparing hierarchical data for Stan

Now that we have a compiled Stan model, we need to supply it data. Referring to the data block in the Stan code above, for this hierarchical model, we need to specify how many data points, what the values of x and y are, the number of different parameters (J_1), and a set of indices the give which of the three \(\theta_1\) values upon which each data point is conditioned. This information is already contained in the tidy data frame, df. The 'x' and 'y' columns contain the data, and the 'trial' column determines the level in the hierarchy. The value of J_1 is given by the number of unique entries in the 'trial' column, and the number of data points is given by the total number of rows in the data frame.

The bebi103.hmc.df_to_datadict_hier() function converts a tidy data frame into a dictionary of data that can be passed to a Stan program for a hierarchical model. Furthermore, since the 'trial' column could have strings or other entries in them, these need to be converted to integers for use in Stan. A column is added to Stan with a suffix _stan to give the values of that column that are used in the Stan program. This is for reference when analyzing the results.

[4]:

# df may be either Polars or Pandas data frame
data, df = bebi103.hmc.df_to_datadict_hier(
    df, level_cols="trial", data_cols=["x", "y"]
)

# Take a look at the updated data frame
df.head()

[4]:

shape: (5, 4)

x	y	trial	trial_stan
f64	f64	i64	i64
1.762886	11.955616	1	1
4.245374	12.896902	1	1
2.495358	6.278253	1	1
4.847436	8.901546	1	1
5.585134	11.772688	1	1

The data frame has been appropriately updated (the trials were numbers, so there is no difference). Let’s look at the data dictionary.

[5]:

data

[5]:

{'N': 63,
 'J_1': 3,
 'index_1': array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]),
 'x': array([ 1.76288583,  4.24537447,  2.49535766,  4.84743608,  5.58513412,
         2.46538959,  2.53372498,  3.03556003,  4.1428636 ,  3.37291658,
         1.80757816,  3.86468666,  2.0406895 ,  6.10018145,  3.70259613,
         4.6946018 , -0.83931887,  3.45762554,  4.36495728,  9.02739122,
        -0.90483667,  0.63091768,  2.7511249 ,  0.84468249, -1.40832682,
         0.97071434,  1.65312163,  4.88902665,  4.74467292,  2.00083667,
         1.5374102 ,  5.36558851,  6.03757711,  1.71903739,  1.56783325,
         0.25157748,  6.76167578, -0.67226419,  3.36289739,  2.17013057,
         4.03836253,  2.72486749,  3.65771395,  4.59028926,  1.40245352,
         4.34252071,  4.1146886 ,  0.79938752,  0.72094663,  3.19538556,
         0.58176677, -0.01108818, -0.10906546,  2.37188306,  4.88494459,
        -0.35811241, -0.85531719,  2.62235591,  0.71487182,  1.38188117,
         1.82837862,  0.16000891,  4.81184826]),
 'y': array([11.95561574, 12.89690196,  6.27825302,  8.90154585, 11.77268772,
         8.9978501 ,  6.16475102, 10.64184364, 10.74399897,  8.71908946,
         9.63231343, 13.00236894,  9.47857277, 10.97263645, 11.46520631,
        11.92533427, 10.40189925, 12.30126673, 11.13663316, 16.78665286,
         4.11289192,  5.32907863,  7.82030839,  3.75238271,  3.89596763,
         1.37792381,  6.55489775,  9.7609403 , 10.03251758,  6.26665374,
         5.62577914,  5.84795532,  8.14624991,  9.12145582,  5.94521903,
         4.34204882, 12.30037908, -0.52873034,  4.3659075 ,  9.05879288,
         8.84628724,  9.95807796,  5.07513802, 10.37096513,  3.73211647,
         5.8037337 ,  5.02628689,  3.61882617,  6.98712067, 11.21360732,
         4.31478889,  4.40895773,  4.73065621,  8.95402277,  8.31905359,
         7.80920144,  5.92093593,  4.84087256,  3.43491574,  5.78404006,
         8.04955064,  5.29564007,  8.64743679])}

The dictionary has all of the entries necessary to pass into the Stan program while sampling.

Sampling

We can now go about our sampling and get the results. We will again suppress logging. Note that this is in general not a good idea. I do it here to manage space in the documentation.

[6]:

with bebi103.hmc.disable_logging():
    samples = sm.sample(
        data=data, chains=4, iter_sampling=1000, adapt_delta=0.95, seed=3252
    )

Converting samples to various formats

The variable samples is a cmdstanpy.CmdStanMCMC instance. The bebi103.hmc.convert_samples() function converts from this format to several others.

'dict': Dictionary of expectands and HMC diagnostics. Each item is a 2D Numpy array. The first dimension indexes the Markov chains and the second indexes the sequential states within each chain.
'arviz': An xarray.DataTree instance as used by ArviZ. Only valid if samples is a cmdstanpy.CmdStanMCMC instance.
'pandas': A Pandas data frame where each row of the data frame is a chain/draw with expectand values and HMC diagnostic results.
'polars': A Polars data frame where each row of the data frame is a chain/draw with expectand values and HMC diagnostic results.

Conversion is often unnecessary, as visualization, diagnostics, and analysis functions typically do that internally. Nonetheless, for ease of visualizing the samples, we convert the samples to a Polars data frame here.

[7]:

df_samples = bebi103.hmc.convert_samples(samples, 'polars')

# Take a look
df_samples.head()

[7]:

shape: (5, 219)

chain__	draw__	lp__	accept_stat__	stepsize__	treedepth__	n_leapfrog__	divergent__	energy__	theta[1]	theta[2]	tau[1]	tau[2]	theta_1[1,1]	theta_1[2,1]	theta_1[3,1]	theta_1[1,2]	theta_1[2,2]	theta_1[3,2]	sigma[1]	sigma[2]	rho	Tau[1,1]	Tau[2,1]	Tau[1,2]	Tau[2,2]	Sigma[1,1]	Sigma[2,1]	Sigma[1,2]	Sigma[2,2]	x_prc[1]	x_prc[2]	x_prc[3]	x_prc[4]	x_prc[5]	x_prc[6]	x_prc[7]	…	log_lik[27]	log_lik[28]	log_lik[29]	log_lik[30]	log_lik[31]	log_lik[32]	log_lik[33]	log_lik[34]	log_lik[35]	log_lik[36]	log_lik[37]	log_lik[38]	log_lik[39]	log_lik[40]	log_lik[41]	log_lik[42]	log_lik[43]	log_lik[44]	log_lik[45]	log_lik[46]	log_lik[47]	log_lik[48]	log_lik[49]	log_lik[50]	log_lik[51]	log_lik[52]	log_lik[53]	log_lik[54]	log_lik[55]	log_lik[56]	log_lik[57]	log_lik[58]	log_lik[59]	log_lik[60]	log_lik[61]	log_lik[62]	log_lik[63]
i64	i64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	…	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64
1	1	-157.36242	0.996946	0.216174	4.0	15.0	0.0	160.96684	3.3400944	9.3800556	0.799549	4.8560164	3.6098261	3.0214278	1.7652373	10.70204	6.9489694	5.8949729	2.051773	2.5774174	0.657098	0.639279	0.0	0.0	23.580895	4.2097725	3.4749157	3.4749157	6.6430805	5.5698704	5.1973786	7.220437	0.857876	4.3333721	5.7248875	1.7558562	…	-4.490788	-3.2506	-3.229742	-3.630459	-5.039638	-3.645568	-3.486727	-3.884245	-3.488522	-5.167867	-4.978146	-4.463282	-3.337886	-3.44238	-4.543225	-4.83689	-3.309076	-3.281724	-3.703284	-3.317581	-4.541385	-3.456033	-3.916834	-5.78315	-3.780804	-3.232168	-4.719223	-3.935118	-3.94908	-5.473484	-4.662725	-4.166747	-3.581278	-3.475598	-4.907618	-3.845557	-4.864691
1	2	-161.0091	0.962797	0.216174	4.0	15.0	0.0	165.22433	2.8179074	6.5020438	1.3747677	8.0005074	2.7597237	1.8844103	1.6541406	9.1836204	5.9835658	5.7517063	2.20051	2.6583775	0.669681	1.8899861	0.0	0.0	64.008118	4.8422441	3.9174915	3.9174915	7.0669707	5.1286782	4.2231406	3.2370389	3.8320794	1.9770098	5.8897723	0.797052	…	-6.16563	-4.307867	-6.243407	-3.668701	-3.312097	-3.805297	-5.720665	-4.197534	-3.427182	-5.124652	-5.945818	-3.795556	-4.564292	-3.425371	-5.15271	-5.656274	-3.748868	-3.431989	-4.28331	-3.56417	-3.396519	-3.416867	-4.828997	-4.999331	-3.398694	-4.096841	-9.312179	-4.426276	-4.251297	-4.072744	-5.300062	-3.693385	-4.255173	-3.468082	-3.445644	-4.263431	-3.408547
1	3	-160.63286	0.94051	0.216174	4.0	15.0	0.0	169.56701	2.8103211	8.1555405	0.534681	0.940767	3.4422601	3.0288367	2.5018904	10.90685	7.1918296	7.68738	1.8942131	2.4381071	0.498989	0.2858833	0.0	0.0	0.885043	3.5880431	2.3044796	2.3044796	5.9443665	5.0211365	3.923145	4.2023807	2.3357392	0.866446	4.765728	2.2710564	…	-3.319997	-4.647367	-3.972264	-3.267477	-5.15114	-5.981165	-8.731551	-4.074365	-8.38317	-4.903477	-3.807537	-3.878422	-3.615925	-3.255246	-5.639008	-3.26244	-3.358559	-3.56321	-3.898647	-3.240369	-4.365215	-4.901795	-4.120195	-4.393892	-4.511024	-3.534596	-5.071555	-3.36541	-4.184606	-4.448638	-3.327942	-6.98695	-3.282434	-3.821043	-3.929123	-3.643508	-4.136507
1	4	-159.6237	1.0	0.216174	4.0	15.0	0.0	164.11557	2.5845624	7.9602053	0.7508871	10.797538	3.4237137	2.4164793	1.3936102	10.402738	6.0391642	6.216457	2.244996	2.7509113	0.731762	0.563831	0.0	0.0	116.58682	5.0400069	4.5192065	4.5192065	7.5675127	0.658077	2.9620842	2.9080832	2.5016567	5.1648933	3.8338198	1.3743562	…	-3.636001	-4.319272	-3.311477	-3.534741	-3.963338	-3.577486	-4.743639	-5.501458	-4.221635	-5.429224	-3.394738	-5.035972	-4.475535	-3.277713	-4.479338	-3.77753	-4.559502	-3.319845	-3.542236	-6.885093	-4.259644	-3.510838	-3.558921	-3.528215	-3.480809	-4.772193	-4.25625	-4.761282	-3.982735	-3.572086	-3.681185	-5.139564	-3.371119	-3.614484	-4.614607	-3.470661	-6.557853
1	5	-162.14088	0.956536	0.216174	4.0	31.0	0.0	166.96483	2.5329305	2.1067779	1.9855133	4.3335428	3.0175406	2.4525505	1.9305323	10.027476	6.8589144	5.9362293	1.8618187	2.5786185	0.405575	3.9422629	0.0	0.0	18.779593	3.4663687	1.9471308	1.9471308	6.6492731	3.3602981	4.4163121	4.66739	2.5490173	5.3635836	3.0283696	2.575362	…	-4.32196	-3.488587	-5.542598	-4.370237	-6.015344	-4.239383	-3.816495	-4.096576	-3.357709	-4.403343	-3.553117	-4.498793	-3.610253	-3.452824	-4.823412	-4.981389	-4.824607	-4.038119	-6.253942	-3.509316	-3.400077	-3.730457	-3.540407	-3.659544	-4.367968	-3.49186	-3.540384	-3.918586	-3.363623	-6.034796	-4.809338	-5.737223	-4.761152	-4.265911	-3.754131	-3.360941	-3.707925

Variables/parameters naming convention

Note how multdimensional variables are named in the above data frame. They are one-indexed, as they are in Stan, and use brackets with different axes of the index separated by a comma with no spaces. Note also that the chain__ and draw__ columns index (again, with one-indexing) the MCMC chain and the index of the sample, respectively. The other columns that end with a double-underscore are diagnostic results outputted by Stan. This nomenclature is used throughout the bebi103 MCMC parsing functions.

Some functions in the bebi103.hmc and bebi103.viz modules use parameters as an argument. In those cases, the entries of multidimensional properties are expected to be entered one-by-one. The bebi103.hms.inferential_calibration() function uses var_names. For multidimensional parameters, only the base name (with no indexing) is expected. This var_names format is congruent with that of ArviZ. As an example, consider a scalar parameter a, a vector parameter b with three entries, and a 2×2 matrix parameter c. The following shows how the respective keyword arguments should be used.

parameters=['a', 'b[0]', 'b[1]', 'b[2]', 'c[0,0]', 'c[0,1]', 'c[1,0]', 'c[1,1]']
var_names=['a', 'b', 'c']

Checking diagnostics

We may check the following diagnostic features, the first four of which are specific to HMC, and the remaining are checks of expectand samples and apply to any MCMC sampling algorithm.

Divergences.
Tree depth.
Energy fraction of missing information (E-FMI).
Average acceptance rate proxy.
The empirical generalized Pareto shape parameter, xi-hat.
Chain variance (if zero, chain is stuck).
The split Gelman-Rubin statistic, Rhat (\(\hat{R}\)).
Incremental empirical integrated autocorrelation time, tau-hat.
The effective sample size (ESS).

The functions to check these diagnostics are ported from Michael Betancourt’s ``mcmc_diagnostics` package <https://github.com/betanalpha/mcmc_diagnostics/>`__.

We can start by checking the HMC diagnostics using the bebi103.hmc.check_hmc_diagnostics().

[8]:

bebi103.hmc.check_hmc_diagnostics(samples)

There were a total of 96 divergences.

Divergent Hamiltonian transitions result from unstable numerical
trajectories.  These instabilities are often due to degenerate target
geometry, especially "pinches".  If there are only a small number of
divergences then running with adept_delta larger than 0.801 may reduce
the instabilities at the cost of more expensive Hamiltonian transitions.

[8]:

{'divergences': array([37., 16., 29., 14.]),
 'divergences_success': np.False_,
 'average_accept_proxy': array([0.89070892, 0.95977102, 0.91658264, 0.95293607]),
 'average_accept_proxy_success': np.True_,
 'efmi': array([3.95195624, 3.55787171, 3.55201796, 3.81144591]),
 'efmi_success': np.True_,
 'treedepth': array([0, 0, 0, 0]),
 'treedepth_success': np.True_}

The function returned a dictionary with reports on the four diagnostics. Keys in the dictionary suffixed with _success give a Boolean that is true if the diagnostic check with successfull (no issues), and false otherwise. It also prints an error message describing any issues what that problems they may be symtpoms of.

We can also check the expectand diagnostics using bebi103.hmc.check_expectand_diagnostics(). Here, the omit keyword argument is useful to omit variables are are not interested in checking (things like posterior retrodictive checks and log likelihood values). We can also omit specific entries in arrays using the omit_array_entry kwarg. In this case, the off-diagonal terms in the matrix Tau are both always zero, so we want to omit them specifically.

[9]:

bebi103.hmc.check_expectand_diagnostics(
    samples,
    omit=('log_lik', '*_prc'),
    omit_array_entry=('Tau[1,2]', 'Tau[2,1]')
)

Parameter Tau[1,1] triggered a tail xi_hat warning in 4 of 4 chains.

Parameter Tau[2,2] triggered a tail xi_hat warning in 2 of 4 chains.

The expectands Tau[1,1], Tau[2,2] triggered tail xi_hat warnings. Large
tail xi_hats suggest that the expectand might not be sufficiently
integrable.

[9]:

{'xi_hat': {'theta[1]': array([[-0.06406819,  0.02575865],
         [-0.04141672, -0.04163014],
         [-0.04206061, -0.01458971],
         [-0.14504202, -0.03686235]]),
  'theta[2]': array([[-0.15676144, -0.17537877],
         [-0.17568799, -0.14660007],
         [-0.31254057, -0.12975208],
         [-0.17332633, -0.16164093]]),
  'tau[1]': array([[-1.08448144, -0.01804568],
         [-1.11955268, -0.04177346],
         [-1.16150183,  0.06176076],
         [-1.02984185,  0.02986211]]),
  'tau[2]': array([[-0.82596774, -0.1456114 ],
         [-0.80736914, -0.1195331 ],
         [-0.7150093 , -0.04665074],
         [-0.74280958, -0.17126554]]),
  'theta_1[1,1]': array([[-0.60717534, -0.3084544 ],
         [-0.39286467, -0.35486797],
         [-0.38839594, -0.45971009],
         [-0.40811596, -0.30726834]]),
  'theta_1[2,1]': array([[-0.41023326, -0.27759831],
         [-0.36111068, -0.38269738],
         [-0.54895149, -0.38793389],
         [-0.40442007, -0.31134057]]),
  'theta_1[3,1]': array([[-0.28318456, -0.43060379],
         [-0.31456768, -0.34375411],
         [-0.54930373, -0.47566234],
         [-0.4313902 , -0.41081455]]),
  'theta_1[1,2]': array([[-0.35064029, -0.42306747],
         [-0.32514512, -0.37308936],
         [-0.45954631, -0.44530911],
         [-0.43746364, -0.4228291 ]]),
  'theta_1[2,2]': array([[-0.35406251, -0.35619478],
         [-0.4088378 , -0.3380721 ],
         [-0.34903268, -0.30799486],
         [-0.29108407, -0.23994918]]),
  'theta_1[3,2]': array([[-0.36430767, -0.3481737 ],
         [-0.30379087, -0.35520482],
         [-0.4505651 , -0.37775902],
         [-0.29405291, -0.32829952]]),
  'sigma[1]': array([[-0.43347761, -0.30696585],
         [-0.40164285, -0.37790329],
         [-0.59150548, -0.30074291],
         [-0.45426694, -0.2732337 ]]),
  'sigma[2]': array([[-0.57019948, -0.28951455],
         [-0.41802794, -0.33270163],
         [-0.40344634, -0.24982322],
         [-0.44237774, -0.38719475]]),
  'rho': array([[-0.42012432, -0.47418907],
         [-0.35394344, -0.5699361 ],
         [-0.29761298, -0.41961411],
         [-0.30206558, -0.5073984 ]]),
  'Tau[1,1]': array([[-1.67689506,  0.50732468],
         [-1.79220542,  0.45378113],
         [-1.8095312 ,  0.58034932],
         [-1.68472627,  0.51054968]]),
  'Tau[2,2]': array([[-1.24223598,  0.24510736],
         [-1.22109721,  0.25741512],
         [-1.08999125,  0.29357303],
         [-1.14430308,  0.17524331]]),
  'Sigma[1,1]': array([[-0.48674844, -0.25424876],
         [-0.45673761, -0.32872146],
         [-0.64105847, -0.24116751],
         [-0.50615014, -0.20393671]]),
  'Sigma[2,1]': array([[-0.59641295, -0.3212759 ],
         [-0.48826989, -0.28959868],
         [-0.46156844, -0.29958642],
         [-0.45520004, -0.24706323]]),
  'Sigma[1,2]': array([[-0.59641295, -0.3212759 ],
         [-0.48826989, -0.28959868],
         [-0.46156844, -0.29958642],
         [-0.45520004, -0.24706323]]),
  'Sigma[2,2]': array([[-0.621347  , -0.23127941],
         [-0.47136766, -0.2811046 ],
         [-0.45332644, -0.19820795],
         [-0.49597539, -0.33307061]])},
 'variance': {'theta[1]': array([ True,  True,  True,  True]),
  'theta[2]': array([ True,  True,  True,  True]),
  'tau[1]': array([ True,  True,  True,  True]),
  'tau[2]': array([ True,  True,  True,  True]),
  'theta_1[1,1]': array([ True,  True,  True,  True]),
  'theta_1[2,1]': array([ True,  True,  True,  True]),
  'theta_1[3,1]': array([ True,  True,  True,  True]),
  'theta_1[1,2]': array([ True,  True,  True,  True]),
  'theta_1[2,2]': array([ True,  True,  True,  True]),
  'theta_1[3,2]': array([ True,  True,  True,  True]),
  'sigma[1]': array([ True,  True,  True,  True]),
  'sigma[2]': array([ True,  True,  True,  True]),
  'rho': array([ True,  True,  True,  True]),
  'Tau[1,1]': array([ True,  True,  True,  True]),
  'Tau[2,2]': array([ True,  True,  True,  True]),
  'Sigma[1,1]': array([ True,  True,  True,  True]),
  'Sigma[2,1]': array([ True,  True,  True,  True]),
  'Sigma[1,2]': array([ True,  True,  True,  True]),
  'Sigma[2,2]': array([ True,  True,  True,  True])},
 'rhat': {'theta[1]': np.float64(1.005735347552244),
  'theta[2]': np.float64(1.0009711268625767),
  'tau[1]': np.float64(1.004183225365101),
  'tau[2]': np.float64(1.001752365371224),
  'theta_1[1,1]': np.float64(1.0033418032900017),
  'theta_1[2,1]': np.float64(1.001902947186569),
  'theta_1[3,1]': np.float64(1.0010474373326692),
  'theta_1[1,2]': np.float64(1.0033486995287977),
  'theta_1[2,2]': np.float64(1.003655595165959),
  'theta_1[3,2]': np.float64(1.000561948409255),
  'sigma[1]': np.float64(1.0018146535959138),
  'sigma[2]': np.float64(1.0025658133160193),
  'rho': np.float64(1.0006876409513898),
  'Tau[1,1]': np.float64(1.0043090478153083),
  'Tau[2,2]': np.float64(1.0017849185222076),
  'Sigma[1,1]': np.float64(1.0019913550578095),
  'Sigma[2,1]': np.float64(1.0020377777005973),
  'Sigma[1,2]': np.float64(1.0020377777005973),
  'Sigma[2,2]': np.float64(1.0026639800236024)},
 'inc_tau_hat': {'theta[1]': array([4.30666881, 2.00042797, 3.31952974, 2.20558727]),
  'theta[2]': array([2.07141614, 2.0236037 , 1.71538034, 1.70916537]),
  'tau[1]': array([3.29965222, 2.56772997, 3.13167848, 2.9714344 ]),
  'tau[2]': array([2.25714359, 2.20091899, 1.44958691, 2.40865864]),
  'theta_1[1,1]': array([3.717011  , 2.20873393, 2.69686387, 2.93115179]),
  'theta_1[2,1]': array([1.65673954, 1.79238647, 1.51444166, 1.45038624]),
  'theta_1[3,1]': array([2.01818712, 1.77895265, 2.15865078, 1.72818407]),
  'theta_1[1,2]': array([1.97806512, 1.8516646 , 1.74918713, 1.5852617 ]),
  'theta_1[2,2]': array([2.00217512, 1.42955673, 1.54426858, 1.55814193]),
  'theta_1[3,2]': array([2.31422079, 1.59015171, 1.27462379, 1.76544694]),
  'sigma[1]': array([1.89571822, 1.37151761, 1.57690572, 2.67122296]),
  'sigma[2]': array([2.07239877, 1.34627942, 1.47538985, 1.58913442]),
  'rho': array([2.27981051, 1.2499321 , 1.72648151, 1.75757444]),
  'Tau[1,1]': array([2.26722911, 2.08539472, 2.38249504, 1.6074432 ]),
  'Tau[2,2]': array([1.9953888 , 2.17328045, 1.29061879, 2.20770191]),
  'Sigma[1,1]': array([1.9111683 , 1.40767124, 1.61168233, 3.04731922]),
  'Sigma[2,1]': array([3.29005206, 1.64675198, 2.24420552, 3.52453044]),
  'Sigma[1,2]': array([3.29005206, 1.64675198, 2.24420552, 3.52453044]),
  'Sigma[2,2]': array([2.12936557, 1.37730611, 1.49880584, 1.66440716])},
 'ess_hat': {'theta[1]': array([232.19802683, 499.89303161, 301.24748926, 453.39398424]),
  'theta[2]': array([482.76151767, 494.16790458, 582.96109302, 585.08089301]),
  'tau[1]': array([303.06224169, 389.44905153, 319.31758158, 336.5378015 ]),
  'tau[2]': array([443.0378307 , 454.35565932, 689.8517032 , 415.16883355]),
  'theta_1[1,1]': array([269.03337113, 452.7480593 , 370.80106702, 341.16281596]),
  'theta_1[2,1]': array([603.59518051, 557.91539147, 660.30935715, 689.47151394]),
  'theta_1[3,1]': array([495.49419359, 562.12850775, 463.25232922, 578.64206439]),
  'theta_1[1,2]': array([505.54452932, 540.05460883, 571.6941204 , 630.81067413]),
  'theta_1[2,2]': array([499.45681017, 699.51753568, 647.55574907, 641.7900597 ]),
  'theta_1[3,2]': array([432.11088786, 628.8708122 , 784.54521878, 566.42880624]),
  'sigma[1]': array([527.50455836, 729.11933095, 634.1533231 , 374.36036357]),
  'sigma[2]': array([482.53261601, 742.78785133, 677.78696107, 629.27338804]),
  'rho': array([438.63294598, 800.04346093, 579.21269127, 568.96594401]),
  'Tau[1,1]': array([441.06702632, 479.52552527, 419.72805077, 622.10596174]),
  'Tau[2,2]': array([501.15546432, 460.1338953 , 774.82213041, 452.95970283]),
  'Sigma[1,1]': array([523.24015679, 710.39314701, 620.46966686, 328.15728466]),
  'Sigma[2,1]': array([303.9465584 , 607.25598637, 445.59198809, 283.72573823]),
  'Sigma[1,2]': array([303.9465584 , 607.25598637, 445.59198809, 283.72573823]),
  'Sigma[2,2]': array([469.62344784, 726.05501125, 667.19782605, 600.81452611])},
 'xi_hat_success': False,
 'variance_success': True,
 'rhat_success': True,
 'inc_tau_hat_success': True,
 'ess_hat_success': True}

Again, the function returns a dictionary with the results for each diagnostic calculation and Booleans for success.

We can check all diagnostics at once using the bebi103.hmc.check_all_diagnostics() function.

[10]:

bebi103.hmc.check_all_diagnostics(
    samples,
    omit=('log_lik', '*_prc'),
    omit_array_entry=('Tau[1,2]', 'Tau[2,1]'),
)

There were a total of 96 divergences.

Divergent Hamiltonian transitions result from unstable numerical
trajectories.  These instabilities are often due to degenerate target
geometry, especially "pinches".  If there are only a small number of
divergences then running with adept_delta larger than 0.801 may reduce
the instabilities at the cost of more expensive Hamiltonian transitions.

Parameter Tau[1,1] triggered a tail xi_hat warning in 4 of 4 chains.

Parameter Tau[2,2] triggered a tail xi_hat warning in 2 of 4 chains.

The expectands Tau[1,1], Tau[2,2] triggered tail xi_hat warnings. Large
tail xi_hats suggest that the expectand might not be sufficiently
integrable.

[10]:

The return value of this function is a number that, when converted to binary, each digit in the code stands for whether or not a test passed. A digit of zero indicates the test passed. The ordering of the tests goes:

bit 0: divergence
bit 1: treedepth
bit 2: E-FMI
bit 3: average acceptance proxy
bit 4: tail xi_hat
bit 5: zero variance
bit 6: split Rhat
bit 7: incremental tau_hat
bit 8: minimum ess_hat

For example, a warning code of 17 has a binary representation of 000010001, which means that xi-hat and divergence diagnostics failed. We can decode the warning code.

[11]:

bebi103.hmc.decode_warning_code(17)

divergence warning
xi_hat warning

Visualizing results

The bebi103 package allows for three visualizations of MCMC results.

Trace plots: bebi103.viz.trace()
Parallel coordinate plots: bebi103.viz.parcoord()
Corner plot: bebi103.viz.corner()

Each of these takes as their first six arguments:

samples: The MCMC samples, as a cmdstanpy.CmdStanMCMC instance, dictionary, data frame, etc. (Note that corner() also accepts other data type for samples, such as in display of bootstrap replicates.)
parameters: List of the names of variables to include in the plot.
palette: List of colors to use.
omit: List of variable names or regular expression patterns for variables to omit in the plot.

samples is the only required argument.

Trace plots

We will begin with a trace plot of the hyperparameters θ and τ.

[12]:

bokeh.io.show(
    bebi103.viz.trace(samples, parameters=["theta[1]", "theta[2]", "tau[1]", "tau[2]"])
)

The traces are color-coded by chain.

Parallel coordinate plots

For a parallel coordinate plot, we want to omit all four entries of the matrices Tau and Sigma because the entries are either zero or determined uniquely by other parameters. We can specify regular expressions in the omit list to conveniently omit all of the entries in these matrix parameters.

[13]:

bokeh.io.show(
    bebi103.viz.parcoord(samples, omit=["Tau", re.compile("Sigma.*"), "*_prc", "log_lik"])
)

In the plot, samples that did not result in a divergence are semi-transparent gray by default. Samples that had a divergence are colored in orange and shown with thicker lines. Because the scale of respective parameters may be different, it is often easer to rescale each parameter for the plot. We can do this with the transformation kwarg, and setting it to 'minmax' sets the minimum of each parameter value to zero and maximum to one.

[14]:

bokeh.io.show(
    bebi103.viz.parcoord(
        samples,
        omit=["Tau", re.compile("Sigma.*"), "*_prc", "log_lik"],
        transformation="minmax",
    )
)

This plot is more revealing. We see (especially if we zoom in) that most of the divergences are going through small tau[1]. This is indicative of the “funnel” behavior often encountered in hierarchical models.

Corner plots

Corner plots display scatter plots of samples from all pairs of variables, as well as histograms or ECDFs for individual variables. We can make a corner plot for the hyperparameters and also for ρ.

[15]:

bokeh.io.show(
    bebi103.viz.corner(
        samples, parameters=["theta[1]", "theta[2]", "rho", "tau[1]", "tau[2]"]
    )
)

These are very useful plots for interpreting a posterior, and also for diagnosing problems with sampling. The divergent samples appear in orange. The clustering of divergences for small tau[1] is clear.

The diagonal can also be plotted with ECDFs.

[16]:

bokeh.io.show(
    bebi103.viz.corner(
        samples,
        parameters=["theta[1]", "theta[2]", "rho", "tau[1]", "tau[2]"],
        plot_ecdf=True,
    )
)

Inferential calibration

In a principled Bayesian workflow, inferential calibration is a useful tool to identify pathologies in a generative model and in the sampler’s ability to draw samples from it.

Briefly, the procedure is:

Draw a parameter set θ out of the prior.
Use θ to draw a data set y out of the likelihood.
Perform MCMC sampling of the posterior using y as if it were the actual measured data set. Draw L MCMC samples of the parameters.
Do steps 1-3 N times (hundreds to a thousand is usually a good number).

For each of these calculations, you should run diagnostics to make sure the effective sample size, Rhat, etc., are within acceptable ranges. You can also compute z-scores (a measure of how well the MCMC samples reveal the ground truth), shrinkage (how informative the data can be), and rank statistics (used to check sampler performance). Please see the Talts, et al. paper for details on these quantities and their interpretation. Our focus here is how to use the bebi103.hmc.() function to perform the analysis and compute these quantities.

Example model

We have already seen in the example so far that we have divergences and the sampler does not seem to be able to sample small values of tau[1]. We will therefore write a new Stan program where we have noncentered parameters. The Stan model is:

data {
  // Total number of data points
  int N;

  // Number of entries in each level of the hierarchy
  int J_1;

  //Index array to keep track of hierarchical structure
  array[N] int index_1;

  // The measurements
  array[N] real x;
  array[N] real y;
}


transformed data {
  // Data are two-dimensional, so store in a vector
  array[N] vector[2] xy;
  for (i in 1:N) {
    xy[i, 1] = x[i];
    xy[i, 2] = y[i];
  }
}


parameters {
  // Hyperparameters level 0
  vector[2] theta;

  // How hyperparameters vary
  vector<lower=0>[2] tau;

  // Parameters
  vector<lower=0>[2] sigma;
  real<lower=-1, upper=1> rho;

  // Noncentered parameters
  array[J_1] vector[2] theta_1_noncentered;
}


transformed parameters {
  // Covariance matrix for likelihood
  matrix[2, 2] Sigma = [
    [sigma[1]^2,                 rho * sigma[1] * sigma[2]],
    [rho * sigma[1] * sigma[2],  sigma[2]^2               ]
  ];

  // Center parameters
  array[J_1] vector[2] theta_1;
  for (i in 1:J_1) {
    theta_1[i] = theta + tau .* theta_1_noncentered[i];
  }
}


model {
  // Hyperpriors
  theta ~ normal(5, 5);
  tau ~ normal(0, 10);

  // Priors
  theta_1_noncentered ~ multi_normal([0, 0], [[1, 0], [0, 1]]);
  sigma ~ normal(0, 10);
  rho ~ uniform(-1, 1);

  // Likelihood
  for (i in 1:N) {
    xy[i] ~ multi_normal(theta_1[index_1[i]], Sigma);
  }
}


generated quantities {
  array[N] real x_prc;
  array[N] real y_prc;
  array[N] real log_lik;

  {
    vector[2] xy_prc;

    for (i in 1:N) {
      xy_prc = multi_normal_rng(theta_1[index_1[i]], Sigma);
      log_lik[i] = multi_normal_lpdf(xy_prc | theta_1[index_1[i]], Sigma);
      x_prc[i] = xy_prc[1];
      y_prc[i] = xy_prc[2];
    }
  }
}

Let’s go ahead and compile the model and draw some samples just to see if it alleviated the problems with divergences. We will use extra warmup to make sure the chains are achieving stationarity.

[17]:

with bebi103.hmc.disable_logging():
    sm_nc = cmdstanpy.CmdStanModel(stan_file="sample_model_noncentered.stan")

    samples_nc = sm_nc.sample(
        data=data,
        chains=4,
        iter_warmup=2000,
        iter_sampling=1000,
        adapt_delta=0.95,
        seed=3252,
    )

samples_nc = az.from_cmdstanpy(
    samples_nc, posterior_predictive=["x_prc", "y_prc"], log_likelihood="log_lik"
)

And we can check a corner plot to check for divergences.

[18]:

bokeh.io.show(
    bebi103.viz.corner(
        samples_nc, parameters=["theta[1]", "theta[2]", "rho", "tau[1]", "tau[2]"]
    )
)

We still have a few divergences, which I will leave unrectified so that we can demonstrate diagnostics results in inferential calibration in what follows.

Building an inferential calibration calculation

To build a simulation-based calibration, we need to write a Stan program to generate data sets based on parameters drawn out of the prior, as prior predictive model.

data {
  // Total number of data points
  int N;

  // Number of entries in each level of the hierarchy
  int J_1;

  //Index array to keep track of hierarchical structure
  array[N] int index_1;
}


generated quantities {
  vector[2] theta;
  vector[2] tau;
  vector[2] sigma;

  for (i in 1:2) {
    theta[i] = normal_rng(5, 5);
    tau[i] = abs(normal_rng(0, 10));
    sigma[i] = abs(normal_rng(0, 10));
  }

  real rho = uniform_rng(-1, 1);

  array[J_1] vector[2] theta_1;
  for (i in 1:J_1) {
    for (j in 1:2) {
      theta_1[i, j] = normal_rng(theta[j], tau[j]);
    }
  }

  array[N] real x;
  array[N] real y;

  {
    matrix[2, 2] Sigma = [
      [sigma[1]^2,                 rho * sigma[1] * sigma[2]],
      [rho * sigma[1] * sigma[2],  sigma[2]^2               ]
    ];

    vector[2] xy;

    for (i in 1:N) {
      xy = multi_normal_rng(theta_1[index_1[i]], Sigma);
      x[i] = xy[1];
      y[i] = xy[2];
    }
  }
}

To use it in inferential calibration, we need to compile it.

[19]:

with bebi103.hmc.disable_logging():
    sm_prior_pred = cmdstanpy.CmdStanModel(
        stan_file="sample_model_prior_predictive.stan"
    )

Running inferential calibration

Finally, we are ready to perform inferential calibration using the bebi103.hmc.inferential_calibration() function. We need to tell the function

What Stan model to use for the prior predictive model.
Which Stan model to use for posterior sampling.
Which data dictionary to feed into the prior predictive model.
Which data dictionary to feed into the posterior sampling mode.
What the names of the measured data are. For multidimensional data, do not include indexing. In the present case, the measured data are ['x', 'y'].
The names of the posterior retrodictive variables. These are not used in shrinkage, z-score, rank statistic, and diagnostic calculations and are included in the omit and omit_array_entry kwargs.
The name of the log likelihood variable. This is ignored in the calculations and included in the omit kwarg.
The data types of the measured data.
Any kwargs to use in posterior sampling.
Any kwargs to pass to bebi103.hmc.check_all_diagnostics().
How many cores you want to use for the calculation.
How many calibration calculations you want to do. You should do at least 400; 1000 is a good number.

Inferential calibration calculations are lengthy, so to keep the runtime of this documentation manageable, we will load a pre-computed inferential calibration result if it is available. For reference, this calculation took about six hours on 16 cores on an AWS c5 instance.

[20]:

try:
    df_ic = pl.read_csv('inferential_calibration_results.csv', comment_prefix='#')
except:
    df_ic = bebi103.hmc.inferential_calibration(
        prior_predictive_stan_file='sample_model_prior_predictive.stan',
        posterior_stan_file='sample_model_noncentered.stan',
        prior_predictive_model_data=data,
        posterior_model_data=data,
        measured_data=['x', 'y'],
        measured_data_dtypes={'x': float, 'y': float},
        sampling_kwargs={'iter_warmup': 2000, 'adapt_delta': 0.95},
        diagnostic_check_kwargs=dict(
            omit=('log_lik', '*_prc'),
            omit_array_entry=('Tau[1,2]', 'Tau[2,1]'),
        ),
        output_csv='inferential_calibration_results.csv',
        cores=8,
        N=1000,
        progress_bar=True,
    )

The output of the inferential calibration calculation is a tidy data frame with the results for each of the 1000 trials.

[21]:

df_ic.sort(by=['trial', 'parameter'])

[21]:

shape: (13_000, 26)

trial	parameter	ground_truth	rank_statistic	mean	sd	shrinkage	z_score	Rhat	ESS	ESS_per_iter	tail_xi_hat_left	tail_xi_hat_right	inc_tau_hat	variance	n_divergences	n_bad_ebfmi	n_max_treedepth	average_accept_proxy	warning_code	prior_predictive_seed	posterior_seed	L	prior_predictive_prefix	samples_prefix	error
i64	str	f64	i64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	i64	i64	i64	f64	i64	i64	i64	i64	str	str	str
0	"rho"	-0.154069	604	-0.024102	0.124772	0.955467	1.041641	1.001715	2642.254985	0.660564	-0.289968	-0.380227	0.000399	0.015568	25	0	0	0.942442	257	297456205	1049904557	4000	null	null	"no error"
0	"sigma[1]"	0.451398	1231	0.477237	0.045503	0.999949	0.567839	1.001352	2813.647741	0.703412	-0.453938	-0.269356	0.000362	0.002071	25	0	0	0.942442	257	297456205	1049904557	4000	null	null	"no error"
0	"sigma[2]"	10.17805	668	11.17185	1.045518	0.967509	0.950534	1.001016	2966.440847	0.74161	-0.272726	-0.269964	0.000337	1.093108	25	0	0	0.942442	257	297456205	1049904557	4000	null	null	"no error"
0	"tau[1]"	0.529298	1691	0.928202	0.859715	0.979039	0.463995	1.001753	790.653705	0.197663	-0.792786	-0.11534	0.001399	0.73911	25	0	0	0.942442	257	297456205	1049904557	4000	null	null	"no error"
0	"tau[2]"	3.6318703	2515	3.506743	3.086189	0.732919	-0.040544	1.005851	1441.21152	0.360303	-1.258229	-0.134348	0.000848	9.52456	25	0	0	0.942442	257	297456205	1049904557	4000	null	null	"no error"
…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…	…
999	"theta_1[1,2]"	23.294449	1759	23.832732	3.504183	0.905783	0.153611	1.001902	2440.008467	0.610002	-0.407182	-0.28614	0.000457	12.279297	0	0	0	0.96791	0	1812196502	734172752	4000	null	null	"no error"
999	"theta_1[2,1]"	-4.330547	649	-0.319537	4.035135	0.870648	0.994021	1.000207	5350.124304	1.337531	-0.400866	-0.309748	0.000202	16.282313	0	0	0	0.96791	0	1812196502	734172752	4000	null	null	"no error"
999	"theta_1[2,2]"	20.703637	3452	17.608484	2.881193	0.935649	-1.074261	1.000118	5190.296919	1.297574	-0.296152	-0.320009	0.000198	8.301275	0	0	0	0.96791	0	1812196502	734172752	4000	null	null	"no error"
999	"theta_1[3,1]"	45.416861	3885	36.570224	4.76149	0.813384	-1.857956	1.000244	4691.256925	1.172814	-0.316922	-0.310625	0.000225	22.671782	0	0	0	0.96791	0	1812196502	734172752	4000	null	null	"no error"
999	"theta_1[3,2]"	12.344925	154	18.421358	3.364857	0.905566	1.805852	1.000597	4943.744802	1.235936	-0.390981	-0.286643	0.000212	11.322262	0	0	0	0.96791	0	1812196502	734172752	4000	null	null	"no error"

For each parameter for each trial, the ground truth is given (as drawn from the prior), as well as a rank statistic, mean, standard deviation, shrinkage, and z-score. MCMC diagnostics are also reported.

Interpreting results

We will not go into much detail about interpreting results here, directing you instead to the Talts, et al. paper. We will instead do a quick run through some of the results and suggested visualization procedures, highlighting some of the useful functions in bebi103 along the way.

We can first check what diagnostic warnings we encountered. There are thirteen parameters, so we can count the warning codes and divide by 13 to see how many MCMC calculations had problems.

[22]:

df_ic.group_by("warning_code").agg(pl.len() // 13).sort(by='warning_code')

[22]:

shape: (20, 2)

warning_code	len
i64	u32
0	556
1	245
2	68
3	35
17	4
…	…
338	1
345	1
347	1
354	3
355	1

Most of the samples either had no diagnostic warnings (warning code 0) or only divergences (warning code 1). To help interpret the rest of the warning codes, we can use the bebi103.hmc.decode_warning_code() function. For example, we can check warning code 17.

[23]:

bebi103.hmc.decode_warning_code(17)

divergence warning
xi_hat warning

This means that we had divergence and ξ-hat warningss. All together, the diagnostics tell us there are some problems with our sampling as we have it set up. Perhaps we should take more iterations, increase adapt_delta, and also increase max_treedepth in doing the calculations.

We should also check the z-score and shrinkage. We will plot the z-score vs. shrinkage for each parameter, coloring the glyphs by the warning code. This coloring enables us to see if any of the MCMC diagnostic issues might be affecting shrinkage and z-scores.

[24]:

points = hv.Points(
    data=df_ic.to_dict(),
    kdims=["shrinkage", ("z_score", "z-score")],
    vdims=[("parameter", "param"), "warning_code"],
).groupby(
    "parameter"
).opts(
    color="warning_code",
    frame_height=125,
    frame_width=125,
    size=2,
    tools=["hover"],
).layout(
)

# Show plot disabling logging to avoid Bokeh warnings
with bebi103.hmc.disable_logging():
    bokeh.io.show(hv.render(points))

The shrinkage and z-scores are good for all parameters (close to one and between –5 and 5, respectively), but the shrinkage is poor for the hyperparameters. This does not appear to be dependent on MCMC diagnostic issues. Rather, this tells us that with the size of data sets we are considering, we cannot learn much about the hyperparameters.

Finally, as suggested by Talts, et al., we can make an SBC rank ECDF plot using the bebi103.viz.sbc_rank_ecdf().

[25]:

p = bebi103.viz.sbc_rank_ecdf(
    df_ic, show_legend=True, marker_kwargs={"size": 2},
)
p.legend.spacing = 0
p.legend.label_text_font_size = '8pt'

bokeh.io.show(p)

The ECDF difference of the rank statistic falling within the theoretical 99-percentile envelope is indicative of good calibration of the sampler and lack of bias. The parameter sigma[1] has a low-biased rank statistic, suggesting that the samples of sigma[1] will tend to be biased toward larger values.

Running inferential calibration from the command line

Inferential calibration may also be run from the command line. In order to do that, the user must specify a .toml file with specifications for the inferential calibration run. Below are contents of an example file. To run inferential calibration, enter the following in the command line.

infcal path_to_specification.toml

Here are contents of an example specification file.

# ============================================================
# Stan files
# ============================================================
prior_predictive_stan_file = "sample_model_prior_predictive.stan"
posterior_stan_file = "sample_model_noncentered.stan"


# ============================================================
# Data contents for prior predictive
# ============================================================
prior_predictive_model_data.N = 63
prior_predictive_model_data.J_1 = 3
prior_predictive_model_data.index_1 = [
    1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2,
    2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
    2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3
]


# ============================================================
# Data contents for posterior predictive
# ============================================================
posterior_model_data.N = 63
posterior_model_data.J_1 = 3
posterior_model_data.index_1 = [
    1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2,
    2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
    2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3
]


# ============================================================
# Parameters of interest
# ============================================================
var_names = ["rho", "sigma", "tau", "theta"]
ignored_var_names = ["*_prc", "log_lik"]


# ============================================================
# Which variables are "measured"
# ============================================================
measured_data = ["x", "y"]


# ============================================================
# Data types of measured variables
# ============================================================
measured_data_dtypes.x = "float"
measured_data_dtypes.y = "float"


# ============================================================
# Inferential calibration settings
# ============================================================
N = 1000
cores = 8
n_prior_draws_for_sd = 1000
samples_dir = "inferential_calibration_samples"
remove_sample_files = true
progress_bar = true
output_csv = "inferential_calibration_results.csv"
overwrite = true


# ============================================================
# Sampling settings
# ============================================================
sampling_kwargs.chains = 4
sampling_kwargs.iter_sampling = 1000
sampling_kwargs.iter_warmup = 2000
sampling_kwargs.adapt_delta = 0.95
sampling_kwargs.max_treedepth = 10
sampling_kwargs.thin = 1


# ============================================================
# Diagnostic kwargs
# ============================================================
diagnostic_check_kwargs.adapt_target = 0.801
diagnostic_check_kwargs.adapt_target_fraction = 0.9
diagnostic_check_kwargs.efmi_rule_of_thumb = 0.2
diagnostic_check_kwargs.max_treedepth = 10
diagnostic_check_kwargs.min_ess_hat_per_chain = 100
diagnostic_check_kwargs.xi_hat_threshold = 0.25

Using stanfunctions

The bebi103 package includes functions that can be included in the functions block of a Stan program. As an example, let us analyze data that are inverse Gaussian distributed, a distribution that is not natively included in Stan. The bebi103 package has functionality for sampling out of this distribution and also for using it in models.

The generative model we will use is as follows

\begin{align} \log_{10}\mu \sim \text{Norm}(-2, 2),\\[1em] \log_{10}\lambda \sim \text{Norm}(-2, 2),\\[1em] y_i \sim \text{InvGaussian}(\mu, \lambda)\;\forall i. \end{align}

First, we’ll generate the dat for some contrived parameters.

[26]:

rng = np.random.default_rng(seed=3252)

# Parameters
mu = 6.5
lam = 11.25

# Generated data
N = 300
y = rng.wald(mu, lam, size=N)

For our Stan model, we include the inv_gaussian.stanfunctions, which will be available when we build the model using cmdstanpy with the appropriate keyword argument for the include path for the Stan compiler, as we will do in a moment. The Stan code is below.

functions {
    #include inv_gaussian.stanfunctions
}


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


parameters {
    real log10_mu;
    real log10_lambda;
}


transformed parameters {
    real mu = 10^log10_mu;
    real lambda = 10^log10_lambda;
}

model {
    log10_mu ~ normal(-2, 2);
    log10_lambda ~ normal(-2, 2);
    y ~ inv_gaussian(mu, lambda);
}


generated quantities {
    array[N] real y_prc;
    for (i in 1:N) {
        y_prc[i] = inv_gaussian_rng(mu, lambda);
    }
}

Now, to compile the model, we need to tell the Stan compiler where to find the inverse Gaussian Stan functions. We do that with the stanc_options keyword argument below.

[27]:

sm = cmdstanpy.CmdStanModel(
    stan_file="inv_gauss.stan",
    stanc_options={"include-paths": bebi103.hmc.include_path()}
)

20:45:23 - cmdstanpy - INFO - compiling stan file /Users/bois/Dropbox/git/bebi103/doc/user_guide/inv_gauss.stan to exe file /Users/bois/Dropbox/git/bebi103/doc/user_guide/inv_gauss
20:45:29 - cmdstanpy - INFO - compiled model executable: /Users/bois/Dropbox/git/bebi103/doc/user_guide/inv_gauss

Let’s take some samples and see how we captured the parameters we used to generate the data.

[28]:

# Perform sampling
with bebi103.hmc.disable_logging():
    samples = sm.sample(data=dict(y=y, N=N))

# Check diagnostics
bebi103.hmc.check_all_diagnostics(samples)

# Make a corner plot
bokeh.io.show(bebi103.viz.corner(samples, parameters=["mu", "lambda"]))


All HMC diagnostics are consistent with accurate Markov chain Monte
Carlo.

All expectands checked appear to be behaving well enough for reliable
Markov chain Monte Carlo estimation.

The parameter \(\mu\) to generate the data was 6.5 and \(\lambda\) was 11.25, which are well within the credible region. We can also check to see how the model predicts the measurements using the predictive_ecdf() function of the bebi103.viz module.

[29]:

# Put data in right structure for predictive_ecdf()
y_prc = np.concatenate(bebi103.hmc.extract_variable(samples, 'y_prc'))

bokeh.io.show(bebi103.viz.predictive_ecdf(y_prc, data=y, y_axis_label="y"))

In addition to functions enabling use of an inverse Gaussian distribution, the bebi103 package has Stan functions useful for working with Gaussian processes when the x-variables are one-dimensional. The contents of the gp_one_dimensional.stanfunctions file are below.

/*
  These functions are for GPs with univariate x. The functions
  gp_posterior_mstar and gp_posterior_sigmastar_cholesky are general.
  The remaining functions are useful for computing derivatives of
  GPs that use exponentiated quadratic (a.k.a. squared exponential or
  radial basis function), Matern 3/2, and Matern 5/2 kernels.
 */

/**
  * Return posterior m* for a model with a Normal likelihood and GP
  * prior for a given Cholesky decomposition, Ly, of the matrix Ky,
  * and K*.
  *
  * @param y y-values of data
  * @param Ly Cholesky decomposition of the matrix Ky
  * @param Kstar Matrix K*
  * @return Posterior m*
  */
vector gp_posterior_mstar(vector y, matrix Ly, matrix Kstar) {
  // Get sizes
  int N = size(y);
  int Nstar = cols(Kstar);

  // Compute xi = inv(Ky) . y, which is solution xi to Ky . xi = y.
  vector[N] z = mdivide_left_tri_low(Ly, y);
  vector[N] xi = mdivide_right_tri_low(z', Ly)';

  // Compute mean vector mstar
  vector[Nstar] mstar = Kstar' * xi;

  return mstar;
}


/**
  * Return Cholesky decomposition of posterior Σ* for a model with a
  * Normal likelihood and GP prior.
  *
  * @param y y-values of data
  * @param Ly Cholesky decomposition of the matrix Ky
  * @param Kstar Matrix K*
  * @param Kstarstar Matrix K**
  * @param delta Small value to add to the diagonal of Σ* to ensure
  *              numerical positive definiteness
  * @return Cholesky decomposition of posterior Σ*
  */
matrix gp_posterior_sigmastar_cholesky(
    vector y,
    matrix Ly,
    matrix Kstar,
    matrix Kstarstar,
    real delta) {

  // Get sizes
  int N = size(y);
  int Nstar = cols(Kstar);

  // Compute Xi = inv(Ky) . Kstar, which is the solution Xi to Ky . Xi = Kstar.
  matrix[N, Nstar] Z = mdivide_left_tri_low(Ly, Kstar);
  matrix[N, Nstar] Xi = mdivide_right_tri_low(Z', Ly)';

  // Compute Sigma_star (plus a small number of the diagonal to ensure pos. def.)
  matrix[Nstar, Nstar] Sigmastar = Kstarstar - Kstar' * Xi
                                  + diag_matrix(rep_vector(delta, Nstar));

  // Compute and return Cholesky decomposition
  matrix[Nstar, Nstar] Lstar = cholesky_decompose(Sigmastar);

  return Lstar;
}


  /**
   * Return exponentiated quadratic (a.k.a. squared exponential, or SE) kernel
   * differentiated by the first variable.
   *
   * @param x1 Value of first variable
   * @param x2 Value of second variable
   * @param alpha α hyperparameter of SE kernel
   * @param rho ρ hyperparameter of SE kernel
   * @return Partial first derivative of SE kernel with respect to the
   *         first variable
   */
  real d1_exp_quad_kernel(real x1, real x2, real alpha, real rho) {
     real x_diff = x1 - x2;

     return -(alpha / rho)^2 * x_diff * exp(-x_diff^2 / 2.0 / rho^2);
  }


/**
  * Return mixed second derivative of exponentiated quadratic
  * (a.k.a. squared exponential, or SE) kernel
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of SE kernel
  * @param rho ρ hyperparameter of SE kernel
  * @return Mixed second derivative of squared exponential (SE) kernel
  */
real d1_d2_exp_quad_kernel(real x1, real x2, real alpha, real rho) {
    real rho2 = rho^2;
    real term1 = x1 - x2 + rho;
    real term2 = x2 - x1 + rho;

    return (alpha / rho2)^2 * term1 * term2 * exp(-(x1 - x2)^2 / 2.0 / rho2);
}

  /**
   * Return Matern 3/2 kernel differentiated by the first
   * variable.
   *
   * @param x1 Value of first variable
   * @param x2 Value of second variable
   * @param alpha α hyperparameter of Matern kernel
   * @param rho ρ hyperparameter of Matern kernel
   * @return Partial first derivative of Matern 3/2 kernel with respect
   * to the first variable
   */
   real d1_matern32_kernel(real x1, real x2, real alpha, real rho) {
    real x_diff = x1 - x2;

    return -3.0 * (alpha / rho)^2 * x_diff * exp(-sqrt(3.0 * x_diff^2 / rho^2));
 }


 /**
  * Return mixed second derivative of Matern 3/2 kernel
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of Matern kernel
  * @param rho ρ hyperparameter of Matern kernel
  * @return Mixed second derivative of Matern 3/2 kernel
  */
real d1_d2_matern32_kernel(real x1, real x2, real alpha, real rho) {
  real exp_arg = sqrt(3.0 * ((x1 - x2) / rho)^2);

  return 3.0 * (alpha / rho)^2 * (1.0 - exp_arg) * exp(-exp_arg);
}


/**
   * Return Matern 5/2 kernel differentiated by the first
   * variable.
   *
   * @param x1 Value of first variable
   * @param x2 Value of second variable
   * @param alpha α hyperparameter of Matern kernel
   * @param rho ρ hyperparameter of Matern kernel
   * @return Partial first derivative of Matern 5/2 kernel with respect
   * to the first variable
   */
   real d1_matern52_kernel(real x1, real x2, real alpha, real rho) {
    real x_diff = x1 - x2;
    real exp_arg = sqrt(5 * (x_diff / rho)^2);

    return -5.0 * (alpha / rho)^2 / 3.0 * x_diff * (1.0 + exp_arg) * exp(-exp_arg);
 }


 /**
  * Return mixed second derivative of Matern 5/2 kernel
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of Matern kernel
  * @param rho ρ hyperparameter of Matern kernel
  * @return Mixed second derivative of Matern 5/2 kernel
  */
real d1_d2_matern52_kernel(real x1, real x2, real alpha, real rho) {
  real x_diff = x1 - x2;
  real x_diff2 = x_diff^2;
  real rho2 = rho^2;
  real exp_arg = sqrt(5 * x_diff^2 / rho2);

  return 5.0 * (alpha / rho2)^2 / 3.0 * ((1.0 + exp_arg) * rho2 - 5.0 * x_diff2) * exp(-exp_arg);
}


/**
  * Return first derivative of the covariance matrix of the exponentiated quadratic
  * (a.k.a. squared exponential, or SE) kernel with respect to the first variable.
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of SE kernel
  * @param rho ρ hyperparameter of SE kernel
  * @return First derivative of the covariance matrix of the squared
  * exponential kernel with respect to the first variable.
  */
matrix d1_cov_exp_quad(array[] real x1, array[] real x2, real alpha, real rho) {
  int m = size(x1);
  int n = size(x2);
  matrix[m, n] d1_K;

  for (i in 1:m) {
    for (j in 1:n) {
      d1_K[i, j] = d1_exp_quad_kernel(x1[i], x2[j], alpha, rho);
    }
  }

  return d1_K;
}


/**
  * Return mixed second derivative of the covariance matrix of the
  * exponentiated quadratic (a.k.a. squared exponential, or SE) kernel.
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of SE kernel
  * @param rho ρ hyperparameter of SE kernel
  * @return Mixed second derivative of the covariance matrix of the
  * squared exponential kernel.
  */
matrix d1_d2_cov_exp_quad(array[] real x1, array[] real x2, real alpha, real rho) {
  int m = size(x1);
  int n = size(x2);
  matrix[m, n] d1_d2_K;

  for (i in 1:m) {
    for (j in 1:n) {
      d1_d2_K[i, j] = d1_d2_exp_quad_kernel(x1[i], x2[j], alpha, rho);
    }
  }

  return d1_d2_K;
}


/**
  * Return first derivative of the covariance matrix of the Matern
  * kernel with nu = 3/2 with respect to the first variable.
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of Matern 3/2 kernel
  * @param rho ρ hyperparameter of Matern 3/2 kernel
  * @return First derivative of the covariance matrix of the Matern 3/2
  * kernel with respect to the first variable.
  */
  matrix d1_cov_matern32(array[] real x1, array[] real x2, real alpha, real rho) {
    int m = size(x1);
    int n = size(x2);
    matrix[m, n] d1_K;

    for (i in 1:m) {
      for (j in 1:n) {
        d1_K[i, j] = d1_matern32_kernel(x1[i], x2[j], alpha, rho);
      }
    }

    return d1_K;
  }


/**
  * Return mixed second derivative of the covariance matrix of the Matern
  * kernel with nu = 3/2.
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of Matern 3/2 kernel
  * @param rho ρ hyperparameter of Matern 3/2 kernel
  * @return Mixed second derivative of the covariance matrix of the
  * Matern 3/2 kernel.
  */
  matrix d1_d2_cov_matern32(array[] real x1, array[] real x2, real alpha, real rho) {
    int m = size(x1);
    int n = size(x2);
    matrix[m, n] d1_d2_K;

    for (i in 1:m) {
      for (j in 1:n) {
        d1_d2_K[i, j] = d1_d2_matern32_kernel(x1[i], x2[j], alpha, rho);
      }
    }

    return d1_d2_K;
  }


/**
  * Return first derivative of the covariance matrix of the Matern
  * kernel with nu = 5/2 with respect to the first variable.
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of Matern 5/2 kernel
  * @param rho ρ hyperparameter of Matern 5/2 kernel
  * @return First derivative of the covariance matrix of the Matern 5/2
  * kernel with respect to the first variable.
  */
  matrix d1_cov_matern52(array[] real x1, array[] real x2, real alpha, real rho) {
    int m = size(x1);
    int n = size(x2);
    matrix[m, n] d1_K;

    for (i in 1:m) {
      for (j in 1:n) {
        d1_K[i, j] = d1_matern52_kernel(x1[i], x2[j], alpha, rho);
      }
    }

    return d1_K;
  }


/**
  * Return mixed second derivative of the covariance matrix of the Matern
  * kernel with nu = 5/2.
  *
  * @param x1 Value of first variable
  * @param x2 Value of second variable
  * @param alpha α hyperparameter of Matern 5/2 kernel
  * @param rho ρ hyperparameter of Matern 5/2 kernel
  * @return Mixed second derivative of the covariance matrix of the
  * Matern 5/2 kernel.
  */
  matrix d1_d2_cov_matern52(array[] real x1, array[] real x2, real alpha, real rho) {
    int m = size(x1);
    int n = size(x2);
    matrix[m, n] d1_d2_K;

    for (i in 1:m) {
      for (j in 1:n) {
        d1_d2_K[i, j] = d1_d2_matern52_kernel(x1[i], x2[j], alpha, rho);
      }
    }

    return d1_d2_K;
  }

Cleaning up

CmdStanPy leaves .hpp, ,o, .d, executable files, and sampling output. To conveniently remove them, use bebi103.hmc.clean_cmdstanpy().

[30]:

bebi103.hmc.clean_cmdstan()