Bayesian Zero-Inflated Count Regression

Model

A count-model class beyond the plain Poisson/NegBin of the count regression example: a mixture for data with far more zeros than any Poisson can produce. Many real count processes have two kinds of zero — structural (the unit was never "at risk": a park group that never fished) and sampling (at risk but happened to catch none). Zero-inflated models separate them via two linked regressions: a count model for how many and a logistic model for who is a structural zero. ZIP uses a Poisson at-risk component; ZINB uses a Negative Binomial, adding dispersion $r$ to also absorb overdispersion among the at-risk counts.

y_i = \begin{cases} 0 & \text{with prob } \pi_i \quad (\text{structural zero}) \\ \text{Poisson}(\lambda_i)\ \text{or NegBin}(\lambda_i, r) & \text{with prob } 1-\pi_i \quad (\text{at risk}) \end{cases}

\log\lambda_i = x_i'\beta \quad (\text{count: how many}), \qquad \text{logit}\,\pi_i = z_i'\gamma \quad (\text{zero: who is structural})

$\beta$ — count-model coefficients (log link, on the at-risk regime)
$\gamma$ — zero-inflation coefficients (logit link, who is a structural zero)
$r$ — NB dispersion (ZINB only); $P(y{=}0)=\pi + (1-\pi)f(0)$

Sampler — Gibbs with data augmentation

The mixture is awkward to sample directly, but data augmentation makes it trivial: introduce a latent regime indicator $s_i$ . Any $y_i > 0$ must be at-risk ( $s_i=0$ ); only the zeros are ambiguous, drawn as Bernoulli with probability $\pi_i / (\pi_i + (1-\pi_i)f(0))$ . Conditioning on $s$ decouples the mixture into two ordinary GLMs we already have — a logistic regression for $\gamma$ (the logit) and a count regression for $\beta$ (the Poisson/NB) — each updated by RW-Metropolis. The augmentation is the trick.

Block	Draw	Update
(1)	$s_i \mid \cdot$	Latent structural-zero indicator — Bernoulli (only ambiguous where $y_i=0$ )
(2)	$\beta\ (+\log r) \mid \cdot$	Count regression on the at-risk subset $\{s_i=0\}$ — RW-Metropolis
(3)	$\gamma \mid \cdot$	Logistic regression of $s$ on $Z$ — RW-Metropolis

Notebooks

Applied to the UCLA/PyMC fish example (250 park groups interviewed about a day's fishing): the catch is exactly the pathology zero-inflated models are built for — 57% zeros and a long overdispersed tail (mean 3.3, variance $\approx 135$ , max 149). The data demand zero-inflation (a plain Poisson cannot manufacture that many zeros) and ZINB fixes both the zeros and the overdispersion ( $r \approx 0.35$ , far from Poisson's $r=\infty$ ). Substantively: among those fishing, each child sharply cuts the expected catch — to roughly a third ( $e^{-1}\approx 0.37$ ) — and a camper raises it ~2.3×; bigger groups are less likely to be structural non-fishers. The zero-inflation and overdispersion stories are different and partly confounded — ZINB lets the data apportion the excess zeros between them. Cross-checked three ways: from-scratch augmentation Gibbs, PyMC (built-in zero-inflated likelihoods, regime indicator marginalised analytically), and R pscl::zeroinfl (ML) — all reproducing the canonical UCLA coefficients.

Python ZIP & ZINB from scratch — data-augmentation Gibbs, the fish example (Python) Open → Python PyMC cross-check — pm.ZeroInflatedPoisson / ZeroInflatedNegativeBinomial via NUTS (Python) Open → R R cross-check — pscl::zeroinfl maximum likelihood (R) Open →

Downloads

zip_gibbs.py From-scratch ZIP & ZINB — data-augmentation Gibbs sampler (NumPy / scipy) fish.csv UCLA/PyMC fish zero-inflation example — 250 park groups

References

