Vectorization for multivariate normal mixed effects model

Modeling
1

Dear Stan community,

I’d love to be able to vectorize this model to improve efficiency, i.e. writing "y ~ " instead of a for loop. However, there are a couple of things that are making it complicated.

This is a bivariate longitudinal mixed effects model with 3 strata (diagnostic groups), I subjects, and N total observations. So for observation n, id[n] is the subject and group[id[n]] is the stratum that the subject is in. The fixed effects “alpha” are common to every observation in the stratum, so I have written “alpha[group[id[n]]]”, but the random effects “beta” are different for every subject, so I have written “beta[id[n]]”. Would there a way to vectorize the likelihood in this situation, perhaps writing something like the following?

y ~ normal(w * alpha[group[id]] + x * beta[id]);

Second, the random effects beta have a specified covariance structure, multivariate normal, but the fixed effects alpha do not. I need to be able to model them separately.

By the way, is the way I’ve specified the covariance optimal?

  • vector of sd’s sigma
  • cholesky_factor_corr Rho
  • diag_pre_multiply(sigma, L_Rho) * diag_pre_multiply(sigma, L_Rho)’

or is there a better option?

Here is the relevant part of the model:

data{
  
  int<lower=1> I; // number of subjects
  int<lower=1> N; // total number of observations, each n is ij
  int<lower=1> id[N]; // id of nth observation
  int<lower=1> group[N]; // diagnostic group of subject i
  vector[2] y[N]; // bivariate outcome
  matrix[N, 6] X; // covariates 
  
}

parameters{

vector[8] alpha[3]; // fixed effects
vector[4] beta[I]; // random effects
vector[4] mu_beta[3]; // means of random effects
vector<lower=0>[4] sigma_beta[3]; // sd's of random effects
cholesky_factor_corr[4] L_Rho_beta[3]; // correlation matrices of random effects
vector<lower=0>[2] sigma_epsilon[3]; // sd's of errors
cholesky_factor_corr[2] L_Rho_epsilon[3]; // correlation matrices of errors

}

model{

// Likelihood

for(n in 1:N){
  y[n] ~ multi_normal(
    [row(X, n) * append_row(beta[id[n]][1:2], alpha[group[id[n]]][1:4]),
    row(X, n) * append_row(beta[id[n]][3:4], alpha[group[id[n]]][5:8])],
    diag_pre_multiply(sigma_epsilon[group[id[n]]], L_Rho_epsilon[group[id[n]]]) * 
    diag_pre_multiply(sigma_epsilon[group[id[n]]], L_Rho_epsilon[group[id[n]]])'  
  );
}

// Priors

for(i in 1:I){
  beta[i] ~ multi_normal(
    mu_beta[group[i]], 
    diag_pre_multiply(sigma_beta[group[i]], L_Rho_beta[group[i]]) *
    diag_pre_multiply(sigma_beta[group[i]], L_Rho_beta[group[i]])'    
  );
}

// rest of priors not included
2

See my approach here (I think that’s the version in that repo that’s most similar to your use case). Note that the likelihood could be accelerated even further using the Gaussian sufficient stats trick, but I haven’t gotten around to adding that yet.

But remember: best to get the model working correctly with smaller simulated data before delving into optimization for larger data.

1 Like
3

Thanks for the reply, Mike. I would love to check out your work. Unfortunately it seems like the link is broken. Could you provide another one?

4

Oops, sorry; link fixed now.