Treatment effects in Python: propensity score matching, IPW and AIPW

People who join a training programme differ from people who do not. If those with low earnings are the ones who sign up, a simple comparison of later earnings mixes the effect of training with the earnings gap that was already there. When the characteristics that drive both participation and outcomes are observed, treatment-effects estimators remove that gap. They compare like with like by modelling the outcome, the probability of treatment, or both, or by matching each person to similar people in the other group.

What treatment-effects estimators estimate

Each person has two potential outcomes: y(1) with treatment and y(0) without. Only one is observed. The average treatment effect (ATE) is E[y(1) − y(0)] over the whole population; the average treatment effect on the treated (ATET) is the same average among those who were treated. They differ when the effect varies with characteristics that also affect who is treated.

Two assumptions identify both. Unconfoundedness: given the covariates x, treatment is as good as random, so no unobserved factor drives both treatment and outcomes. Overlap: every combination of covariates has some chance of each treatment status. The first cannot be tested; it rests on knowing how treatment was assigned. The second can be checked by looking at the estimated propensity scores, the probabilities p(x) of treatment given x. Rosenbaum and Rubin (1983) show that conditioning on the propensity score alone is enough to remove the bias from observed covariates.

• Regression adjustment (RA) fits an outcome model for each treatment group and averages the predicted difference over the sample.
• Inverse-probability weighting (IPW) weights each treated person by 1/p(x) and each control by 1/(1 − p(x)), so both groups resemble the full population.
• Doubly robust estimators combine the two. Augmented IPW (AIPW, Robins, Rotnitzky and Zhao 1994) and IPW regression adjustment (IPWRA) are consistent if either the outcome model or the propensity model is correct.
• Matching imputes each missing potential outcome from the nearest people in the other group, on the covariates or on the estimated propensity score. Abadie and Imbens (2006) show that with several continuous covariates matching carries a bias that can dominate in large samples, that matching with a fixed number of neighbours is not efficient, and how to estimate its variance consistently; biasadj= applies a regression correction for the bias. Abadie and Imbens (2016) give the variance correction for matching on an estimated propensity score.

In OpenEconometrics

oe.teffects covers all six estimators with one call shape: oe.teffects(data=df, y="earnings", treatment="training", x=["prior", "age"], method="aipw", estimand="ate"). The outcome covariates go in x; the treatment-model covariates go in tx and default to x. Each method corresponds to a Stata teffects subcommand:

• estimand="ate" (default), "atet" or "pomeans" (potential-outcome means; not for matching).
• RA, IPW, IPWRA and AIPW also handle treatments with more than two levels, through a multinomial logit treatment model; control= and tlevel= choose the levels.
• Their standard errors come from the stacked estimating equations of the treatment model, the outcome models and the means, so the estimation of the propensity score and the outcome models is accounted for. covariance="cluster" with cluster= clusters them.
• Matching standard errors follow Abadie and Imbens (2006), with the 2016 correction for psmatch.
• result.extra["overlap"] summarizes the estimated propensity scores by treatment group, the numbers behind Stata's teoverlap plot.

The treatment-effects family guide gives every estimating equation and variance formula.

Supported options

Options of oe.teffects

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

Example

The script simulates 2,000 workers. Prior earnings (in thousands of dollars) have mean 20 and standard deviation 5, and age has mean 35 and standard deviation 8. Workers with low prior earnings are more likely to take a training course, through a logit propensity model. Later earnings depend linearly on both covariates, and the individual effect of training is 2 − 0.2 · (prior − 20): larger for those who earned less. The population ATE is 2.0; the ATET is higher because trainees earned less. The script prints the AIPW summary, a table of estimators and, for comparison, the sample means of the true individual effects.

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


def simulate(n=2000, seed=21):
    """Job training chosen more often by people with low prior earnings (in $1,000).
    Individual effect: 2 - 0.2 * (prior - 20), so the ATE is 2 and the ATET is larger."""
    rng = np.random.default_rng(seed)
    prior = 20 + 5 * rng.normal(size=n)
    age = 35 + 8 * rng.normal(size=n)
    p = 1 / (1 + np.exp(-(-0.3 - 0.16 * (prior - 20) + 0.05 * (age - 35))))
    training = rng.binomial(1, p)
    effect = 2 - 0.2 * (prior - 20)
    earnings = 22 + 0.6 * (prior - 20) + 0.1 * (age - 35) + effect * training + rng.normal(scale=3, size=n)
    return pd.DataFrame({"earnings": earnings, "training": training, "prior": prior, "age": age,
                         "effect": effect})


