doc/examples/double_ml_bonus_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# DML: Bonus Data\n",
    "This example shows\n",
    "TODO: Add a general description!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import doubleml as dml\n",
    "from doubleml.datasets import fetch_bonus\n",
    "\n",
    "from sklearn.linear_model import Lasso, LogisticRegression\n",
    "from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['figure.figsize'] = 14, 6\n",
    "sns.set()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load bonus data using the dml datasets module\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data = dml.datasets.fetch_bonus()\n",
    "dml_data.data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify learner and estimate causal parameter: PLR model with random forest as learner\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set machine learning methods for m & g\n",
    "ml_g = RandomForestRegressor()\n",
    "ml_m = RandomForestRegressor()\n",
    "n_folds = 2\n",
    "n_rep = 10\n",
    "\n",
    "np.random.seed(3141)\n",
    "dml_plr_rf = dml.DoubleMLPLR(dml_data,\n",
    "                             ml_g,\n",
    "                             ml_m,\n",
    "                             n_folds,\n",
    "                             n_rep,\n",
    "                             'partialling out',\n",
    "                             'dml2')\n",
    "\n",
    "# set some hyperparameters for the learners\n",
    "pars = {'n_estimators': 500,\n",
    "        'max_features': 'sqrt',\n",
    "        'max_depth': 5}\n",
    "\n",
    "dml_plr_rf.set_ml_nuisance_params('ml_g', 'tg', pars)\n",
    "dml_plr_rf.set_ml_nuisance_params('ml_m', 'tg', pars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_plr_rf.fit()\n",
    "dml_plr_rf.summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data with polynomial features\n",
    "dml_data_lasso = dml.datasets.fetch_bonus(polynomial_features=True)\n",
    "print(dml_data_lasso)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify learner and estimate causal parameter: PLR model with Lasso as learner\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set machine learning methods for m & g\n",
    "ml_g = Lasso()\n",
    "ml_m = Lasso()\n",
    "n_folds = 2\n",
    "n_rep = 10\n",
    "\n",
    "np.random.seed(3141)\n",
    "dml_plr_lasso = dml.DoubleMLPLR(dml_data_lasso,\n",
    "                                ml_g,\n",
    "                                ml_m,\n",
    "                                n_folds,\n",
    "                                n_rep,\n",
    "                                'partialling out',\n",
    "                                'dml2')\n",
    "\n",
    "# set some hyperparameters for the learners\n",
    "dml_plr_lasso.set_ml_nuisance_params('ml_g', 'tg', {'alpha': 0.0005})\n",
    "dml_plr_lasso.set_ml_nuisance_params('ml_m', 'tg', {'alpha': 0.0026})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_plr_lasso.fit()\n",
    "dml_plr_lasso.summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify learner and estimate causal parameter: IRM model with random forest as learner\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set machine learning methods for m & g\n",
    "ml_g = RandomForestRegressor()\n",
    "ml_m = RandomForestClassifier()\n",
    "n_folds = 2\n",
    "n_rep = 10\n",
    "\n",
    "np.random.seed(3141)\n",
    "dml_irm_rf = dml.DoubleMLIRM(dml_data,\n",
    "                             ml_g,\n",
    "                             ml_m,\n",
    "                             n_folds,\n",
    "                             n_rep,\n",
    "                             'ATE',\n",
    "                             'dml2')\n",
    "\n",
    "# set some hyperparameters for the learners\n",
    "pars = {'n_estimators': 500,\n",
    "        'max_features': 'sqrt',\n",
    "        'max_depth': 5}\n",
    "\n",
    "dml_irm_rf.set_ml_nuisance_params('ml_g0', 'tg', pars)\n",
    "dml_irm_rf.set_ml_nuisance_params('ml_g1', 'tg', pars)\n",
    "dml_irm_rf.set_ml_nuisance_params('ml_m', 'tg', pars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_irm_rf.fit()\n",
    "dml_irm_rf.summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify learner and estimate causal parameter: IRM model with Lasso as learner\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set machine learning methods for m & g\n",
    "ml_g = Lasso()\n",
    "ml_m = LogisticRegression()\n",
    "np.random.seed(1234)\n",
    "n_folds = 2\n",
    "n_rep = 10\n",
    "\n",
    "np.random.seed(3141)\n",
    "dml_irm_lasso = dml.DoubleMLIRM(dml_data_lasso,\n",
    "                                ml_g,\n",
    "                                ml_m,\n",
    "                                n_folds,\n",
    "                                n_rep,\n",
    "                                'ATE',\n",
    "                                'dml2')\n",
    "\n",
    "# set some hyperparameters for the learners\n",
    "dml_irm_lasso.set_ml_nuisance_params('ml_g0', 'tg', {'alpha': 0.0019})\n",
    "dml_irm_lasso.set_ml_nuisance_params('ml_g1', 'tg', {'alpha': 0.0073})\n",
    "dml_irm_lasso.set_ml_nuisance_params('ml_m', 'tg', {'C': 0.0001})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_irm_lasso.fit()\n",
    "dml_irm_lasso.summary"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}