I am having trouble with the chains properly mixing in my model (modified from ZINB-WaVE in Stan for scRNA-seq analysis — What do you mean "heterogeneity"?). In particular, the regression parameters for the bernoulli component of the zero-inflated negative binomial model (beta_pi) is really badly mixing (see plots below). Can you suggest a possible fix/es?

data {
int<lower=0> N;
int<lower=1> P;
vector[P] x[N];
int y[N];
}
parameters {
vector[P] beta_mu;
vector[P] beta_pi;
real<lower=0> phi_rec;
}
model {
real mu;
real pi_;
phi_rec ~ normal(0, 1);
for (n in 1:N){
mu = dot_product(beta_mu, x[n]);
pi_ = dot_product(beta_pi, x[n]);
if (y[n] > 0) {
target += bernoulli_logit_lpmf(0 | pi_) + neg_binomial_2_log_lpmf(y[n] | mu, (1/phi_rec)^2);
}
else {
target += log_sum_exp(bernoulli_logit_lpmf(1 | pi_),
bernoulli_logit_lpmf(0 | pi_) + neg_binomial_2_log_lpmf(y[n] | mu, (1/phi_rec)^2));
}
}
}

What are the values of the covariates? In particular, is the covariate associated with beta_pi[2] entirely positive? If so, then it looks to me like what’s happening is that the model sees no zero inflation as a plausible model configuration. Thus, the chains for beta_pi[2] are free to wander to arbitrarily low values (and they don’t have a prior on them, so wander they will!). You can regularize this using some reasonable weakly informative prior on your coefficients. That will give you a trustworthy fit, and if it shows that the model is confident that there is no zero-inflation, then you can drop the zero-inflation component from the model.

If the covariate associated with beta_pi[2]contains both positive and negative values, then it is virtually certain that you have complete separation along that covariate (unless the covariate values are tiny (on the order of 10^{-15}). Complete separation is a bit tricky in this context, because even if you have zeros over positive values of the covariate these could be attributed to the negative binomial part of the model. But at a minimum, you should find yourself in a situation where there are no nonzero outcomes over negative values of the covariate.

Have you considered relaxing the prior on ‘phi_rec’? As I understand from the code, this is actually the coefficient of variation of your NB distribution, as the shape of an NB is indeed 1/CV^2. I think your prior is quite restrictive as I imagine most of the prior mass for ‘1/phi_rec^2’ is way above 1 and maybe even beyond 5 (beyond which an NB is very Poisson-like). Maybe this is also causing trouble with your regression coefficients for the zero inflation component, which I imagine is trying to compensate for the lack of flexibility in your prior for the shape of your NB distribution.

Broadly speaking, you could opt for one of two flavours, I think:

define a prior directly for the shape parameter with enough probability mass below 5, e.g. a half-normal with sd of 3 to 5, or

define a prior for the CV that makes sense in terms of the shape parameter; this you should explore yourself, for instance by simulating values of phi_rec from your prior of choice (e.g., in R - not sure what you’re working with) and then plotting a histogram of 1 / phi_rec^2. In the past, when I specified a half-normal prior for CV, I made sure that the sd of that prior was sufficiently large for what I expected in terms of shape parameter.