Causal Inference II(b) — Partial Pooling Across Subgroups

The Electric Company by grade: four effects, and whether they really differ

The covariate-adjustment notebook reported the Electric Company experiment within each grade, because its randomization was blocked by grade. Those four estimates are strikingly different:

grade	ANCOVA effect	SE
1	8.79	2.61
2	4.27	1.36
3	1.91	0.77
4	1.70	0.71

Read literally, the programme helps first-graders five times as much as fourth-graders. But look at the standard errors: the largest estimate is also the least precise, which is the signature of a number that owes something to luck. Grade 1 has the widest interval and the biggest effect — exactly the combination that regresses toward the mean on replication.

The frequentist analysis has two options and both are unattractive. Estimate each grade separately (no pooling) and accept that grade 1's 8.79 is noisy. Estimate one common effect (complete pooling) and assert the grades are identical, which the design gives no reason to believe. There is no principled middle.

A hierarchical model supplies the middle, and two things the published page could not produce:

1. Partially pooled grade effects — each grade shrunk toward the overall mean by an amount the data decides, not the analyst.
2. A posterior for the between-grade standard deviation $\sigma_\tau$ — which answers the question the frequentist version cannot even pose: is the variation across grades real, or is it noise?

Python/PyMC lead. This is the second example in the Bayesian causal group, and where the machinery starts doing structural work rather than sensitivity analysis.

1. The data and the two unattractive options

import numpy as np, pandas as pd, matplotlib.pyplot as plt, warnings
warnings.filterwarnings("ignore")
import pymc as pm, arviz as az
import statsmodels.api as sm
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; GREY="#a0aec0"; PURP="#6b46c1"

d = pd.read_csv("electric.csv", index_col=0)
d["grade"] = d["grade"].astype(int)
print(f"{len(d)} classrooms, grades {sorted(d.grade.unique())}, {int(d.treatment.sum())} treated")

# --- no pooling: ANCOVA within each grade, exactly as the frequentist page reports it
rows=[]
for g in sorted(d.grade.unique()):
    s = d[d.grade==g]
    X = sm.add_constant(np.column_stack([s.treatment.values, s.pre_test.values]))
    fit = sm.OLS(s.post_test.values, X).fit()
    rows.append((g, len(s), fit.params[1], fit.bse[1]))
nop = pd.DataFrame(rows, columns=["grade","n","effect","se"])
print("\nNO POOLING -- ANCOVA fitted separately in each grade:")
print(nop.to_string(index=False, float_format=lambda v: f"{v:7.3f}"))

# --- complete pooling: one effect for all grades
X = sm.add_constant(np.column_stack([d.treatment.values, d.pre_test.values]))
pooled = sm.OLS(d.post_test.values, X).fit()
print(f"\nCOMPLETE POOLING -- one common effect: {pooled.params[1]:.3f}  (SE {pooled.bse[1]:.3f})")
print("\nNeither is satisfactory. No pooling treats grade 1's noisy 8.79 as a real quantity; complete pooling")
print("asserts the four grades share one effect, which nothing in the design guarantees.")

g++ not available, if using conda: `conda install gxx`

192 classrooms, grades [np.int64(1), np.int64(2), np.int64(3), np.int64(4)], 96 treated

NO POOLING -- ANCOVA fitted separately in each grade:
 grade  n  effect      se
     1 42   8.787   2.612
     2 68   4.266   1.359
     3 40   1.910   0.776
     4 42   1.701   0.685

COMPLETE POOLING -- one common effect: 4.734  (SE 1.160)

Neither is satisfactory. No pooling treats grade 1's noisy 8.79 as a real quantity; complete pooling
asserts the four grades share one effect, which nothing in the design guarantees.

2. The hierarchical model

The middle option is to let the grade effects come from a common distribution whose spread is estimated:

$$y_i = \alpha_{g[i]} + \tau_{g[i]}\,T_i + \beta\,x_i + \varepsilon_i, \qquad \tau_g \sim \mathcal{N}(\mu_\tau,\ \sigma_\tau^2).$$

