Model specification of a zero-inflated negative binomial random effect model using brms

Dear Stan Users,

I am currently using stan via brms to model a zero-inflated negative binomial model with a random intercept and a random slope. As I am quite new to stan and brms, I am struggling to write down the model equation from the brms generated stan code. I hope you can support me here, as I would like to learn to understand and interpret the model specifications in stan / generated by brms.

The model I am running in R is defined as follows:


outcome ~ status + offset(log(N)) + (1+status|author)



The status is a dichotomous variable indicating disease status - author indicates the original study the observation belongs to, thus, observations are nested in studies.

The stan code that was generated using this specification looks as follows:


// generated with brms 2.14.4

functions {

/* turn a vector into a matrix of defined dimension

* stores elements in row major order

* Args:

*   X: a vector

*   N: first dimension of the desired matrix

*   K: second dimension of the desired matrix

* Returns:

*   a matrix of dimension N x K

*/

matrix as_matrix(vector X, int N, int K) {

matrix[N, K] Y;

for (i in 1:N) {

Y[i] = to_row_vector(X[((i - 1) * K + 1):(i * K)]);

}

return Y;

}

/* compute correlated group-level effects

* Args:

*   z: matrix of unscaled group-level effects

*   SD: vector of standard deviation parameters

*   L: cholesky factor correlation matrix

* Returns:

*   matrix of scaled group-level effects

*/

matrix scale_r_cor(matrix z, vector SD, matrix L) {

// r is stored in another dimension order than z

return transpose(diag_pre_multiply(SD, L) * z);

}

/* zero-inflated negative binomial log-PDF of a single response

* Args:

*   y: the response value

*   mu: mean parameter of negative binomial distribution

*   phi: shape parameter of negative binomial distribution

*   zi: zero-inflation probability

* Returns:

*   a scalar to be added to the log posterior

*/

real zero_inflated_neg_binomial_lpmf(int y, real mu, real phi,

real zi) {

if (y == 0) {

return log_sum_exp(bernoulli_lpmf(1 | zi),

bernoulli_lpmf(0 | zi) +

neg_binomial_2_lpmf(0 | mu, phi));

} else {

return bernoulli_lpmf(0 | zi) +

neg_binomial_2_lpmf(y | mu, phi);

}

}

/* zero-inflated negative binomial log-PDF of a single response

* logit parameterization of the zero-inflation part

* Args:

*   y: the response value

*   mu: mean parameter of 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_logit_lpmf(int y, real mu,

real phi, real zi) {

if (y == 0) {

return log_sum_exp(bernoulli_logit_lpmf(1 | zi),

bernoulli_logit_lpmf(0 | zi) +

neg_binomial_2_lpmf(0 | mu, phi));

} else {

return bernoulli_logit_lpmf(0 | zi) +

neg_binomial_2_lpmf(y | mu, phi);

}

}

/* zero-inflated negative binomial log-PDF of a single response

* log parameterization for the negative binomial part

* Args:

*   y: the response value

*   eta: linear predictor for negative binomial distribution

*   phi: shape parameter of negative binomial distribution

*   zi: zero-inflation probability

* Returns:

*   a scalar to be added to the log posterior

*/

real zero_inflated_neg_binomial_log_lpmf(int y, real eta,

real phi, real zi) {

if (y == 0) {

return log_sum_exp(bernoulli_lpmf(1 | zi),

bernoulli_lpmf(0 | zi) +

neg_binomial_2_log_lpmf(0 | eta, phi));

} else {

return bernoulli_lpmf(0 | zi) +

neg_binomial_2_log_lpmf(y | eta, phi);

}

}

/* 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);

}

}

// zero_inflated negative binomial log-CCDF and log-CDF functions

real zero_inflated_neg_binomial_lccdf(int y, real mu, real phi, real hu) {

return bernoulli_lpmf(0 | hu) + neg_binomial_2_lccdf(y | mu, phi);

}

real zero_inflated_neg_binomial_lcdf(int y, real mu, real phi, real hu) {

return log1m_exp(zero_inflated_neg_binomial_lccdf(y | mu, phi, hu));

}

}

data {

int<lower=1> N;  // total number of observations

int Y[N];  // response variable

int<lower=1> K;  // number of population-level effects

matrix[N, K] X;  // population-level design matrix

vector[N] offsets;

// data for group-level effects of ID 1

int<lower=1> N_1;  // number of grouping levels

int<lower=1> M_1;  // number of coefficients per level

int<lower=1> J_1[N];  // grouping indicator per observation

// group-level predictor values

vector[N] Z_1_1;

vector[N] Z_1_2;

int<lower=1> NC_1;  // number of group-level correlations

int prior_only;  // should the likelihood be ignored?

}

transformed data {

int Kc = K - 1;

matrix[N, Kc] Xc;  // centered version of X without an intercept

vector[Kc] means_X;  // column means of X before centering

for (i in 2:K) {

means_X[i - 1] = mean(X[, i]);

Xc[, i - 1] = X[, i] - means_X[i - 1];

}

}

parameters {

vector[Kc] b;  // population-level effects

real Intercept;  // temporary intercept for centered predictors

real<lower=0> shape;  // shape parameter

real<lower=0,upper=1> zi;  // zero-inflation probability

vector<lower=0>[M_1] sd_1;  // group-level standard deviations

matrix[M_1, N_1] z_1;  // standardized group-level effects

cholesky_factor_corr[M_1] L_1;  // cholesky factor of correlation matrix

}

transformed parameters {

matrix[N_1, M_1] r_1;  // actual group-level effects

// using vectors speeds up indexing in loops

vector[N_1] r_1_1;

vector[N_1] r_1_2;

// compute actual group-level effects

r_1 = scale_r_cor(z_1, sd_1, L_1);

r_1_1 = r_1[, 1];

r_1_2 = r_1[, 2];

}

model {

// likelihood including all constants

if (!prior_only) {

// initialize linear predictor term

vector[N] mu = Intercept + Xc * b + offsets;

for (n in 1:N) {

// add more terms to the linear predictor

mu[n] += r_1_1[J_1[n]] * Z_1_1[n] + r_1_2[J_1[n]] * Z_1_2[n];

}

for (n in 1:N) {

target += zero_inflated_neg_binomial_log_lpmf(Y[n] | mu[n], shape, zi);

}

}

// priors including all constants

target += normal_lpdf(b | 0, 5);

target += normal_lpdf(Intercept | 0, 5);

target += gamma_lpdf(shape | 0.01, 0.01);

target += beta_lpdf(zi | 1, 1);

target += student_t_lpdf(sd_1 | 3, 0, 2.5)

- 2 * student_t_lccdf(0 | 3, 0, 2.5);

target += std_normal_lpdf(to_vector(z_1));

target += lkj_corr_cholesky_lpdf(L_1 | 1);

}

generated quantities {

// actual population-level intercept

real b_Intercept = Intercept - dot_product(means_X, b);

// compute group-level correlations

corr_matrix[M_1] Cor_1 = multiply_lower_tri_self_transpose(L_1);

vector<lower=-1,upper=1>[NC_1] cor_1;

// extract upper diagonal of correlation matrix

for (k in 1:M_1) {

for (j in 1:(k - 1)) {

cor_1[choose(k - 1, 2) + j] = Cor_1[j, k];

}

}

}



What I understood from the code and some research in the stan forums, the model uses the log alternative paramterization with \eta = log(\mu) and shape parameter \phi.
However, I am not sure how to understand the part below “// add more terms to the linear predictor” and how the various priors are specified. Can anyone explain to me, what happens here?

If I write down the model with what I currently understand, It would look as follows:

\begin{align} p(y_n|\mu, \phi,\theta) &= \left\{ \begin{array}{ll} \theta + (1 - \theta) * \mathsf{NegativeBinomial2}(y_n|\mu, \phi) & \mbox{ if } y_n = 0, \mbox{ and} \\[3pt] (1-\theta) * \mathsf{NegativeBinomial2}(y_n|\mu, \phi) & \mbox{ if } y_n > 0, \end{array} \right.\\ \\ \log(\mu) &= \alpha + X\beta + offset(\log(N)), \\ \alpha &= \alpha_0 + \alpha_c, \\ \beta &= \beta_0 + \beta_c,\\ \alpha &\sim \text{Normal}(0, 5), \\ \beta &\sim \text{Normal}(0, 5), \\ \phi &\sim \text{Gamma}(0.01, 0.01), \\ \theta &\sim \text{Beta}(1,1) \end{align}

I would really appreciate some help or some hints on how to get a better intuition for stan code.

Sven

1 Like

Hi,
great question! brms can hide a lot of complex stuff from you, so it is good to make sure you understand it.

Here, brms encodes the varying intercept and varying effect. The idea is that for each level of author, you have an intercept and an effect. Those have a joint normal distribution. So we’ll add index i to index observations. You then have

\begin{align} \log(\mu_i) &= \alpha + X_i\beta + offset(\log(N_i)) + \gamma_{\mathrm{author[i]}, 1} + \gamma_{\mathrm{author[i]}, 2} Z_i \\ \gamma_k &\sim MVN(0, \Sigma) \\ \Sigma &= \begin{bmatrix} \sigma_1 \\ \sigma_2 \end{bmatrix} \begin{bmatrix} 1 & \rho \\ \rho & 1 \end{bmatrix} \begin{bmatrix} \sigma_1 & \sigma_2 \end{bmatrix} \end{align}

where \mathrm{author[i]} is the group index for author and Z_i is the predictor for the varying intercept (status in your case).

The formula \Sigma is just a decomposition of a covariance matrix into a correlation matrix and a vector of standard deviations.

You can check out Chapter 5 in BDA 3 (available online at http://www.stat.columbia.edu/~gelman/book/BDA3.pdf) for a bit more background.

Does that make sense?

1 Like

Hi Martin,

thanks a lot for your reply! I will try to dig into the chapter you recommended.

Looking at your equation, I feel that my lines

\alpha = \alpha_0 + \alpha_c,
\beta = \beta_0 + \beta_c,\\

are not necessary, am I correct?

Further, you only have Z_i as the predictor for the varying intercept, but the model says

      // add more terms to the linear predictor

mu[n] += r_1_1[J_1[n]] * Z_1_1[n] + r_1_2[J_1[n]] * Z_1_2[n];


with Z_1_1 and Z_1_2, both earlier specified as group level predictor values. Can you elaborate why here are two Zs and only one Z in the formula?

If I understand you correctly, theno only partially. brms will by default do some centering of the intercept (\alpha), so that bit is probably important, but the “effects” (\beta) are AFAIK not modified in any way.
Good question! I would expect Z_1_1 to correspond to the varying intercept per author and so to be a vector of ones (you can check this by inspecting the result of make_standata - I hope I am not wrong). You can definitely write your formula, with \gamma_{\mathrm{author[i]}, 1}Z_{1,i} + \gamma_{\mathrm{author[i]}, 2} Z_{2,i} or even treat it as a vector product to have just \mathbf{\gamma}_{\mathrm{author[i]}} Z_i. Definitely many options and it is more a question of taste than hard rules :-)