Direct-Regression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Direct Regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {
    "tags": [
     "hide_cell"
    ]
   },
   "outputs": [],
   "source": [
    "# import libraries\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.linear_model import Lasso, Ridge, LinearRegression\n",
    "from sklearn.ensemble import GradientBoostingRegressor\n",
    "from sklearn.metrics import r2_score\n",
    "import pickle\n",
    "import time\n",
    "from reed import drop_missing_treatment_or_outcome, Model, transform_outcome\n",
    "\n",
    "# set global notebook options\n",
    "pd.options.display.max_columns = 200\n",
    "pd.options.display.max_rows = 500\n",
    "pd.options.display.max_colwidth = 100\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {
    "tags": [
     "overview"
    ]
   },
   "source": [
    "## Response Model\n",
    "\n",
    "How well can we predict outcomes $Y$ conditional on treatment $T$ and other covariates $Z$?\n",
    "\n",
    "### Treatent variables\n",
    "\n",
    "   - **reduhl**\tCompleted re-education based on highest level of attainment\n",
    "   - **redudl**\tCompleted re-education based on detailed qualifications\n",
    "   - **redufl**\tCompleted re-education using highest lvl and detailed qualifications.\n",
    "\n",
    "### Outcome variables\n",
    "   - Mental health in 2019 (**mh**). This is the transformed mental health scores from the aggregation of mental health items of the SF-36 Health Survey, as reported by the individual in 2019. It ranges from 0 to 100, with higher scores indicating better mental health.  \n",
    "   - Working hours in 2019 (**wkhr**) records the total number of hours the individual works in all jobs in a week on average. Working hours are set to 0 for those not working. \n",
    "   - Hourly Wages in 2019 (**rlwage**) records the average hourly wage for the individual’s main job in 2019. Hourly wages are set to 0 for those not working and set to missing for those reporting working more than 100 hours a week. \n",
    "   \n",
    "#### Columns explicitly excluded\n",
    "   - **xwaveid** (unique identifier)\n",
    "   - **p_rcom*** (timing of completion of re-education, proxies treatment) TODO think about how we would include this\n",
    "   - **p_cotrl** (first avail 2003)\n",
    "   - **p_rdf*** (first avail 2012)\n",
    "   \n",
    "### (Nested) cross-validate to evaluate model performance\n",
    "![image.png](images/nested_cross_val.png)\n",
    "\n",
    "### Bootstraped cross-validation to estimate parameter uncertainty\n",
    "![image.png](images/bootstrap_validation.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "#### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "optimisation_metric = 'neg_mean_squared_error'\n",
    "evaluation_metrics = ('r2','neg_mean_squared_error')\n",
    "log_outcome=True\n",
    "standardize_outcome=True\n",
    "load_from_cache=False\n",
    "exclude_patterns = [\n",
    "    '^reduhl$', '^rehllt$', '^redudl$', '^redufl$', '^redllt$', '^refllt$',\n",
    "    '^rlwage$', '^mh$', '^mhbm$', '^wkhr$', '^y_', '^p_rcom','^p_rdf','^p_cotrl',\n",
    "    '^xwaveid$','p_rcom18','^aedcq', '^abnfsty','^aedcqfpt','^aedqstdy'\n",
    "]\n",
    "data_load_func = lambda filepath: pd.read_csv(filepath, index_col='xwaveid')\n",
    "\n",
    "configuration_name = 'default'\n",
    "outcome = 'y_wsce'#'y_wsce'\n",
    "treatment = 'redufl'\n",
    "test=False\n",
    "data_file = \"data/all_lasso_selected_100_restricted.csv\"\n",
    "release = \"restricted\" # only used to determine which metadata file to read. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters that depend on those set above (which may have been inserted by Papermill)\n",
    "if test:\n",
    "    inner_cv = 2\n",
    "    outer_cv = 2\n",
    "    bootstrap_samples = 3\n",
    "\n",
    "else:\n",
    "    inner_cv = 5\n",
    "    outer_cv = 10\n",
    "    bootstrap_samples = 50\n",
    "\n",
    "cross_val_cache = f\"data/cross-val-{configuration_name}.pkl\"#\"data/cross-val-all.pkl\"\n",
    "bootstrap_cache = f\"data/bootstrap-{configuration_name}.pkl\"#\"data/bootstrap-all.pkl\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### Load Data \n",
    "   - drop rows missing the specified treatment or outcome"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "data = data_load_func(data_file)\n",
    "drop_missing_treatment_or_outcome(data, treatment, outcome)\n",
    "data[outcome] = transform_outcome(data[outcome],log_outcome, standardize_outcome)\n",
    "\n",
    "    \n",
    "plt.hist(data[outcome])\n",
    "plt.xlabel(outcome)\n",
    "plt.ylabel(\"count\")\n",
    "plt.title(\"Distribution of outcomes\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Set up models\n",
    "Specify which models to use and the hyper-parameter space to search over for each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from sklearn.linear_model import Lasso, Ridge\n",
    "from sklearn.ensemble import GradientBoostingRegressor\n",
    "from direct_regression import importance_from_coef\n",
    "\n",
    "def construct_models():\n",
    "    models = [\n",
    "        Model('ridge',Ridge(), \n",
    "              parameters = {\n",
    "                  'alpha':np.logspace(-1,4,30)\n",
    "              },\n",
    "              importance_func=importance_from_coef\n",
    "        ),\n",
    "        Model('lasso',Lasso(),\n",
    "              parameters = {\n",
    "                  'alpha':np.logspace(-2,4,30)\n",
    "              },\n",
    "              importance_func=importance_from_coef\n",
    "        ), \n",
    "        Model('gbr',GradientBoostingRegressor(n_iter_no_change=20, max_depth=2),\n",
    "              parameters = {\n",
    "                'max_features':np.linspace(0,1,6)[1:],\n",
    "                'learning_rate':np.logspace(-3,0,10),\n",
    "                'min_samples_leaf':np.logspace(0,3,10).astype(int)\n",
    "              }\n",
    "        ),\n",
    "    ]\n",
    "    return models\n",
    "\n",
    "# used for quick tests\n",
    "def construct_test_models(): \n",
    "    models = [\n",
    "        Model('Ridge',Ridge(),\n",
    "              parameters = {'alpha':np.logspace(1,4,2)},\n",
    "              importance_func=importance_from_coef\n",
    "        ),\n",
    "        Model('gbr',GradientBoostingRegressor(n_iter_no_change=20, max_depth=2),\n",
    "              parameters = {\n",
    "                'max_features':np.linspace(0,1,4)[1:],\n",
    "              }\n",
    "        ),\n",
    "    ]\n",
    "    return models\n",
    "\n",
    "model_init = construct_test_models if test else construct_models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "### Prepare data for modeling\n",
    "   - split into treated/control\n",
    "   - impute missing values and scale\n",
    "   - separate features from outcomes&treatments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from direct_regression import seperate_and_transform_data\n",
    "X0, X1, y0, y1, X, y, t, features = seperate_and_transform_data(data, treatment, outcome)\n",
    "print(\"Control data dimensions: \",X0.shape)\n",
    "print(\"Treated data dimensions:\",X1.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "### Compute unconditional/unadjusted estimate of treatment effect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from direct_regression import print_unconditional_effects\n",
    "print_unconditional_effects(data, treatment, y0, y1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Propensity model\n",
    "\n",
    "- Run primarily to check for violations over overlap assumptions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {
    "tags": [
     "hide_cell"
    ]
   },
   "outputs": [],
   "source": [
    "from reed import visualise_propensity_model_performance\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.linear_model import LogisticRegression, LogisticRegressionCV\n",
    "\n",
    "X_train, X_test, t_train, t_test = train_test_split(X,t,stratify=t)\n",
    "pmodel = LogisticRegressionCV(Cs = np.logspace(-5,-1,20), penalty='l2', max_iter=10000)\n",
    "pmodel.fit(X_train, t_train);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "visualise_propensity_model_performance(t_test, pmodel.predict_proba(X_test)[:,1],'l2-logistic',bins=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "def propensity_coefficients(model, features):\n",
    "    if hasattr(model,'best_estimator_'):\n",
    "        model = model.best_estimator_\n",
    "    coef = model.coef_[0]\n",
    "    \n",
    "    with open(f'data/metadata_{release}.pkl','rb') as f:\n",
    "        meta = pickle.load(f)\n",
    "        labels = [meta.column_names_to_labels.get(f) for f in features]\n",
    "    return pd.DataFrame({'label': labels, \"c\":coef,'abs':np.abs(coef)},index=features).sort_values(by='abs',ascending=False).head(50)\n",
    "print(\"Propensity Model Coefficients\")      \n",
    "propensity_coefficients(pmodel, features).head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {
    "tags": [
     "hide_cell"
    ]
   },
   "outputs": [],
   "source": [
    "from direct_regression import nested_cross_val\n",
    "\n",
    "models0, models1, results = nested_cross_val(\n",
    "    model_init,\n",
    "    cross_val_cache,\n",
    "    X0, X1, y0, y1,\n",
    "    optimisation_metric,\n",
    "    evaluation_metrics,\n",
    "    innercv=inner_cv,\n",
    "    outercv=outer_cv,\n",
    "    load_from_cache=load_from_cache\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "### Report estimate ATE and model performance\n",
    "\n",
    "  - Mean and Std of prediction performance for each model (both treatment & control surface)\n",
    "  - Mean and Std of average treatment effect for each model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "from direct_regression import visualise_ate\n",
    "metrics, ate_vals = visualise_ate(results,X,evaluation_metrics);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "from direct_regression import plot_ate_distribution\n",
    "plot_ate_distribution(ate_vals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "from direct_regression import plot_hyperparam_distributions\n",
    "for model, (results0, results1) in results.items():\n",
    "    plot_hyperparam_distributions(results0,f\"{model}-control\")\n",
    "    plot_hyperparam_distributions(results1,f\"{model}-treated\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "### Visualise models\n",
    "- Features responsible for treatment effect heterogeneity & functional form (with uncertainty)\n",
    "      - coefficeints for linear models\n",
    "      - TODO permutation importance & partial dependence curves for non-linear models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from direct_regression import display_feature_importance\n",
    "display_feature_importance(models0, models1, results, features);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "### Compare against OLS on basic variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "basic = pd.read_csv(\"data/basic_variables.csv\",index_col='xwaveid')\n",
    "drop_missing_treatment_or_outcome(basic, treatment, outcome)\n",
    "basic[outcome] = transform_outcome(basic[outcome],log_outcome, standardize_outcome)\n",
    "\n",
    "    \n",
    "X0b, X1b, y0b, y1b, Xb, yb, tb, featuresb = seperate_and_transform_data(basic, treatment, outcome)\n",
    "\n",
    "def construct_basic_model():\n",
    "    models = [\n",
    "        Model('OLS',LinearRegression(),importance_func=importance_from_coef)\n",
    "    ]\n",
    "    return models\n",
    "\n",
    "modelsb0, modelsb1, resultsb = nested_cross_val(\n",
    "    construct_basic_model,\n",
    "    \"data/tmp.pkl\",\n",
    "    X0b, X1b, y0b, y1b,\n",
    "    optimisation_metric,\n",
    "    evaluation_metrics,\n",
    "    innercv=inner_cv,\n",
    "    outercv=outer_cv,\n",
    "    load_from_cache=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import r2_score\n",
    "y0b = resultsb['OLS'][0]['estimator'][0].predict(X0b)\n",
    "y1b = resultsb['OLS'][1]['estimator'][0].predict(X1b)\n",
    "fig, ax = plt.subplots(1,2, figsize=(15,5))\n",
    "ax[0].scatter(y0, y0b,s=10,alpha=0.1)\n",
    "ax[0].set_xlabel('y0')\n",
    "ax[0].set_ylabel('$\\\\hat{y}0$');\n",
    "ax[0].set_title(f'Control (OLS-basic), $R^2={r2_score(y0,y0b):.2f}$')\n",
    "\n",
    "ax[1].scatter(y1, y1b,s=10,alpha=0.1)\n",
    "ax[1].set_xlabel('y1')\n",
    "ax[1].set_ylabel('$\\\\hat{y}1$');\n",
    "ax[1].set_title(f'Treated (OLS-basic), $R^2={r2_score(y1,y1b):.2f}$');\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "if 'ridge' in results.keys():\n",
    "    y0a = results['ridge'][0]['estimator'][4].predict(X0)\n",
    "    y1a = results['ridge'][1]['estimator'][4].predict(X1)\n",
    "    fig, ax = plt.subplots(1,2, figsize=(15,5))\n",
    "    ax[0].scatter(y0, y0a,s=10,alpha=0.1)\n",
    "    ax[0].set_xlabel('y0')\n",
    "    ax[0].set_ylabel('$\\\\hat{y}0$');\n",
    "    ax[0].set_title(f'Control (ridge), $R^2={r2_score(y0,y0a):.2f}$')\n",
    "\n",
    "    ax[1].scatter(y1, y1a,s=10,alpha=0.1)\n",
    "    ax[1].set_xlabel('y1')\n",
    "    ax[1].set_ylabel('$\\\\hat{y}1$');\n",
    "    ax[1].set_title(f'Treated (ridge), $R^2={r2_score(y1,y1a):.2f}$');\n",
    "\n",
    "# think about why these r2 scores are higher than on the holdout set (and whether they are likely overfitting)\n",
    "# despite there being so few variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {
    "tags": [
     "hide_cell"
    ]
   },
   "outputs": [],
   "source": [
    "from direct_regression import bootstrapped_cross_val\n",
    "def extract_params(estimator):\n",
    "    return estimator.coef_\n",
    "\n",
    "\n",
    "bootstrap_results = bootstrapped_cross_val(\n",
    "    model_init,\n",
    "    bootstrap_cache,\n",
    "    X0, X1, y0, y1,\n",
    "    optimisation_metric,\n",
    "    extract_params,\n",
    "    inner_cv=inner_cv,\n",
    "    load_from_cache=load_from_cache,\n",
    "    samples=bootstrap_samples\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30",
   "metadata": {},
   "source": [
    "####  Average treatment effects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from direct_regression import compute_ate\n",
    "metrics, tau_estimatesb = compute_ate(bootstrap_results,X)\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_ate_distribution(tau_estimatesb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33",
   "metadata": {},
   "source": [
    "#### Distribution of hyper-parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from direct_regression import plot_hyperparam_distributions\n",
    "for model, (results0, results1) in bootstrap_results.items():\n",
    "    plot_hyperparam_distributions(results0,f\"{model}-control\")\n",
    "    plot_hyperparam_distributions(results1,f\"{model}-treated\")"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}