Everything hinges on $\sigma_\tau$, and it is estimated rather than assumed. If $\sigma_\tau \to 0$ the model collapses to complete pooling; if $\sigma_\tau \to \infty$ it becomes no pooling. Any value in between produces shrinkage, and — this is the part worth watching — the shrinkage is not uniform. A grade with a precise estimate barely moves; a grade with a noisy one is pulled hard toward the mean. The data decide which is which.

We use a non-centred parameterisation, which is the standard fix for the funnel geometry that hierarchical models produce when $\sigma_\tau$ is small.

g_idx = (d.grade.values - 1).astype(int); G = 4
T = d.treatment.values.astype(float); x = d.pre_test.values; y = d.post_test.values
xc = x - x.mean()

with pm.Model() as hier:
    mu_tau    = pm.Normal("mu_tau", 0, 10)
    sigma_tau = pm.HalfNormal("sigma_tau", 5)
    z         = pm.Normal("z", 0, 1, shape=G)                 # non-centred
    tau       = pm.Deterministic("tau", mu_tau + sigma_tau*z)
    alpha     = pm.Normal("alpha", 100, 30, shape=G)
    beta      = pm.Normal("beta", 0, 5)
    sigma     = pm.HalfNormal("sigma", 20)
    mu = alpha[g_idx] + tau[g_idx]*T + beta*xc
    pm.Normal("obs", mu, sigma, observed=y)
    idata = pm.sample(2000, tune=2000, chains=4, cores=1, target_accept=0.95,
                      progressbar=False, random_seed=0)

sm_ = az.summary(idata, var_names=["mu_tau","sigma_tau","tau","beta"])
cols = [c for c in ["mean","sd"] + [c for c in sm_.columns if c.startswith("hdi")] + ["ess_bulk","r_hat"] if c in sm_.columns]
print(sm_[cols].to_string())
print(f"\nmax R-hat = {float(pd.to_numeric(sm_['r_hat'], errors='coerce').max()):.4f};  divergences = {int(idata.sample_stats.diverging.values.sum())}")

Initializing NUTS using jitter+adapt_diag...

Sequential sampling (4 chains in 1 job)

NUTS: [mu_tau, sigma_tau, z, alpha, beta, sigma]

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

There was 1 divergence after tuning. Increase `target_accept` or reparameterize.

            mean     sd  ess_bulk r_hat
mu_tau      3.86   1.93      3567  1.00
sigma_tau   2.75   1.93      2330  1.00
tau[0]      6.13   2.15      4566  1.00
tau[1]      4.08   1.53      7466  1.00
tau[2]      3.08   1.86      7406  1.00
tau[3]      2.64   1.94      5754  1.00
beta       0.796  0.056      3805  1.00

max R-hat = 1.0000;  divergences = 1

3. What the shrinkage did

The comparison that matters is the unpooled estimate against the partially pooled one, grade by grade, read alongside the standard error that produced it.

tp = idata.posterior["tau"].values.reshape(-1, G)
part = tp.mean(0); plo = np.percentile(tp,2.5,axis=0); phi = np.percentile(tp,97.5,axis=0)
comp = pooled.params[1]

out = pd.DataFrame({"grade": nop.grade, "n": nop.n,
                    "no pooling": nop.effect.values, "SE": nop.se.values,
                    "partial pooling": part, "2.5%": plo, "97.5%": phi,
                    "shrinkage": nop.effect.values - part})
print(out.to_string(index=False, float_format=lambda v: f"{v:8.3f}"))
print(f"\ncomplete pooling would put every grade at {comp:.3f}")

mu_hat = idata.posterior["mu_tau"].values.mean()
out2 = out.assign(**{"distance from mu_tau": out["no pooling"] - mu_hat})
print("\nShrinkage is NOT simply 'the noisiest estimate moves most'. It is imprecision times DISTANCE")
print(f"from the common mean (mu_tau = {mu_hat:.2f}), and the two effects can pull against each other:")
print(out2[["grade","SE","distance from mu_tau","shrinkage"]].to_string(index=False, float_format=lambda v: f"{v:9.3f}"))
print()
print(f"  grade 1: far above the mean ({out2['distance from mu_tau'][0]:+.1f}) AND imprecise (SE {out.SE[0]:.2f}) -> moves most, {out.shrinkage[0]:+.2f}")
print(f"  grade 2: already AT the mean ({out2['distance from mu_tau'][1]:+.1f}) -> barely moves, {out.shrinkage[1]:+.2f}, despite SE {out.SE[1]:.2f}")
print(f"  grade 4: precise (SE {out.SE[3]:.2f}) but far BELOW the mean -> still pulled up {-out.shrinkage[3]:.2f}")
print()
print("Note the direction too: grades 3 and 4 move UP, not down. Shrinkage is toward the common mean, not")
print("toward zero -- a distinction that matters when the subgroup estimates straddle the average.")

