Bayesian data analysis - CmdStanPy demos

Author: Aki Vehtari, Tuomas Sivula, Co-authors: Lassi Meronen, Pellervo Ruponen, Osvaldo A Martin

Last modified 2025-Sep.

License: CC-BY

This notebook contains several examples of how to use Stan in python with CmdStanPy. This notebook assumes basic knowledge of Bayesian inference and MCMC. The Stan models are stored in separate .stan-files. The examples are related to Bayesian data analysis course by Aki Vehtari.

Setting up the CmdStanPy environment

We begin by importing the CmdStanPy module as well as other useful modules.

We also import some utilities by Michael Betancourt and introduced in PyStan workflow case study in module check_utility.

import os, sys

import arviz as az
import numpy as np
import matplotlib.pyplot as plt

# import stan interface
from cmdstanpy import CmdStanModel

# add utilities directory to path
# so we can import check_utility
util_path = os.path.abspath(os.path.join(os.path.pardir, 'utilities_and_data'))
if util_path not in sys.path and os.path.exists(util_path):
    sys.path.insert(0, util_path)
from check_utility import check_div, check_energy, check_treedepth

az.style.use('arviz-doc')
az.rcParams["stats.ci_prob"] = 0.95  # set default hdi to 95%

Bernoulli model

Toy data with sequence of failures (0) and successes (1). We would like to learn about the unknown probability of success. In following, we denote this probability with theta.

data = dict(N=10, y=[0,1,0,0,1,1,1,0,1,0])

Bernoulli model with a Beta(1,1) (uniform) prior

with open('bern.stan') as file:
    print(file.read())

// Bernoulli model
data {
  int<lower=0> N;
  array[N] int<lower=0, upper=1> y;
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}

Given the Stan program we then use the compile_model method of check_utility module to compile the Stan program into a C++ executable. This utility function automatically saves a cached version of the compiled model to the disk for possible future use.

model = CmdStanModel(stan_file='bern.stan')

Getting the model again with the utility function uses the cached version automatically.

del model
model = model = CmdStanModel(stan_file='bern.stan')

Sample from the posterior, show the summary and plot the histogram of the posterior draws. In addition, plot estimated posterior density with arviz posterior plot and 95% credible interval. We recommend explicitly specifying the seed of Stan’s random number generator, as we have done here, so that we can reproduce these exactly results in the future, at least when using the same machine, operating system, and interface. This is especially helpful for the more subtle pathologies that may not always be found, which results in seemingly stochastic behavior.

fit = model.sample(data=data, seed=194838)
idata = az.from_cmdstanpy(fit)


ax = az.plot_posterior(fit, var_names=['theta'], round_to=2)
ax.set_xlabel(r'Probability of success (theta)')
ax.set_ylabel('Relative density');

Binomial model

Instead of sequence of 0’s and 1’s, we can summarize the data with the number of experiments and the number successes:

data = dict(N=10, y=7)

And then we use Binomial model with Beta(1,1) prior for the probability of success.

with open('binom.stan') as file:
    print(file.read())

// Binomial model with beta(1,1,) prior
data {
  int<lower=0> N;
  int<lower=0> y;
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ binomial(N,theta);
}

Sample from the posterior and plot the posterior. The histogram should look similar as in the Bernoulli case, except that now the number of successes is 7 instead of 5.

model = CmdStanModel(stan_file='binom.stan')
fit = model.sample(data=data, seed=194838)
samples = az.from_cmdstanpy(fit)


ax = az.plot_posterior(fit, var_names=['theta'], round_to=2)
ax.set_xlabel(r'Probability of success (theta)')
ax.set_ylabel('Relative density');

Now we re-run the model with a new data (now number of successes being 70 out of 100). The previously compiled Stan program is re-used making the re-use faster.

data = dict(N=100, y=70)
fit = model.sample(data=data, seed=194838)
samples = az.from_cmdstanpy(fit)

