Count regression from scratch — Poisson & Negative Binomial¶
Models, algorithms (MLE vs Bayes), and the overdispersion comparison¶
From-scratch module: count_reg_mcmc.py. Single-level (no hierarchy) count GLMs, fit two ways — maximum likelihood (Newton/IRLS) and Bayesian (random-walk Metropolis) — so we can compare point estimates, standard errors, and the Poisson-vs-NegBin choice.
| Section | Content |
|---|---|
| Models | Poisson (equidispersion) and Negative Binomial (overdispersion) |
| Algorithms | MLE via Newton/IRLS; Bayes via RW-Metropolis; what we compare |
| Section 1 | Poisson on equidispersed data — MLE ≈ Bayes ≈ truth |
| Section 2 | Overdispersed data — Poisson understates uncertainty; NegBin fixes it |
Models¶
Poisson regression — the baseline count GLM with a log link: $$y_i \sim \text{Poisson}(\mu_i),\qquad \mu_i = \exp(x_i'\beta),\qquad \mathbb{E}[y_i]=\text{Var}[y_i]=\mu_i.$$ Its defining property is equidispersion: conditional mean = conditional variance. Real counts are often more variable than this.
Negative Binomial (NB2) — relaxes that by adding a dispersion parameter $r$ (the "size"): $$y_i \sim \text{NB}(\mu_i, r),\qquad \mu_i=\exp(x_i'\beta),\qquad \text{Var}[y_i]=\mu_i + \frac{\mu_i^2}{r}.$$ NB2 is a gamma–Poisson mixture: $y_i\mid\lambda_i\sim\text{Poisson}(\lambda_i)$, $\lambda_i\sim\text{Gamma}(r,\mu_i/r)$. As $r\to\infty$ the extra variance vanishes and NB2 → Poisson. Finite $r$ ⇒ overdispersion.
Diagnostic. The Pearson dispersion $\hat\phi=\frac{1}{n-k}\sum_i (y_i-\hat\mu_i)^2/\hat\mu_i$ is ≈ 1 under Poisson and > 1 under overdispersion. (Note: the marginal variance of $y$ always exceeds its mean because $\mu_i$ varies with $x_i$ — dispersion must be judged conditionally, hence the Pearson statistic.)
Algorithms¶
Maximum likelihood (Newton–Raphson / IRLS). For the Poisson the canonical log link makes the observed and expected information coincide, so Newton's method is iteratively reweighted least squares: $$\text{score } g=X'(y-\mu),\quad \text{information } \mathcal I = X'\,\text{diag}(\mu)\,X,\quad \beta\leftarrow\beta+\mathcal I^{-1}g.$$ Asymptotic standard errors come from $\widehat{\text{Var}}(\hat\beta)=\mathcal I^{-1}$. The NB2 optimises $(\beta,\log r)$ jointly (BFGS), with SEs from the inverse Hessian.
Bayesian (random-walk Metropolis). Prior $\beta\sim N(0,\,\sigma_p^2 I)$ (vague). Proposals are drawn from $N\!\big(\beta_{\text{cur}},\, s^2 (\mathcal I+\text{prior prec})^{-1}\big)$ — i.e. scaled by the posterior curvature at the MLE, with $s=2.38/\sqrt{d}$ (Roberts' optimal scale) giving ~25–40% acceptance. For NB2 the step also covers $\log r$.
What we compare. With a vague prior and large $n$, the Bernstein–von Mises theorem says the posterior ≈ $N(\hat\beta_{\text{MLE}},\,\mathcal I^{-1})$ — so Bayes posterior mean ≈ MLE and posterior SD ≈ asymptotic SE. We verify this, then show where Poisson and NegBin diverge (the standard errors under overdispersion).
import numpy as np, matplotlib.pyplot as plt
from count_reg_mcmc import (simulate_counts, poisson_mle, poisson_rwm,
negbin_mle, negbin_rwm, summary)
def pearson_dispersion(X, y, beta):
mu = np.exp(X @ beta)
return np.sum((y - mu) ** 2 / mu) / (len(y) - X.shape[1])
print('count_reg_mcmc loaded')
count_reg_mcmc loaded
1. Poisson on equidispersed data¶
Simulate $n=2000$ Poisson counts with $\beta=(0.5,\,0.8,\,-0.4)$. We expect Pearson dispersion ≈ 1, and MLE ≈ Bayes ≈ truth.
beta_true = np.array([0.5, 0.8, -0.4])
Xp, yp = simulate_counts(2000, beta_true, r=None, seed=1)
b_mle, se_mle, _ = poisson_mle(Xp, yp)
fit = poisson_rwm(yp, Xp, R=8000, burn=2000, seed=2)
b_bayes, sd_bayes = fit['beta'].mean(0), fit['beta'].std(0)
print(f"Pearson dispersion = {pearson_dispersion(Xp, yp, b_mle):.2f} (~1 => equidispersed)\n")
print(f"{'coef':>6}{'truth':>8}{'MLE':>9}{'SE':>8}{'Bayes':>9}{'postSD':>8}")
for j in range(3):
print(f"{'b'+str(j):>6}{beta_true[j]:>8.2f}{b_mle[j]:>9.3f}{se_mle[j]:>8.3f}{b_bayes[j]:>9.3f}{sd_bayes[j]:>8.3f}")
print(f"\nMetropolis acceptance = {fit['accept']:.2f}")
Pearson dispersion = 1.02 (~1 => equidispersed)
coef truth MLE SE Bayes postSD
b0 0.50 0.531 0.019 0.531 0.018
b1 0.80 0.807 0.014 0.807 0.014
b2 -0.40 -0.387 0.014 -0.387 0.014
Metropolis acceptance = 0.32
# MLE (point ± 1.96 SE) vs Bayes posterior (density) for each coefficient
fig, ax = plt.subplots(1, 3, figsize=(13, 3.6))
for j in range(3):
d = fit['beta'][:, j]
ax[j].hist(d, bins=50, density=True, color='lightsteelblue', label='Bayes posterior')
ax[j].axvline(b_mle[j], color='firebrick', lw=2, label='MLE')
ax[j].axvspan(b_mle[j]-1.96*se_mle[j], b_mle[j]+1.96*se_mle[j], color='firebrick', alpha=0.15)
ax[j].axvline(beta_true[j], color='k', ls='--', lw=1.5, label='truth')
ax[j].set_title(f'b{j}'); ax[j].set_yticks([])
ax[0].legend(fontsize=8)
plt.suptitle('Poisson: Bayes posterior ≈ MLE ± SE (large-n agreement)')
plt.tight_layout(); plt.savefig('count_poisson.png', dpi=120, bbox_inches='tight'); plt.show()
1. Results¶
| coef | truth | MLE (SE) | Bayes (postSD) |
|---|---|---|---|
| b0 | 0.50 | 0.531 (0.019) | 0.531 (0.018) |
| b1 | 0.80 | 0.807 (0.014) | 0.807 (0.014) |
| b2 | −0.40 | −0.387 (0.014) | −0.387 (0.014) |
Pearson dispersion ≈ 1.02 (equidispersed, as it should be). MLE and Bayes coincide — same point estimates and the asymptotic SE equals the posterior SD to three decimals — the Bernstein–von Mises agreement, with the posterior densities sitting right on the MLE ± SE bands. Metropolis acceptance ≈ 0.32 (healthy).
2. Overdispersed data: Poisson vs Negative Binomial¶
Now generate overdispersed counts (true $r=2$, so $\text{Var}=\mu+\mu^2/2$). The mean function is unchanged, but the variance is much larger. Fit Poisson first (it will look fine for $\beta$ but understate uncertainty), then NegBin.
Xo, yo = simulate_counts(2000, beta_true, r=2.0, seed=3)
# Poisson fit (mis-specified — ignores overdispersion)
bp, sep, _ = poisson_mle(Xo, yo)
disp = pearson_dispersion(Xo, yo, bp)
# Negative Binomial fit (MLE + Bayes)
bn, lrn, sen = negbin_mle(Xo, yo)
nb = negbin_rwm(yo, Xo, R=8000, burn=2000, seed=4)
print(f"Pearson dispersion under Poisson = {disp:.2f} (>>1 => Poisson is wrong)\n")
print(f"{'coef':>6}{'truth':>8}{'Pois MLE':>11}{'Pois SE':>9}{'NB MLE':>10}{'NB SE':>9}{'NB Bayes':>11}{'NB postSD':>11}")
for j in range(3):
print(f"{'b'+str(j):>6}{beta_true[j]:>8.2f}{bp[j]:>11.3f}{sep[j]:>9.3f}{bn[j]:>10.3f}{sen[j]:>9.3f}"
f"{nb['beta'].mean(0)[j]:>11.3f}{nb['beta'].std(0)[j]:>11.3f}")
print(f"\ndispersion r: truth 2.0 NB MLE {np.exp(lrn):.2f} NB Bayes {nb['r'].mean():.2f} "
f"[{np.percentile(nb['r'],2.5):.2f}, {np.percentile(nb['r'],97.5):.2f}]")
print(f"slope SE inflation NegBin/Poisson = {sen[1]/sep[1]:.2f}x (Poisson understated it)")
Pearson dispersion under Poisson = 2.17 (>>1 => Poisson is wrong)
coef truth Pois MLE Pois SE NB MLE NB SE NB Bayes NB postSD
b0 0.50 0.496 0.019 0.491 0.025 0.491 0.026
b1 0.80 0.723 0.015 0.748 0.025 0.749 0.025
b2 -0.40 -0.461 0.014 -0.432 0.023 -0.432 0.022
dispersion r: truth 2.0 NB MLE 2.10 NB Bayes 2.10 [1.83, 2.39]
slope SE inflation NegBin/Poisson = 1.73x (Poisson understated it)
fig, ax = plt.subplots(1, 2, figsize=(12, 4.2))
# (a) SE comparison: Poisson understates, NegBin corrects
x = np.arange(3); w = 0.38
ax[0].bar(x - w/2, sep, w, color='steelblue', label='Poisson SE')
ax[0].bar(x + w/2, sen[:3], w, color='firebrick', label='NegBin SE')
ax[0].set_xticks(x); ax[0].set_xticklabels(['b0','b1','b2']); ax[0].set_ylabel('standard error')
ax[0].set_title('Poisson understates SEs under overdispersion'); ax[0].legend()
# (b) posterior of the dispersion r
ax[1].hist(nb['r'], bins=50, density=True, color='lightsteelblue')
ax[1].axvline(2.0, color='k', ls='--', lw=1.5, label='true r=2')
ax[1].axvline(nb['r'].mean(), color='firebrick', lw=2, label=f"post. mean {nb['r'].mean():.2f}")
ax[1].set_xlabel('dispersion r'); ax[1].set_yticks([]); ax[1].set_title('NegBin recovers the dispersion'); ax[1].legend()
plt.tight_layout(); plt.savefig('count_negbin.png', dpi=120, bbox_inches='tight'); plt.show()
2. Results¶
| coef | truth | Poisson MLE (SE) | NegBin MLE (SE) | NegBin Bayes (postSD) |
|---|---|---|---|---|
| b0 | 0.50 | 0.496 (0.019) | 0.491 (0.025) | 0.491 (0.026) |
| b1 | 0.80 | 0.723 (0.015) | 0.748 (0.025) | 0.749 (0.025) |
| b2 | −0.40 | −0.461 (0.014) | −0.432 (0.023) | −0.432 (0.022) |
Pearson dispersion ≈ 2.17 (≫1) flags that Poisson is wrong. Both models give roughly the right $\beta$ (the mean function is still consistent under Poisson), but the Poisson standard errors are ~1.7× too small — it would produce overconfident intervals and spurious significance. The Negative Binomial recovers the dispersion ($r\approx2.1$, CrI [1.83, 2.39], covering the truth 2) and reports honest, wider SEs. NegBin MLE and Bayes agree (acc ≈ 0.22).
Takeaways¶
- MLE ↔ Bayes: with a vague prior and $n=2000$ they're indistinguishable (Bernstein–von Mises) — the from-scratch RW-Metropolis reproduces the IRLS/Newton fit and its SEs.
- Poisson vs NegBin: always check the Pearson dispersion. Under overdispersion Poisson keeps $\beta$ roughly right but understates uncertainty; NegBin adds one parameter $r$ and restores calibrated inference.
- Next steps (separate notebooks): a PyMC cross-check and an R (
glm/MASS::glm.nb) benchmark, then a real count dataset.
3. Real data: the Epil epilepsy trial (overdispersed counts)¶
Data (epil.csv). Thall & Vail (1990); BUGS Vol I "Epil". A randomized, double-blind, placebo-controlled trial of the anti-epileptic progabide: 59 patients (31 progabide / 28 placebo), an 8-week baseline seizure count, then 4 two-week visits with seizure counts — 236 observations.
| column | meaning |
|---|---|
seizures (y) |
2-week seizure count — the response (0–102, mean ≈ 8.3) |
Trt |
1 = progabide, 0 = placebo (the effect of interest) |
Base |
8-week baseline seizures (entered as log(Base/4)) |
Age |
age in years (entered as log(Age)) |
V4 |
4th-visit indicator |
Counts this skewed are almost always overdispersed — exactly where the Poisson/NegBin choice bites.
import pandas as pd
ep = pd.read_csv('epil.csv')
y_e = ep['seizures'].values.astype(float)
X_e = np.column_stack([np.ones(len(ep)), np.log(ep['Base']/4), ep['Trt'], np.log(ep['Age']), ep['V4']])
enames = ['intercept', 'log(Base/4)', 'Trt', 'log(Age)', 'V4']
bp_e, sep_e, _ = poisson_mle(X_e, y_e)
bn_e, lrn_e, sen_e = negbin_mle(X_e, y_e)
nb_e = negbin_rwm(y_e, X_e, R=8000, burn=2000, seed=8)
print(f"n={len(y_e)} Pearson dispersion (Poisson) = {pearson_dispersion(X_e, y_e, bp_e):.2f} (>>1 => overdispersed)\n")
print(f"{'coef':>12}{'Pois MLE':>11}{'(SE)':>8}{'NB MLE':>10}{'(SE)':>8}{'NB Bayes':>11}{'(SD)':>8}")
for j, nm in enumerate(enames):
print(f"{nm:>12}{bp_e[j]:>11.3f}{sep_e[j]:>8.3f}{bn_e[j]:>10.3f}{sen_e[j]:>8.3f}"
f"{nb_e['beta'].mean(0)[j]:>11.3f}{nb_e['beta'].std(0)[j]:>8.3f}")
print(f"\ndispersion r: NB MLE {np.exp(lrn_e):.2f} NB Bayes {nb_e['r'].mean():.2f} "
f"[{np.percentile(nb_e['r'],2.5):.2f}, {np.percentile(nb_e['r'],97.5):.2f}] (acc {nb_e['accept']:.2f})")
ti = enames.index('Trt')
print(f"\nTREATMENT (progabide):")
print(f" Poisson : Trt={bp_e[ti]:+.3f} SE={sep_e[ti]:.3f} z={bp_e[ti]/sep_e[ti]:+.2f} -> rate ratio {np.exp(bp_e[ti]):.2f}")
print(f" NegBin : Trt={bn_e[ti]:+.3f} SE={sen_e[ti]:.3f} z={bn_e[ti]/sen_e[ti]:+.2f} -> rate ratio {np.exp(bn_e[ti]):.2f}")
n=236 Pearson dispersion (Poisson) = 4.71 (>>1 => overdispersed)
coef Pois MLE (SE) NB MLE (SE) NB Bayes (SD)
intercept -2.340 0.404 -1.184 0.817 -1.138 0.854
log(Base/4) 1.224 0.033 1.060 0.064 1.067 0.063
Trt -0.017 0.048 -0.233 0.101 -0.224 0.103
log(Age) 0.579 0.110 0.369 0.233 0.350 0.245
V4 -0.160 0.055 -0.152 0.115 -0.145 0.122
dispersion r: NB MLE 2.63 NB Bayes 2.59 [1.91, 3.32] (acc 0.28)
TREATMENT (progabide):
Poisson : Trt=-0.017 SE=0.048 z=-0.35 -> rate ratio 0.98
NegBin : Trt=-0.233 SE=0.101 z=-2.31 -> rate ratio 0.79
fig, ax = plt.subplots(1, 2, figsize=(12, 4.2))
# (a) coefficient estimates +-1.96 SE: Poisson vs NegBin (SEs widen; Trt shifts)
yp = np.arange(5)
ax[0].errorbar(bp_e, yp+0.12, xerr=1.96*sep_e, fmt='o', color='steelblue', capsize=3, label='Poisson')
ax[0].errorbar(bn_e, yp-0.12, xerr=1.96*sen_e[:5], fmt='s', color='firebrick', capsize=3, label='NegBin')
ax[0].axvline(0, color='gray', lw=.8, ls=':')
ax[0].set_yticks(yp); ax[0].set_yticklabels(enames); ax[0].invert_yaxis()
ax[0].set_xlabel('coefficient (95% CI)'); ax[0].set_title('Poisson vs NegBin: wider, shifted estimates'); ax[0].legend()
# (b) dispersion r posterior
ax[1].hist(nb_e['r'], bins=50, density=True, color='lightsteelblue')
ax[1].axvline(nb_e['r'].mean(), color='firebrick', lw=2, label=f"post. mean {nb_e['r'].mean():.2f}")
ax[1].set_xlabel('dispersion r'); ax[1].set_yticks([]); ax[1].set_title('Strong overdispersion: r ≈ 2.6 (finite)'); ax[1].legend()
plt.tight_layout(); plt.savefig('count_epil.png', dpi=120, bbox_inches='tight'); plt.show()
3. Posterior probability that treatment reduces seizures¶
A key payoff of the Bayesian fit: the posterior gives P(Trt < 0 | data) = P(rate ratio < 1) — the probability that progabide lowers the seizure rate — instead of a yes/no significance verdict.
| model | P(reduces seizures) | median rate ratio | 95% CrI on RR |
|---|---|---|---|
| NegBin (well-specified) | 0.99 | 0.79 | [0.65, 0.98] |
| Poisson (normal/BvM approx) | 0.63 | 0.98 | — |
The NegBin model is ~99% sure progabide reduces seizures (about a 21% rate cut); the Poisson model is barely above a coin flip — its near-null treatment estimate is the single-outlier artifact (patient #49) diagnosed below.
# NegBin RW-Metropolis draws (nb_e) + Poisson MLE (BvM normal posterior):
trt_nb = nb_e["beta"][:, enames.index("Trt")] # posterior draws of the Trt coefficient
P_reduces = (trt_nb < 0).mean() # = P(rate ratio < 1) -> 0.99
3. Results — Epil¶
| coef | Poisson MLE (SE) | NegBin MLE (SE) | NegBin Bayes (SD) |
|---|---|---|---|
| intercept | −2.340 (0.404) | −1.184 (0.817) | −1.14 (0.85) |
| log(Base/4) | 1.224 (0.033) | 1.060 (0.064) | 1.07 (0.06) |
| Trt | −0.017 (0.048) | −0.233 (0.101) | −0.224 (0.103) |
| log(Age) | 0.579 (0.110) | 0.369 (0.233) | 0.35 (0.25) |
| V4 | −0.160 (0.055) | −0.152 (0.115) | −0.15 (0.12) |
Pearson dispersion ≈ 4.71 — Poisson is badly mis-specified. NegBin recovers a clean r ≈ 2.6 (Bayes [1.9, 3.3]), and MLE ≈ Bayes (acc ≈ 0.28).
This real example is richer than the synthetic one — two things happen at once:
- SEs roughly double going Poisson → NegBin (e.g. log(Base/4) 0.033 → 0.064; log(Age) 0.110 → 0.233) — the usual overdispersion correction.
- Point estimates also move, unlike the textbook case. Most strikingly the treatment effect: Poisson gives Trt ≈ 0 (rate ratio 0.98, z = −0.35, "no effect"), while NegBin gives Trt ≈ −0.23 (rate ratio 0.79, a 21% seizure reduction, z = −2.3). The two models weight observations differently — NegBin down-weights the high-variance high-count patients (e.g. the outlier patient 49) — so they disagree on more than just uncertainty.
The honest caveat (and the bridge to the next project). Epil has 4 correlated visits per patient, which this single-level model ignores. So the single-level NegBin's apparent treatment significance is not yet trustworthy — the effective sample size is inflated by pseudo-replication. The correct analysis adds patient-level random effects (a hierarchical / mixed Poisson-NB model — the original Breslow–Clayton 1993 analysis). That is the natural hierarchical-count follow-up; the BUGS "Epil" example is exactly that model.
Takeaway. On real counts, always check the Pearson dispersion; under overdispersion Poisson understates SEs and (with influential points) can shift estimates; NegBin restores calibrated single-level inference — but repeated measures still demand a hierarchical model.
# Why do Poisson and NegBin disagree on the treatment? Influence of high-count patients.
def trt_fit(dd):
Xd=np.column_stack([np.ones(len(dd)),np.log(dd['Base']/4),dd['Trt'],np.log(dd['Age']),dd['V4']])
yd=dd['seizures'].values.astype(float)
return poisson_mle(Xd,yd)[0][2], negbin_mle(Xd,yd)[0][2]
print('arm balance (mean baseline):', {int(k):round(v,1) for k,v in ep.groupby('Trt')['Base'].mean().items()})
print(f"\n{'subset':<22}{'Poisson Trt':>13}{'NegBin Trt':>12}")
for lab,sub in [('full sample',ep),('drop patient 49',ep[ep.patient!=49]),
('drop top-5% counts',ep[ep.seizures<=ep.seizures.quantile(.95)])]:
p,n=trt_fit(sub); print(f"{lab:<22}{p:>13.3f}{n:>12.3f}")
r_hat=nb_e['r'].mean()
print(f"\noverdispersion at mu=8: Poisson Var=8 (SD 2.8) vs NegBin Var={8+64/r_hat:.0f} (SD {(8+64/r_hat)**.5:.1f})")
arm balance (mean baseline): {0: 30.8, 1: 31.6}
subset Poisson Trt NegBin Trt
full sample -0.017 -0.233
drop patient 49 -0.231 -0.284
drop top-5% counts -0.164 -0.217
overdispersion at mu=8: Poisson Var=8 (SD 2.8) vs NegBin Var=33 (SD 5.7)
3. Analysis — what the Epil results mean¶
Coefficients (from the NegBin fit, the trustworthy single-level model), on the log-rate scale.
- log(Base/4) ≈ 1.06 — near-proportional: a patient's follow-up seizure rate tracks their own 8-week baseline almost 1-for-1. Baseline is by far the dominant predictor — patients regress toward their personal level, not a common mean.
- Trt ≈ −0.23 → rate ratio e^(−0.23) = 0.79: progabide is associated with roughly a 21% reduction in seizure frequency.
- log(Age) ≈ 0.37 — older patients have somewhat higher rates (weak; wide CI).
- V4 ≈ −0.15 — a mild downward drift by the 4th visit.
Why Poisson and NegBin disagree on the treatment — and it comes down to one patient. The gap (Poisson Trt ≈ 0 vs NegBin −0.23) is not noise. A Negative-Binomial GLM down-weights high-count observations (its score carries a factor r/(r+μ)); Poisson weights every observation equally. Patient 49 — baseline 151, randomized to progabide yet with persistently high counts — receives outsized weight under Poisson and single-handedly cancels the treatment signal: dropping just that one patient moves the Poisson estimate from −0.02 to −0.23, matching NegBin, while NegBin barely budges (it had already down-weighted him). The arms are balanced on baseline (≈30.8 vs 31.6), so this is not confounding — it is variance-weighting. Poisson's apparent "no effect" is an artifact of an influential outlier; NegBin (and leave-one-out Poisson) agree the effect is ≈ −0.23.
The overdispersion is large. Pearson dispersion 4.71, r ≈ 2.6 ⇒ at the average count (μ ≈ 8) the SD is ≈ 5.7 — double Poisson's 2.8 (4× the variance). It reflects genuine patient-to-patient heterogeneity (beyond baseline/age), within-patient correlation across the 4 visits, and temporal clustering of seizures.
Treatment conclusion — suggestive, not definitive. Single-level NegBin gives RR 0.79, 95% CI [0.65, 0.96] — a ~21% reduction that is robust to the outlier. But the four visits per patient are correlated, so the effective sample size is well below 236 and this CI is too narrow — the significance is overstated. The correct analysis adds patient random effects (Breslow–Clayton 1993; the BUGS Epil model), where the progabide effect typically becomes marginal and specification-sensitive — a classic cautionary result. Bottom line: overdispersion modelling already overturns the naive Poisson "null", but a defensible treatment claim needs the hierarchical model — the next project.
3. Data vs fitted, by treatment arm¶
The plot below makes the treatment contrast visual. Left: observed mean seizures per 2-week visit (points ± SE) vs the NegBin-fitted means (dashed), for placebo and progabide separately — the fit captures the arm gap and the visit-4 dip, while the raw arms nearly overlap at visits 1–3 (the effect is hidden in noisy, baseline-confounded counts). Right: the covariate-adjusted treatment effect — predicted rate for a typical patient (mean baseline & age) under each arm with 95% CrIs: progabide ≈ 5.4 vs placebo ≈ 6.9, the ~20% reduction (RR ≈ 0.80, Bayesian). The gap between the barely-separated raw means and the clear adjusted difference is exactly why the covariate-adjusted, overdispersion-aware model matters.
ep['fitted'] = np.exp(X_e @ bn_e) # NegBin fitted mean per observation
obs = ep.groupby(['visit','Trt'])['seizures'].mean().unstack(1)
sem = ep.groupby(['visit','Trt'])['seizures'].sem().unstack(1)
fit = ep.groupby(['visit','Trt'])['fitted'].mean().unstack(1)
v = obs.index.values; cols = {0:'steelblue', 1:'firebrick'}; lab = {0:'placebo', 1:'progabide'}
fig, ax = plt.subplots(1, 2, figsize=(12, 4.4))
for t in (0, 1):
ax[0].errorbar(v, obs[t], yerr=sem[t], fmt='o', color=cols[t], capsize=3, label=f'{lab[t]} (observed)')
ax[0].plot(v, fit[t], '--', color=cols[t], lw=2, label=f'{lab[t]} (NegBin fitted)')
ax[0].set_xlabel('visit'); ax[0].set_ylabel('mean seizures / 2 weeks'); ax[0].set_xticks(v)
ax[0].set_title('Observed vs fitted, by arm'); ax[0].legend(fontsize=8)
B = nb_e['beta']; mlb = np.log(ep['Base']/4).mean(); mla = np.log(ep['Age']).mean()
p0 = np.exp(B[:,0] + B[:,1]*mlb + B[:,3]*mla); p1 = p0 * np.exp(B[:,2]) # placebo, progabide
m = [p0.mean(), p1.mean()]
lo = [np.percentile(p0,2.5), np.percentile(p1,2.5)]; hi = [np.percentile(p0,97.5), np.percentile(p1,97.5)]
ax[1].bar([0,1], m, yerr=[np.array(m)-lo, np.array(hi)-m], color=['steelblue','firebrick'], capsize=6, width=0.55)
ax[1].set_xticks([0,1]); ax[1].set_xticklabels(['placebo','progabide']); ax[1].set_ylabel('predicted seizures / 2 weeks')
ax[1].set_title(f'Adjusted treatment effect (RR {np.exp(B[:,2].mean()):.2f})')
plt.tight_layout(); plt.savefig('count_epil_fit.png', dpi=120, bbox_inches='tight'); plt.show()
