Selection Models — Missing Not at Random

Python · PyMC · R  ·  Download selection module

Modelling Why the Data Are Missing

The first three projects assumed the missingness was ignorable. This one crosses the line. When whether a value is observed depends on the value itself, no imputation from the data model can help — the observed data are formally silent about the distribution of what is missing. The remedy is to stop treating missingness as a nuisance and model it. The selection-model factorisation writes the joint density as p(y) p(r∣y)p(y)\,p(r\mid y): a model for the outcome times a model for being observed that is allowed to depend on the outcome, including where it is missing.

yi=xi⊤β+εi,Si∗=wi⊤γ+ui,ri=1{Si∗>0},(εu)∼N ⁣(0,(σ2λλ1))y_i=x_i^{\top}\beta+\varepsilon_i,\qquad S_i^{*}=w_i^{\top}\gamma+u_i,\qquad r_i=\mathbf 1\{S_i^{*}>0\},\qquad \begin{pmatrix}\varepsilon\\u\end{pmatrix}\sim N\!\left(0,\begin{pmatrix}\sigma^2&\lambda\\\lambda&1\end{pmatrix}\right)

Heckman's version is two equations with correlated errors — an outcome equation seen only when a latent selection index clears zero. The correlation ρ\rho between the two errors is the non-ignorability: at ρ=0\rho=0 selection is ignorable and OLS on the observed sample is fine; at ρ≠0\rho\neq0 it is biased by the classic inverse-Mills term — the expected value of the outcome error given that the observation was selected, which is non-zero exactly when the two errors are correlated. Omitting it is what makes OLS on the selected sample wrong; estimating it is what the correction does, and its near-collinearity with the regressors is what makes ρ\rho so hard to pin down. The sampler is data augmentation again — augment the latent index truncated by the response indicator, augment the missing outcomes, and every remaining block is conjugate.

What it corrects

On simulated data the machinery works and the failure it corrects is stark. With ρ=0.7\rho=0.7 the selected points sit systematically above the population at low xx, and OLS returns an intercept of 1.36 against a truth of 1.00 and a slope of 0.69 against 0.80. More data of the same kind would not help — the bias is structural, not sampling noise. The Heckman sampler recovers (0.94, 0.83, −0.49) against a true (1.0, 0.8, −0.5). Set ρ=0\rho=0 and it reduces to OLS exactly, which is the right behaviour: the correction disappears when there is nothing to correct.

simulated, ρ=0.7\rho=0.7interceptslope x1x_1slope x2x_2
truth1.000.80−0.50
Heckman0.940.83−0.49
OLS on the selected sample1.360.69−0.49

What it rests on

The catch is ρ\rho itself. It is identified partly by functional form — the normality assumption alone can produce an estimate — which makes an exclusion restriction (a variable driving selection but not the outcome) close to essential. Removing it from the same data widens the ρ\rho posterior by a factor of 2.5 and shifts its centre. What remains is an estimate resting on an assumption the data cannot check, which is the defining feature of MNAR analysis rather than a defect of this model.

When the Data Cannot Tell

The real data make the point better than the simulation does. On the Mroz labour-supply problem — 753 women, 428 with an observed wage — the estimated selection correlation is ρ=+0.00\rho=+0.00 with a 95% interval of [−0.27, +0.25]. The interval spans zero, so the correction barely moves the returns to education from the OLS value of about 11% per year. The honest verdict is not "there is no selection" but "the data cannot tell".

Mroz — returns to educationestimateselection correlation ρ\rho
naive OLS on the 428 workers0.107—
Bayesian Heckman (from scratch)0.108+0.00, CrI [−0.27, +0.25]
Heckman maximum likelihood (R)0.1310.03, SE 0.15

Two engines, two symptoms

And the two engines respond to that weak identification differently, which is worth seeing. R's sampleSelection maximum-likelihood fit returns ρ=0.03\rho=0.03 with a standard error of 0.15 — thoroughly insignificant — yet still shifts the returns to education from 0.107 to 0.131, a 22% move, because the inverse-Mills term is nearly collinear with education. The Bayesian sampler, regularised by its prior, instead stays near the OLS value at 0.108. Same weak-identification problem, two different symptoms: maximum likelihood moves the estimate on the strength of a correlation it cannot pin down, and the posterior declines to. Neither is wrong, and the disagreement is more informative than either number alone.

Notebooks