ax = az.plot_posterior(fit, var_names=['theta'], round_to=2)
ax.set_xlabel(r'Probability of success (theta)')
ax.set_ylabel('Relative density');

Explicit transformation of variables

In the above examples the probability of success θ was declared as:

real<lower=0,upper=1> theta;

Stan makes automatic transformation of the variable to the unconstrained space using logit transofrmation for interval constrained and log transformation for half constraints.

The following example shows how we can also make an explicit transformation and use binomial_logit function which takes the unconstrained parameter as an argument and uses logit transformation internally. This form can be useful for better numerical stability.

with open('binomb.stan') as file:
    print(file.read())

// Binomial model with a roughly uniform prior for
// the probability of success (theta)
data {
  int<lower=0> N;
  int<lower=0> y;
}
parameters {
  real alpha;
}
transformed parameters {
  real theta;
  theta = inv_logit(alpha);
}
model {
  // roughly auniform prior for the number of successes
  alpha ~ normal(0,1.5);
  y ~ binomial_logit(N,alpha);
}

model = CmdStanModel(stan_file='binomb.stan')

Here we have used Gaussian prior in the unconstrained space, which produces close to uniform prior for theta.

Sample from the posterior and plot the posterior. The histogram should look similar as with the previous model

data = dict(N=100, y=70)
fit = model.sample(data=data, seed=194838)
samples = az.from_cmdstanpy(fit)

ax = az.plot_posterior(fit, var_names=['theta'], round_to=2)
ax.set_xlabel(r'Probability of success (theta)')
ax.set_ylabel('Relative density');

Comparison of two groups with Binomial

An experiment was performed to estimate the effect of beta-blockers on mortality of cardiac patients. A group of patients were randomly assigned to treatment and control groups:

• out of 674 patients receiving the control, 39 died
• out of 680 receiving the treatment, 22 died

Data:

data = dict(N1=674, y1=39, N2=680, y2=22)

To analyse whether the treatment is useful, we can use Binomial model for both groups and compute odds-ratio. If the odds-ratio is less than 1, it means that the treatment is useful.

with open('binom2.stan') as file:
    print(file.read())

//  Comparison of two groups with Binomial
data {
  int<lower=0> N1;
  int<lower=0> y1;
  int<lower=0> N2;
  int<lower=0> y2;
}
parameters {
  real<lower=0,upper=1> theta1;
  real<lower=0,upper=1> theta2;
}
model {
  theta1 ~ beta(1,1);
  theta2 ~ beta(1,1);
  y1 ~ binomial(N1,theta1);
  y2 ~ binomial(N2,theta2);
}
generated quantities {
  real oddsratio;
  oddsratio = (theta2/(1-theta2))/(theta1/(1-theta1));
}

Sample from the posterior and plot the posterior for the odds-ratio. We plot the posterior both as a histogram and and as an estimated probability density function. Using ArviZ, we can easily see the probability of the odds-ratio being less than 1.

model = CmdStanModel(stan_file='binom2.stan')
fit = model.sample(data=data, seed=194838)
samples = az.from_cmdstanpy(fit)

# With az.plot_posterior we use parameter ref_val = 1
# since treatment is beneficial if oddsratio is smaller than 1.
ax = az.plot_posterior(fit, var_names=['oddsratio'], round_to=2, ref_val=1)
ax.set_xlabel('oddsratio')
ax.set_ylabel('Relative density')
plt.savefig('odd-ratio.pdf')

Linear Gaussian model

The following file has Kilpisjärvi summer month temperatures 1952-2013:

data_path = os.path.abspath(
    os.path.join(
        os.path.pardir,
        'utilities_and_data',
        'kilpisjarvi-summer-temp.csv'
    )
)
d = np.loadtxt(data_path, dtype=np.double, delimiter=';', skiprows=1)
x = d[:, 0]
y = d[:, 4]
N = len(x)
xpred = 2016

Plot the data

plt.figure(figsize=(8, 4))
plt.scatter(x, y)
plt.xlabel('year')
plt.ylabel('summer temperature at Kilpisjarvi');

