Hierarchical Bayesian Models

A look at hierarchical models, a kind of model that allows you to share information between different clusters of data. We'll start with a simple example and then apply it to a real-world situation.

Recently I’ve been learning about multilevel models. These are a kind of model where, by clustering the data in some way, you can ‘share’ information between different pieces of the model. That makes them an excellent way to get the most out of the often limited data that you might have available. Another advantage of these models is that it provides a way to account for additional uncertainty, in much the same way that a mixture model does. Let’s go through and see how to construct such a model (specifically, a varying-intercepts model) and cap it off with a real-world application!

Example - Heights

Let’s start off with a relatively simple example. We’re going to look at the simulated heights of people in different countries. In this situation, there is a natural hierarchical nature to the data: each height belongs to a person and each student belongs to a country. We might naturally expect that the country that a person is part of will have some influence on their score. There are a few ways that we could approach this:

1. Ignore the countries We could just ignore the fact that the people are in different countries and group them all together when making our model. This approach is called complete pooling.
2. Fit separate models Another approach would be to fit a separate model to each country completely independently. This is the no pooling approach.

At least in this case, both of those seem unsatisfactory. It seems very likely that there is some similarity between the countries - if you learn the mean for one country, that might influence your belief about the next country you visit. So how can we incorporate that information? Essentially, we’re going to fit a model for each country, but each one will share a common prior. This prior will be influenced by the data, and since it is shared between each country that will allow the information from one country to influence the model for the others.

This approach, where we share some information between the different clusters (countries, in this case) is called partial pooling. In the case where there is a hierarchical structure to the data, we call it a hierarchical model.

Let’s try this on some simulated data to see how it works before we apply it to some real data on diploma exam results in Alberta.

Simulation

First let’s set up the situation. For simplicity, we’ll assume that each person’s height is drawn from a normal distribution, and that the mean of that distribution is specific to each country. However, we’ll also assume that the means for the different countries are drawn from a common distribution as well. This will show up in the model as a shared prior.

Let’s use this process to generate some data, and then we’ll use the same model to recover it.

# Set up the environment
library(ggplot2)
library(cmdstanr)
library(posterior)
library(readxl)

options(repr.plot.width = 17, repr.plot.height = 8)

set.seed(500)

# country parameters
NUM_COUNTRIES <- 4
country_means <- rnorm(NUM_COUNTRIES, 176, 10)
country_sigma <- rexp(1, 1)

# now generate some people
NUM_PEOPLE <- 20
country_data <- data.frame(
    country_index = as.factor(rep(1:NUM_COUNTRIES, each = NUM_PEOPLE)),
    height = rnorm(
        n = NUM_COUNTRIES * NUM_PEOPLE,
        mean = rep(country_means, each = NUM_PEOPLE),
        sd = 3
    )
)

ggplot(country_data, aes(height)) +
    geom_density(aes(group = country_index, colour = country_index)) +
    geom_vline(data = data.frame(country = as.factor(1:NUM_COUNTRIES), mean = country_means), mapping = aes(xintercept = mean, colour = country)) +
    labs(x = "Height (cm)", y = "Density", colour = "Country") +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major.y = element_line(size = 0.5, linetype = 'dashed', colour = "grey")
    )

To recover the values, let’s use the following model.

The way that we’re going to share the information between the countries is through the $\bar\alpha$ parameter. Because this mean value for the prior for each country must be the same for each of the different countries, it provides a way for that data to flow between them as the model is fit.

Now that we have the data and know the generative process behind it, let’s ensure that we can recover those same values using the model.

stan_model_code <- "
data {
    int<lower = 0> N_COUNTRIES; // number of countries
    int<lower = 0> N; // number of people
    vector[N] heights;
    array[N] int<lower = 0, upper = N_COUNTRIES> country_id; // the country id corresponding to each height
}
parameters {
    vector[N_COUNTRIES] mu_country;
    real<lower = 1e-3> sigma; // sd for the heights
    real<lower = 1e-3> sigma_country; // sd for the country mean values
    real alpha_bar;
}
model {
    heights ~ normal(mu_country[country_id], sigma);
    mu_country ~ normal(alpha_bar, 10);
    alpha_bar ~ normal(176, 1);
    sigma_country ~ exponential(1);
    sigma ~ exponential(1);
}
"
stan_model_file <- write_stan_file(stan_model_code)
stan_model <- cmdstan_model(stan_model_file)

