95_Losses.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp losses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Losses\n",
    "> All the losses used in SA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from abc import ABC, abstractmethod\n",
    "from typing import Callable, Tuple\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose that we have:\n",
    "$$\n",
    "t_i = \\mu + \\xi_i\n",
    "$$\n",
    "and $\\xi_i\\sim p(\\xi_i|\\theta)$. Then $\\xi_i|\\mu\\sim p_\\mu(\\xi_i|\\theta)$ where $p_\\mu(\\xi_i|\\theta)$ is simply the distribution $p(\\xi_i|\\theta)$ shifted to the left by $\\mu$.\n",
    "\n",
    "In the event that the event is censored ($e_i=0$), we know that $t_i < \\mu + \\xi_i$ since the 'death' offset of $\\xi_i$ is not observed. \n",
    "\n",
    "Therefore we may write the likelihood of \n",
    "$$\n",
    "\\begin{aligned}\n",
    "p(t_i, e_i|\\mu) =& \\left(p(t_i-\\mu)\\right)^{e_i} \\left(\\int_{t_i}^\\infty p(t-\\mu) dt\\right)^{1-e_i}\\\\\n",
    "\\log p(t_i, e_i|\\mu) =& e_i \\log p(t-\\mu) + (1 - e_i) \\log \\left(1 - \\int_{-\\infty}^{t_i} p(t-\\mu) dt \\right)\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Loss(ABC):\n",
    "    @abstractmethod\n",
    "    def __call__(event:torch.Tensor, *args):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "LossType = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class AFTLoss(Loss):\n",
    "    @staticmethod\n",
    "    def __call__(event:torch.Tensor, log_pdf: torch.Tensor, log_icdf: torch.Tensor) -> torch.Tensor:\n",
    "        lik = event * log_pdf + (1 - event) * log_icdf\n",
    "        return -lik.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-0.0400)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N = 5\n",
    "event = torch.randint(0, 2, (N,))\n",
    "log_pdf = torch.randn((N,))\n",
    "log_cdf = -torch.rand((N,))\n",
    "\n",
    "aft_loss = AFTLoss()\n",
    "aft_loss(event, log_pdf, log_cdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def _aft_loss(\n",
    "    log_pdf: torch.Tensor, log_cdf: torch.Tensor, e: torch.Tensor\n",
    ") -> torch.Tensor:\n",
    "    lik = e * log_pdf + (1 - e) * log_cdf\n",
    "    return -lik.mean()\n",
    "\n",
    "\n",
    "def aft_loss(log_prob, e):\n",
    "    log_pdf, log_cdf = log_prob\n",
    "    return _aft_loss(log_pdf, log_cdf, e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the following loss function to infer our model. See [here](./SAT#Likelihood-Function) for theory.\n",
    "$$\n",
    "-\\log L = \\sum_{i=1}^N \\Lambda(t_i) - d_i \\log \\lambda(t_i)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class HazardLoss(Loss):\n",
    "    @staticmethod\n",
    "    def __call__(event: torch.Tensor, logλ: torch.Tensor, Λ: torch.Tensor) -> torch.Tensor:\n",
    "        log_lik = event * logλ - Λ\n",
    "        return -log_lik.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def _hazard_loss(logλ: torch.Tensor, Λ: torch.Tensor, e: torch.Tensor) -> torch.Tensor:\n",
    "    log_lik = e * logλ - Λ\n",
    "    return -log_lik.mean()\n",
    "\n",
    "\n",
    "def hazard_loss(\n",
    "    hazard: Tuple[torch.Tensor, torch.Tensor], e: torch.Tensor\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    parameters:\n",
    "    - hazard: log hazard and Cumulative hazard\n",
    "    - e: torch.Tensor of 1 if death event occured and 0 otherwise\n",
    "    \"\"\"\n",
    "    logλ, Λ = hazard\n",
    "    return _hazard_loss(logλ, Λ, e)"
   ]
  },
  {
   "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()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}