Bayesian Modeling using Stan

Stan

Stan is a domain-specific programming language
Can specify probabilistic models
Has a collection of algorithms
Can perform Bayesian statistical analysis
Can be used with R, Python, Julia

Documentation and Installation

Stan is well documented and resources can be found at https://mc-stan.org/docs/
Installation instructions specific to your domain is available here for R and Python
You will need a C++ toolchain installed to install it properly on your machine
Once installed, you can almost entirely work with your choice of domain (R or Python) software using Stan’s domain interface (and perhaps a notepad)

Workflow

Basic Stan usage has two components:
- A .stan file containing all the information about the model, the parameters you want to know about, and the data you want to fit your model on
- The corresponding functions in R or Python to run the Stan model and get posterior samples and analysis.

R functions

A typical R Script for running a Stan code requires the cmdstanr package and its dependencies.
After writing the .stan file, we mae use of the cmdstan_model function in R. Doing so results in
- Translating the Stan code into C++
- Compiles and runs the C++ code
- Draws posterior samples for the required parameters and store them

R functions

## Load or prep data
schools_data <- list(
  J = 8,
  y = c(28,  8, -3,  7, -1,  1, 18, 12),
  sigma = c(15, 10, 16, 11,  9, 11, 10, 18)
)

## Run the Stan Function
library(cmdstanr)
file2 <- file.path("school.stan")    # Reads from a file
mod <- cmdstan_model(file2) # Compiles the model given in the file

fit <- mod$sample(
  data = schools_data,
  seed = 123,
  chains = 4,
  parallel_chains = 4,
  refresh = 500 # draws posterior samples
)

fit$summary()


# summarise all variables with default and additional summary measures
fit$summary(
  variables = NULL,
  posterior::default_summary_measures(),
  extra_quantiles = ~posterior::quantile2(., probs = c(.0275, .975))
)

# default is a 3-D draws_array object from the posterior package
# iterations x chains x variables
draws_arr <- fit$draws() # or format="array"

# Diagnostics
# this is a draws_array object from the posterior package
str(fit$sampler_diagnostics())

Python functions

A typical Python Script for running a Stan code requires the cmdstanpy module.
After writing the .stan file, we mae use of the stan.build and posterior.sample functions in Python. Doing so requires importing the stan module results in
- Translating the Stan code into C++
- Compiles and runs the C++ code
- Draws posterior samples for the required parameters and store them

Python functions

import os

from cmdstanpy import CmdStanModel

stan_file = os.path.join(filepath, 'school.stan')  # put the filepath in '...' before using

model = CmdStanModel(stan_file=stan_file)  # Compiles model

data_file = os.path.join(filepath, 'school.data.json') # put the filepath in '...' before using

fit = model.sample(data=data_file)

fit.summary()

print(fit.diagnose())

RStan and PyStan

There are earlier packages RStan and PyStan
The primary advantages of using these is that you can use other packages developed using these in your works, such as the R package brms
But Cmd variants do not require compilation by R or Python, instead use their own compilers
They also have less overhead and therefore is generally faster

The Stan Component

As is evident from the workflow, a major portion is creating the contents of the .stan file
That is where you specify what model you want to run, on what data, and what parameters you want to draw posteriors for
This requires knowledge about Stan programming language
The language is simple and intuitive, with data types similar to any other prevalent programming language

Components of the Stan File

There are three (or four) main components for a Stan file

data {
  
}

parameters {
  
}

transformed parameters {
  
}

model {
  
}

The Data Component

The data component does not include the actual data
The actual data is loaded in R and supplied to Stan as a list

## Load or prep data
schools_data <- list(
  J = 8,
  y = c(28,  8, -3,  7, -1,  1, 18, 12),
  sigma = c(15, 10, 16, 11,  9, 11, 10, 18)
)
## Run the Stan Function
library(cmdstanr)
file2 <- file.path("school.stan")    # Reads from a file
mod <- cmdstan_model(file2) # Compiles the model given in the file

fit <- mod$sample(
  data = schools_data,
  seed = 123,
  chains = 4,
  parallel_chains = 4,
  refresh = 500 # draws posterior samples
)

Data Component in Stan File

Contains the data types of the variables in the actual data
You tell Stan what variables you expect in the data and what are they like

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

Data Types in Stan

A complete list is available here
Most important ones for us are
- int for integers
- real for real numbers
- vector[N] for real column vector of size N
- matrix[M,N] for real matrix of size M x N
- array[M,N,P] for arrays of size M x N x P
For arrays, you have declare the data type that fills them, such as array[N] int

Constrained Data Types

You can put lower and/or upper bounds on integers and real numbers
This can be done by
- int<lower=l> for integers bigger or equal to l
- real<upper=u> for reals smaller or equal to u
- real<lower=l,upper=u> for reals in between l and u

Revisiting the Data Part

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

This is telling Stan that
- There are three entries in data named J, y, and sigma
- The variable J is a non-negative integer
- The variable y is a real vector of size J
- The variable sigma is a non-negative real valued vector of size J

