Zero-inflated counts — ZIP & ZINB from scratch

Data-augmentation Gibbs; the UCLA/PyMC "fish" example

Module: zip_gibbs.py. A new count-model class beyond count_reg (plain Poisson/NegBin): 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.

---

Model

Each observation comes from one of two regimes: $$y_i=\begin{cases}0 & \text{with prob }\pi_i\quad(\text{structural zero})\\[2pt] \text{Poisson}(\lambda_i)\ \text{or NegBin}(\lambda_i,r) & \text{with prob }1-\pi_i\quad(\text{at risk})\end{cases}$$ so $P(y_i=0)=\pi_i+(1-\pi_i)f(0)$ and $P(y_i=k>0)=(1-\pi_i)f(k)$. Two linked regressions: $$\log\lambda_i=x_i'\beta\ \ (\text{count: how many}),\qquad \text{logit}\,\pi_i=z_i'\gamma\ \ (\text{zero: who is a structural zero}).$$ ZIP uses a Poisson at-risk component; ZINB uses a Negative Binomial, adding dispersion $r$ to also absorb over-dispersion among the at-risk counts.

Algorithm — Gibbs with data augmentation

The mixture is awkward directly, but trivial once we augment a latent regime indicator $s_i$:

1. $s_i$ — structural-zero indicator. If $y_i>0$ then $s_i=0$ (must be at risk). If $y_i=0$, draw $s_i\sim\text{Bernoulli}\!\big(\tfrac{\pi_i}{\pi_i+(1-\pi_i)f(0)}\big)$.
2. $\beta$ (and $\log r$ for ZINB) — a plain count regression on the at-risk subset $\{s_i=0\}$: RW-Metropolis (Poisson/NB log-likelihood).
3. $\gamma$ — a plain logistic regression of $s$ on $Z$: RW-Metropolis. Conditioning on $s$ decouples the mixture into the two ordinary GLMs we already know — the augmentation is the trick.

import numpy as np, pandas as pd, matplotlib.pyplot as plt
from scipy.stats import nbinom
from zip_gibbs import simulate_zip, zip_gibbs, summary, _pois_mle

# ---- Section 1: synthetic recovery (ZIP then ZINB) ----
bt=(0.5,0.8,-0.4); gm=(-0.4,1.1)
y,X,Z,st=simulate_zip(n=4000,beta=bt,gamma=gm,seed=1)
oz=zip_gibbs(y,X,Z,R=8000,burn=2000,seed=1)
y2,X2,Z2,_=simulate_zip(n=4000,beta=bt,gamma=gm,r=2.0,seed=2)
on=zip_gibbs(y2,X2,Z2,R=10000,burn=3000,seed=2,negbin=True)
print('ZIP  synthetic (zeros %.0f%%, structural %.0f%%):'%(100*(y==0).mean(),100*st.mean()))
for nm,m,s,*_ in summary(oz['beta'],['cnt_int','cnt_b1','cnt_b2'])+summary(oz['gamma'],['zero_int','zero_g1']):
    print('   %-9s %7.3f (%.3f)'%(nm,m,s))
print('   truth beta (0.5,0.8,-0.4)  gamma (-0.4,1.1)')
print('ZINB synthetic (true r=2):')
for nm,m,s,*_ in summary(on['beta'],['cnt_int','cnt_b1','cnt_b2'])+summary(on['gamma'],['zero_int','zero_g1']):
    print('   %-9s %7.3f (%.3f)'%(nm,m,s))
print('   r %.3f (%.3f)  truth 2.0'%(on['r'].mean(),on['r'].std()))

ZIP  synthetic (zeros 58%, structural 43%):
   cnt_int     0.504 (0.021)
   cnt_b1      0.790 (0.014)
   cnt_b2     -0.415 (0.014)
   zero_int   -0.309 (0.050)
   zero_g1     1.116 (0.058)
   truth beta (0.5,0.8,-0.4)  gamma (-0.4,1.1)
ZINB synthetic (true r=2):
   cnt_int     0.537 (0.032)
   cnt_b1      0.796 (0.024)
   cnt_b2     -0.363 (0.023)
   zero_int   -0.312 (0.068)
   zero_g1     1.118 (0.071)
   r 2.381 (0.237)  truth 2.0

