{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bf82f39a",
   "metadata": {},
   "source": [
    "# Sample Code to Run Kaplan-Meier Estimator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "075dad06",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6649f9d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "\n",
    "seed = 0\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48327bc9",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "09c2011b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "       pid  event  time  num_age  num_kappa  num_lambda  num_flc_grp  \\\n",
      "0        0      1    85       97      5.700       4.860           10   \n",
      "1        1      1  1281       92      0.870       0.683            1   \n",
      "2        2      1    69       94      4.360       3.850           10   \n",
      "3        3      1   115       92      2.420       2.220            9   \n",
      "4        4      1  1039       93      1.320       1.690            6   \n",
      "...    ...    ...   ...      ...        ...         ...          ...   \n",
      "7869  7869      0  4997       52      1.210       1.610            6   \n",
      "7870  7870      0  3652       52      0.858       0.581            1   \n",
      "7871  7871      0  2507       54      1.700       1.720            8   \n",
      "7872  7872      0  4982       53      1.710       2.690            9   \n",
      "7873  7873      0  3995       50      1.190       1.250            4   \n",
      "\n",
      "      num_creatinine  num_sample_yr fac_sex  fac_chapter fac_sample_yr  \\\n",
      "0                1.7           1997       F  Circulatory          1997   \n",
      "1                0.9           2000       F    Neoplasms          2000   \n",
      "2                1.4           1997       F  Circulatory          1997   \n",
      "3                1.0           1996       F  Circulatory          1996   \n",
      "4                1.1           1996       F  Circulatory          1996   \n",
      "...              ...            ...     ...          ...           ...   \n",
      "7869             1.0           1995       F      missing          1995   \n",
      "7870             0.8           1999       F      missing          1999   \n",
      "7871             NaN           2002       F      missing          2002   \n",
      "7872             NaN           1995       F      missing          1995   \n",
      "7873             0.7           1998       F      missing          1998   \n",
      "\n",
      "     fac_mgus  \n",
      "0           0  \n",
      "1           0  \n",
      "2           0  \n",
      "3           0  \n",
      "4           0  \n",
      "...       ...  \n",
      "7869        0  \n",
      "7870        0  \n",
      "7871        0  \n",
      "7872        0  \n",
      "7873        0  \n",
      "\n",
      "[7874 rows x 13 columns]\n"
     ]
    }
   ],
   "source": [
    "# Prerequisite: pip install SurvSet\n",
    "from SurvSet.data import SurvLoader\n",
    "\n",
    "dataset_name = \"flchain\"\n",
    "loader = SurvLoader()\n",
    "df, _ = loader.load_dataset(ds_name=dataset_name).values()\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c4abd2",
   "metadata": {},
   "source": [
    "## Run Kaplan-Meier Estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f7c08980",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cenreg.model.kaplan_meier\n",
    "\n",
    "km = cenreg.model.kaplan_meier.KaplanMeierDistribution()\n",
    "km.fit(df[\"time\"].values, df[\"event\"].values)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88bf98e7",
   "metadata": {},
   "source": [
    "## Plot Kaplan-Meier Curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7aebf393",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20004 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig = plt.figure()\n",
    "plt.xlabel('time')\n",
    "plt.ylabel('survival probability')\n",
    "plt.ylim([0.0,1.05])\n",
    "x = km.bins\n",
    "y = km.survival_rates[:-1]\n",
    "plt.stairs(y, x, baseline=None)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cenreg",
   "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