Lambert, D. (1992). Zero-inflated Poisson regression, with an application to defects in manufacturing. Technometrics 34(1), 1–14. — the paper that introduced the ZIP model
Greene, W. H. (1994). Accounting for excess zeros and sample selection in Poisson and negative binomial regression models. NYU Working Paper EC-94-10. — the zero-inflated Negative Binomial (ZINB) extension
Tanner, M. A. & Wong, W. H. (1987). The calculation of posterior distributions by data augmentation. Journal of the American Statistical Association 82(398), 528–540. — the latent-indicator data-augmentation strategy behind the from-scratch sampler
Cameron, A. C. & Trivedi, P. K. (2013). Regression Analysis of Count Data (2nd ed.). Cambridge University Press. — zero-inflation, overdispersion, and the count-model context
Zeileis, A., Kleiber, C. & Jackman, S. (2008). Regression models for count data in R. Journal of Statistical Software 27(8), 1–25. — the pscl::zeroinfl implementation used in the R cross-check

Gibbs Sampler — Source Code

"""
Zero-inflated counts (ZIP and ZINB) -- from-scratch Gibbs with data augmentation.
The PyMC / UCLA "fish" example and BUGS Vol II "Hearts" family. The new idea vs count_reg_mcmc.py
(plain Poisson/NegBin) is a MIXTURE likelihood with a latent regime indicator:

  with prob pi_i        : structural zero   (y_i = 0 for sure -- e.g. group never fished)
  with prob 1 - pi_i    : "at risk"         (y_i ~ Poisson(lambda_i) or NegBin -- may still be 0)

so   P(y=0) = pi + (1-pi) f(0),   P(y=k>0) = (1-pi) f(k).
Two linked regressions:
  log lambda_i = x_i' beta     (count model, on the at-risk regime)
  logit pi_i   = z_i' gamma    (zero-inflation model, who is a structural zero)

Sampler (Gibbs with augmentation)
---------------------------------
  1. s_i  -- latent structural-zero indicator. If y_i>0 then s_i=0 (must be at-risk). If y_i=0,
            P(s_i=1) = pi_i / (pi_i + (1-pi_i) f(0)); draw Bernoulli.
  2. beta -- count regression on the at-risk subset {s_i=0}: RW-Metropolis (Poisson / NB loglik).
  3. gamma-- logistic regression of s on Z: RW-Metropolis.
  4. (ZINB only) log r -- NB dispersion, updated jointly with beta as a block.

Priors: beta, gamma ~ N(0, prior_sd^2); log r ~ N(0, 4).
"""
import numpy as np
from scipy.special import gammaln, expit


def simulate_zip(n=3000, beta=(0.5, 0.8, -0.4), gamma=(-0.4, 1.1), r=None, seed=0):
    """Zero-inflated counts. X=[1,x1,x2] for the count part, Z=[1,z1] for the zero part."""
    rng = np.random.default_rng(seed)
    beta = np.asarray(beta, float); gamma = np.asarray(gamma, float)
    x1 = rng.normal(size=n); x2 = rng.normal(size=n); z1 = rng.normal(size=n)
    X = np.column_stack([np.ones(n), x1, x2]); Z = np.column_stack([np.ones(n), z1])
    pi = expit(Z @ gamma); lam = np.exp(X @ beta)
    struct = rng.random(n) < pi                       # structural zeros
    if r is None:
        cnt = rng.poisson(lam)                        # at-risk counts (Poisson)
    else:
        cnt = rng.poisson(lam * rng.gamma(r, 1.0 / r, size=n))  # NB via gamma mixing
    y = np.where(struct, 0, cnt)
    return y, X, Z, struct


def _pois_mle(X, y, maxit=100):
    b = np.zeros(X.shape[1])
    for _ in range(maxit):
        mu = np.exp(X @ b)
        b += np.linalg.solve((X * mu[:, None]).T @ X + 1e-8 * np.eye(X.shape[1]), X.T @ (y - mu))
    return b


def _nb_log0(lam, r):
    """log P(y=0) under NegBin(mean=lam, size=r)."""
    return r * (np.log(r) - np.log(r + lam))