To analyse whether the average summer month temperature is rising, we use a linear model with Gaussian model for the unexplained variation.

Gaussian linear model with adjustable priors

The following Stan code allows also setting hyperparameter values as data allowing easier way to use different priors in different analyses:

with open('lin.stan') as file:
    print(file.read())

// Gaussian linear model with adjustable priors
data {
  int<lower=0> N; // number of data points
  vector[N] x; //
  vector[N] y; //
  real xpred; // input location for prediction
  real pmualpha; // prior mean for alpha
  real psalpha;  // prior std for alpha
  real pmubeta;  // prior mean for beta
  real psbeta;   // prior std for beta
}
parameters {
  real alpha;
  real beta;
  real<lower=0> sigma;
}
transformed parameters {
  vector[N] mu;
  mu = alpha + beta*x;
}
model {
  alpha ~ normal(pmualpha, psalpha);
  beta ~ normal(pmubeta, psbeta);
  y ~ normal(mu, sigma);
}
generated quantities {
  real ypred;
  vector[N] log_lik;
  ypred = normal_rng(alpha + beta*xpred, sigma);
  for (i in 1:N)
    log_lik[i] = normal_lpdf(y[i] | mu[i], sigma);
}

Create a list with data and priors:

data = dict(
    N = N,
    x = x,
    y = y,
    xpred = xpred,
    pmualpha = y.mean(),     # centered
    psalpha  = 100,          # weakly informative prior
    pmubeta  = 0,            # a priori increase and decrease as likely
    psbeta   = (.1--.1)/6.0  # avg temp probably does not increase more than 1
                             # degree per 10 years
)

Run Stan

model = CmdStanModel(stan_file='lin.stan')
fit_lin = model.sample(data=data, seed=194838)
idata_lin = az.from_cmdstanpy(fit_lin)

Check the n_eff and Rhat

az.summary(idata_lin, round_to=2)

	mean	sd	hdi_2.5%	hdi_97.5%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	-28.45	15.13	-59.73	0.37	0.42	0.32	1326.36	1303.04	1.01
beta	0.02	0.01	0.00	0.03	0.00	0.00	1326.10	1289.75	1.01
sigma	1.13	0.11	0.94	1.34	0.00	0.00	1601.76	1794.23	1.00
mu[0]	8.73	0.28	8.18	9.26	0.01	0.00	1643.19	1914.22	1.00
mu[1]	8.75	0.27	8.21	9.26	0.01	0.00	1661.56	1992.92	1.00
...	...	...	...	...	...	...	...	...	...
mu[58]	9.84	0.26	9.33	10.33	0.01	0.00	1746.94	2061.76	1.00
mu[59]	9.86	0.26	9.34	10.36	0.01	0.00	1722.75	2040.03	1.00
mu[60]	9.88	0.27	9.36	10.40	0.01	0.00	1700.51	2008.37	1.00
mu[61]	9.90	0.27	9.37	10.44	0.01	0.00	1680.57	1998.57	1.00
ypred	9.96	1.18	7.62	12.21	0.02	0.01	3755.17	3813.69	1.00

66 rows × 9 columns

Check the treedepth, E-BFMI, and divergences

check_treedepth(idata_lin)
check_energy(idata_lin)
check_div(idata_lin)

68 of 4000 iterations saturated the maximum tree depth of 10 (1.7%)
Run again with max_depth set to a larger value to avoid saturation

We see that several iterations saturated the maximum treedepth. The reason for this is a very strong posterior correlation between alpha and beta as shown in the next plot. This doesn’t invalidate the results, but leads to suboptimal performance. We’ll later look at alternative which reduces the posterior correlation. The following figure shows the strong posterior correlation in the current case.

az.plot_pair(idata_lin, var_names=['alpha', 'beta']);

Compute the probability that the summer temperature is increasing. The beta-parameter represents the slope of the temperature change. Hence, by checking the probability that beta > 0, we get the probability that the temperature is rising.

