Hierarchical Linear Mixture Model


I have implemented a stan hierarchical model with level 1 to be a linear model and level 2 Gaussian mixture model. The models has convergence problem. Any comments on improving the model.

            data {
              int<lower=0> J; # No of subjects
              int<lower=0> N; # No of observations
              int<lower=1,upper=J> RID[N]; 
              vector[N] x;  # Cognitive Score
              vector[N] y; # Output Label - Disease Stage
              vector[J] a;  
              real mu_a;
              vector[J] b;
              real mu_b;
              real<lower=0,upper=1> sigma_a;
              real<lower=0,upper=1> sigma_b;
             # Gaussian Parameters for level 2
              vector[2] mu;
              real<lower=0> sigma[2];
              real<lower=0, upper=1> theta;
            transformed parameters {
              vector[N] y_hat;
              for(i in 1:N)
                y_hat[i] = a[RID[i]] + x[i] * b[RID[i]];
            model {
              sigma_a ~ uniform(0, 1);
              a ~ normal(mu_a, sigma_a);
              sigma_b ~ uniform(0, 1);
              b ~ normal(mu_b, sigma_b);
             sigma ~ normal(0, 1);
             mu ~ normal(0, 1);
             theta ~ beta(5, 5);
             target += log_mix(theta,
                                 normal_lpdf(y_hat | mu[1], sigma[1]),
                                 normal_lpdf(y_hat | mu[2], sigma[2]));


There is a lot of stuff that could be wrong - I wrote a list of strategies to try at https://www.martinmodrak.cz/2018/02/19/taming-divergences-in-stan-models/ - it is aimed at divergences, but a lot of it works as well for high Rhat, low BFMI or low n_eff scenarios.