{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7325c0cb",
   "metadata": {},
   "source": [
    "# Sample Code to Run TS-Brier with Frank Copula"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5d7b84bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ae19496a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import torch\n",
    "\n",
    "seed = 0\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d95eeace",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d9d1e139",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   pid  event  time   num_age  num_num_co  num_edu  num_scoma  num_hday  \\\n",
      "0    0      0  2029  62.84998           0     11.0        0.0         1   \n",
      "1    1      1     4  60.33899           2     12.0       44.0         3   \n",
      "2    2      1    47  52.74698           2     12.0        0.0         4   \n",
      "3    3      1   133  42.38498           2     11.0        0.0         1   \n",
      "4    4      0  2029  79.88495           1      NaN       26.0         3   \n",
      "\n",
      "     num_sps  num_surv2m  ...        fac_dzgroup         fac_dzclass  \\\n",
      "0  33.898438    0.262939  ...        Lung Cancer              Cancer   \n",
      "1  52.695312    0.001000  ...          Cirrhosis  COPD/CHF/Cirrhosis   \n",
      "2  20.500000    0.790894  ...          Cirrhosis  COPD/CHF/Cirrhosis   \n",
      "3  20.097656    0.698975  ...        Lung Cancer              Cancer   \n",
      "4  23.500000    0.634888  ...  ARF/MOSF w/Sepsis            ARF/MOSF   \n",
      "\n",
      "   fac_num_co  fac_race  fac_diabetes  fac_dementia      fac_ca  fac_dnr  \\\n",
      "0           0     other             0             0  metastatic   no dnr   \n",
      "1           2     white             0             0          no  missing   \n",
      "2           2     white             0             0          no   no dnr   \n",
      "3           2     white             0             0  metastatic   no dnr   \n",
      "4           1     white             0             0          no   no dnr   \n",
      "\n",
      "             fac_sfdm2  fac_income  \n",
      "0              missing    $11-$25k  \n",
      "1     <2 mo. follow-up    $11-$25k  \n",
      "2     <2 mo. follow-up  under $11k  \n",
      "3  no(M2 and SIP pres)  under $11k  \n",
      "4  no(M2 and SIP pres)     missing  \n",
      "\n",
      "[5 rows x 38 columns]\n"
     ]
    }
   ],
   "source": [
    "# Prerequisite: pip install SurvSet\n",
    "from SurvSet.data import SurvLoader\n",
    "\n",
    "dataset_name = \"support2\"\n",
    "loader = SurvLoader()\n",
    "df, _ = loader.load_dataset(ds_name=dataset_name).values()\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28505de2",
   "metadata": {},
   "source": [
    "## Encode Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b4f78464",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   pid  event  time   num_age  num_num_co   num_edu  num_scoma  num_hday  \\\n",
      "0    0      0  2029  0.534663    0.000000  0.354839       0.00  0.000000   \n",
      "1    1      1     4  0.504702    0.222222  0.387097       0.44  0.013605   \n",
      "2    2      1    47  0.414111    0.222222  0.387097       0.00  0.020408   \n",
      "3    3      1   133  0.290468    0.222222  0.354839       0.00  0.000000   \n",
      "4    4      0  2029  0.737930    0.111111  0.387097       0.26  0.013605   \n",
      "\n",
      "    num_sps  num_surv2m  ...  fac_dnr_missing  fac_sfdm2_Coma or Intub  \\\n",
      "0  0.340431    0.271080  ...              0.0                      0.0   \n",
      "1  0.530323    0.001031  ...              1.0                      0.0   \n",
      "2  0.205077    0.815379  ...              0.0                      0.0   \n",
      "3  0.201012    0.720614  ...              0.0                      0.0   \n",
      "4  0.235383    0.654543  ...              0.0                      0.0   \n",
      "\n",
      "   fac_sfdm2_SIP>=30  fac_sfdm2_adl>=4 (>=5 if sur)  \\\n",
      "0                0.0                            0.0   \n",
      "1                0.0                            0.0   \n",
      "2                0.0                            0.0   \n",
      "3                0.0                            0.0   \n",
      "4                0.0                            0.0   \n",
      "\n",
      "   fac_sfdm2_no(M2 and SIP pres)  fac_sfdm2_missing  fac_income_$25-$50k  \\\n",
      "0                            0.0                1.0                  0.0   \n",
      "1                            0.0                0.0                  0.0   \n",
      "2                            0.0                0.0                  0.0   \n",
      "3                            1.0                0.0                  0.0   \n",
      "4                            1.0                0.0                  0.0   \n",
      "\n",
      "   fac_income_>$50k  fac_income_under $11k  fac_income_missing  \n",
      "0               0.0                    0.0                 0.0  \n",
      "1               0.0                    0.0                 0.0  \n",
      "2               0.0                    1.0                 0.0  \n",
      "3               0.0                    1.0                 0.0  \n",
      "4               0.0                    0.0                 1.0  \n",
      "\n",
      "[5 rows x 68 columns]\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "\n",
    "# rescale numerical features into [0, 1]\n",
    "cols_cont = [col for col in df.columns if col.startswith(\"num_\")]\n",
    "if len(cols_cont) > 0:\n",
    "    df[cols_cont] = df[cols_cont].fillna(df[cols_cont].median())\n",
    "    mm = MinMaxScaler()\n",
    "    df[cols_cont] = mm.fit_transform(df[cols_cont])\n",
    "\n",
    "# encode categorical columns\n",
    "cols_cat = []\n",
    "for col in df.columns:\n",
    "    if col.startswith(\"fac_\"):\n",
    "        cols_cat.append(col)\n",
    "df = pd.get_dummies(df, columns=cols_cat, drop_first=True, dtype=\"float32\")\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e22218",
   "metadata": {},
   "source": [
    "## Split Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c8ba9e12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(7284, 68) (1821, 68)\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "df_train, df_test = train_test_split(df, random_state=42, test_size=0.2) \n",
    "print(df_train.shape, df_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71aaaee4",
   "metadata": {},
   "source": [
    "## Train with Neural Network (MLP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c4d5ce80",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cenreg.utils\n",
    "from cenreg.pytorch.datamodule import SurvDataModule\n",
    "from cenreg.pytorch.loss_cjd import Brier\n",
    "from cenreg.pytorch.mlp import MLP\n",
    "\n",
    "# drop the columns that are not needed\n",
    "x_train = df_train.drop(columns=[\"pid\", \"time\", \"event\"]).values\n",
    "x_test = df_test.drop(columns=[\"pid\", \"time\", \"event\"]).values\n",
    "\n",
    "# preparation for PyTorch training\n",
    "sdm = SurvDataModule(128)\n",
    "train_dataloader = sdm.train_dataloader(x_train, df_train[\"time\"].values, df_train[\"event\"].values)\n",
    "num_bins = 32\n",
    "bins = torch.tensor(cenreg.utils.create_bins(df[\"time\"].max(), 0.0, num_bins))\n",
    "loss_fn = Brier(bins, 2)\n",
    "model = MLP(x_train.shape[1], num_bins*2, 64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2e5f8205",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5, Train Loss: 0.6275973048126489\n",
      "Epoch 10, Train Loss: 0.6248274539646349\n",
      "Epoch 15, Train Loss: 0.6236271492221899\n",
      "Epoch 20, Train Loss: 0.6239773129162035\n",
      "Epoch 25, Train Loss: 0.6240868835072768\n",
      "Epoch 30, Train Loss: 0.6224435128663716\n",
      "Epoch 35, Train Loss: 0.6235858191523636\n",
      "Epoch 40, Train Loss: 0.6254984548217372\n"
     ]
    }
   ],
   "source": [
    "optimizer_mlp = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "for epoch in range(40):\n",
    "    loss_sum = 0.0\n",
    "    for i, data in enumerate(train_dataloader):\n",
    "        x, t, e = data\n",
    "\n",
    "        optimizer_mlp.zero_grad()\n",
    "        f_pred = model(x)\n",
    "        loss = loss_fn.loss(f_pred, t, e).mean()\n",
    "        loss_sum += loss.item()\n",
    "        loss.backward()\n",
    "        optimizer_mlp.step()\n",
    "    if (epoch + 1) % 5 == 0:\n",
    "        print(f\"Epoch {epoch+1}, Train Loss: {loss_sum / len(train_dataloader)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d38d4439",
   "metadata": {},
   "source": [
    "## Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3429b83c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y_pred.shape torch.Size([116, 64])\n",
      "cjd_pred.shape (1821, 2, 32)\n"
     ]
    }
   ],
   "source": [
    "test_dataloader = sdm.test_dataloader(x_test, df_test[\"time\"].values, df_test[\"event\"].values)\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for i, data in enumerate(test_dataloader):\n",
    "        x, _, _ = data\n",
    "        y_pred = model(x)\n",
    "        break  # assuming that the first batch contains all the test data\n",
    "print(\"y_pred.shape\", f_pred.shape)\n",
    "cjd_pred = y_pred.reshape(-1, num_bins, 2).permute(0, 2, 1).detach().cpu().numpy()\n",
    "print(\"cjd_pred.shape\", cjd_pred.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10e25bd2",
   "metadata": {},
   "source": [
    "## Postprocess the density estimates into CDFs (equivalently, survival functions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "62b4f42e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=0, mean loss=0.001079196\n",
      "epoch=50, mean loss=0.000004016\n",
      "epoch=100, mean loss=0.000002269\n",
      "epoch=150, mean loss=0.000001589\n",
      "epoch=200, mean loss=0.000001186\n",
      "epoch=250, mean loss=0.000000928\n",
      "epoch=300, mean loss=0.000000749\n",
      "epoch=350, mean loss=0.000000618\n",
      "epoch=400, mean loss=0.000000520\n",
      "epoch=450, mean loss=0.000000445\n"
     ]
    }
   ],
   "source": [
    "import cenreg.model\n",
    "import cenreg.pytorch.copula_torch\n",
    "from cenreg.distribution.linear import LinearCDF\n",
    "from cenreg.pytorch.cjd2F import MseModel\n",
    "\n",
    "# Create a Frank copula object\n",
    "copula = cenreg.pytorch.copula_torch.create(\"frank\", 5.0)\n",
    "\n",
    "# initialize the model\n",
    "init_F = cenreg.model.cjd2surv(cjd_pred)\n",
    "init_f = init_F[:,:,1:] - init_F[:,:,:-1]  # compute the difference to get f\n",
    "F_softmax = MseModel(cjd_pred, copula, init_f = init_f)\n",
    "optimizer_mse = torch.optim.Adam(F_softmax.parameters(), lr=0.1)\n",
    "\n",
    "# train the model\n",
    "num_epochs = 500\n",
    "with torch.enable_grad():\n",
    "    model.train()\n",
    "    for epoch in range(num_epochs):\n",
    "        optimizer_mse.zero_grad()\n",
    "        loss = F_softmax.loss().mean()\n",
    "        loss.backward()\n",
    "        optimizer_mse.step()\n",
    "        if epoch % 50 == 0:\n",
    "            print(\"epoch=%d, mean loss=%.9f\" % (epoch, loss.item()))\n",
    "\n",
    "# Evaluate the model\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    F_pred = F_softmax.forward().detach().cpu().numpy()\n",
    "\n",
    "# Create a list of LinearCDF objects for each risk\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": "0b4164d6",
   "metadata": {},
   "source": [
    "## Evaluate Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a3cbdc55",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CJD-Brier 0.6394180830320167\n",
      "CJD-Logarithmic 2.379232\n",
      "CJD-KS 0.1807599546171742\n",
      "NLL-SC 6.533535194087778\n"
     ]
    }
   ],
   "source": [
    "import cenreg.metric.cdf\n",
    "import cenreg.metric.cjd\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",
    "bins_np = bins.detach().cpu().numpy()\n",
    "\n",
    "# Compute Brier score on CJD representation\n",
    "cjd_brier = cenreg.metric.cjd.brier(observed_times, events, 2, cjd_pred, bins_np)\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_np)\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_np)\n",
    "print(\"CJD-KS\", cjd_ks)\n",
    "\n",
    "# Compute NLL-SC metric\n",
    "copula_np = cenreg.model.copula_np.FrankCopula(5.0)\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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5a6a98f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24180 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[1,:-1], marker='o')  # Plot the 1st survival function\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
}
