Diff-in-Diff 2x2 with Covariates.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1e898411",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-12-08T15:10:36.386845Z",
     "start_time": "2022-12-08T15:10:35.661343Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import statsmodels.formula.api as smf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "987b2358",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-12-08T15:10:36.482666Z",
     "start_time": "2022-12-08T15:10:36.457972Z"
    }
   },
   "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>date</th>\n",
       "      <th>unit</th>\n",
       "      <th>region</th>\n",
       "      <th>treated_unit</th>\n",
       "      <th>unit_fe</th>\n",
       "      <th>time_fe</th>\n",
       "      <th>reg_ps</th>\n",
       "      <th>reg_fe</th>\n",
       "      <th>reg_trend</th>\n",
       "      <th>treated</th>\n",
       "      <th>post</th>\n",
       "      <th>y0</th>\n",
       "      <th>y1</th>\n",
       "      <th>tau</th>\n",
       "      <th>y</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2020-01-01</td>\n",
       "      <td>1</td>\n",
       "      <td>N</td>\n",
       "      <td>0</td>\n",
       "      <td>-2.226872</td>\n",
       "      <td>0.288617</td>\n",
       "      <td>0.4</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>11.061745</td>\n",
       "      <td>12.061745</td>\n",
       "      <td>1.0</td>\n",
       "      <td>12.693680</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2021-01-01</td>\n",
       "      <td>1</td>\n",
       "      <td>N</td>\n",
       "      <td>0</td>\n",
       "      <td>-2.226872</td>\n",
       "      <td>0.058307</td>\n",
       "      <td>0.4</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>13.831435</td>\n",
       "      <td>14.831435</td>\n",
       "      <td>1.0</td>\n",
       "      <td>13.429646</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2020-01-01</td>\n",
       "      <td>2</td>\n",
       "      <td>W</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.134820</td>\n",
       "      <td>0.288617</td>\n",
       "      <td>0.8</td>\n",
       "      <td>-1</td>\n",
       "      <td>7</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>9.153797</td>\n",
       "      <td>10.153797</td>\n",
       "      <td>1.0</td>\n",
       "      <td>8.953857</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2021-01-01</td>\n",
       "      <td>2</td>\n",
       "      <td>W</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.134820</td>\n",
       "      <td>0.058307</td>\n",
       "      <td>0.8</td>\n",
       "      <td>-1</td>\n",
       "      <td>7</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>15.923487</td>\n",
       "      <td>16.923487</td>\n",
       "      <td>1.0</td>\n",
       "      <td>15.930876</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2020-01-01</td>\n",
       "      <td>3</td>\n",
       "      <td>S</td>\n",
       "      <td>0</td>\n",
       "      <td>2.322880</td>\n",
       "      <td>0.288617</td>\n",
       "      <td>0.2</td>\n",
       "      <td>5</td>\n",
       "      <td>-1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>17.611497</td>\n",
       "      <td>18.611497</td>\n",
       "      <td>1.0</td>\n",
       "      <td>17.887161</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        date  unit region  treated_unit   unit_fe   time_fe  reg_ps  reg_fe  \\\n",
       "0 2020-01-01     1      N             0 -2.226872  0.288617     0.4       3   \n",
       "0 2021-01-01     1      N             0 -2.226872  0.058307     0.4       3   \n",
       "1 2020-01-01     2      W             0 -0.134820  0.288617     0.8      -1   \n",
       "1 2021-01-01     2      W             0 -0.134820  0.058307     0.8      -1   \n",
       "2 2020-01-01     3      S             0  2.322880  0.288617     0.2       5   \n",
       "\n",
       "   reg_trend  treated  post         y0         y1  tau          y  \n",
       "0          3        0     0  11.061745  12.061745  1.0  12.693680  \n",
       "0          3        0     1  13.831435  14.831435  1.0  13.429646  \n",
       "1          7        0     0   9.153797  10.153797  1.0   8.953857  \n",
       "1          7        0     1  15.923487  16.923487  1.0  15.930876  \n",
       "2         -1        0     0  17.611497  18.611497  1.0  17.887161  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "date = pd.to_datetime([\"2020-01-01\", \"2021-01-01\"])\n",
    "units = np.array(range(1, 500+1))\n",
    "\n",
    "region = np.array([\"S\", \"N\", \"E\", \"W\"])\n",
    "reg_ps = dict(zip(region,    [0.2, 0.4, 0.6, 0.8]))\n",
    "reg_fe = dict(zip(region,    [5, 3, 1, -1]))\n",
    "reg_trend = dict(zip(region, [-1, 3, 6, 7]))\n",
    "\n",
    "np.random.seed(1)\n",
    "unit_region = pd.Series(np.random.choice(region, len(units)))\n",
    "treated_unit = np.random.binomial(1, unit_region.map(reg_ps))\n",
    "\n",
    "df = pd.DataFrame(dict(\n",
    "    date = np.tile(date, len(units)),\n",
    "    unit = np.repeat(units, len(date)),\n",
    "    region = np.repeat(unit_region, len(date)),\n",
    "    treated_unit = np.repeat(treated_unit, len(date)),\n",
    "    \n",
    "    unit_fe = np.repeat(np.random.normal(0, 2, size=len(units)), len(date)),\n",
    "    time_fe = np.tile(np.random.normal(size=len(date)), len(units)),\n",
    "    \n",
    "    reg_ps = pd.Series(np.repeat(unit_region, len(date))).map(reg_ps),\n",
    "    reg_fe = pd.Series(np.repeat(unit_region, len(date))).map(reg_fe),\n",
    "    reg_trend = pd.Series(np.repeat(unit_region, len(date))).map(reg_trend))\n",
    ").assign(\n",
    "    treated = lambda d: (d[\"treated_unit\"] == 1).astype(int),\n",
    "    post = lambda d: (d[\"date\"] >= \"2021-01-01\").astype(int)\n",
    ").assign(\n",
    "    y0 = lambda d: (10\n",
    "                    + d[\"unit_fe\"] \n",
    "                    + d[\"time_fe\"] \n",
    "                    + d[\"reg_fe\"]\n",
    "                    + d[\"reg_trend\"]*(d[\"date\"] >= \"2021-01-01\")),\n",
    ").assign(\n",
    "    y1 = lambda d: d[\"y0\"] + 1\n",
    ").assign(\n",
    "    tau = lambda d: d[\"y1\"] - d[\"y0\"],\n",
    "    y = lambda d: np.where(d[\"treated\"]*d[\"post\"] == 1, d[\"y1\"], d[\"y0\"]) + np.random.normal(0,1,len(d))\n",
    ")\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9b7e16b0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-12-08T15:10:43.779644Z",
     "start_time": "2022-12-08T15:10:43.747868Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table class=\"simpletable\">\n",
       "<tr>\n",
       "        <td></td>          <th>coef</th>     <th>std err</th>      <th>t</th>      <th>P>|t|</th>  <th>[0.025</th>    <th>0.975]</th>  \n",
       "</tr>\n",
       "<tr>\n",
       "  <th>Intercept</th>    <td>   13.4804</td> <td>    0.170</td> <td>   79.434</td> <td> 0.000</td> <td>   13.147</td> <td>   13.813</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post</th>         <td>    2.0901</td> <td>    0.240</td> <td>    8.709</td> <td> 0.000</td> <td>    1.619</td> <td>    2.561</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>treated</th>      <td>   -2.1725</td> <td>    0.241</td> <td>   -8.998</td> <td> 0.000</td> <td>   -2.646</td> <td>   -1.699</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post:treated</th> <td>    3.6291</td> <td>    0.341</td> <td>   10.628</td> <td> 0.000</td> <td>    2.959</td> <td>    4.299</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<class 'statsmodels.iolib.table.SimpleTable'>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smf.ols(\"y~post*treated\", data=df).fit().summary().tables[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cf240050",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-12-08T15:10:48.794807Z",
     "start_time": "2022-12-08T15:10:48.763423Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table class=\"simpletable\">\n",
       "<tr>\n",
       "         <td></td>           <th>coef</th>     <th>std err</th>      <th>t</th>      <th>P>|t|</th>  <th>[0.025</th>    <th>0.975]</th>  \n",
       "</tr>\n",
       "<tr>\n",
       "  <th>Intercept</th>      <td>   13.1802</td> <td>    0.232</td> <td>   56.871</td> <td> 0.000</td> <td>   12.725</td> <td>   13.635</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>C(region)[T.N]</th> <td>    0.3098</td> <td>    0.242</td> <td>    1.281</td> <td> 0.201</td> <td>   -0.165</td> <td>    0.785</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>C(region)[T.S]</th> <td>    0.6852</td> <td>    0.243</td> <td>    2.823</td> <td> 0.005</td> <td>    0.209</td> <td>    1.162</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>C(region)[T.W]</th> <td>   -0.9844</td> <td>    0.246</td> <td>   -3.994</td> <td> 0.000</td> <td>   -1.468</td> <td>   -0.501</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post</th>           <td>    2.0901</td> <td>    0.235</td> <td>    8.886</td> <td> 0.000</td> <td>    1.629</td> <td>    2.552</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>treated</th>        <td>   -1.6090</td> <td>    0.252</td> <td>   -6.389</td> <td> 0.000</td> <td>   -2.103</td> <td>   -1.115</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post:treated</th>   <td>    3.6291</td> <td>    0.335</td> <td>   10.844</td> <td> 0.000</td> <td>    2.972</td> <td>    4.286</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<class 'statsmodels.iolib.table.SimpleTable'>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smf.ols(\"y~post*treated + C(region)\", data=df).fit().summary().tables[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "43aeb07e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-12-08T15:10:56.854680Z",
     "start_time": "2022-12-08T15:10:56.824697Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table class=\"simpletable\">\n",
       "<tr>\n",
       "           <td></td>              <th>coef</th>     <th>std err</th>      <th>t</th>      <th>P>|t|</th>  <th>[0.025</th>    <th>0.975]</th>  \n",
       "</tr>\n",
       "<tr>\n",
       "  <th>Intercept</th>           <td>   11.3287</td> <td>    0.239</td> <td>   47.379</td> <td> 0.000</td> <td>   10.859</td> <td>   11.798</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>C(region)[T.N]</th>      <td>    1.7889</td> <td>    0.290</td> <td>    6.175</td> <td> 0.000</td> <td>    1.220</td> <td>    2.357</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>C(region)[T.S]</th>      <td>    4.1519</td> <td>    0.291</td> <td>   14.287</td> <td> 0.000</td> <td>    3.582</td> <td>    4.722</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>C(region)[T.W]</th>      <td>   -1.4337</td> <td>    0.295</td> <td>   -4.858</td> <td> 0.000</td> <td>   -2.013</td> <td>   -0.855</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post</th>                <td>    5.7932</td> <td>    0.338</td> <td>   17.132</td> <td> 0.000</td> <td>    5.130</td> <td>    6.457</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post:C(region)[T.N]</th> <td>   -2.9581</td> <td>    0.410</td> <td>   -7.220</td> <td> 0.000</td> <td>   -3.762</td> <td>   -2.154</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post:C(region)[T.S]</th> <td>   -6.9333</td> <td>    0.411</td> <td>  -16.870</td> <td> 0.000</td> <td>   -7.740</td> <td>   -6.127</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post:C(region)[T.W]</th> <td>    0.8986</td> <td>    0.417</td> <td>    2.153</td> <td> 0.032</td> <td>    0.080</td> <td>    1.718</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>treated</th>             <td>   -0.2784</td> <td>    0.225</td> <td>   -1.235</td> <td> 0.217</td> <td>   -0.721</td> <td>    0.164</td>\n",
       "</tr>\n",
       "<tr>\n",
       "  <th>post:treated</th>        <td>    0.9678</td> <td>    0.319</td> <td>    3.036</td> <td> 0.002</td> <td>    0.342</td> <td>    1.593</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<class 'statsmodels.iolib.table.SimpleTable'>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smf.ols(\"y~post*(treated + C(region))\", data=df).fit().summary().tables[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b0c370ab",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-12-08T15:11:15.199914Z",
     "start_time": "2022-12-08T15:11:15.156952Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N 0.9226567466899882\n",
      "W 1.1623121008046138\n",
      "S 0.9169824580119377\n",
      "E 0.9421142952058243\n"
     ]
    }
   ],
   "source": [
    "for region in df[\"region\"].unique():\n",
    "    att = smf.ols(\"y~post*treated\", data=df.query(\"region==@region\")).fit().params[\"post:treated\"]\n",
    "    print(region, att)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a0c983a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root] *",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.7.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}