print('Pr(beta > 0) = {}'.format(np.mean(idata_lin.posterior['beta'] > 0)))

Pr(beta > 0) = <xarray.DataArray 'beta' ()> Size: 8B
array(0.99625)

Plot the data, the model fit and prediction for year 2016.

fig, ax = plt.subplots(1, 1, figsize=(12, 5))

az.plot_hdi(x, idata_lin.posterior['mu'], ax=ax)
ax.plot(x, idata_lin.posterior['mu'].median(["chain", "draw"]), color="C1")

ax.scatter(x, y, 10, color="C0")
ax.set_xlabel('year')
ax.set_ylabel('summer temp. @Kilpisjarvi')
ax.set_xlim((1951.5, 2013.5))


az.plot_posterior(fit_lin, var_names=['beta','sigma','ypred'], round_to=2, kind='hist', bins=35, figsize=(12,3));

For the beta parameter (the slope), let’s also plot the relative density with reference value of 0. Now we can easily evaluate the probability for slope being positive.

ax = az.plot_posterior(fit_lin, var_names=['beta'], round_to=2, ref_val=0, kind='kde')
ax.set_ylabel('Relative density');

Gaussian linear model with standardized data

In the above we used the unnormalized data and as x values are far away from zero, this will cause very strong posterior dependency between alpha and beta. The strong posterior dependency can be removed by normalizing the data to have zero mean. The following Stan code makes it in Stan. In generated quantities we do correspnding transformation back to the original scale.

with open('lin_std.stan') as file:
    print(file.read())

// Gaussian linear model with standardized data
data {
  int<lower=0> N; // number of data points
  vector[N] x; //
  vector[N] y; //
  real xpred; // input location for prediction
}
transformed data {
  vector[N] x_std;
  vector[N] y_std;
  real xpred_std;
  x_std = (x - mean(x)) / sd(x);
  y_std = (y - mean(y)) / sd(y);
  xpred_std = (xpred - mean(x)) / sd(x);
}
parameters {
  real alpha;
  real beta;
  real<lower=0> sigma_std;
}
transformed parameters {
  vector[N] mu_std;
  mu_std = alpha + beta*x_std;
}
model {
  alpha ~ normal(0, 1);
  beta ~ normal(0, 1);
  y_std ~ normal(mu_std, sigma_std);
}
generated quantities {
  vector[N] mu;
  real<lower=0> sigma;
  real ypred;
  vector[N] log_lik;
  mu = mu_std*sd(y) + mean(y);
  sigma = sigma_std*sd(y);
  ypred = normal_rng((alpha + beta*xpred_std)*sd(y)+mean(y), sigma*sd(y));
  for (i in 1:N)
    log_lik[i] = normal_lpdf(y[i] | mu[i], sigma);
}

Run Stan

model = CmdStanModel(stan_file='lin_std.stan')
fit_lin_std = model.sample(data=data, seed=194838)
idata_lin_std = az.from_cmdstanpy(fit_lin_std)

Check the n_eff and Rhat

az.summary(idata_lin_std, round_to=2)

	mean	sd	hdi_2.5%	hdi_97.5%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	-0.00	0.12	-0.23	0.26	0.00	0.00	3372.10	2780.75	1.0
beta	0.32	0.12	0.07	0.54	0.00	0.00	3399.85	2796.37	1.0
sigma_std	0.98	0.09	0.80	1.17	0.00	0.00	3526.27	2769.26	1.0
mu_std[0]	-0.54	0.24	-0.98	-0.04	0.00	0.00	3333.99	2430.86	1.0
mu_std[1]	-0.52	0.23	-0.95	-0.04	0.00	0.00	3327.94	2337.85	1.0
...	...	...	...	...	...	...	...	...	...
mu[59]	9.89	0.27	9.41	10.45	0.00	0.00	3432.56	2825.51	1.0
mu[60]	9.91	0.27	9.41	10.47	0.00	0.00	3431.50	2846.56	1.0
mu[61]	9.93	0.28	9.42	10.49	0.00	0.00	3432.51	2889.90	1.0
sigma	1.13	0.11	0.93	1.35	0.00	0.00	3526.24	2769.26	1.0
ypred	9.97	1.34	7.41	12.75	0.02	0.02	3528.99	3629.67	1.0