stan_model_data <- list(
    N_COUNTRIES = NUM_COUNTRIES,
    N = nrow(country_data),
    heights = country_data$height,
    country_id = as.integer(country_data$country_index)
)
fit <- stan_model$sample(
    data = stan_model_data,
    chains = 4
)

Running MCMC with 4 sequential chains...

All 4 chains finished successfully.
Mean chain execution time: 0.0 seconds.
Total execution time: 0.5 seconds.

fit$summary()

From the summary table, we can see that our posterior values are close to the true values. Let’s see the same thing graphically.

height_draws <- as_draws_matrix(fit)
head(height_draws)

country_mean_results <- height_draws[, grep('mu_country', colnames(height_draws))]
results_df <- data.frame(
    country_id = 1:NUM_COUNTRIES,
    mean = apply(country_mean_results, 2, mean),
    lower = apply(country_mean_results, 2, function(row) quantile(row, 0.025)),
    upper = apply(country_mean_results, 2, function(row) quantile(row, 0.975)),
    type = "Posterior"
)

posterior_alpha_bar <- mean(height_draws[, 'alpha_bar'])

ggplot(results_df, aes(country_id, mean)) +
    geom_point(aes(colour = type)) +
    geom_errorbar(aes(ymin = lower, ymax = upper)) +
    geom_point(data = data.frame(country_id = 1:NUM_COUNTRIES, mean = country_means, type = "Actual"), mapping = aes(colour = type)) +
    geom_hline(mapping = aes(yintercept = posterior_alpha_bar), linetype = 'dashed') +
    labs(x = "Country", y = "Height (cm)", colour = "Type") +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major.y = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

One thing to note is that we have a larger error than we would expect if we fit each country individually. The reason for that is that a hierarchical model like this tends to ‘shrink’ estimates towards the mean. You can see this here by the fact that the values of the posterior tend to be closer to the mean than the actual values. While it may seem crazy to deliberately choose a model that will tend to do worse on the data, in fact this is just part of the overfitting / underfitting tradeoff - by biasing estimates to the mean, we will do worse on our training set but better on outside data.

So now that we’ve seen a relatively simple example, let’s try to build our way towards our eventual goal of applying this idea to school test scores.

Example - School Test Scores

Imagine that you have a number of different schools in a school division, and you’re interested in how the students in those schools are doing on a test. We believe that each school has an effect on the students, and we also believe that the school division has an effect on the school. Here we have two nested pools - each student belongs to a school, and each school belongs to a district.

Simulation

Imagine that you have nine schools taking part in some sort of standardized exam. Each school has some number of students that are taking the exam, and they each get a grade. The grade they get is a combination of luck, their own properties, and some influence of the school. The influence of the school is itself dependent on the school district. Let’s try to generate some data, then create a model to recover the original values. By doing so, we’re ensuring that our model is fit for purpose.

Since this model is a bit more complex, let’s take it one step at a time.

First, let’s start with a single school. The distribution for the school’s students will be drawn from a beta distribution. The beta is the perfect distribution for this since it’s a distribution of probabilities, which is what we’re looking for. The beta distribution has two parameters, $\alpha$ and $\beta$. The mean is $\alpha / (\alpha + \beta)$, and the variance is $\alpha\beta / (\alpha + \beta)^2 (\alpha + \beta + 1)$. If we’re thinking of the beta distribution as being related to a Binomial random process (i.e. a set of coin flips), then we can think of $\alpha$ as the number of successes and $\beta$ as the number of failures. This is not quite correct, but it’s pretty close.

So let’s say that we want the school average to be about 75%. There are a few different values of $\alpha$ and $\beta$ that we could choose; we just need the mean $\alpha / (\alpha + \beta) = 0.75$. For instance, $\alpha = 3$ and $\beta = 1$ would work. So would any multiple of the those: 6 and 2, 4 and $\frac{4}{3}$, &c.

alpha_base <- 3
beta_base <- 1

multiples <- seq.int(from = 1, to = 4)
plot_df <- data.frame(p = numeric(), density = numeric(), type = character())

probs <- seq(from = 0, to = 1, length.out = 100)
for (multiple in multiples) {
    alpha <- alpha_base * multiple
    beta <- beta_base * multiple
    data <- data.frame(p = probs, density = dbeta(probs, alpha, beta), type = paste("alpha = ", alpha, " beta = ", beta))
    plot_df <- rbind(plot_df, data)
}

ggplot(plot_df, aes(p, density, group = type, colour = type)) +
    geom_line() +
    labs(x = "p", y = "Density", colour = "Parameter Values") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

Note that although the mean for each of these is at $0.75$, that’s not the same as having the peak value be there.