fig,ax=plt.subplots(1,2,figsize=(13,4.4))
gg = out.grade.values
ax[0].errorbar(gg-0.06, out["no pooling"], yerr=1.96*out.SE, fmt="o", color=RED, capsize=3, label="no pooling (per-grade ANCOVA)")
ax[0].errorbar(gg+0.06, out["partial pooling"], yerr=[out["partial pooling"]-out["2.5%"], out["97.5%"]-out["partial pooling"]],
               fmt="s", color=BLUE, capsize=3, label="partial pooling (hierarchical)")
ax[0].axhline(comp, color=GREEN, ls="--", lw=1.2, label=f"complete pooling {comp:.2f}")
ax[0].axhline(0, color="k", lw=.6); ax[0].set_xticks(gg); ax[0].set_xlabel("grade")
ax[0].set_ylabel("effect on post-test reading"); ax[0].set_title("Pulled toward the mean by imprecision AND distance")
ax[0].legend(fontsize=8)
for i,row in out.iterrows():
    ax[1].plot([0,1],[row["no pooling"], row["partial pooling"]], "o-", color=[BLUE,ORANGE,GREEN,PURP][i], label=f"grade {int(row.grade)}")
ax[1].axhline(comp, color=GREY, ls="--", lw=1)
ax[1].set_xticks([0,1]); ax[1].set_xticklabels(["no pooling","partial pooling"]); ax[1].set_ylabel("effect")
ax[1].set_title("Where each grade ended up"); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.show()

 grade  n  no pooling       SE  partial pooling     2.5%    97.5%  shrinkage
     1 42       8.787    2.612            6.135    2.401   10.572      2.652
     2 68       4.266    1.359            4.078    1.061    7.132      0.187
     3 40       1.910    0.776            3.080   -0.841    6.496     -1.171
     4 42       1.701    0.685            2.643   -1.468    6.069     -0.942

complete pooling would put every grade at 4.734

Shrinkage is NOT simply 'the noisiest estimate moves most'. It is imprecision times DISTANCE
from the common mean (mu_tau = 3.86), and the two effects can pull against each other:
 grade        SE  distance from mu_tau  shrinkage
     1     2.612                 4.930      2.652
     2     1.359                 0.410      0.187
     3     0.776                -1.946     -1.171
     4     0.685                -2.155     -0.942

  grade 1: far above the mean (+4.9) AND imprecise (SE 2.61) -> moves most, +2.65
  grade 2: already AT the mean (+0.4) -> barely moves, +0.19, despite SE 1.36
  grade 4: precise (SE 0.69) but far BELOW the mean -> still pulled up 0.94

Note the direction too: grades 3 and 4 move UP, not down. Shrinkage is toward the common mean, not
toward zero -- a distinction that matters when the subgroup estimates straddle the average.

4. Is the variation across grades real?

This is the question the frequentist analysis cannot ask. Comparing four point estimates tells you they differ; it cannot tell you whether they differ by more than sampling noise would produce if the underlying effects were identical.

The posterior for $\sigma_\tau$ answers it directly. A posterior concentrated near zero says the grades are effectively interchangeable and the apparent spread was noise. A posterior clearly away from zero says the programme genuinely works differently by grade — which would be a substantive finding about the intervention, not a statistical one.