129 rows × 9 columns

We get now much better effective sample sizes n_eff.

Check the treedepth, E-BFMI, and divergences

check_treedepth(idata_lin_std)
check_energy(idata_lin_std)
check_div(idata_lin_std)

We see no warning message, everything is fine now! The next figure shows that with the standardized data there is not much posterior correlation:

az.plot_pair(idata_lin_std, var_names=['alpha', 'beta']);

Compute the probability that the summer temperature is increasing.

print(f'Pr(beta > 0) = {np.mean(idata_lin_std.posterior['beta'] > 0).item()}')

Pr(beta > 0) = 0.99525

Plot the data, the model fit and prediction for year 2016.

fig, ax = plt.subplots(1, 1, figsize=(12, 5))

az.plot_hdi(x, idata_lin_std.posterior['mu'], ax=ax)
ax.plot(x, idata_lin_std.posterior['mu'].median(["chain", "draw"]), color="C1")

ax.scatter(x, y, 10, color="C0")
ax.set_xlabel('year')
ax.set_ylabel('summer temp. @Kilpisjarvi')
ax.set_xlim((1951.5, 2013.5))


az.plot_posterior(fit_lin_std, var_names=['beta','sigma','ypred'], round_to=2, kind='hist', bins=35, figsize=(12,3));

Now we can also plot ridgeplot for mu samples. Ridge plot better illustrates the fact that each year is represented by distribution for mu parameter.

# Make complementary ridge plot
axes = az.plot_forest(idata_lin_std,
                           kind='ridgeplot',
                           var_names=['mu'],
                           combined=True,
                           r_hat = False,
                           ess = True,
                           ridgeplot_overlap=5,
                           ridgeplot_alpha=0.7,
                           colors='lightblue',
                           figsize=(9, 7))
# Let's adjust labels (default arviz labels are not so good for large amount of distributions)
axes[0].set_yticklabels(
    [str(year) if year % 2 == 0 else '' for year in range(2013, 1952-1, -1)])
axes[0].set_ylabel('Year')
axes[0].set_xlabel('Average temperature');

Linear Student’s t model.

The temperatures used in the above analyses are averages over three months, which makes it more likely that they are normally distributed, but there can be extreme events in the weather and we can check whether more robust Student’s t observation model would give different results. For the Student’s t distribution we also need the degree of freedom, nu, as an additional parameter.

with open('lin_t.stan') as file:
    print(file.read())

// Linear student-t model
data {
  int<lower=0> N; // number of data points
  vector[N] x; //
  vector[N] y; //
  real xpred; // input location for prediction
}
parameters {
  real alpha;
  real beta;
  real<lower=0> sigma;
  real<lower=1, upper=80> nu;
}
transformed parameters {
  vector[N] mu;
  mu = alpha + beta*x;
}
model {
  nu ~ gamma(2, 0.1); // Juarez and Steel(2010)
  y ~ student_t(nu, mu, sigma);
}
generated quantities {
  real ypred;
  vector[N] log_lik;
  ypred = normal_rng(alpha + beta*xpred, sigma);
  for (i in 1:N)
    log_lik[i] = student_t_lpdf(y[i] | nu, mu[i], sigma);
}

model = CmdStanModel(stan_file='lin_t.stan')
fit_lin_t = model.sample(data=data, seed=194838)
idata_lin_t = az.from_cmdstanpy(fit_lin_t)

Check the n_eff and Rhat

az.summary(idata_lin_t, round_to=2)

	mean	sd	hdi_2.5%	hdi_97.5%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	-32.78	15.78	-65.53	-4.78	0.41	0.30	1506.18	1771.05	1.0
