TFCox Simulated Example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 210,
   "id": "2c8b42f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.layers import Input, Dense, Dropout, BatchNormalization, concatenate, Lambda, Multiply\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.optimizers import Adam, RMSprop,SGD,Nadam, Adagrad, Adadelta\n",
    "from tensorflow.keras.regularizers import l1,l2,l1_l2\n",
    "from tensorflow.keras.initializers import Constant ,Orthogonal, RandomNormal, VarianceScaling, Ones, Zeros\n",
    "from tensorflow.keras.constraints import Constraint, UnitNorm\n",
    "from keras.callbacks import Callback, TerminateOnNaN, ModelCheckpoint\n",
    "from sksurv.metrics import concordance_index_censored as concordance\n",
    "import math\n",
    "import time\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "id": "15e1b84f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#This is the TFCox class\n",
    "\n",
    "class TFCox():\n",
    "    def __init__(self, seed=42,batch_norm=False,l1_ratio=1,lbda=0.0001,\n",
    "                 max_it=50,learn_rate=0.001,stop_if_nan=True,stop_at_value=False, cscore_metric=False,suppress_warnings=True,verbose=0):\n",
    "        \n",
    "        self.max_it = max_it\n",
    "        self.tnan = stop_if_nan\n",
    "        self.tcscore = stop_at_value\n",
    "        self.lr=learn_rate\n",
    "        self.cscore=cscore_metric\n",
    "        np.random.seed(seed)\n",
    "        tf.random.set_seed(seed)\n",
    "        \n",
    "        self.l1r = l1_ratio\n",
    "        self.lbda=lbda\n",
    "        self.bnorm = batch_norm\n",
    "        self.verbose=verbose\n",
    "        if suppress_warnings == True:\n",
    "            import warnings\n",
    "            warnings.filterwarnings('ignore')\n",
    "       \n",
    "    def coxloss(self, state):\n",
    "        \n",
    "        def loss(y_true, y_pred):  \n",
    "\n",
    "                return -K.mean((y_pred - K.log(tf.math.cumsum(K.exp(y_pred),reverse=True,axis=0)+0.0001))*state,axis=0)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def cscore_metric(self, state):\n",
    "        def loss(y_true,y_pred):\n",
    "            con = 0\n",
    "            dis = 0\n",
    "            for a in range(len(y_pred)):\n",
    "                for b in range(a+1,len(y_pred)):                                       \n",
    "                        if (y_pred[a]>y_pred[b])  & (y_pred[a]*state[a]!=0):\n",
    "                            con+=1\n",
    "                            \n",
    "                        elif (y_pred[a]<y_pred[b])  & (y_pred[a]*state[a]!=0):\n",
    "                            dis+=1\n",
    "            return     con/(con+dis)\n",
    "        return loss\n",
    " \n",
    "    \n",
    "    def fit(self, X,state,time):\n",
    "        from tensorflow.python.framework.ops import disable_eager_execution\n",
    "        disable_eager_execution()\n",
    "        K.clear_session()\n",
    "       \n",
    "        \n",
    "        \n",
    "        self.time = np.array(time)  \n",
    "        self.newindex = pd.DataFrame(self.time).sort_values(0).index\n",
    "        self.X = (pd.DataFrame(np.array(X)).reindex(self.newindex))                      \n",
    "        self.state = np.array(pd.DataFrame(np.array(state)).reindex(self.newindex))\n",
    "        self.time  = np.array(pd.DataFrame(np.array(time)).reindex(self.newindex))                       \n",
    "        inputsx = Input(shape=(self.X.shape[1],)) \n",
    "        state = Input(shape=(1,))\n",
    "        \n",
    "        if self.bnorm==True:\n",
    "            out = BatchNormalization()(inputsx)\n",
    "            out = Dense(1,activation='linear',\n",
    "                    kernel_regularizer=l1_l2(self.lbda*self.l1r,self.lbda*(1-self.l1r)),\n",
    "                   use_bias=False)(out)\n",
    "        else:\n",
    "            out = Dense(1,activation='linear',\n",
    "                    kernel_regularizer=l1_l2(self.lbda*self.l1r,self.lbda*(1-self.l1r)),\n",
    "                   use_bias=False)(inputsx)\n",
    "\n",
    "        \n",
    "        model = Model(inputs=[inputsx, state], outputs=out)\n",
    "        if (self.tcscore != False) or (self.cscore==True) :\n",
    "            model.compile(optimizer=Adam(self.lr) ,\n",
    "                          loss=self.coxloss(state) , metrics=[self.cscore_metric(state)],\n",
    "                          experimental_run_tf_function=False)\n",
    "        else:\n",
    "            model.compile(optimizer=Adam(self.lr) ,\n",
    "                          loss=self.coxloss(state) ,\n",
    "                          experimental_run_tf_function=False)\n",
    "        \n",
    "        self.model=model\n",
    "        if self.verbose==1:\n",
    "            print(self.model.summary())\n",
    "\n",
    "        self.loss_history_ = []\n",
    "        for its in range(self.max_it):\n",
    "            self.temp_weights = self.model.get_weights()\n",
    "           \n",
    "            tr = self.model.train_on_batch([self.X, self.state],np.zeros(self.state.shape))\n",
    "           \n",
    "            self.loss_history_.append(tr) \n",
    "            \n",
    "            if self.verbose == 1:\n",
    "                if (self.tcscore != False) or (self.cscore==True) :\n",
    "                    print('loss:', self.loss_history_[-1][0],' C-score: ',self.loss_history_[-1][1] )\n",
    "                else:\n",
    "                    print('loss:', self.loss_history_[-1] )\n",
    "            \n",
    "            if self.tcscore != False:\n",
    "                if self.loss_history_[-1][1]>=self.tcscore:\n",
    "                    print('Terminated early because concordance >=' +str(self.tcscore)+ ' as set by stop_at_value flag.')\n",
    "                    break\n",
    "            if (self.tcscore != False) or (self.cscore==True) :\n",
    "                if (math.isnan(self.loss_history_[-1][0]) or math.isinf(self.loss_history_[-1][0])) and self.tnan:\n",
    "                    self.model.set_weights(self.temp_weights)\n",
    "                    print('Terminated because weights == nan or inf, reverted to last valid weight set')\n",
    "                    break\n",
    "            else:\n",
    "                if (math.isnan(self.loss_history_[-1]) or math.isinf(self.loss_history_[-1])) and self.tnan:\n",
    "                    self.model.set_weights(self.temp_weights)\n",
    "                    print('Terminated because weights == nan or inf, reverted to last valid weight set')\n",
    "                    break\n",
    "            \n",
    "        self.beta_ = self.model.get_weights()[-1]\n",
    "\n",
    "    def predict(self,X):\n",
    "        preds = self.model.predict([X,np.zeros(len(X))])\n",
    "\n",
    "        return preds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "596c6cc2",
   "metadata": {},
   "source": [
    "## Creating the simulated Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 224,
   "id": "808cd0ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Uniform ranodmly distributed time data (no censoring)\n",
    "y_time = np.random.rand(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 225,
   "id": "9c42bdcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Creating simulated X data from the time data \n",
    "X = (y_time * (np.random.rand(5000,1)-0.5)).transpose()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 226,
   "id": "efd0f1f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Adding noise to the X data\n",
    "X = X +( np.random.rand(1000,5000)-0.5)*5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 227,
   "id": "2b589d9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Randomly censoring 75% of the samples  (censored =0, uncensored = 1)\n",
    "y_state=np.zeros(1000)\n",
    "y_state[np.random.choice(1000,250)] =1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 228,
   "id": "beecde3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Removing a random amount of time from the censored samples\n",
    "for a in range(len(y_time)):\n",
    "    if y_state[a] == 0:\n",
    "        y_time[a] = y_time[a] - np.random.rand(1)[0]*y_time[a]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0686f4a",
   "metadata": {},
   "source": [
    "## Running the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "id": "fbf24882",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Creating a simple 80-20 test train split\n",
    "train_index = np.random.choice(1000,800,replace=False)\n",
    "test_index = [x for x in range(1000) if x not in train_index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 230,
   "id": "94ee220c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train concordance lambda= 0 : (0.9994821685878963, 44393, 23, 0, 0)\n",
      "test concordance: lambda= 0 : (0.7963263101026472, 1474, 377, 0, 0)\n",
      "train concordance lambda= 0.0001 : (0.9995046829971181, 44394, 22, 0, 0)\n",
      "test concordance: lambda= 0.0001 : (0.7855213398163156, 1454, 397, 0, 0)\n",
      "train concordance lambda= 0.001 : (0.9994596541786743, 44392, 24, 0, 0)\n",
      "test concordance: lambda= 0.001 : (0.8227984873041599, 1523, 328, 0, 0)\n",
      "train concordance lambda= 0.01 : (0.9978611311239193, 44321, 95, 0, 0)\n",
      "test concordance: lambda= 0.01 : (0.8854673149648838, 1639, 212, 0, 0)\n",
      "train concordance lambda= 0.1 : (0.9538904899135446, 42368, 2048, 0, 0)\n",
      "test concordance: lambda= 0.1 : (0.7277147487844409, 1347, 504, 0, 0)\n",
      "train concordance lambda= 1 : (0.6059528097982709, 26914, 17502, 0, 0)\n",
      "test concordance: lambda= 1 : (0.5807671528903295, 1075, 776, 0, 0)\n",
      "train concordance lambda= 10 : (0.506686779538905, 22505, 21911, 0, 0)\n",
      "test concordance: lambda= 10 : (0.42409508373851973, 785, 1066, 0, 0)\n"
     ]
    }
   ],
   "source": [
    "#Running the TFCox model (default L1_ratio = 1) in a small loop for different values of lambda\n",
    "\n",
    "for a in [0,0.0001,0.001,0.01,0.1,1,10]:\n",
    "    cox = TFCox(lbda=a)\n",
    "    cox.fit(X[train_index],y_state[train_index],y_time[train_index])\n",
    "    train_pred = cox.predict(X[train_index])\n",
    "    test_pred = cox.predict(X[test_index])\n",
    "    \n",
    "    print('train concordance', 'lambda=',a,':' ,concordance(y_state[train_index].astype(bool),y_time[train_index],train_pred.flatten()))\n",
    "\n",
    "    print('test concordance:', 'lambda=',a,':'  ,concordance(y_state[test_index].astype(bool),y_time[test_index],test_pred.flatten()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0440be4f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cleanpython",
   "language": "python",
   "name": "cleanpython"
  },
  "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}