{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python: Average Potential Outcome (APO) Models\n",
    "\n",
    "In this example, we illustrate how the [DoubleML](https://docs.doubleml.org/stable/index.html) package can be used to estimate average potential outcomes (APOs) in an interactive regression model (see [DoubleMLIRM](https://docs.doubleml.org/stable/guide/models.html#interactive-regression-model-irm)).\n",
    "\n",
    "The goal is to estimate the average potential outcome\n",
    "\n",
    " $$\\theta_0 =\\mathbb{E}[Y(d)]$$\n",
    "\n",
    "for a given treatment level $d$ and and discrete valued treatment $D$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "from sklearn.linear_model import LogisticRegression, LinearRegression\n",
    "from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor\n",
    "\n",
    "import doubleml as dml\n",
    "from doubleml.datasets import make_irm_data_discrete_treatments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Generating Process (DGP)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At first, let us generate data according to the [make_irm_data_discrete_treatments](https://docs.doubleml.org/dev/api/generated/doubleml.datasets.make_irm_data_discrete_treatments.html#doubleml.datasets.make_irm_data_discrete_treatments) data generating process. The process generates data with a continuous\n",
    "treatment variable and contains the true individual treatment effects (ITEs) with respect to option of not getting treated.\n",
    "\n",
    "According to the continuous treatment variable, the treatment is discretized into multiple levels, based on quantiles. Using the *oracle* ITEs, enables the comparison to the true APOs and averate treatment effects (ATEs) for the different levels of the treatment variable.\n",
    "\n",
    "**Remark:** The average potential outcome model does not require an underlying continuous treatment variable. The model will work identically if the treatment variable is discrete by design."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "n_obs = 3000\n",
    "n_levels = 5\n",
    "linear = True\n",
    "n_rep = 10\n",
    "\n",
    "np.random.seed(42)\n",
    "data_apo = make_irm_data_discrete_treatments(n_obs=n_obs,n_levels=n_levels, linear=linear)\n",
    "\n",
    "y0 = data_apo['oracle_values']['y0']\n",
    "cont_d = data_apo['oracle_values']['cont_d']\n",
    "ite = data_apo['oracle_values']['ite']\n",
    "d = data_apo['d']\n",
    "potential_level = data_apo['oracle_values']['potential_level']\n",
    "level_bounds = data_apo['oracle_values']['level_bounds']\n",
    "\n",
    "average_ites = np.full(n_levels + 1, np.nan)\n",
    "apos = np.full(n_levels + 1, np.nan)\n",
    "mid_points = np.full(n_levels, np.nan)\n",
    "\n",
    "for i in range(n_levels + 1):\n",
    "    average_ites[i] = np.mean(ite[d == i]) * (i > 0)\n",
    "    apos[i] = np.mean(y0) + average_ites[i]\n",
    "\n",
    "print(f\"Average Individual effects in each group:\\n{np.round(average_ites,2)}\\n\")\n",
    "print(f\"Average Potential Outcomes in each group:\\n{np.round(apos,2)}\\n\")\n",
    "print(f\"Levels and their counts:\\n{np.unique(d, return_counts=True)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To better grasp the distribution of the treatment effects, let us plot the true APOs and ATEs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a colorblind-friendly palette\n",
    "palette = sns.color_palette(\"colorblind\")\n",
    "\n",
    "df = pd.DataFrame({'cont_d': cont_d, 'ite': ite})\n",
    "df_sorted = df.sort_values('cont_d')\n",
    "\n",
    "mid_points = np.full(n_levels, np.nan)\n",
    "for i in range(n_levels):\n",
    "    mid_points[i] = (level_bounds[i] + level_bounds[i + 1]) / 2\n",
    "\n",
    "df_apos = pd.DataFrame({'mid_points': mid_points, 'treatment effects': apos[1:] - apos[0]})\n",
    "\n",
    "# Create the primary plot with scatter and line plots\n",
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "sns.lineplot(data=df_sorted, x='cont_d', y='ite', color=palette[0], label='ITE', ax=ax1)\n",
    "sns.scatterplot(data=df_apos, x='mid_points', y='treatment effects', color=palette[1], label='Grouped Treatment Effects', ax=ax1)\n",
    "\n",
    "# Add vertical dashed lines at level_bounds\n",
    "for bound in level_bounds:\n",
    "    ax1.axvline(x=bound, color='grey', linestyle='--', alpha=0.7)\n",
    "\n",
    "ax1.set_title('Grouped Effects vs. Continuous Treatment')\n",
    "ax1.set_xlabel('Continuous Treatment')\n",
    "ax1.set_ylabel('Effects')\n",
    "\n",
    "# Create a secondary y-axis for the histogram\n",
    "ax2 = ax1.twinx()\n",
    "\n",
    "# Plot the histogram on the secondary y-axis\n",
    "ax2.hist(df_sorted['cont_d'], bins=30, alpha=0.3, weights=np.ones_like(df_sorted['cont_d']) / len(df_sorted['cont_d']), color=palette[2])\n",
    "ax2.set_ylabel('Density')\n",
    "\n",
    "# Make sure the legend includes all plots\n",
    "lines, labels = ax1.get_legend_handles_labels()\n",
    "ax1.legend(lines, labels, loc='upper left')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As for all [DoubleML](https://docs.doubleml.org/stable/index.html) models, we specify a [DoubleMLData](https://docs.doubleml.org/stable/api/generated/doubleml.DoubleMLData.html#doubleml.DoubleMLData) object to handle the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = data_apo['y']\n",
    "x = data_apo['x']\n",
    "d = data_apo['d']\n",
    "df_apo = pd.DataFrame(\n",
    "    np.column_stack((y, d, x)),\n",
    "    columns=['y', 'd'] + ['x' + str(i) for i in range(data_apo['x'].shape[1])]\n",
    ")\n",
    "\n",
    "dml_data = dml.DoubleMLData(df_apo, 'y', 'd')\n",
    "print(dml_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single Average Potential Outcome Models (APO)\n",
    "\n",
    "Further, we have to specify machine learning algorithms. As in the [DoubleMLIRM](https://docs.doubleml.org/stable/guide/models.html#interactive-regression-model-irm) model, we have to set ``ml_m`` as a classifier and ``ml_g`` as a regressor (since the outcome is continuous). As in the \n",
    "[DoubleMLIRM](https://docs.doubleml.org/stable/guide/models.html#interactive-regression-model-irm) model, the classifier ``ml_m`` is used to estimate the conditional probability of receiving treatment level $d$ given the covariates $X$\n",
    "\n",
    "$$m_{0,d}(X) = \\mathbb{E}[1\\{D=d\\}|X]$$\n",
    "\n",
    "and the regressor ``ml_g`` is used to estimate the conditional expectation of the outcome $Y$ given the covariates $X$ and the treatment $D$\n",
    "\n",
    "$$g_{0}(D, X) = \\mathbb{E}[Y|X,D].$$\n",
    "\n",
    "As the DGP is linear we will use a linear regression model for the regressor and a logistic regression model for the classifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_g = LinearRegression()\n",
    "ml_m = LogisticRegression()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Further, the [DoubleMLAPO](https://docs.doubleml.org/dev/api/generated/doubleml.DoubleMLAPO.html#doubleml.DoubleMLAPO) model requires a specification of the treatment level $a$ for which the APOs should be estimated. In this example, we will loop over all treatment levels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "\n",
    "treatment_levels = np.unique(d)\n",
    "thetas = np.full(n_levels + 1, np.nan)\n",
    "ci = np.full((n_levels + 1, 2), np.nan)\n",
    "\n",
    "for i_level, treatment_level in enumerate(treatment_levels):\n",
    "    dml_obj = dml.DoubleMLAPO(\n",
    "        dml_data,\n",
    "        ml_g,\n",
    "        ml_m,\n",
    "        treatment_level=treatment_level,\n",
    "        n_rep=n_rep,\n",
    "    )\n",
    "\n",
    "    dml_obj.fit()\n",
    "\n",
    "    thetas[i_level] = dml_obj.coef[0]\n",
    "    ci[i_level, :] = dml_obj.confint(level=0.95).values\n",
    "\n",
    "# combine results\n",
    "df_apo_ci = pd.DataFrame(\n",
    "    {'treatment_level': treatment_levels,\n",
    "     'apo': apos,\n",
    "     'theta': thetas,\n",
    "     'ci_lower': ci[:, 0],\n",
    "     'ci_upper': ci[:, 1]}\n",
    ")\n",
    "\n",
    "df_apo_ci"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The tables above displays the estimated values in the ``theta`` column and the corresponding oracle values in the ``apo`` column. \n",
    "\n",
    "Again, let us summarize the results in a plot of the APOs with confidence intervals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting\n",
    "plt.figure(figsize=(10, 6))\n",
    "# Plot Estimate with 95% CI\n",
    "plt.errorbar(df_apo_ci['treatment_level'], df_apo_ci['theta'], \n",
    "             yerr=[df_apo_ci['theta'] - df_apo_ci['ci_lower'], df_apo_ci['ci_upper'] - df_apo_ci['theta']], \n",
    "             fmt='o', capsize=5, capthick=2, ecolor=palette[1], color=palette[0], label='Estimate with 95% CI', zorder=2)\n",
    "# Plot APO as a scatter plot, with zorder set to 2 to be in front\n",
    "plt.scatter(df_apo_ci['treatment_level'], df_apo_ci['apo'], color=palette[2], label='APO', marker='d', zorder=3)\n",
    "\n",
    "plt.title('Estimated APO, Theta, and 95% Confidence Interval by Treatment Level')\n",
    "plt.xlabel('Treatment Level')\n",
    "plt.ylabel('Value')\n",
    "plt.xticks(df_apo_ci['treatment_level'])\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple Average Potential Outcome Models (APOS)\n",
    "\n",
    "Instead of looping over different treatment levels, one can directly use the [DoubleMLAPOS](https://docs.doubleml.org/dev/api/generated/doubleml.DoubleMLAPOS.html#doubleml.DoubleMLAPOS) model which internally combines multiple [DoubleMLAPO](https://docs.doubleml.org/dev/api/generated/doubleml.DoubleMLAPO.html#doubleml.DoubleMLAPO) models. An advantage of this approach is that the model can be parallelized, create joint confidence intervals and allow for a comparison between the average potential outcome levels.\n",
    "\n",
    "### Average Potential Outcome (APOs)\n",
    "\n",
    "As before, we just have to specify the machine learning algorithms and the treatment levels for which the APOs should be estimated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_obj = dml.DoubleMLAPOS(\n",
    "    dml_data,\n",
    "    ml_g,\n",
    "    ml_m,\n",
    "    treatment_levels=treatment_levels,\n",
    "    n_rep=n_rep,\n",
    ")\n",
    "\n",
    "dml_obj.fit()\n",
    "\n",
    "ci_pointwise = dml_obj.confint(level=0.95)\n",
    "\n",
    "df_apos_ci = pd.DataFrame(\n",
    "    {'treatment_level': treatment_levels,\n",
    "     'apo': apos,\n",
    "     'theta': thetas,\n",
    "     'ci_lower': ci_pointwise.values[:, 0],\n",
    "     'ci_upper': ci_pointwise.values[:, 1]}\n",
    ")\n",
    "\n",
    "df_apos_ci"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, let us summarize the results in a plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting\n",
    "plt.figure(figsize=(10, 6))\n",
    "# Plot Estimate with 95% CI\n",
    "plt.errorbar(df_apos_ci['treatment_level'], df_apos_ci['theta'], \n",
    "             yerr=[df_apos_ci['theta'] - df_apos_ci['ci_lower'], df_apos_ci['ci_upper'] - df_apos_ci['theta']], \n",
    "             fmt='o', capsize=5, capthick=2, ecolor=palette[1], color=palette[0], label='Estimate with 95% CI', zorder=2)\n",
    "# Plot APO as a scatter plot, with zorder set to 2 to be in front\n",
    "plt.scatter(df_apos_ci['treatment_level'], df_apos_ci['apo'], color=palette[2], label='APO', marker='d', zorder=3)\n",
    "\n",
    "plt.title('Estimated APO, Theta, and 95% Confidence Interval by Treatment Level')\n",
    "plt.xlabel('Treatment Level')\n",
    "plt.ylabel('Value')\n",
    "plt.xticks(df_apos_ci['treatment_level'])\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For joint confidence intervals, the ``bootstrap`` method can be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_obj.bootstrap(n_rep_boot=2000)\n",
    "ci_joint = dml_obj.confint(level=0.95, joint=True)\n",
    "\n",
    "ci_joint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sensitivity Analysis\n",
    "\n",
    "For [DoubleMLAPO](https://docs.doubleml.org/dev/api/generated/doubleml.DoubleMLAPO.html#doubleml.DoubleMLAPO) and [DoubleMLAPOS](https://docs.doubleml.org/dev/api/generated/doubleml.DoubleMLAPOS.html#doubleml.DoubleMLAPOS) model all methods for sensitivity analysis are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_obj.sensitivity_analysis()\n",
    "print(dml_obj.sensitivity_summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, ``sensitvity_benchmark`` can be used. In this example we benchmark covariate ``x4`` which does not affect treatment $D$ or outcome $Y$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_obj.sensitivity_benchmark(benchmarking_set=['x4'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " For more details on the sensitivity analysis, please refer to the [User Guide](https://docs.doubleml.org/stable/guide/sensitivity.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Causal Contrasts\n",
    "\n",
    "The [DoubleMLAPOS](https://docs.doubleml.org/dev/api/generated/doubleml.DoubleMLAPOS.html#doubleml.DoubleMLAPOS) model also allows for the estimation of causal contrasts. \n",
    "The contrast is defined as the difference in the average potential outcomes between the treatment levels $d_i$ and $d_j$ where\n",
    "\n",
    "$$ \\theta_{0,ij} = \\mathbb{E}[Y(d_i)] - \\mathbb{E}[Y(d_{j})]$$\n",
    "\n",
    "and will be calculated for all defined treatment levels $i$ and reference levels $j$.\n",
    "\n",
    "In this example, we will estimate the causal contrast between the treatment level $0$ and all other treatment levels, as the treatment level $0$ corresponds to no treatment at all whereas the the other levels are based on the treatment dosage.\n",
    "\n",
    "Therefore we have to specify ``reference_levels=0``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "causal_contrast_model = dml_obj.causal_contrast(reference_levels=0)\n",
    "print(causal_contrast_model.summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let us summarize the results in a plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ates = causal_contrast_model.thetas\n",
    "ci_ates = causal_contrast_model.confint(level=0.95)\n",
    "\n",
    "df_ates = pd.DataFrame(\n",
    "    {'treatment_level': treatment_levels[1:],\n",
    "     'ate': ates,\n",
    "     'ci_lower': ci_ates.iloc[:, 0].values,\n",
    "     'ci_upper': ci_ates.iloc[:, 1].values}\n",
    ")\n",
    "\n",
    "# Plotting\n",
    "plt.figure(figsize=(10, 6))\n",
    "# Plot Estimate with 95% CI\n",
    "plt.errorbar(df_ates['treatment_level'], df_ates['ate'],\n",
    "             yerr=[df_ates['ate'] - df_ates['ci_lower'], df_ates['ci_upper'] - df_ates['ate']], \n",
    "             fmt='o', capsize=5, capthick=2, ecolor=palette[1], color=palette[0], label='Estimate with 95% CI', zorder=2)\n",
    "# Plot APO as a scatter plot, with zorder set to 2 to be in front\n",
    "plt.scatter(df_apos_ci['treatment_level'][1:], average_ites[1:], color=palette[2], label='ATE', marker='d', zorder=3)\n",
    "\n",
    "plt.title('Estimated ATE, Theta, and 95% Confidence Interval by Treatment Level')\n",
    "plt.xlabel('Treatment Level')\n",
    "plt.ylabel('Value')\n",
    "plt.xticks(df_apos_ci['treatment_level'])\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The methods ``sensitivity_analysis`` and ``sensitivity_plot`` are also available for the causal contrasts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "causal_contrast_model.sensitivity_analysis()\n",
    "print(causal_contrast_model.sensitivity_summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an example see the ``sensitivity_plot`` for the first causal contrast ``1.0 vs 0.0``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "causal_contrast_model.sensitivity_plot(idx_treatment=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Benchmarking is not available for causal contrasts."
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}