If we’re thinking forward to making our model, it’s not clear to me exactly what we should be thinking of the parameters $\alpha$ and $\beta$ as representing. Intuitively, I’d prefer to think in terms of the mean and variance (or even more generally, the centre and the spread) of the expected data. Luckily, it turns out that there is a reparameterization of the beta in terms of these:

Solving for $\alpha$ and $\beta$,

Let’s recreate the above using this parameterization to convince ourselves that it works. Since we’re fixing our mean, the spread parameter is the only one that needs to change.

mu <- 0.75
spread <- 4 # 'sample size'

plot_df <- data.frame(p = numeric(), density = numeric(), type = character())

probs <- seq(from = 0, to = 1, length.out = 100)
for (sample_size in seq.int(from = 4, to = 16, by = 4)) {
    alpha <- mu * sample_size
    beta <- (1 - mu) * sample_size
    data <- data.frame(p = probs, density = dbeta(probs, alpha, beta), type = paste("mu = ", mu, " nu = ", sample_size))
    plot_df <- rbind(plot_df, data)
}

ggplot(plot_df, aes(p, density, group = type, colour = type)) +
    geom_line() +
    labs(x = "p", y = "Density", colour = "Parameter Values") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

Great! So now we have our parameterization. Let’s make sure that we can model at least this part of the problem. Let’s say that we have a school where the average test score is 75%, and let’s simulate a class of 20 students.

NUM_STUDENTS <- 20
school_mu <- 0.75
school_nu <- 10 # this is *not* the same as the number of students!

alpha <- school_mu * school_nu
beta <- (1 - school_mu) * school_nu

actual_student_scores <- rbeta(NUM_STUDENTS, alpha, beta)

p <- seq(from = 0, to = 1, length.out = 100)
line_data <- data.frame(p = p, density = dbeta(p, alpha, beta))
student_data <- data.frame(p = actual_student_scores, density = dbeta(actual_student_scores, alpha, beta))
ggplot() +
    geom_line(data = line_data, mapping = aes(p, density)) +
    geom_point(data = student_data, mapping = aes(p, density)) +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

So here we can see the theoretical distribution of the students’ scores along with the actual scores. Now let’s use a model to recover those values!

How did we arrive at those priors? We ran some simulations of the model with different parameters to find ones that were reasonable. In a real model we would do this before looking at the data, but in this case that’s not really possible. We’ll just pretend that we don’t know what the data should look like. Also, when checking a model like this, it’s not unreasonable to start with a model that reflects the true underlying process - the idea here is that we are ground-truthing it. If the model can’t get back the correct values even when we know the model is correct and the values are close, then it really isn’t fit for purpose.

NUM_SIM_STUDENTS <- 100
nu <- rnorm(NUM_SIM_STUDENTS, 10, 1)
mu <- rbeta(NUM_SIM_STUDENTS, 12, 4)
alpha <- mu * nu
beta <- (1 - mu) * nu
s <- rbeta(NUM_SIM_STUDENTS, alpha, beta)

plot_df <- data.frame(s = s)
ggplot(plot_df, aes(s)) +
    geom_density() +
    labs(x = "Score", y = "Density") +
    coord_cartesian(xlim = c(0, 1)) +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

ggplot(data.frame(nu = nu), aes(nu)) +
    geom_density() +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    labs(x = "Score", y = "Density") +
    theme_bw() +
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

Before we do anything else, let’s do a prior predictive check. The idea here is to see what kind of values our model is producing before we feed it any data. In large part this should help to constrain our priors - are the data that we’re seeing come out of the model in line with our expectations? Are they plausible? If not, we’ll have to go back and revise our priors. We can actually use Stan to help us with this.

student_score_model_pp_code <- "
data {
    int<lower = 0> N; // number of students
}
parameters {
    real<lower = 0, upper = 1> mu;
    real<lower = 0> nu;
}
 transformed parameters {
    real<lower = 0> a = mu * nu;
    real<lower = 0> b = (1 - mu) * nu;
}
model {
    mu ~ beta(12, 4);
    nu ~ normal(10, 1);
}
generated quantities {
    vector[N] s_sim;
    for (n in 1:N) {
        s_sim[n] = beta_rng(a, b);
    }
}
"

student_score_model_pp_file <- write_stan_file(student_score_model_pp_code)
student_score_pp_model <- cmdstan_model(student_score_model_pp_file)

student_score_model_pp_data <- list(
    N = 100
)