beta	0.02	0.01	0.01	0.04	0.00	0.00	1506.10	1821.18	1.0
sigma	1.08	0.12	0.87	1.32	0.00	0.00	2025.73	1908.44	1.0
nu	23.83	12.97	4.65	50.58	0.31	0.28	1796.00	1916.52	1.0
mu[0]	8.66	0.29	8.12	9.23	0.01	0.01	1771.78	2179.43	1.0
...	...	...	...	...	...	...	...	...	...
mu[58]	9.89	0.26	9.40	10.38	0.01	0.00	1888.60	2545.49	1.0
mu[59]	9.91	0.26	9.41	10.42	0.01	0.00	1864.15	2534.61	1.0
mu[60]	9.93	0.27	9.41	10.44	0.01	0.00	1843.39	2406.09	1.0
mu[61]	9.96	0.28	9.44	10.50	0.01	0.00	1823.61	2476.28	1.0
ypred	10.01	1.14	7.87	12.28	0.02	0.01	3529.88	3925.69	1.0

67 rows × 9 columns

Again without standardization we get smaller n_eff’s, but large enough for practical purposes.

Check the treedepth, E-BFMI, and divergences

check_treedepth(idata_lin_t)
check_energy(idata_lin_t)
check_div(idata_lin_t)

66 of 4000 iterations saturated the maximum tree depth of 10 (1.6%)
Run again with max_depth set to a larger value to avoid saturation

We see again many iterations saturating the maximum tree depth, which is harmful for the efficiency but doesn’t invalidate the results.

Compute the probability that the summer temperature is increasing.

print(f'Pr(beta > 0) = {np.mean(idata_lin_t.posterior['beta'] > 0).item()}')

Pr(beta > 0) = 0.9955

We get similar probability as with Gaussian obervation model.

Plot the data, the model fit, and the marginal posteriors for sigma and nu.

fig, ax = plt.subplots(1, 1, figsize=(12, 5))

az.plot_hdi(x, idata_lin_t.posterior['mu'], ax=ax)
ax.plot(x, idata_lin_t.posterior['mu'].median(["chain", "draw"]), color="C1")

ax.scatter(x, y, 10, color="C0")
ax.set_xlabel('year')
ax.set_ylabel('summer temp. @Kilpisjarvi')
ax.set_xlim((1951.5, 2013.5))


az.plot_posterior(fit_lin_t, var_names=['beta','sigma','ypred', 'nu'], round_to=2, kind='hist', bins=35, figsize=(12,3));

The posterior of the degree of freedom “nu” reveals that Gaussian model is likely to be ok, as Student’s t with nu > 20 is very close to Gaussian.

Pareto-smoothed importance-sampling leave-one-out cross-validation (PSIS-LOO)

We can use leave-one-out cross-validation to compare the expected predictive performance. For the following three lines to execute, the log-likelihood needs to be evaluated in the stan code. For an example, see: - stan file lin.stan - ArviZ API (for Python) - Computing approximate leave-one-out cross-validation usig PSIS-LOO (In R).

# Store inference data objects to dictionary
models = {"linear_model": idata_lin,
         "linear_model_t": idata_lin_t}

# Compare models with loo
az.compare(models)

	rank	elpd_loo	p_loo	elpd_diff	weight	se	dse	warning	scale
linear_model	0	-96.232467	2.591355	0.000000	1.000000e+00	4.973661	0.000000	False	log
linear_model_t	1	-96.715898	2.723147	0.483431	1.110223e-16	5.096291	0.422857	False	log

Loo values are similar for models. Hence, there is no practical difference between Gaussian and Student’s t observation model for this data.

Comparison of k groups with hierarchical models

In following, for each year, we analyze the monthly average temperatures for three summer months: June, July and August. In previous example, we used the average of these three summer months as estimate for yearly average summer temperarute. Now we would like to know, if there is a difference between three summer months: June, July and August.

