{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:20:46.674211Z",
     "start_time": "2019-09-05T05:20:43.134054Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scipy as sp\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\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-09-05T05:20:46.686376Z",
     "start_time": "2019-09-05T05:20:46.676206Z"
    }
   },
   "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-09-05T05:20:46.730266Z",
     "start_time": "2019-09-05T05:20:46.689373Z"
    }
   },
   "outputs": [],
   "source": [
    "def predict_g(X, y, do_kfold):\n",
    "    predictions=np.full_like(y, np.nan, dtype=float)\n",
    "    if do_kfold:        \n",
    "        kf=KFold(n_splits=3, shuffle=True, random_state=42)\n",
    "        for train_index, test_index in kf.split(X):\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-09-05T05:20:46.759190Z",
     "start_time": "2019-09-05T05:20:46.734253Z"
    }
   },
   "outputs": [],
   "source": [
    "def predict_Q(X,y,do_kfold):\n",
    "    predictions=np.full_like(y, np.nan, dtype=float)\n",
    "    if do_kfold:\n",
    "        kf=KFold(n_splits=3, 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-09-05T05:20:46.785117Z",
     "start_time": "2019-09-05T05:20:46.762178Z"
    }
   },
   "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-09-05T05:20:46.809054Z",
     "start_time": "2019-09-05T05:20:46.788109Z"
    }
   },
   "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-09-05T05:20:46.827007Z",
     "start_time": "2019-09-05T05:20:46.812046Z"
    }
   },
   "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",
    "    return Rsqhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:20:46.882856Z",
     "start_time": "2019-09-05T05:20:46.831992Z"
    }
   },
   "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",
    "    return ahat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:47:47.241496Z",
     "start_time": "2019-09-05T05:47:47.212574Z"
    }
   },
   "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",
    "    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['Rsq_hat']<0.05) & (variable_importances['alpha_hat']<0.2), 'labels']=''\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}$', fill='', y='$R^2$', title='kfold= '+ str(do_kfold))\n",
    "            +scale_fill_brewer(type='qual', palette='Set1')\n",
    "            +geom_label(data=variable_importances, mapping=aes(x='alpha_hat', y='Rsq_hat', label='labels'), size=8, nudge_y=0.03)\n",
    "            +annotate(\"text\", x=0.35, y=0.06,label='Bias = $'+str(bias), size=8.5, color='#303030')\n",
    "           )\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-06T02:15:31.433581Z",
     "start_time": "2019-09-06T02:15:31.325869Z"
    }
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'plotnine' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-2-cf78e05284b6>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[0mprint\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mplotnine\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m__version__\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m: name 'plotnine' is not defined"
     ]
    }
   ],
   "source": [
    "print (plotnine.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:20:46.952671Z",
     "start_time": "2019-09-05T05:20:46.933720Z"
    },
    "hidden": true
   },
   "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-09-05T05:20:47.094289Z",
     "start_time": "2019-09-05T05:20:46.956658Z"
    },
    "hidden": true
   },
   "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-09-05T05:20:47.107255Z",
     "start_time": "2019-09-05T05:20:47.097284Z"
    },
    "hidden": true
   },
   "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-09-05T05:20:47.161112Z",
     "start_time": "2019-09-05T05:20:47.110247Z"
    },
    "hidden": 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>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",
       "      <th>0</th>\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",
       "      <th>1</th>\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",
       "      <th>2</th>\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",
       "      <th>3</th>\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",
       "      <th>4</th>\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-09-05T05:20:47.212972Z",
     "start_time": "2019-09-05T05:20:47.165101Z"
    },
    "hidden": 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>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",
       "      <th>0</th>\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",
       "      <th>1</th>\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",
       "      <th>2</th>\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",
       "      <th>3</th>\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",
       "      <th>4</th>\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-09-05T05:20:47.224941Z",
     "start_time": "2019-09-05T05:20:47.216962Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "imbens2=pd.concat([treatment, psid]).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:20:47.297746Z",
     "start_time": "2019-09-05T05:20:47.229928Z"
    },
    "hidden": 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>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",
       "      <th>0</th>\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",
       "      <th>1</th>\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",
       "      <th>2</th>\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",
       "      <th>3</th>\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",
       "      <th>4</th>\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-09-05T05:20:47.383517Z",
     "start_time": "2019-09-05T05:20:47.301736Z"
    },
    "hidden": 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>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",
       "      <th>0</th>\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",
       "      <th>1</th>\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",
       "      <th>2</th>\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",
       "      <th>3</th>\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",
       "      <th>4</th>\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-09-05T05:20:47.427399Z",
     "start_time": "2019-09-05T05:20:47.390497Z"
    },
    "hidden": true,
    "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",
       "      <th>0.0</th>\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",
       "      <th>1.0</th>\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-09-05T05:20:47.484247Z",
     "start_time": "2019-09-05T05:20:47.431388Z"
    },
    "hidden": 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>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",
       "      <th>0</th>\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",
       "      <th>1</th>\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",
       "      <th>2</th>\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",
       "      <th>3</th>\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",
       "      <th>4</th>\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": 46,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:34:54.712776Z",
     "start_time": "2019-09-05T05:34:54.707789Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "covariate_cols=['age', 'education', 'black', 'hispanic', 'married', 'RE75', 'pos74', 'pos75']\n",
    "recent_earnings=['RE74', 'RE75']\n",
    "#what about \"pre-program earnings\"? Is that RE74?\n",
    "pre_program_earnings=['RE74']\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": {
    "heading_collapsed": true
   },
   "source": [
    "## Imbens1 results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:24:26.360608Z",
     "start_time": "2019-09-05T05:24:21.267234Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     covariates  alpha_hat   Rsq_hat                groups\n",
      "0         [age]   0.009190 -0.032642                   age\n",
      "1   [education]   0.005096 -0.020045             education\n",
      "2       [black]   0.001434 -0.011529                 black\n",
      "3    [hispanic]   0.003582 -0.009316              hispanic\n",
      "4     [married]  -0.002287 -0.013836               married\n",
      "5        [RE75]   0.008251  0.002645                  RE75\n",
      "6       [pos74]   0.002922 -0.002045                 pos74\n",
      "7       [pos75]   0.004396 -0.005274                 pos75\n",
      "8  [RE74, RE75]  -0.001743  0.016176       recent_earnings\n",
      "9        [RE74]   0.001291 -0.010482  pre_program_earnings\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\mizani\\palettes.py:434: UserWarning: Warning message:Brewer palette Set1 has a maximum of 9 colorsReturning the palette you asked for withthat many colors\n",
      "  warnings.warn(msg)\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:74: UserWarning: geom_path: Removed 577 rows containing missing values.\n",
      "  warn(msg.format(n1-n2))\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\layer.py:449: UserWarning: geom_point : Removed 2 rows containing missing values.\n",
      "  self.data = self.geom.handle_na(self.data)\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\layer.py:449: UserWarning: geom_label : Removed 2 rows containing missing values.\n",
      "  self.data = self.geom.handle_na(self.data)\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\mizani\\palettes.py:434: UserWarning: Warning message:Brewer palette Set1 has a maximum of 9 colorsReturning the palette you asked for withthat many colors\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35480 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (147029745423)>"
      ]
     },
     "execution_count": 30,
     "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-09-05T05:21:01.579483Z",
     "start_time": "2019-09-05T05:20:53.779983Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     covariates  alpha_hat   Rsq_hat                groups\n",
      "0         [age]  -0.000111  0.007301                   age\n",
      "1   [education]   0.003554  0.021058             education\n",
      "2       [black]   0.001859  0.001030                 black\n",
      "3    [hispanic]   0.004178 -0.002678              hispanic\n",
      "4     [married]   0.006078 -0.000805               married\n",
      "5        [RE75]   0.011156  0.023402                  RE75\n",
      "6       [pos74]   0.006789  0.001314                 pos74\n",
      "7       [pos75]   0.007476  0.001314                 pos75\n",
      "8  [RE74, RE75]   0.004947  0.037792       recent_earnings\n",
      "9        [RE74]   0.009500  0.009418  pre_program_earnings\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\mizani\\palettes.py:434: UserWarning: Warning message:Brewer palette Set1 has a maximum of 9 colorsReturning the palette you asked for withthat many colors\n",
      "  warnings.warn(msg)\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:74: UserWarning: geom_path: Removed 580 rows containing missing values.\n",
      "  warn(msg.format(n1-n2))\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\mizani\\palettes.py:434: UserWarning: Warning message:Brewer palette Set1 has a maximum of 9 colorsReturning the palette you asked for withthat many colors\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34680 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (-9223371889825205097)>"
      ]
     },
     "execution_count": 22,
     "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=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imbens2 results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:46:13.478645Z",
     "start_time": "2019-09-05T05:46:06.958578Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     covariates  alpha_hat   Rsq_hat                 groups\n",
      "0         [age]   0.049994 -0.000348  individual_covariates\n",
      "1   [education]   0.093793  0.002167  individual_covariates\n",
      "2       [black]   0.089747 -0.000026  individual_covariates\n",
      "3    [hispanic]   0.073269  0.000000  individual_covariates\n",
      "4     [married]   0.171808  0.000000  individual_covariates\n",
      "5        [RE75]   0.239140  0.084090  individual_covariates\n",
      "6       [pos74]   0.146788  0.000000  individual_covariates\n",
      "7       [pos75]   0.063540  0.000000  individual_covariates\n",
      "8  [RE74, RE75]   0.342459  0.357890        recent_earnings\n",
      "9        [RE74]   0.197068  0.034693   pre_program_earnings\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:74: UserWarning: geom_path: Removed 569 rows containing missing values.\n",
      "  warn(msg.format(n1-n2))\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26480 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (-9223371889821218707)>"
      ]
     },
     "execution_count": 67,
     "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-09-05T05:21:16.672108Z",
     "start_time": "2019-09-05T05:21:09.744641Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     covariates  alpha_hat       Rsq_hat                groups\n",
      "0         [age]   0.023654 -1.465572e-16                   age\n",
      "1   [education]   0.038416 -1.465572e-16             education\n",
      "2       [black]   0.118367  0.000000e+00                 black\n",
      "3    [hispanic]   0.051676 -1.465572e-16              hispanic\n",
      "4     [married]   0.102664  0.000000e+00               married\n",
      "5        [RE75]   0.222615  1.069351e-01                  RE75\n",
      "6       [pos74]   0.122983 -1.465572e-16                 pos74\n",
      "7       [pos75]   0.050472 -1.465572e-16                 pos75\n",
      "8  [RE74, RE75]   0.268689  3.892966e-01       recent_earnings\n",
      "9        [RE74]   0.182897  3.187298e-02  pre_program_earnings\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\mizani\\palettes.py:434: UserWarning: Warning message:Brewer palette Set1 has a maximum of 9 colorsReturning the palette you asked for withthat many colors\n",
      "  warnings.warn(msg)\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:74: UserWarning: geom_path: Removed 569 rows containing missing values.\n",
      "  warn(msg.format(n1-n2))\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\mizani\\palettes.py:434: UserWarning: Warning message:Brewer palette Set1 has a maximum of 9 colorsReturning the palette you asked for withthat many colors\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32956 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (147029683916)>"
      ]
     },
     "execution_count": 24,
     "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=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imbens 3 results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:46:33.453659Z",
     "start_time": "2019-09-05T05:46:26.847420Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     covariates  alpha_hat   Rsq_hat                 groups\n",
      "0         [age]   0.049994  0.000094  individual_covariates\n",
      "1   [education]   0.093793  0.012082  individual_covariates\n",
      "2       [black]   0.089747  0.006245  individual_covariates\n",
      "3    [hispanic]   0.073269  0.005789  individual_covariates\n",
      "4     [married]   0.171808  0.006218  individual_covariates\n",
      "5        [RE75]   0.239140  0.041961  individual_covariates\n",
      "6       [pos74]   0.146788  0.007062  individual_covariates\n",
      "7       [pos75]   0.063540  0.009703  individual_covariates\n",
      "8  [RE74, RE75]   0.342459  0.042609        recent_earnings\n",
      "9        [RE74]   0.197068  0.053531   pre_program_earnings\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:74: UserWarning: geom_path: Removed 569 rows containing missing values.\n",
      "  warn(msg.format(n1-n2))\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26444 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (-9223371889820991705)>"
      ]
     },
     "execution_count": 68,
     "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": 69,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:47:02.683622Z",
     "start_time": "2019-09-05T05:46:54.795681Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     covariates  alpha_hat       Rsq_hat                 groups\n",
      "0         [age]   0.023654  1.331012e-03  individual_covariates\n",
      "1   [education]   0.038416  5.140678e-03  individual_covariates\n",
      "2       [black]   0.118367  2.417258e-05  individual_covariates\n",
      "3    [hispanic]   0.051676  1.159721e-03  individual_covariates\n",
      "4     [married]   0.102664  3.194241e-04  individual_covariates\n",
      "5        [RE75]   0.222615  6.885221e-02  individual_covariates\n",
      "6       [pos74]   0.122983  2.417258e-05  individual_covariates\n",
      "7       [pos75]   0.050472  1.428739e-16  individual_covariates\n",
      "8  [RE74, RE75]   0.268689  9.934257e-02        recent_earnings\n",
      "9        [RE74]   0.182897  6.373313e-02   pre_program_earnings\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:74: UserWarning: geom_path: Removed 569 rows containing missing values.\n",
      "  warn(msg.format(n1-n2))\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27408 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (-9223371889821002787)>"
      ]
     },
     "execution_count": 69,
     "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=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imbens4 results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:47:17.691084Z",
     "start_time": "2019-09-05T05:47:13.070322Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     covariates  alpha_hat   Rsq_hat                 groups\n",
      "0         [age]   0.205597  0.006522  individual_covariates\n",
      "1   [education]   0.175836  0.058414  individual_covariates\n",
      "2       [black]   0.235781  0.003151  individual_covariates\n",
      "3    [hispanic]   0.019743  0.011947  individual_covariates\n",
      "4     [married]   0.161594 -0.002795  individual_covariates\n",
      "5        [RE75]  -0.036851  0.003848  individual_covariates\n",
      "6       [pos74]  -0.009016 -0.001610  individual_covariates\n",
      "7       [pos75]  -0.040660  0.001604  individual_covariates\n",
      "8  [RE74, RE75]   0.040347 -0.011857        recent_earnings\n",
      "9        [RE74]  -0.020837  0.003069   pre_program_earnings\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:74: UserWarning: geom_path: Removed 571 rows containing missing values.\n",
      "  warn(msg.format(n1-n2))\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\layer.py:449: UserWarning: geom_point : Removed 3 rows containing missing values.\n",
      "  self.data = self.geom.handle_na(self.data)\n",
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\layer.py:449: UserWarning: geom_text : Removed 3 rows containing missing values.\n",
      "  self.data = self.geom.handle_na(self.data)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28512 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (147033753057)>"
      ]
     },
     "execution_count": 70,
     "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": 72,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-05T05:48:03.008387Z",
     "start_time": "2019-09-05T05:47:58.556783Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     covariates  alpha_hat   Rsq_hat                 groups\n",
      "0         [age]   0.243647  0.094917  individual_covariates\n",
      "1   [education]   0.210517  0.146431  individual_covariates\n",
      "2       [black]   0.259567  0.000211  individual_covariates\n",
      "3    [hispanic]   0.052713 -0.000157  individual_covariates\n",
      "4     [married]   0.214431  0.003137  individual_covariates\n",
      "5        [RE75]   0.000613  0.008552  individual_covariates\n",
      "6       [pos74]   0.018241  0.002981  individual_covariates\n",
      "7       [pos75]   0.006299 -0.000157  individual_covariates\n",
      "8  [RE74, RE75]   0.114825  0.006138        recent_earnings\n",
      "9        [RE74]   0.044227 -0.003283   pre_program_earnings\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\light\\Anaconda3\\lib\\site-packages\\plotnine\\geoms\\geom_path.py:74: UserWarning: geom_path: Removed 572 rows containing missing values.\n",
      "  warn(msg.format(n1-n2))\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31988 characters]\n",
      "text/plain": [
       "<Figure size 350x350 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<ggplot: (147034950285)>"
      ]
     },
     "execution_count": 72,
     "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=False)"
   ]
  }
 ],
 "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
}
