{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fbff13a0",
   "metadata": {},
   "source": [
    "# Sample Code to Run TS-LGB\n",
    "This notebook shows how to use TS-LGB, a two-step algorithm for survival analysis with LightGBM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a879879e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0e3cb25d",
   "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": "e90be0f4",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "abbbdc13",
   "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": "b7e52ac2",
   "metadata": {},
   "source": [
    "## Compute bins and label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b1a07ae2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bins: [   0.          168.22580645  336.4516129   504.67741935  672.90322581\n",
      "  841.12903226 1009.35483871 1177.58064516 1345.80645161 1514.03225806\n",
      " 1682.25806452 1850.48387097 2018.70967742 2186.93548387 2355.16129032\n",
      " 2523.38709677 2691.61290323 2859.83870968 3028.06451613 3196.29032258\n",
      " 3364.51612903 3532.74193548 3700.96774194 3869.19354839 4037.41935484\n",
      " 4205.64516129 4373.87096774 4542.09677419 4710.32258065 4878.5483871\n",
      " 5046.77419355 5215.         5383.22580645] (33,)\n",
      "Label: [ 1 15  1 ... 28 58 46] (7874,)\n"
     ]
    }
   ],
   "source": [
    "import cenreg.utils\n",
    "\n",
    "num_bins = 32\n",
    "bins = cenreg.utils.create_bins(df[\"time\"].max(), 0.0, num_bins)\n",
    "print(\"Bins:\", bins, bins.shape)\n",
    "label = cenreg.utils.create_discretized_labels(bins, 2, df[\"time\"].values, df[\"event\"].values)\n",
    "print(\"Label:\", label, label.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "463b8de5",
   "metadata": {},
   "source": [
    "## Encode Categorical Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fa2dd5a8",
   "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.70       4.860           10   \n",
      "1    1      1  1281       92       0.87       0.683            1   \n",
      "2    2      1    69       94       4.36       3.850           10   \n",
      "3    3      1   115       92       2.42       2.220            9   \n",
      "4    4      1  1039       93       1.32       1.690            6   \n",
      "\n",
      "   num_creatinine  num_sample_yr  fac_sex  fac_chapter  fac_sample_yr  \\\n",
      "0             1.7           1997        0            1              2   \n",
      "1             0.9           2000        0           12              5   \n",
      "2             1.4           1997        0            1              2   \n",
      "3             1.0           1996        0            1              1   \n",
      "4             1.1           1996        0            1              1   \n",
      "\n",
      "   fac_mgus  \n",
      "0         0  \n",
      "1         0  \n",
      "2         0  \n",
      "3         0  \n",
      "4         0  \n"
     ]
    }
   ],
   "source": [
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "cols_cat = []\n",
    "le = LabelEncoder()\n",
    "for col in df.columns:\n",
    "    if col.startswith(\"fac_\"):\n",
    "        cols_cat.append(col)\n",
    "        le = le.fit(df[col])\n",
    "        df[col]=le.transform(df[col])\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a4078df",
   "metadata": {},
   "source": [
    "## Split dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ad343a33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(6299, 13) (1575, 13)\n",
      "(6299,) (1575,)\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "df_train, df_test, label_train, label_test = train_test_split(df, label, random_state=42, test_size=0.2) \n",
    "print(df_train.shape, df_test.shape)\n",
    "print(label_train.shape, label_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "817ec528",
   "metadata": {},
   "source": [
    "## Train density estimation model (LightGBM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1006fee1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y_pred [[7.96144450e-03 2.28014262e-03 1.63790535e-02 ... 3.47160459e-03\n",
      "  9.15537136e-02 1.44555474e-05]\n",
      " [2.45960302e-03 4.58311545e-03 2.14122411e-03 ... 1.88605012e-04\n",
      "  1.00814236e-03 1.80684927e-05]\n",
      " [5.08969667e-03 2.11147059e-03 3.63121857e-03 ... 1.72455236e-04\n",
      "  9.23145421e-04 1.65451299e-05]\n",
      " ...\n",
      " [8.01488372e-03 2.30534501e-03 1.17194916e-02 ... 1.88272035e-04\n",
      "  1.00790828e-03 1.80642973e-05]\n",
      " [3.01111284e-03 6.91891441e-02 1.79628984e-03 ... 5.93539067e-03\n",
      "  1.24445083e-02 1.10585316e-04]\n",
      " [3.04142034e-03 1.75346504e-03 9.12410486e-04 ... 5.75605252e-04\n",
      "  2.13013848e-01 1.37248861e-05]] (1575, 62)\n"
     ]
    }
   ],
   "source": [
    "import lightgbm as lgb\n",
    "\n",
    "# drop the columns that are not needed\n",
    "x_train = df_train.drop(columns=[\"pid\", \"time\", \"event\"])\n",
    "x_test = df_test.drop(columns=[\"pid\", \"time\", \"event\"])\n",
    "\n",
    "# train with LightGBM with parameters tuned for flchain dataset\n",
    "lgb_model = lgb.LGBMClassifier(\n",
    "    n_estimators = 279,\n",
    "    max_depth = 3,\n",
    "    min_child_samples = 6,\n",
    "    num_leaves = 26,\n",
    "    criterion = \"log_loss\",\n",
    "    learning_rate = 0.007696701896574052,\n",
    "    lambda_l1 = 3.7252491463357303e-05,\n",
    "    random_state=0,\n",
    "    verbosity=-1)\n",
    "lgb_model.fit(x_train, label_train, categorical_feature=cols_cat)\n",
    "y_pred = lgb_model.predict_proba(x_test)\n",
    "print(\"y_pred\", y_pred, y_pred.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3e7ec2fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F_pred [[0.         0.00792984 0.02421205 0.02551188 0.03161314 0.03192653\n",
      "  0.03256613 0.03274607 0.03379904 0.03453036 0.03720796 0.03891525\n",
      "  0.04138735 0.08424525 0.08732598 0.0883953  0.08940544 0.09252055\n",
      "  0.09413113 0.09540505 0.09978056 0.1154755  0.11832524 0.12142105\n",
      "  0.12668884 0.13371802 0.14039584 0.14960921 0.16797443 0.18434106\n",
      "  0.66088253 0.87522596 1.        ]\n",
      " [0.         0.00227755 0.00425415 0.0055923  0.00722133 0.008572\n",
      "  0.00958212 0.0107335  0.0122836  0.01334901 0.01438471 0.01633658\n",
      "  0.01791809 0.01870086 0.02074726 0.02234813 0.0236001  0.0276752\n",
      "  0.02932456 0.03060169 0.03232558 0.03383638 0.03518583 0.03711157\n",
      "  0.03896298 0.04051644 0.04230976 0.04347664 0.04455178 0.0453818\n",
      "  0.04963866 0.04963866 1.        ]] (1575, 2, 33)\n"
     ]
    }
   ],
   "source": [
    "from cenreg.distribution.linear import LinearCDF\n",
    "import cenreg.model\n",
    "\n",
    "def fill_missing_pred(y_pred, classes, num_class):\n",
    "    list_pred = []\n",
    "    for i in range(num_class):\n",
    "        not_found = True\n",
    "        for j in range(len(classes)):\n",
    "            if i == classes[j]:\n",
    "                list_pred.append(y_pred[:, j])\n",
    "                not_found = False\n",
    "                break\n",
    "        if not_found:\n",
    "            list_pred.append(np.zeros(y_pred.shape[0]))\n",
    "    return np.stack(list_pred, axis=1)\n",
    "\n",
    "if len(lgb_model.classes_) < 2 * num_bins:\n",
    "    y_pred = fill_missing_pred(y_pred, lgb_model.classes_, 2 * num_bins)\n",
    "\n",
    "cjd_pred = y_pred.reshape(-1, num_bins, 2).transpose(0, 2, 1)\n",
    "\n",
    "# Convert the output of LightGBM to the CDF for each risk\n",
    "F_pred = cenreg.model.cjd2surv(cjd_pred)\n",
    "print(\"F_pred\", F_pred[0,:,:], F_pred.shape)\n",
    "\n",
    "# Create a list of LinearCDF objects\n",
    "list_t_dist = []\n",
    "for k in range(2):\n",
    "    list_t_dist.append(LinearCDF(bins, F_pred[:,k,:], apply_cumsum=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01c048ec",
   "metadata": {},
   "source": [
    "## Evaluate Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8a8612a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CJD-Brier 0.7091098104328221\n",
      "CJD-Logarithmic 2.0925483462835563\n",
      "CJD-KS 0.2252175349043452\n",
      "NLL-SC 7.176017453955104\n",
      "Cen-log 0.9788527011807912\n",
      "D-calibration 0.0017084230184280867\n",
      "KM-calibration 0.015210059234563285\n"
     ]
    }
   ],
   "source": [
    "import cenreg.metric.cdf\n",
    "import cenreg.metric.cjd\n",
    "import cenreg.metric.quantile\n",
    "import cenreg.model.copula_np\n",
    "\n",
    "cjd_pred = cjd_pred.reshape(-1, num_bins * 2)\n",
    "observed_times = df_test['time'].values\n",
    "events = df_test[\"event\"].astype(bool).values\n",
    "\n",
    "# Compute Brier score on CJD representation\n",
    "cjd_brier = cenreg.metric.cjd.brier(observed_times, events, 2, cjd_pred, bins)\n",
    "print(\"CJD-Brier\", cjd_brier.mean())\n",
    "\n",
    "# Compute Logarithmic score on CDF representation\n",
    "cjd_logarithmic = cenreg.metric.cjd.negative_loglikelihood(observed_times, events, cjd_pred, bins)\n",
    "print(\"CJD-Logarithmic\", cjd_logarithmic.mean())\n",
    "\n",
    "# Compute KS calibration error on CDF representation\n",
    "cjd_ks = cenreg.metric.cjd.kolmogorov_smirnov_calibration_error(observed_times, events, cjd_pred, bins)\n",
    "print(\"CJD-KS\", cjd_ks)\n",
    "\n",
    "# Compute NLL-SC metric\n",
    "copula_np = cenreg.model.copula_np.IndependenceCopula()\n",
    "survival_copula_np = cenreg.model.copula_np.SurvivalCopula(copula_np)\n",
    "nll_sc = cenreg.metric.cdf.nll_sc(list_t_dist, observed_times, events, survival_copula_np)\n",
    "print(\"NLL-SC\", nll_sc.mean())\n",
    "\n",
    "# Compute cen-log metric\n",
    "nll = cenreg.metric.cdf.negative_loglikelihood(list_t_dist[1], observed_times, events)\n",
    "print(\"Cen-log\", nll.mean())\n",
    "\n",
    "# Compute D-calibration\n",
    "quantiles_cal = np.linspace(0.0, 1.0, 11)\n",
    "dcal = cenreg.metric.quantile.d_calibration(list_t_dist[1], observed_times, events, quantiles_cal)\n",
    "print(\"D-calibration\", dcal)\n",
    "\n",
    "# Compute KM-calibration\n",
    "kmcal = cenreg.metric.cdf.km_calibration(list_t_dist[1], observed_times, events, bins)\n",
    "print(\"KM-calibration\", kmcal)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8fca2660",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24000 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "x = list_t_dist[1].boundaries\n",
    "y = 1.0 - list_t_dist[1].cdf(x)\n",
    "\n",
    "plt.style.use(\"tableau-colorblind10\")\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.title(dataset_name)\n",
    "plt.xlabel(\"Time\")\n",
    "plt.ylabel(\"Survival probability\")\n",
    "plt.ylim(0.0, 1.05)\n",
    "plt.plot(x[:-1], y[50,:-1], marker='o')  # Plot the 50-th survival function\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bee77dd9",
   "metadata": {},
   "source": []
  }
 ],
 "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
}