def zip_gibbs(y, X, Z, R=8000, burn=2000, prior_sd=5.0, seed=0,
              s_beta=None, s_gamma=None, negbin=False, a_r=2.0):
    """Gibbs for ZIP (negbin=False) or ZINB (negbin=True)."""
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float); n, k = X.shape; q = Z.shape[1]
    is0 = (y == 0)
    s_beta = (2.38 / np.sqrt(k)) if s_beta is None else s_beta
    s_gamma = (2.38 / np.sqrt(q)) if s_gamma is None else s_gamma

    # fixed proposal scales from crude full-data fits
    b0 = _pois_mle(X, y); Lb = np.linalg.cholesky(
        np.linalg.inv((X * np.exp(X @ b0)[:, None]).T @ X + np.eye(k) / prior_sd ** 2))
    w = 0.25
    Lg = np.linalg.cholesky(np.linalg.inv(w * Z.T @ Z + np.eye(q) / prior_sd ** 2))

    beta = b0.copy(); gamma = np.zeros(q); logr = 1.0
    keep = R - burn
    B = np.zeros((keep, k)); G = np.zeros((keep, q)); RR = np.zeros(keep)
    acc_b = 0; acc_g = 0
    for it in range(R):
        eta = X @ beta; lam = np.exp(eta); zg = Z @ gamma; pi = expit(zg)
        r = np.exp(logr) if negbin else None
        # 1. augment structural-zero indicator s (only ambiguous where y==0)
        logf0 = _nb_log0(lam, r) if negbin else -lam            # log P(at-risk count = 0)
        p_struct = pi / (pi + (1 - pi) * np.exp(logf0))         # P(s=1 | y=0)
        s = np.zeros(n)
        s[is0] = (rng.random(is0.sum()) < p_struct[is0]).astype(float)
        at = (s == 0)                                           # at-risk subset

        # 2. beta (+logr) on the at-risk subset
        Xa = X[at]; ya = y[at]
        def count_ll(bb, lr):
            e = Xa @ bb
            if not negbin:
                return np.sum(ya * e - np.exp(e)) - 0.5 * np.sum(bb ** 2) / prior_sd ** 2
            rr = np.exp(lr); mu = np.exp(e)
            return np.sum(gammaln(ya + rr) - gammaln(rr) + rr * np.log(rr) + ya * e
                          - (ya + rr) * np.log(rr + mu)) \
                - 0.5 * np.sum(bb ** 2) / prior_sd ** 2 - 0.5 * lr ** 2 / 4.0
        cur = count_ll(beta, logr)
        prop = beta + s_beta * (Lb @ rng.standard_normal(k))
        lrp = logr + (0.15 * rng.standard_normal() if negbin else 0.0)
        if np.log(rng.random()) < count_ll(prop, lrp) - cur:
            beta = prop; logr = lrp; acc_b += 1
        # 3. gamma: logistic regression of s on Z
        def logit_ll(gg):
            e = Z @ gg
            return np.sum(s * e - np.logaddexp(0, e)) - 0.5 * np.sum(gg ** 2) / prior_sd ** 2
        gp = gamma + s_gamma * (Lg @ rng.standard_normal(q))
        if np.log(rng.random()) < logit_ll(gp) - logit_ll(gamma):
            gamma = gp; acc_g += 1
        if it >= burn:
            i = it - burn; B[i] = beta; G[i] = gamma; RR[i] = np.exp(logr)
    out = dict(beta=B, gamma=G, accept_beta=acc_b / R, accept_gamma=acc_g / R)
    if negbin:
        out['r'] = RR
    return out


def summary(draws, names=None):
    d = draws.reshape(draws.shape[0], -1); q = d.shape[1]
    names = names or [f'p{j}' for j in range(q)]
    return [(nm, d[:, j].mean(), d[:, j].std(), *np.percentile(d[:, j], [2.5, 97.5]))
            for j, nm in enumerate(names)]