student_score_pp_fit <- student_score_pp_model$sample(
    data = student_score_model_pp_data,
    chains = 4
)
student_score_pp_fit$summary()

Running MCMC with 4 sequential chains...

All 4 chains finished successfully.
Mean chain execution time: 0.0 seconds.
Total execution time: 0.5 seconds.

From the above summary the values of the parameters mu and nu look correct, as do the values for the simulated student scores. The rhat for each parameter also looks good (close to 1) and the effective sample size (ess_bulk) is also good. There was a minor problem in the fitting where a nan value was generated, but that’s a rare thing and only happens intermittently. Overall, this looks good! Now let’s get the data.

draws <- as_draws_matrix(student_score_pp_fit)
head(draws)

Now that we have the draws, let’s plot a few different classes of 100 students (s_sim) to see what kind of variation the model is producing.

set.seed(2024)
s_sim_draws <- draws[, grep('s_sim', colnames(draws))]
s_sim_means <- mean(apply(s_sim_draws, 2, mean))

p <- ggplot() +
    coord_cartesian(xlim = c(0, 1)) +
    geom_vline(data = data.frame(xintercept = s_sim_means), mapping = aes(xintercept=xintercept)) +
    labs(x = "Score", y = "Density") +
    scale_y_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

for (row in sample(1:4000, size = 100)) {
    data <- data.frame(s = as.vector(s_sim_draws[row, ]))
    p <- p + geom_density(data = data, mapping = aes(s), colour = alpha('black', 0.1))
}

print(p)

What we see here are the results of 100 randomly chosen simulated classes. This seems roughly reasonable to me - the mean is firmly centred around 75% and there’s considerable variation. Looking at this, there is maybe a bit too much variation - I would expect almost all of the weight to be in the range 50% - 100%, and there are a few classes with a significant amount of weight below 50%. On the other hand, it’s probably not the worst thing in the world for this to have too much variation - the data should be able to fix that.

Now let’s actually train the model on some data and see what happens. We’ll have to slightly modify the model we used for our prior predictive check to incorporate the data, but otherwise it will look very similar.

student_score_model_code <- "
data {
    int<lower = 0> N; // number of students
    vector[N] s; // student scores
}
parameters {
    real<lower = 0, upper = 1> mu;
    real<lower = 0> nu;
}
 transformed parameters {
    real<lower = 0> a = mu * nu;
    real<lower = 0> b = (1 - mu) * nu;
}
model {
    s ~ beta(a, b);
    mu ~ beta(12, 4);
    nu ~ normal(10, 1);
}
"

student_score_model_file <- write_stan_file(student_score_model_code)
student_score_model <- cmdstan_model(student_score_model_file)

student_score_model_data <- list(
    N = NUM_STUDENTS,
    s = actual_student_scores
)
student_score_fit <- student_score_model$sample(
    data = student_score_model_data,
    chains = 4
)
student_score_fit$summary()

Running MCMC with 4 sequential chains...

All 4 chains finished successfully.
Mean chain execution time: 0.0 seconds.
Total execution time: 0.6 seconds.

draws <- as_draws_matrix(student_score_fit)

# mu
ggplot(data.frame(mu = draws[, 'mu'])) +
    geom_density(aes( mu )) +
    geom_vline(aes(xintercept = 0.75)) +
    coord_cartesian(xlim = c(0, 1)) +
    labs(x = "mu", y = "Density") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

# nu
ggplot(data.frame(nu = draws[, 'nu'])) +
    geom_density(aes(nu)) +
    geom_vline(aes(xintercept = 10)) +
    labs(x = "nu", y = "Density") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

Great! So this model seems to be doing a good job of recovering the actual values. Now that we have some confidence in this part of the model, let’s extend it a bit more to consider the situation of different schools, each with their own means, which now influence the scores of their students. We’ll assume that the schools each draw their means from a single beta distribution, and we’ll also assument that each student draws their score from a beta distribution whose mean is the school mean.

set.seed(1234)
NUM_SCHOOLS <- 5
school_mu <- 0.75
school_nu <- 50

school_alpha <- school_mu * school_nu
school_beta <- (1 - school_mu) * school_nu

school_means <- rbeta(NUM_SCHOOLS, school_alpha, school_beta)

ps <- seq(from = 0, to = 1, length.out = 100)
density <- dbeta(ps, school_alpha, school_beta)
ggplot() +
    geom_line(data = data.frame(p = ps, density = density), mapping = aes(p, density)) +
    geom_point(data = data.frame(p = school_means, density = dbeta(school_means, school_alpha, school_beta)), mapping = aes(p, density)) +
    labs(x = "School-Specific Mean", y = "Density") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

