experiments/experiments_00_km_plots_tte_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "58997b8a",
   "metadata": {},
   "source": [
    "## Generate Kaplan-Meier plots for standard time-to-event regressors\n",
    "\n",
    " - Figure 1 in the paper\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "733120b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from lifelines import KaplanMeierFitter\n",
    "from synthcity.plugins import Plugins\n",
    "from synthcity.plugins.core.models.time_to_event.loader import get_model_template\n",
    "from synthcity.utils.serialization import load_from_file, save_to_file\n",
    "\n",
    "from datasets import get_dataset\n",
    "\n",
    "out_dir = Path(\"output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e07708",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecdd6580",
   "metadata": {},
   "outputs": [],
   "source": [
    "tte_models = [\n",
    "    \"cox_ph\",\n",
    "    \"weibull_aft\",\n",
    "    \"survival_xgboost\",\n",
    "    \"deephit\",\n",
    "    \"date\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b497b5a1",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.image as mpimg\n",
    "import matplotlib.pyplot as plt\n",
    "from lifelines import KaplanMeierFitter\n",
    "from synthcity.utils.serialization import load_from_file, save_to_file\n",
    "\n",
    "out_dir = Path(\"output\")\n",
    "plt.style.use(\"seaborn-whitegrid\")\n",
    "\n",
    "\n",
    "def plot_tte(title, X, T, E, models, **kwargs):\n",
    "    model_names = \"_\".join([model.name() for model in models])\n",
    "    fig_name = Path(f\"diagrams/tte_km_plot_grouped_{title}_{model_names}.pdf\")\n",
    "    fig_name_png = Path(f\"diagrams/tte_km_plot_grouped_{title}_{model_names}.png\")\n",
    "\n",
    "    fig, axs = plt.subplots(\n",
    "        1, len(models), figsize=(4 * len(models), 4), constrained_layout=True\n",
    "    )\n",
    "\n",
    "    for idx, model in enumerate(models):\n",
    "        model_bkp = out_dir / f\"tte_km_plot_model_{title}_{model.name()}\"\n",
    "        ax = axs[idx]\n",
    "\n",
    "        if model_bkp.exists():\n",
    "            preds = load_from_file(model_bkp)\n",
    "        else:\n",
    "            try:\n",
    "                X = X[T > 0]\n",
    "                E = E[T > 0]\n",
    "                T = T[T > 0]\n",
    "\n",
    "                model.fit(X, T, E)\n",
    "                preds = model.predict(X)\n",
    "                save_to_file(model_bkp, preds)\n",
    "            except BaseException as e:\n",
    "                raise\n",
    "                print(\"model failed\", e)\n",
    "                continue\n",
    "\n",
    "        kmf = KaplanMeierFitter()\n",
    "        kmf.fit(T, E, label=\"Real dataset\")\n",
    "        ax = kmf.plot(ax=ax)\n",
    "\n",
    "        model_name = model.name()\n",
    "        if model_name == \"random_survival_forest\":\n",
    "            model_name = \"rsf\"\n",
    "\n",
    "        syn_kmf = KaplanMeierFitter()\n",
    "        syn_kmf.fit(preds, [1] * len(preds), label=f\"Time-to-event model: {model_name}\")\n",
    "        syn_kmf.plot(ax=ax)\n",
    "\n",
    "        ax.axvline(T[E == 1].max(), color=\"r\", linestyle=\"--\")  # vertical\n",
    "        ax.set_xlabel(\"\", fontsize=14)\n",
    "\n",
    "    axs[0].set_ylabel(\"Event probability\", fontsize=14)\n",
    "    plt.suptitle(f\"Dataset: {title}\", fontsize=18)\n",
    "    fig.text(0.5, -0.05, \"Days to event\", ha=\"center\", fontsize=14)\n",
    "\n",
    "    plt.savefig(fig_name)\n",
    "    plt.savefig(fig_name_png)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f10b1a9d",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for dataset in [\"aids\"]:\n",
    "    models = []\n",
    "    for n in tte_models:\n",
    "        models.append(get_model_template(n)())\n",
    "    df, duration_col, event_col, time_horizons = get_dataset(dataset)\n",
    "\n",
    "    Xcov = df.drop(columns=[duration_col, event_col])\n",
    "    T = df[duration_col]\n",
    "    E = df[event_col]\n",
    "\n",
    "    plot_tte(dataset, Xcov, T, E, models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62d26a69",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py38",
   "language": "python",
   "name": "py38"
  },
  "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}