Prior Predictive Checks

First, let’s see what different prior distributions look like. We’ll use conventional normal priors on \(\beta_0\) and \(\beta_1\). We’ll try out three different normal distributions.

We create a grid of evenly spaced values. We then want to see the likelihood of these xs in the context of three normal distributions with the same mean but different standard deviations. We use the map_dfr function from the purrr package to iterate over the three different distributions and store the result in a data frame.

grid <- seq(-3, 3, length.out = 1000) # evenly spaced values from -3 to 3

b0_prior <- map_dfr(.x = c(0.5, 1, 2), ~ data.frame( # .x represents the three sigmas
  grid = grid,
  b0 = dnorm(grid, mean = 0, sd = .x)),
  .id = "sigma_id")

# create friendlier labels
b0_prior <- b0_prior %>%
  mutate(sigma_id = factor(sigma_id, labels = c("normal(0, 0.5)",
                                                "normal(0, 1)",
                                                "normal(0, 2)")))

ggplot(b0_prior, aes(x = grid, y = b0)) +
  geom_area(fill = "cadetblue4", color = "black", alpha = 0.90) +
  scale_x_continuous(expand = c(0, 0)) +
  scale_y_continuous(expand = c(0, 0), limits = c(0, 0.85)) +
  labs(x = NULL,
       y = "probability density",
       title = latex2exp::TeX("Possible $\\beta_0$ (intercept) priors")) +
  facet_wrap(~sigma_id, nrow = 3)

We can see that the three distributions vary in terms of how spread out they are. We might say that the top distribution with mean 0 and sd 0.5 is more “confident” about the probability of values close to zero, or that it is more skeptical about extreme values. Flatter distributions allocate probability more evenly and are therefore more open to extreme values.

This is all well and good, but looking at the raw probability densities doesn’t tell us much about what the priors assume about the data. We need to look at what kind of data is compatible with the priors. So this time we’ll generate regression lines from the \(\beta_0\) and \(\beta_1\) priors. Remember, all you need to create a line is an intercept and a slope!

The trick this time is to generate intercepts and slopes from the different normal distributions. Again, we’ll stick with three different distributions, but now we use map2_dfr to iterate over a \(\beta_0\) prior and a \(\beta_1\) prior.

b0b1 <- map2_df(.x = c(0.5, 1, 2), .y = c(0.25, 0.5, 1), ~ data.frame(
  b0 = rnorm(100, mean = 0, sd = .x), 
  b1 = rnorm(100, mean = 0, sd = .y)), .id = "sigma_id"
)

# create friendlier labels
b0b1 <- b0b1 %>%
  mutate(sigma_id = factor(sigma_id, labels = c("b0 ~ normal(0, 0.5); b1 ~ normal(0, 0.25)",
                                                "b0 ~ normal(0, 1); b1 ~ normal(0, 0.50)",
                                                "b0 ~ normal(0, 2); b1 ~ normal(0, 1)")))

ggplot(b0b1) +
  geom_abline(aes(intercept = b0, slope = b1), color = "cadetblue4", alpha = 0.75) +
  scale_x_continuous(limits = c(-2, 2)) +
  scale_y_continuous(limits = c(-3, 3)) +
  labs(x = "x",
       y = "y",
       title = latex2exp::TeX("Sampling of lines from $\\beta_0$ and $\\beta_1$ priors")) +
  facet_wrap(~sigma_id, nrow = 3)

The more restrictive set of priors on the top constrains the lines to have intercepts and slopes close to zero, and you can see that the lines vary little from one another. In contrast, the weaker priors allow a much greater variety of intercept/slope combinations.

Later we’ll see how to generate actual data from the priors, but for now let’s go ahead with the priors I’ve provided in the model formula and begin coding the model in Stan!

Data

The data block is where we define our observed variables. We also need to define lengths/dimensions of stuff, which can seem strange if you’re used to R or Python. We first declare an integer variable N to be the number of observations: int N; (note the use of semicolon to denote the end of a line). I’m also declaring an integer K, which is the number of predictors in our model. Note that I’m including the intercept in this count, so we end up with 2 predictors (2 columns in the model matrix).

