60_AFT_models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.aft"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accelerated Failure Time Models\n",
    "> AFT Model theory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can model the time to failure as:\n",
    "$$\n",
    "\\log T_i = \\mu + \\xi_i\n",
    "$$\n",
    "where $\\xi_i\\sim p(\\xi|\\theta)$ and $\\mu$ is the most likely log time of death (the mode of the distribution of $T_i$). We model log death as that way we do not need to restrict $\\mu + \\xi_i$ to be positive.\n",
    "\n",
    "In the censored case, where $t_i$ is the time where an instance was censored, and $T_i$ is the unobserved time of death, we have:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\log T_i &= \\mu(x_i) + \\xi_i > \\log t_i\\\\\n",
    "\\therefore \\xi_i &> \\log t_i - \\mu(x_i)\n",
    "\\end{aligned}\n",
    "$$\n",
    "Note that $\\mu$ is a function of the features $x$. The log likelihood of the data ($\\mathcal{D}$) can then shown to be:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\log p(\\mathcal{D}) = \\sum_{i=1}^N \\mathcal{1}(y_i=1)\\log p(\\xi_i = \\log t_i - \\mu(x_i)) + \\mathcal{1}(y_i=0)\\log p(\\xi_i &> \\log t_i - \\mu(x_i))\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from sklearn.preprocessing import MaxAbsScaler, StandardScaler\n",
    "\n",
    "from torchlife.models.error_dist import get_distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class AFTModel(nn.Module):\n",
    "    \"\"\"\n",
    "    Accelerated Failure Time model\n",
    "    parameters:\n",
    "    - Distribution of which the error is assumed to be\n",
    "    - dim (optional): input dimensionality of variables\n",
    "    - h (optional): number of hidden nodes\n",
    "    \"\"\"\n",
    "    def __init__(self, distribution:str, input_dim:int, h:tuple=()):\n",
    "        super().__init__()\n",
    "        self.logpdf, self.logicdf = get_distribution(distribution)\n",
    "        self.β = nn.Parameter(-torch.rand(1))\n",
    "        self.logσ = nn.Parameter(-torch.rand(1))\n",
    "        \n",
    "        if input_dim > 0:\n",
    "            nodes = (input_dim,) + h + (1,)\n",
    "            self.layers = nn.ModuleList([nn.Linear(a,b, bias=False) \n",
    "                                       for a,b in zip(nodes[:-1], nodes[1:])])\n",
    "\n",
    "        self.eps = 1e-7\n",
    "\n",
    "    def get_mode_time(self, x:torch.Tensor=None):\n",
    "        μ = self.β\n",
    "        if x is not None:\n",
    "            for layer in self.layers[:-1]:\n",
    "                x = F.relu(layer(x))\n",
    "            μ = self.β + self.layers[-1](x)\n",
    "\n",
    "        σ = torch.exp(self.logσ)\n",
    "        return μ, σ\n",
    "    \n",
    "    def forward(self, t:torch.Tensor, x:torch.Tensor=None):\n",
    "        μ, σ = self.get_mode_time(x)\n",
    "        ξ = torch.log(t + self.eps) - μ\n",
    "        logpdf = self.logpdf(ξ, σ)\n",
    "        logicdf = self.logicdf(ξ, σ)\n",
    "        return logpdf, logicdf\n",
    "    \n",
    "    def survival_function(self, t:np.array, t_scaler, x:np.array=None, x_scaler=None):\n",
    "        if len(t.shape) == 1:\n",
    "            t = t[:,None]\n",
    "        t = t_scaler.transform(t)\n",
    "        t = torch.Tensor(t)\n",
    "        if x is not None:\n",
    "            if len(x.shape) == 1:\n",
    "                x = x[None, :]\n",
    "            if len(x) == 1:\n",
    "                x = np.repeat(x, len(t), axis=0)\n",
    "            x = x_scaler.transform(x)\n",
    "            x = torch.Tensor(x)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            # calculate cumulative hazard according to above\n",
    "            _, Λ = self(t, x)\n",
    "            return torch.exp(Λ)\n",
    "            \n",
    "    def plot_survival_function(self, t:np.array, t_scaler, x:np.array=None, x_scaler=None):\n",
    "        surv_fun = self.survival_function(t, t_scaler, x, x_scaler)\n",
    "        \n",
    "        # plot\n",
    "        plt.figure(figsize=(12,5))\n",
    "        plt.plot(t, surv_fun)\n",
    "        plt.xlabel('Time')\n",
    "        plt.ylabel('Survival Probability')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Modelling based on **only** time and (death) event variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from torchlife.data import create_dl\n",
    "# import pandas as pd\n",
    "\n",
    "# url = \"https://raw.githubusercontent.com/vincentarelbundock/Rdatasets/master/csv/survival/flchain.csv\"\n",
    "# df = pd.read_csv(url).iloc[:,1:]\n",
    "# df.rename(columns={'futime':'t', 'death':'e'}, inplace=True)\n",
    "\n",
    "# cols = [\"age\", \"sample.yr\", \"kappa\"]\n",
    "# db, t_scaler, x_scaler = create_dl(df[['t', 'e'] + cols])\n",
    "\n",
    "# death_rate = 100*df[\"e\"].mean()\n",
    "# print(f\"Death occurs in {death_rate:.2f}% of cases\")\n",
    "# print(df.shape)\n",
    "# df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # hide\n",
    "# from fastai.basics import Learner\n",
    "# from torchlife.losses import aft_loss\n",
    "\n",
    "# model = AFTModel(\"Gumbel\", t_scaler, x_scaler)\n",
    "# learner = Learner(db, model, loss_func=aft_loss)\n",
    "# # wd = 1e-4\n",
    "# learner.lr_find(start_lr=1, end_lr=10)\n",
    "# learner.recorder.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# learner.fit(epochs=10, lr=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.plot_survival_function(np.linspace(0, df[\"t\"].max(), 100), df.loc[0, cols])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_index.ipynb.\n",
      "Converted 10_SAT.ipynb.\n",
      "Converted 20_KaplanMeier.ipynb.\n",
      "Converted 30_overall_model.ipynb.\n",
      "Converted 50_hazard.ipynb.\n",
      "Converted 55_hazard.PiecewiseHazard.ipynb.\n",
      "Converted 59_hazard.Cox.ipynb.\n",
      "Converted 60_AFT_models.ipynb.\n",
      "Converted 65_AFT_error_distributions.ipynb.\n",
      "Converted 80_data.ipynb.\n",
      "Converted 90_model.ipynb.\n",
      "Converted 95_Losses.ipynb.\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from nbdev.export import *\n",
    "notebook2script()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}