Hello,
I’m attempting to adapt some brms code to fit a mixed-effect location scale model where the variance of the location random effects is a function of predictors. Here’s the code I’m trying to adapt:
data {
int<lower=1> N; // number of observations
vector[N] Y; // response variable
int<lower=1> K; // number of population-level effects
matrix[N, K] X; // population-level design matrix
int<lower=1> K_sigma; // number of population-level effects
matrix[N, K_sigma] X_sigma; // population-level design matrix
// data for group-level effects of ID 1
int<lower=1> N_1; // number of grouping levels
int<lower=1> M_1; // number of coefficients per level
int<lower=1> J_1[N]; // grouping indicator per observation
// group-level predictor values
vector[N] Z_1_1;
vector[N] Z_1_sigma_2;
int<lower=1> NC_1; // number of group-level correlations
int prior_only; // should the likelihood be ignored?
}
transformed data {
int Kc = K - 1;
matrix[N, Kc] Xc; // centered version of X without an intercept
vector[Kc] means_X; // column means of X before centering
int Kc_sigma = K_sigma - 1;
matrix[N, Kc_sigma] Xc_sigma; // centered version of X_sigma without an intercept
vector[Kc_sigma] means_X_sigma; // column means of X_sigma before centering
for (i in 2:K) {
means_X[i - 1] = mean(X[, i]);
Xc[, i - 1] = X[, i] - means_X[i - 1];
}
for (i in 2:K_sigma) {
means_X_sigma[i - 1] = mean(X_sigma[, i]);
Xc_sigma[, i - 1] = X_sigma[, i] - means_X_sigma[i - 1];
}
}
parameters {
vector[Kc] b; // population-level effects
real Intercept; // temporary intercept for centered predictors
vector[Kc_sigma] b_sigma; // population-level effects
real Intercept_sigma; // temporary intercept for centered predictors
vector<lower=0>[M_1] sd_1; // group-level standard deviations
matrix[M_1, N_1] z_1; // standardized group-level effects
cholesky_factor_corr[M_1] L_1; // cholesky factor of correlation matrix
}
transformed parameters {
matrix[N_1, M_1] r_1; // actual group-level effects
// using vectors speeds up indexing in loops
vector[N_1] r_1_1;
vector[N_1] r_1_sigma_2;
// compute actual group-level effects
r_1 = (diag_pre_multiply(sd_1, L_1) * z_1)’;
r_1_1 = r_1[, 1];
r_1_sigma_2 = r_1[, 2];
}
model {
// initialize linear predictor term
vector[N] mu = Intercept + Xc * b;
// initialize linear predictor term
vector[N] sigma = Intercept_sigma + Xc_sigma * b_sigma;
for (n in 1:N) {
// add more terms to the linear predictor
mu[n] += r_1_1[J_1[n]] * Z_1_1[n];
}
for (n in 1:N) {
// add more terms to the linear predictor
sigma[n] += r_1_sigma_2[J_1[n]] * Z_1_sigma_2[n];
}
for (n in 1:N) {
// apply the inverse link function
sigma[n] = exp(sigma[n]);
}
// priors including all constants
target += student_t_lpdf(Intercept | 3, 2, 10);
target += student_t_lpdf(Intercept_sigma | 3, 0, 10);
target += student_t_lpdf(sd_1 | 3, 0, 10)
- 2 * student_t_lccdf(0 | 3, 0, 10);
target += normal_lpdf(to_vector(z_1) | 0, 1);
target += lkj_corr_cholesky_lpdf(L_1 | 1);
// likelihood including all constants
if (!prior_only) {
target += normal_lpdf(Y | mu, sigma);
}
}
generated quantities {
// actual population-level intercept
real b_Intercept = Intercept - dot_product(means_X, b);
// actual population-level intercept
real b_sigma_Intercept = Intercept_sigma - dot_product(means_X_sigma, b_sigma);
// compute group-level correlations
corr_matrix[M_1] Cor_1 = multiply_lower_tri_self_transpose(L_1);
vector<lower=-1,upper=1>[NC_1] cor_1;
// extract upper diagonal of correlation matrix
for (k in 1:M_1) {
for (j in 1:(k - 1)) {
cor_1[choose(k - 1, 2) + j] = Cor_1[j, k];
}
}
}
I’m thinking I need to make the first element of the 2x1 vector sd_1 a function of the predictors before computing computing the random effects (r_1). Since this code uses a Cholesky decomposition to compute the random effects, I’m not sure how to go about this. Does anybody have any ideas?
Thanks!
Madeline
