Heckman selection model in Python: ML and two-step estimation

Some outcomes are observed only for part of the sample: wages only for people who work, repayment only for loans that were approved, test scores only for students who sat the exam. If the reasons for being in the sample are related to the unobserved part of the outcome, a regression on the observed cases is biased for the population relationship, however many observations it uses. Gronau (1974) showed the problem for wage comparisons. Heckman (1979) treated it as an omitted-variable problem and gave a simple correction.

What the Heckman model estimates

The model has an outcome equation and a selection equation whose errors are jointly normal with correlation ρ:

yᵢ = xᵢ′β + u₁ᵢ                     observed only when sᵢ = 1
sᵢ = 1[zᵢ′γ + u₂ᵢ > 0]              sd(u₁) = σ,  sd(u₂) = 1,  corr(u₁, u₂) = ρ
E[yᵢ | xᵢ, sᵢ = 1] = xᵢ′β + ρσ λ(zᵢ′γ),    λ(c) = φ(c) / Φ(c)

The last line shows the bias. Among selected observations, the outcome error does not have mean zero: its conditional mean is ρσ times the inverse Mills ratio λ, which depends on the selection regressors. Least squares on the selected sample omits that term. If ρ = 0, nothing is omitted; otherwise the slopes are biased whenever x is correlated with λ. β describes the relationship in the whole population, for example wage offers rather than the wages of those who accepted one.

• Two-step estimator (Heckman 1979). A probit of sᵢ on zᵢ gives γ and hence λ for every selected observation. A regression of y on x and λ over the selected sample then gives β, and the coefficient on λ estimates ρσ. The standard errors must allow for the estimated λ and for the heteroskedasticity the selection creates.
• Maximum likelihood. Both equations are estimated jointly under bivariate normality. It is efficient when that assumption holds and reports ρ and σ with standard errors and a test of ρ = 0.

Identification. The selection equation should contain at least one variable that affects selection but not the outcome, such as young children or other household income in a wage equation. Without one, the model is identified only through the nonlinearity of λ, which is often close to linear over the range of the data. λ is then nearly collinear with x, and the estimates become imprecise and fragile (Puhani 2000).

Both estimators rely on normality. The correction is only as good as the selection equation and the distributional assumption behind it.

In OpenEconometrics

oe.heckman fits both estimators. The model in the example corresponds to the Stata command heckman log_wage educ age, select(works = educ age kids other_inc), and method="twostep" to the same command with the twostep option:

oe.heckman(data=df, y="log_wage", x=["educ", "age"], select="works", select_x=["educ", "age", "kids", "other_inc"])

• select names a 0/1 column. y may be missing where select is 0; values recorded there are ignored. The outcome regressors must be complete in every row.
• select_x lists the selection regressors; a constant is always added. Include the outcome regressors and at least one excluded variable.
• method="ml" (the default) or "twostep".
• For ML, covariance may be "nonrobust" (the observed information, the default), "opg", "robust" or "cluster" with cluster= naming one or two columns. ML accepts frequency, analytic, sampling and importance weights; sampling weights need a robust or cluster covariance. The two-step estimator has its own covariance, accepts only "nonrobust" and takes frequency weights only.
• ML reports the outcome equation, the selection equation as select:<term>, and /athrho and /lnsigma, the parameters it estimates; metrics give rho, sigma, lambda = ρσ, n_selected and n_censored. The two-step fit reports the coefficient on λ as mills:lambda.
• tests["rho"] is the likelihood-ratio test of independent equations under nonrobust and opg covariances and a Wald test of athrho = 0 under robust and cluster.

For a binary outcome observed with selection, oe.heckprobit fits the probit version. The limited-dependent-variable guide gives the likelihood, the two-step covariance and the conventions.

Supported options

Options of oe.heckman

Generated from oe.capabilities() of openecon 0.3.20a1. Stata names identify the corresponding workflow, not numerical parity.

Example

The script simulates 3,000 women. A woman works when 0.8 + 0.12(educ − 12) − 0.6·kids − 0.3·other_inc + u₂ > 0, and her log wage is 1.0 + 0.08·educ + 0.01·age + u₁, observed only if she works. The errors have σ = 0.5 and ρ = 0.5, so λ = ρσ = 0.25. Young children and other household income affect participation but not wages, which identifies the model.

import numpy as np
import pandas as pd
import openecon as oe


def simulate(n=3000, rho=0.5, sigma=0.5, seed=13):
    """Wages are observed only for women who work; the two equations' errors are correlated."""
    rng = np.random.default_rng(seed)
    educ = rng.normal(12, 2.5, size=n)
    age = rng.integers(25, 61, size=n).astype(float)
    kids = rng.poisson(0.8, size=n).astype(float)          # children under six
    other_inc = rng.exponential(1.0, size=n)               # other household income
    u_wage, u_select = rng.multivariate_normal(
        [0, 0], [[sigma**2, rho * sigma], [rho * sigma, 1]], size=n).T
    works = (0.8 + 0.12 * (educ - 12) - 0.6 * kids - 0.3 * other_inc + u_select > 0).astype(float)
    log_wage = np.where(works == 1, 1.0 + 0.08 * educ + 0.01 * age + u_wage, np.nan)
    return pd.DataFrame({"log_wage": log_wage, "works": works, "educ": educ, "age": age,
                         "kids": kids, "other_inc": other_inc})


