{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "funded-italic",
   "metadata": {},
   "source": [
    "## Treatment effects under conditional independence\n",
    "\n",
    "Author: Josef Perktold\n",
    "\n",
    "This notebook illustrates the basic usage of the new treatment effect functionality in statsmodels.\n",
    "\n",
    "The main class is `statsmodels.treatment.treatment_effects.TreatmentEffect`.\n",
    "\n",
    "\n",
    "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.\n",
    "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.\n",
    "Estimation and inference are valid under conditional independence or ignorability.\n",
    "\n",
    "The outcome model is currently limited to a linear model based on OLS.\n",
    "Treatment is currently restricted to binary treatment which can be either Logit or Probit.\n",
    "\n",
    "The example follows Cattaneo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "adult-capture",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from numpy.testing import assert_allclose\n",
    "import pandas as pd \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",
    "# Load data for example\n",
    "cur_dir = os.path.abspath(os.path.dirname(res_st.__file__))\n",
    "file_name = 'cataneo2.csv'\n",
    "file_path = os.path.join(cur_dir, file_name)\n",
    "dta_cat = pd.read_csv(file_path)\n",
    "\n",
    "methods = ['ra', 'ipw', 'aipw', 'aipw_wls', 'ipw_ra']\n",
    "methods_st = [\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",
    "    ]\n",
    "\n",
    "# allow wider display of data frames\n",
    "pd.set_option('display.width', 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "32e76526",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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>bweight</th>\n",
       "      <th>mmarried</th>\n",
       "      <th>mhisp</th>\n",
       "      <th>fhisp</th>\n",
       "      <th>foreign</th>\n",
       "      <th>alcohol</th>\n",
       "      <th>deadkids</th>\n",
       "      <th>mage</th>\n",
       "      <th>medu</th>\n",
       "      <th>fage</th>\n",
       "      <th>...</th>\n",
       "      <th>prenatal</th>\n",
       "      <th>birthmonth</th>\n",
       "      <th>lbweight</th>\n",
       "      <th>fbaby</th>\n",
       "      <th>prenatal1</th>\n",
       "      <th>mbsmoke_</th>\n",
       "      <th>mmarried_</th>\n",
       "      <th>fbaby_</th>\n",
       "      <th>prenatal1_</th>\n",
       "      <th>mage2</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3459</td>\n",
       "      <td>married</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>24</td>\n",
       "      <td>14</td>\n",
       "      <td>28</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>0</td>\n",
       "      <td>No</td>\n",
       "      <td>Yes</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>576.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3260</td>\n",
       "      <td>notmarried</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>20</td>\n",
       "      <td>10</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>0</td>\n",
       "      <td>No</td>\n",
       "      <td>Yes</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>400.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3572</td>\n",
       "      <td>married</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>22</td>\n",
       "      <td>9</td>\n",
       "      <td>30</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>No</td>\n",
       "      <td>Yes</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>484.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2948</td>\n",
       "      <td>married</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>26</td>\n",
       "      <td>12</td>\n",
       "      <td>30</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>No</td>\n",
       "      <td>Yes</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>676.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2410</td>\n",
       "      <td>married</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>20</td>\n",
       "      <td>12</td>\n",
       "      <td>21</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>Yes</td>\n",
       "      <td>Yes</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>400.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 28 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   bweight    mmarried  mhisp  fhisp  foreign  alcohol  deadkids  mage  medu  fage  ...  prenatal  birthmonth  lbweight  fbaby prenatal1 mbsmoke_  mmarried_  fbaby_  prenatal1_  mage2\n",
       "0     3459     married      0      0        0        0         0    24    14    28  ...         1          12         0     No       Yes        0          1       0           1  576.0\n",
       "1     3260  notmarried      0      0        1        0         0    20    10     0  ...         1           7         0     No       Yes        0          0       0           1  400.0\n",
       "2     3572     married      0      0        1        0         0    22     9    30  ...         1           3         0     No       Yes        0          1       0           1  484.0\n",
       "3     2948     married      0      0        0        0         0    26    12    30  ...         1           1         0     No       Yes        0          1       0           1  676.0\n",
       "4     2410     married      0      0        0        0         0    20    12    21  ...         1           3         1    Yes       Yes        0          1       1           1  400.0\n",
       "\n",
       "[5 rows x 28 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dta_cat.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "talented-bidding",
   "metadata": {},
   "source": [
    "### Create TreatmentEffect instance and compute ipw\n",
    "\n",
    "The TreatmentEffect class requires \n",
    "- a OLS model instance for the outcome model, \n",
    "- a results instance of the selection model and \n",
    "- a treatment indicator variable.\n",
    "\n",
    "In the following example we use Probit as the selection model. Using Logit is also supported.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "controlling-literacy",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimization terminated successfully.\n",
      "         Current function value: 0.439575\n",
      "         Iterations 6\n"
     ]
    }
   ],
   "source": [
    "# treatment selection model\n",
    "formula = 'mbsmoke_ ~ mmarried_ + mage + mage2 + fbaby_ + medu'\n",
    "res_probit = Probit.from_formula(formula, dta_cat).fit()\n",
    "\n",
    "# outcome model\n",
    "formula_outcome = 'bweight ~ prenatal1_ + mmarried_ + mage + fbaby_'\n",
    "mod = OLS.from_formula(formula_outcome, dta_cat)\n",
    "\n",
    "# treatment indicator variable\n",
    "tind = np.asarray(dta_cat['mbsmoke_'])\n",
    "\n",
    "teff = TreatmentEffect(mod, tind, results_select=res_probit)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecbd49a4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "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.\n",
    "\n",
    "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.\n",
    "\n",
    "\n",
    "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.\n",
    "Method `ipw` uses the selection model but not the outcome model.\n",
    "Method `ra` uses the outcome model but not the selection model.\n",
    "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.\n",
    "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.\n",
    "\n",
    "The internal gmm results are attached to the treatment results as `results_gmm`.\n",
    "\n",
    "By default the treatment effect methods computes average treatment effect, where average is take over the sample observations.\n",
    "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`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "former-mirror",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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.6891     25.817     -8.936      0.000    -281.289    -180.089\n",
       "POM0        3403.4632      9.571    355.586      0.000    3384.704    3422.223\n",
       "POM1        3172.7741     24.001    132.193      0.000    3125.733    3219.815\n",
       "=============================================================================="
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = teff.ipw()\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "metropolitan-burns",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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.689070</td>\n",
       "      <td>25.816758</td>\n",
       "      <td>-8.935633</td>\n",
       "      <td>4.048542e-19</td>\n",
       "      <td>-281.288985</td>\n",
       "      <td>-180.089154</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>POM0</th>\n",
       "      <td>3403.463163</td>\n",
       "      <td>9.571412</td>\n",
       "      <td>355.586324</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>3384.703540</td>\n",
       "      <td>3422.222785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>POM1</th>\n",
       "      <td>3172.774093</td>\n",
       "      <td>24.001059</td>\n",
       "      <td>132.193085</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>3125.732881</td>\n",
       "      <td>3219.815305</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
       "ATE   -230.689070  25.816758   -8.935633  4.048542e-19     -281.288985      -180.089154\n",
       "POM0  3403.463163   9.571412  355.586324  0.000000e+00     3384.703540      3422.222785\n",
       "POM1  3172.774093  24.001059  132.193085  0.000000e+00     3125.732881      3219.815305"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res.summary_frame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "documentary-invasion",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                               _IPWGMM Results                                \n",
      "==============================================================================\n",
      "Dep. Variable:                      y   Hansen J:                    3.988e-09\n",
      "Model:                        _IPWGMM   Prob (Hansen J):                   nan\n",
      "Method:                           GMM                                         \n",
      "Date:                Fri, 20 Mar 2026                                         \n",
      "Time:                        11:18:46                                         \n",
      "No. Observations:                4642                                         \n",
      "==============================================================================\n",
      "                 coef    std err          z      P>|z|      [0.025      0.975]\n",
      "------------------------------------------------------------------------------\n",
      "p 0         -230.6891     25.817     -8.936      0.000    -281.289    -180.089\n",
      "p 1         3403.4632      9.571    355.586      0.000    3384.704    3422.223\n",
      "p 2           -1.5583      0.461     -3.380      0.001      -2.462      -0.655\n",
      "p 3           -0.6485      0.055    -11.711      0.000      -0.757      -0.540\n",
      "p 4            0.1744      0.036      4.836      0.000       0.104       0.245\n",
      "p 5           -0.0033      0.001     -4.921      0.000      -0.005      -0.002\n",
      "p 6           -0.2176      0.050     -4.390      0.000      -0.315      -0.120\n",
      "p 7           -0.0864      0.010     -8.630      0.000      -0.106      -0.067\n",
      "==============================================================================\n"
     ]
    }
   ],
   "source": [
    "print(res.results_gmm.summary())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9410bb5",
   "metadata": {},
   "source": [
    "**average treatment effect on the treated**\n",
    "\n",
    "see more below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d00a2846",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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         -225.1796     23.658     -9.518      0.000    -271.549    -178.811\n",
       "POM0        3362.8393     14.198    236.855      0.000    3335.012    3390.667\n",
       "POM1        3137.6597     19.071    164.526      0.000    3100.281    3175.038\n",
       "=============================================================================="
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "teff.ipw(effect_group=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "464fbfec",
   "metadata": {},
   "source": [
    "**average treatment effect on the untreated**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "df5a0f1d",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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         -231.8782     27.699     -8.371      0.000    -286.168    -177.588\n",
       "POM0        3412.9116      9.283    367.634      0.000    3394.716    3431.107\n",
       "POM1        3181.0334     26.120    121.786      0.000    3129.840    3232.227\n",
       "=============================================================================="
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "teff.ipw(effect_group=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3071337",
   "metadata": {},
   "source": [
    "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`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0795f217",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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         -239.6392     23.824    -10.059      0.000    -286.333    -192.945\n",
       "POM0        3403.2423      9.525    357.288      0.000    3384.573    3421.911\n",
       "POM1        3163.6031     21.864    144.698      0.000    3120.751    3206.455\n",
       "=============================================================================="
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_ra = teff.ra()\n",
    "res_ra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "vocational-crawford",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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>-239.639211</td>\n",
       "      <td>23.824021</td>\n",
       "      <td>-10.058722</td>\n",
       "      <td>8.408247e-24</td>\n",
       "      <td>-286.333435</td>\n",
       "      <td>-192.944988</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>POM0</th>\n",
       "      <td>3403.242272</td>\n",
       "      <td>9.525207</td>\n",
       "      <td>357.288006</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>3384.573209</td>\n",
       "      <td>3421.911335</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>POM1</th>\n",
       "      <td>3163.603060</td>\n",
       "      <td>21.863509</td>\n",
       "      <td>144.697867</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>3120.751371</td>\n",
       "      <td>3206.454750</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             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"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_ra.summary_frame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "posted-basin",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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>-223.545262</td>\n",
       "      <td>23.794008</td>\n",
       "      <td>-9.395023</td>\n",
       "      <td>5.720507e-21</td>\n",
       "      <td>-270.180660</td>\n",
       "      <td>-176.909864</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>POM0</th>\n",
       "      <td>3361.204984</td>\n",
       "      <td>14.465009</td>\n",
       "      <td>232.367989</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>3332.854088</td>\n",
       "      <td>3389.555880</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>POM1</th>\n",
       "      <td>3137.659722</td>\n",
       "      <td>19.070923</td>\n",
       "      <td>164.525844</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>3100.281400</td>\n",
       "      <td>3175.038045</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
       "ATE   -223.545262  23.794008   -9.395023  5.720507e-21     -270.180660      -176.909864\n",
       "POM0  3361.204984  14.465009  232.367989  0.000000e+00     3332.854088      3389.555880\n",
       "POM1  3137.659722  19.070923  164.525844  0.000000e+00     3100.281400      3175.038045"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ra2 = teff.ipw_ra(effect_group=1, return_results=True)\n",
    "ra2.summary_frame()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "banned-rogers",
   "metadata": {},
   "source": [
    "## All methods in TreatmentEffect\n",
    "\n",
    "The following computes and prints ATE and POM for all methods.\n",
    "(We include the call to TreatmentEffect as a reminder.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "innovative-garage",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ipw\n",
      "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
      "ATE   -230.689070  25.816758   -8.935633  4.048542e-19     -281.288985      -180.089154\n",
      "POM0  3403.463163   9.571412  355.586324  0.000000e+00     3384.703540      3422.222785\n",
      "POM1  3172.774093  24.001059  132.193085  0.000000e+00     3125.732881      3219.815305\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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",
    "The results in statsmodels are very close to the results in Stata because both packages use the same approach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "compact-pixel",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "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_st:\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": "markdown",
   "id": "physical-newman",
   "metadata": {},
   "source": [
    "### Treatment effects without inference\n",
    "\n",
    "It is possible to compute POM and ATE without computing standard errors and inferential statistics. In this case the GMM model is not computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "preceding-dividend",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ra\n",
      "(np.float64(-239.63921146434177), np.float64(3403.2422719354927), np.float64(3163.603060471151))\n",
      "\n",
      " ipw\n",
      "(np.float64(-230.6886377952617), np.float64(3403.4627086845567), np.float64(3172.774070889295))\n",
      "\n",
      " aipw\n",
      "(np.float64(-230.98920111257803), np.float64(3403.355253173835), np.float64(3172.366052061257))\n",
      "\n",
      " aipw_wls\n",
      "(np.float64(-227.19561818674947), np.float64(3403.250650975787), np.float64(3176.0550327890373))\n",
      "\n",
      " ipw_ra\n",
      "(np.float64(-229.96707793511632), np.float64(3403.3356393074127), np.float64(3173.3685613722964))\n"
     ]
    }
   ],
   "source": [
    "for m in methods:\n",
    "    print(\"\\n\", m)\n",
    "    res = getattr(teff, m)(return_results=False)\n",
    "    print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "personal-marble",
   "metadata": {},
   "source": [
    "## Treatment effect on the treated\n",
    "\n",
    "Treatment effects on subgroups are not available for `aipw` and `aipw-wls`.\n",
    "\n",
    "`effect_group` choses the group for which treatement effect and potential outcomes are computed\n",
    "Options are \n",
    "\"all\" for sample average treatment effect, \n",
    "`1` for average treatment effect on the treated and \n",
    "`0` for average treatment effect on the untreated.\n",
    "\n",
    "Note: The row labels in the pandas dataframe, POM and ATE, are the same even for treatment effect on subgroups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "lyric-producer",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ra\n",
      "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
      "ATE   -223.301651  22.742195   -9.818826  9.342545e-23     -267.875534      -178.727767\n",
      "POM0  3360.961373  12.757489  263.450069  0.000000e+00     3335.957154      3385.965592\n",
      "POM1  3137.659722  19.070923  164.525844  0.000000e+00     3100.281400      3175.038045\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ipw\n",
      "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
      "ATE   -225.179608  23.658119   -9.518069  1.764270e-21     -271.548669      -178.810546\n",
      "POM0  3362.839334  14.197866  236.855264  0.000000e+00     3335.012028      3390.666640\n",
      "POM1  3137.659726  19.070923  164.525844  0.000000e+00     3100.281404      3175.038049\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ipw_ra\n",
      "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
      "ATE   -223.545262  23.794008   -9.395023  5.720507e-21     -270.180660      -176.909864\n",
      "POM0  3361.204984  14.465009  232.367989  0.000000e+00     3332.854088      3389.555880\n",
      "POM1  3137.659722  19.070923  164.525844  0.000000e+00     3100.281400      3175.038045\n"
     ]
    }
   ],
   "source": [
    "for m in methods:\n",
    "    if m.startswith(\"aipw\"):\n",
    "        continue\n",
    "    res = getattr(teff, m)(effect_group=1)\n",
    "    print(\"\\n\", m)\n",
    "    print(res.summary_frame())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "speaking-functionality",
   "metadata": {},
   "source": [
    "### Treatment effect on the untreated\n",
    "\n",
    "Similar to ATT, we can compute average treatment effect on the untreated by using option `effect_group=0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "southern-startup",
   "metadata": {
    "execution": {
 
 
 
 
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ra\n",
      "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
      "ATE   -243.375488  24.902030   -9.773319  1.465697e-22     -292.182569      -194.568406\n",
      "POM0  3412.911593   9.283454  367.633804  0.000000e+00     3394.716358      3431.106829\n",
      "POM1  3169.536106  23.128805  137.038471  0.000000e+00     3124.204480      3214.867731\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ipw\n",
      "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
      "ATE   -231.878176  27.699436   -8.371224  5.702294e-17     -286.168073      -177.588279\n",
      "POM0  3412.911593   9.283454  367.633804  0.000000e+00     3394.716357      3431.106829\n",
      "POM1  3181.033418  26.119760  121.786472  0.000000e+00     3129.839629      3232.227206\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ipw_ra\n",
      "             coef    std err           z         P>|z|  Conf. Int. Low  Conf. Int. Upp.\n",
      "ATE   -231.125972  28.813022   -8.021580  1.043933e-15     -287.598458      -174.653487\n",
      "POM0  3412.911593   9.283454  367.633804  0.000000e+00     3394.716358      3431.106829\n",
      "POM1  3181.785621  27.301318  116.543297  0.000000e+00     3128.276021      3235.295221\n"
     ]
    }
   ],
   "source": [
    "for m in methods:\n",
    "    if m.startswith(\"aipw\"):\n",
    "        # not available\n",
    "        continue\n",
    "    res = getattr(teff, m)(effect_group=0)\n",
    "    print(\"\\n\", m)\n",
    "    print(res.summary_frame())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a88cd1ff",
   "metadata": {},
   "source": [
    "The docstring for the TreatmentEffect class and it's methods can be obtained using help\n",
    "\n",
    "`help(teff)`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.14.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
