Predicting from a zero inflated negative binomial model

Modeling
1

I’m fitting a ZINB model with some code initially generated in brms.

I’m using this function for the likelihood:

/* zero-inflated negative binomial log-PDF of a single response
   * log parameterization for the negative binomial part
   * logit parameterization of the zero-inflation part
   * Args:
   *   y: the response value
   *   eta: linear predictor for negative binomial distribution
   *   phi: shape parameter of negative binomial distribution
   *   zi: linear predictor for zero-inflation part
   * Returns:
   *   a scalar to be added to the log posterior
   */
real zero_inflated_neg_binomial_log_logit_lpmf(int y, real eta,
                                                 real phi, real zi) {
    if (y == 0) {
      return log_sum_exp(bernoulli_logit_lpmf(1 | zi),
                         bernoulli_logit_lpmf(0 | zi) +
                         neg_binomial_2_log_lpmf(0 | eta, phi));
    } else {
      return bernoulli_logit_lpmf(0 | zi) +
             neg_binomial_2_log_lpmf(y | eta, phi);
    }
  }

The likelhood is specified as::

for (n in 1:N) {
      target += zero_inflated_neg_binomial_log_logit_lpmf(Y[n] | mu[n], shape, zi[n]);
    }

I’m trying to get samples for post-fit analysis in generated quantities, but I’m not sure how to translate the lpmf functions to rng’s, particularly when the inputs are set to 0 and 1.

So far, I have tried:

for (n in 1 : N) {
    if (Y[n] == 0) {
      y_pred[n] = log_sum_exp(bernoulli_logit_rng(zi[n]),
                              bernoulli_logit_rng(1- zi[n]) +
                               neg_binomial_2_log_rng(mu[n], shape)); // This seems incorrect
    } else {
      y_pred[n] = bernoulli_logit_rng(1-zi[n]) +
        neg_binomial_2_log_rng(mu[n], shape);
    }
  }

The predictions are pretty wild and much more variable than when predicting from a brms object with the posterior_epred function. Am I way off the mark with the approach?

I’ve unsuccessfully looked through the brms code to see how the predict functions work.

2

I’ve seen similar approaches using the following:

for (n in 1 : N) {
    zero=bernoulli_logit_rng(zi[n]);
            y_pred[n]=(1-zero)*neg_binomial_2_log_rng(mu[n], shape);
    }
3

Looks good. bernoulli_logit_rng(zi[n]) decides for zero or not with probability inv_logit(zi[n]).
If zero (-inflated), we are done, then y_pred[n] = 0. If it’s a non-zero we have to sample the
the remaining probability, here some negative binomial. (1-zero) is a clever way of saving
an if-clause.