Load data for example

Treatment effects under conditional independence

Author: Josef Perktold

This notebook illustrates the basic usage of the new treatment effect functionality in statsmodels.

The main class is statsmodels.treatment.treatment_effects.TreatmentEffect.

This class estimates treatment effect and potential outcome using 5 different methods, ipw, ra, aipw, aipw-wls, ipw-ra. The last three methods require both a treatment or selection model and an outcome model. Standard errors and inference are based on the joint GMM representation of selection or treatment model, outcome model and effect functions. The approach for inference follows Stata, however Stata support a wider range of models. Estimation and inference are valid under conditional independence or ignorability.

The outcome model is currently limited to a linear model based on OLS. Treatment is currently restricted to binary treatment which can be either Logit or Probit.

The example follows Cattaneo.

import os

import numpy as np
import pandas as pd
from numpy.testing import assert_allclose
from statsmodels.discrete.discrete_model import Probit
from statsmodels.regression.linear_model import OLS
from statsmodels.treatment.tests.results import results_teffects as res_st
from statsmodels.treatment.treatment_effects import TreatmentEffect

# Load data for example
cur_dir = os.path.abspath(os.path.dirname(res_st.__file__))
file_name = "cataneo2.csv"
file_path = os.path.join(cur_dir, file_name)
dta_cat = pd.read_csv(file_path)

methods = ["ra", "ipw", "aipw", "aipw_wls", "ipw_ra"]
methods_st = [
    ("ra", res_st.results_ra),
    ("ipw", res_st.results_ipw),
    ("aipw", res_st.results_aipw),
    ("aipw_wls", res_st.results_aipw_wls),
    ("ipw_ra", res_st.results_ipwra),
]

# allow wider display of data frames
pd.set_option("display.width", 500)

dta_cat.head()

Create TreatmentEffect instance and compute ipw

The TreatmentEffect class requires

• a OLS model instance for the outcome model,
• a results instance of the selection model and
• a treatment indicator variable.

In the following example we use Probit as the selection model. Using Logit is also supported.

# treatment selection model
formula = "mbsmoke_ ~ mmarried_ + mage + mage2 + fbaby_ + medu"
res_probit = Probit.from_formula(formula, dta_cat).fit()

# outcome model
formula_outcome = "bweight ~ prenatal1_ + mmarried_ + mage + fbaby_"
mod = OLS.from_formula(formula_outcome, dta_cat)

# treatment indicator variable
tind = np.asarray(dta_cat["mbsmoke_"])

teff = TreatmentEffect(mod, tind, results_select=res_probit)

After creating the TreatmentEffect instance, we can call any of the 5 methods to compute potential outcomes, POM0, POM1, and average treatment effect, ATE. POM0 is the potential outcome for the no treatment group, POM1 is the potential outcome for the treatment group, treatment effect is POM1 - POM0.

For example teff.ipw() computes POM and ATE using inverse probability weighting. The probability of treatment is also commonly called the propensity score. The summary of the estimation includes standard errors and confidence interval for POM and ATE.

Standard errors and other inferential statistics are based on the Generalized Method of Moments (GMM) representation of the selection and outcome models and the moment conditions for the results statistic. Method ipw uses the selection model but not the outcome model. Method ra uses the outcome model but not the selection model. The doubly robust estimators aipw, aipw-wls, ipw-ra include both selection and outcome models, where at least one of those two has to be correctly specified to get consistent estimates of the treatment effect. The moment conditions for the target variables, POM0, POM1, and ATE are based on POM0 and ATE. The remaining POM1 is computed as a linear combination of POM0 and ATE.

The internal gmm results are attached to the treatment results as results_gmm.

By default the treatment effect methods computes average treatment effect, where average is take over the sample observations. Option effect_group can be used to compute either average treatment effect on the treated, ATT, using effect_group=1 or average treatment effect on the non-treated using effect_group=0.

res = teff.ipw()
res

res.summary_frame()

print(res.results_gmm.summary())

average treatment effect on the treated

see more below

teff.ipw(effect_group=1)

average treatment effect on the untreated

teff.ipw(effect_group=0)

Other methods to compute ATE work in the same or similar way as for ipw for example regression adjustment ra and double robust ipw_ra.

res_ra = teff.ra()
res_ra

res_ra.summary_frame()

ra2 = teff.ipw_ra(effect_group=1, return_results=True)
ra2.summary_frame()

All methods in TreatmentEffect

The following computes and prints ATE and POM for all methods. (We include the call to TreatmentEffect as a reminder.)

teff = TreatmentEffect(mod, tind, results_select=res_probit)

for m in methods:
    res = getattr(teff, m)()
    print("\n", m)
    print(res.summary_frame())

Results in Stata

The results in statsmodels are very close to the results in Stata because both packages use the same approach.

for m, st in methods_st:
    print("\n", m)
    res = pd.DataFrame(
        st.table[:2, :6], index=["ATE", "POM0"], columns=st.table_colnames[:6]
    )
    print(res)

Treatment effects without inference

It is possible to compute POM and ATE without computing standard errors and inferential statistics. In this case the GMM model is not computed.

for m in methods:
    print("\n", m)
    res = getattr(teff, m)(return_results=False)
    print(res)

Treatment effect on the treated

Treatment effects on subgroups are not available for aipw and aipw-wls.

effect_group choses the group for which treatement effect and potential outcomes are computed Options are "all" for sample average treatment effect, 1 for average treatment effect on the treated and 0 for average treatment effect on the untreated.

Note: The row labels in the pandas dataframe, POM and ATE, are the same even for treatment effect on subgroups.

for m in methods:
    if m.startswith("aipw"):
        continue
    res = getattr(teff, m)(effect_group=1)
    print("\n", m)
    print(res.summary_frame())

Treatment effect on the untreated

Similar to ATT, we can compute average treatment effect on the untreated by using option effect_group=0.

for m in methods:
    if m.startswith("aipw"):
        # not available
        continue
    res = getattr(teff, m)(effect_group=0)
    print("\n", m)
    print(res.summary_frame())

The docstring for the TreatmentEffect class and it's methods can be obtained using help

help(teff)