Skip to content

Instantly share code, notes, and snippets.

@josef-pkt

josef-pkt/ex_treatment_effect.ipynb

Created February 8, 2022 17:28
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save josef-pkt/8a8755d2930cb23f7282cf683232bf17 to your computer and use it in GitHub Desktop.
Save josef-pkt/8a8755d2930cb23f7282cf683232bf17 to your computer and use it in GitHub Desktop.
Download ZIP
treatment effect ATE and POM under conditional independence
Display the source blob
Display the rendered blob
Raw
ex_treatment_effect.ipynb
{
"cells": [
{
"cell_type": "markdown",
"id": "funded-italic",
"metadata": {},
"source": [
"## Treatment effects under conditional independence\n",
"\n",
"example for new `statsmodels.treatment.treatment_effects.TreatmentEffect`"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "adult-capture",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Optimization terminated successfully.\n",
" Current function value: 0.439575\n",
" Iterations 6\n"
]
}
],
"source": [
"import os\n",
"import numpy as np\n",
"from numpy.testing import assert_allclose\n",
"import pandas as pd\n",
"import pytest\n",
"\n",
"from statsmodels.regression.linear_model import OLS\n",
"from statsmodels.discrete.discrete_model import Probit\n",
"from statsmodels.treatment.treatment_effects import (\n",
" TreatmentEffect\n",
" )\n",
"\n",
"from statsmodels.treatment.tests.results import results_teffects as res_st\n",
"\n",
"\n",
"cur_dir = os.path.abspath(os.path.dirname(res_st.__file__))\n",
"\n",
"file_name = 'cataneo2.csv'\n",
"file_path = os.path.join(cur_dir, file_name)\n",
"\n",
"dta_cat = pd.read_csv(file_path)\n",
"\n",
"formula = 'mbsmoke_ ~ mmarried_ + mage + mage2 + fbaby_ + medu'\n",
"res_probit = Probit.from_formula(formula, dta_cat).fit()\n",
"\n",
"methods = [\n",
" (\"ra\", res_st.results_ra),\n",
" (\"ipw\", res_st.results_ipw),\n",
" (\"aipw\", res_st.results_aipw),\n",
" (\"aipw_wls\", res_st.results_aipw_wls),\n",
" (\"ipw_ra\", res_st.results_ipwra),\n",
" ]"
]
},
{
"cell_type": "markdown",
"id": "talented-bidding",
"metadata": {},
"source": [
"### create TreatmentEffect instance and compute ipw"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "controlling-literacy",
"metadata": {},
"outputs": [],
"source": [
"formula_outcome = 'bweight ~ prenatal1_ + mmarried_ + mage + fbaby_'\n",
"mod = OLS.from_formula(formula_outcome, dta_cat)\n",
"tind = np.asarray(dta_cat['mbsmoke_'])\n",
"teff = TreatmentEffect(mod, tind, results_select=res_probit)"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "former-mirror",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<class 'statsmodels.treatment.treatment_effects.TreatmentEffectResults'>\n",
" Test for Constraints \n",
"==============================================================================\n",
" coef std err z P>|z| [0.025 0.975]\n",
"------------------------------------------------------------------------------\n",
"ATE -230.6886 25.817 -8.936 0.000 -281.289 -180.089\n",
"POM0 3403.4627 9.571 355.586 0.000 3384.703 3422.222\n",
"POM1 3172.7741 24.001 132.193 0.000 3125.733 3219.815\n",
"=============================================================================="
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res = teff.ipw()\n",
"res"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "metropolitan-burns",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>coef</th>\n",
" <th>std err</th>\n",
" <th>z</th>\n",
" <th>P&gt;|z|</th>\n",
" <th>Conf. Int. Low</th>\n",
" <th>Conf. Int. Upp.</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>ATE</th>\n",
" <td>-230.688638</td>\n",
" <td>25.816745</td>\n",
" <td>-8.935621</td>\n",
" <td>4.048989e-19</td>\n",
" <td>-281.288528</td>\n",
" <td>-180.088748</td>\n",
" </tr>\n",
" <tr>\n",
" <th>POM0</th>\n",
" <td>3403.462709</td>\n",
" <td>9.571412</td>\n",
" <td>355.586256</td>\n",
" <td>0.000000e+00</td>\n",
" <td>3384.703085</td>\n",
" <td>3422.222332</td>\n",
" </tr>\n",
" <tr>\n",
" <th>POM1</th>\n",
" <td>3172.774071</td>\n",
" <td>24.001048</td>\n",
" <td>132.193145</td>\n",
" <td>0.000000e+00</td>\n",
" <td>3125.732880</td>\n",
" <td>3219.815261</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" coef std err z P>|z| Conf. Int. Low \\\n",
"ATE -230.688638 25.816745 -8.935621 4.048989e-19 -281.288528 \n",
"POM0 3403.462709 9.571412 355.586256 0.000000e+00 3384.703085 \n",
"POM1 3172.774071 24.001048 132.193145 0.000000e+00 3125.732880 \n",
"\n",
" Conf. Int. Upp. \n",
"ATE -180.088748 \n",
"POM0 3422.222332 \n",
"POM1 3219.815261 "
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res.summary_frame()"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "bibliographic-internship",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array([-2.30688638e+02, 3.40346271e+03, -1.55825526e+00, -6.48482138e-01,\n",
" 1.74432697e-01, -3.25591262e-03, -2.17596162e-01, -8.63630870e-02])"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res.start_params"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "improving-inspector",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array([0., 0., 0., 0., 0., 0., 0., 0.])"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res.results_gmm.params - res.start_params"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "democratic-monroe",
"metadata": {},
"outputs": [],
"source": [
"pd.set_option('display.width', 500)"
]
},
{
"cell_type": "markdown",
"id": "banned-rogers",
"metadata": {},
"source": [
"## all methods in TreatmentEffect"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "innovative-garage",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" ra\n",
" coef std err z P>|z| Conf. Int. Low Conf. Int. Upp.\n",
"ATE -239.639211 23.824021 -10.058722 8.408247e-24 -286.333435 -192.944988\n",
"POM0 3403.242272 9.525207 357.288006 0.000000e+00 3384.573209 3421.911335\n",
"POM1 3163.603060 21.863509 144.697867 0.000000e+00 3120.751371 3206.454750\n",
"\n",
" ipw\n",
" coef std err z P>|z| Conf. Int. Low Conf. Int. Upp.\n",
"ATE -230.688638 25.816745 -8.935621 4.048989e-19 -281.288528 -180.088748\n",
"POM0 3403.462709 9.571412 355.586256 0.000000e+00 3384.703085 3422.222332\n",
"POM1 3172.774071 24.001048 132.193145 0.000000e+00 3125.732880 3219.815261\n",
"\n",
" aipw\n",
" coef std err z P>|z| Conf. Int. Low Conf. Int. Upp.\n",
"ATE -230.989648 26.214445 -8.811541 1.234375e-18 -282.369017 -179.610280\n",
"POM0 3403.355674 9.568514 355.682783 0.000000e+00 3384.601731 3422.109616\n",
"POM1 3172.366025 24.427402 129.869153 0.000000e+00 3124.489197 3220.242854\n",
"\n",
" aipw_wls\n",
" coef std err z P>|z| Conf. Int. Low Conf. Int. Upp.\n",
"ATE -227.195618 27.372036 -8.300282 1.038645e-16 -280.843822 -173.547414\n",
"POM0 3403.250651 9.596571 354.631943 0.000000e+00 3384.441717 3422.059585\n",
"POM1 3176.055033 25.654642 123.800406 0.000000e+00 3125.772859 3226.337206\n",
"\n",
" ipw_ra\n",
" coef std err z P>|z| Conf. Int. Low Conf. Int. Upp.\n",
"ATE -229.967078 26.629411 -8.635830 5.830196e-18 -282.159765 -177.774391\n",
"POM0 3403.335639 9.571288 355.577620 0.000000e+00 3384.576260 3422.095018\n",
"POM1 3173.368561 24.871955 127.588224 0.000000e+00 3124.620425 3222.116697\n"
]
}
],
"source": [
"teff = TreatmentEffect(mod, tind, results_select=res_probit)\n",
"\n",
"for m, _ in methods:\n",
" res = getattr(teff, m)()\n",
" print(\"\\n\", m)\n",
" print(res.summary_frame())"
]
},
{
"cell_type": "markdown",
"id": "broadband-export",
"metadata": {},
"source": [
"## results in Stata\n",
"\n",
"from unit tests"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "compact-pixel",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" ra\n",
" b se z pvalue ll ul\n",
"ATE -239.639211 23.824021 -10.058722 8.408247e-24 -286.333435 -192.944988\n",
"POM0 3403.242272 9.525207 357.288005 0.000000e+00 3384.573209 3421.911335\n",
"\n",
" ipw\n",
" b se z pvalue ll ul\n",
"ATE -230.688638 25.815244 -8.936140 4.030006e-19 -281.285586 -180.091690\n",
"POM0 3403.462709 9.571369 355.587873 0.000000e+00 3384.703170 3422.222247\n",
"\n",
" aipw\n",
" b se z pvalue ll ul\n",
"ATE -230.989201 26.210565 -8.812828 1.220276e-18 -282.360964 -179.617438\n",
"POM0 3403.355253 9.568472 355.684297 0.000000e+00 3384.601393 3422.109114\n",
"\n",
" aipw_wls\n",
" b se z pvalue ll ul\n",
"ATE -227.195618 27.347936 -8.307597 9.765984e-17 -280.796587 -173.594649\n",
"POM0 3403.250651 9.596622 354.630065 0.000000e+00 3384.441618 3422.059684\n",
"\n",
" ipw_ra\n",
" b se z pvalue ll ul\n",
"ATE -229.967078 26.626676 -8.636718 5.785117e-18 -282.154403 -177.779752\n",
"POM0 3403.335639 9.571260 355.578657 0.000000e+00 3384.576315 3422.094963\n"
]
}
],
"source": [
"for m, st in methods:\n",
" print(\"\\n\", m)\n",
" res = pd.DataFrame(st.table[:2, :6], index = [\"ATE\", \"POM0\"], columns=st.table_colnames[:6])\n",
" print(res)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "valued-contrary",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "physical-newman",
"metadata": {},
"source": [
"### treatment effects without inference"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "preceding-dividend",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" ra\n",
"(-239.63921146433722, 3403.2422719354868, 3163.6030604711495)\n",
"\n",
" ipw\n",
"(-230.68863779526237, 3403.4627086845567, 3172.7740708892948)\n",
"\n",
" aipw\n",
"(-230.98920111257803, 3403.355253173835, 3172.366052061257)\n",
"\n",
" aipw_wls\n",
"(-227.19561818674902, 3403.2506509757864, 3176.0550327890373)\n",
"\n",
" ipw_ra\n",
"(-229.96707793512905, 3403.335639307418, 3173.368561372289)\n"
]
}
],
"source": [
"\n",
"\n",
"for m, st in methods:\n",
" print(\"\\n\", m)\n",
" res = getattr(teff, m)(return_results=False)\n",
" print(res)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "acoustic-summer",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment