12-Doubly-Robust-Estimation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Don't Put All your Eggs in One Basket\n",
    "\n",
    "We've learned how to use linear regression and propensity score weighting to estimate \\\\(E[Y|Y=1] - E[Y|Y=0] | X\\\\). But which one should we use and when? When in doubt, just use both! Doubly Robust Estimation is a way of combining propensity score and linear regression in a way you don't have to rely on any single one of them. \n",
    "\n",
    "To see how this works, let's consider the mindset experiment. It is a randomised study conducted in U.S. public high schools which aims at finding the impact of a growth mindset. The way it works is that students receive from the school a seminary to instil in them a growth mindset. Then, they follow up the students on college years to measure how well they performed academically. This measurement was compiled into an achievement score and standardised. The real data on this study is not publicly available in order to preserve students' privacy. However, we have a simulated dataset with the same statistical properties provided by [Athey and Wager](https://arxiv.org/pdf/1902.07409.pdf), so we will use that instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from matplotlib import style\n",
    "from matplotlib import pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "style.use(\"fivethirtyeight\")\n",
    "pd.set_option(\"display.max_columns\", 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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>schoolid</th>\n",
       "      <th>intervention</th>\n",
       "      <th>achievement_score</th>\n",
       "      <th>...</th>\n",
       "      <th>school_ethnic_minority</th>\n",
       "      <th>school_poverty</th>\n",
       "      <th>school_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>259</th>\n",
       "      <td>73</td>\n",
       "      <td>1</td>\n",
       "      <td>1.480828</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.515202</td>\n",
       "      <td>-0.169849</td>\n",
       "      <td>0.173954</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3435</th>\n",
       "      <td>76</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.987277</td>\n",
       "      <td>...</td>\n",
       "      <td>-1.310927</td>\n",
       "      <td>0.224077</td>\n",
       "      <td>-0.426757</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9963</th>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.152340</td>\n",
       "      <td>...</td>\n",
       "      <td>0.875012</td>\n",
       "      <td>-0.724801</td>\n",
       "      <td>0.761781</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4488</th>\n",
       "      <td>67</td>\n",
       "      <td>0</td>\n",
       "      <td>0.358336</td>\n",
       "      <td>...</td>\n",
       "      <td>0.315755</td>\n",
       "      <td>0.054586</td>\n",
       "      <td>1.862187</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2637</th>\n",
       "      <td>16</td>\n",
       "      <td>1</td>\n",
       "      <td>1.360920</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.033161</td>\n",
       "      <td>-0.982274</td>\n",
       "      <td>1.591641</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 13 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      schoolid  intervention  achievement_score  ...  school_ethnic_minority  \\\n",
       "259         73             1           1.480828  ...               -0.515202   \n",
       "3435        76             0          -0.987277  ...               -1.310927   \n",
       "9963         4             0          -0.152340  ...                0.875012   \n",
       "4488        67             0           0.358336  ...                0.315755   \n",
       "2637        16             1           1.360920  ...               -0.033161   \n",
       "\n",
       "      school_poverty  school_size  \n",
       "259        -0.169849     0.173954  \n",
       "3435        0.224077    -0.426757  \n",
       "9963       -0.724801     0.761781  \n",
       "4488        0.054586     1.862187  \n",
       "2637       -0.982274     1.591641  \n",
       "\n",
       "[5 rows x 13 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_csv(\"./data/learning_mindset.csv\")\n",
    "data.sample(5, random_state=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although the study was randomised, it doesn't seem to be the case that this data is free from confounding. One possible reason for this is that the treatment variable is measured by the student's receipt of the seminar. So, although the opportunity to participate was random, participation is not. We are here dealing with a case of non-compliance here. One evidence of this is how the student's success expectation is correlated with the participation in the seminar. Students with higher self-reported participation are more likely to have joined the growth mindset seminar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "success_expect\n",
       "1    0.271739\n",
       "2    0.265957\n",
       "3    0.294118\n",
       "4    0.271617\n",
       "5    0.311070\n",
       "6    0.354287\n",
       "7    0.362319\n",
       "Name: intervention, dtype: float64"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.groupby(\"success_expect\")[\"intervention\"].mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we will use linear methods (linear regression and logistic regression), we need to convert the categorical variables to dummies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10391, 32)\n"
     ]
    }
   ],
   "source": [
    "categ = [\"ethnicity\", \"gender\", \"school_urbanicity\"]\n",
    "cont = [\"school_mindset\", \"school_achievement\", \"school_ethnic_minority\", \"school_poverty\", \"school_size\"]\n",
    "\n",
    "data_with_categ = pd.concat([\n",
    "    data.drop(columns=categ), # dataset without the categorical features\n",
    "    pd.get_dummies(data[categ], columns=categ, drop_first=False)# dataset without categorical converted to dummies\n",
    "], axis=1)\n",
    "\n",
    "print(data_with_categ.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to understand how doubly robust estimation works.\n",
    "\n",
    "## Doubly Robust Estimation\n",
    "\n",
    "![img](./data/img/doubly-robust/double.png)\n",
    "\n",
    "Instead of deriving the estimator, I'll first show it to you and only then tell why it is awesome.\n",
    "\n",
    "$\n",
    "\\hat{ATE} = \\frac{1}{N}\\sum \\bigg( \\dfrac{T_i(Y_i - \\hat{\\mu_1}(X_i))}{\\hat{P}(X_i)} + \\hat{\\mu_1}(X_i) \\bigg) - \\frac{1}{N}\\sum \\bigg( \\dfrac{(1-T_i)(Y_i - \\hat{\\mu_0}(X_i))}{1-\\hat{P}(X_i)} + \\hat{\\mu_0}(X_i) \\bigg)\n",
    "$\n",
    "\n",
    "where \\\\(\\hat{P}(x)\\\\) is an estimation of the propensity score (using logistic regression, for example), \\\\(\\hat{\\mu_1}(x)\\\\) is an estimation of \\\\(E[Y|X, T=1]\\\\) (using linear regression, for example), and \\\\(\\hat{\\mu_1}(x)\\\\) is an estimation of \\\\(E[Y|X, T=0]\\\\). As you might have already guessed, the first part of the doubly robust estimator estimates \\\\(E[Y_1]\\\\) and the second part estimates \\\\(E[Y_0]\\\\). Let's examine the first part, as all the intuition will also apply to the second part by analogy.\n",
    "\n",
    "But first, let's see it in practice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression, LinearRegression\n",
    "\n",
    "def doubly_robust(df, X, T, Y):\n",
    "    ps = LogisticRegression(C=1e6).fit(df[X], df[T]).predict_proba(df[X])[:, 1]\n",
    "    mu0 = LinearRegression().fit(df.query(f\"{T}==0\")[X], df.query(f\"{T}==0\")[Y]).predict(df[X])\n",
    "    mu1 = LinearRegression().fit(df.query(f\"{T}==1\")[X], df.query(f\"{T}==1\")[Y]).predict(df[X])\n",
    "    return (\n",
    "        np.mean(df[T]*(df[Y] - mu1)/ps + mu1) -\n",
    "        np.mean((1-df[T])*(df[Y] - mu0)/(1-ps) + mu0)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3882218480614833"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "T = 'intervention'\n",
    "Y = 'achievement_score'\n",
    "X = data_with_categ.columns.drop(['schoolid', T, Y])\n",
    "\n",
    "doubly_robust(data_with_categ, X, T, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Doubly robust estimator is saying that we should expect individuals who attended the mindset seminar to be 0.388 standard deviations above their untreated fellows, in terms of achievements. Once again, we can use bootstrap to construct confidence intervals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import Parallel, delayed # for parallel processing\n",
    "\n",
    "np.random.seed(88)\n",
    "# run 1000 bootstrap samples\n",
    "bootstrap_sample = 1000\n",
    "ates = Parallel(n_jobs=4)(delayed(doubly_robust)(data_with_categ.sample(frac=1, replace=True), X, T, Y)\n",
    "                          for _ in range(bootstrap_sample))\n",
    "ates = np.array(ates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ATE 95% CI: (0.3530692065036196, 0.42163132128223224)\n"
     ]
    }
   ],
   "source": [
    "print(f\"ATE 95% CI:\", (ates.mean() - 1.96*ates.std(), ates.mean() + 1.96*ates.std()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZwAAAEGCAYAAABRvCMcAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAc+ElEQVR4nO3de5wcZZ3v8c93JjcTIEHdXCRwABljAGHlcLhm3Qgu4OUAAioqhGBW12NEEFwCyEpcRAFR4CzI0UMkwQuCCAtGUELWyAsE9ihGBCJ2DNkkkGTQkME4gRDyO39UTejpzLUvT8/0fN+vV72666mnqn5P93T95ql6uloRgZmZWa011TsAMzMbGpxwzMwsCSccMzNLwgnHzMyScMIxM7MknHDMzCwJJxyzQU7SEkk31mjbcyUt726+BvubL+n+Wm3f6ssJZ4iSNFHSS5LWSRqel02XFL1M8/O63S2/rod9ziypu1nS7yWdV4P2Tc73Mb2f6y2XNLfa8fSXpJVFr9PLktZKuk/SP0oaVlL9JODcPm63v6/LVcBh/Qi9TySdJqmrLwGeDXyg2vuzgaH0D9eGjo8BPwHeCpwA3A78EphUVOcc4KPA/ygq21z0/NPAj0q2+9de9vsqMDl//jrgGOA6Sa0R8Z3+NKCeJI2IiC013s0VwDVkn9OJwLvIEsAZko6NiHaAiNhQ7R1LagIUEZuATdXefncioi3VvqwOIsLTEJvIerbPAMcD5wP3dVPvYmBlN8sCOK2f+50JbO2i/NfAtUXzAj4HrAC2AH8EzilZZ2fgm8DzwEvAr4BjSuIrnlbm5ZPJkuSfyJLnCuCf82VLulhvT2B6/vy9wIP5/j4N7Ap8F1iVb+tp4DyyA3VHHPOB+8l6IM8C7fn+39jLa7USuLiL8r8FXgG+UFS2BLixaH4a8BDwl3z6LXBsL6/LXGA58CHg98BWYP+O8qJtd9T7SP7avZS3b6/SOiVxT+vi9Sye5he/Xv38W1gJ/CtwLbABWE+WmJvr/Vnz1HnyKbWh6RhgDHAP8B1guqS9UwehzHRgKlnvqsOngEuBy4H9gK8Cl0uaVVTn28CxwGnA28kOsAslvTVfflD+eDJZr62jl/YNYCxZb2EqMAtYky87iezg9bV8nUnA6qJ9fg24Ml/v34GRwO+AE4F985i/SJZYix1CdpA9DngPcEAef79FxFLgZ8AHu1ouqRm4G3iU7DU4iCwBtOdVuntdAN5E9trPzNvzX92EMSmv9yHg78iS/79LUh+b8UuyhN2xrUlkp9K60pe/BYCzgLXAocBnyHrnM/oYj6VS74znKf0E3AlcXTR/D/DlLur11sN5iex0S/H0oR72OzNfr6PuK/n8VSX1VgNXlpRdDazIn++Tr/eekjqPAd/On0/O60wvqfNbYG4PMS4vXc5r/5Gf3ofX9lpgUdH8/LytY4vKjsm319LDdlbSRQ8nX3Y50F40v4S8h0PW69qh3UV1u3td5gLbgD26KC/t4QSwT1HZW/Kyd3W1Tl62vYeTz5+WHX52iG8+nXs4Pf4tFL1Wd5fU+SlwS8rPlafeJ/dwhhhJk4D3AQuKiucDZ3ZxMbo3nyc7xVM8/aSXdV4tqnsg2amZGR0X6iXtQnZQfKBkvV8Ae0oaTfbfN13UeYDsv+CeXANcJOlRSVdIekcv9Yv9Z/GMpCZJF0haKulPkjYBnwT+W8l6T0XnaxMP5Y9T+7HvTrsmO3jvICJeAG4Efibp3jy+KX3c7vqIWNWHes9HxPaRahHxB7JTlPt2v0r/9fFvocPSkjrPAhOqGY9Vzgln6JlFdhH6V5K2StoKfJ/sovTx/dzW+ohYXjL1eoG5qO5TEXEL2X+sF0kaWVytZLW+nK7p9kBctO+byBLC/yE7lXOvpO/2Yduw44CI84ALgX8D/oEsid4IjOjj9sq1P9m1jC5FxMeB/w4sAv4eeELSP/Vhu70N+OhJ8fuzjR3fr+EVbLsvfwulAzgCH98GHL8hQ0g+8ugfgS+zY8/ku8An6hTaVrIkODIiXiS7pvL3JXXeATwT2cisJ4vKiv1d0bKOA1Bz6c4iYm1E3BQRM8gS8Efz/6Y71tthnW68A/hpRMyLiN/k//W3dFFvatH2AY7IH5f1cT/bSfpbsmtXt/ZULyKeiIivR8S7gXm89t52+7r0w99IenNRTG8B3sBr7WkFxufXkzocRGdb8nW7jaOPfws2iHhY9NByHLAH8M3SUyeSbgIWSdozIlb2cXtjJU0sKXspIjb2tFLROiPITqudAyzODzAAXwG+JqlAdn3iKOB/AbMBIuKPkn4IfCP/z/2/8uX7k52ig+wUzybgGElPAi9HxAv594TuIRtRNopsoMBqstFckI3eO1LSHmQX2nsacvw0cLqkd5KdwplBdtH6hZJ6Adws6WLg9cD1wE8iotDT6wTslL9WxcOi55CNlPt6VytI2gf4OPDjvF1vIkvEj/X0uvQSR6l24CZJnyXrbfwb2eCJji9s/hwYDVwqaR5Zspldso1n8sfjJT0IbO6md9zj34INMvW+iOQp3QTcBTzczbJmYB3wpaKy3gYNdDUt7GH/M0vqvkJ24LkOeENRPQH/nC97hWxIbOlQ2F14bVj0y5QMi87rzCjaxsq87HrgD2TDmP9Mds1pv6J1DiYbpr2ZHYfxTi7Z/ljgNuDFfFvXk42oWllUZz7ZgfhzZKOoNpMN2vibXt6rlUWv05b8vbmPrIfaXFJ3Ca8NGpgE3EHWM3gZeA74v3QetNDV6zKXkgv9XZXz2rDo0/IYXwb+A3hzyXofy9+3zcC9wKkUDRrI61xDNoR5Gz0Pi+7tb2ElJQMsyE5tLqn3Z85T50n5m2NmNZDfmWFyRLyr3rGY1Zuv4ZiZWRJOOGZmloRPqZmZWRLu4ZiZWRJ1GRbd1tbmbpWZWYMbO3Zspy/puodjZmZJOOGYmVkSTjgJFAq9faG8cbitjcltbUyp2+qEY2ZmSTjhmJlZEr55p5lZLyKCTZs2sW3btnqHUlWjRo2ira2t94rdaGpqYqeddqKvP/bqhGNm1otNmzYxcuRIRoyo9U8dpTVy5EhGjRpV9vpbtmxh06ZN7Lzzzn2q71NqZma92LZtW8Mlm2oYMWJEv3p9TjhmZpaEE46Z2SBwww03cPjhh3PYYYfxjW98Y3v5V77yFaZOncq0adOYNm0a9913HwCPPPIIRxxxBO985ztZsWIFABs3buSkk06iu3tovvLKK8ydO5eDDjqIww8/nKOOOopFixYB8La3vY0///nPFbXB13DM6mj+038ta72ZU8ZUORIbyJ566iluvvlmFi9ezIgRIzj55JM59thjefObs1/6/tSnPsVZZ53VaZ3rrruO73znO6xatYp58+Zx2WWX8dWvfpVzzz2324v8l112GevWrePhhx9m5MiRtLa28tBDD1WtHe7hmJn107hx4zpN3Zk/f36nemeffXZZ+/vDH/7AwQcfzOjRoxk2bBhHHnkkCxcu7HGd4cOHs3nzZtrb2xk+fDjPPPMMa9euZdq0aV3Wb29vZ8GCBVx55ZWMHDkSgPHjx/P+97+/rJi74oRjZjbATZ06lV/+8pds2LCB9vZ2Fi1axJo1a7Yv/9a3vsURRxzB7Nmz2bhxIwCf/exnOeecc7jhhhv4+Mc/zqWXXsrnP//5bvexYsUKJk+ezC677FKzdjjhmJkNcFOmTOHss8/mxBNP5OSTT2b//fdn2LDsisisWbNYunQpDz74IBMnTtyeVA444ADuv/9+Fi5cyMqVK5k4cSIRwZlnnsknPvEJWltbk7fDCcfMbBCYMWMGDzzwAPfeey+77rrr9us348ePp7m5maamJmbMmMFjjz3Wab2I4KqrruL888/niiuu4MILL+SDH/wg3/zmNzvV23vvvVmzZg1/+ctfatYGJxwzs37auHFjp6k7M2fO7FTv2muvLXufzz//PACrV6/mxz/+MaeccgoA69at215n4cKFTJ06tdN63//+9znmmGMYN24c7e3tNDU10dTURHt7e6d6o0eP5vTTT2fOnDls2bJl+7ZvvfXWsmMu5VFqZmaDwIwZM9iwYQPDhg3jqquu2j5Y4Qtf+AJPPPEEAHvssQfXXHPN9nXa29u55ZZbuPPOOwGYPXs2M2bMYPjw4cybN2+HfVx88cV86Utf4tBDD2XUqFGMHj2aiy66qGptUHfjsWtpqP3iZ6FQoKWlpd5hJOG29s9gGRY91N/XtrY2xo4dW6eIauell16q6NY20PNr41/8NDOzunDCMTOzJJxwzMwsCSccMzNLoteEI+nbklolPVFU9lVJv5f0uKQ7JY0rWnahpOWSnpZ0bK0CNzNLpampaftQYXvNli1baGrqe7+lL8Oi5wPXATcXlS0CLoyIrZKuAC4E5kjaFzgV2A94E3C/pLdExKt9jsjMbIDZaaed2LRpE5s3b653KFX14osvVnQrm45f/OyrXhNORDwgac+SsvuKZh8BTsmfnwD8ICJeBp6RtBw4BHi4zxGZmQ0wkvr8q5aDSWtrK7vvvnuy/VXjGs7HgHvz57sBq4uWrcnLzMxsiKvoTgOSPg9sBb7XUdRFtR6/5FkoFCoJYdAYKu0Et7U/1rc2l7ffpvRnqf2+NqZqt7WnLwiXnXAknQG8Dzg6XrtdwRqguH82GXiu3OAaxVD/lnajqkZbJ2wr704DLS2+00CtuK21U9YpNUnHAXOA4yOi+A5wdwOnShopaS+gBfjPysM0M7PBrtcejqRbgOnAGyWtAS4hG5U2EliU/1TpIxHxyYh4UtJtwFNkp9pme4SamZlB30apfbiL4h1vM/pa/cuAyyoJyszMGo/vNGBmZkk44ZiZWRJOOGZmloQTjpmZJeGEY2ZmSTjhmJlZEk44ZmaWhBOOmZkl4YRjZmZJOOGYmVkSTjhmZpaEE46ZmSXhhGNmZkk44ZiZWRJOOGZmloQTjpmZJeGEY2ZmSTjhmJlZEk44ZmaWhBOOmZkl4YRjZmZJOOGYmVkSTjhmZpZErwlH0rcltUp6oqjs9ZIWSSrkj7vm5ZL0vyUtl/S4pINqGbyZmQ0efenhzAeOKym7AFgcES3A4nwe4N1ASz59ArihOmGamdlg12vCiYgHgA0lxScAC/LnC4ATi8pvjswjwDhJk6oVrJmZDV7lXsOZEBFrAfLH8Xn5bsDqonpr8jIzMxvihlV5e+qiLHpaoVAoVDmEgWmotBPc1v5Y39pc3n6bXq1ov2Xt0+9rQ6p2W1taWrpdVm7CWS9pUkSszU+Ztebla4Ddi+pNBp4rN7hGUSgUhkQ7wW3trwnb/lrWei0tYyrab3/5fW1MqdtabsK5GzgDuDx/vKuo/NOSfgAcCrR1nHozs+qZ/3R5iWrmlLSJyqxYrwlH0i3AdOCNktYAl5AlmtskzQJWAR/Iq98DvAdYDrQDZ9YgZjMzG4R6TTgR8eFuFh3dRd0AZlcalJmZNR7facDMzJJwwjEzsySccMzMLAknHDMzS8IJx8zMknDCMTOzJJxwzMwsCSccMzNLoto37zSzAaycW+L4djhWLe7hmJlZEk44ZmaWhBOOmZkl4YRjZmZJOOGYmVkSHqVmVqKvI7nWtzZ3+sVOj+Yy65l7OGZmloQTjpmZJeGEY2ZmSTjhmJlZEk44ZmaWhBOOmZkl4YRjZmZJ+Hs4ZlVSzp2YzYaSino4kj4r6UlJT0i6RdIoSXtJelRSQdKtkkZUK1gzMxu8yk44knYDPgMcHBH7A83AqcAVwNUR0QK8AMyqRqBmZja4VXoNZxjwOknDgNHAWuAo4PZ8+QLgxAr3YWZmDaDshBMRzwJXAavIEk0b8GtgY0RszautAXarNEgzMxv8FBHlrSjtCvwI+BCwEfhhPn9JROyT19kduCci3la8bltb2/adFgqF8iI3q5E71jXXO4QB5aSJr9Y7BBtEWlpatj8fO3asipdVMkrtXcAzEfE8gKQ7gCOAcZKG5b2cycBzfQ2uURUKhSHRTmiMthbfAbon61vXM2H8hBpHU38tLWMa4n3tK7e1diq5hrMKOEzSaEkCjgaeAn4OnJLXOQO4q7IQzcysEVRyDedRssEBjwG/y7f1LWAOcK6k5cAbgHlViNPMzAa5ir74GRGXAJeUFK8ADqlku2Zm1nh8pwEz69H8p/+6w6+b9oV/AdVK+V5qZmaWhBOOmZkl4YRjZmZJOOGYmVkSTjhmZpaEE46ZmSXhhGNmZkk44ZiZWRL+4qc1LP/ks9nA4h6OmZkl4YRjZmZJOOGYmVkSTjhmZpaEE46ZmSXhhGNmZkl4WLQNCh7ibDb4uYdjZmZJOOGYmVkSTjhmZpaEE46ZmSXhhGNmZkk44ZiZWRIVJRxJ4yTdLun3kpZJOlzS6yUtklTIH3etVrBmZjZ4VdrDuRb4aUS8FTgQWAZcACyOiBZgcT5vZmZDXNkJR9IuwDuAeQARsSUiNgInAAvyaguAEysN0szMBr9Kejh7A88DN0n6jaQbJY0BJkTEWoD8cXwV4jQzs0FOEVHeitLBwCPAkRHxqKRrgReBsyJiXFG9FyKi03Wctra27TstFApl7d+GljvWNdc7BOunkya+Wu8QrA5aWlq2Px87dqyKl1VyL7U1wJqIeDSfv53ses16SZMiYq2kSUBrX4NrVIVCYUi0E2rX1gnbBt691Na3rmfC+An1DiOJctra0jKmRtHUlj+vtVP2KbWIWAesljQlLzoaeAq4GzgjLzsDuKuiCM3MrCFUerfos4DvSRoBrADOJEtit0maBawCPlDhPsxsECr3Dt8zpwzOnpH1rqKEExFLgYO7WHR0Jds1M7PG4zsNmJlZEk44ZmaWhBOOmZkl4YRjZmZJOOGYmVkSTjhmZpaEE46ZmSXhhGNmZkk44ZiZWRJOOGZmloQTjpmZJeGEY2ZmSVR6t2izfin3DsJmNvi5h2NmZkk44ZiZWRJOOGZmloQTjpmZJeGEY2ZmSTjhmJlZEk44ZmaWhBOOmZkl4YRjZmZJOOGYmVkSFSccSc2SfiNpYT6/l6RHJRUk3SppROVhmpnZYFeNHs7ZwLKi+SuAqyOiBXgBmFWFfZiZ2SBXUcKRNBl4L3BjPi/gKOD2vMoC4MRK9mFmZo2h0h7ONcD5wLZ8/g3AxojYms+vAXarcB9mZtYAyv55AknvA1oj4teSpncUd1E1etpOoVAoN4RBpdHaece65m6WNMO6FUljqaf1revrHUIyqdpaaHo1yX56jKHBPq89qXZbW1paul1Wye/hHAkcL+k9wChgF7IezzhJw/JezmTguXKDaxSFQqHh2jlhW9e/a7O+dT0Txk9IHE19uK210dIyJsl+utOIn9fupG5r2afUIuLCiJgcEXsCpwL/EREfBX4OnJJXOwO4q+Iozcxs0KvF93DmAOdKWk52TWdeDfZhZmaDTFV+YjoilgBL8ucrgEOqsV0zM2scvtOAmZkl4YRjZmZJOOGYmVkSTjhmZpaEE46ZmSXhhGNmZkk44ZiZWRJOOGZmloQTjpmZJVGVOw2YmVXL/Ke7vjFsT2ZOqe8NP61v3MMxM7Mk3MMxs0GvnF4RuGeUmns4ZmaWhBOOmZkl4YRjZmZJOOGYmVkSTjhmZpaEE46ZmSXhhGNmZkk44ZiZWRJOOGZmloQTjpmZJeFb25jZkNXVLXHWtzYzYVvPt8rxLXHKU3YPR9Lukn4uaZmkJyWdnZe/XtIiSYX8cdfqhWtmZoNVJafUtgLnRcRU4DBgtqR9gQuAxRHRAizO583MbIgr+5RaRKwF1ubP/yJpGbAbcAIwPa+2AFgCzKkoSqupcu+0a2bWH4qIyjci7Qk8AOwPrIqIcUXLXoiITqfV2tratu+0UChUvH/rn2XLlnWe33X/OkViNrCtXf5Up/lJ++wLwEkTX61HOINCS0vL9udjx45V8bKKBw1I2gn4EXBORLwoqbdVug2uURUKhQHVztJYNlSxh7O+dT0Txk+o2vYGMre1MRW3tbs2t7Q0xqCB1MemioZFSxpOlmy+FxF35MXrJU3Kl08CWisL0czMGkElo9QEzAOWRcTXixbdDZyRPz8DuKv88MzMrFFUckrtSOB04HeSluZlFwGXA7dJmgWsAj5QWYhmZtYIKhml9iDQ3QWbo8vdrpmZNSbfaWAIWrp0aeeC1w2cAQ1mA8nqZY93mt996gF1iqQxOOEMQdOnT+80f82jz9YnELMB7msz391p3p+VyvjmnWZmloQTjpmZJeFTakPQgQceWO8QzAaFyVPeVu8QGooTzhD0i1/8otO876Vm1rXP3fzTeofQUHxKzczMknAPp4G4p2JmA5l7OGZmloQTjpmZJeGEY2ZmSfgazhB065fP7zT/oYuurFMkZgObPyvV5YQzBD181/c6zftDZNY1f1aqy6fUzMwsCSccMzNLwqfUBqBaf5/mgxdcUdPtmzUKf1aqywlnCDri/afVOwSzQcGflepywjEz66dyzkLMnDKmBpEMLr6GY2ZmSTjhmJlZEj6lZmaWQOqb6w7EU3ju4ZiZWRLu4QxB5xy6W6f5ax59tk6RmA1s/qxUV80SjqTjgGuBZuDGiLi82vtIOVKkku7w+tZmJmzzb9WYWTp9OWaVHptqfRquJqfUJDUD1wPvBvYFPixp31rsy8zMBgdFRPU3Kh0OzI2IY/P5CwEi4isAbW1t1d+pmZkNKGPHjlXxfK0GDewGrC6aX5OXmZnZEFWrhKMuytyrMTMbwmo1aGANsHvR/GTguY6Z0m6WmZk1vlr1cP4f0CJpL0kjgFOBu2u0LzMzGwRqknAiYivwaeBnwDLgtoh4shb7qidJx0l6WtJySRd0sfyTkn4naamkB0tH6knaQ9ImSZ9LF3V5KmmrpAMkPSzpybzOqLTR90+5bZU0XNKCfNmyjsEyA1lvbS2qd4qkkHRwUdmF+XpPSzo2TcTlK7etkv5B0q/z9/XXko5KF3V5Knlf8/LaHJsiwlMZE9n3i/4I7A2MAH4L7FtSZ5ei58cDPy1Z/iPgh8Dn6t2eWrWV7LTt48CB+fwbgOZ6t6lGbf0I8IP8+WhgJbBnvdtUSVvzejsDDwCPAAfnZfvm9UcCe+XbGdTvaw9tfTvwpvz5/sCz9W5PrdpatKwmxybf2qZ8hwDLI2JFRGwBfgCcUFwhIl4smh1D0cAJSScCK4DB0POrpK3HAI9HxG/zen+OiFcTxFyuStoawBhJw4DXAVuA4roDTa9tzV0KXAm8VFR2AllyfTkingGW59sbqMpua0T8JiI6rkE/CYySNLLWAVegkve1pscmJ5zy9Wnot6TZkv5I9sZ+Ji8bA8wBvpggzmoou63AW4CQ9DNJj0k6v+bRVqaStt4O/BVYC6wCroqIDbUNtyK9tlXS24HdI2Jhf9cdYCppa7GTgd9ExMvVD7Fqym5rrY9NTjjl69PQ74i4PiLeTPYmXpwXfxG4OiI21TC+aqqkrcOAacBH88f3Szq6VoFWQSVtPQR4FXgT2Wmm8yTtXatAq6DHtkpqAq4GzuvvugNQJW3tqLMfcAXwT1WPrroqaWtNj02+eWf5ehz63YUfADfkzw8FTpF0JTAO2CbppYi4riaRVq6Stq4BfhERfwKQdA9wELC4BnFWQyVt/QjZ9ZxXgFZJDwEHk52eGIh6a+vOZNcslkgCmAjcLen4Pqw70JTd1oj4laTJwJ3AjIj4Y6KYy1XJ+1rbY1O9L3AN1oksWa8g+0+248LcfiV1Woqe/0/gV11sZy4Df9BA2W0FdgUeI7uIPgy4H3hvvdtUo7bOAW4i+w9zDPAUcEC921RJW0vqL+G1C+n70XnQwAoG9qCBSto6Lq9/cr3bUeu2lpRX/djkHk6ZImKrpI6h383AtyPiSUn/SnYAuhv4tKR3Aa8ALwBn1C/i8lXS1oh4QdLXyb6bFcA9EfGTujSkDyp8X68nSzhPkCWdmyLi8eSN6KM+trW7dZ+UdBtZUt0KzI4BPBikkraSfcVjH+BfJP1LXnZMRLTWNuryVNjWmqrJzTvNzMxKedCAmZkl4YRjZmZJOOGYmVkSTjhmZpaEE46ZmSXhhGNmZkk44ZiZWRJOOGZmlsT/BzjHilGtykArAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "sns.distplot(ates, kde=False)\n",
    "plt.vlines(ates.mean()-1.96*ates.std(), 0, 20, linestyles=\"dotted\")\n",
    "plt.vlines(ates.mean()+1.96*ates.std(), 0, 20, linestyles=\"dotted\", label=\"95% CI\")\n",
    "plt.title(\"ATE Bootstrap Distribution\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's examine why the doubly robust estimator is so great. It is called doubly robust because it only requires one of the models, \\\\(\\hat{P}(x)\\\\) or \\\\(\\hat{\\mu}(x)\\\\), to be correctly specified. To see this, take the first part that estimates \\\\(E[Y_1]\\\\) and take a good look at it.\n",
    "\n",
    "$\n",
    "\\hat{E}[Y_1] = \\frac{1}{N}\\sum \\bigg( \\dfrac{T_i(Y_i - \\hat{\\mu_1}(X_i))}{\\hat{P}(X_i)} + \\hat{\\mu_1}(X_i) \\bigg)\n",
    "$\n",
    "\n",
    "First, assume that \\\\(\\hat{\\mu_1}(x)\\\\) is correct. If the propensity score model is wrong, we wouldn't need to worry. Because if \\\\(\\hat{\\mu_1}(x)\\\\) is correct, then \\\\(E[T_i(Y_i - \\hat{\\mu_1}(X_i))]=0\\\\). That is because the multiplication by \\\\(T_i\\\\) selects only the treated and the residual of \\\\(\\hat{\\mu_1}\\\\) on the treated have, by definition, mean zero. This causes the whole thing to reduce to \\\\(\\hat{\\mu_1}(X_i)\\\\), which is correctly estimated \\\\(E[Y_1]\\\\) by assumption. So, you see that by being correct, \\\\(\\hat{\\mu_1}(X_i)\\\\) wipes out the relevance of the propensity score model. We can apply the same reasoning to understand the estimator of \\\\(E[Y_0]\\\\). \n",
    "\n",
    "But don't take my word for it. Let the code show you the way! In the following estimator, I've replaced the logistic regression that estimates the propensity score by a random uniform variable that goes from 0.1 to 0.9 (I don't want very small weights to blow up my propensity score variance). Since this is random, there is no way it is a good propensity score model, but we will see that the doubly robust estimator still manages to produce an estimation that is very close to when the propensity score was estimated with logistic regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression, LinearRegression\n",
    "\n",
    "def doubly_robust_wrong_ps(df, X, T, Y):\n",
    "    # wrong PS model\n",
    "    np.random.seed(654)\n",
    "    ps = np.random.uniform(0.1, 0.9, df.shape[0])\n",
    "    mu0 = LinearRegression().fit(df.query(f\"{T}==0\")[X], df.query(f\"{T}==0\")[Y]).predict(df[X])\n",
    "    mu1 = LinearRegression().fit(df.query(f\"{T}==1\")[X], df.query(f\"{T}==1\")[Y]).predict(df[X])\n",
    "    return (\n",
    "        np.mean(df[T]*(df[Y] - mu1)/ps + mu1) -\n",
    "        np.mean((1-df[T])*(df[Y] - mu0)/(1-ps) + mu0)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3797180320806519"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doubly_robust_wrong_ps(data_with_categ, X, T, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we use bootstrap, we can see that the variance is slightly higher than when the propensity score was estimated with a logistic regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import Parallel, delayed # for parallel processing\n",
    "\n",
    "np.random.seed(88)\n",
    "parallel_fn = delayed(doubly_robust_wrong_ps)\n",
    "wrong_ps = Parallel(n_jobs=4)(parallel_fn(data_with_categ.sample(frac=1, replace=True), X, T, Y)\n",
    "                              for _ in range(bootstrap_sample))\n",
    "wrong_ps = np.array(wrong_ps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ATE 95% CI: (0.3398327292509242, 0.43263125686601944)\n"
     ]
    }
   ],
   "source": [
    "print(f\"ATE 95% CI:\", (wrong_ps.mean() - 1.96*wrong_ps.std(), wrong_ps.mean() + 1.96*wrong_ps.std()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This covers the case that the propensity model is wrong but the outcome model is correct. What about the other situation? Let's again take a good look at the first part of the estimator, but let's rearrange some terms\n",
    "\n",
    "$\n",
    "\\hat{E}[Y_1] = \\frac{1}{N}\\sum \\bigg( \\dfrac{T_i(Y_i - \\hat{\\mu_1}(X_i))}{\\hat{P}(X_i)} + \\hat{\\mu_1}(X_i) \\bigg)\n",
    "$\n",
    "\n",
    "$\n",
    "\\hat{E}[Y_1] = \\frac{1}{N}\\sum \\bigg( \\dfrac{T_iY_i}{\\hat{P}(X_i)} - \\dfrac{T_i\\hat{\\mu_1}(X_i)}{\\hat{P}(X_i)} + \\hat{\\mu_1}(X_i) \\bigg)\n",
    "$\n",
    "\n",
    "$\n",
    "\\hat{E}[Y_1] = \\frac{1}{N}\\sum \\bigg( \\dfrac{T_iY_i}{\\hat{P}(X_i)} - \\bigg(\\dfrac{T_i}{\\hat{P}(X_i)} - 1\\bigg) \\hat{\\mu_1}(X_i) \\bigg)\n",
    "$\n",
    "\n",
    "$\n",
    "\\hat{E}[Y_1] = \\frac{1}{N}\\sum \\bigg( \\dfrac{T_iY_i}{\\hat{P}(X_i)} - \\bigg(\\dfrac{T_i - \\hat{P}(X_i)}{\\hat{P}(X_i)}\\bigg) \\hat{\\mu_1}(X_i) \\bigg)\n",
    "$\n",
    "\n",
    "Now, assume that the propensity score \\\\(\\hat{P}(X_i)\\\\) is correctly specified. In this case, \\\\(E[T_i - \\hat{P}(X_i)]=0\\\\), which wipes out the part dependent on \\\\(\\hat{\\mu_1}(X_i)\\\\). This makes the doubly robust estimator reduce to the propensity score weighting estimator \\\\(\\frac{T_iY_i}{\\hat{P}(X_i)}\\\\), which is correct by assumption. So even if the \\\\(\\hat{\\mu_1}(X_i)\\\\) is wrong, the estimator will still be correct, provided that the propensity score is correctly specified.\n",
    "\n",
    "Once again, if you believe more in code than in formulas, here it is the practical verification. In the code below, we've replaced both regression models with a random normal variable. So, no doubt that \\\\(\\hat{\\mu}(X_i)\\\\) is not correctly specified. Still, we will see that doubly robust estimation still manages to recover the same \\\\(\\hat{ATE\\\\) of about 0.38 that we've seen before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression, LinearRegression\n",
    "\n",
    "def doubly_robust_wrong_model(df, X, T, Y):\n",
    "    np.random.seed(654)\n",
    "    ps = LogisticRegression(C=1e6).fit(df[X], df[T]).predict_proba(df[X])[:, 1]\n",
    "    \n",
    "    # wrong mu(x) model\n",
    "    mu0 = np.random.normal(0, 1, df.shape[0])\n",
    "    mu1 = np.random.normal(0, 1, df.shape[0])\n",
    "    return (\n",
    "        np.mean(df[T]*(df[Y] - mu1)/ps + mu1) -\n",
    "        np.mean((1-df[T])*(df[Y] - mu0)/(1-ps) + mu0)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3745648055762825"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doubly_robust_wrong_model(data_with_categ, X, T, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One again, we can use bootstrap and see that the variance is just slightly higher."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import Parallel, delayed # for parallel processing\n",
    "\n",
    "np.random.seed(88)\n",
    "parallel_fn = delayed(doubly_robust_wrong_model)\n",
    "wrong_mux = Parallel(n_jobs=4)(parallel_fn(data_with_categ.sample(frac=1, replace=True), X, T, Y)\n",
    "                               for _ in range(bootstrap_sample))\n",
    "wrong_mux = np.array(wrong_mux)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ATE 95% CI: (0.34243822810645, 0.4346749772942791)\n"
     ]
    }
   ],
   "source": [
    "print(f\"ATE 95% CI:\", (wrong_mux.mean() - 1.96*wrong_mux.std(), wrong_mux.mean() + 1.96*wrong_mux.std()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In practice, what ends up happening is that neither the propensity score nor the outcome model are 100% correct. They are both wrong, but in different ways. Doubly robust estimation can combine those two wrong models to make them less wrong. \n",
    "\n",
    "## Keys Ideias\n",
    "Here, we saw a simple way of combining linear regression with the propensity score to produce a doubly robust estimator. This estimator bears that name because it only requires one of the models to be correct. If the propensity score model is correct, we will be able to identify the causal effect even if the outcome model is wrong. On the flip side, if the outcome model is correct, we will also be able to identify the causal effect even if the propensity score model is wrong.\n",
    "\n",
    "## References\n",
    "\n",
    "I like to think of this entire series as a tribute to Joshua Angrist, Alberto Abadie and Christopher Walters for their amazing Econometrics class. Most of the ideas here are taken from their classes at the American Economic Association. Watching them is what is keeping me sane during this tough year of 2020.\n",
    "* [Cross-Section Econometrics](https://www.aeaweb.org/conference/cont-ed/2017-webcasts)\n",
    "* [Mastering Mostly Harmless Econometrics](https://www.aeaweb.org/conference/cont-ed/2020-webcasts)\n",
    "\n",
    "I'll also like to reference the amazing books from Angrist. They have shown me that Econometrics, or 'Metrics as they call it, is not only extremely useful but also profoundly fun.\n",
    "\n",
    "* [Mostly Harmless Econometrics](https://www.mostlyharmlesseconometrics.com/)\n",
    "* [Mastering 'Metrics](https://www.masteringmetrics.com/)\n",
    "\n",
    "My final reference is Miguel Hernan and Jamie Robins' book. It has been my trustworthy companion in the most thorny causal questions I had to answer.\n",
    "\n",
    "* [Causal Inference Book](https://www.hsph.harvard.edu/miguel-hernan/causal-inference-book/)\n",
    "\n",
    "The data that we used was taken from the article [Estimating Treatment Effects with Causal Forests: An Application](https://arxiv.org/pdf/1902.07409.pdf), by Susan Athey and Stefan Wager. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}