{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e401eb2c",
   "metadata": {},
   "source": [
    "# R: DoubleML for Difference-in-Differences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we demonstrate, how `DoubleML` can be used in combination with the [did package for R](https://bcallaway11.github.io/did/index.html) in order to estimate group-time average treatment effects in difference-in-difference (DiD) models with multiple periods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "library(DoubleML)\n",
    "library(did)\n",
    "library(mlr3)\n",
    "library(mlr3learners)\n",
    "\n",
    "# suppress messages during fitting\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\")\n",
    "\n",
    "set.seed(1234)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19aaa906",
   "metadata": {},
   "source": [
    "# Demo Example from `did`\n",
    "\n",
    "We will demonstrate the use of `DoubleML` for DiD in the [introductory example](https://bcallaway11.github.io/did/articles/did-basics.html) of the `did` package. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5d9b7ac",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "# Generate data, original code available at https://github.com/bcallaway11/did/blob/master/vignettes/did-basics.Rmd\n",
    "time.periods <- 4\n",
    "sp <- reset.sim()\n",
    "sp$te <- 0\n",
    "\n",
    "set.seed(1814)\n",
    "\n",
    "# generate dataset with 4 time periods\n",
    "time.periods <- 4\n",
    "\n",
    "# add dynamic effects\n",
    "sp$te.e <- 1:time.periods\n",
    "\n",
    "# generate data set with these parameters\n",
    "# here, we dropped all units who are treated in time period 1 as they do not help us recover ATT(g,t)'s.\n",
    "dta <- build_sim_dataset(sp)\n",
    "\n",
    "# How many observations remained after dropping the ``always-treated'' units\n",
    "nrow(dta)\n",
    "#This is what the data looks like\n",
    "head(dta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f508a08b",
   "metadata": {},
   "source": [
    "### Comparison to `did` package\n",
    "\n",
    "By default, estimation in `did` is based on (unpenalized) linear and logistic regression. Let's start with this default model first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "945902b5",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "# estimate group-time average treatment effects using att_gt method\n",
    "example_attgt <- att_gt(yname = \"Y\",\n",
    "                        tname = \"period\",\n",
    "                        idname = \"id\",\n",
    "                        gname = \"G\",\n",
    "                        xformla = ~X,\n",
    "                        data = dta\n",
    "                        )\n",
    "\n",
    "# summarize the results\n",
    "summary(example_attgt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed9b4f0e",
   "metadata": {},
   "source": [
    "### Using ML for DiD: Integrating `DoubleML` in `did`\n",
    "\n",
    "As described in our [Section on DiD models in the user guide](https://docs.doubleml.org/stable/guide/models.html#difference-in-differences-models-did), [Sant'Anna and Zhao (2020)](https://linkinghub.elsevier.com/retrieve/pii/S0304407620301901) have developed a doubly robust DiD model which is compatible with ML-based estimation. As this doubly robust model is internally used in `did`, it is possible to use `DoubleML` here to obtain valid point estimates and confidence intervals. For this, we need to write a wrapper around a `DoubleMLIRM` model and pass it to `did` as a custom estimation approach. Once this is implemented, we can use all the nice features and advantages of the `did` package.\n",
    "\n",
    "For now, let's abstract from using fancy ML algorithms to keep the comparison to the classic `did` implementation simple. Hence, we will use linear and logistic regression for the nuisance compontents in the DiD model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0e92c95",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "# DoubleML wrapper for did\n",
    "set.seed(1234)\n",
    "doubleml_did_linear <- function(y1, y0, D, covariates,\n",
    "                         ml_g = lrn(\"regr.lm\"),\n",
    "                         ml_m = lrn(\"classif.log_reg\"),\n",
    "                         n_folds = 10, n_rep = 1, ...) {\n",
    "  \n",
    "  # warning if n_rep > 1 to handle mapping from psi to inf.func\n",
    "  if (n_rep > 1) {\n",
    "    warning(\"n_rep > 1 is not supported.\")\n",
    "  }\n",
    "  # Compute difference in outcomes\n",
    "  delta_y <- y1 - y0\n",
    "  # Prepare data backend\n",
    "  dml_data = DoubleML::double_ml_data_from_matrix(X = covariates, y = delta_y, d = D)\n",
    "  # Compute the ATT\n",
    "  dml_obj = DoubleML::DoubleMLIRM$new(dml_data, ml_g = ml_g, ml_m = ml_m, score = \"ATTE\", n_folds = n_folds)\n",
    "  dml_obj$fit()\n",
    "  att = dml_obj$coef[1]\n",
    "  # Return results\n",
    "  inf.func <- dml_obj$psi[, 1, 1]\n",
    "  output <- list(ATT = att, att.inf.func = inf.func)\n",
    "  return(output)\n",
    "}\n",
    "\n",
    "example_attgt_dml_linear <- att_gt(yname = \"Y\",\n",
    "                        tname = \"period\",\n",
    "                        idname = \"id\",\n",
    "                        gname = \"G\",\n",
    "                        xformla = ~X,\n",
    "                        data = dta,\n",
    "                        est_method = doubleml_did_linear)\n",
    "\n",
    "\n",
    "summary(example_attgt_dml_linear)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "344bfbf4",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "source": [
    "Any differences from the default `did` implementation arise due to sampling randomness, because `DoubleML` uses cross-fitting internally, which is not necessary if classical parametric estimation methods are used.\n",
    "\n",
    "Next, let's demonstrate how we can use more complex ML learners. For this, we just have to pass another `mlr3` learner through the wrapper, for example a random forest. Please note that the original data generating process is linear, such that we don't expect random forest to lead to better results than the linear learners. We provide a variant of the wrapper that includes an evaluation of the nuisance predictions at the end of this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23e26476",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "# DoubleML wrapper for did with random forest learner\n",
    "set.seed(1234)\n",
    "\n",
    "doubleml_did_rf <- function(y1, y0, D, covariates,\n",
    "  ml_g = lrn(\"regr.ranger\"),\n",
    "  ml_m = lrn(\"classif.ranger\"),\n",
    "  n_folds = 10, n_rep = 1, ...) {\n",
    "\n",
    "# warning if n_rep > 1 to handle mapping from psi to inf.func\n",
    "if (n_rep > 1) {\n",
    "warning(\"n_rep > 1 is not supported.\")\n",
    "}\n",
    "# Compute difference in outcomes\n",
    "delta_y <- y1 - y0\n",
    "# Prepare data backend\n",
    "dml_data = DoubleML::double_ml_data_from_matrix(X = covariates, y = delta_y, d = D)\n",
    "# Compute the ATT\n",
    "dml_obj = DoubleML::DoubleMLIRM$new(dml_data, ml_g = ml_g, ml_m = ml_m, score = \"ATTE\", n_folds = n_folds)\n",
    "dml_obj$fit()\n",
    "att = dml_obj$coef[1]\n",
    "# Return results\n",
    "inf.func <- dml_obj$psi[, 1, 1]\n",
    "output <- list(ATT = att, att.inf.func = inf.func)\n",
    "return(output)\n",
    "}\n",
    "\n",
    "example_attgt_dml_rf <- att_gt(yname = \"Y\",\n",
    " tname = \"period\",\n",
    " idname = \"id\",\n",
    " gname = \"G\",\n",
    " xformla = ~X,\n",
    " data = dta,\n",
    " est_method = doubleml_did_rf)\n",
    "\n",
    "summary(example_attgt_dml_rf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25d43c1d",
   "metadata": {},
   "source": [
    "We can see that the results are not dramatically different from the results before. We can observe from the larger standard errors that the default random forest learners seems to be a less precise prediction rule."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "676c1a49",
   "metadata": {},
   "source": [
    "### Exploiting the Functionalities of `did`\n",
    "\n",
    "The `did` package offers various tools for multi-period DiD models, for example plotting the group-time average treatment effects, which can be exploited just as in the native `did` usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8b9596e",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "# Plot group-time average treatment effects\n",
    "ggdid(example_attgt_dml_linear)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4e701ec",
   "metadata": {},
   "source": [
    "It's also possible to calculate aggregated effect estimates. Please note that, the results are again very close to those in [the original notebook](https://bcallaway11.github.io/did/articles/did-basics.html#aggregating-group-time-average-treatment-effects)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f163147",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "agg.simple <- aggte(example_attgt_dml_linear, type = \"simple\")\n",
    "summary(agg.simple)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb050302",
   "metadata": {},
   "source": [
    "### Details on Predictive Performance\n",
    "\n",
    "We can add an evaluation functionality to the wrapper to assess how the predictive performance from the linear and logistic regression differ from that of the random forest learner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8bb468f",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "library(mlr3measures)\n",
    "\n",
    "# Add a wrapper that computes the RMSE and accuracy for the nuisance components, can be customized by providing custom measures\n",
    "eval_preds = function(y, d, predictions, params_names, custom_measures = NULL) {\n",
    "  measures_res = list()\n",
    "\n",
    "  if (!is.null(custom_measures)) {\n",
    "    # Alternatively provide a named list with custom evaluation functions\n",
    "    measure_funcs = list()\n",
    "  } else {\n",
    "    measure_funcs = list()\n",
    "    measure_funcs[['ml_m']] = mlr3measures::acc\n",
    "    measure_funcs[['ml_g0']] = mlr3measures::rmse\n",
    "  }\n",
    "\n",
    "  for (param_name in params_names) {\n",
    "    preds = predictions[[param_name]][, 1, 1]\n",
    "\n",
    "  if (param_name == \"ml_m\") {\n",
    "    obs = d\n",
    "    # map probability predictions to binary\n",
    "    preds = as.factor(ifelse(preds > 0.5, 1, 0))\n",
    "    obs = as.factor(preds)\n",
    "  }\n",
    "\n",
    "  else if (param_name == \"ml_g0\") {\n",
    "    obs = y[d == 0]\n",
    "    preds = preds[d == 0]\n",
    "  }\n",
    "\n",
    "  if (param_name == \"ml_g1\") {\n",
    "    next\n",
    "  }\n",
    "\n",
    "  else {\n",
    "    measure_func = measure_funcs[[param_name]]\n",
    "    measure_pred = measure_func(obs, preds)\n",
    "\n",
    "    measures_res[[param_name]] = measure_pred\n",
    "  }\n",
    "\n",
    "  }\n",
    "    return(measures_res)\n",
    "}\n",
    "\n",
    "# evaluate learner performance: linear models\n",
    "doubleml_did_eval_linear <- function(y1, y0, D, covariates,\n",
    "  ml_g = lrn(\"regr.lm\"),\n",
    "  ml_m = lrn(\"classif.log_reg\"),\n",
    "  n_folds = 10, n_rep = 1, ...) {\n",
    "\n",
    "  # warning if n_rep > 1 to handle mapping from psi to inf.func\n",
    "  if (n_rep > 1) {\n",
    "    warning(\"n_rep > 1 is not supported.\")\n",
    "  }\n",
    "  # Compute difference in outcomes\n",
    "  delta_y <- y1 - y0\n",
    "  # Prepare data backend\n",
    "  dml_data = DoubleML::double_ml_data_from_matrix(X = covariates, y = delta_y, d = D)\n",
    "  # Compute the ATT\n",
    "  dml_obj = DoubleML::DoubleMLIRM$new(dml_data, ml_g = ml_g, ml_m = ml_m, score = \"ATTE\", n_folds = n_folds)\n",
    "  dml_obj$fit(store_predictions = TRUE)\n",
    "  att = dml_obj$coef[1]\n",
    "  # Return results\n",
    "  inf.func <- dml_obj$psi[, 1, 1]\n",
    "\n",
    "  # Evaluate learner performance\n",
    "  predictions = dml_obj$predictions\n",
    "  params_names = dml_obj$params_names()\n",
    "  eval_predictions = eval_preds(delta_y, D, predictions, params_names)\n",
    "  print(eval_predictions)\n",
    "\n",
    "  output <- list(ATT = att, att.inf.func = inf.func)\n",
    "  return(output)\n",
    "  }\n",
    "\n",
    "  library(mlr3measures)\n",
    "\n",
    "# evaluate learner performance: random forest\n",
    "doubleml_did_eval_rf <- function(y1, y0, D, covariates,\n",
    "  ml_g = lrn(\"regr.ranger\"),\n",
    "  ml_m = lrn(\"classif.ranger\"),\n",
    "  n_folds = 10, n_rep = 1, ...) {\n",
    "\n",
    "  # warning if n_rep > 1 to handle mapping from psi to inf.func\n",
    "  if (n_rep > 1) {\n",
    "    warning(\"n_rep > 1 is not supported.\")\n",
    "  }\n",
    "  # Compute difference in outcomes\n",
    "  delta_y <- y1 - y0\n",
    "  # Prepare data backend\n",
    "  dml_data = DoubleML::double_ml_data_from_matrix(X = covariates, y = delta_y, d = D)\n",
    "  # Compute the ATT\n",
    "  dml_obj = DoubleML::DoubleMLIRM$new(dml_data, ml_g = ml_g, ml_m = ml_m, score = \"ATTE\", n_folds = n_folds)\n",
    "  dml_obj$fit(store_predictions = TRUE)\n",
    "  att = dml_obj$coef[1]\n",
    "  # Return results\n",
    "  inf.func <- dml_obj$psi[, 1, 1]\n",
    "\n",
    "  # Evaluate learner performance\n",
    "  predictions = dml_obj$predictions\n",
    "  params_names = dml_obj$params_names()\n",
    "  eval_predictions = eval_preds(delta_y, D, predictions, params_names)\n",
    "  print(eval_predictions)\n",
    "\n",
    "  output <- list(ATT = att, att.inf.func = inf.func)\n",
    "  return(output)\n",
    "  }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0bca913",
   "metadata": {},
   "source": [
    "Running the evaluation wrappers helps to see that the random forest learner has a higher RMSE for predicting the outcome $E[\\Delta Y|D=1,X]$. Both models predict individuals' treatment (group) status with an accuracy of $1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de15e86b",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "# Run estimation with evaluation: Linear model\n",
    "set.seed(1234)\n",
    "example_attgt_dml_eval_linear <- att_gt(yname = \"Y\",\n",
    "                      tname = \"period\",\n",
    "                      idname = \"id\",\n",
    "                      gname = \"G\",\n",
    "                      xformla = ~X,\n",
    "                      data = dta,\n",
    "                      est_method = doubleml_did_eval_linear,\n",
    "                      print_details = TRUE)\n",
    "\n",
    "\n",
    "summary(example_attgt_dml_eval_linear)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13fe0be0",
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "# Run estimation with evaluation: Linear model\n",
    "set.seed(1234)\n",
    "example_attgt_dml_eval_rf <- att_gt(yname = \"Y\",\n",
    "                      tname = \"period\",\n",
    "                      idname = \"id\",\n",
    "                      gname = \"G\",\n",
    "                      xformla = ~X,\n",
    "                      data = dta,\n",
    "                      est_method = doubleml_did_eval_rf,\n",
    "                      print_details = TRUE)\n",
    "\n",
    "\n",
    "summary(example_attgt_dml_eval_rf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5734056",
   "metadata": {},
   "source": [
    "### Acknowledgements and Final Remarks\n",
    "\n",
    "We'd like to thank the authors of the `did` package for R for maintaining a flexible interface for multi-period DiD models.\n",
    "\n",
    "We'd like to note that the implementation presented is here is very similar to the one implemented in the Python package. For more details, we would like to reference to the [DiD](https://docs.doubleml.org/stable/examples/py_double_ml_did.html) and [pretesting](https://docs.doubleml.org/stable/examples/py_double_ml_did_pretest.html) examples in Python."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76593791",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "Callaway, Brantly, and Pedro HC Sant’Anna. \"Difference-in-differences with multiple time periods.\" Journal of Econometrics 225.2 (2021): 200-230.\n",
    "\n",
    "Sant’Anna, Pedro HC, and Jun Zhao. \"Doubly robust difference-in-differences estimators.\" Journal of Econometrics 219.1 (2020): 101-122."
   ]
  }
 ],
 "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": "4.2.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}