Now we supply actual data. We write matrix[N, K] to tell Stan that x is a \(N \times K\) matrix. y is easier - just a vector of length \(N\).

// STAN CODE
data {
  int N; // number of observations
  int K; // number of predictors + intercept
  matrix[N, K] x; // x vector
  vector[N] y; // y vector
}

Parameters

Now we tell Stan the parameters in our model. These are the unobserved variables that we want to estimate. In this example, we have 3 parameters: \(\beta_0\), \(\beta_1\) and \(\sigma\). Like before, we first tell Stan the type of data this parameter will contain - in this case, \(\beta_0\) and \(\beta_1\) are contained in a vector of length \(K\) that we will sensibly call beta. \(\sigma\) will just be a single real value (“real” means that the number can have a decimal point). We use <lower=0> to constrain it to be positive because it is impossible to have negative standard deviation.

// STAN CODE
parameters {
  vector[K] beta; // intercept and slope parameters
  real<lower=0> sigma; // constrained to be positive
}

Model

The model block is where the action is at. We start with our priors. There should be a prior for each parameter declared above.

// STAN CODE
model {
  vector[N] mu; // declaring a mu vector
  
  // priors
  beta ~ normal(0, 1); // declares the same prior for intercept and slope
  sigma ~ exponential(1); // using an exponential prior on sigma
  
  // likelihood
  for(i in 1:N) {
    mu[i] = x[i] * beta; // * is matrix multiplication in this context, where x[i] is the first row of x
  }
  
  y ~ normal(mu, sigma);
}

The part above where we define \(\mu\) (mu[i]) can seem a bit strange. You may be wondering how the intercept and slope factor into this calculation. The compact notation hides some of the magic, which really isn’t very magical since it’s just a dot product. In matrix form, it looks like this:

\[\mu_i=\begin{bmatrix} 1&x_i\\ \end{bmatrix} \begin{bmatrix} \beta_0\\ \beta_1 \end{bmatrix}\]

where \(1\) is the intercept and \(x_i\) is the \(i^{th}\) value of the variable \(x\). Computing this dot product for each participant i gives you an expected value mu for each participant.

Full Stan Code

The full Stan code looks like this:

// STAN CODE
data {
  int N; // number of observations
  int K; // number of predictors + intercept
  matrix[N, K] x; // matrix of predictors
  vector[N] y; // y vector
}

parameters {
  vector[K] beta; // intercept and slope parameters
  real<lower=0> sigma; // constrained to be positive
}

model {
  vector[N] mu; // declaring a mu vector
  
  // priors
  beta ~ normal(0, 1); // declares the same prior for intercept and slope
  sigma ~ exponential(1); // using an exponential prior on sigma
  
  // likelihood
  for(i in 1:N) {
    mu[i] = x[i] * beta; // * is matrix multiplication in this context
  }
  
  y ~ normal(mu, sigma);
}

Saving this file will give it a .stan file extension. It’s a good idea to save it in your project directory. I’ll name this model "mod1.stan".

Running the Model

Now we go back to R to run the model. If you haven’t already, load up rstan.

library(rstan)

We need to put the data in a list for Stan. Everything that we declared in the data block of our Stan code should be entered into this list. We’ll also standardize the variables, so that they match up with our priors.

stan_dat <- list(
  N = nrow(dat), # using nrow() to get the number of rows (observations)
  K = 2, # (1) intercept, (2) slope
  x = matrix(c(rep(1, nrow(dat)), scale(dat$x)), ncol = 2), # model matrix (x standardized)
  y = as.vector(scale(dat$y)) # standardized outcome. convert to vector to avoid issues with scale function
)

Finally, we use the function stan to run the model. We set chains and cores to 4, which will allow us to run 4 Markov chains in parallel.

stan_fit1 <- stan(file = "mod1.stan",
                  data = stan_dat,
                  chains = 4, cores = 4)

Evaluating the Model

We can take a look at the parameters in the console by printing the model fit. We get posterior means, standard errors, and quantiles for each parameter. We also get things called n_eff and Rhat. These are indicators of how well Stan’s engine explored the parameter space (if this is cryptic, that’s ok), It’s enough for now to know that when Rhat is 1, things are good.

stan_fit1