The Model and Parameters Part

These two parts give the description of the statistical model being used and any parameters it has
The parameters part tells Stan what the model parameters are called and what their data type is
The model part tells Stan what is the distribution of the variables and the priors for the parameters

Model for This Problem

We want to fit the following model: \[ Y_j \sim N(\theta_j,\sigma_j)\] \[ \theta_j \sim N(\mu, \tau)\] \[\pi(\mu,\tau) \propto 1\]

The parameters are then \(\theta_j\), \(\mu\), and \(\tau\)

Parameter Part for This Problem

The parameter part in the Stan script then should like:

parameters {
  vector[J] theta;
  real mu;
  real<lower=0> tau;
}

Model Part

There are two ways to define the model for Stan
- Using distributional expression
- Using log posterior

Using Distributional Expression

model {
  for(j in 1:J){
    y[j] ~ normal(theta[j],sigma[j]);
    theta[j] ~ normal(mu,tau);
  }
}

Declaring a parameter and not declaring any distribution for them results in Stan assigning flat prior to them

Model Part with Proper Priors

If we were to use \[\mu \sim N(0,100),\] \[\tau \sim IG(0.1,0.1),\] the code would look like:

model {
  for(j in 1:J){
    y[j] ~ normal(theta[j],sigma[j]);
    theta[j] ~ normal(mu,tau);
  }
  mu ~ normal(0,100);
  tau ~ inv_gamma(0.1,0.1);
}

Vectorization in Stan

Due to vectorization in Stan, you can do away with the for loop, and it will still understand:

model {
  y ~ normal(theta,sigma);
  theta ~ normal(mu,tau);
  mu ~ normal(0,100);
  tau ~ inv_gamma(0.1,0.1);
}

Available distributions in Stan

A complete list can be obtained here
Almost all standard distributions are available
Their truncated versions are also available
Multivariate densities such as multivariate normal, multivariate student’s t is available
Distributions on correlation and covariance matrices are also available

Model Part with log-pdf

We can alternatively define the model part with the log-pdf

model {
  target += normal_lpdf(y | theta, sigma);
  target += normal_lpdf(theta | mu, tau);
  target += normal_lpdf(mu | 0, 100);
  target += inv_gamma_lpdf(tau | 0.1, 0.1);
}

Differences

The distribution statement based models are faster
The log-pdf based version computes the exact log-posterior including the constants
The log-pdf version can be used to model custom distributions

Custom Densities

Consider that we wish to draw posteriors from a Gaussian mixture model for some parameter \(\mu\) with density \[f(\mu) = 0.3 \phi\left(\frac{\mu+1}{2}\right) + 0.7 \phi\left(\frac{\mu-1}{3}\right).\]
This can be achieved with the following code:

parameters {
  real mu;
}

model {
  target += log_sum_exp(log(0.3) + normal_lpdf(mu | -1, 2^2),
                        log(0.7) + normal_lpdf(mu | 1, 3^2));
}

log_sum_exp is an inbuilt function in Stan

Transformed Parameters

Often parameters are transformed to get better MCMC sampling performance
For example, in our model, we had \[Y_j \sim N(\theta_j,\sigma_j),\] \[\theta_j \sim N(\mu, \tau).\]
This is equivalent to writing \[Y_j \sim N(\mu + \tau\eta_j,\sigma_j),\] \[\eta_j \sim N(0,1).\]
Often the second formulation gives a more stable Markov chain than the former
Here the parameters are \(\mu\), \(\tau\), and \(\eta_j\). But we want inference about the means, which we can call \(\theta_j = \mu + \tau\eta_j\)

Transformed Parameters in Stan

The last model with transformed parameters can be incorporated into Stan with the following code:

data {
  int<lower=0> J;
  array[J] real y;
  array[J] real<lower=0> sigma;
}
parameters {
  vector[J] eta;
  real mu;
  real<lower=0> tau;
}
transformed parameters {
  vector[J] theta;
  theta = mu + tau*eta;
}
model {
  y ~ normal(theta,sigma);
  eta ~ normal(0,1);
}

Revisiting the R Code

## Load or prep data
schools_data <- list(
  J = 8,
  y = c(28,  8, -3,  7, -1,  1, 18, 12),
  sigma = c(15, 10, 16, 11,  9, 11, 10, 18)
)

## Run the Stan Function
library(cmdstanr)
file2 <- file.path("school.stan")    # Reads from a file
mod <- cmdstan_model(file2, # Compiles the model given in the file
                     compile_model_methods = FALSE,
                     force_recompile = TRUE) 

fit <- mod$sample(
  data = schools_data,
  seed = 123,
  chains = 2,
  parallel_chains = 1,
  iter_warmup = 1000,
  iter_sampling = 2000,
  save_warmup = TRUE,
  thin = 1,
  refresh = 500 # draws posterior samples
)

Running MCMC with 2 sequential chains...

