{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "(weibull_aft)=\n", "\n", "# Reparameterizing the Weibull Accelerated Failure Time Model\n", "\n", ":::{post} January 17, 2023\n", ":tags: censored, survival analysis, weibull\n", ":category: intermediate, how-to\n", ":author: Junpeng Lao, George Ho, Chris Fonnesbeck\n", ":::" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Running on PyMC v5.28.0+58.gf58491a3\n" ] } ], "source": [ "import arviz as az\n", "import numpy as np\n", "import pymc as pm\n", "import pytensor.tensor as pt\n", "\n", "print(f\"Running on PyMC v{pm.__version__}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ":::{include} ../extra_installs.md\n", ":::" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "# These dependencies need to be installed separately from PyMC\n", "import statsmodels.api as sm" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "%config InlineBackend.figure_format = 'retina'\n", "# These seeds are for sampling data observations\n", "RANDOM_SEED = 8927\n", "np.random.seed(RANDOM_SEED)\n", "# Set a seed for reproducibility of posterior results\n", "seed: int = sum(map(ord, \"aft_weibull\"))\n", "rng: np.random.Generator = np.random.default_rng(seed=seed)\n", "az.style.use(\"arviz-variat\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Dataset\n", "\n", "The {ref}`previous example notebook on Bayesian parametric survival analysis ` introduced two different accelerated failure time (AFT) models: Weibull and log-linear. In this notebook, we present three different parameterizations of the Weibull AFT model.\n", "\n", "The data set we'll use is the `flchain` R data set, which comes from a medical study investigating the effect of serum free light chain (FLC) on lifespan. Read the full documentation of the data by running:\n", "\n", "`print(sm.datasets.get_rdataset(package='survival', dataname='flchain').__doc__)`." ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "# Fetch and clean data\n", "data = (\n", " sm.datasets.get_rdataset(package=\"survival\", dataname=\"flchain\")\n", " .data.sample(500) # Limit ourselves to 500 observations\n", " .reset_index(drop=True)\n", ")" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "y = data.futime.values\n", "censored = ~data[\"death\"].values.astype(bool)" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([ 975, 2272, 138, 4262, 4928])" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "y[:5]" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([False, True, False, True, True])" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "censored[:5]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using `pm.Potential`\n", "\n", "We have an unique problem when modelling censored data. Strictly speaking, we don't have any _data_ for censored values: we only know the _number_ of values that were censored. How can we include this information in our model?\n", "\n", "One way do this is by making use of `pm.Potential`. The [PyMC2 docs](https://pymc-devs.github.io/pymc/modelbuilding.html#the-potential-class) explain its usage very well. Essentially, declaring `pm.Potential('x', logp)` will add `logp` to the log-likelihood of the model. \n", "\n", "However, `pm.Potential` only effect probability based sampling this excludes using `pm.sample_prior_predictice` and `pm.sample_posterior_predictive`. We can overcome these limitations by using `pm.Censored` instead. We can model our right-censored data by defining the `upper` argument of `pm.Censored`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Parameterization 1\n", "\n", "This parameterization is an intuitive, straightforward parameterization of the Weibull survival function. This is probably the first parameterization to come to one's mind." ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "# normalize the event time between 0 and 1\n", "y_norm = y / np.max(y)" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "# If censored then observed event time else maximum time\n", "right_censored = [x if x > 0 else np.max(y_norm) for x in y_norm * censored]" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "with pm.Model() as model_1:\n", " alpha_sd = 1.0\n", "\n", " mu = pm.Normal(\"mu\", mu=0, sigma=1)\n", " alpha_raw = pm.Normal(\"a0\", mu=0, sigma=0.1)\n", " alpha = pm.Deterministic(\"alpha\", pt.exp(alpha_sd * alpha_raw))\n", " beta = pm.Deterministic(\"beta\", pt.exp(mu / alpha))\n", " beta_backtransformed = pm.Deterministic(\"beta_backtransformed\", beta * np.max(y))\n", "\n", " latent = pm.Weibull.dist(alpha=alpha, beta=beta)\n", " y_obs = pm.Censored(\"Censored_likelihood\", latent, upper=right_censored, observed=y_norm)" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "87c3404576f84e2983289b926051d759", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/2000 [00:00" ] }, "metadata": { "image/png": { "height": 860, "width": 2423 } }, "output_type": "display_data" } ], "source": [ "az.plot_trace_dist(idata_param1, var_names=[\"alpha\", \"beta\"]);" ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
meansdeti89_lbeti89_ubess_bulkess_tailr_hatmcse_meanmcse_sd
alpha0.970.060.871.073351.302885.051.00.000.00
beta2.860.352.363.482504.322343.841.00.010.01
beta_backtransformed14652.301804.0912126.5617870.692504.322343.841.035.9828.83
\n", "
" ], "text/plain": [ " mean sd eti89_lb eti89_ub ess_bulk \\\n", "alpha 0.97 0.06 0.87 1.07 3351.30 \n", "beta 2.86 0.35 2.36 3.48 2504.32 \n", "beta_backtransformed 14652.30 1804.09 12126.56 17870.69 2504.32 \n", "\n", " ess_tail r_hat mcse_mean mcse_sd \n", "alpha 2885.05 1.0 0.00 0.00 \n", "beta 2343.84 1.0 0.01 0.01 \n", "beta_backtransformed 2343.84 1.0 35.98 28.83 " ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "az.summary(idata_param1, var_names=[\"alpha\", \"beta\", \"beta_backtransformed\"], round_to=2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Parameterization 2\n", "\n", "Note that, confusingly, `alpha` is now called `r`, and `alpha` denotes a prior; we maintain this notation to stay faithful to the original implementation in Stan. In this parameterization, we still model the same parameters `alpha` (now `r`) and `beta`.\n", "\n", "For more information, see [this Stan example model](https://github.com/stan-dev/example-models/blob/5e9c5055dcea78ad756a6fb9b3ff9a77a0a4c22b/bugs_examples/vol1/kidney/kidney.stan) and [the corresponding documentation](https://www.mrc-bsu.cam.ac.uk/wp-content/uploads/WinBUGS_Vol1.pdf)." ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [], "source": [ "with pm.Model() as model_2:\n", " alpha = pm.Normal(\"alpha\", mu=0, sigma=1)\n", " r = pm.Gamma(\"r\", alpha=2, beta=1)\n", " beta = pm.Deterministic(\"beta\", pt.exp(-alpha / r))\n", " beta_backtransformed = pm.Deterministic(\"beta_backtransformed\", beta * np.max(y))\n", "\n", " latent = pm.Weibull.dist(alpha=r, beta=beta)\n", " y_obs = pm.Censored(\"Censored_likelihood\", latent, upper=right_censored, observed=y_norm)" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "0a3fb09c7e4c4bf3ad95af5606c70452", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/2000 [00:00" ] }, "metadata": { "image/png": { "height": 860, "width": 2423 } }, "output_type": "display_data" } ], "source": [ "az.plot_trace_dist(idata_param2, var_names=[\"r\", \"beta\"]);" ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
meansdeti89_lbeti89_ubess_bulkess_tailr_hatmcse_meanmcse_sd
r0.950.080.831.082634.92481.781.00.000.00
beta2.970.452.343.742406.42152.161.00.010.01
beta_backtransformed15217.572309.3312008.0719184.762406.42152.161.047.9147.51
\n", "
" ], "text/plain": [ " mean sd eti89_lb eti89_ub ess_bulk \\\n", "r 0.95 0.08 0.83 1.08 2634.9 \n", "beta 2.97 0.45 2.34 3.74 2406.4 \n", "beta_backtransformed 15217.57 2309.33 12008.07 19184.76 2406.4 \n", "\n", " ess_tail r_hat mcse_mean mcse_sd \n", "r 2481.78 1.0 0.00 0.00 \n", "beta 2152.16 1.0 0.01 0.01 \n", "beta_backtransformed 2152.16 1.0 47.91 47.51 " ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "az.summary(idata_param2, var_names=[\"r\", \"beta\", \"beta_backtransformed\"], round_to=2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Parameterization 3\n", "\n", "In this parameterization, we model the log-linear error distribution with a Gumbel distribution instead of modelling the survival function directly. For more information, see [this blog post](http://austinrochford.com/posts/2017-10-02-bayes-param-survival.html)." ] }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [], "source": [ "logtime = np.log(y)" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [], "source": [ "# If censored then observed event time else maximum time\n", "right_censored = [x if x > 0 else np.max(logtime) for x in logtime * censored]" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [], "source": [ "with pm.Model() as model_3:\n", " s = pm.HalfNormal(\"s\", tau=3.0)\n", " gamma = pm.Normal(\"gamma\", mu=0, sigma=5)\n", "\n", " latent = pm.Gumbel.dist(mu=gamma, beta=s)\n", " y_obs = pm.Censored(\"Censored_likelihood\", latent, upper=right_censored, observed=logtime)" ] }, { "cell_type": "code", "execution_count": 21, "metadata": { "scrolled": false }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "51e63c11c9464bcb923f878e070ddd68", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/6000 [00:00" ] }, "metadata": { "image/png": { "height": 860, "width": 2423 } }, "output_type": "display_data" } ], "source": [ "az.plot_trace_dist(idata_param3);" ] }, { "cell_type": "code", "execution_count": 23, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
meansdeti89_lbeti89_ubess_bulkess_tailr_hatmcse_meanmcse_sd
gamma9.480.219.149.833328.874221.941.00.00.0
s3.540.163.293.813256.124251.451.00.00.0
\n", "
" ], "text/plain": [ " mean sd eti89_lb eti89_ub ess_bulk ess_tail r_hat mcse_mean \\\n", "gamma 9.48 0.21 9.14 9.83 3328.87 4221.94 1.0 0.0 \n", "s 3.54 0.16 3.29 3.81 3256.12 4251.45 1.0 0.0 \n", "\n", " mcse_sd \n", "gamma 0.0 \n", "s 0.0 " ] }, "execution_count": 23, "metadata": {}, "output_type": "execute_result" } ], "source": [ "az.summary(idata_param3, round_to=2)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Authors\n", "\n", "- Originally collated by [Junpeng Lao](https://junpenglao.xyz/) on Apr 21, 2018. See original code [here](https://github.com/junpenglao/Planet_Sakaar_Data_Science/blob/65447fdb431c78b15fbeaef51b8c059f46c9e8d6/PyMC3QnA/discourse_1107.ipynb).\n", "- Authored and ported to Jupyter notebook by [George Ho](https://eigenfoo.xyz/) on Jul 15, 2018.\n", "- Updated for compatibility with PyMC v5 by Chris Fonnesbeck on Jan 16, 2023.\n", "- Updated to replace `pm.Potential` with `pm.Censored` by Jonathan Dekermanjian on Nov 25, 2024.\n", "- Updated by Osvaldo Martin in April 2026." ] }, { "cell_type": "code", "execution_count": 24, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Last updated: Sat, 25 Apr 2026\n", "\n", "Python implementation: CPython\n", "Python version : 3.14.4\n", "IPython version : 9.12.0\n", "\n", "arviz : 1.1.0\n", "numpy : 2.4.4\n", "pymc : 5.28.0+58.gf58491a3\n", "pytensor : 2.38.0+133.g80cc113b5\n", "statsmodels: 0.14.6\n", "\n", "Watermark: 2.6.0\n", "\n" ] } ], "source": [ "%load_ext watermark\n", "%watermark -n -u -v -iv -w" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ ":::{include} ../page_footer.md\n", ":::" ] } ], "metadata": { "kernelspec": { "display_name": "arviz_1", "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.14.4" }, "myst": { "substitutions": { "extra_dependencies": "statsmodels" } }, "toc": { "base_numbering": 1, "nav_menu": {}, "number_sections": true, "sideBar": true, "skip_h1_title": false, "title_cell": "Table of Contents", "title_sidebar": "Contents", "toc_cell": false, "toc_position": {}, "toc_section_display": true, "toc_window_display": false }, "varInspector": { "cols": { "lenName": 16, "lenType": 16, "lenVar": 40 }, "kernels_config": { "python": { "delete_cmd_postfix": "", "delete_cmd_prefix": "del ", "library": "var_list.py", "varRefreshCmd": "print(var_dic_list())" }, "r": { "delete_cmd_postfix": ") ", "delete_cmd_prefix": "rm(", "library": "var_list.r", "varRefreshCmd": "cat(var_dic_list()) " } }, "types_to_exclude": [ "module", "function", "builtin_function_or_method", "instance", "_Feature" ], "window_display": false } }, "nbformat": 4, "nbformat_minor": 4 }