I am trying to fit a two-part latent variable model in stan, which includes fitting one two-part model and two latent variables. For imputing one of the latent variables - sensation-seeking, I will be using another two-part model, comprised of a Bernoulli and a gamma part. For the other latent variable - financial literacy, I will be using purely Bernoulli manifest variables. Our main two-part model, which depicts my response variable, consists of a Bernoulli part and a log-normal part.
I am well aware of the fact that my model specification is over-complicated, and it takes ages to even run a single iteration in stan. So I am wondering whether there is any way that can help speed up my code. Is the slow iteration mainly caused by Gamma distribution? Any suggestions would be appreciated!
data {
int<lower=0> n; //number of data items
int<lower=0> p;//number of predictors (fixed covariates)
real<lower = 0> s_cont[n]; //a 1 X n array of outcome variable
matrix[n,p] x_cov; //fixed covariates matrix
int<lower=0> m_1; //number of manifest variables for sensation seeking
int<lower=0> m_2; //number of manifest variables for financial literacy
int<lower=0> m; //number of manifest variables
real<lower=0> u1_i[n, m_1]; //a m_1 X n array of sensation seeking manifests
int<lower=0> u2_i[n, m_2]; //a m_2 X n array of financia literacy manifests
}
transformed data{
matrix[n,p] Q;
matrix[p,p] R;
matrix[p,p] R_inverse;
Q = qr_Q(x_cov)[,1:p] * n;
R = qr_R(x_cov)[1:p,] / n;
R_inverse = inverse(R);
}
parameters {
real mu1; // constant term in two-part model's binary model
real mu2; // constant term in two-part model's continuous model
vector[p] alpha1_tilde; // factor loadings in two-part model's binary model
vector[p] alpha2_tilde; // factor loadings in two-part model's continuous model
matrix[p,m_1] alpha3_tilde; //factor loadings for fixed covariates in sensation seeking model's binary part
matrix[p,m_1] alpha4_tilde; //factor loadings for fixed covariates in sensation seeking model's continuous part
real<lower = 0.05> sigma; //std. of error term in two-part model's continuous part
vector[2] beta1; //latent factor loadings in two-part model's binary part
vector[2] beta2; //latent factor loadings in two-part model's continuous part
vector[m_1] mu3;
vector[m_1] mu4;
vector[m_2] mu5;
vector[m_1] lambda_11; //factor loadings for sensation seeking's binary part
vector[m_1] lambda_12; //factor loadings for sensation seeking's continuous part
vector<lower=0.05>[m_1] u1_shape; //shape parameter for gamma distribution
vector[m_2] lambda_2; //factor loadings for financial literacy
vector[n] omega1_i; //sensation seeking indexes
vector[n] omega2_i; //financial literacy indexes
vector<lower=0.05>[2] tau; //covariance matrix of latent variables
}
transformed parameters {
vector[p] alpha1 = R_inverse * alpha1_tilde; // factor loadings in two-part model's binary part
vector[p] alpha2 = R_inverse * alpha2_tilde; // factor loadings in two-part model's continuous part
matrix[p,m_1] alpha3 = R_inverse * alpha3_tilde ; //factor loadings for fixed covariates in our sensation seeking model's binary part
matrix[p,m_1] alpha4 = R_inverse * alpha4_tilde ; //factor loadings for fixed covariates in sensation seeking model's continuous part
matrix[2,n] omega_i; //latent factors
vector[n] x_beta1;
vector[n] x_beta2;
vector[m_1] lambda1_fix;
vector[m_2] lambda2_fix;
vector[n] u_temp11[m_1];
vector[n] u_temp12[m_1];
lambda1_fix[1] = 1;
lambda2_fix[1] = 1;
for(i in 1:m_1-1){
lambda1_fix[i+1] = lambda_11[i+1];
}
for(i in 1:m_2-1){
lambda2_fix[i+1] = lambda_2[i+1];
}
omega_i = [omega1_i', omega2_i'];
x_beta1 = mu1 + Q * alpha1_tilde + (omega_i') * beta1; // X*B_1: n X 1
x_beta2 = mu2 + Q * alpha2_tilde + (omega_i') * beta2; // X*B_2: n X 1
for (i in 1:m_1) {
u_temp11[i] = mu3[i] + Q * alpha3_tilde[,i] + omega1_i * lambda1_fix[i]; //X*B_3: n X 1: overall, (m1*n) X 1
u_temp12[i] = mu4[i] + Q * alpha4_tilde[,i] + omega1_i * lambda_12[i];
}
}
model{
vector[m_2] u_temp2[n]; // logit(Prob. of success)
for(i in 1:n){
omega_i[,i] ~ normal(0, tau);
}
for(i in 1:n) {
for (j in 1:m_1) {
if (u1_i[i,j]==0)
target += bernoulli_logit_lpmf(0|u_temp11[j,i]);
else
target += bernoulli_logit_lpmf(1|u_temp11[j,i])
+ gamma_lpdf( u1_i[i,j] |u1_shape[j],u1_shape[j]/exp(u_temp12[j,i]));
}
}
for(i in 1:n){
u_temp2[i] = mu5 + lambda2_fix*(omega2_i[i]);
}
for(i in 1:n){
for(j in 1:m_2){
target += bernoulli_logit_lpmf(u2_i[i,j]| u_temp2[i,j]);
}
}
for(i in 1:n){
if(s_cont[i]==0)
target += bernoulli_logit_lpmf(0|x_beta1[i]);
else
target += bernoulli_logit_lpmf(1|x_beta1[i])
+ lognormal_lpdf(s_cont[i]|x_beta2[i],sigma);
}
mu1 ~ normal(0,2);
mu2 ~ normal(0,2);
sigma ~ inv_gamma(0.01,0.01);
alpha1 ~ normal(0,5);
alpha2 ~ normal(0,5);
beta1 ~ normal(0,5);
beta2 ~ normal(0,5);
tau ~ inv_gamma(0.01,0.01);
mu3 ~ normal(0,2);
mu4 ~ normal(0,2);
mu5 ~ normal(0,2);
lambda_11 ~ normal(0,5);
lambda_12 ~ normal(0,5);
lambda_2 ~ normal(0,5);
to_vector(alpha3) ~ normal(0,5);
to_vector(alpha4) ~ normal(0,5);
to_vector(u1_shape) ~ gamma(0.01,0.01);
}