---

Section 2 — the fish data (UCLA / PyMC zero-inflation example)

Source. The UCLA OARC teaching dataset for zero-inflated count models: 250 groups leaving a state park, interviewed by wildlife biologists about their day's fishing. It is simulated (synthetic-but-realistic) — the columns xb and zg are the true linear predictors used to generate it (the log-rate and the zero-logit), so the data carries embedded ground truth, which is what makes it such a clean teaching example.

Response. count = fish caught — range 0–149, mean 3.30, variance ≈135: massively over-dispersed, with 57% zeros.

Covariates (standard spec count ~ child + camper | persons):

variable	meaning	spread	raw association
persons	people in the group (1–4)	57 / 70 / 57 / 66	bigger groups fish more: % zero falls 65% (1) → 48% (4)
child	children in the group (0–3)	132 / 75 / 33 / 10	0 kids → mean 5.19; ≥1 kid → 1.18
camper	brought a camper (0/1)	103 / 147	no camper 1.52; camper 4.54 (~3×)

child is always ≤ persons (children are a subset of the group). Two columns are not used: livebait (216/250 used it — little variation) and nofish — and notably nofish barely correlates with a zero catch (r ≈ 0.05), so despite its name it does not flag the structural non-fishers; that latent "didn't really fish" regime is what the model infers via $\pi_i$, not something read off a column.

The catch is exactly the pathology zero-inflated models are built for — far more zeros than any Poisson can make, and a long over-dispersed tail:

d=pd.read_csv('fish.csv'); y=d['count'].to_numpy(float)
print('n=%d   zeros=%.0f%%   mean=%.2f   variance=%.1f   max=%d'%(len(d),100*(y==0).mean(),y.mean(),y.var(),int(y.max())))
Xc=np.column_stack([np.ones(len(d)),d['child'],d['camper']])    # count part:  child, camper
Zc=np.column_stack([np.ones(len(d)),d['persons']])             # zero  part:  persons

zip_=zip_gibbs(y,Xc,Zc,R=12000,burn=4000,seed=3)
znb =zip_gibbs(y,Xc,Zc,R=12000,burn=4000,seed=3,negbin=True)
cn=['intercept','child','camper']; zn=['intercept','persons']
print('\nZIP  count model (log lambda):')
for nm,(_,m,s,*_q) in zip(cn,summary(zip_['beta'],cn)): print('   %-9s %7.3f (%.3f)'%(nm,m,s))
print('ZIP  zero model (logit pi):')
for nm,(_,m,s,*_q) in zip(zn,summary(zip_['gamma'],zn)): print('   %-9s %7.3f (%.3f)'%(nm,m,s))
print('\nZINB count model:')
for nm,(_,m,s,*_q) in zip(cn,summary(znb['beta'],cn)): print('   %-9s %7.3f (%.3f)'%(nm,m,s))
print('ZINB zero model:')
for nm,(_,m,s,*_q) in zip(zn,summary(znb['gamma'],zn)): print('   %-9s %7.3f (%.3f)'%(nm,m,s))
print('ZINB dispersion r = %.3f (%.3f)  -> strong over-dispersion (small r)'%(znb['r'].mean(),znb['r'].std()))

n=250   zeros=57%   mean=3.30   variance=134.8   max=149

ZIP  count model (log lambda):
   intercept   1.600 (0.088)
   child      -1.048 (0.098)
   camper      0.830 (0.096)
ZIP  zero model (logit pi):
   intercept   1.307 (0.385)
   persons    -0.574 (0.170)

ZINB count model:
   intercept   1.357 (0.278)
   child      -1.517 (0.210)
   camper      0.865 (0.287)
ZINB zero model:
   intercept   2.167 (1.109)
   persons    -2.440 (1.025)
ZINB dispersion r = 0.350 (0.062)  -> strong over-dispersion (small r)

