{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python: Policy Learning with Trees\n",
    "\n",
    "In this simple example, we illustrate how the [DoubleML](https://docs.doubleml.org/stable/index.html) package can be used to estimate deterministic binary treatment policies."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "\n",
    "The data will be generated with a simple data generating process to enable us to know the true policy cut-offs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import pandas as pd\n",
    "import doubleml as dml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we consider a treatment effect that depends on two covariates which might correspond to opposite effects depending on age or income.\n",
    "For simplicity, the treatment effect within each group is generated to be constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def group_effect(x):\n",
    "    if x[0] <= -0.3:\n",
    "        te = 2.2\n",
    "    elif (x[0] >= 0.2) & (x[1] >= 0.4):\n",
    "        te = 2\n",
    "    else:\n",
    "        te = -2\n",
    "    return te"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data is generated as\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "Y_i & = g(W_1,W_2)T_i + \\langle W_i,\\gamma_0\\rangle + \\epsilon_i \\\\\n",
    "T_i & = \\langle W_i,\\beta_0\\rangle +\\eta_i,\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $W_i\\sim\\mathcal{N}(0,I_{d_w})$ and $\\epsilon_i,\\eta_i\\sim\\mathcal{U}[0,1]$.\n",
    "The coefficient vectors $\\gamma_0$ and $\\beta_0$ both have small random support which values are drawn independently from $\\mathcal{U}[0,1]$.\n",
    "Further, $g(w_1,w_2)$ defines the conditional treatment effect, which is generated depending on $\\{w_1,w_2\\}$.\n",
    "\n",
    "$$g(w_1) = \\begin{cases}2.2\\quad &\\text{for } w_1\\le -0.3\\\\\n",
    "2\\quad &\\text{for } w_1\\ge  0.2 \\land w_2\\ge 0.4\\\\\n",
    "-2\\quad &\\text{otherwise. } \\\\\n",
    " \\end{cases}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_synthetic_group_data(n_samples=200, n_w=10, support_size=5):\n",
    "    \"\"\"\n",
    "    Creates a simple synthetic example for group effects.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    n_samples : int\n",
    "        Number of samples.\n",
    "        Default is ``200``.\n",
    "\n",
    "    n_w : int\n",
    "        Dimension of covariates.\n",
    "        Default is ``10``.\n",
    "\n",
    "    support_size : int\n",
    "        Number of relevant covariates.\n",
    "        Default is ``5``.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "     data : pd.DataFrame\n",
    "            A data frame.\n",
    "\n",
    "    \"\"\"\n",
    "    # Outcome support\n",
    "    support_w = np.random.choice(np.arange(n_w), size=support_size, replace=False)\n",
    "    coefs_w = np.random.uniform(0, 1, size=support_size)\n",
    "    # Define the function to generate the noise\n",
    "    epsilon_sample = lambda n: np.random.uniform(-1, 1, size=n_samples)\n",
    "    # Treatment support\n",
    "    # Assuming the matrices gamma and beta have the same number of non-zero components\n",
    "    support_t = np.random.choice(np.arange(n_w), size=support_size, replace=False)\n",
    "    coefs_t = np.random.uniform(0, 1, size=support_size)\n",
    "    # Define the function to generate the noise\n",
    "    eta_sample = lambda n: np.random.uniform(-1, 1, size=n_samples)\n",
    "\n",
    "    # Generate controls, covariates, treatments and outcomes\n",
    "    w = np.random.normal(0, 1, size=(n_samples, n_w))\n",
    "    # Group treatment effect\n",
    "    te = np.apply_along_axis(group_effect, axis=1, arr=w)\n",
    "    # Define treatment\n",
    "    log_odds = np.dot(w[:, support_t], coefs_t) + eta_sample(n_samples)\n",
    "    t_sigmoid = 1 / (1 + np.exp(-log_odds))\n",
    "    t = np.array([np.random.binomial(1, p) for p in t_sigmoid])\n",
    "    # Define the outcome\n",
    "    y = te * t + np.dot(w[:, support_w], coefs_w) + epsilon_sample(n_samples)\n",
    "\n",
    "    # Now we build the dataset\n",
    "    y_df = pd.DataFrame({'y': y})\n",
    "    t_df = pd.DataFrame({'t': t})\n",
    "    w_df = pd.DataFrame(data=w, index=np.arange(w.shape[0]), \n",
    "                        columns=[f'w_{i}' for i in range(1, w.shape[1] + 1)])\n",
    "\n",
    "    data = pd.concat([y_df, t_df, w_df], axis=1)\n",
    "    covariates = list(w_df.columns.values)\n",
    "\n",
    "    return data, covariates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will consider a quite small number of covariates to ensure fast calcualtion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DGP constants\n",
    "np.random.seed(42)\n",
    "n_samples = 500\n",
    "n_w = 10\n",
    "support_size = 5\n",
    "\n",
    "# Create data\n",
    "data, covariates = create_synthetic_group_data(n_samples=n_samples, n_w=n_w, support_size=support_size)\n",
    "data_dml_base = dml.DoubleMLData(data,\n",
    "                                 y_col='y',\n",
    "                                 d_cols='t',\n",
    "                                 x_cols=covariates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive Regression Model (IRM)\n",
    "The first step is to fit a [DoubleML IRM Model](https://docs.doubleml.org/stable/guide/models.html#interactive-regression-models-irm) to the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First stage estimation\n",
    "from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor\n",
    "randomForest_reg = RandomForestRegressor(n_estimators=200, random_state=42)\n",
    "randomForest_class = RandomForestClassifier(n_estimators=200, random_state=42)\n",
    "\n",
    "np.random.seed(42)\n",
    "\n",
    "dml_irm = dml.DoubleMLIRM(data_dml_base,\n",
    "                          ml_g=randomForest_reg,\n",
    "                          ml_m=randomForest_class,\n",
    "                          trimming_threshold=0.01,\n",
    "                          n_folds=5)\n",
    "print(\"Training IRM Model\")\n",
    "dml_irm.fit(store_predictions=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Policy Learning with Trees\n",
    "Next, we specify the covariates as a [DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) against which to learn the policy. This can either be all covariates $w_i$, or if domain knowledge is available we can use a subset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = data[[\"w_1\",\"w_2\"]].copy()\n",
    "print(features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To estimate a Policy just call the ``policy_tree()`` method and supply the DataFrame with the ``features`` and the ``depth`` of the desired tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy_tree = dml_irm.policy_tree(features=features, depth=3)\n",
    "policy_tree.plot_tree();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The splits in this classification tree reflect the heterogeneity in the data generating process above. Exemplatory, we see that the first split is estimated at $w_1 = -0.292$. In the DGP, at $w_1 = -0.3$ the treatment effect jumps from $-2$ to $2.2$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, it is also possible to estimate the tree on all covariates, which in this case results in the same splits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy_tree_2 = dml_irm.policy_tree(features=data[covariates], depth=3)\n",
    "policy_tree_2.plot_tree();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The policytree `predict()` method is useful to estimate the optimal treatments on new data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_data, _ = create_synthetic_group_data(n_samples=n_samples, n_w=n_w, support_size=support_size)\n",
    "pred_df = policy_tree.predict(features=new_data[[\"w_1\",\"w_2\"]])\n",
    "print(pred_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These predicted optimal treatments maximize the overall Average Treament Effect on the treated within the new data. We showcase that with the help of the `gate()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dml_new = dml.DoubleMLData(new_data,\n",
    "                                y_col=\"y\",\n",
    "                                d_cols=\"t\")\n",
    "\n",
    "dml_irm_new = dml.DoubleMLIRM(data_dml_new,\n",
    "                              ml_g=randomForest_reg,\n",
    "                              ml_m=randomForest_class,\n",
    "                              trimming_threshold=0.01,\n",
    "                              n_folds=5)\n",
    "dml_irm_new.fit()\n",
    "print(\"ATTE under the original treatment policy:\")\n",
    "print(dml_irm_new.gate(groups=new_data[\"t\"].to_frame()).confint())\n",
    "print(\"\\nATTE under the estimated optimal treatment policy:\")\n",
    "print(dml_irm_new.gate(groups=pred_df[\"pred_treatment\"].to_frame()).confint())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the original policy in the underlying DGP, we achieve a significant ATTE improvement."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dml_dev",
   "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.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}