Chain 1 Iteration:    1 / 3000 [  0%]  (Warmup) 
Chain 1 Iteration:  500 / 3000 [ 16%]  (Warmup) 
Chain 1 Iteration: 1000 / 3000 [ 33%]  (Warmup) 
Chain 1 Iteration: 1001 / 3000 [ 33%]  (Sampling) 
Chain 1 Iteration: 1500 / 3000 [ 50%]  (Sampling) 
Chain 1 Iteration: 2000 / 3000 [ 66%]  (Sampling) 
Chain 1 Iteration: 2500 / 3000 [ 83%]  (Sampling) 
Chain 1 Iteration: 3000 / 3000 [100%]  (Sampling) 
Chain 1 finished in 0.0 seconds.
Chain 2 Iteration:    1 / 3000 [  0%]  (Warmup) 
Chain 2 Iteration:  500 / 3000 [ 16%]  (Warmup) 
Chain 2 Iteration: 1000 / 3000 [ 33%]  (Warmup) 
Chain 2 Iteration: 1001 / 3000 [ 33%]  (Sampling) 
Chain 2 Iteration: 1500 / 3000 [ 50%]  (Sampling) 
Chain 2 Iteration: 2000 / 3000 [ 66%]  (Sampling) 
Chain 2 Iteration: 2500 / 3000 [ 83%]  (Sampling) 
Chain 2 Iteration: 3000 / 3000 [100%]  (Sampling) 
Chain 2 finished in 0.0 seconds.

Both chains finished successfully.
Mean chain execution time: 0.0 seconds.
Total execution time: 0.3 seconds.

Outputs from R Code

fit$summary()

# A tibble: 19 × 10
   variable    mean  median    sd   mad      q5   q95  rhat ess_bulk ess_tail
   <chr>      <dbl>   <dbl> <dbl> <dbl>   <dbl> <dbl> <dbl>    <dbl>    <dbl>
 1 lp__     -4.80   -4.54   2.61  2.49  -9.49   -1.12 1.00     1275.    2189.
 2 eta[1]    0.388   0.413  0.943 0.926 -1.20    1.90 1.000    3748.    2724.
 3 eta[2]   -0.0262 -0.0399 0.884 0.832 -1.49    1.45 1.000    4037.    2745.
 4 eta[3]   -0.202  -0.204  0.940 0.907 -1.70    1.37 1.00     3692.    2817.
 5 eta[4]   -0.0615 -0.0660 0.871 0.824 -1.49    1.39 1.00     3221.    2362.
 6 eta[5]   -0.350  -0.348  0.855 0.837 -1.75    1.10 1.00     3401.    2749.
 7 eta[6]   -0.210  -0.229  0.866 0.864 -1.58    1.23 1.000    3570.    2793.
 8 eta[7]    0.333   0.341  0.877 0.875 -1.14    1.75 1.00     3069.    2866.
 9 eta[8]    0.0255  0.0272 0.921 0.903 -1.46    1.54 1.00     3747.    2999.
10 mu        8.21    8.03   5.30  4.78   0.0831 16.7  1.00     1098.     486.
11 tau       6.69    5.32   5.78  4.76   0.540  17.9  1.00     1201.     886.
12 theta[1] 11.7    10.4    8.51  7.00   0.389  28.1  1.00     2096.    1537.
13 theta[2]  7.91    7.86   6.18  5.57  -2.01   18.1  1.00     4999.    3363.
14 theta[3]  6.33    6.87   7.87  6.61  -7.54   18.2  1.00     3401.    2930.
15 theta[4]  7.55    7.73   6.60  5.95  -3.34   17.9  1.000    4738.    3124.
16 theta[5]  5.34    5.64   6.21  5.90  -5.45   14.8  1.00     3850.    3105.
17 theta[6]  6.29    6.66   6.79  6.05  -5.75   16.7  1.00     4376.    2814.
18 theta[7] 10.7    10.2    6.82  6.07   0.800  23.1  1.00     3120.    2876.
19 theta[8]  8.41    8.28   7.75  6.30  -3.71   21.1  1.00     3285.    2128.

Outputs from R Code (Contd.)

# default is a 3-D draws_array object from the posterior package
# iterations x chains x variables
draws_arr <- fit$draws() # or format="array"
str(draws_arr)

 'draws_array' num [1:2000, 1:2, 1:19] -2.492 -0.735 -6.395 -3.985 -5.681 ...
 - attr(*, "dimnames")=List of 3
  ..$ iteration: chr [1:2000] "1" "2" "3" "4" ...
  ..$ chain    : chr [1:2] "1" "2"
  ..$ variable : chr [1:19] "lp__" "eta[1]" "eta[2]" "eta[3]" ...

# Diagnostics
# this is a draws_array object from the posterior package
str(fit$sampler_diagnostics())

 'draws_array' num [1:2000, 1:2, 1:6] 3 3 3 3 4 3 3 3 3 3 ...
 - attr(*, "dimnames")=List of 3
  ..$ iteration: chr [1:2000] "1" "2" "3" "4" ...
  ..$ chain    : chr [1:2] "1" "2"
  ..$ variable : chr [1:6] "treedepth__" "divergent__" "energy__" "accept_stat__" ...