data_path = os.path.abspath(
    os.path.join(
        os.path.pardir,
        'utilities_and_data',
        'kilpisjarvi-summer-temp.csv'
    )
)
# Load data
d = np.loadtxt(data_path, dtype=np.double, delimiter=';', skiprows=1)
# summer months are numbered from 1 to 3
x = np.tile(np.arange(1, 4), d.shape[0]) # create array with repeating numbers 1,2,3 for each year
y = d[:, 1:4].ravel()
N = len(x)
data = dict(
    N = N,
    K = 3,  # 3 groups
    x = x,  # group indicators
    y = y)  # observations

Common variance (ANOVA) model

with open('grp_aov.stan') as file:
    print(file.read())

// Comparison of k groups with common variance (ANOVA)
data {
  int<lower=0> N; // number of data points
  int<lower=0> K; // number of groups
  array[N] int<lower=1, upper=K> x; // group indicator
  vector[N] y; //
}
parameters {
  vector[K] mu;        // group means
  real<lower=0> sigma; // common std
}
model {
  y ~ normal(mu[x], sigma);
}

Fit the model

model = CmdStanModel(stan_file='grp_aov.stan')
fit = model.sample(data=data, seed=194838)
idata_aov = az.from_cmdstanpy(fit)

Check the n_eff and Rhat values for sampled model.

az.summary(idata_aov, round_to=2)

	mean	sd	hdi_2.5%	hdi_97.5%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
mu[0]	7.54	0.19	7.16	7.91	0.0	0.0	3877.60	3090.03	1.0
mu[1]	10.96	0.20	10.57	11.35	0.0	0.0	4554.89	2982.75	1.0
mu[2]	9.44	0.19	9.04	9.79	0.0	0.0	4786.56	2876.07	1.0
sigma	1.53	0.08	1.38	1.70	0.0	0.0	4315.71	3366.99	1.0

Check the treedepth, E-BFMI, and divergences

check_treedepth(idata_aov)
check_energy(idata_aov)
check_div(idata_aov)

Plot group mean distributions and matrix of probabilities that one mu is larger than other.

# matrix of probabilities that one mu is larger than other
mu = az.extract(idata_aov, var_names=["mu"]).T
ps = np.zeros((3, 3))
for k1 in range(3):
    for k2 in range(k1+1, 3):
        ps[k1, k2] = np.mean(mu[:, k1] > mu[:, k2])
        ps[k2, k1] = 1 - ps[k1, k2]
print("Matrix of probabilities that one mu is larger than other:")
print(ps.round(2))

months = ['June', 'July', 'August']
ax = az.plot_forest(fit,
                    kind='ridgeplot',
                    var_names=['mu'],
                    combined=True,
                    ridgeplot_overlap=1,
                    r_hat=False,
                    ess=False,
                    figsize=(12, 5))
ax[0].set_yticklabels(months[::-1]) # Order of months if reversed, since arviz flips the order of ticks in plot
ax[0].set_xlabel('Average temperature');

Matrix of probabilities that one mu is larger than other:
[[0. 0. 0.]
 [1. 0. 1.]
 [1. 0. 0.]]

Common variance and hierarchical prior for mean.

Results do not differ much from the previous, because there is only few groups and quite much data per group, but this works as an example of a hierarchical model.

with open('grp_prior_mean.stan') as file:
    print(file.read())

// Comparison of k groups with common variance and
// hierarchical prior for the mean
data {
    int<lower=0> N; // number of data points
    int<lower=0> K; // number of groups
    array[N] int<lower=1, upper=K> x; // group indicator
    vector[N] y; //
}
parameters {
    real mu0;             // prior mean
    real<lower=0> sigma0; // prior std
    vector[K] mu;         // group means
    real<lower=0> sigma;  // common std
}
model {
  mu0 ~ normal(10,10);      // weakly informative prior
  sigma0 ~ cauchy(0,4);     // weakly informative prior
  mu ~ normal(mu0, sigma0); // population prior with unknown parameters
  sigma ~ cauchy(0,4);      // weakly informative prior
  y ~ normal(mu[x], sigma);
}

Fit the model