Python Selection bias shown against a known population, the Heckman correction recovering the truth and collapsing to OLS when ρ = 0, the exclusion restriction sharpening ρ by a factor of 2.5, the Mroz wage problem where ρ spans zero, and a PyMC likelihood cross-check (Python) Open → R Independent R implementation via sampleSelection — maximum-likelihood and two-step Heckman against naive OLS on Mroz, with the insignificant ρ and the 22% shift it nonetheless produces (R) Open →

Downloads

md_selection.py Bayesian Heckman selection model by data augmentation — truncated draws of the latent selection index, augmentation of the unobserved outcomes, and conjugate blocks for the coefficients and the error covariance reparametrised so that σ² = λ² + τ² (NumPy / SciPy) mroz.csv Mroz (1987) labour supply — 753 married women, 428 with an observed wage, the canonical sample-selection dataset

Selection Module — Source Code

"""
selection.py -- HECKMAN SELECTION MODELS for nonignorable (MNAR) missingness (from scratch).

Backs the notebooks in  "Selection Models -- Missing Not at Random".

Projects 1-3 assumed IGNORABLE missingness (MCAR/MAR): the reason data are missing carries no
information beyond the observed data, so imputing from the data model is valid. This project crosses
into NONIGNORABLE (MNAR) territory, where whether a value is observed depends on the value ITSELF.
The canonical case is sample selection: a wage is seen only for people who work, and the same
unobserved traits (drive, ability) push both the decision to work and the wage -- so the observed
wages are a biased sample of all potential wages. No ignorable method can undo that; we must MODEL
the selection.

The SELECTION-MODEL factorisation writes the joint density of the outcome y and the response
indicator r as  p(y) p(r | y)  -- a model for the outcome times a model for BEING OBSERVED that is
allowed to depend on the (possibly missing) outcome. Heckman's Gaussian version is two equations
with correlated errors:

    outcome:    y_i   = x_i' beta  + eps_i                          (observed only when r_i = 1)
    selection:  S_i*  = w_i' gamma + u_i ,   r_i = 1{ S_i* > 0 }
    (eps_i, u_i) ~ N(0, [[sigma^2, lambda],[lambda, 1]])            var(u)=1 identifies the probit

The correlation rho = lambda/sigma is the crux: rho = 0 means selection is ignorable (MAR) and OLS
on the selected sample is fine; rho != 0 means the missingness is nonignorable and OLS is biased by
the classic inverse-Mills-ratio term. rho is what the model estimates -- but it is only weakly
identified by functional form, so an EXCLUSION RESTRICTION (a covariate in w but not in x) and a
frank sensitivity analysis are essential. This is the deep point of MNAR: the correction rests on an
assumption the data cannot fully check.

CONNECTION TO TOBIT. Tobit ("Bayesian Tobit -- Censored Gaussian Regression") is a
selection model too -- but one where selection is a
DETERMINISTIC function of the outcome (y observed iff its latent value clears a threshold), i.e. the
same equation drives value and observation (rho = 1, w = x). Heckman generalises it to a SEPARATE
selection equation with its own covariates and a free correlation. Tobit is Heckman with the
selection welded to the outcome.

The sampler is data augmentation, reusing the Albert-Chib truncated-normal machinery: augment the
latent index S_i* (truncated by r_i) and the missing outcomes y_i (for r=0), then draw beta, gamma
and the covariance (reparametrised as lambda = cov and tau^2 = var(eps|u), so sigma^2 = lambda^2 +
tau^2 and rho = lambda/sigma) from conjugate conditionals.
"""

import numpy as np
from scipy.special import ndtr as Phi, ndtri as Phinv


def _rtn(mean, sd, positive, rng):
    """draw N(mean, sd^2) truncated to (0, inf) where `positive` is True, else to (-inf, 0).
    `positive` is a boolean array (per element)."""
    positive = np.asarray(positive); alpha = Phi(-mean / sd); uu = rng.random(mean.shape)
    p = np.where(positive, alpha + uu * (1 - alpha), uu * alpha)
    return mean + sd * Phinv(np.clip(p, 1e-12, 1 - 1e-12))


