Feedback on my model

Modeling
1

Hi all,

Looking to get some feedback about my stan model that I am implementing in rstan. It is quite a complex model and probably a bit beyond my knowledge. It is joint longitudinal model with imputation for missing data. It does work, but it takes a significant time to run and significant memory requirements (too much for my 64gb of ram!). I am wondering if there are any obvious signs to make the model more efficient either regarding speed or memory. Thank you!

functions {
  // LINEAR PREDICTOR FOR THE LONGITUDINAL SUBMODEL
  vector linear_predictor(int[] ID, vector tt, matrix X1, real i_beta, real s_beta, vector beta, matrix bi){
    int N = num_elements(tt);
    vector[N] out = i_beta + bi[ID,1] + s_beta * tt + rows_dot_product(bi[ID,2],tt) + X1 * beta;
    return out;
  } 
}

data {
  int N;
  int n;
  int nX1;
  int nX2;
  int<lower=0,upper=1> R[N];            
  int<lower=1> N_obs;                 
  int<lower=1> N_miss;                   
  int obs_index[N_obs];                 
  int miss_index[N_miss];               
  int<lower=0,upper=1> y[N_obs];       
  vector[N] tt;
  matrix[N,nX1] X1;
  int<lower=1,upper=n> ID[N];
  vector[n] time;
  vector[n] status;
  matrix[n,nX2] X2;
  int X1_nMiss;
  int X2_nMiss;
  int X1_row_index[X1_nMiss];
  int X1_col_index[X1_nMiss];
  int X2_row_index[X2_nMiss]; 
  int X2_col_index[X2_nMiss];
}

parameters {
  real i_beta;
  real s_beta;
  vector[nX1] beta;
  vector[nX2] gamma;
  vector[2] alpha;
  real<lower=0> phi;
  vector<lower=0>[2] sigma_bi;
  cholesky_factor_corr[2] L_bi;  
  vector[nX1] mu_X1;
  cholesky_factor_cov[nX1] L_X1;
  vector[X1_nMiss] X1_imp;
  vector[nX2] mu_X2;
  cholesky_factor_cov[nX2] L_X2;
  vector[X2_nMiss] X2_imp;
  vector<lower=0, upper=1>[N_miss] y_miss;  
  real delta0;
  real delta1;
  matrix[n,2] bi;
}

transformed parameters {
  matrix[N, nX1] X1_full = X1;
  matrix[n, nX2] X2_full = X2;

  for(i in 1:X1_nMiss){ X1_full[X1_row_index[i], X1_col_index[i]] = X1_imp[i]; }
  for(i in 1:X2_nMiss){ X2_full[X2_row_index[i], X2_col_index[i]] = X2_imp[i]; }
}

model {
  vector[N] linpred = linear_predictor(ID, tt, X1_full, i_beta, s_beta, beta, bi);
  vector[n] logHaz;
  vector[n] cumHaz;

  // LONGITUDINAL SUBMODEL
  // OBSERVED
  target += bernoulli_logit_lpmf(y | linpred[obs_index]);

  // SURVIVAL SUBMODEL
  for(i in 1:n){
     logHaz[i] = log(phi) + (phi - 1) * log(time[i]) + X2_full[i,] * gamma + alpha[1] * bi[i,1] + alpha[2] * bi[i,2];
     cumHaz[i] = pow(time[i], phi) * exp(X2_full[i,] * gamma + alpha[1] * bi[i,1] + alpha[2] * bi[i,2]);
     target += status[i] * logHaz[i] - cumHaz[i];
  }

  // RANDOM EFFECTS
 for(i in 1:n){
  bi[i] ~ multi_normal_cholesky(rep_vector(0,2), diag_pre_multiply(sigma_bi, L_bi));
}

  // MISSINGNESS MODEL 
  for(i in 1:N_obs){ target += bernoulli_logit_lpmf(R[obs_index[i]] | delta0 + delta1 * y[i]); }
  for(j in 1:N_miss){ target += bernoulli_logit_lpmf(R[miss_index[j]] | delta0 + delta1 * y_miss[j]); }

  // MISSING (as latent probabilities)
  for(j in 1:N_miss){
     target += log_mix(y_miss[j], 
                       bernoulli_logit_lpmf(1 | linpred[miss_index[j]]),
                       bernoulli_logit_lpmf(0 | linpred[miss_index[j]]));
  }

  for(i in 1:N){ X1_full[i] ~ multi_normal_cholesky(mu_X1, L_X1); }
  for(i in 1:n){ X2_full[i] ~ multi_normal_cholesky(mu_X2, L_X2); }

  // PRIORS
  i_beta ~ normal(0, 5);
  s_beta ~ normal(0, 5);
  beta ~ normal(0, 5);
  gamma ~ normal(0, 5);
  alpha ~ normal(0, 5);
  phi ~ gamma(2, 0.1);
    
  sigma_bi ~ exponential(1);
  L_bi ~ lkj_corr_cholesky(2);
  
  y_miss ~ beta(1, 1);
  mu_X1 ~ normal(0, 1); 
  L_X1 ~ lkj_corr_cholesky(2);
  mu_X2 ~ normal(0, 1);  
  L_X2 ~ lkj_corr_cholesky(2);
  X1_imp ~ normal(0, 1);
  X2_imp ~ normal(0, 1);
  y_miss ~ beta(1, 1);
  delta0 ~ normal(0, 2);
  delta1 ~ normal(0, 2);

}

Thank you in advance for any help or guidance

2

There seems to be quite a bit going on there.

Can you give a brief description of the basics of the model, and how the actual models extends or expands on it? It will probably understanding the essential mechanics and review the implementation.