I have some data of 200 study participants belonging to two groups who completed a motor task twice, once for each hand. We derived 20 features from each task’s recording.
The purpose is to identify which features differ between the 2 groups, and secondly whether there are any differences between the features derived from the 2 different hand replications.
The data is setup as 1 row per task (so 400 rows of 200 participants x 2 replicates) with 20 columns.
I want to model each feature as a function of group & hand, so I’m using the seemingly unrelated regressions example from the user guide with just a single intercept as the linear predictor being group and hand specific.
I’ve already added the Onion initialisation as it wasn’t sampling well without it (kept throwing the error about rejected proposals because of a cholesky factor being 0).
With the onion it now samples, but I get both the treedepth warning (21%) and divergences (5%), along with many of the correlation matrix having extremely poor Rhat (235 params have Rhat > 1.1).
Are there any more optimisations I can use?
NB: I’ve had to hand-copy the model definition as it’s being run on a secure server so there might be typos.
functions {
array[] vector create_shapes(int K, real eta) {
real alpha = eta + (K-2) / 2.0;
array[2] vector[K-1] shape;
shape[1, 1] = alpha;
shape[2, 1] = alpha;
for (k in 2:(K-1)) {
alpha -= 0.5;
shape[1,k] = k / 2.0;
shape[2,k] = alpha;
}
return shape;
}
matrix lkj_onion(int K, row_vector l, vector R2, data array[] vector shape) {
matrix[K,K] L = rep_matrix(0, K, K);
{
int start = 1;
int end = 2;
L[1,1] = 1.0;
L[2, 1] = 2.0 * R2[1] - 1.0;
L[2,2] = sqrt(1.0 - square(L[2,1]));
for (k in 2:(K-1)) {
int kp1 = k+1;
row_vector[k] l_row = segment(l, start, k);
real scale = sqrt(R2[k] / dot_self(l_row));
L[kp1, 1:k] = l_row[1:k] * scale;
L[kp1, kp1] = sqrt(1.0 - R2[k]);
start = end + 1;
end = start + k-1;
}
}
return L;
}
}
data {
real<lower=0> eta;
int<lower=0> N_obs;
int<lower=0> N_feats;
int<lower=0> N_groups;
int<lower=0> N_hands;
array[N_obs] vector[N_feats] y;
array[N_obs] int<lower=0> group;
array[N_obs] int<lower=0> handid;
}
transformed data {
array[2] vector[N_feats-1] shapes = create_shapes(N_feats, eta);
}
parameters {
// Onion parameters
row_vector[choose(N_feats, 2) -1] l; // do NOT init with 0 for all elements
vector<lower=0, upper=1>[N_feats-1] R2; // first element is not really a R^2 but is on (0,1)
array[N_groups, N_hands] vector[N_feats] alpha;
vector<lower=0>[N_feats] sigma;
}
transformed parameters {
// Onion LKJ, implies Lcorr ~ lkj(eta)
cholesky_factor_corr[N_feats] Lcorr = lkj_onion(N_feats, l, R2, shapes);
}
model {
for (i in 1:N_groups) {
for (j in 1:N_hands) {
alpha[i,j] = normal(0, 1);
}
}
sigma ~ cauchy(0, 2.5);
// Onion, implies Lcorr ~ lkj_corr_cholesky(eta)
R2 ~ beta(shapes[1], shapes[2]);
l ~ std_normal();
array[N_obs] vector[N_feats] mu;
for (i in 1:N_obs) {
mu[i] = alpha[group[i], handid[i]];
}
y ~ multi_normal_cholesky(mu, diag_pre_multiply(sigma, Lcorr));
}
generated quantities {
matrix[N_feats, N_feats] Rho;
Rho = multiply_lower_tri_self_transpose(Lcorr);
}