Now we’ll simulate some students for each of these schools.

STUDENT_NU <- 30

student_alpha <- school_means * STUDENT_NU
student_beta <- (1 - school_means) * STUDENT_NU

ps <- seq(from = 0, to = 1, length.out = 100)
plot_df <- data.frame(p = numeric(), density = numeric(), school = integer())
for (school_id in 1:NUM_SCHOOLS) {
    alpha <- student_alpha[school_id]
    beta <- student_beta[school_id]
    school_data <- data.frame(
        p = ps,
        density = dbeta(ps, alpha, beta),
        school = school_id
    )
    plot_df <- rbind(plot_df, school_data)
}
plot_df$school <- as.factor(plot_df$school)

ggplot(plot_df, aes(p, density, group = school, colour = school)) +
    geom_line() +
    labs(x = "p", y = 'Density', colour = "School") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

Here we can see the distribution of scores from which each student will draw their scores.

set.seed(1234)
STUDENTS_PER_SCHOOL <- 20
student_scores <- data.frame(s = numeric(), school = integer())

for (school in 1:NUM_SCHOOLS) {
    alpha <- student_alpha[school]
    beta <- student_beta[school]
    school_data <- data.frame(
        s = rbeta(STUDENTS_PER_SCHOOL, alpha, beta),
        school = school
    )
    student_scores <- rbind(student_scores, school_data)
}
student_scores$school <- as.factor(student_scores$school)

ggplot(student_scores, aes(s, colour = school, group = school)) +
    geom_density() +
    coord_cartesian(xlim = c(0, 1)) +
    labs(x = 'Score', y = 'Density', colour = "School") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

plot_df <- data.frame(true_mean = numeric(), empirical_mean = numeric())
for (school_id in 1:NUM_SCHOOLS) {
    school_data <- student_scores[student_scores$school == school_id, ]
    plot_df <- rbind(plot_df, data.frame(true_mean = school_means[school_id], empirical_mean = mean(school_data$s)))
}

empirical_school_means <- plot_df$empirical_mean

ggplot(plot_df, aes(true_mean, empirical_mean)) +
    geom_point() +
    geom_abline(intercept = 0, slope = 1, linetype = 'dashed') +
    labs(x = "True Mean", y = "Empirical Mean") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

First, let’s run an unpooled model - that is, one where we estimate each school’s values separately. That model looks like

student_school_unpooled_model_code <- "
data {
    int<lower = 0> N_students; // number of students
    int<lower = 0> N_schools; // number of schools
    vector[N_students] s; // student scores
    array[N_students] int school_id; // the school id for each student
}
parameters {
    array[N_schools] real<lower = 0, upper = 1> mu_school;
    array[N_schools] real<lower = 0> nu_school;
}
 transformed parameters {
    array[N_schools] real<lower = 0> a_school;
    array[N_schools] real<lower = 0> b_school;
    for (i in 1:N_schools) {
        a_school[i] = mu_school[i] * nu_school[i];
        b_school[i] = (1 - mu_school[i]) * nu_school[i];
    }
}
model {
    for (student_index in 1:N_students) {
        s[student_index] ~ beta(a_school[school_id[student_index]], b_school[school_id[student_index]]);
    }
    mu_school ~ beta(3, 1);
    nu_school ~ normal(30, 4);
}
"

student_school_unpooled_model_file <- write_stan_file(student_school_unpooled_model_code)
student_school_unpooled_model <- cmdstan_model(student_school_unpooled_model_file)

student_school_model_data <- list(
    N_students = nrow(student_scores),
    N_schools = NUM_SCHOOLS,
    s = student_scores$s,
    school_id = as.integer(student_scores$school)
)
student_school_unpooled_fit <- student_school_unpooled_model$sample(
    data = student_school_model_data,
    chains = 4
)
student_school_unpooled_fit$summary()

Running MCMC with 4 sequential chains...

All 4 chains finished successfully.
Mean chain execution time: 0.2 seconds.
Total execution time: 0.9 seconds.

unpooled_student_school_draws <- as_draws_matrix(student_school_unpooled_fit$draws())

# school means
unpooled_school_means_draws <- unpooled_student_school_draws[, grep('mu_school', colnames(unpooled_student_school_draws))]
unpooled_school_mean_means <- apply(unpooled_school_means_draws, 2, mean)
unpooled_school_mean_lower <- apply(unpooled_school_means_draws, 2, function(col) quantile(col, 0.025))
unpooled_school_mean_upper <- apply(unpooled_school_means_draws, 2, function(col) quantile(col, 0.975))

