How to fit zero-inflated SAR negative binomial models¶

import arviz as az
import matplotlib.pyplot as plt
import numpy as np

from neighbayes.dgp import simulate_sar_zinb
from neighbayes.models import SARZINB

az.style.use("arviz-darkgrid")
plt.rcParams["figure.figsize"] = (10, 6)

# Simulate data from the ZINB DGP
# Use n=50 (50x50 grid → 2500 observations) for reliable parameter recovery.
# Spatial parameters (ρ, λ) require n ≥ 900 for reliable recovery in
# cross-sectional count models. The model raises a UserWarning when n < 900.
data = simulate_sar_zinb(
    n=50,  # 50x50 grid → 2500 observations
    rho=0.4,  # Count SAR parameter
    lam=0.3,  # Selection SAR parameter
    alpha=5.0,  # NB dispersion (larger = less overdispersion)
    beta=np.array([1.0, 0.6]),  # Count coefficients [intercept, slope]
    gamma=np.array([0.3, 1.0]),  # Selection coefficients [intercept, slope]
    target_pi=0.7,  # Target corridor activation probability
    seed=42,
)

y = data["y"]
X = data["X"]
Z = data["Z"]
W_graph = data["W_graph"]
W_sel_graph = data["W_sel_graph"]

print(f"Observations: {len(y)}")
print(f"Zero fraction: {(y == 0).mean():.2%}")
print(f"Mean (y>0): {y[y > 0].mean():.1f}")
print(f"Max: {y.max():.0f}")
print("\nTrue parameters:")
print(f"  ρ (count SAR):  {data['params_true']['rho']}")
print(f"  λ (sel SAR):    {data['params_true']['lam']}")
print(f"  α (dispersion): {data['params_true']['alpha']}")
print(f"  β (count coef): {data['params_true']['beta']}")
print(f"  γ (sel coef):   {data['params_true']['gamma']}")

# Create and fit the ZINB-SAR model
model = SARZINB(
    y=y,
    X=X,
    Z=Z,
    W=W_graph,
    W_sel=W_sel_graph,
)

idata = model.fit(
    draws=2000,
    tune=1000,
    chains=4,
    random_seed=123,
    n_jobs=-1,
    progressbar=True,
)

# Posterior summary
summary = az.summary(idata, var_names=["rho", "lam", "alpha", "beta", "gamma"])
print(summary)

# Compare with true values
print("\n--- Parameter Recovery ---")
print(
    f"ρ (count SAR):  true={data['params_true']['rho']:.2f}, "
    f"posterior mean={float(idata.posterior['rho'].mean()):.3f} "
    f"± {float(idata.posterior['rho'].std()):.3f}"
)
print(
    f"λ (sel SAR):    true={data['params_true']['lam']:.2f}, "
    f"posterior mean={float(idata.posterior['lam'].mean()):.3f} "
    f"± {float(idata.posterior['lam'].std()):.3f}"
)
print(
    f"α (dispersion): true={data['params_true']['alpha']:.2f}, "
    f"posterior mean={float(idata.posterior['alpha'].mean()):.3f} "
    f"± {float(idata.posterior['alpha'].std()):.3f}"
)

# Note: spatial parameters (ρ, λ) require large samples (n ≥ 900)
# for reliable recovery in cross-sectional count models. The posterior
# is typically wide and may be attenuated toward zero at moderate n.
# The model raises a UserWarning when n < 900.

# Convergence diagnostics
az.plot_trace(idata, var_names=["rho", "lam", "alpha"], compact=False)
plt.tight_layout()
plt.show()

# R-hat and ESS
diagnostics = az.summary(idata, var_names=["rho", "lam", "alpha"])
print(diagnostics[["r_hat", "ess_bulk", "ess_tail"]])

# Zero attribution
attribution = model.zero_attribution()

print(f"Number of zero observations: {len(attribution['zero_indices'])}")
print(f"\nMean P(structural | y=0): {attribution['structural_prob'].mean():.3f}")
print(f"Mean P(sampling | y=0):   {attribution['sampling_prob'].mean():.3f}")

# Histogram of structural zero probabilities
fig, ax = plt.subplots(1, 2, figsize=(12, 4))
ax[0].hist(attribution["structural_prob"], bins=20, edgecolor="white")
ax[0].set_xlabel("P(structural zero | y=0)")
ax[0].set_ylabel("Count")
ax[0].set_title("Structural Zero Probabilities")

ax[1].hist(attribution["sampling_prob"], bins=20, edgecolor="white", color="C1")
ax[1].set_xlabel("P(sampling zero | y=0)")
ax[1].set_ylabel("Count")
ax[1].set_title("Sampling Zero Probabilities")
plt.tight_layout()
plt.show()

# Corridor activation probabilities
pi = model.corridor_probabilities()

print(f"Mean corridor probability: {pi.mean():.3f}")
print(f"Min corridor probability:   {pi.min():.3f}")
print(f"Max corridor probability:   {pi.max():.3f}")

# Compare with true d
d_true = data["d"]
print(f"\nTrue activation rate: {d_true.mean():.3f}")
print(f"Estimated activation rate: {pi.mean():.3f}")

# Scatter: pi vs observed y
fig, ax = plt.subplots(figsize=(8, 5))
ax.scatter(pi, y, alpha=0.5, s=20)
ax.set_xlabel("Estimated P(active)")
ax.set_ylabel("Observed count y")
ax.set_title("Corridor Activation Probability vs Observed Count")
ax.axhline(0, color="grey", linestyle="--", alpha=0.5)
plt.tight_layout()
plt.show()

# Fitted mean counts
fitted = model.fitted_values()

print(f"Mean fitted count: {fitted.mean():.2f}")
print(f"Mean observed count: {y.mean():.2f}")

# Scatter: fitted vs observed
fig, ax = plt.subplots(figsize=(8, 5))
ax.scatter(fitted, y, alpha=0.5, s=20)
ax.plot(
    [0, max(y.max(), fitted.max())],
    [0, max(y.max(), fitted.max())],
    "r--",
    alpha=0.5,
    label="45° line",
)
ax.set_xlabel("Fitted E[y]")
ax.set_ylabel("Observed y")
ax.set_title("Fitted vs Observed Counts")
ax.legend()
plt.tight_layout()
plt.show()

# Posterior distributions with true values
fig, axes = plt.subplots(1, 3, figsize=(14, 4))

# ρ (count SAR)
az.plot_posterior(
    idata, var_names=["rho"], ref_val=data["params_true"]["rho"], ax=axes[0]
)
axes[0].set_title("ρ (Count SAR)")

# λ (selection SAR)
az.plot_posterior(
    idata, var_names=["lam"], ref_val=data["params_true"]["lam"], ax=axes[1]
)
axes[1].set_title("λ (Selection SAR)")

# α (NB dispersion)
az.plot_posterior(
    idata, var_names=["alpha"], ref_val=data["params_true"]["alpha"], ax=axes[2]
)
axes[2].set_title("α (NB Dispersion)")

plt.tight_layout()
plt.show()

# Note: spatial parameters (ρ, λ) require large samples (n ≥ 900)
# for reliable recovery in cross-sectional count models. The posterior
# is typically wide and may be attenuated toward zero at moderate n.