{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python: Difference-in-Differences Pre-Testing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Remark:**\n",
    "*This notebook is based on the deprecated version of the DiD implementation of DoubleML. Please check out the [DiD Section](https://docs.doubleml.org/dev/guide/models.html#difference-in-differences-models-did) to find out about the current implementation.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example illustrates how to use the Difference-in-Differences implmentation `DoubleMLDID` of the [DoubleML](https://docs.doubleml.org/stable/index.html) package can be used to pre-test the parallel trends assumptions.\n",
    "The example is based on the great implmentation of the [did-package](https://cran.r-project.org/web/packages/did/vignettes/did-basics.html) in `R`. \n",
    "You can find further references and a detailed guide on pre-testing with the `did`-package at the [did-package pre-testing documentation](https://cran.r-project.org/web/packages/did/vignettes/pre-testing.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from doubleml import DoubleMLDIDData, DoubleMLDID\n",
    "from lightgbm import LGBMClassifier, LGBMRegressor"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At first, we will generate some data to evaluate the corresponding effect estimates. The DGP is slightly adjusted from [did-package pre-testing documentation](https://cran.r-project.org/web/packages/did/vignettes/pre-testing.html) to consider conditional non-linear parallel trends.\n",
    "\n",
    "We will consider $n_{obs}=1000$ units for $2n_t$ time periods $t\\in\\{-n_t, .., n_t - 1\\}$, where the treatment is applied just before $t=0$.\n",
    "The data generating process will take the following form\n",
    "\n",
    "$$\n",
    "Y_{it} = \\theta_t + \\eta_i + D_{i}\\mu_{it} + g(X_{it}) + \\epsilon_{it}.\n",
    "$$\n",
    "Here $\\theta_t = t$ denotes a time fixed effect, $\\eta_i\\sim \\mathcal{N}(0,1)$ is a unit fixed effect, $X_{it}\\sim\\mathcal{N}(0, I_4)\\in \\mathbb{R}^4$ denote 'pre-treatment' covariates for the time period $t$ and $\\epsilon_{it}\\sim \\mathcal{N}(0,1)$ are independent error terms.\n",
    "Further, $\\mu_{it} \\sim \\mathcal{N}(\\tau,1)$ denote time varying individual effects with $\\tau = (\\underbrace{0,\\dots,0}_{n_t\\ times}, n_t, \\dots, 1)^T$ beeing the average effects (coinciding with the ATTE since $\\mu_{it}$ is independent of $X$.) The treatment indicator $D_{i}$ denotes whether unit $i$ was treated between $t=-1$ and $t=0$ and is generated as\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "f_{ps}(X) &= 0.75\\cdot(-X_1 + 0.5X_2 - 0.25X_3 - 0.1X_4X_3 + \\cos(5X_2))\\\\\n",
    "D_{i} &= \\frac{\\exp(f_{ps}(X_{i0}))}{1+\\exp(f_{ps}(X_{i0}))}\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Finally, the direct effect of the covariates is given via\n",
    "\n",
    "$$\n",
    "g(X) = \\exp(X_2) - \\sin(5X_3) + 2X_4.\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "n_obs = 1000\n",
    "n_time_periods = 5\n",
    "\n",
    "time_periods = (np.arange(2*n_time_periods) - n_time_periods)\n",
    "\n",
    "# fixed effects\n",
    "theta = np.zeros(shape=(n_obs, 2*n_time_periods)) + time_periods \n",
    "eta = np.random.normal(loc=0, scale=1, size=(n_obs,1))\n",
    "\n",
    "# covariates\n",
    "X = np.random.normal(loc=0, scale=1, size=(n_obs, 4, 2*n_time_periods))\n",
    "\n",
    "# treatment effects\n",
    "mu_means = np.concatenate((np.zeros(n_time_periods) , np.arange(n_time_periods, 0, -1)))\n",
    "mu = np.random.normal(loc=0, scale=1, size=(n_obs, 2*n_time_periods)) + mu_means\n",
    "\n",
    "# treatment\n",
    "f_ps = 0.75*(-X[:, 0, 0] + 0.5*X[:, 1, 0] - 0.25*X[:, 2, 0] - 0.1*X[:, 3, 0]*X[:, 2, 0] + np.cos(5*X[:, 1, 0]))\n",
    "ps = (np.exp(f_ps) / (1 + np.exp(f_ps))).reshape(-1,1)\n",
    "u = np.random.uniform(low=0, high=1, size=(n_obs,1))\n",
    "treatment = np.ones(shape=(n_obs, 2*n_time_periods)) * (ps >= u)\n",
    "\n",
    "\n",
    "# outcome\n",
    "g_X = (np.exp(X[:, 1, :]) - np.sin(5*X[:, 2, :]) + 2*X[:, 3, :])\n",
    "epsilon = np.random.normal(loc=0, scale=1, size=(n_obs, 2*n_time_periods))\n",
    "Y = theta + eta + treatment*mu + g_X + epsilon"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate a corresponding dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate dataframe\n",
    "Y_df = np.hstack(Y)\n",
    "treatment_df = np.hstack(treatment)\n",
    "X_df = np.transpose(np.hstack(X))\n",
    "time_df = np.hstack(np.ones(shape=(n_obs, 2*n_time_periods))*time_periods)\n",
    "individual_df = np.hstack(np.ones(shape=(n_obs, 2*n_time_periods))*np.arange(n_obs).reshape(-1,1))\n",
    "\n",
    "data = pd.DataFrame(np.column_stack((Y_df,treatment_df, X_df, time_df, individual_df)),\n",
    "                    columns = ['y'] + ['d'] + [f'X_{i}' for i in range(4)] + ['t'] + ['i'])\n",
    "print(data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will just rely on [LightGBM](https://lightgbm.readthedocs.io/en/latest/index.html) with a small number of trees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_estimators = 50\n",
    "ml_g = LGBMRegressor(n_estimators=n_estimators, num_leaves=5, verbose=-1)\n",
    "ml_m = LGBMClassifier(n_estimators=n_estimators, num_leaves=5, verbose=-1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, construct the corresponding `DoubleMLData` for each time period and fit a `DoubleMLDID` model.\n",
    "Remark thatwe have to include the covariates from both time periods (and they are either pre-treatment or independent of previous outcomes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(np.nan, \n",
    "                  index=range(2*n_time_periods-1), \n",
    "                  columns=['lower', 'effect', 'upper'])\n",
    "df['time'] = time_periods[1:]\n",
    "df['true_effect'] = mu_means[1:]\n",
    "\n",
    "np.random.seed(42)\n",
    "for t_idx, t in enumerate(time_periods[1:]):\n",
    "    if t <= 0:\n",
    "        t_diff = t-1\n",
    "    else:\n",
    "        # compare to outcome before treatment\n",
    "        t_diff = -1\n",
    "    # outcome as the difference for each model\n",
    "    y_diff = data[data['t'] == t]['y'].values - data[data['t'] == t_diff]['y'].values\n",
    "    covariates = np.column_stack((data[data['t'] == t][[f'X_{i}' for i in range(4)]].values, data[data['t'] == t_diff][[f'X_{i}' for i in range(4)]].values))\n",
    "    dml_data = DoubleMLDIDData.from_arrays(x=covariates,\n",
    "                                        y=y_diff,\n",
    "                                        d=data[data['t'] == t]['d'].values)\n",
    "    dml_did = DoubleMLDID(dml_data,\n",
    "                          ml_g=ml_g,\n",
    "                          ml_m=ml_m)\n",
    "    dml_did.fit()\n",
    "\n",
    "    df.at[t_idx, 'effect'] = dml_did.coef\n",
    "    confint = dml_did.confint(level=0.95)\n",
    "    df.at[t_idx, 'lower'] = confint['2.5 %'].iloc[0]\n",
    "    df.at[t_idx, 'upper'] = confint['97.5 %'].iloc[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a look at the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams['figure.figsize'] = 10., 7.5\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "errors = np.full((2, 2*n_time_periods - 1), np.nan)\n",
    "errors[0, :] = df['effect'] -   df['lower']\n",
    "errors[1, :] = df['upper'] - df['effect']\n",
    "\n",
    "plt.errorbar(df['time'], df['effect'], fmt='o', yerr=errors, color='#1F77B4',\n",
    "             ecolor='#1F77B4', label='Estimated Effect (with CI)')\n",
    "ax.plot(time_periods[1:], df['effect'], linestyle='--', color='#1F77B4', linewidth=1)\n",
    "\n",
    "# add horizontal line\n",
    "ax.axhline(y=0, color='r', linestyle='--', linewidth=1)\n",
    "\n",
    "# add true effect\n",
    "ax.scatter(x=df['time'], y=df['true_effect'], c='#FF7F0E', label='True Effect')\n",
    "\n",
    "plt.xlabel('Exposure')\n",
    "plt.legend()\n",
    "_ = plt.ylabel('Effect and 95%-CI')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}