unpooled_means_plot_df <- data.frame(school = 1:NUM_SCHOOLS, mean = unpooled_school_mean_means, lower = unpooled_school_mean_lower, upper = unpooled_school_mean_upper)

unpooled_school_means_plot <- ggplot() +
    geom_errorbar(data = unpooled_means_plot_df, mapping = aes(school, ymin = lower, ymax = upper)) +
    geom_point(data = unpooled_means_plot_df, aes(school, mean)) +
    geom_point(data = data.frame(school = 1:NUM_SCHOOLS, mean = school_means), mapping = aes(school, mean, colour = "True Mean")) +
    geom_point(data = data.frame(school = 1:NUM_SCHOOLS, mean = empirical_school_means), mapping = aes(school, mean, colour = 'Empirical Mean')) +
    labs(x = "School", y = "Mean Score", colour = "Type") +
    scale_y_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )
print(unpooled_school_means_plot)

This looks pretty good! We’re doing a good job of recovering these values. Now let’s share information among the schools by using a hierarchical model.

We’ll use the following model to recover the values.

student_school_model_code <- "
data {
    int<lower = 0> N_students; // number of students
    int<lower = 0> N_schools; // number of schools
    vector[N_students] s; // student scores
    array[N_students] int school_id; // the school id for each student
}
parameters {
    array[N_schools] real<lower = 0, upper = 1> mu_school;
    array[N_schools] real<lower = 0> nu_school_raw; // non-centred parameterization; implies nu_school ~ normal(nu_bar, sigma_bar)
    real<lower = 0> alpha_bar;
    real<lower = 0> beta_bar;
    real<lower = 0> nu_bar;
    real<lower = 0> sigma_nu;
}
 transformed parameters {
    array[N_schools] real<lower = 0> nu_school;
    array[N_schools] real<lower = 0> a_school;
    array[N_schools] real<lower = 0> b_school;
    for (i in 1:N_schools) {
        nu_school[i] = nu_bar + sigma_nu * nu_school_raw[i];
        a_school[i] = mu_school[i] * nu_school[i];
        b_school[i] = (1 - mu_school[i]) * nu_school[i];
    }
}
model {
    for (student_index in 1:N_students) {
        s[student_index] ~ beta(a_school[school_id[student_index]], b_school[school_id[student_index]]);
    }
    mu_school ~ beta(alpha_bar, beta_bar);
    nu_school_raw ~ std_normal();
    alpha_bar ~ normal(3, 0.5);
    beta_bar ~ normal(1, 0.3);
    nu_bar ~ normal(30, 5);
    sigma_nu ~ exponential(1);
}
"

student_school_model_file <- write_stan_file(student_school_model_code)
student_school_model <- cmdstan_model(student_school_model_file)

student_school_fit <- student_school_model$sample(
    data = student_school_model_data,
    chains = 4,
    iter_warmup = 1000,
    adapt_delta = 0.95
)
student_school_fit$summary()

Running MCMC with 4 sequential chains...

All 4 chains finished successfully.
Mean chain execution time: 0.4 seconds.
Total execution time: 1.9 seconds.

student_school_draws <- as_draws_matrix(student_school_fit$draws())

# # school means
school_means_draws <- student_school_draws[, grep('mu_school', colnames(student_school_draws))]
school_mean_means <- apply(school_means_draws, 2, mean)
school_mean_lower <- apply(school_means_draws, 2, function(col) quantile(col, 0.025))
school_mean_upper <- apply(school_means_draws, 2, function(col) quantile(col, 0.975))

means_plot_df <- data.frame(school = 1:NUM_SCHOOLS, mean = school_mean_means, lower = school_mean_lower, upper = school_mean_upper)

prior_mean <- mean(student_school_draws[, 'alpha_bar'] / (student_school_draws[, 'alpha_bar'] + student_school_draws[, 'beta_bar']))

pooled_school_means_plot <- ggplot() +
    geom_errorbar(data = means_plot_df, mapping = aes(school, ymin = lower, ymax = upper)) +
    geom_point(data = means_plot_df, aes(school, mean)) +
    geom_point(data = data.frame(school = 1:NUM_SCHOOLS, mean = school_means), mapping = aes(school, mean, colour = 'True Mean')) +
    geom_point(data = data.frame(school = 1:NUM_SCHOOLS, mean = unpooled_school_mean_means), aes(school, mean, colour = "Unpooled Mean")) +
    geom_hline(mapping = aes(yintercept = prior_mean), linetype = 'dashed') +
    labs(x = "School", y = "Value", colour = "Mean Type") +
    scale_y_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )
