{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.568697Z",
     "start_time": "2019-10-12T20:04:00.913074Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import adjustText\n",
    "import pathlib\n",
    "from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier\n",
    "from sklearn.metrics import accuracy_score, mean_squared_error\n",
    "from sklearn.model_selection import KFold, StratifiedKFold\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from plotnine import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.592682Z",
     "start_time": "2019-10-12T20:04:09.576685Z"
    }
   },
   "outputs": [],
   "source": [
    "def make_data_imbens(df):\n",
    "    df_new=df.drop(['nodegree'], axis=1)\n",
    "    df_new['pos74']=(df_new['RE74']>0).astype(int)\n",
    "    df_new['pos75']=(df_new['RE75']>0).astype(int)\n",
    "    return df_new\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.608682Z",
     "start_time": "2019-10-12T20:04:09.596682Z"
    }
   },
   "outputs": [],
   "source": [
    "def predict_g(X, y, do_kfold):\n",
    "    np.random.seed(42)\n",
    "    predictions=np.full_like(y, np.nan, dtype=float)\n",
    "    if do_kfold:        \n",
    "        kf=StratifiedKFold(n_splits=10, shuffle=True, random_state=42)\n",
    "        for train_index, test_index in kf.split(X,y):\n",
    "            clf=RandomForestClassifier(random_state=42, oob_score=False, n_estimators=10, max_depth=3, class_weight='balanced')\n",
    "            clf.fit(X[train_index],y[train_index])\n",
    "            predictions[test_index]=clf.predict_proba(X[test_index])[:,1]\n",
    "    #print('oob_score: '+ str(clf.oob_score_))\n",
    "    else:\n",
    "        clf=RandomForestClassifier(random_state=42, oob_score=False, n_estimators=10, max_depth=3, class_weight='balanced')\n",
    "        clf.fit(X,y)\n",
    "        predictions=clf.predict_proba(X)[:,1]\n",
    "    assert np.isnan(predictions).sum()==0\n",
    "    return predictions\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.624680Z",
     "start_time": "2019-10-12T20:04:09.608682Z"
    }
   },
   "outputs": [],
   "source": [
    "def predict_Q(X,y,do_kfold):\n",
    "    np.random.seed(42)\n",
    "    predictions=np.full_like(y, np.nan, dtype=float)\n",
    "    if do_kfold:\n",
    "        kf=KFold(n_splits=10, shuffle=True, random_state=42)\n",
    "        for train_index, test_index in kf.split(X):\n",
    "            reg=RandomForestRegressor(random_state=42, n_estimators=10, max_depth=3)\n",
    "            reg.fit(X[train_index], y[train_index])\n",
    "            predictions[test_index]=reg.predict(X[test_index])\n",
    "    else:\n",
    "        reg=RandomForestRegressor(random_state=42, n_estimators=10, max_depth=3)\n",
    "        reg.fit(X, y)\n",
    "        predictions=reg.predict(X)\n",
    "    assert np.isnan(predictions).sum()==0\n",
    "    return predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.636680Z",
     "start_time": "2019-10-12T20:04:09.624680Z"
    }
   },
   "outputs": [],
   "source": [
    "def calc_delta(alpha, g, bias):\n",
    "    alpha_term=(1/alpha)-1\n",
    "    bias_term=sp.special.digamma(g*alpha_term+1)\\\n",
    "        -sp.special.digamma((1-g)*alpha_term)\\\n",
    "        -sp.special.digamma(g*alpha_term)\\\n",
    "        +sp.special.digamma((1-g)*alpha_term+1)\n",
    "    delta=bias/np.mean(bias_term)\n",
    "    return delta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.644679Z",
     "start_time": "2019-10-12T20:04:09.636680Z"
    }
   },
   "outputs": [],
   "source": [
    "def calc_beta_shapes(g, alpha):\n",
    "    alpha_shape=g*((1/alpha)-1)\n",
    "    beta_shape=(1-g)*((1/alpha)-1)\n",
    "    #print(alpha_shape, beta_shape)\n",
    "    return alpha_shape, beta_shape\n",
    "    \n",
    "\n",
    "def calc_Rsq(alpha, delta, calculated_params, treatment, response):\n",
    "    alpha_shape, beta_shape=calc_beta_shapes(calculated_params['g'], alpha)\n",
    "    Rsq_num=delta**2*np.mean(sp.special.polygamma(1, alpha_shape+treatment)\\\n",
    "                             +sp.special.polygamma(1, beta_shape+(1-treatment)))\n",
    "    Rsq_den=mean_squared_error(response, calculated_params['Q'])\n",
    "    #print(Rsq_num, Rsq_den)\n",
    "    return Rsq_num/Rsq_den\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.656679Z",
     "start_time": "2019-10-12T20:04:09.644679Z"
    }
   },
   "outputs": [],
   "source": [
    "def calc_Rsq_hat(covariate_col, response_col, df, Q, do_kfold):\n",
    "    y=df[response_col].values\n",
    "    Qhat=predict_Q(df.drop([response_col]+covariate_col, axis=1).values, y, do_kfold)\n",
    "    Rsqhat=(mean_squared_error(y, Qhat)-mean_squared_error(y, Q))/(mean_squared_error(y, Qhat))\n",
    "    if Rsqhat<0:\n",
    "        Rsqhat=0\n",
    "    return Rsqhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.668679Z",
     "start_time": "2019-10-12T20:04:09.656679Z"
    }
   },
   "outputs": [],
   "source": [
    "def calc_alpha_hat(covariate_col, response_col, treatment_col, df, g, do_kfold):\n",
    "    #call y t\n",
    "    t=df[treatment_col].values\n",
    "    ghat=predict_g(df.drop([response_col, treatment_col]+covariate_col, axis=1).values, t, do_kfold)\n",
    "#     ahat_numerator=np.var(g)-np.var(ghat)\n",
    "#     ahat_denominator=np.mean(ghat*(1-ghat))\n",
    "    ahat=1-(np.mean(g*(1-g))/np.mean(ghat*(1-ghat)))\n",
    "    #ahat=ahat_numerator/ahat_denominator\n",
    "    if ahat<0:\n",
    "        ahat=0\n",
    "    return ahat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.708678Z",
     "start_time": "2019-10-12T20:04:09.672678Z"
    }
   },
   "outputs": [],
   "source": [
    "def plot_sensitivity_graph(df, treatment_col, response_col, covariate_cols, covariate_groups, bias, do_kfold):\n",
    "    #Calculate g and Q\n",
    "    g=predict_g(df.drop([response_col, treatment_col], axis=1).values, df[treatment_col].values, do_kfold)\n",
    "    # important to solve numerical precision issue with scipy's implementation of digamma function\n",
    "    g[g<0.001] = 0.001\n",
    "    g[g>0.999] = 0.999\n",
    "    Q=predict_Q(df.drop([response_col], axis=1).values, df[response_col].values, do_kfold)\n",
    "    calculated_params=pd.DataFrame({'g': g, 'Q': Q})\n",
    "    \n",
    "    #Calculate alpha, delta, Rsq\n",
    "    alpha_series=pd.Series(np.arange(0.0001,1,0.0001))\n",
    "    #alpha_series=pd.Series(np.arange(0.001,1,0.001))\n",
    "    delta_series=alpha_series.apply(calc_delta, g=calculated_params['g'], bias=bias)\n",
    "    sensitivity_df=pd.DataFrame({'alphas':alpha_series, 'deltas':delta_series})\n",
    "    Rsq=[]\n",
    "    for n in range(len(alpha_series)):\n",
    "        Rsq.append(calc_Rsq(alpha_series[n], delta_series[n], calculated_params, df[treatment_col], df[response_col]))\n",
    "    sensitivity_df['Rsq']=Rsq\n",
    "    \n",
    "    #Plot observed co-variates\n",
    "    alpha_hat=[]\n",
    "    Rsq_hat=[]\n",
    "    \n",
    "    for covar in covariate_cols:\n",
    "        alpha_hat.append(calc_alpha_hat(covar, response_col, treatment_col, df, calculated_params['g'], do_kfold))\n",
    "        Rsq_hat.append(calc_Rsq_hat(covar, response_col, df, calculated_params['Q'], do_kfold))\n",
    "    \n",
    "    variable_importances=pd.DataFrame({'covariates':covariate_cols, 'alpha_hat': alpha_hat, 'Rsq_hat':Rsq_hat, 'groups':covariate_groups})\n",
    "    \n",
    "    print(variable_importances)\n",
    "    \n",
    "    #variable_importances['labels']=variable_importances['covariates'].str[0]\n",
    "    #variable_importances.loc[variable_importances['groups']!='individual_covariates', 'labels']=''\n",
    "    #variable_importances.loc[(variable_importances['Rsq_hat']<0.05) & (variable_importances['alpha_hat']<0.2), 'labels']=''\n",
    "    print(sensitivity_df)\n",
    "    \n",
    "    return (ggplot(data=sensitivity_df, mapping=aes(x='alphas', y='Rsq'))\n",
    "            +geom_line(color='#585858', size=1)\n",
    "            +theme_light()\n",
    "            +geom_point(data=variable_importances, mapping=aes(x='alpha_hat', y='Rsq_hat', fill='groups'), color=\"black\", alpha=0.6, size=2.5)\n",
    "            +theme(figure_size=(3.5,3.5), legend_key = element_blank(), axis_title=element_text(size=10), axis_text=element_text(color='black'), plot_title=element_text(size=10))\n",
    "            +scale_x_continuous(expand=[0, 0, 0, 0], limits=(-0.01,0.43))\n",
    "            #+scale_y_continuous(expand=[0, 0, 0, 0], limits=(-0.02, min(0.75, max(Rsq+Rsq_hat))))\n",
    "            +scale_y_continuous(expand=[0, 0, 0, 0], limits=(-0.02, 0.6))\n",
    "            +labs(x=r'$\\mathregular{\\alpha}$ (treatment)', fill='', y='partial $R^2$ (outcome)')\n",
    "            +scale_fill_brewer(type='qual', palette='Set1', labels=['Individual covariates', 'Pre-program earnings', 'Recent earnings'])\n",
    "            #+geom_label(data=variable_importances, mapping=aes(x='alpha_hat', y='Rsq_hat', label='labels'), size=8, nudge_y=0.03, adjust_text={'expand_text': (2, 2), 'expand_points': (2,2), 'force_text': (0.1, 0.25), 'force_points':(0.2, 0.5)})\n",
    "            +annotate(\"text\", x=0.35, y=0.065,label='Bias = $'+str(bias), size=8.5, color='#303030')\n",
    "           )\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.716677Z",
     "start_time": "2019-10-12T20:04:09.708678Z"
    }
   },
   "outputs": [],
   "source": [
    "col_names=['treatment', 'age', 'education', 'black', 'hispanic', 'married', 'nodegree', 'RE74', 'RE75', 'RE78']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.940669Z",
     "start_time": "2019-10-12T20:04:09.716677Z"
    }
   },
   "outputs": [],
   "source": [
    "control=pd.read_csv(pathlib.Path.cwd().parents[0].joinpath('data', 'nswre74_control.txt'), header=None, sep=r\"\\s\\s\", names=col_names, engine='python')\n",
    "treatment=pd.read_csv(pathlib.Path.cwd().parents[0].joinpath('data', 'nswre74_treated.txt'), header=None, sep=r\"\\s\\s\", names=col_names, engine='python')\n",
    "psid=pd.read_csv(pathlib.Path.cwd().parents[0].joinpath('data', 'psid_controls.txt'), header=None, sep=r\"\\s\\s\", names=col_names, engine='python')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:09.964668Z",
     "start_time": "2019-10-12T20:04:09.944669Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(260, 10) (185, 10) (2490, 10)\n"
     ]
    }
   ],
   "source": [
    "print(control.shape, treatment.shape, psid.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:10.040666Z",
     "start_time": "2019-10-12T20:04:09.972668Z"
    }
   },
   "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>treatment</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>black</th>\n",
       "      <th>hispanic</th>\n",
       "      <th>married</th>\n",
       "      <th>nodegree</th>\n",
       "      <th>RE74</th>\n",
       "      <th>RE75</th>\n",
       "      <th>RE78</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>23.0</td>\n",
       "      <td>10.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>26.0</td>\n",
       "      <td>12.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>12383.68</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>22.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.0</td>\n",
       "      <td>18.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>10740.08</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.0</td>\n",
       "      <td>45.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>11796.47</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   treatment   age  education  black  hispanic  married  nodegree  RE74  RE75  \\\n",
       "0        0.0  23.0       10.0    1.0       0.0      0.0       1.0   0.0   0.0   \n",
       "1        0.0  26.0       12.0    0.0       0.0      0.0       0.0   0.0   0.0   \n",
       "2        0.0  22.0        9.0    1.0       0.0      0.0       1.0   0.0   0.0   \n",
       "3        0.0  18.0        9.0    1.0       0.0      0.0       1.0   0.0   0.0   \n",
       "4        0.0  45.0       11.0    1.0       0.0      0.0       1.0   0.0   0.0   \n",
       "\n",
       "       RE78  \n",
       "0      0.00  \n",
       "1  12383.68  \n",
       "2      0.00  \n",
       "3  10740.08  \n",
       "4  11796.47  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imbens1=pd.concat([control,treatment]).reset_index(drop=True)\n",
    "imbens1.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:10.108663Z",
     "start_time": "2019-10-12T20:04:10.044667Z"
    }
   },
   "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>treatment</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>black</th>\n",
       "      <th>hispanic</th>\n",
       "      <th>married</th>\n",
       "      <th>RE74</th>\n",
       "      <th>RE75</th>\n",
       "      <th>RE78</th>\n",
       "      <th>pos74</th>\n",
       "      <th>pos75</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>23.0</td>\n",
       "      <td>10.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>26.0</td>\n",
       "      <td>12.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>12383.68</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>22.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.0</td>\n",
       "      <td>18.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>10740.08</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.0</td>\n",
       "      <td>45.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>11796.47</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   treatment   age  education  black  hispanic  married  RE74  RE75      RE78  \\\n",
       "0        0.0  23.0       10.0    1.0       0.0      0.0   0.0   0.0      0.00   \n",
       "1        0.0  26.0       12.0    0.0       0.0      0.0   0.0   0.0  12383.68   \n",
       "2        0.0  22.0        9.0    1.0       0.0      0.0   0.0   0.0      0.00   \n",
       "3        0.0  18.0        9.0    1.0       0.0      0.0   0.0   0.0  10740.08   \n",
       "4        0.0  45.0       11.0    1.0       0.0      0.0   0.0   0.0  11796.47   \n",
       "\n",
       "   pos74  pos75  \n",
       "0      0      0  \n",
       "1      0      0  \n",
       "2      0      0  \n",
       "3      0      0  \n",
       "4      0      0  "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imbens1=make_data_imbens(imbens1)\n",
    "imbens1.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:10.132662Z",
     "start_time": "2019-10-12T20:04:10.112663Z"
    }
   },
   "outputs": [],
   "source": [
    "imbens2=pd.concat([treatment, psid]).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:10.196659Z",
     "start_time": "2019-10-12T20:04:10.140662Z"
    }
   },
   "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>treatment</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>black</th>\n",
       "      <th>hispanic</th>\n",
       "      <th>married</th>\n",
       "      <th>RE74</th>\n",
       "      <th>RE75</th>\n",
       "      <th>RE78</th>\n",
       "      <th>pos74</th>\n",
       "      <th>pos75</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>9930.0460</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>22.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3595.8940</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.0</td>\n",
       "      <td>30.0</td>\n",
       "      <td>12.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>24909.4500</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>27.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>7506.1460</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>1.0</td>\n",
       "      <td>33.0</td>\n",
       "      <td>8.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>289.7899</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   treatment   age  education  black  hispanic  married  RE74  RE75  \\\n",
       "0        1.0  37.0       11.0    1.0       0.0      1.0   0.0   0.0   \n",
       "1        1.0  22.0        9.0    0.0       1.0      0.0   0.0   0.0   \n",
       "2        1.0  30.0       12.0    1.0       0.0      0.0   0.0   0.0   \n",
       "3        1.0  27.0       11.0    1.0       0.0      0.0   0.0   0.0   \n",
       "4        1.0  33.0        8.0    1.0       0.0      0.0   0.0   0.0   \n",
       "\n",
       "         RE78  pos74  pos75  \n",
       "0   9930.0460      0      0  \n",
       "1   3595.8940      0      0  \n",
       "2  24909.4500      0      0  \n",
       "3   7506.1460      0      0  \n",
       "4    289.7899      0      0  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imbens2=make_data_imbens(imbens2)\n",
    "imbens2.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:10.304656Z",
     "start_time": "2019-10-12T20:04:10.200660Z"
    }
   },
   "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>treatment</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>black</th>\n",
       "      <th>hispanic</th>\n",
       "      <th>married</th>\n",
       "      <th>RE74</th>\n",
       "      <th>RE75</th>\n",
       "      <th>change_earnings</th>\n",
       "      <th>pos74</th>\n",
       "      <th>pos75</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>9930.0460</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>22.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3595.8940</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.0</td>\n",
       "      <td>30.0</td>\n",
       "      <td>12.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>24909.4500</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>27.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>7506.1460</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>1.0</td>\n",
       "      <td>33.0</td>\n",
       "      <td>8.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>289.7899</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   treatment   age  education  black  hispanic  married  RE74  RE75  \\\n",
       "0        1.0  37.0       11.0    1.0       0.0      1.0   0.0   0.0   \n",
       "1        1.0  22.0        9.0    0.0       1.0      0.0   0.0   0.0   \n",
       "2        1.0  30.0       12.0    1.0       0.0      0.0   0.0   0.0   \n",
       "3        1.0  27.0       11.0    1.0       0.0      0.0   0.0   0.0   \n",
       "4        1.0  33.0        8.0    1.0       0.0      0.0   0.0   0.0   \n",
       "\n",
       "   change_earnings  pos74  pos75  \n",
       "0        9930.0460      0      0  \n",
       "1        3595.8940      0      0  \n",
       "2       24909.4500      0      0  \n",
       "3        7506.1460      0      0  \n",
       "4         289.7899      0      0  "
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imbens3=pd.concat([treatment, psid]).reset_index(drop=True)\n",
    "imbens3['change_earnings']=imbens3['RE78']-imbens3['RE74']\n",
    "imbens3=make_data_imbens(imbens3)\n",
    "imbens3=imbens3.drop('RE78', axis=1)\n",
    "imbens3.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:10.344654Z",
     "start_time": "2019-10-12T20:04:10.304656Z"
    },
    "scrolled": true
   },
   "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>age</th>\n",
       "      <th>education</th>\n",
       "      <th>black</th>\n",
       "      <th>hispanic</th>\n",
       "      <th>married</th>\n",
       "      <th>nodegree</th>\n",
       "      <th>RE74</th>\n",
       "      <th>RE75</th>\n",
       "      <th>RE78</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>treatment</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0.0</td>\n",
       "      <td>242</td>\n",
       "      <td>242</td>\n",
       "      <td>242</td>\n",
       "      <td>242</td>\n",
       "      <td>242</td>\n",
       "      <td>242</td>\n",
       "      <td>242</td>\n",
       "      <td>242</td>\n",
       "      <td>242</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1.0</td>\n",
       "      <td>148</td>\n",
       "      <td>148</td>\n",
       "      <td>148</td>\n",
       "      <td>148</td>\n",
       "      <td>148</td>\n",
       "      <td>148</td>\n",
       "      <td>148</td>\n",
       "      <td>148</td>\n",
       "      <td>148</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           age  education  black  hispanic  married  nodegree  RE74  RE75  \\\n",
       "treatment                                                                   \n",
       "0.0        242        242    242       242      242       242   242   242   \n",
       "1.0        148        148    148       148      148       148   148   148   \n",
       "\n",
       "           RE78  \n",
       "treatment        \n",
       "0.0         242  \n",
       "1.0         148  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imbens4=pd.concat([treatment, psid]).reset_index(drop=True)\n",
    "imbens4=imbens4[~((imbens4['RE74']>=5000) | (imbens4['RE75']>=5000))]\n",
    "imbens4=imbens4.reset_index(drop=True)\n",
    "#Does this match imbens' number of datapoints?\n",
    "imbens4.groupby('treatment').count()\n",
    "#yes!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:10.440651Z",
     "start_time": "2019-10-12T20:04:10.352655Z"
    }
   },
   "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>treatment</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>black</th>\n",
       "      <th>hispanic</th>\n",
       "      <th>married</th>\n",
       "      <th>RE74</th>\n",
       "      <th>RE75</th>\n",
       "      <th>RE78</th>\n",
       "      <th>pos74</th>\n",
       "      <th>pos75</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>9930.0460</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>22.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3595.8940</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.0</td>\n",
       "      <td>30.0</td>\n",
       "      <td>12.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>24909.4500</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>27.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>7506.1460</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>1.0</td>\n",
       "      <td>33.0</td>\n",
       "      <td>8.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>289.7899</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   treatment   age  education  black  hispanic  married  RE74  RE75  \\\n",
       "0        1.0  37.0       11.0    1.0       0.0      1.0   0.0   0.0   \n",
       "1        1.0  22.0        9.0    0.0       1.0      0.0   0.0   0.0   \n",
       "2        1.0  30.0       12.0    1.0       0.0      0.0   0.0   0.0   \n",
       "3        1.0  27.0       11.0    1.0       0.0      0.0   0.0   0.0   \n",
       "4        1.0  33.0        8.0    1.0       0.0      0.0   0.0   0.0   \n",
       "\n",
       "         RE78  pos74  pos75  \n",
       "0   9930.0460      0      0  \n",
       "1   3595.8940      0      0  \n",
       "2  24909.4500      0      0  \n",
       "3   7506.1460      0      0  \n",
       "4    289.7899      0      0  "
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imbens4=make_data_imbens(imbens4)\n",
    "imbens4.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:04:10.452650Z",
     "start_time": "2019-10-12T20:04:10.440651Z"
    }
   },
   "outputs": [],
   "source": [
    "covariate_cols=['age', 'education', 'black', 'hispanic', 'married', 'RE74', 'RE75', 'pos74', 'pos75']\n",
    "recent_earnings=['RE75', 'pos75']\n",
    "#what about \"pre-program earnings\"? Is that RE74?\n",
    "pre_program_earnings=['RE74', 'RE75', 'pos74', 'pos75']\n",
    "\n",
    "covariates_for_graphs=[[x] for x in covariate_cols]\n",
    "covariates_for_graphs.append(recent_earnings)\n",
    "covariates_for_graphs.append(pre_program_earnings)\n",
    "\n",
    "covariate_groups=['individual_covariates']*len(covariate_cols)+['recent_earnings', 'pre_program_earnings']\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imbens1 results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:05:04.080513Z",
     "start_time": "2019-10-12T20:04:10.456650Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                    covariates  alpha_hat   Rsq_hat                 groups\n",
      "0                        [age]   0.001604  0.000000  individual_covariates\n",
      "1                  [education]   0.006287  0.003081  individual_covariates\n",
      "2                      [black]   0.000329  0.030177  individual_covariates\n",
      "3                   [hispanic]   0.002467  0.000894  individual_covariates\n",
      "4                    [married]   0.000409  0.001636  individual_covariates\n",
      "5                       [RE74]   0.002401  0.008671  individual_covariates\n",
      "6                       [RE75]   0.009153  0.000000  individual_covariates\n",
      "7                      [pos74]   0.006828  0.007203  individual_covariates\n",
      "8                      [pos75]   0.008426  0.006909  individual_covariates\n",
      "9                [RE75, pos75]   0.008612  0.000000        recent_earnings\n",
      "10  [RE74, RE75, pos74, pos75]   0.005429  0.015254   pre_program_earnings\n",
      "      alphas        deltas        Rsq\n",
      "0     0.0001  2.426543e+06  53.746893\n",
      "1     0.0002  1.213150e+06  26.870871\n",
      "2     0.0003  8.086860e+05  17.912197\n",
      "3     0.0004  6.064538e+05  13.432861\n",
      "4     0.0005  4.851145e+05  10.745259\n",
      "...      ...           ...        ...\n",
      "9994  0.9995  1.214000e-01   0.005947\n",
      "9995  0.9996  9.711028e-02   0.005947\n",
      "9996  0.9997  7.282543e-02   0.005947\n",
      "9997  0.9998  4.854543e-02   0.005947\n",
      "9998  0.9999  2.427029e-02   0.005947\n",
      "\n",
      "[9999 rows x 3 columns]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:75: PlotnineWarning: geom_path: Removed 5787 rows containing missing values.\n",
      "  warn(msg.format(n1-n2), PlotnineWarning)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29804 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (94728667519)>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plot_sensitivity_graph(imbens1, treatment_col='treatment', response_col='RE78', covariate_cols=covariates_for_graphs, covariate_groups=covariate_groups, bias=1000, do_kfold=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:05:04.084512Z",
     "start_time": "2019-10-12T20:05:04.080513Z"
    }
   },
   "outputs": [],
   "source": [
    "# plot_sensitivity_graph(imbens1, treatment_col='treatment', response_col='RE78', covariate_cols=covariates_for_graphs, covariate_groups=covariate_groups, bias=1000, do_kfold=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imbens2 results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:05:57.480019Z",
     "start_time": "2019-10-12T20:05:04.088515Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                    covariates  alpha_hat       Rsq_hat                 groups\n",
      "0                        [age]   0.036161  2.042829e-04  individual_covariates\n",
      "1                  [education]   0.055530  4.018286e-04  individual_covariates\n",
      "2                      [black]   0.090657  1.366872e-16  individual_covariates\n",
      "3                   [hispanic]   0.050661  1.366872e-16  individual_covariates\n",
      "4                    [married]   0.126502  1.366872e-16  individual_covariates\n",
      "5                       [RE74]   0.186596  1.512373e-02  individual_covariates\n",
      "6                       [RE75]   0.243562  8.049797e-02  individual_covariates\n",
      "7                      [pos74]   0.182654  1.366872e-16  individual_covariates\n",
      "8                      [pos75]   0.100062  1.366872e-16  individual_covariates\n",
      "9                [RE75, pos75]   0.217591  1.050862e-01        recent_earnings\n",
      "10  [RE74, RE75, pos74, pos75]   0.318959  4.400202e-01   pre_program_earnings\n",
      "      alphas         deltas       Rsq\n",
      "0     0.0001  288657.539263  2.659410\n",
      "1     0.0002  144314.335311  1.335425\n",
      "2     0.0003   96199.933994  0.894140\n",
      "3     0.0004   72142.733335  0.673530\n",
      "4     0.0005   57708.412940  0.541189\n",
      "...      ...            ...       ...\n",
      "9994  0.9995       0.014442  0.023060\n",
      "9995  0.9996       0.011552  0.023060\n",
      "9996  0.9997       0.008663  0.023060\n",
      "9997  0.9998       0.005775  0.023060\n",
      "9998  0.9999       0.002887  0.023060\n",
      "\n",
      "[9999 rows x 3 columns]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:75: PlotnineWarning: geom_path: Removed 5703 rows containing missing values.\n",
      "  warn(msg.format(n1-n2), PlotnineWarning)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29204 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (94729245802)>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plot_sensitivity_graph(imbens2, treatment_col='treatment', response_col='RE78', covariate_cols=covariates_for_graphs, covariate_groups=covariate_groups, bias=1000, do_kfold=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:05:57.487577Z",
     "start_time": "2019-10-12T20:05:57.480019Z"
    }
   },
   "outputs": [],
   "source": [
    "#plot_sensitivity_graph(imbens2, treatment_col='treatment', response_col='RE78', covariate_cols=covariates_for_graphs, covariate_groups=covariate_groups, bias=1000, do_kfold=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imbens 3 results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:06:59.850725Z",
     "start_time": "2019-10-12T20:05:57.487577Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                    covariates  alpha_hat   Rsq_hat                 groups\n",
      "0                        [age]   0.036161  0.000000  individual_covariates\n",
      "1                  [education]   0.055530  0.002154  individual_covariates\n",
      "2                      [black]   0.090657  0.000000  individual_covariates\n",
      "3                   [hispanic]   0.050661  0.000000  individual_covariates\n",
      "4                    [married]   0.126502  0.000000  individual_covariates\n",
      "5                       [RE74]   0.186596  0.035885  individual_covariates\n",
      "6                       [RE75]   0.243562  0.049587  individual_covariates\n",
      "7                      [pos74]   0.182654  0.000000  individual_covariates\n",
      "8                      [pos75]   0.100062  0.000000  individual_covariates\n",
      "9                [RE75, pos75]   0.217591  0.072572        recent_earnings\n",
      "10  [RE74, RE75, pos74, pos75]   0.318959  0.120793   pre_program_earnings\n",
      "      alphas         deltas       Rsq\n",
      "0     0.0001  288657.539263  2.585219\n",
      "1     0.0002  144314.335311  1.298170\n",
      "2     0.0003   96199.933994  0.869196\n",
      "3     0.0004   72142.733335  0.654740\n",
      "4     0.0005   57708.412940  0.526091\n",
      "...      ...            ...       ...\n",
      "9994  0.9995       0.014442  0.022417\n",
      "9995  0.9996       0.011552  0.022417\n",
      "9996  0.9997       0.008663  0.022417\n",
      "9997  0.9998       0.005775  0.022417\n",
      "9998  0.9999       0.002887  0.022417\n",
      "\n",
      "[9999 rows x 3 columns]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:75: PlotnineWarning: geom_path: Removed 5703 rows containing missing values.\n",
      "  warn(msg.format(n1-n2), PlotnineWarning)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29088 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (94730332813)>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plot_sensitivity_graph(imbens3, treatment_col='treatment', response_col='change_earnings', covariate_cols=covariates_for_graphs, covariate_groups=covariate_groups, bias=1000, do_kfold=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:06:59.858746Z",
     "start_time": "2019-10-12T20:06:59.850725Z"
    }
   },
   "outputs": [],
   "source": [
    "#plot_sensitivity_graph(imbens3, treatment_col='treatment', response_col='change_earnings', covariate_cols=covariates_for_graphs, covariate_groups=covariate_groups, bias=1000, do_kfold=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imbens4 results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:07:39.378314Z",
     "start_time": "2019-10-12T20:06:59.859752Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                    covariates  alpha_hat   Rsq_hat                 groups\n",
      "0                        [age]   0.189506  0.000000  individual_covariates\n",
      "1                  [education]   0.158524  0.058437  individual_covariates\n",
      "2                      [black]   0.207957  0.020464  individual_covariates\n",
      "3                   [hispanic]   0.029249  0.015469  individual_covariates\n",
      "4                    [married]   0.206164  0.026970  individual_covariates\n",
      "5                       [RE74]   0.005032  0.008125  individual_covariates\n",
      "6                       [RE75]   0.000000  0.029789  individual_covariates\n",
      "7                      [pos74]   0.000000  0.028513  individual_covariates\n",
      "8                      [pos75]   0.000000  0.027947  individual_covariates\n",
      "9                [RE75, pos75]   0.106004  0.017131        recent_earnings\n",
      "10  [RE74, RE75, pos74, pos75]   0.000000  0.018052   pre_program_earnings\n",
      "      alphas         deltas        Rsq\n",
      "0     0.0001  917484.249622  15.677525\n",
      "1     0.0002  458696.246011   7.851108\n",
      "2     0.0003  305766.911474   5.242409\n",
      "3     0.0004  229302.244205   3.938139\n",
      "4     0.0005  183423.443844   3.155641\n",
      "...      ...            ...        ...\n",
      "9994  0.9995       0.045902   0.052166\n",
      "9995  0.9996       0.036718   0.052166\n",
      "9996  0.9997       0.027536   0.052166\n",
      "9997  0.9998       0.018355   0.052166\n",
      "9998  0.9999       0.009177   0.052166\n",
      "\n",
      "[9999 rows x 3 columns]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:75: PlotnineWarning: geom_path: Removed 5726 rows containing missing values.\n",
      "  warn(msg.format(n1-n2), PlotnineWarning)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29724 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (-9223371942124466663)>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plot_sensitivity_graph(imbens4, treatment_col='treatment', response_col='RE78', covariate_cols=covariates_for_graphs, covariate_groups=covariate_groups, bias=1000, do_kfold=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:07:39.386315Z",
     "start_time": "2019-10-12T20:07:39.378314Z"
    }
   },
   "outputs": [],
   "source": [
    "#plot_sensitivity_graph(imbens4, treatment_col='treatment', response_col='RE78', covariate_cols=covariates_for_graphs, covariate_groups=covariate_groups, bias=1000, do_kfold=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-12T20:07:39.394464Z",
     "start_time": "2019-10-12T20:07:39.386315Z"
    }
   },
   "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.7.3"
  },
  "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": 2
}