st = idata.posterior["sigma_tau"].values.ravel()
mt = idata.posterior["mu_tau"].values.ravel()
print(f"sigma_tau posterior: median {np.median(st):.2f}, 95% interval [{np.percentile(st,2.5):.2f}, {np.percentile(st,97.5):.2f}]")
print(f"  P(sigma_tau < 1 point) = {(st<1).mean():.3f}")
print(f"  P(sigma_tau < 2 points) = {(st<2).mean():.3f}")
print(f"\nmu_tau (average grade effect): {mt.mean():.2f}, 95% [{np.percentile(mt,2.5):.2f}, {np.percentile(mt,97.5):.2f}]")
print(f"  P(mu_tau > 0) = {(mt>0).mean():.4f}")
print(f"  compare complete pooling: {pooled.params[1]:.2f} (SE {pooled.bse[1]:.2f}) -- a much tighter statement")
print("  The hierarchical interval is WIDER than the pooled one, and that is correct rather than a loss.")
print("  Complete pooling gets its precision by ASSUMING the grades share one effect; the hierarchical model")
print("  pays for not assuming it, propagating the uncertainty in sigma_tau into the average.")

# does grade 1 really beat grade 4?
d14 = tp[:,0] - tp[:,3]
print(f"\ngrade 1 minus grade 4: posterior mean {d14.mean():+.2f}, 95% [{np.percentile(d14,2.5):+.2f}, {np.percentile(d14,97.5):+.2f}]")
print(f"  P(grade 1 effect > grade 4 effect) = {(d14>0).mean():.3f}")
print(f"  the unpooled gap was {nop.effect.values[0]-nop.effect.values[3]:+.2f} points, which looked decisive")

fig,ax=plt.subplots(1,2,figsize=(13,4.2))
ax[0].hist(st,bins=60,color=PURP,alpha=.75); ax[0].axvline(np.median(st),color=RED,lw=1.5,label=f"median {np.median(st):.2f}")
ax[0].set_xlabel("sigma_tau  (between-grade sd of the effect)"); ax[0].set_ylabel("posterior draws")
ax[0].set_title("Is the grade-to-grade variation real?"); ax[0].legend(fontsize=8)
ax[1].hist(d14,bins=60,color=BLUE,alpha=.75); ax[1].axvline(0,color="k",lw=1)
ax[1].axvline(nop.effect.values[0]-nop.effect.values[3],color=RED,ls="--",lw=1.5,label="unpooled gap")
ax[1].set_xlabel("grade 1 effect minus grade 4 effect"); ax[1].set_ylabel("posterior draws")
ax[1].set_title("The comparison that looked decisive"); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.show()

sigma_tau posterior: median 2.41, 95% interval [0.13, 7.44]

  P(sigma_tau < 1 point) = 0.193
  P(sigma_tau < 2 points) = 0.408

mu_tau (average grade effect): 3.86, 95% [-0.24, 7.53]
  P(mu_tau > 0) = 0.9702
  compare complete pooling: 4.73 (SE 1.16) -- a much tighter statement
  The hierarchical interval is WIDER than the pooled one, and that is correct rather than a loss.
  Complete pooling gets its precision by ASSUMING the grades share one effect; the hierarchical model
  pays for not assuming it, propagating the uncertainty in sigma_tau into the average.

grade 1 minus grade 4: posterior mean +3.49, 95% [-0.86, +10.02]
  P(grade 1 effect > grade 4 effect) = 0.888
  the unpooled gap was +7.09 points, which looked decisive

5. Summary

Three results, none of them available from the frequentist analysis of the same 192 classrooms:

• Partial pooling, with the amount decided by the data. Each grade is pulled toward the common mean by its imprecision times its distance from that mean — which is why grade 2 barely moves despite a middling standard error (it was already at the mean) while grade 4 is pulled up almost a full point despite being the most precise. Note the direction as well: grades 3 and 4 move up. Shrinkage is toward the common mean, not toward zero.
• A posterior for the between-grade spread, which turns "do the grades differ?" from an eyeball comparison of four numbers into a probability statement.
• A defensible answer to the subgroup question. The unpooled gap between grade 1 and grade 4 looked large. The posterior for that difference says how much of it survives once both estimates are allowed to be noisy — which is the honest way to read a subgroup analysis, and the reason pre-registered subgroup claims are so often disappointed on replication.

The general lesson is that partial pooling is the correct default whenever an experiment is blocked or stratified, which describes most field experiments. The frequentist page had to choose between two extremes; the hierarchical model treats the choice itself as a parameter.

Cross-links. This is the hierarchical machinery of the Bayesian arc applied to a causal estimand rather than a descriptive one. The next example in this group scales the same idea from four grades to seventy-nine schools, where the clustering is not a design feature but the thing that broke the standard errors in the first place.