print(pooled_school_means_plot)

So as we can see, pooling didn’t have too much of an effect; there’s a minimal change from the pooled mean (black) and the unpooled mean (blue).

Now we’re ready to tackle a more complex hierarchical model with multiple layers! To start off, let’s say we have some districts. Each district will have some schools, and each school will have some students. Each of the students will have a score that is the school mean plus some standard deviation, drawn from a beta distribution. In turn, each of the school means for a district will be drawn from the same distribution. Finally, each of the district means will be drawn from a different shared distribution.

set.seed(1234)

NUM_DISTRICTS <- 3
SCHOOLS_PER_DISTRICT <- 3
TOTAL_SCHOOLS <- NUM_DISTRICTS * SCHOOLS_PER_DISTRICT
STUDENTS_PER_SCHOOL <- 20
TOTAL_STUDENTS <- TOTAL_SCHOOLS * STUDENTS_PER_SCHOOL

district_mu <- 0.75
district_nu <- 30
district_alpha <- district_mu * district_nu
district_beta <- (1 - district_mu) * district_nu

district_means <- rbeta(NUM_DISTRICTS, district_alpha, district_beta)
ps <- seq(from = 0, to = 1, length.out = 100)
district_distribution_df <- data.frame(
    p = ps,
    density = dbeta(ps, district_alpha, district_beta)
)
districts_df <- data.frame(
    p = district_means,
    density = dbeta(district_means, district_alpha, district_beta)
)
ggplot(NULL, aes(p, density)) +
    geom_line(data = district_distribution_df) +
    geom_point(data = districts_df) +
    labs(x = "p", y = "Density") +
    scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +
    theme(
        panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
    )

So now we have the district means coming from a single distribution. Within that, each district will have schools, each of whose means comes from a distribution whose mean is the district mean!

set.seed(2024)
ps <- seq(from = 0, to = 1, length.out = 1e3)
school_params <- data.frame(district = integer(), school_id = integer(), school_mu = numeric())
district_nu <- 40
for (district_id in 1:NUM_DISTRICTS) {
    district_mu <- district_means[district_id]
    district_alpha <- district_mu * district_nu
    district_beta <- (1 - district_mu) * district_nu
    school_data <- data.frame(
        district = rep(district_id, SCHOOLS_PER_DISTRICT),
        school_id = district_id + 1:SCHOOLS_PER_DISTRICT,
        school_mu = rbeta(SCHOOLS_PER_DISTRICT, district_alpha, district_beta)
    )
    school_params <- rbind(school_params, school_data)

    district_distribution_df <- data.frame(p = ps, density = dbeta(ps, district_alpha, district_beta))
    school_means_df <- data.frame(p = school_data$school_mu, density = dbeta(school_data$school_mu, district_alpha, district_beta))
    school_plot <- ggplot(NULL, aes(p, density)) +
        geom_line(data = district_distribution_df) +
        geom_point(data = school_means_df) +
        labs(x = "p", y = "Density", title = paste("Distribution and school means for district ", district_id)) +
        scale_x_continuous(labels = scales::percent_format(accuracy = 1)) +
        theme_bw() +
        theme(
            panel.grid.major = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
        )
    print(school_plot)
}

Now that we have the school information, we can generate a simulated group of students for each school. Each student will draw their score from a distribution whose mean is the school mean.

set.seed(2024)
ps <- seq(from = 0, to = 1, length.out = 1e3)
student_data <- data.frame(district_id = integer(), school_id = integer(), score = numeric())
student_nu <- 30
for (school_id in 1:TOTAL_SCHOOLS) {
    district_id <- school_params[school_id, 'district']
    school_mu <- school_params[school_id, 'school_mu']
    student_alpha <- school_mu * student_nu
    student_beta <- (1 - school_mu) * student_nu
    school_scores <- data.frame(
        district_id = rep(district_id, STUDENTS_PER_SCHOOL),
        school_id <- rep(school_id, STUDENTS_PER_SCHOOL),
        score = rbeta(STUDENTS_PER_SCHOOL, student_alpha, student_beta)
    )
    student_data <- rbind(student_data, school_scores)
}
student_data$student_id <- as.factor(seq_len(nrow(student_data)))
student_data$school_id <- as.factor(student_data$school_id)
student_data$district_id <- as.factor(student_data$district_id)