# how well does each model reproduce the count distribution? (the zero-inflation diagnostic)
from scipy.stats import poisson as poi
bp=_pois_mle(Xc,y); lam_p=np.exp(Xc@bp)                                   # plain Poisson MLE
lam_zip=np.exp(Xc@zip_['beta'].mean(0)); pi_zip=1/(1+np.exp(-Zc@zip_['gamma'].mean(0)))
lam_znb=np.exp(Xc@znb['beta'].mean(0)); pi_znb=1/(1+np.exp(-Zc@znb['gamma'].mean(0))); r=znb['r'].mean()
cs=np.arange(0,11)
obs=np.array([(y==c).mean() for c in cs])
pP =np.array([poi.pmf(c,lam_p).mean() for c in cs])
pZ =np.array([(pi_zip*(c==0)+(1-pi_zip)*poi.pmf(c,lam_zip)).mean() for c in cs])
pN =np.array([(pi_znb*(c==0)+(1-pi_znb)*nbinom.pmf(c,r,r/(r+lam_znb))).mean() for c in cs])
print('predicted P(y=0): observed %.2f | Poisson %.2f | ZIP %.2f | ZINB %.2f'%(obs[0],pP[0],pZ[0],pN[0]))

fig,ax=plt.subplots(1,2,figsize=(12,4.4)); w=.2
ax[0].bar(cs-1.5*w,obs,w,label='observed',color='lightgray',edgecolor='k')
ax[0].bar(cs-0.5*w,pP,w,label='Poisson',color='firebrick',alpha=.8)
ax[0].bar(cs+0.5*w,pZ,w,label='ZIP',color='steelblue',alpha=.8)
ax[0].bar(cs+1.5*w,pN,w,label='ZINB',color='seagreen',alpha=.8)
ax[0].set_xlabel('fish caught'); ax[0].set_ylabel('probability'); ax[0].set_xticks(cs)
ax[0].set_title('Count distribution: observed vs fitted'); ax[0].legend()
# zoom on the tail (counts 1..10) where Poisson dies
ax[1].plot(cs[1:],obs[1:],'ko-',label='observed',lw=1)
ax[1].plot(cs[1:],pP[1:],'s--',color='firebrick',label='Poisson')
ax[1].plot(cs[1:],pZ[1:],'^--',color='steelblue',label='ZIP')
ax[1].plot(cs[1:],pN[1:],'d--',color='seagreen',label='ZINB')
ax[1].set_xlabel('fish caught'); ax[1].set_ylabel('probability'); ax[1].set_xticks(cs[1:])
ax[1].set_title('Positive tail (zoom)'); ax[1].legend()
plt.tight_layout(); plt.savefig('zip_fish.png',dpi=120,bbox_inches='tight'); plt.show()

predicted P(y=0): observed 0.57 | Poisson 0.22 | ZIP 0.54 | ZINB 0.58

Results & interpretation

• The data demand zero-inflation. Observed zeros = 57%; a plain Poisson predicts only ~22% (see printout) — it cannot manufacture that many zeros, and it cannot reach the long tail (max catch 149). ZIP fixes the zeros; ZINB fixes both the zeros and the over-dispersion ($r\approx0.35$, far from Poisson's $r=\infty$).
• Count model (how many fish, among those fishing): children reduce the catch (β_child ≈ −1.0 to −1.5: each child cuts the expected catch to roughly a third down to a fifth, $e^{-1}\approx0.37$ to $e^{-1.5}\approx0.22$) and a camper raises it ($e^{0.85}\approx2.3\times$ — proxies for a serious fishing trip).
• Zero model (who is a structural zero): more persons ⇒ lower $\pi$ (zero_persons < 0), i.e. bigger groups are less likely to be non-fishers — large parties came to fish.
• ZIP vs ZINB matters. ZIP still under-fits the positive tail because the at-risk counts are themselves over-dispersed; ZINB's NB component captures it. This mirrors the count_reg Poisson-vs-NegBin lesson, now layered on top of the zero mixture.

Takeaways

• Augmentation turns a mixture into two familiar GLMs. Drawing the regime indicator $s$ decouples the zero-inflated likelihood into a logistic regression (for $\gamma$) and a count regression (for $\beta$) — each handled by the samplers from earlier projects.
• Two zero stories. Zero-inflation (a separate non-fishing population) and over-dispersion (heavy tails among fishers) are different and partly confounded; ZINB lets the data apportion the excess zeros between them.
• Cross-checks: PyMC (pm.ZeroInflatedPoisson/NegativeBinomial) and R (pscl::zeroinfl) follow — both reproduce these coefficients.