## Inference for Stan model: mod1.
## 4 chains, each with iter=2000; warmup=1000; thin=1; 
## post-warmup draws per chain=1000, total post-warmup draws=4000.
## 
##            mean se_mean   sd    2.5%     25%     50%     75%   97.5% n_eff Rhat
## beta[1]    0.00    0.00 0.07   -0.13   -0.05    0.00    0.05    0.13  4608    1
## beta[2]    0.09    0.00 0.07   -0.05    0.04    0.09    0.14    0.22  4188    1
## sigma      1.00    0.00 0.05    0.91    0.97    1.00    1.03    1.10  3928    1
## lp__    -106.17    0.03 1.22 -109.31 -106.75 -105.88 -105.27 -104.74  2311    1
## 
## Samples were drawn using NUTS(diag_e) at Sun Mar 21 15:40:57 2021.
## For each parameter, n_eff is a crude measure of effective sample size,
## and Rhat is the potential scale reduction factor on split chains (at 
## convergence, Rhat=1).

But really, the best way to interpret the model is to see it. There are many ways to plot the samples produced in the model. One of the simplest ways is to use the bayesplot library.

library(bayesplot)

color_scheme_set("teal")

mcmc_areas(stan_fit1, pars = c("beta[1]", "beta[2]", "sigma"), prob = 0.89) +
  scale_y_discrete(expand = c(0, 0))

For regression problems, it’s also good to see the regression lines in the posterior distribution. This tells us what kind of lines are compatible with the data and the priors. We use rstan’s extract function to get samples from the posterior.

samples <- as.data.frame(rstan::extract(stan_fit1, pars = c("beta[1]", "beta[2]"))) %>%
  rename(intercept = beta.1., slope = beta.2.) # friendlier variable names

ggplot(samples) +
  geom_abline(aes(intercept = intercept, slope = slope), color = "cadetblue4", alpha = 0.20) +
  geom_abline(aes(intercept = mean(intercept), slope = mean(slope)), color = "black", size = 2) +
  scale_x_continuous(limits = c(-1, 1)) +
  scale_y_continuous(limits = c(-0.45, 0.45)) +
  labs(x = "x",
       y = "y",
       title = "Samples from beta posteriors",
       subtitle = latex2exp::TeX("$N_{samples} = 4000$"))

Although the average posterior regression line has positive slope, it’s clear that many lines, even some with negative slope, are compatible.

It doesn’t matter anyway because we know this model is bad - it ignores the clustering. Dealing with the clustering will be the objective for the remainder of this post.

Data

The data block shouldn’t look too daunting compared to before. What’s new is that we have a number of observations int N_obs and a number of participants int N_pts. We also have a vector of participant ids int pid[N_obs].

// STAN CODE
data {
  int N_obs; // number of observations
  int N_pts; // number of participants
  int K; // number of predictors + intercept
  int pid[N_obs]; // participant id vector
  matrix[N_obs, K] x; // matrix of predictors
  real y[N_obs]; // y vector
}

At first it may seem weird that pid has a length of N_obs; you may wonder why it’s not the length of N_pts. This is because this vector will have the participant identifier for each participant in the dataset. (Look at the column pid in actual data to see what I mean).

Parameters

Plenty different here. We start with vector[K] beta_p[N_pts], which describes a vector of vectors. Specifically, we’re declaring an intercept and slope for each participant (N_pts). We then have vector<lower=0>[K] sigma_p which describes the SD for the participant intercepts and slopes. Now for the hyper-priors, vector[K] beta will hold the means for our intercept and slope hyper-priors and corr_matrix[K] Omega is the \(2 \times 2\) correlation matrix that will be in the multivariate normal prior. Finally, real<lower=0> sigma is the population error that will be in the likelihood of \(y\).

// STAN CODE
parameters {
  vector[K] beta_p[N_pts]; // ind participant intercept and slope coefficients by group
  vector<lower=0>[K] sigma_p; // sd for intercept and slope
  vector[K] beta; // intercept and slope hyper-priors
  corr_matrix[K] Omega; // correlation matrix
  real<lower=0> sigma; // population sigma
}

Model