model = CmdStanModel(stan_file='grp_prior_mean.stan')
fit = model.sample(data=data, seed=194838)
idata_grp = az.from_cmdstanpy(fit)

Plot group mean distributions and matrix of probabilities that one mu is larger than other. The probabilities are practically either 1 or 0 since the distributions are not overlapping almost at all. This can be visually seen from ridge plot below.

# matrix of probabilities that one mu is larger than other
mu = az.extract(idata_grp, var_names=["mu"]).T
ps = np.zeros((3, 3))
for k1 in range(3):
    for k2 in range(k1+1, 3):
        ps[k1, k2] = np.mean(mu[:, k1] > mu[:, k2])
        ps[k2, k1] = 1 - ps[k1, k2]
print("Matrix of probabilities that one mu is larger than other:")
print(ps.round(2))

months = ['June', 'July', 'August']
ax = az.plot_forest(fit,
                    kind='ridgeplot',
                    var_names=['mu'],
                    combined=True,
                    ridgeplot_overlap=1,
                    r_hat=False,
                    ess=False,
                    figsize=(12, 5))
ax[0].set_yticklabels(months[::-1])
ax[0].set_xlabel('Average temperature');

Matrix of probabilities that one mu is larger than other:
[[0. 0. 0.]
 [1. 0. 1.]
 [1. 0. 0.]]

Unequal variance and hierarchical prior for mean and variance

Results do not differ much from the previous, because there is only few groups and quite much data per group, but this works as an example of a hierarchical model.

with open('grp_prior_mean_var.stan') as file:
    print(file.read())

// Comparison of k groups with unequal variance and
// hierarchical priors for the mean and the variance
data {
  int<lower=0> N; // number of data points
  int<lower=0> K; // number of groups
  array[N] int<lower=1, upper=K> x;  // group indicator
  vector[N] y; //
}
parameters {
  real mu0;                 // prior mean
  real<lower=0> musigma0;   // prior std
  vector[K] mu;             // group means
  real lsigma0;             // prior mean
  real<lower=0> lsigma0s;   // prior std
  vector<lower=0>[K] sigma; // group stds
}
model {
  mu0 ~ normal(10, 10);       // weakly informative prior
  musigma0 ~ cauchy(0,10);    // weakly informative prior
  mu ~ normal(mu0, musigma0); // population prior with unknown parameters
  lsigma0 ~ normal(0,1);      // weakly informative prior
  lsigma0s ~ normal(0,1);     // weakly informative prior
  sigma ~ cauchy(lsigma0, lsigma0s); // population prior with unknown parameters
  y ~ normal(mu[x], sigma[x]);
}

Fit the model

model = CmdStanModel(stan_file='grp_prior_mean_var.stan')
fit = model.sample(data=data, seed=194838)
idata_grp_mv = az.from_cmdstanpy(fit)

Plot group mean distributions and matrix of probabilities that one mu is larger than other:

samples = az.extract(idata_grp_mv, var_names=["mu0", "mu", "sigma"])
print(f"std(mu0): {np.std(samples['mu0']).item():.2f}")
mu = samples['mu'].T

# matrix of probabilities that one mu is larger than other
ps = np.zeros((3, 3))
for k1 in range(3):
    for k2 in range(k1+1, 3):
        ps[k1, k2] = np.mean(mu[:, k1] > mu[:, k2])
        ps[k2, k1] = 1 - ps[k1, k2]
print("Matrix of probabilities that one mu is larger than other:")
print(ps)

months = ['June', 'July', 'August']
ax = az.plot_forest(fit,
                        kind='ridgeplot',
                        var_names=['mu'],
                        combined=True,
                        ridgeplot_overlap=1,
                        r_hat=False,
                        ess=False,
                        figsize=(12, 5))
ax[0].set_yticklabels(months[::-1])
ax[0].set_xlabel('Average temperature');

std(mu0): 2.43
Matrix of probabilities that one mu is larger than other:
[[0. 0. 0.]
 [1. 0. 1.]
 [1. 0. 0.]]