student_plot <- ggplot(student_data, aes(student_id, score)) +
    geom_point(aes(colour = school_id, shape = district_id)) +
    labs(x = "Student", y = "Score", colour = "School", shape = "District") +
    scale_y_continuous(labels = scales::percent_format(accuracy = 1)) +
    theme_bw() +  # Set background to white
    theme(
        panel.grid.major.x = element_blank(),
        panel.grid.minor.x = element_blank(),
        panel.grid.major.y = element_line(size = 0.5, linetype = 'dashed', colour = "grey"),
        axis.ticks.x = element_blank(),
        axis.text.x = element_blank(),
    )
for (district_id in 1:(NUM_DISTRICTS - 1)) {
    print(district_id)
    student_plot <- student_plot +
        geom_vline(xintercept = SCHOOLS_PER_DISTRICT * STUDENTS_PER_SCHOOL * district_id + 0.5)
}
print(student_plot)

[1] 1
[1] 2

Great! Now that we have the data and a good handle on the generation process, let’s look what this looks like as a model. Again, there are a few levels here. There will be a shared district-level set of parameters as well as a shared school-level set of parameters influenced by the district. The student will in turn have their own parameters, again inheriting from the school parameters. There’s a lot here, but let’s see what it looks like in Stan!

# NB: this is the non-centred parameterization to avoid divergent transitions
student_school_district_model_code <- "
data {
    int<lower = 0> N_students; // number of students
    int<lower = 0> N_schools; // number of schools
    int<lower = 0> N_districts; // number of districts
    vector[N_students] s; // student scores
    array[N_students] int school_id; // the school id for each student
    array[N_schools] int district_id; // the district for each school
}
parameters {
    array[N_districts] real<lower = 0, upper = 1> mu_district;
    array[N_districts] real<lower = 0> nu_district_raw; // standard normal -> implies nu_school ~ normal(nu_bar, sigma_nu)
    array[N_schools] real<lower = 0, upper = 1> mu_school;
    array[N_schools] real<lower = 0> nu_school;
    real<lower = 0> alpha_bar;
    real<lower = 0> beta_bar;
    real<lower = 0> nu_bar;
    real<lower = 0> sigma_nu;
}
 transformed parameters {
    // array[N_schools] real<lower = 0> nu_school;
    array[N_schools] real<lower = 0> a_school;
    array[N_schools] real<lower = 0> b_school;
    array[N_districts] real<lower = 0> a_district;
    array[N_districts] real <lower = 0> b_district;
    array[N_districts] real <lower = 0> nu_district;
    for (i in 1:N_schools) {
        // nu_school[i] = nu_bar + sigma_nu * nu_school_raw[i];
        a_school[i] = mu_school[i] * nu_school[i];
        b_school[i] = (1 - mu_school[i]) * nu_school[i];
    }
    for (district_index in 1:N_districts) {
        nu_district[district_index] = nu_bar + sigma_nu * nu_district_raw[district_index];
        a_district[district_index] = mu_district[district_index] * nu_district[district_index];
        b_district[district_index] = (1 - mu_district[district_index]) * nu_district[district_index];
    }
}
model {
    for (student_index in 1:N_students) {
        s[student_index] ~ beta(a_school[school_id[student_index]], b_school[school_id[student_index]]);
    }
    for (school_index in 1:N_schools) {
        mu_school[school_index] ~ beta(a_district[district_id[school_index]], b_district[district_id[school_index]]);
    }
    mu_district ~ beta(alpha_bar, beta_bar);
    nu_district_raw ~ std_normal();

    // hyperpriors
    alpha_bar ~ normal(3, 0.5);
    beta_bar ~ normal(1, 0.3);
    nu_bar ~ normal(40, 5);
    sigma_nu ~ exponential(1);
}
"

student_school_district_model_file <- write_stan_file(student_school_district_model_code)
student_school_district_model <- cmdstan_model(student_school_district_model_file)

student_school_district_model_data <- list(
    N_students = nrow(student_data),
    N_schools = length(unique(student_data$school_id)),
    N_districts = length(unique(student_data$district_id)),
    s = student_data$score,
    school_id = student_data$school_id,
    district_id = school_params$district
)

student_school_district_fit <- student_school_district_model$sample(
    data = student_school_district_model_data,
    chains = 4,
    iter_warmup = 1000,
    adapt_delta = 0.95
)
student_school_district_fit$summary()

Running MCMC with 4 sequential chains...

All 4 chains finished successfully.
Mean chain execution time: 0.9 seconds.
Total execution time: 3.9 seconds.