There are two novelties in this model compared to the simple regression. First, assigning a multivariate normal prior to beta_p, our participant intercepts and slopes, requires some unfamiliar notation. We use the multi_normal command. The first argument is beta because beta is our hyper-prior vector for the mean intercept and slope. What’s this quad_form_diag thing? It is a compact and efficient way of expressing diag(sigma_p) * Omega * diag(sigma_p).

The likelihood looks more or less the same as before. The main difference is that we multiply x by the beta_p parameter. The bracket indexing can be a bit confusing. The expression beta_p[pid[i]] is like saying “the intercept and slope for the \(i^{th}\) participant id”.

// STAN CODE
model {
  vector[N_obs] mu;
  
  // priors
  beta ~ normal(0, 1);
  Omega ~ lkj_corr(2);
  sigma_p ~ exponential(1);
  sigma ~ exponential(1);
  beta_p ~ multi_normal(beta, quad_form_diag(Omega, sigma_p));
  
  // likelihood
  for(i in 1:N_obs) {
    mu[i] = x[i] * (beta_p[pid[i]]); // * is matrix multiplication in this context
  }
  
  y ~ normal(mu, sigma);
}

Full Stan Code

Here is the full Stan code. I’ll save the file as “mod2.stan”.

// STAN CODE
data {
  int N_obs; // number of observations
  int N_pts; // number of participants
  int K; // number of predictors + intercept
  int pid[N_obs]; // participant id vector
  matrix[N_obs, K] x; // matrix of predictors
  real y[N_obs]; // y vector
}

parameters {
  vector[K] beta_p[N_pts]; // ind participant intercept and slope coefficients by group
  vector<lower=0>[K] sigma_p; // sd for intercept and slope
  vector[K] beta; // intercept and slope hyper-priors
  corr_matrix[K] Omega; // correlation matrix
  real<lower=0> sigma; // population sigma
}

model {
  vector[N_obs] mu;
  
  // priors
  beta ~ normal(0, 1);
  Omega ~ lkj_corr(2);
  sigma_p ~ exponential(1);
  sigma ~ exponential(1);
  beta_p ~ multi_normal(beta, quad_form_diag(Omega, sigma_p));
  
  // likelihood
  for(i in 1:N_obs) {
    mu[i] = x[i] * (beta_p[pid[i]]); // * is matrix multiplication in this context
  }
  
  y ~ normal(mu, sigma);
}

Visualizing Correlated Slopes and Intercepts

If you look further down the list of parameters in the model output, you’ll see the four \(\Omega\), Omega parameters. These describe the correlation matrix for participant intercepts and slopes. A negative correlation here, tells us that as a participant’s intercept increases, their slope decreases. Let’s plot the participant-specific intercepts and slopes to see this.

We use extract to get the beta_p parameters from the model. Then, using the apply function, we calculate the average of the samples for each beta_p.

beta_p_samples <- rstan::extract(stan_fit2, pars = "beta_p")
beta_p <- as.data.frame(apply(beta_p_samples$beta_p, 2:3, mean)) %>%
  rename(intercept = V1, slope = V2)

beta_p %>%
  ggplot(aes(x = intercept, y = slope)) +
  geom_point(color = "cadetblue4", alpha = 0.75, size = 3) +
  labs(title = "Participant intercept and slope are negatively correlated")

We want to infer beyond the participants in our sample though, so let’s look at the posterior samples for our \(\beta_0\) and \(\beta_1\) parameters. Like with the simple linear regression, we’re going to visualize all the regression lines in the posterior distribution - only now partitioned by intercept. We’ll split the samples into three equal groups based on their intercept. This is of course an arbitrary choice and we’re not implying that there are three distinct groups here. We just want to see how the slope varies by intercept and it’s a lot easier to do this by splitting the samples.

beta_sim <- as.data.frame(rstan::extract(stan_fit2, pars = c("beta[1]", "beta[2]"))) %>%
  rename(intercept = beta.1., slope = beta.2.)

