doc/examples/R_double_ml_pension.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cooperative-usage",
   "metadata": {
    "papermill": {
     "duration": 0.040312,
     "end_time": "2021-03-23T13:44:30.492247",
     "exception": false,
     "start_time": "2021-03-23T13:44:30.451935",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# R: Impact of 401(k) on Financial Wealth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "identified-cruise",
   "metadata": {
    "papermill": {
     "duration": 0.036965,
     "end_time": "2021-03-23T13:44:30.565334",
     "exception": false,
     "start_time": "2021-03-23T13:44:30.528369",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "In this real-data example, we illustrate how the [DoubleML](https://docs.doubleml.org/stable/index.html) package can be used to estimate the effect of 401(k) eligibility and participation on accumulated assets. The 401(k) data set has been analyzed in several studies, among others [Chernozhukov et al. (2018)](https://arxiv.org/abs/1608.00060).\n",
    "\n",
    "401(k) plans are pension accounts sponsored by employers. The key problem in determining the effect of participation in 401(k) plans on accumulated assets is saver heterogeneity coupled with the fact that the decision to enroll in a 401(k) is non-random. It is generally recognized that some people have a higher preference for saving than others. It also seems likely that those individuals with high unobserved preference for saving would be most likely to choose to participate in tax-advantaged retirement savings plans and would tend to have otherwise high amounts of accumulated assets. The presence of unobserved savings preferences with these properties then implies that conventional estimates that do not account for saver heterogeneity and endogeneity of participation will be biased upward, tending to overstate the savings effects of 401(k) participation.\n",
    "\n",
    "One can argue that eligibility for enrolling in a 401(k) plan in this data can be taken as exogenous after conditioning on a few observables of which the most important for their argument is income. The basic idea is that, at least around the time 401(k)’s initially became available, people were unlikely to be basing their employment decisions on whether an employer offered a 401(k) but would instead focus on income and other aspects of the job. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adequate-trunk",
   "metadata": {
    "papermill": {
     "duration": 0.035502,
     "end_time": "2021-03-23T13:44:30.636343",
     "exception": false,
     "start_time": "2021-03-23T13:44:30.600841",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Data\n",
    "\n",
    "The preprocessed data can be fetched by calling [fetch_401k()](https://docs.doubleml.org/r/stable/reference/fetch_401k.html). The arguments `polynomial_features` and `instrument` can be used to replicate the models used in [Chernozhukov et al. (2018)](https://arxiv.org/abs/1608.00060). Note that an internet connection is required for loading the data. We start with a baseline specification of the regression model and reload the data later in case we want to use another specification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "underlying-lemon",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:44:30.715038Z",
     "iopub.status.busy": "2021-03-23T13:44:30.712600Z",
     "iopub.status.idle": "2021-03-23T13:44:31.163126Z",
     "shell.execute_reply": "2021-03-23T13:44:31.161559Z"
    },
    "papermill": {
     "duration": 0.491422,
     "end_time": "2021-03-23T13:44:31.163324",
     "exception": false,
     "start_time": "2021-03-23T13:44:30.671902",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load required packages for this tutorial\n",
    "library(DoubleML)\n",
    "library(mlr3)\n",
    "library(mlr3learners)\n",
    "library(data.table)\n",
    "library(ggplot2)\n",
    "\n",
    "# suppress messages during fitting\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "\n",
    "# load data as a data.table\n",
    "data = fetch_401k(return_type = \"data.table\", instrument = TRUE)\n",
    "dim(data)\n",
    "str(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sophisticated-forum",
   "metadata": {
    "papermill": {
     "duration": 0.036035,
     "end_time": "2021-03-23T13:44:31.235882",
     "exception": false,
     "start_time": "2021-03-23T13:44:31.199847",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "See the \"Details\" section on the description of the data set, which can be accessed by typing [help(fetch_401k)](https://docs.doubleml.org/r/stable/reference/fetch_401k.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "expired-problem",
   "metadata": {
    "papermill": {
     "duration": 0.036944,
     "end_time": "2021-03-23T13:44:31.669355",
     "exception": false,
     "start_time": "2021-03-23T13:44:31.632411",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The data consist of 9,915 observations at the household level drawn from the 1991 Survey of Income and Program Participation (SIPP).  All the variables are referred to 1990. We use net financial assets (*net\\_tfa*) as the outcome variable, $Y$,  in our analysis. The net financial assets are computed as the sum of IRA balances, 401(k) balances, checking accounts, saving bonds, other interest-earning accounts, other interest-earning assets, stocks, and mutual funds less non mortgage debts. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "circular-electric",
   "metadata": {
    "papermill": {
     "duration": 0.035954,
     "end_time": "2021-03-23T13:44:31.741375",
     "exception": false,
     "start_time": "2021-03-23T13:44:31.705421",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Among the $9915$ individuals, $3682$ are eligible to participate in the program. The variable *e401* indicates eligibility and *p401* indicates participation, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "latin-copying",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:44:31.820240Z",
     "iopub.status.busy": "2021-03-23T13:44:31.818777Z",
     "iopub.status.idle": "2021-03-23T13:44:32.849032Z",
     "shell.execute_reply": "2021-03-23T13:44:32.847733Z"
    },
    "papermill": {
     "duration": 1.07102,
     "end_time": "2021-03-23T13:44:32.849241",
     "exception": false,
     "start_time": "2021-03-23T13:44:31.778221",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "hist_e401 = ggplot(data, aes(x = e401, fill = factor(e401))) +\n",
    "            geom_bar() + theme_minimal() + \n",
    "            ggtitle(\"Eligibility, 401(k)\") +\n",
    "            theme(legend.position = \"bottom\", plot.title = element_text(hjust = 0.5),\n",
    "                  text = element_text(size = 20)) \n",
    "hist_e401"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "lesbian-blood",
   "metadata": {},
   "outputs": [],
   "source": [
    "hist_p401 = ggplot(data, aes(x = p401, fill = factor(p401))) +\n",
    "            geom_bar() + theme_minimal() + \n",
    "            ggtitle(\"Participation, 401(k)\") +\n",
    "            theme(legend.position = \"bottom\", plot.title = element_text(hjust = 0.5),\n",
    "                  text = element_text(size = 20)) \n",
    "hist_p401"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "streaming-flood",
   "metadata": {
    "papermill": {
     "duration": 0.039347,
     "end_time": "2021-03-23T13:44:32.932448",
     "exception": false,
     "start_time": "2021-03-23T13:44:32.893101",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Eligibility is highly associated with financial wealth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "focal-boring",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:44:33.016258Z",
     "iopub.status.busy": "2021-03-23T13:44:33.013504Z",
     "iopub.status.idle": "2021-03-23T13:44:33.681630Z",
     "shell.execute_reply": "2021-03-23T13:44:33.681007Z"
    },
    "papermill": {
     "duration": 0.711317,
     "end_time": "2021-03-23T13:44:33.681783",
     "exception": false,
     "start_time": "2021-03-23T13:44:32.970466",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dens_net_tfa = ggplot(data, aes(x = net_tfa, color = factor(e401), fill = factor(e401)) ) + \n",
    "                    geom_density() + xlim(c(-20000, 150000)) + \n",
    "                    facet_wrap(.~e401)  + theme_minimal() + \n",
    "                    theme(legend.position = \"bottom\", text = element_text(size = 20))\n",
    "                    \n",
    "dens_net_tfa"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "third-turkey",
   "metadata": {
    "papermill": {
     "duration": 0.038905,
     "end_time": "2021-03-23T13:44:33.759588",
     "exception": false,
     "start_time": "2021-03-23T13:44:33.720683",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "As a first estimate, we calculate the unconditional average predictive effect (APE) of 401(k) eligibility on accumulated assets. This effect corresponds to the average treatment effect if 401(k) eligibility would be assigned to individuals in an entirely randomized way. The unconditional APE of e401 is about $19559$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cultural-netherlands",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:44:33.843625Z",
     "iopub.status.busy": "2021-03-23T13:44:33.842058Z",
     "iopub.status.idle": "2021-03-23T13:44:33.866334Z",
     "shell.execute_reply": "2021-03-23T13:44:33.864827Z"
    },
    "papermill": {
     "duration": 0.067968,
     "end_time": "2021-03-23T13:44:33.866504",
     "exception": false,
     "start_time": "2021-03-23T13:44:33.798536",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "APE_e401_uncond = data[e401==1, mean(net_tfa)] - data[e401==0, mean(net_tfa)]\n",
    "round(APE_e401_uncond, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "intelligent-threshold",
   "metadata": {
    "papermill": {
     "duration": 0.040553,
     "end_time": "2021-03-23T13:44:33.950269",
     "exception": false,
     "start_time": "2021-03-23T13:44:33.909716",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Among the $3682$ individuals that  are eligible, $2594$ decided to participate in the program. The unconditional APE of p401 is about $27372$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "objective-spokesman",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:44:34.036419Z",
     "iopub.status.busy": "2021-03-23T13:44:34.034858Z",
     "iopub.status.idle": "2021-03-23T13:44:34.062248Z",
     "shell.execute_reply": "2021-03-23T13:44:34.060319Z"
    },
    "papermill": {
     "duration": 0.07273,
     "end_time": "2021-03-23T13:44:34.062417",
     "exception": false,
     "start_time": "2021-03-23T13:44:33.989687",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "APE_p401_uncond = data[p401==1, mean(net_tfa)] - data[p401==0, mean(net_tfa)]\n",
    "round(APE_p401_uncond, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "charitable-mailing",
   "metadata": {
    "papermill": {
     "duration": 0.040806,
     "end_time": "2021-03-23T13:44:34.144049",
     "exception": false,
     "start_time": "2021-03-23T13:44:34.103243",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "As discussed, these estimates are biased since they do not account for saver heterogeneity and endogeneity of participation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "integral-surgery",
   "metadata": {
    "papermill": {
     "duration": 0.040491,
     "end_time": "2021-03-23T13:44:34.224987",
     "exception": false,
     "start_time": "2021-03-23T13:44:34.184496",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## The `DoubleML` package"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "religious-cable",
   "metadata": {
    "papermill": {
     "duration": 0.041623,
     "end_time": "2021-03-23T13:44:34.307306",
     "exception": false,
     "start_time": "2021-03-23T13:44:34.265683",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's use the package [DoubleML](https://docs.doubleml.org/stable/index.html) to estimate the average treatment effect of 401(k) eligibility, i.e. `e401`, and participation, i.e. `p401`, on net financial assets `net_tfa`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "classified-lodging",
   "metadata": {
    "papermill": {
     "duration": 0.042605,
     "end_time": "2021-03-23T13:47:34.483687",
     "exception": false,
     "start_time": "2021-03-23T13:47:34.441082",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Estimating the Average Treatment Effect of 401(k) Eligibility on Net Financial Assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "french-pizza",
   "metadata": {
    "papermill": {
     "duration": 0.043776,
     "end_time": "2021-03-23T13:47:34.573553",
     "exception": false,
     "start_time": "2021-03-23T13:47:34.529777",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We first look at the treatment effect of `e401` on net total financial assets. We give estimates of the ATE in the linear model\n",
    "\n",
    "\\begin{equation*}\n",
    "Y = D \\alpha + f(X)'\\beta+ \\epsilon,\n",
    "\\end{equation*}\n",
    "where $f(X)$ is a dictonary applied to the raw regressors. $X$ contains variables on marital status, two-earner status, defined benefit pension status, IRA participation, home ownership, family size, education, age, and income. \n",
    "\n",
    "In the following, we will consider two different models, \n",
    "\n",
    "* a basic model specification that includes the raw regressors, i.e., $f(X) = X$, and \n",
    "\n",
    "* a flexible model specification, where $f(X)$ includes the raw regressors $X$ and the orthogonal polynomials of degree 2 for the variables family size education, age, and income. \n",
    "\n",
    "We will use the basic model specification whenever we use nonlinear methods, for example regression trees or random forests, and use the flexible model for linear methods such as the lasso. There are, of course, multiple ways how the model can be specified even more flexibly, for example including interactions of variable and higher order interaction. However, for the sake of simplicity we stick to the specification above. Users who are interested in varying the model can manipulate the model formula below (`formula_flex`), for example implementing the orignal specification in Chernozhukov et al. (2018). \n",
    "\n",
    "In the first step, we report estimates of the average treatment effect (ATE) of 401(k) eligibility on net financial assets both in the partially linear regression (PLR) model and in the interactive regression model (IRM) allowing for heterogeneous treatment effects. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "living-stopping",
   "metadata": {},
   "source": [
    "### The Data Backend: `DoubleMLData`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cosmetic-emergency",
   "metadata": {},
   "source": [
    "To start our analysis, we initialize the data backend, i.e., a new instance of a [DoubleMLData](https://docs.doubleml.org/r/stable/reference/DoubleMLData.html) object. Here, we manually implement the regression model by using R's formula interface. A shortcut would be to directly specify the options `polynomial_features` and `instrument` when calling [fetch_401k()](https://docs.doubleml.org/r/stable/reference/fetch_401k.html).$^{**}$\n",
    "\n",
    "To implement both models (basic and flexible), we generate two data backends: `data_dml_base` and `data_dml_flex`.\n",
    "\n",
    "$^{**}$ Note that the model specification using `polynomial_features` differs from the one used in our example. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "apparent-israeli",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:47:34.664674Z",
     "iopub.status.busy": "2021-03-23T13:47:34.662690Z",
     "iopub.status.idle": "2021-03-23T13:47:34.941823Z",
     "shell.execute_reply": "2021-03-23T13:47:34.939905Z"
    },
    "papermill": {
     "duration": 0.326212,
     "end_time": "2021-03-23T13:47:34.941985",
     "exception": false,
     "start_time": "2021-03-23T13:47:34.615773",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set up basic model: Specify variables for data-backend\n",
    "features_base = c(\"age\", \"inc\", \"educ\", \"fsize\",\n",
    "                  \"marr\", \"twoearn\", \"db\", \"pira\", \"hown\")\n",
    "\n",
    "# Initialize DoubleMLData (data-backend of DoubleML)\n",
    "data_dml_base = DoubleMLData$new(data,\n",
    "                                 y_col = \"net_tfa\",\n",
    "                                 d_cols = \"e401\",\n",
    "                                 x_cols = features_base)\n",
    "data_dml_base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "narrative-senior",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:47:34.664674Z",
     "iopub.status.busy": "2021-03-23T13:47:34.662690Z",
     "iopub.status.idle": "2021-03-23T13:47:34.941823Z",
     "shell.execute_reply": "2021-03-23T13:47:34.939905Z"
    },
    "papermill": {
     "duration": 0.326212,
     "end_time": "2021-03-23T13:47:34.941985",
     "exception": false,
     "start_time": "2021-03-23T13:47:34.615773",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set up a model according to regression formula with polynomials\n",
    "formula_flex = formula(\" ~ -1 + poly(age, 2, raw=TRUE) +\n",
    "                        poly(inc, 2, raw=TRUE) + poly(educ, 2, raw=TRUE) +\n",
    "                        poly(fsize, 2, raw=TRUE) + marr + twoearn +\n",
    "                        db + pira + hown\")\n",
    "features_flex = data.frame(model.matrix(formula_flex, data))\n",
    "\n",
    "model_data = data.table(\"net_tfa\" = data[, net_tfa],\n",
    "                        \"e401\" = data[, e401],\n",
    "                        features_flex)\n",
    "\n",
    "# Initialize DoubleMLData (data-backend of DoubleML)\n",
    "data_dml_flex = DoubleMLData$new(model_data,\n",
    "                                 y_col = \"net_tfa\",\n",
    "                                 d_cols = \"e401\")\n",
    "\n",
    "data_dml_flex"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "strong-digest",
   "metadata": {
    "papermill": {
     "duration": 0.044143,
     "end_time": "2021-03-23T13:47:35.029558",
     "exception": false,
     "start_time": "2021-03-23T13:47:34.985415",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Partially Linear Regression Model (PLR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "polar-personal",
   "metadata": {
    "papermill": {
     "duration": 0.043027,
     "end_time": "2021-03-23T13:47:35.118878",
     "exception": false,
     "start_time": "2021-03-23T13:47:35.075851",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We start using lasso to estimate the function $g_0$ and $m_0$ in the following PLR model:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "upper-surfing",
   "metadata": {
    "papermill": {
     "duration": 0.044546,
     "end_time": "2021-03-23T13:47:35.207536",
     "exception": false,
     "start_time": "2021-03-23T13:47:35.162990",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "\\begin{eqnarray}\n",
    "& Y = D\\theta_0 + g_0(X) + \\zeta, &\\quad E[\\zeta \\mid D,X]= 0,\\\\\n",
    "& D = m_0(X) +  V, &\\quad E[V \\mid X] = 0.\n",
    "\\end{eqnarray}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "diagnostic-regression",
   "metadata": {},
   "source": [
    "To estimate the causal parameter $\\theta_0$ here, we use double machine learning with 3-fold cross-fitting. \n",
    " \n",
    "Estimation of the nuisance components $g_0$ and $m_0$, is based on the lasso with cross-validated choice of the penalty term , $\\lambda$, as provided by the [glmnet package](https://glmnet.stanford.edu/reference/cv.glmnet.html). We load the learner by using the [mlr3](https://mlr3.mlr-org.com/) function [lrn()](https://mlr3.mlr-org.com/reference/Learner.html). Hyperparameters and options can be set during instantiation of the learner. Here we specify that the lasso should use that value of $\\lambda$ that minimizes the cross-validated mean squared error which is based on 5-fold cross validation.\n",
    "\n",
    "In order to use a learner, the underlying R packages have to be installed. In this case, the package [glmnet package](https://glmnet.stanford.edu/reference/cv.glmnet.html) needs to be installed. Moreover, installation of the package [mlr3learners](https://mlr3learners.mlr-org.com/) is required.\n",
    "\n",
    "We start by estimation the ATE in the basic model and then repeat the estimation in the flexible model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "serious-wrestling",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:47:35.302091Z",
     "iopub.status.busy": "2021-03-23T13:47:35.299667Z",
     "iopub.status.idle": "2021-03-23T13:47:58.447363Z",
     "shell.execute_reply": "2021-03-23T13:47:58.446093Z"
    },
    "papermill": {
     "duration": 23.196892,
     "end_time": "2021-03-23T13:47:58.447523",
     "exception": false,
     "start_time": "2021-03-23T13:47:35.250631",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Initialize learners\n",
    "set.seed(123)\n",
    "lasso = lrn(\"regr.cv_glmnet\", nfolds = 5, s = \"lambda.min\")\n",
    "lasso_class = lrn(\"classif.cv_glmnet\", nfolds = 5, s = \"lambda.min\")\n",
    "\n",
    "# Initialize DoubleMLPLR model\n",
    "dml_plr_lasso = DoubleMLPLR$new(data_dml_base, \n",
    "                                ml_g = lasso,\n",
    "                                ml_m = lasso_class,\n",
    "                                n_folds = 3)\n",
    "dml_plr_lasso$fit()\n",
    "dml_plr_lasso$summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "settled-appointment",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize learners\n",
    "set.seed(123)\n",
    "lasso = lrn(\"regr.cv_glmnet\", nfolds = 5, s = \"lambda.min\")\n",
    "lasso_class = lrn(\"classif.cv_glmnet\", nfolds = 5, s = \"lambda.min\")\n",
    "\n",
    "# Initialize DoubleMLPLR model\n",
    "dml_plr_lasso = DoubleMLPLR$new(data_dml_flex, \n",
    "                                ml_g = lasso,\n",
    "                                ml_m = lasso_class,\n",
    "                                n_folds = 3)\n",
    "dml_plr_lasso$fit()\n",
    "dml_plr_lasso$summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "intimate-angel",
   "metadata": {
    "papermill": {
     "duration": 0.046627,
     "end_time": "2021-03-23T13:47:58.987016",
     "exception": false,
     "start_time": "2021-03-23T13:47:58.940389",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Alternatively, we can repeat this procedure with other machine learning methods, for example a random forest learner as provided by the [ranger](https://github.com/imbs-hl/ranger) package for R. The website of the [mlr3extralearners](https://mlr3extralearners.mlr-org.com/articles/learners/list_learners.html) package has a searchable list of all learners that are available in the [mlr3verse](https://mlr3verse.mlr-org.com/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "handled-wholesale",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:47:59.087352Z",
     "iopub.status.busy": "2021-03-23T13:47:59.086372Z",
     "iopub.status.idle": "2021-03-23T13:49:02.075109Z",
     "shell.execute_reply": "2021-03-23T13:49:02.075650Z"
    },
    "papermill": {
     "duration": 63.042388,
     "end_time": "2021-03-23T13:49:02.075853",
     "exception": false,
     "start_time": "2021-03-23T13:47:59.033465",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Random Forest\n",
    "randomForest = lrn(\"regr.ranger\", max.depth = 7,\n",
    "                   mtry = 3, min.node.size = 3)\n",
    "randomForest_class = lrn(\"classif.ranger\", max.depth = 5,\n",
    "                         mtry = 4, min.node.size = 7)\n",
    "\n",
    "set.seed(123)\n",
    "dml_plr_forest = DoubleMLPLR$new(data_dml_base,\n",
    "                                 ml_g = randomForest,\n",
    "                                 ml_m = randomForest_class,\n",
    "                                 n_folds = 3)\n",
    "dml_plr_forest$fit() \n",
    "dml_plr_forest$summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "early-privacy",
   "metadata": {},
   "source": [
    "Now, let's use a regression tree as provided by the R package [rpart](https://github.com/bethatkinson/rpart)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "challenging-cloud",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:49:02.526135Z",
     "iopub.status.busy": "2021-03-23T13:49:02.524567Z",
     "iopub.status.idle": "2021-03-23T13:49:04.092172Z",
     "shell.execute_reply": "2021-03-23T13:49:04.090774Z"
    },
    "papermill": {
     "duration": 1.627688,
     "end_time": "2021-03-23T13:49:04.092395",
     "exception": false,
     "start_time": "2021-03-23T13:49:02.464707",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Trees\n",
    "trees = lrn(\"regr.rpart\", cp = 0.0047, minsplit = 203)\n",
    "trees_class = lrn(\"classif.rpart\", cp = 0.0042, minsplit = 104)\n",
    "\n",
    "set.seed(123)\n",
    "dml_plr_tree = DoubleMLPLR$new(data_dml_base,\n",
    "                               ml_g = trees,\n",
    "                               ml_m = trees_class,\n",
    "                               n_folds = 3)\n",
    "dml_plr_tree$fit()\n",
    "dml_plr_tree$summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "modified-player",
   "metadata": {},
   "source": [
    "We can also experiment with extreme gradient boosting as provided by [xgboost](https://xgboost.readthedocs.io/en/latest/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hollywood-heaven",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:49:04.206928Z",
     "iopub.status.busy": "2021-03-23T13:49:04.206179Z",
     "iopub.status.idle": "2021-03-23T13:49:05.194244Z",
     "shell.execute_reply": "2021-03-23T13:49:05.192417Z"
    },
    "papermill": {
     "duration": 1.048317,
     "end_time": "2021-03-23T13:49:05.194441",
     "exception": false,
     "start_time": "2021-03-23T13:49:04.146124",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Boosted trees\n",
    "boost = lrn(\"regr.xgboost\",\n",
    "            objective = \"reg:squarederror\",\n",
    "            eta = 0.1, nrounds = 35)\n",
    "boost_class = lrn(\"classif.xgboost\",\n",
    "                  objective = \"binary:logistic\", eval_metric = \"logloss\",\n",
    "                  eta = 0.1, nrounds = 34)\n",
    "\n",
    "set.seed(123)\n",
    "dml_plr_boost = DoubleMLPLR$new(data_dml_base,\n",
    "                                ml_g = boost,\n",
    "                                ml_m = boost_class,\n",
    "                                n_folds = 3)\n",
    "dml_plr_boost$fit()\n",
    "dml_plr_boost$summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "progressive-welsh",
   "metadata": {
    "papermill": {
     "duration": 0.055456,
     "end_time": "2021-03-23T13:49:05.305048",
     "exception": false,
     "start_time": "2021-03-23T13:49:05.249592",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's sum up the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "trained-least",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:49:05.421244Z",
     "iopub.status.busy": "2021-03-23T13:49:05.419710Z",
     "iopub.status.idle": "2021-03-23T13:49:05.487573Z",
     "shell.execute_reply": "2021-03-23T13:49:05.486132Z"
    },
    "papermill": {
     "duration": 0.127098,
     "end_time": "2021-03-23T13:49:05.487776",
     "exception": false,
     "start_time": "2021-03-23T13:49:05.360678",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "confints = rbind(dml_plr_lasso$confint(), dml_plr_forest$confint(),\n",
    "                 dml_plr_tree$confint(), dml_plr_boost$confint())\n",
    "estimates = c(dml_plr_lasso$coef, dml_plr_forest$coef,\n",
    "              dml_plr_tree$coef, dml_plr_boost$coef)\n",
    "result_plr = data.table(\"model\" = \"PLR\", \n",
    "                        \"ML\" = c(\"glmnet\", \"ranger\", \"rpart\", \"xgboost\"), \n",
    "                        \"Estimate\" = estimates,\n",
    "                        \"lower\" = confints[,1],\n",
    "                        \"upper\" = confints[,2])\n",
    "result_plr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tough-funeral",
   "metadata": {},
   "outputs": [],
   "source": [
    "g_ci = ggplot(result_plr, aes(x = ML, y = Estimate, color = ML)) + \n",
    "        geom_point() +\n",
    "        geom_errorbar(aes(ymin = lower, ymax = upper, color = ML))  +\n",
    "        geom_hline(yintercept = 0, color = \"grey\") +\n",
    "        theme_minimal() + ylab(\"Coefficients and 0.95- confidence interval\") + \n",
    "        xlab(\"\") + \n",
    "        theme(axis.text.x = element_text(angle = 90), legend.position = \"none\",\n",
    "              text = element_text(size = 20))\n",
    "\n",
    "g_ci"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "patient-straight",
   "metadata": {
    "papermill": {
     "duration": 0.056069,
     "end_time": "2021-03-23T13:49:05.604569",
     "exception": false,
     "start_time": "2021-03-23T13:49:05.548500",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Interactive Regression Model (IRM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "immune-investor",
   "metadata": {
    "papermill": {
     "duration": 0.055247,
     "end_time": "2021-03-23T13:49:05.714403",
     "exception": false,
     "start_time": "2021-03-23T13:49:05.659156",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Next, we consider estimation of average treatment effects when treatment effects are fully heterogeneous:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "basic-architecture",
   "metadata": {
    "papermill": {
     "duration": 0.057453,
     "end_time": "2021-03-23T13:49:05.828088",
     "exception": false,
     "start_time": "2021-03-23T13:49:05.770635",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "\\begin{eqnarray}\n",
    "& Y = g_0(D,X) + U, &\\quad E[U\\mid X,D] = 0,\\\\\n",
    "& D = m_0(X) + V, &\\quad E[V\\mid X] = 0.\n",
    "\\end{eqnarray}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "separate-information",
   "metadata": {
    "papermill": {
     "duration": 0.055141,
     "end_time": "2021-03-23T13:49:05.940426",
     "exception": false,
     "start_time": "2021-03-23T13:49:05.885285",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "To reduce the disproportionate impact of extreme propensity score weights in the interactive model\n",
    "we trim the propensity scores which are close to the bounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "monetary-shepherd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:49:06.055682Z",
     "iopub.status.busy": "2021-03-23T13:49:06.054164Z",
     "iopub.status.idle": "2021-03-23T13:49:24.871519Z",
     "shell.execute_reply": "2021-03-23T13:49:24.869890Z"
    },
    "papermill": {
     "duration": 18.876848,
     "end_time": "2021-03-23T13:49:24.871740",
     "exception": false,
     "start_time": "2021-03-23T13:49:05.994892",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "set.seed(123)\n",
    "# Initialize DoubleMLIRM model\n",
    "dml_irm_lasso = DoubleMLIRM$new(data_dml_flex,\n",
    "                                ml_g = lasso,\n",
    "                                ml_m = lasso_class,\n",
    "                                trimming_threshold = 0.01,\n",
    "                                n_folds = 3)\n",
    "dml_irm_lasso$fit()\n",
    "dml_irm_lasso$summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "attractive-space",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:49:25.191466Z",
     "iopub.status.busy": "2021-03-23T13:49:25.189452Z",
     "iopub.status.idle": "2021-03-23T13:50:24.952570Z",
     "shell.execute_reply": "2021-03-23T13:50:24.951094Z"
    },
    "papermill": {
     "duration": 59.82937,
     "end_time": "2021-03-23T13:50:24.953063",
     "exception": false,
     "start_time": "2021-03-23T13:49:25.123693",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Initialize Learner\n",
    "randomForest = lrn(\"regr.ranger\")\n",
    "randomForest_class = lrn(\"classif.ranger\")\n",
    "\n",
    "# Random Forest\n",
    "set.seed(123)\n",
    "dml_irm_forest = DoubleMLIRM$new(data_dml_base,\n",
    "                                 ml_g = randomForest,\n",
    "                                 ml_m = randomForest_class,\n",
    "                                 trimming_threshold = 0.01,\n",
    "                                 n_folds = 3)\n",
    "\n",
    "# Set nuisance-part specific parameters\n",
    "dml_irm_forest$set_ml_nuisance_params(\n",
    "    \"ml_g0\", \"e401\", list(max.depth = 6, mtry = 4, min.node.size = 7))\n",
    "dml_irm_forest$set_ml_nuisance_params(\n",
    "    \"ml_g1\", \"e401\", list(max.depth = 6, mtry = 3, min.node.size = 5))\n",
    "dml_irm_forest$set_ml_nuisance_params(\n",
    "    \"ml_m\", \"e401\", list(max.depth = 6, mtry = 3, min.node.size = 6))\n",
    "\n",
    "dml_irm_forest$fit()\n",
    "dml_irm_forest$summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "inappropriate-shirt",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Learner\n",
    "trees = lrn(\"regr.rpart\")\n",
    "trees_class = lrn(\"classif.rpart\")\n",
    "\n",
    "# Trees\n",
    "set.seed(123)\n",
    "dml_irm_tree = DoubleMLIRM$new(data_dml_base,\n",
    "                          ml_g = trees,\n",
    "                          ml_m = trees_class,\n",
    "                          trimming_threshold = 0.01,\n",
    "                          n_folds = 3)\n",
    "\n",
    "# Set nuisance-part specific parameters\n",
    "dml_irm_tree$set_ml_nuisance_params(\n",
    "    \"ml_g0\", \"e401\", list(cp = 0.0016, minsplit = 74))\n",
    "dml_irm_tree$set_ml_nuisance_params(\n",
    "    \"ml_g1\", \"e401\", list(cp = 0.0018, minsplit = 70))\n",
    "dml_irm_tree$set_ml_nuisance_params(\n",
    "    \"ml_m\", \"e401\", list(cp = 0.0028, minsplit = 167))\n",
    "\n",
    "dml_irm_tree$fit()\n",
    "dml_irm_tree$summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "protective-communication",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Learners\n",
    "boost = lrn(\"regr.xgboost\", objective = \"reg:squarederror\")\n",
    "boost_class = lrn(\"classif.xgboost\", objective = \"binary:logistic\", eval_metric = \"logloss\")\n",
    "\n",
    "# Boosted Trees\n",
    "set.seed(123)\n",
    "dml_irm_boost = DoubleMLIRM$new(data_dml_base,\n",
    "                                ml_g = boost,\n",
    "                                ml_m = boost_class,\n",
    "                                trimming_threshold = 0.01,\n",
    "                                n_folds = 3)\n",
    "\n",
    "# Set nuisance-part specific parameters\n",
    "if (compareVersion(as.character(packageVersion(\"DoubleML\")), \"0.2.1\") > 0) {\n",
    "    dml_irm_boost$set_ml_nuisance_params(\n",
    "        \"ml_g0\", \"e401\", list(nrounds = 8, eta = 0.1))\n",
    "    dml_irm_boost$set_ml_nuisance_params(\n",
    "        \"ml_g1\", \"e401\", list(nrounds = 29, eta = 0.1))\n",
    "    dml_irm_boost$set_ml_nuisance_params(\n",
    "        \"ml_m\", \"e401\", list(nrounds = 23, eta = 0.1))\n",
    "} else {\n",
    "    # behavior of set_ml_nuisance_params() changed in https://github.com/DoubleML/doubleml-for-r/pull/89\n",
    "    dml_irm_boost$set_ml_nuisance_params(\n",
    "        \"ml_g0\", \"e401\", list(nrounds = 8, eta = 0.1, objective = \"reg:squarederror\", verbose=0))\n",
    "    dml_irm_boost$set_ml_nuisance_params(\n",
    "        \"ml_g1\", \"e401\", list(nrounds = 29, eta = 0.1, objective = \"reg:squarederror\", verbose=0))\n",
    "    dml_irm_boost$set_ml_nuisance_params(\n",
    "        \"ml_m\", \"e401\", list(nrounds = 23, eta = 0.1, objective = \"binary:logistic\", eval_metric = \"logloss\", verbose=0))\n",
    "}\n",
    "\n",
    "dml_irm_boost$fit()\n",
    "dml_irm_boost$summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "duplicate-sport",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:50:25.095222Z",
     "iopub.status.busy": "2021-03-23T13:50:25.093677Z",
     "iopub.status.idle": "2021-03-23T13:50:25.130873Z",
     "shell.execute_reply": "2021-03-23T13:50:25.129138Z"
    },
    "papermill": {
     "duration": 0.108326,
     "end_time": "2021-03-23T13:50:25.131041",
     "exception": false,
     "start_time": "2021-03-23T13:50:25.022715",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "confints = rbind(dml_irm_lasso$confint(), dml_irm_forest$confint(),\n",
    "                 dml_irm_tree$confint(), dml_irm_boost$confint())\n",
    "estimates = c(dml_irm_lasso$coef, dml_irm_forest$coef,\n",
    "              dml_irm_tree$coef, dml_irm_boost$coef)\n",
    "result_irm = data.table(\"model\" = \"IRM\", \n",
    "                        \"ML\" = c(\"glmnet\", \"ranger\", \"rpart\", \"xgboost\"), \n",
    "                        \"Estimate\" = estimates,\n",
    "                        \"lower\" = confints[,1],\n",
    "                        \"upper\" = confints[,2])\n",
    "result_irm\n",
    "\n",
    "g_ci = ggplot(result_irm, aes(x = ML, y = Estimate, color = ML)) + \n",
    "        geom_point() +\n",
    "        geom_errorbar(aes(ymin = lower, ymax = upper, color = ML))  +\n",
    "        geom_hline(yintercept = 0, color = \"grey\") +\n",
    "        theme_minimal() + ylab(\"Coefficients and 0.95- confidence interval\") + \n",
    "        xlab(\"\") + \n",
    "        theme(axis.text.x = element_text(angle = 90), legend.position = \"none\",\n",
    "              text = element_text(size = 20))\n",
    "\n",
    "g_ci"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "checked-paradise",
   "metadata": {
    "papermill": {
     "duration": 0.067511,
     "end_time": "2021-03-23T13:50:25.267772",
     "exception": false,
     "start_time": "2021-03-23T13:50:25.200261",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "These estimates that flexibly account for confounding are\n",
    "substantially attenuated relative to the baseline estimate (*19559*) that does not account for confounding. They suggest much smaller causal effects of 401(k) eligiblity on financial asset holdings. The best model with lowest RMSE in both equations is the PLR model estimated via lasso. It gives the following estimate:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "attended-wisconsin",
   "metadata": {
    "papermill": {
     "duration": 0.06959,
     "end_time": "2021-03-23T13:50:25.563570",
     "exception": false,
     "start_time": "2021-03-23T13:50:25.493980",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Local Average Treatment Effects of 401(k) Participation on Net Financial Assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "million-steel",
   "metadata": {
    "papermill": {
     "duration": 0.068243,
     "end_time": "2021-03-23T13:50:25.700083",
     "exception": false,
     "start_time": "2021-03-23T13:50:25.631840",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Interactive IV Model (IIVM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "historic-bleeding",
   "metadata": {
    "papermill": {
     "duration": 0.070268,
     "end_time": "2021-03-23T13:50:25.842409",
     "exception": false,
     "start_time": "2021-03-23T13:50:25.772141",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "In the examples above, we estimated the average treatment effect of *eligibility* on financial asset holdings. Now, we consider estimation of local average treatment effects (LATE) of *participation* using eligibility as an instrument for the participation decision. Under appropriate assumptions, the LATE identifies the treatment effect for so-called compliers, i.e., individuals who would only participate if eligible and otherwise not participate in the program. \n",
    "\n",
    "As before, $Y$ denotes the outcome `net_tfa`, and $X$ is the vector of covariates. We use `e401` as a binary instrument for the treatment variable `p401`. Here the structural equation model is:\n",
    "\n",
    "\\begin{eqnarray}\n",
    "& Y = g_0(Z,X) + U, &\\quad E[U\\mid Z,X] = 0,\\\\\n",
    "& D = r_0(Z,X) + V, &\\quad E[V\\mid Z, X] = 0,\\\\\n",
    "& Z = m_0(X) + \\zeta, &\\quad E[\\zeta \\mid X] = 0.\n",
    "\\end{eqnarray}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ignored-muslim",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize DoubleMLData with an instrument\n",
    "\n",
    "# Basic model\n",
    "data_dml_base_iv = DoubleMLData$new(data,\n",
    "                                    y_col = \"net_tfa\",\n",
    "                                    d_cols = \"p401\",\n",
    "                                    x_cols = features_base,\n",
    "                                    z_cols = \"e401\")\n",
    "data_dml_base_iv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "starting-ridge",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:50:25.988471Z",
     "iopub.status.busy": "2021-03-23T13:50:25.986744Z",
     "iopub.status.idle": "2021-03-23T13:50:26.049870Z",
     "shell.execute_reply": "2021-03-23T13:50:26.048200Z"
    },
    "papermill": {
     "duration": 0.139252,
     "end_time": "2021-03-23T13:50:26.050042",
     "exception": false,
     "start_time": "2021-03-23T13:50:25.910790",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Flexible model\n",
    "model_data = data.table(\"net_tfa\" = data[, net_tfa],\n",
    "                        \"e401\" = data[, e401],\n",
    "                        \"p401\" = data[, p401],\n",
    "                        features_flex)\n",
    "data_dml_flex_iv = DoubleMLData$new(model_data,\n",
    "                                    y_col = \"net_tfa\",\n",
    "                                    d_cols = \"p401\",\n",
    "                                    z_cols = \"e401\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rubber-thursday",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:50:26.199177Z",
     "iopub.status.busy": "2021-03-23T13:50:26.197932Z",
     "iopub.status.idle": "2021-03-23T13:50:47.781356Z",
     "shell.execute_reply": "2021-03-23T13:50:47.779723Z"
    },
    "papermill": {
     "duration": 21.662252,
     "end_time": "2021-03-23T13:50:47.781590",
     "exception": false,
     "start_time": "2021-03-23T13:50:26.119338",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "set.seed(123)\n",
    "dml_iivm_lasso = DoubleMLIIVM$new(data_dml_flex_iv,\n",
    "                                  ml_g = lasso, \n",
    "                                  ml_m = lasso_class,\n",
    "                                  ml_r = lasso_class,\n",
    "                                  n_folds = 3,\n",
    "                                  trimming_threshold = 0.01,\n",
    "                                  subgroups = list(always_takers = FALSE,\n",
    "                                                   never_takers = TRUE))\n",
    "dml_iivm_lasso$fit()\n",
    "dml_iivm_lasso$summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dental-apparel",
   "metadata": {
    "papermill": {
     "duration": 0.071304,
     "end_time": "2021-03-23T13:50:48.725067",
     "exception": false,
     "start_time": "2021-03-23T13:50:48.653763",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Again, we repeat the procedure for the other machine learning methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "verbal-george",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-23T13:50:48.876339Z",
     "iopub.status.busy": "2021-03-23T13:50:48.874661Z",
     "iopub.status.idle": "2021-03-23T13:52:00.072131Z",
     "shell.execute_reply": "2021-03-23T13:52:00.069435Z"
    },
    "papermill": {
     "duration": 71.274438,
     "end_time": "2021-03-23T13:52:00.072442",
     "exception": false,
     "start_time": "2021-03-23T13:50:48.798004",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Initialize Learner\n",
    "randomForest = lrn(\"regr.ranger\")\n",
    "randomForest_class = lrn(\"classif.ranger\")\n",
    "\n",
    "# Random Forest\n",
    "set.seed(123)\n",
    "dml_iivm_forest = DoubleMLIIVM$new(data_dml_base_iv,\n",
    "                              ml_g = randomForest, \n",
    "                              ml_m = randomForest_class,\n",
    "                              ml_r = randomForest_class,\n",
    "                              n_folds = 3,\n",
    "                              trimming_threshold = 0.01,\n",
    "                              subgroups = list(always_takers = FALSE,\n",
    "                                               never_takers = TRUE))\n",
    "\n",
    "# Set nuisance-part specific parameters\n",
    "dml_iivm_forest$set_ml_nuisance_params(\n",
    "    \"ml_g0\", \"p401\",\n",
    "    list(max.depth = 6, mtry = 4, min.node.size = 7))\n",
    "dml_iivm_forest$set_ml_nuisance_params(\n",
    "    \"ml_g1\", \"p401\", \n",
    "    list(max.depth = 6, mtry = 3, min.node.size = 5))\n",
    "dml_iivm_forest$set_ml_nuisance_params(\n",
    "    \"ml_m\", \"p401\",\n",
    "    list(max.depth = 6, mtry = 3, min.node.size = 6))\n",
    "dml_iivm_forest$set_ml_nuisance_params(\n",
    "    \"ml_r1\", \"p401\",\n",
    "    list(max.depth = 4, mtry = 7, min.node.size = 6))\n",
    "\n",
    "dml_iivm_forest$fit()\n",
    "dml_iivm_forest$summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "twelve-cambodia",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Learner\n",
    "trees = lrn(\"regr.rpart\")\n",
    "trees_class = lrn(\"classif.rpart\")\n",
    "\n",
    "# Trees\n",
    "set.seed(123)\n",
    "dml_iivm_tree = DoubleMLIIVM$new(data_dml_base_iv,\n",
    "                                 ml_g = trees, \n",
    "                                 ml_m = trees_class,\n",
    "                                 ml_r = trees_class,\n",
    "                                 n_folds = 3,\n",
    "                                 trimming_threshold = 0.01,\n",
    "                                 subgroups = list(always_takers = FALSE,\n",
    "                                                  never_takers = TRUE))\n",
    "\n",
    "# Set nuisance-part specific parameters\n",
    "dml_iivm_tree$set_ml_nuisance_params(\n",
    "    \"ml_g0\", \"p401\",\n",
    "    list(cp = 0.0016, minsplit = 74))\n",
    "dml_iivm_tree$set_ml_nuisance_params(\n",
    "    \"ml_g1\", \"p401\",\n",
    "    list(cp = 0.0018, minsplit = 70))\n",
    "dml_iivm_tree$set_ml_nuisance_params(\n",
    "    \"ml_m\", \"p401\",\n",
    "    list(cp = 0.0028, minsplit = 167))\n",
    "dml_iivm_tree$set_ml_nuisance_params(\n",
    "    \"ml_r1\", \"p401\",\n",
    "    list(cp = 0.0576, minsplit = 55))\n",
    "\n",
    "dml_iivm_tree$fit()\n",
    "dml_iivm_tree$summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cardiac-cotton",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Learner\n",
    "boost = lrn(\"regr.xgboost\", objective = \"reg:squarederror\")\n",
    "boost_class = lrn(\"classif.xgboost\", objective = \"binary:logistic\", eval_metric = \"logloss\")\n",
    "\n",
    "# Boosted Trees\n",
    "set.seed(123)\n",
    "dml_iivm_boost = DoubleMLIIVM$new(data_dml_base_iv,\n",
    "                              ml_g = boost, \n",
    "                              ml_m = boost_class,\n",
    "                              ml_r = boost_class,\n",
    "                              n_folds = 3,\n",
    "                              trimming_threshold = 0.01,\n",
    "                              subgroups = list(always_takers = FALSE,\n",
    "                                               never_takers = TRUE))\n",
    "\n",
    "# Set nuisance-part specific parameters\n",
    "if (compareVersion(as.character(packageVersion(\"DoubleML\")), \"0.2.1\") > 0) {\n",
    "    dml_iivm_boost$set_ml_nuisance_params(\n",
    "        \"ml_g0\", \"p401\",\n",
    "        list(nrounds = 9, eta = 0.1))\n",
    "    dml_iivm_boost$set_ml_nuisance_params(\n",
    "        \"ml_g1\", \"p401\",\n",
    "        list(nrounds = 33, eta = 0.1))\n",
    "    dml_iivm_boost$set_ml_nuisance_params(\n",
    "        \"ml_m\", \"p401\",\n",
    "        list(nrounds = 12, eta = 0.1))\n",
    "    dml_iivm_boost$set_ml_nuisance_params(\n",
    "        \"ml_r1\", \"p401\",\n",
    "        list(nrounds = 25, eta = 0.1))\n",
    "} else {\n",
    "    # behavior of set_ml_nuisance_params() changed in https://github.com/DoubleML/doubleml-for-r/pull/89\n",
    "    dml_iivm_boost$set_ml_nuisance_params(\n",
    "        \"ml_g0\", \"p401\",\n",
    "        list(nrounds = 9, eta = 0.1, objective = \"reg:squarederror\", verbose=0))\n",
    "    dml_iivm_boost$set_ml_nuisance_params(\n",
    "        \"ml_g1\", \"p401\",\n",
    "        list(nrounds = 33, eta = 0.1, objective = \"reg:squarederror\", verbose=0))\n",
    "    dml_iivm_boost$set_ml_nuisance_params(\n",
    "        \"ml_m\", \"p401\",\n",
    "        list(nrounds = 12, eta = 0.1, objective = \"binary:logistic\", eval_metric = \"logloss\", verbose=0))\n",
    "    dml_iivm_boost$set_ml_nuisance_params(\n",
    "        \"ml_r1\", \"p401\",\n",
    "        list(nrounds = 25, eta = 0.1, objective = \"binary:logistic\", eval_metric = \"logloss\", verbose=0))\n",
    "}\n",
    "\n",
    "dml_iivm_boost$fit()\n",
    "dml_iivm_boost$summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "nearby-montana",
   "metadata": {},
   "outputs": [],
   "source": [
    "confints = rbind(dml_iivm_lasso$confint(), dml_iivm_forest$confint(),\n",
    "                 dml_iivm_tree$confint(), dml_iivm_boost$confint())\n",
    "estimates = c(dml_iivm_lasso$coef, dml_iivm_forest$coef,\n",
    "              dml_iivm_tree$coef, dml_iivm_boost$coef)\n",
    "result_iivm = data.table(\"model\" = \"IIVM\", \n",
    "                        \"ML\" = c(\"glmnet\", \"ranger\", \"rpart\", \"xgboost\"), \n",
    "                        \"Estimate\" = estimates,\n",
    "                        \"lower\" = confints[,1],\n",
    "                        \"upper\" = confints[,2])\n",
    "result_iivm\n",
    "\n",
    "g_ci = ggplot(result_iivm, aes(x = ML, y = Estimate, color = ML)) + \n",
    "        geom_point() +\n",
    "        geom_errorbar(aes(ymin = lower, ymax = upper, color = ML))  +\n",
    "        geom_hline(yintercept = 0, color = \"grey\") +\n",
    "        theme_minimal() + ylab(\"Coefficients and 0.95- confidence interval\") + \n",
    "        xlab(\"\") + \n",
    "        theme(axis.text.x = element_text(angle = 90), legend.position = \"none\",\n",
    "              text = element_text(size = 20))\n",
    "\n",
    "g_ci"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tropical-texas",
   "metadata": {},
   "source": [
    "## Summary of Results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "labeled-command",
   "metadata": {},
   "source": [
    "To sum up, let's merge all our results so far and illustrate them in a plot. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "legislative-wesley",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_result = rbindlist(list(result_plr, result_irm, result_iivm))\n",
    "summary_result[, model := factor(model, levels = c(\"PLR\", \"IRM\", \"IIVM\"))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "finnish-pontiac",
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "g_all = ggplot(summary_result, aes(x = ML, y = Estimate, color = ML)) + \n",
    "        geom_point() +\n",
    "        geom_errorbar(aes(ymin = lower, ymax = upper, color = ML))  +\n",
    "        geom_hline(yintercept = 0, color = \"grey\") +\n",
    "        theme_minimal() + ylab(\"Coefficients and 0.95- confidence interval\") + \n",
    "        xlab(\"\") + \n",
    "        theme(axis.text.x = element_text(angle = 90), legend.position = \"none\",\n",
    "              text = element_text(size = 20)) + \n",
    "        facet_wrap(model ~., ncol = 1)\n",
    "\n",
    "g_all"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "connected-coach",
   "metadata": {
    "papermill": {
     "duration": 0.079866,
     "end_time": "2021-03-23T13:52:00.452482",
     "exception": false,
     "start_time": "2021-03-23T13:52:00.372616",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We report results based on four ML methods for estimating the nuisance functions used in\n",
    "forming the orthogonal estimating equations. We find again that the estimates of the treatment effect are stable across ML methods. The estimates are highly significant, hence we would reject the hypothesis\n",
    "that 401(k) participation has no effect on financial wealth."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fatal-standing",
   "metadata": {},
   "source": [
    "______\n",
    "\n",
    "**Acknowledgement**\n",
    "\n",
    "We would like to thank [Jannis Kueck](https://www.bwl.uni-hamburg.de/en/statistik/team/kueck.html) for sharing [the kaggle notebook](https://www.kaggle.com/janniskueck/pm5-401k). The pension data set has been analyzed in several studies, among others [Chernozhukov et al. (2018)](https://arxiv.org/abs/1608.00060).\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "R",
   "language": "R",
   "name": "ir"
  },
  "language_info": {
   "codemirror_mode": "r",
   "file_extension": ".r",
   "mimetype": "text/x-r-source",
   "name": "R",
   "pygments_lexer": "r",
   "version": "3.6.1"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 503.595131,
   "end_time": "2021-03-23T13:52:50.313398",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2021-03-23T13:44:26.718267",
   "version": "2.3.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}