if __name__ == "__main__":
    df = simulate()
    x = ["prior", "age"]
    aipw = oe.teffects(data=df, y="earnings", treatment="training", x=x, method="aipw")
    print(aipw.summary())
    naive = oe.ols(data=df, y="earnings", x=["training"], covariance="HC1").coefficients[1]
    rows = [("difference in means", "-", naive)]
    for method, estimand, extra in [("ra", "ate", {}), ("ipw", "ate", {}), ("aipw", "ate", {}),
                                    ("psmatch", "ate", {}), ("ipwra", "atet", {}),
                                    ("nnmatch", "atet", {"biasadj": x})]:
        fit = oe.teffects(data=df, y="earnings", treatment="training", x=x, method=method,
                          estimand=estimand, **extra)
        rows.append((method, estimand, fit.coefficients[0]))
    print("\nMethod               Estimand  Estimate  Std. error  95% CI")
    for method, estimand, c in rows:
        print(f"{method:<20} {estimand:<8} {c.estimate:9.3f}  {c.std_error:10.3f}  [{c.ci_low:.3f}, {c.ci_high:.3f}]")
    print(f"Mean true effect: all {df.effect.mean():.3f}, trained {df.effect[df.training == 1].mean():.3f}")

With openecon 0.3.20a1 from PyPI the script prints:

Treatment-effects estimation: augmented IPW — earnings
Observations: 2000  |  Covariance: robust  |  Confidence: 95%

Term                Estimate  Std. error        z     P>|stat|  CI lower  CI upper
------------------  --------  ----------  -------  -----------  --------  --------
[ATE]
ATE:r1vs0.training   1.89193    0.154357  12.2569  1.54365e-34    1.5894   2.19447
[POmean]
POmean:0.training    22.1618    0.114203  194.055            0   21.9379   22.3856

n_levels: 2  |  n_control: 1157  |  n_treated: 843  |  treatment_log_likelihood: -1207.37
propensity_min: 0.0327554  |  propensity_max: 0.931165

Method               Estimand  Estimate  Std. error  95% CI
difference in means  -            0.478       0.175  [0.134, 0.821]
ra                   ate          1.876       0.154  [1.574, 2.178]
ipw                  ate          1.903       0.160  [1.590, 2.215]
aipw                 ate          1.892       0.154  [1.589, 2.194]
psmatch              ate          2.019       0.177  [1.672, 2.366]
ipwra                atet         2.369       0.164  [2.046, 2.691]
nnmatch              atet         2.527       0.207  [2.122, 2.932]
Mean true effect: all 1.966, trained 2.363

Reading the output

• The raw comparison is badly biased. Trainees earn only 0.478 more, with an interval from 0.134 to 0.821, because they started with lower earnings. The true average effect in this sample is 1.966.
• The ATE estimators agree with each other and with the truth. RA, IPW, AIPW and propensity-score matching give 1.876, 1.903, 1.892 and 2.019. Every interval contains 1.966. Here both the outcome and the propensity models are correctly specified, so all four are consistent; double robustness matters when one of them is not.
• The ATET is larger, as it should be. IPWRA estimates 2.369 against a true 2.363 among the trained. Nearest-neighbour matching with bias adjustment gives 2.527 with an interval from 2.122 to 2.932.
• Matching is less precise. The matching standard errors, 0.177 and 0.207, exceed the 0.154 to 0.164 of the model-based estimators. Matching with a fixed number of neighbours, here one, is not efficient (Abadie and Imbens 2006).
• Overlap is adequate. The estimated propensity scores run from 0.033 to 0.931, well away from 0 and 1. The POmean row, 22.16, is the estimated mean earnings had nobody been trained.

Limitations

• The common oe.predict and oe.margins adapters do not cover oe.teffects; use its own result tables and family helpers.
• Cluster-robust covariance for oe.teffects allows at most 1 clustering dimension.
• Matching (nnmatch, psmatch) needs a binary treatment and the ate or atet estimand, accepts only covariance="robust" and takes frequency weights by bounded row duplication, not sampling weights.
• Any estimated treatment probability below pstolerance (default 1e-5), or a treatment model that separates the groups perfectly, raises overlap_violation.
• A caliper that leaves any observation without enough matches raises caliper_violation instead of dropping the observation.
• Stata's hetprobit treatment model and wnls outcome option are not implemented. nnmatch on several covariates scales with the product of the group sizes, and a search beyond about 4e9 distance evaluations raises matching_too_large.
• The cluster sandwich of ra, ipw, ipwra and aipw applies no small-sample factor.
• No numerical parity with Stata is claimed: the inventory's stata_parity_validated flag is False.

References

• Abadie, A. and Imbens, G. W. (2006). Large sample properties of matching estimators for average treatment effects. Econometrica, 74(1), 235–267. doi:10.1111/j.1468-0262.2006.00655.x
• Abadie, A. and Imbens, G. W. (2016). Matching on the estimated propensity score. Econometrica, 84(2), 781–807. doi:10.3982/ECTA11293
• Robins, J. M., Rotnitzky, A. and Zhao, L. P. (1994). Estimation of regression coefficients when some regressors are not always observed. Journal of the American Statistical Association, 89(427), 846–866. doi:10.1080/01621459.1994.10476818
• Rosenbaum, P. R. and Rubin, D. B. (1983). The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1), 41–55. doi:10.1093/biomet/70.1.41