Selection Models — Missing Not at Random¶
When whether you see a value depends on the value itself¶
The first three projects assumed ignorable missingness: the reason data are missing added nothing beyond the observed data, so imputing from the data model was valid. This project crosses into non-ignorable (MNAR) territory, where a value is missing because of what it would have been. The archetype is sample selection: a wage is observed only for people who work, and the same unobserved traits — drive, ability — raise both the chance of working and the wage. The wages we see are then a biased sample of all potential wages, and no ignorable method can undo it — the correction must model the selection.
Heckman's model is two equations with correlated errors: $$\text{outcome: } y_i=x_i'\beta+\varepsilon_i\ \text{(seen only if }r_i=1),\qquad \text{selection: } r_i=\mathbb 1\{w_i'\gamma+u_i>0\},\qquad (\varepsilon_i,u_i)\sim N\!\Big(0,\begin{bmatrix}\sigma^2&\lambda\\\lambda&1\end{bmatrix}\Big).$$ The correlation $\rho=\lambda/\sigma$ is everything: $\rho=0$ means selection is ignorable and ordinary least squares on the observed sample is fine; $\rho\neq0$ means the missingness is non-ignorable and OLS is biased. The catch — the essence of MNAR — is that $\rho$ is only weakly identified, leaning on an exclusion restriction (a covariate driving selection but not the outcome) and, ultimately, an assumption the data cannot fully check.
This is Tobit's bigger sibling. The Tobit model (Bayesian Tobit — Censored Gaussian Regression) is a selection model where selection is a deterministic function of the outcome (observed iff the latent value clears a threshold) — one equation, $\rho=1$. Heckman frees the selection into its own equation with its own covariates and a free correlation. We build the augmentation sampler from scratch, show it corrects the bias, confront the identification problem, run the classic Mroz wage data, and cross-check in PyMC.
import numpy as np, pandas as pd, matplotlib.pyplot as plt
import selection as SL
rng = np.random.default_rng(4)
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; PURP="#6b46c1"; GREY="#718096"
print("MNAR: a value is missing because of what it would have been. Selection models add an equation for BEING")
print("observed, correlated with the outcome; the correlation rho is the non-ignorability -- and the hard-to-pin part.")
MNAR: a value is missing because of what it would have been. Selection models add an equation for BEING observed, correlated with the outcome; the correlation rho is the non-ignorability -- and the hard-to-pin part.
1. Selection bias, and the naive fix that isn't¶
Simulate a wage-like outcome observed only for a selected subsample, with the selection error correlated with the outcome error ($\rho=0.7$ — those who work have higher unobserved wage components). Ordinary least squares on the observed cases is biased: the selected sample is not representative. The plot shows the full population, the selected subsample, and how the OLS line through the selected points misses the true relationship.
beta=np.array([1.0,0.8,-0.5]); gamma=np.array([0.3,0.5,0.7])
yobs,r,X,W,yfull=SL.simulate_heckman(6000,beta,gamma,rho=0.7,sigma=1.0,rng=rng)
sel=r>0.5; nb,nse=SL.naive_ols(yobs,X)
fig,ax=plt.subplots(figsize=(7.5,4.2)); xg=X[:,1]
ax.scatter(xg[~sel],yfull[~sel],s=6,color="#d9d9d9",label="missing (not selected)")
ax.scatter(xg[sel],yfull[sel],s=6,color=BLUE,alpha=.5,label="observed (selected)")
xs=np.linspace(xg.min(),xg.max(),50)
ax.plot(xs,beta[0]+beta[1]*xs,color=GREEN,lw=2.5,label="true relationship")
ax.plot(xs,nb[0]+nb[1]*xs,color=RED,lw=2.5,ls="--",label="OLS on selected (biased)")
ax.set_xlabel(r"$x_1$"); ax.set_ylabel("outcome y"); ax.set_title(r"Selection bias: OLS on the observed sample misses the truth ($\rho=0.7$)")
ax.legend(frameon=False,fontsize=8); plt.tight_layout(); plt.show()
print(f"true intercept {beta[0]:.2f}, slope {beta[1]:.2f}; OLS on selected gives intercept {nb[0]:.2f}, slope {nb[1]:.2f}.")
print("The selected points sit systematically above the population at low x, so the fitted line is pulled off -- and")
print("more data of the same kind would not help: the bias is structural, not sampling noise.")
true intercept 1.00, slope 0.80; OLS on selected gives intercept 1.36, slope 0.69. The selected points sit systematically above the population at low x, so the fitted line is pulled off -- and more data of the same kind would not help: the bias is structural, not sampling noise.
2. The Heckman correction recovers the truth¶
The selection-model sampler augments the latent selection index and the missing outcomes, and estimates the outcome coefficients, the selection coefficients, and the correlation $\rho$ together. It should recover $\beta$ and detect $\rho=0.7$; and when we regenerate the data with $\rho=0$ (ignorable), it should agree with OLS and report $\rho\approx0$.
res=SL.selection_gibbs(yobs,r,X,W,rng,draws=3000,burn=1500)
bm=res["beta"].mean(0); bl,bh=np.percentile(res["beta"],[2.5,97.5],axis=0)
print("rho = 0.7 (non-ignorable):")
print(f" beta true {beta} Heckman {bm.round(2)} naive-OLS {nb.round(2)}")
print(f" rho estimated {res['rho'].mean():.2f} 95% CrI [{np.percentile(res['rho'],2.5):.2f}, {np.percentile(res['rho'],97.5):.2f}]")
# ignorable case
yo0,r0,X0,W0,_=SL.simulate_heckman(6000,beta,gamma,rho=0.0,sigma=1.0,rng=rng)
res0=SL.selection_gibbs(yo0,r0,X0,W0,rng,draws=2000,burn=1000); nb0,_=SL.naive_ols(yo0,X0)
fig,ax=plt.subplots(1,2,figsize=(12,3.6))
for k,(ttl,truth,hk,ol) in enumerate([(r"$\rho=0.7$ (non-ignorable)",beta,bm,nb),
(r"$\rho=0$ (ignorable)",beta,res0["beta"].mean(0),nb0)]):
xi=np.arange(3); ax[k].scatter(xi-.12,truth,color=GREEN,s=80,marker="_",label="truth",zorder=3)
ax[k].scatter(xi,hk,color=BLUE,s=50,label="Heckman"); ax[k].scatter(xi+.12,ol,color=RED,s=50,marker="s",label="naive OLS")
ax[k].set_xticks(xi); ax[k].set_xticklabels(["intercept","slope x1","slope x2"]); ax[k].set_title(ttl); ax[k].axhline(0,color="k",lw=.5)
ax[0].legend(frameon=False,fontsize=8); ax[0].set_ylabel("coefficient"); plt.tight_layout(); plt.show()
print(f"Non-ignorable: Heckman recovers beta and finds rho={res['rho'].mean():.2f}; OLS is off. Ignorable: rho={res0['rho'].mean():+.2f},")
print("and Heckman and OLS agree. The model reduces to OLS exactly when there is no selection on unobservables.")
rho = 0.7 (non-ignorable): beta true [ 1. 0.8 -0.5] Heckman [ 0.94 0.83 -0.49] naive-OLS [ 1.36 0.69 -0.49] rho estimated 0.76 95% CrI [0.71, 0.81]
Non-ignorable: Heckman recovers beta and finds rho=0.76; OLS is off. Ignorable: rho=+0.03, and Heckman and OLS agree. The model reduces to OLS exactly when there is no selection on unobservables.
3. The catch — $\rho$ needs an exclusion restriction¶
MNAR corrections are not free. The correlation $\rho$ is identified partly by functional form alone, which is fragile; it becomes trustworthy only with an exclusion restriction — a variable that drives selection but not the outcome (here $z$, in $w$ but not $x$). We refit the same data with the exclusion restriction and without it (selection uses only the outcome's covariates), and compare how sharply $\rho$ is pinned down.
res_ex = SL.selection_gibbs(yobs,r,X,W,rng,draws=3000,burn=1500) # W has the exclusion variable z
res_no = SL.selection_gibbs(yobs,r,X,X,rng,draws=3000,burn=1500) # W = X : no exclusion restriction
fig,ax=plt.subplots(figsize=(7.5,3.6))
ax.hist(res_ex["rho"],bins=40,color=GREEN,alpha=.7,density=True,label=f"with exclusion (sd {res_ex['rho'].std():.2f})")
ax.hist(res_no["rho"],bins=40,color=ORANGE,alpha=.6,density=True,label=f"without exclusion (sd {res_no['rho'].std():.2f})")
ax.axvline(0.7,color="k",ls="--",lw=1.5,label="true rho 0.7"); ax.set_xlabel(r"posterior of selection correlation $\rho$")
ax.set_title("An exclusion restriction sharpens the (otherwise fragile) selection correlation"); ax.legend(frameon=False,fontsize=8)
plt.tight_layout(); plt.show()
print(f"with exclusion: rho posterior mean {res_ex['rho'].mean():.2f}, sd {res_ex['rho'].std():.2f}")
print(f"without exclusion: rho posterior mean {res_no['rho'].mean():.2f}, sd {res_no['rho'].std():.2f}"
f" ({res_no['rho'].std()/res_ex['rho'].std():.1f}x wider, leaning on functional form alone)")
print("Without a genuine exclusion restriction the selection correction rests on the normality assumption -- exactly")
print("the kind of untestable crutch that makes MNAR analysis a matter of judgement, not just computation.")
with exclusion: rho posterior mean 0.76, sd 0.02 without exclusion: rho posterior mean 0.65, sd 0.05 (2.5x wider, leaning on functional form alone) Without a genuine exclusion restriction the selection correction rests on the normality assumption -- exactly the kind of untestable crutch that makes MNAR analysis a matter of judgement, not just computation.
4. Real data — the Mroz wage-selection problem¶
The Mroz data (753 married women, 428 working) is the selection benchmark: the log wage is observed only for the 428 who work, and labour-force participation is the selection. We model $\log \text{wage}\sim\text{educ}+\text{exper}+\text{exper}^2$, with participation additionally depending on age, young children and non-wife income (the exclusion restrictions). The question: are the returns to education biased by looking only at working women?
These very same women run through three projects in this portfolio, each modelling a different facet of one labour-supply decision. In Bayesian Binary Logit the target is whether she works — labour-force participation itself, exactly the selection equation $r_i=\mathbb 1\{w_i'\gamma+u_i>0\}$ used here. In Bayesian Tobit — Censored Gaussian Regression the target is hours worked, left-censored at zero for the 325 non-working women. And here the target is the wage, missing entirely for those same 325. Censoring keeps hours in the model at a known boundary (zero); selection removes the wage altogether — the textbook distinction between the Type I Tobit (censoring, one equation) and the Type II Tobit / Heckman model (selection, two equations). The participation model of Bayesian Binary Logit is the second equation of this one: one dataset, three complementary views of the same decision.
d=pd.read_csv("mroz.csv"); n=len(d); ex=d["exper"].to_numpy()/10; ed=d["educ"].to_numpy()
Xm=np.column_stack([np.ones(n),ed,ex,ex**2])
Wm=np.column_stack([np.ones(n),ed,ex,ex**2,d["age"]/10,d["kids5"],d["kids618"],d["nwifeinc"]/10])
ym=d["lwage"].to_numpy(); rm=d["lfp"].to_numpy(float)
rH=SL.selection_gibbs(ym,rm,Xm,Wm,rng,draws=4000,burn=2000); nbm,nsem=SL.naive_ols(ym,Xm)
rho_m=rH["rho"]; rlo,rhi=np.percentile(rho_m,[2.5,97.5]); educ_H=rH["beta"][:,1]
print(f"returns to education: Heckman {educ_H.mean():.3f} naive-OLS {nbm[1]:.3f}")
print(f"selection correlation rho = {rho_m.mean():+.2f} 95% CrI [{rlo:+.2f}, {rhi:+.2f}] -> {'includes 0: no strong evidence of selection on unobservables' if rlo<0<rhi else 'excludes 0'}")
fig,ax=plt.subplots(1,2,figsize=(12,3.8))
ax[0].hist(rho_m,bins=40,color=PURP,alpha=.8,density=True); ax[0].axvline(0,color="k",lw=1.5,ls="--")
ax[0].set_xlabel(r"selection correlation $\rho$"); ax[0].set_title("Mroz: rho is barely distinguishable from 0")
ax[1].scatter(rho_m[::4],educ_H[::4],s=5,color=GREY,alpha=.4); ax[1].axhline(nbm[1],color=RED,lw=1.5,label="OLS returns")
ax[1].set_xlabel(r"$\rho$"); ax[1].set_ylabel("returns to education"); ax[1].set_title("Returns to education vs assumed selection"); ax[1].legend(frameon=False,fontsize=8)
plt.tight_layout(); plt.show()
print("On Mroz the data do not clearly identify selection: rho's interval spans zero, so the correction barely moves")
print("the returns to education (~0.11, ~11% per year) from the OLS value. That is the honest MNAR verdict here -- not")
print("'there is no selection' but 'the data cannot tell', and the joint plot shows how the answer WOULD shift if we")
print("assumed a particular rho. Pinning rho down is an assumption, not an estimate -- the theme of Project 5.")
returns to education: Heckman 0.108 naive-OLS 0.107 selection correlation rho = +0.00 95% CrI [-0.27, +0.25] -> includes 0: no strong evidence of selection on unobservables
On Mroz the data do not clearly identify selection: rho's interval spans zero, so the correction barely moves the returns to education (~0.11, ~11% per year) from the OLS value. That is the honest MNAR verdict here -- not 'there is no selection' but 'the data cannot tell', and the joint plot shows how the answer WOULD shift if we assumed a particular rho. Pinning rho down is an assumption, not an estimate -- the theme of Project 5.
5. Cross-check in PyMC¶
We write the Heckman likelihood directly in PyMC via pm.Potential: for a selected case the contribution is the normal density of the observed $y$ times the conditional probability of being selected given $y$; for a non-selected case it is the probability of not being selected. Sampling $\beta,\gamma,\sigma,\rho$ should match the from-scratch augmentation on the simulated data.
import pymc as pm, pytensor.tensor as pt
sel=r>0.5; ys=yobs[sel]; Xs=X[sel]; Ws=W[sel]; Wn=W[~sel]
def logphi(z): return -0.5*(z**2+np.log(2*np.pi))
with pm.Model() as mod:
b=pm.Normal("beta",0,10,shape=3); g=pm.Normal("gamma",0,10,shape=3)
sigma=pm.HalfNormal("sigma",5); rho=pm.Uniform("rho",-0.99,0.99)
e=(ys-Xs@b)/sigma
sel_ll = logphi(e)-pt.log(sigma) + pm.math.log(pm.math.invprobit((Ws@g+rho*e)/pt.sqrt(1-rho**2)))
non_ll = pm.math.log(pm.math.invprobit(-(Wn@g)))
pm.Potential("lik", pt.sum(sel_ll)+pt.sum(non_ll))
idata=pm.sample(800,tune=1200,chains=4,target_accept=0.9,random_seed=6,progressbar=False)
bp=idata.posterior["beta"].mean(("chain","draw")).values; rp=float(idata.posterior["rho"].mean())
print(" from-scratch PyMC")
print(f"beta {bm.round(2)} {bp.round(2)}")
print(f"rho {res['rho'].mean():.2f} {rp:.2f}")
print("The explicit Heckman likelihood in PyMC and the from-scratch data-augmentation sampler agree on the")
print("coefficients and the selection correlation -- two routes to the same MNAR correction.")
g++ not available, if using conda: `conda install gxx`
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta, gamma, sigma, rho]
Sampling 4 chains for 1_200 tune and 800 draw iterations (4_800 + 3_200 draws total) took 8 seconds.
from-scratch PyMC beta [ 0.94 0.83 -0.49] [ 0.94 0.83 -0.49] rho 0.76 0.76 The explicit Heckman likelihood in PyMC and the from-scratch data-augmentation sampler agree on the coefficients and the selection correlation -- two routes to the same MNAR correction.
6. Summary¶
Non-ignorable missingness breaks every method from the earlier projects, because the observed data are a biased window on the full data. Selection models meet it head-on: an equation for the outcome and a correlated equation for being observed, with the correlation $\rho$ measuring the non-ignorability. When $\rho\neq0$, ordinary least squares on the observed cases is biased and the Heckman correction recovers the truth — the simulation confirmed both, and the model collapses to OLS exactly when $\rho=0$. But the correction is only as good as its identification: without an exclusion restriction $\rho$ leans on functional form alone, and on the real Mroz data $\rho$'s posterior spanned zero — the data could not say whether selection mattered, leaving the returns to education near their OLS value but for an untestable reason.
That last point is the heart of MNAR: the correction depends on assumptions the data cannot check. The connections: this is Congdon's selection-model treatment of dropout (BMCD 11.1–11.2), and it generalises Tobit (Bayesian Tobit — Censored Gaussian Regression), which is the special case where selection is welded to the outcome ($\rho=1$, $w=x$); the augmentation is the same Albert–Chib truncated-normal machinery, now on a bivariate latent. Selection models factor the joint as $p(y)\,p(r\mid y)$; the next project, pattern-mixture models, takes the complementary factorisation $p(y\mid r)\,p(r)$ and makes the untestable assumption explicit as a sensitivity parameter — turning the discomfort of this project's wide $\rho$ into a formal analysis.