Hi!
Below is the full code for my multilevel VAR(1) model in which I eliminated as many loops and vectorized as much as I can see is possible :). I kept the loops for alpha[i] and beta[i], but that is because I want only I (number of participants) values for those in my output :).
// Learn more about model development with Stan at:
//
// http://mc-stan.org/users/interfaces/rstan.html
// https://github.com/stan-dev/rstan/wiki/RStan-Getting-Started
//
// The input data is a vector 'y' of length 'N'.
data {
int N; // number of observations (number of rows in the panel) this is equal to N*T due to long format
int K; // number of dimensions of the outcome variable
int I; // number of individuals
int T; // The greatest number of time periods for any individual
int<lower = 1, upper = I> individual[N]; // integer vector the same length
// as the panel, indicating which individual each row corresponds to
int<lower = 1, upper = T> time[N]; // integer vector with individual-time
// (not absolute time). Each individual starts at 1
vector[K] Y[N]; // Dependent variables
vector[K] X[N]; // lagged predictor
int Use_obs[N-I]; // indicator that indicates which rows are NOT individual's first observations.
//matrix[N, K] Y; // the outcome matrix, each individual stacked on each
// other, observations ordered earliest to latest
}
// The parameters accepted by the model.
parameters {
vector<lower = 0>[K] tau; // scale vector for residuals
corr_matrix[K] Omega; // Correlation matrix for residuals
vector[K*K] beta_hat_location; // means of the VAR(1) parameters
vector<lower = 0>[K*K] beta_hat_scale; // Sd's of the AR and Cross-lagged products
vector[K] alpha_hat_location; // means for the K DV's
vector<lower = 0>[K] alpha_hat_scale;// sd's of the means of predictors
corr_matrix[K] Omega_alpha;// Correlations between means of predictors
corr_matrix[K*K] Omega_beta;// Correlations between AR and Cross-lagged products
// individual-level parameters
matrix[I,K*K] z_beta;
matrix[I,K] z_alpha;
}
transformed parameters {
matrix[K, K] beta[I]; // individual VAR(1) coefficient matrix
vector[K] alpha[I]; // individual means matrix
matrix[K, K] L_alpha;
matrix[(K*K), (K*K)] L_beta;
L_alpha = cholesky_decompose(quad_form_diag(Omega_alpha, alpha_hat_scale));
L_beta = cholesky_decompose(quad_form_diag(Omega_beta, beta_hat_scale));
for(i in 1:I) { // Can't remove this loop, because I need I alpha and beta's in output
alpha[i] = alpha_hat_location + L_alpha * z_alpha[i]';
// implies: alpha[i] ~ multi_normal(alpha_hat_location, alpha_hat_scale)
beta[i] = to_matrix((beta_hat_location + L_beta * z_beta[i]'), K, K);
}
}
// The model to be estimated.
model {
// hyperpriors
alpha_hat_location ~ normal(4, 4);
alpha_hat_scale ~ cauchy(0, 2);
beta_hat_location ~ normal(0, .5);
beta_hat_scale ~ cauchy(0, .5);
Omega_alpha ~ lkj_corr(3);
Omega_beta ~ lkj_corr(5); // old value 3
// hierarchical priors
// non-centered parameterization
to_vector(z_alpha) ~ std_normal();
to_vector(z_beta) ~ std_normal();
// Priors for error variance; tau is in sd's
tau ~ normal(1, 2);
Omega ~ lkj_corr(3);
{
vector[K] Xc[N];
vector[K] Mean[N];
vector[K] alpha_vec[N];
vector[N] ar1_vec;
vector[N] ar2_vec;
vector[N] cl1_vec;
vector[N] cl2_vec;
//vector[K] cl_vec[N];
vector[K] Mean_Vec[N-I];
vector[K] Y_Vec[N-I];
alpha_vec = alpha[individual]; // store individual means in a vector, with each pair of means appearing multiple times. Once for each row that
// belong to an individual. Arguments are: individual means (vector[K] alpha[I]), indictor array listing which individual belong to each row
// (int individual[N]). Results in vector[K] alpha_vec[N] array of all mean pairs for the rows. Here K = number of variables, I = number of
// distinct individuals, and N = number of rows in data (#people * #observations pp (T)).
ar1_vec = to_vector(beta[individual,1,1]);// separately store all the individual AR-parameters for variable 1 belonging to the rows.
// Since mulitple rows belong to the same person, values are repeated. Imput is the individual matrices holding the VAR-parameters (AR and
// cross-lagged) for each individual (matrix[K, K] beta[I];), indictor array listing which individual belong to each row (int individual[N]).
// Results in vector[K] ar1_vec[N] array of all AR-parameters for variable 1, for each of the rows.
ar2_vec = to_vector(beta[individual,K,K]); // same as above vut for AR-parameter of variable 2
cl1_vec = to_vector(beta[individual,1,K]);// same as above vut for cross-lagged effect of variable 2 on variable 1
cl2_vec = to_vector(beta[individual,K,1]);// same as above vut for cross-lagged effect of variable 1 on variable 2
Xc[,1] = to_array_1d(to_vector(X[1:N,1]) - to_vector(alpha_vec[1:N,1]));
Xc[,2] = to_array_1d(to_vector(X[1:N,2]) - to_vector(alpha_vec[1:N,2]));
// Here I want to determine the predicted mean for each row (I'm doing so per column). This is done by taking the individual means for that row,
// and adding 1) variable 1 multiplied with its AR-parameter, and 2) variable 2 multiplied with the cross-lagged effect of var 2 on var1.
// so I want vector[K] Mean[,1] = vector[K] alpha_vec[,1] + (vector[K] ar1_vec multiplied elementwise with vector[K] Xc[,1]) +
// vector[K] cl1_vec multiplied elementwise with vector[K] Xc[,2]). This is then adding a bunch of similar sized vectors to get a vector that
// is the sum of those vectors.
Mean[,1] = to_array_1d(to_vector(alpha_vec[1:N,1]) + to_vector(ar1_vec .* to_vector(Xc[1:N,1])) + to_vector(cl1_vec .* to_vector(Xc[1:N,2])));
Mean[,2] = to_array_1d(to_vector(alpha_vec[1:N,2]) + to_vector(ar2_vec .* to_vector(Xc[1:N,2])) + to_vector(cl2_vec .* to_vector(Xc[1:N,1])));
// Don't use first observations for a participant, due to lag, no predictors there
Y_Vec = Y[Use_obs];
Mean_Vec = Mean[Use_obs];
Y_Vec ~ multi_normal(Mean_Vec, quad_form_diag(Omega, tau));
}
}
This code took about 4 minutes of the time needed to construct Mean by looping over all 14000 rows. so I guess looping 14000 times takes 4 minues ;).
I’m now slowly parallel-izing this using reduce_sum. Starting simple by only moving the final sampling from the multivariate normal into the reduce_sum function. I’ll play around with the grainsize there, and slowly move more into the code that is send to the slices. So the first code for parallelized multilevel VAR(1) is:
// The reduce_sum function
functions {
real part_sum(vector[] y_slice,
int start,
int end,
vector[] mu,
matrix Omega,
vector tau
) {
return multi_normal_lpdf(y_slice | mu[start:end], quad_form_diag(Omega, tau));
}
}
data {
int N; // number of observations (number of rows in the panel) this is equal to N*T due to long format
int K; // number of dimensions of the outcome variable
int I; // number of individuals
int T; // The greatest number of time periods for any individual
int<lower = 1, upper = I> individual[N]; // integer vector the same length
// as the panel, indicating which individual each row corresponds to
int<lower = 1, upper = T> time[N]; // integer vector with individual-time
// (not absolute time). Each individual starts at 1
vector[K] Y[N]; // The data with I*T (=N) observations on K variables
vector[K] X[N]; // lagged predictor
int Use_obs[N-I]; // indicator that indicates which rows are NOT individual's first observations. I want to
// remove the first observations sinc ethose don't have alagged predictor value (nothing came before)
int grainsize;
}
// The parameters accepted by the model.
parameters {
vector<lower = 0>[K] tau; // scale vector for residuals
corr_matrix[K] Omega; // Correlation matrix for correlations between residuals
vector[K*K] beta_hat_location; // a K*K matrix of grand-mean (across-individual) AR-parameters (on the diagonal)
// and cross-lagged parameters.
vector<lower = 0>[K*K] beta_hat_scale; // Sd's for AR and Cross-lagged effects
vector[K] alpha_hat_location; // grand-mean (across-individuals) on the K DV's
vector<lower = 0>[K] alpha_hat_scale;// sd of the means
corr_matrix[K] Omega_alpha; // correlations between means
corr_matrix[K*K] Omega_beta;// correlations between AR- and cross-lagged effects
// individual-level parameters.
matrix[I,K*K] z_beta;// use for the non-centered parametrization
matrix[I,K] z_alpha; // use for the non-centered parametrization
}
transformed parameters {
matrix[K, K] beta[I]; // Individual VAR(1) coefficient matrices (with individual AR- and cross-lagged effects)
vector[K] alpha[I]; // Individual means on the K DV's
matrix[K, K] L_alpha; // cholesky decomposition of correlation matrix of the means
matrix[(K*K), (K*K)] L_beta; // cholesky decomposition of correlation matrix of the AR/Cross-lagged effects
L_alpha = cholesky_decompose(quad_form_diag(Omega_alpha, alpha_hat_scale));
L_beta = cholesky_decompose(quad_form_diag(Omega_beta, beta_hat_scale));
for(i in 1:I) {
alpha[i] = alpha_hat_location + L_alpha * z_alpha[i]';// non=centered parametrization
// implies: alpha[i] ~ multi_normal(alpha_hat_location, alpha_hat_scale)
beta[i] = to_matrix((beta_hat_location + L_beta * z_beta[i]'), K, K); // non=centered parametrization
}
}
// The model to be estimated.
model {
// hyperpriors
alpha_hat_location ~ normal(4, 4);
alpha_hat_scale ~ cauchy(0, 2);
beta_hat_location ~ normal(0, .5);
beta_hat_scale ~ cauchy(0, .5);
Omega_alpha ~ lkj_corr(3);
Omega_beta ~ lkj_corr(5);
// hierarchical priors
to_vector(z_alpha) ~ std_normal();
to_vector(z_beta) ~ std_normal();
// Priors for error variance; tau is in sd's
tau ~ normal(1, 2);
Omega ~ lkj_corr(3);
{
vector[K] Xc[N];
vector[K] Mean[N];
vector[K] alpha_vec[N];
vector[N] ar1_vec;
vector[N] ar2_vec;
vector[N] cl1_vec;
vector[N] cl2_vec;
//vector[K] cl_vec[N];
vector[K] Mean_Vec[N-I];
vector[K] Y_Vec[N-I];
alpha_vec = alpha[individual]; // store individual means in a vector, with each pair of means appearing multiple times. Once for each row that
// belong to an individual. Arguments are: individual means (vector[K] alpha[I]), indictor array listing which individual belong to each row
// (int individual[N]). Results in vector[K] alpha_vec[N] array of all mean pairs for the rows. Here K = number of variables, I = number of
// distinct individuals, and N = number of rows in data (#people * #observations pp (T)).
ar1_vec = to_vector(beta[individual,1,1]);// separately store all the individual AR-parameters for variable 1 belonging to the rows.
// Since mulitple rows belong to the same person, values are repeated. Imput is the individual matrices holding the VAR-parameters (AR and
// cross-lagged) for each individual (matrix[K, K] beta[I];), indictor array listing which individual belong to each row (int individual[N]).
// Results in vector[K] ar1_vec[N] array of all AR-parameters for variable 1, for each of the rows.
ar2_vec = to_vector(beta[individual,K,K]); // same as above vut for AR-parameter of variable 2
cl1_vec = to_vector(beta[individual,1,K]);// same as above vut for cross-lagged effect of variable 2 on variable 1
cl2_vec = to_vector(beta[individual,K,1]);// same as above vut for cross-lagged effect of variable 1 on variable 2
Xc[,1] = to_array_1d(to_vector(X[1:N,1]) - to_vector(alpha_vec[1:N,1]));
Xc[,2] = to_array_1d(to_vector(X[1:N,2]) - to_vector(alpha_vec[1:N,2]));
// Here I want to determine the predicted mean for each row (I'm doing so per column). This is done by taking the individual means for that row,
// and adding 1) variable 1 multiplied with its AR-parameter, and 2) variable 2 multiplied with the cross-lagged effect of var 2 on var1.
// so I want vector[K] Mean = vector[K] alpha_vec[,1] + (vector[K] ar1_vec multiplied elementwise with vector[K] Xc[,1]) +
// (vector[K] cl1_vec multiplied elementwise with vector[K] Xc[,2]). This is then adding a bunch of similar sized vectors to get a vector that
// is the sum of those vectors.
Mean[,1] = to_array_1d(to_vector(alpha_vec[1:N,1]) + to_vector(ar1_vec .* to_vector(Xc[1:N,1])) + to_vector(cl1_vec .* to_vector(Xc[1:N,2])));
Mean[,2] = to_array_1d(to_vector(alpha_vec[1:N,2]) + to_vector(ar2_vec .* to_vector(Xc[1:N,2])) + to_vector(cl2_vec .* to_vector(Xc[1:N,1])));
// Don't use first observations for a participant, due to lag, no predictors there
Y_Vec = Y[Use_obs];
Mean_Vec = Mean[Use_obs];
target += reduce_sum(part_sum,Y_Vec,grainsize,Mean_Vec,Omega,tau); // Feed DV's, Mean, Correlation matrix and
}
}