if __name__ == "__main__":
    df = simulate()
    spec = dict(data=df, y="log_wage", x=["educ", "age"], select="works",
                select_x=["educ", "age", "kids", "other_inc"])
    ml = oe.heckman(**spec)
    print(ml.summary())
    print()
    workers = df[df["works"] == 1]
    ols = oe.ols(data=workers, y="log_wage", x=["educ", "age"], covariance="HC1")
    two = oe.heckman(**spec, method="twostep")
    print(f"{'estimator':<22}{'educ':>8}{'s.e.':>8}{'rho':>8}{'sigma':>8}{'lambda':>8}")
    print(f"{'OLS on workers':<22}{ols.coefficients[1].estimate:>8.4f}{ols.coefficients[1].std_error:>8.4f}")
    for label, fit in [("Heckman ML", ml), ("Heckman two-step", two)]:
        educ = next(c for c in fit.coefficients if c.term == "educ")
        m = fit.metrics
        print(f"{label:<22}{educ.estimate:>8.4f}{educ.std_error:>8.4f}"
              f"{m['rho']:>8.3f}{m['sigma']:>8.3f}{m['lambda']:>8.3f}")

With openecon 0.3.20a1 from PyPI the script prints:

Heckman selection model — log_wage
Observations: 3000  |  Covariance: nonrobust  |  Confidence: 95%

Term                 Estimate  Std. error         z      P>|stat|     CI lower    CI upper
----------------  -----------  ----------  --------  ------------  -----------  ----------
[log_wage]                                                                                
Intercept             1.02336   0.0982622   10.4146   2.12761e-25     0.830768     1.21595
educ                0.0783991  0.00552163   14.1986   9.35266e-46    0.0675769   0.0892213
age                 0.0102444  0.00120434   8.50623    1.7968e-17   0.00788392   0.0126048
[select]                                                                                  
select:Intercept    -0.689395    0.160679  -4.29052   1.78256e-05     -1.00432   -0.374471
select:educ          0.118457   0.0101489   11.6719   1.77409e-31    0.0985651    0.138348
select:age        0.000896083  0.00235889  0.379875      0.704038  -0.00372725  0.00551942
select:kids         -0.572559   0.0303715  -18.8519   2.83672e-79    -0.632086   -0.513032
select:other_inc    -0.292848   0.0260121  -11.2581   2.11185e-29    -0.343831   -0.241865
/athrho              0.493356    0.089159   5.53344   3.14005e-08     0.318608    0.668105
/lnsigma            -0.691846   0.0261431  -26.4639  2.52752e-154    -0.743086   -0.640607

log_likelihood: -2805.52  |  aic: 5631.05  |  bic: 5691.11  |  rho: 0.456876
sigma: 0.500651  |  lambda: 0.228735  |  n_selected: 1562  |  n_censored: 1438
Wald chi2 test of the outcome slopes: chi2(2) = 274.869, p = 2.055e-60
LR test of independent equations (rho = 0): chi2(1) = 29.053, p = 7.042e-08

estimator                 educ    s.e.     rho   sigma  lambda
OLS on workers          0.0650  0.0047
Heckman ML              0.0784  0.0055   0.457   0.501   0.229
Heckman two-step        0.0796  0.0057   0.494   0.506   0.250

Reading the output

• OLS on workers is biased. It estimates the return to education at 0.0650, about three standard errors below the true 0.08. Education raises participation, so women with little education who work anyway tend to have favourable unobservables, and because ρ > 0 those also raise their wages. That flattens the education gradient among workers.
• Maximum likelihood recovers the outcome equation. The return to education is 0.0784 (95% interval 0.0676 to 0.0892) and the age effect 0.0102; both intervals contain the true values. The selection equation is also close: 0.118 for education, −0.573 for children, −0.293 for other income and 0.0009 (p = 0.70) for age, whose true value is 0.
• The selection parameters are estimated with some error. ρ is 0.457, the hyperbolic tangent of /athrho = 0.4934; transforming the /athrho interval gives about 0.31 to 0.58 for ρ, which contains 0.5. σ = exp(/lnsigma) is 0.501, and λ = ρσ is 0.229 against 0.25.
• Selection matters. The likelihood-ratio test of independent equations gives χ²(1) = 29.05 with p = 7.0e-08, so the hypothesis ρ = 0 is rejected.
• The two-step estimator agrees. It gives 0.0796 with a standard error of 0.0057, against 0.0055 for ML, and ρ = 0.494. With strong excluded variables, as here, the two agree closely; a large gap between them would be a reason to check the specification.

Limitations

• Cluster-robust covariance for oe.heckman allows at most 2 clustering dimensions.
• The selection indicator must be given explicitly, the selection equation always has a constant, and the outcome regressors must be complete in every row, including nonselected ones.
• Offsets and Stata's first, mills() and nshazard() options are not implemented.
• The two-step estimator offers only its nonrobust covariance and accepts only frequency weights.
• The likelihood is not globally concave: the reported maximum is the one reached from the two-step starting values, and a correlation that runs to ±1 raises boundary_solution.
• No numerical parity with Stata is claimed: the inventory's stata_parity_validated flag is False.

References

• Gronau, R. (1974). Wage comparisons—a selectivity bias. Journal of Political Economy, 82(6), 1119–1143. doi:10.1086/260267
• Heckman, J. J. (1979). Sample selection bias as a specification error. Econometrica, 47(1), 153–161. doi:10.2307/1912352
• Puhani, P. (2000). The Heckman correction for sample selection and its critique. Journal of Economic Surveys, 14(1), 53–68. doi:10.1111/1467-6419.00104