notebook2_R.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dancing with censored data: How to survive with explainable survival analysis?\n",
    "**ML in PL Conference 2023**\n",
    "\n",
    "\n",
    "Mateusz Krzyziński, Mikołaj Spytek (**MI2.AI**)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agenda\n",
    "\n",
    "- 9:00 - 9:15 - Introduction, technicalities\n",
    "- 9:15 - 10:00 - ML in survival analysis, creating models `(Python and R)`\n",
    "- 10:00 - 10:10 - *break*\n",
    "\n",
    "- 10:10 - 10:20 - XAI in survival analysis\n",
    "- 10:20 - 10:45 - SurvLIME `(Python and R)`\n",
    "- 10:45 - 11:30 - SurvSHAP(t) and its aggregations `(Python and R)`\n",
    "- 11:30 - 11:40 - *break*\n",
    "\n",
    "- 11:40 - 11:50 - creating explainers for any models `(R with Python via reticulate)`\n",
    "- 11:50 - 12:15 - global performance & global variable importance explanations `(R)`\n",
    "- 12:15 - 12:30 - local variable dependence explanations `(R)`\n",
    "- 12:30 - 12:45 - global variable dependence explanations `(R)`\n",
    "- 12:45 - 13:00 - summary and Q&A\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why we need explanations for survival models?\n",
    "\n",
    "**Complex machine learning survival models are used more and more often (also in healthcare and scientific research) because**:\n",
    "- Cox PH model cannot model complex dependencies and has the strong assumption of proportional hazards that is often not met,\n",
    "- ML models are more flexible and can model complex dependencies, but they are often treated as black boxes.\n",
    "\n",
    "\n",
    "<img src=\"images/need_quotes.png\" width=\"1000\">\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "**More about medical context**\n",
    "\n",
    "The **overoptimistic** use of AI models in medicine **→** The need for a method of **validation** other than just performance validation.\n",
    "> There exists a notable tendency among some researchers steeped in ML to report only the model performance metrics. XAI techniques offer a means to validate and explore survival models, effectively complementing performance validation. They enable to shift of the focus towards model reliability and accountability rather than solely optimizing model performance metrics. This shift is crucial in mitigating the risk of over-optimistic assessments of complex models, thereby promoting responsible and informed model usage.\n",
    "\n",
    "The **complexity** and **lack of interpretability** of AI models **→** The need for a method of **examining** models that would make it possible to understand their operation.\n",
    "> Researchers rely on survival models to analyze the effectiveness of new therapies, explore genetic markers associated with diseases, or evaluate the impact of lifestyle factors on health outcomes. Given the vast implications of these models, ensuring that their predictions are not only accurate but also interpretable becomes a matter of utmost importance\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## XAI as model exploration stack\n",
    "\n",
    "- There are many different types of (survival) models, and each of them has its own structure. \n",
    "- Altough model-specific approaches are possible, it is much more convenient to use a model-agnostic approach. \n",
    "- In this approach, the model is wrapped in an abstract object, and the explanation methods are applied to this object (they are usually based on the analysis of the model's predictions). \n",
    "\n",
    "<img src=\"images/piramide.png\" width=\"800\">\n",
    "\n",
    "**Read more:** https://ema.drwhy.ai\n",
    "\n",
    "\n",
    "## Tools\n",
    "- In R, we will use the [`survex`](https://modeloriented.github.io/survex/) package that follows similar principles.\n",
    "<br/>\n",
    "<img src=\"images/survex.png\" width=\"800\"><br/>\n",
    "Spytek, M., Krzyziński, M., Langbein, S. H., Baniecki, H., Wright, M. N., & Biecek, P. (2023). **survex: an R package for explaining machine learning survival models**. arXiv preprint arXiv:2308.16113. https://arxiv.org/abs/2308.16113\n",
    "\n",
    "- In Python, we will use the [`survlimepy`](https://github.com/imatge-upc/SurvLIMEpy) and [`survshap`](https://github.com/MI2DataLab/survshap) packages. \n",
    "<br/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "library(survival)\n",
    "library(survex)\n",
    "library(ranger)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "options(jupyter.plot_scale=1, repr.plot.width = 6, repr.plot.height = 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "df <- read.csv(\"datasets/lung_dataset.csv\")\n",
    "df <- df[complete.cases(df),]\n",
    "head(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "cph_model <- coxph(Surv(time, status) ~ ., data = df, x = TRUE, model = TRUE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "set.seed(123)\n",
    "rsf_model <- ranger(Surv(time, status) ~ ., data = df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating explainers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For some models explainers are created fully automatically, extracting all necessary information along with the predict functions from the model object.\n",
    "\n",
    "For example, for Cox PH model created with `survival` package, we have to only set `x = TRUE` and `model = TRUE` in `survival::coxph` function when creating model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "cph_explainer <- explain(cph_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For other models, we have to provide additional informations to `explain` function, for example dataset. \n",
    "\n",
    "Of course, dataset always can be provided manually. It can be useful when we want to use test set or specific subset of data for creating explanations (for example to answer different questions - context in explanations is crucial).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "rsf_explainer <- explain(rsf_model,\n",
    "                        data = df[3:9], \n",
    "                        y = Surv(df$time, df$status))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "#?explain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "cph_explainer$times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SurvLIME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- based on LIME (Local Interpretable Model-agnostic Explanations) method \n",
    "- local explanation method (for one observation & prediction)\n",
    "- **main idea:** approximate the black-box model locally with a simple interpretable model - apply the Cox proportional hazards model to approximate the black-box model\n",
    "- **the explanation:** coefficients of the locally fitted Cox model\n",
    "- those coefficients are calculated based on distance between cumulative hazard functions of black-box model and surrogate model\n",
    "$\\min_{\\mathbf{b}}\\sum_{k=1}^{N}w_{k}\\sum_{j=0}^{m}v_{kj}^{2}\\left(  \\ln\n",
    "H_{j}(\\mathbf{x}_{k})-\\ln H_{0j}-\\mathbf{b}^{\\mathrm{T}}\\mathbf{x}_{k}\\right)\n",
    "^{2}\\left(  t_{j+1}-t_{j}\\right)$\n",
    "\n",
    "<img src=\"images/survlime.png\" width=\"600\"/>\n",
    "\n",
    "Kovalev, M. S., Utkin, L. V., & Kasimov, E. M. (2020). **SurvLIME: A method for explaining machine learning survival models**. Knowledge-Based Systems, 203, 106164. https://doi.org/10.1016/j.knosys.2020.106164\n",
    "Pachón-García, C., Hernández-Pérez, C., Delicado, P., & Vilaplana, V. (2024). **SurvLIMEpy: A Python package implementing SurvLIME**. Expert Systems with Applications, 237, 121620. https://doi.org/10.1016/j.eswa.2023.121620\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "patientA <- df[42, 3:9]\n",
    "patientA "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "set.seed(42)\n",
    "survlime_cph_patientA <- predict_parts(cph_explainer, patientA, type = \"survlime\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Default results (coefficients)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survlime_cph_patientA, type = \"coefficients\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "survlime_cph_patientA$result "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local importance (coefficients multiplied by features values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survlime_cph_patientA, type = \"local_importance\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters\n",
    "- kernel width - how to calculate weights for observations generated as neighbours \n",
    "- number of neighbours - how many neighbours to generate\n",
    "- which variables should be treated as categorical (different perturbation method)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "set.seed(42)\n",
    "survlime_cph_patientA_kw5 <- predict_parts(cph_explainer, patientA, type = \"survlime\", kernel_width = 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survlime_cph_patientA_kw5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "set.seed(42)\n",
    "survlime_cph_patientA_catvars <- predict_parts(cph_explainer, patientA, type = \"survlime\",\n",
    "                                    categorical_variables = c(\"sex\", \"ph.ecog\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survlime_cph_patientA_catvars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problems for non-linear models\n",
    "- Predictions from surrogate model are often far from predictions of black-box model. The reason is that the surrogate model is too simple to approximate more complex black-box models. \n",
    "- Explanations are not stable. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "set.seed(42)\n",
    "survlime_rsf_patientA <- predict_parts(rsf_explainer, patientA, type = \"survlime\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survlime_rsf_patientA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "set.seed(12)\n",
    "survlime_rsf_patientA <- predict_parts(rsf_explainer, patientA, type = \"survlime\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survlime_rsf_patientA)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SurvSHAP(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- based on SHAP (SHapley Additive exPlanations) method \n",
    "- local explanation method (for one observation & prediction)\n",
    "- **main idea:** allow for time-dependent explainability better suited to complex models\n",
    "- **the explanation:** curves showing the attribution of each feature to the prediction over time (can be aggregated to single numbers)\n",
    "\n",
    "- **formula:**\n",
    "$$\\phi_{t}(\\mathbf{x}_*, d) = \\frac{1}{|\\Pi|} \\sum_{\\pi \\in \\Pi} e_{t, \\mathbf{x}_*}^{\\mathrm{before}(\\pi, d) \\cup \\{d\\}} - e_{t, \\mathbf{x}_*}^{\\mathrm{before}(\\pi, d)},$$ \n",
    "where $\\Pi$ is a set of all permutations of $p$ variables and $\\mathrm{before}(\\pi, d)$ denotes a subset of predictors that are before $d$ in the ordering $\\pi \\in \\Pi$.\n",
    "\n",
    "- **local variable importance:**\n",
    "$    \\psi(\\textbf{x}_*, d) = \\int_0^{t_m} \\left| \\phi_{t}(\\textbf{x}_*, d)\\right| \\, \\mathrm{d}w(t) $\n",
    "\n",
    "<img src=\"images/survshap.png\" width=\"800\"/>\n",
    "\n",
    "\n",
    "Krzyziński, M., Spytek, M., Baniecki, H., & Biecek, P. (2023). **SurvSHAP(t): Time-dependent explanations of machine learning survival models**. Knowledge-Based Systems, 262, 110234. https://doi.org/10.1016/j.knosys.2022.110234"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "set.seed(42)\n",
    "survshap_cph_patientA <- predict_parts(cph_explainer, patientA, type = \"survshap\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_cph_patientA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "survshap_rsf_patientA <- predict_parts(rsf_explainer, patientA, type = \"survshap\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_patientA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "print(survshap_rsf_patientA$aggregate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters\n",
    "- number of samples in the background dataset \n",
    "- aggregation method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "survshap_rsf_patientA_N50 <- predict_parts(rsf_explainer, patientA, type = \"survshap\", N=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_patientA_N50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using other forms of predictions   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "survshap_rsf_patientA_shap_risk <- predict_parts(rsf_explainer, patientA, type = \"shap\", output_type = \"risk\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_patientA_shap_risk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "survshap_rsf_patientA_shap_chf <- predict_parts(rsf_explainer, patientA, type = \"survshap\", output_type = \"chf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_patientA_shap_chf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SurvSHAP(t) aggregations \n",
    "- global (model) explanations based on local explanations\n",
    "- can be calculated for each model but it is computationally expensive \n",
    "- for tree-based models can be calculated faster using TreeSHAP algorithm (`treeshap` package)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "survshap_rsf_model <- model_survshap(rsf_explainer, rsf_explainer$data, calculation_method = \"treeshap\", verbose = FALSE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_model, geom = \"beeswarm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_model, geom = \"profile\", variable = \"pat.karno\", color_variable = \"sex\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_model, geom = \"curves\", variable = \"pat.karno\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": [
    "plot(survshap_rsf_model, geom = \"curves\", variable = \"pat.karno\", boxplot = TRUE, coef = 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "r"
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "R",
   "language": "R",
   "name": "ir"
  },
  "language_info": {
   "codemirror_mode": "r",
   "file_extension": ".r",
   "mimetype": "text/x-r-source",
   "name": "R",
   "pygments_lexer": "r",
   "version": "4.3.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}