Zero-inflated counts — PyMC cross-check (ZIP & ZINB)¶

Companion to zip_python.ipynb (from-scratch augmentation Gibbs). PyMC has the mixture likelihoods built in (pm.ZeroInflatedPoisson, pm.ZeroInflatedNegativeBinomial) and marginalises the latent regime indicator analytically, sampling $\beta,\gamma,r$ with NUTS. We check it reproduces the from-scratch fit on the fish data.

Parameterisation note. PyMC's psi is the probability of the count (at-risk) regime, i.e. $\psi_i=1-\pi_i$. We model the structural-zero logit $\text{logit}\,\pi_i=z_i'\gamma$ and pass psi = 1 - sigmoid(Z@gamma), so $\gamma$ matches the from-scratch sampler directly.

In [1]:
import numpy as np, pandas as pd, pymc as pm, pytensor.tensor as pt, arviz as az
from zip_gibbs import zip_gibbs

d=pd.read_csv('fish.csv'); y=d['count'].to_numpy()
Xc=np.column_stack([np.ones(len(d)),d['child'],d['camper']])
Zc=np.column_stack([np.ones(len(d)),d['persons']])

def fit(negbin):
    with pm.Model() as m:
        beta=pm.Normal('beta',0,5,shape=3); gamma=pm.Normal('gamma',0,5,shape=2)
        mu=pt.exp(pt.dot(Xc,beta)); psi=1-pm.math.sigmoid(pt.dot(Zc,gamma))   # psi = P(at risk)
        if negbin:
            r=pm.Exponential('r',1.0)
            pm.ZeroInflatedNegativeBinomial('y',psi=psi,mu=mu,alpha=r,observed=y)
        else:
            pm.ZeroInflatedPoisson('y',psi=psi,mu=mu,observed=y)
        idata=pm.sample(2000,tune=2000,chains=4,target_accept=0.95,random_seed=1,progressbar=False)
    return idata

idz=fit(False); idn=fit(True)
def show(idata,tag,negbin):
    po=idata.posterior; b=po['beta'].mean(('chain','draw')).values; g=po['gamma'].mean(('chain','draw')).values
    vn=['beta','gamma']+(['r'] if negbin else [])
    mxrhat=float(az.summary(idata,var_names=vn)['r_hat'].max())
    print('=== %s ==='%tag)
    print('  count: intercept %.3f  child %.3f  camper %.3f'%(b[0],b[1],b[2]))
    print('  zero : intercept %.3f  persons %.3f'%(g[0],g[1]))
    if negbin: print('  r = %.3f'%float(po['r'].mean()))
    print('  max r_hat %.3f'%mxrhat)
show(idz,'PyMC ZIP',False); show(idn,'PyMC ZINB',True)

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [beta, gamma]

Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 6 seconds.

C:\Users\user\anaconda3\envs\pymc-env\Lib\site-packages\threadpoolctl.py:1226: RuntimeWarning: 
Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
the same time. Both libraries are known to be incompatible and this
can cause random crashes or deadlocks on Linux when loaded in the
same Python program.
Using threadpoolctl may cause crashes or deadlocks. For more
information and possible workarounds, please see
    https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md

  warnings.warn(msg, RuntimeWarning)
Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [beta, gamma, r]

Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 13 seconds.

=== PyMC ZIP ===
  count: intercept 1.592  child -1.047  camper 0.838
  zero : intercept 1.323  persons -0.581
  max r_hat 1.000
=== PyMC ZINB ===
  count: intercept 1.332  child -1.507  camper 0.884
  zero : intercept 2.461  persons -2.936
  r = 0.345
  max r_hat 1.000
In [2]:
# side-by-side with the from-scratch Gibbs
gz=zip_gibbs(y.astype(float),Xc,Zc,R=12000,burn=4000,seed=3)
gn=zip_gibbs(y.astype(float),Xc,Zc,R=12000,burn=4000,seed=3,negbin=True)
poz=idz.posterior; pon=idn.posterior
print('%-22s %10s %10s'%('parameter','from-scratch','PyMC'))
labs=['count_intercept','count_child','count_camper']
for j,l in enumerate(labs): print('%-22s %10.3f %10.3f'%('ZIP  '+l,gz['beta'].mean(0)[j],float(poz['beta'].mean(('chain','draw')).values[j])))
for j,l in enumerate(['zero_intercept','zero_persons']): print('%-22s %10.3f %10.3f'%('ZIP  '+l,gz['gamma'].mean(0)[j],float(poz['gamma'].mean(('chain','draw')).values[j])))
for j,l in enumerate(labs): print('%-22s %10.3f %10.3f'%('ZINB '+l,gn['beta'].mean(0)[j],float(pon['beta'].mean(('chain','draw')).values[j])))
print('%-22s %10.3f %10.3f'%('ZINB r',gn['r'].mean(),float(pon['r'].mean())))

parameter              from-scratch       PyMC
ZIP  count_intercept        1.600      1.592
ZIP  count_child           -1.048     -1.047
ZIP  count_camper           0.830      0.838
ZIP  zero_intercept         1.307      1.323
ZIP  zero_persons          -0.574     -0.581
ZINB count_intercept        1.357      1.332
ZINB count_child           -1.517     -1.507
ZINB count_camper           0.865      0.884
ZINB r                      0.350      0.345

Results¶

  • PyMC's built-in zero-inflated likelihoods (regime indicator marginalised, sampled by NUTS) reproduce the from-scratch augmentation-Gibbs coefficients: child ≈ −1.0 (ZIP) / −1.5 (ZINB), camper ≈ +0.85, persons (zero) < 0, and ZINB r ≈ 0.35. r̂ ≈ 1.00.
  • Same conclusion two ways: explicitly sampling the latent structural-zero indicator (from-scratch) vs analytically summing it out (PyMC) give the same posterior — the augmentation is a computational device, not a different model.
  • The ZINB count coefficients are tighter/larger in magnitude than ZIP because, once the NB soaks up the over-dispersion, the model attributes more of the low catches to children rather than noise.