beta_sim %>%
  mutate(int_grp = factor(ntile(intercept, 3), labels = c("[0-33%)", "[33%-66%)", "[66%-100%)"))) %>%
  group_by(int_grp) %>%
  mutate(av_int = mean(intercept),
         av_slope = mean(slope)) %>%
  ungroup() %>%
  ggplot() +
  geom_abline(aes(intercept = intercept, slope = slope), color = "cadetblue4", alpha = 0.20) +
  geom_abline(aes(intercept = av_int, slope = av_slope), color = "black", size = 2) +
  scale_x_continuous(limits = c(-1, 1)) +
  scale_y_continuous(limits = c(-0.75, 0.75)) +
  labs(x = "x",
       y = "y",
       title = "Relation between x and y as a function of intercept.",
       subtitle = "Intercept percentile") +
  facet_wrap(~int_grp) +
  theme(plot.subtitle = element_text(hjust = 0.5))

Yes, the slope seems to be getting shallower as the intercept increases.

Although I find the ‘many lines’ approach to be appealing, it is more common to present figures displaying means and compatibility intervals (Bayesian equivalent of a confidence interval). We start by creating a sequence of x values over a specified range. Since we’re dealing in standardized units, we’ll make it from -1 to 1. Then for each value of x in this sequence we calculate the expected value of y using the beta samples.

I again make use of dplyr’s map_dfr function to iterate over each value of x and I also bin the intercept into the three groups from before. After this step, we have a large dataframe of 100 x values for each of 4000 samples. Then again for each value of x we calculate the mean (mu) of the samples and its lower and upper bound for the compatiblity interval. Here I adopt McElreath’s convention of 89% compatability interval, but there’s nothing more special about this value than, say, 95%.

seq(-1, 1, length.out = 100) %>% # create a sequence of evenly spaced numbers from -1 to 1
  map_dfr(~ data.frame(y = beta_sim$intercept + .x * beta_sim$slope,
                       x = .x,
                       int = factor(ntile(beta_sim$intercept, 3), levels = c(1, 2, 3),
                                    labels = c("[0-33%)", "[33%-66%)", "[66%-100%)")))) %>%
  group_by(x, int) %>%
  summarise(mu = mean(y),
            lower = quantile(y, probs = 0.055),
            upper = quantile(y, probs = 0.945)) %>%
  ungroup() %>%
  ggplot() +  
  geom_ribbon(aes(x = x, ymin = lower, ymax = upper), fill = "cadetblue4", alpha = 0.50) +
  geom_line(aes(x = x, y = mu), color = "black", size = 2) +
  labs(title = "Relation between x and mu as a function of intercept",
       subtitle = "Intercept percentile",
       caption = "Shaded area = 89% compatibility interval") +
  facet_wrap(~int) +
  theme(plot.subtitle = element_text(hjust = 0.5))

Now that’s a publishable figure!

Full Stan Code

Here’s the full non-centered parameterization code.

// STAN CODE
data {
  int N_obs; // number of observations
  int N_pts; // number of participants
  int K; // number of predictors + intercept
  int pid[N_obs]; // participant id vector
  matrix[N_obs, K] x; // matrix of predictors
  real y[N_obs]; // y vector
}

parameters {
  matrix[K, N_pts] z_p; // matrix of intercepts and slope
  vector<lower=0>[K] sigma_p; // sd for intercept and slope
  vector[K] beta; // intercept and slope hyper-priors
  cholesky_factor_corr[K] L_p; // Cholesky correlation matrix
  real<lower=0> sigma; // population sigma
}

transformed parameters {
  matrix[K, N_pts] z; // non-centered version of beta_p
  z = diag_pre_multiply(sigma_p, L_p) * z_p; 
}

model {
  vector[N_obs] mu;
  
  // priors
  beta ~ normal(0, 1);
  sigma ~ exponential(1);
  sigma_p ~ exponential(1);
  L_p ~ lkj_corr_cholesky(2);
  to_vector(z_p) ~ normal(0, 1);
  
  // likelihood
  for(i in 1:N_obs) {
    mu[i] = beta[1] + z[1, pid[i]] + (beta[2] + z[2, pid[i]]) * x[i, 2];
  }
  y ~ normal(mu, sigma);
}

generated quantities {
  matrix[2, 2] Omega;
  Omega = multiply_lower_tri_self_transpose(L_p);
}

I’ll save it to the working directory as “mod2-nc.stan”.