def simulate_heckman(n, beta, gamma, rho, sigma, rng):
    """Two-equation Heckman data with an exclusion restriction (z enters selection, not outcome).
    Returns observed y (NaN when not selected), r, X, W, and the full latent y for checking."""
    x1 = rng.standard_normal(n); x2 = rng.standard_normal(n); z = rng.standard_normal(n)
    X = np.column_stack([np.ones(n), x1, x2]); W = np.column_stack([np.ones(n), x1, z])
    u = rng.standard_normal(n)
    eps = rho * sigma * u + sigma * np.sqrt(1 - rho ** 2) * rng.standard_normal(n)   # corr(eps,u)=rho
    y = X @ np.asarray(beta, float) + eps; r = (W @ np.asarray(gamma, float) + u > 0).astype(float)
    return np.where(r > 0.5, y, np.nan), r, X, W, y


def selection_gibbs(yobs, r, X, W, rng, draws=3000, burn=1500, b_sd=10.0, g_sd=10.0):
    """Bayesian Heckman selection model by data augmentation. yobs:(n,) with NaN for r=0; r:(n,) 0/1;
    X,W design matrices (include intercept; W should have an exclusion restriction). Returns posterior
    draws of beta, gamma, sigma, and the selection correlation rho."""
    r = np.asarray(r, float); X = np.asarray(X, float); W = np.asarray(W, float)
    n, px = X.shape; pw = W.shape[1]; sel = r > 0.5
    yc = np.where(sel, np.nan_to_num(yobs), np.nan)
    # init
    b0, *_ = np.linalg.lstsq(X[sel], yobs[sel], rcond=None); beta = b0
    gamma = np.zeros(pw); gamma[0] = 0.0
    resid0 = yobs[sel] - X[sel] @ beta; tau2 = resid0 @ resid0 / sel.sum(); lam = 0.0
    yc[~sel] = (X[~sel] @ beta)
    XtXi = np.linalg.inv(X.T @ X + (1 / b_sd ** 2) * np.eye(px)); Lx = np.linalg.cholesky(XtXi)
    WtWi = np.linalg.inv(W.T @ W + (1 / g_sd ** 2) * np.eye(pw)); Lw = np.linalg.cholesky(WtWi)
    B = np.empty((draws, px)); G = np.empty((draws, pw)); SIG = np.empty(draws); RHO = np.empty(draws)
    for it in range(draws + burn):
        sig2 = lam ** 2 + tau2; eps = yc - X @ beta
        # ---- augment latent selection index S* | y, truncated by r ----
        mS = W @ gamma + (lam / sig2) * eps; vS = tau2 / sig2
        S = _rtn(mS, np.sqrt(vS), sel, rng); u = S - W @ gamma
        # ---- gamma | .  (probit-style, correlation-adjusted response, error var vS) ----
        resp_g = S - (lam / sig2) * eps
        gm = WtWi @ (W.T @ resp_g); gamma = gm + np.sqrt(vS) * (Lw @ rng.standard_normal(pw))
        u = S - W @ gamma
        # ---- augment missing outcomes y | u  (r=0) ----
        yc[~sel] = X[~sel] @ beta + lam * u[~sel] + np.sqrt(tau2) * rng.standard_normal((~sel).sum())
        # ---- beta | .  (regress y - lam*u on X, error var tau2) ----
        resp_b = yc - lam * u
        bm = XtXi @ (X.T @ resp_b); beta = bm + np.sqrt(tau2) * (Lx @ rng.standard_normal(px))
        eps = yc - X @ beta
        # ---- lambda, tau2 | .  (regress eps on u through the origin) ----
        uu = u @ u + 1e-6; ue = u @ eps; lam_hat = ue / uu
        rss = eps @ eps - lam_hat * ue
        tau2 = max(rss, 1e-6) / rng.chisquare(max(n - 1, 1))
        lam = lam_hat + np.sqrt(tau2 / uu) * rng.standard_normal()
        if it >= burn:
            i = it - burn; sig = np.sqrt(lam ** 2 + tau2)
            B[i] = beta; G[i] = gamma; SIG[i] = sig; RHO[i] = lam / sig
    return dict(beta=B, gamma=G, sigma=SIG, rho=RHO)


def naive_ols(yobs, X):
    """OLS on the selected sample only -- biased when the selection correlation rho != 0."""
    sel = ~np.isnan(yobs); Xs = X[sel]
    XtXi = np.linalg.inv(Xs.T @ Xs); b = XtXi @ (Xs.T @ yobs[sel])
    resid = yobs[sel] - Xs @ b; s2 = resid @ resid / (sel.sum() - Xs.shape[1])
    return b, np.sqrt(np.diag(s2 * XtXi))

References