Predicting future medical diagnoses with RNNs using Fast AI API from scratch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predicting future medical diagnoses with RNNs using Fast AI API from scratch\n",
    "*Full pytorch implementation Doctor AI paper using Electronic Health Records*<br>\n",
    "by: Sparkle Russell-Puleri"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the first [part one](https://towardsdatascience.com/using-electronic-health-records-ehr-for-predicting-future-diagnosis-codes-using-gated-recurrent-bcd0de7d7436) of this tutorial we created a rough template of the [Doctor AI: Predicting Clinical Events via Recurrent Neural Networks paper(2016)](https://arxiv.org/abs/1511.05942) by Edwared Choi et.al. In this tutorial we took it a step further using the Fast.ai buttoms up approach. This code is fully functional and details on how the data was processed can be accesed in [part one](https://towardsdatascience.com/using-electronic-health-records-ehr-for-predicting-future-diagnosis-codes-using-gated-recurrent-bcd0de7d7436)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import torch \n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checking for GPU availability\n",
    "This model was trained on a GPU enabled system...highly recommended."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on CPU!\n"
     ]
    }
   ],
   "source": [
    "# check if GPU is available\n",
    "if(torch.cuda.is_available()):\n",
    "    print('Training on GPU!')\n",
    "else: \n",
    "    print('Training on CPU!')\n",
    "    \n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### About the data set:\n",
    "This study will utilize the [MIMIC III](https://mimic.physionet.org/) electronic health record (EHR) dataset, which is comprised of over 58,000 hospital admissions for 38,645 adults and 7,875 neonates. This dataset is a collection of de-identified intensive care unit stays at the Beth Israel Deaconess Medical Center from June 2001- October 2012. A detailed walkthrough of the data pre-processing steps used can be found in [part one](https://towardsdatascience.com/using-electronic-health-records-ehr-for-predicting-future-diagnosis-codes-using-gated-recurrent-bcd0de7d7436)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seqs = np.array(pickle.load(open('data/Jan19.seqs','rb')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data pre-processed datasets will be loaded and split into a train, test and validation set at a `75%:15%:10%` ratio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(seqFile, labelFile, test_frac=0.15, valid_frac=0.10):\n",
    "    sequences = np.array(pickle.load(open(seqFile,'rb')))\n",
    "    labels = np.array(pickle.load(open(labelFile,'rb')))\n",
    "    \n",
    "    dataSize = len(labels)\n",
    "    idx = np.random.permutation(dataSize)\n",
    "    nTest = int(np.ceil(test_frac * dataSize))\n",
    "    nValid = int(np.ceil(valid_frac * dataSize))\n",
    "\n",
    "    test_idx = idx[:nTest]\n",
    "    valid_idx = idx[nTest:nTest+nValid]\n",
    "    train_idx = idx[nTest+nValid:]\n",
    "\n",
    "    train_x = sequences[train_idx]\n",
    "    train_y = labels[train_idx]\n",
    "    test_x = sequences[test_idx]\n",
    "    test_y = labels[test_idx]\n",
    "    valid_x = sequences[valid_idx]\n",
    "    valid_y = labels[valid_idx]\n",
    "\n",
    "    train_x = [sorted(seq) for seq in train_x]\n",
    "    train_y = [sorted(seq) for seq in train_y]\n",
    "    valid_x = [sorted(seq) for seq in valid_x]\n",
    "    valid_y = [sorted(seq) for seq in valid_y]\n",
    "    test_x = [sorted(seq) for seq in test_x]\n",
    "    test_y = [sorted(seq) for seq in test_y]\n",
    "\n",
    "    train = (train_x, train_y)\n",
    "    test = (test_x, test_y)\n",
    "    valid = (valid_x, valid_y)\n",
    "    return (train, test, valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Padding sequences: to address variable length sequences\n",
    "Using the artifical EHR data created in part one we pad the sequences to the length of the longest sequence in each minibatch. To help explain this in greater depth let's take a look at the `Artificial EHR data` created in part one. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detailed explanation using artificially generated EHR data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "artificalData_seqs = np.array(pickle.load(open('../GRU_EHR/ArtificialEHR_Data.encodedDxs','rb')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here you can see that we have an array with two list, each list representing a unique patient. Now, within each list are a series of lists, each representing a unique visit. Finally, the encoded numericals represent the diagnosis codes assigned during each unique visit. It is key to note that given the uniqueness of each patient's condition, there are `variable length` sequences for both the visits and diagnosis codes assigned. Because EHR data is longitudinal in nature and we are often interested in understand a patient's risk or progression over time. When using tabular data processing these nested time-dependent `variable length` sequences can get complicated quickly. Recall the following image from part one, detailing the mapping of each visit date to the diagnosis codes assigned during that visit.\n",
    "\n",
    "<img src=\"img/step4b.png\" style=\"height:500px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([list([[0, 1, 1], [0, 1, 1, 2, 3], [0, 1, 1, 4, 5, 6]]),\n",
       "       list([[7, 8, 9], [7, 8, 9, 10]])], dtype=object)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "artificalData_seqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[[0, 1, 1], [0, 1, 1, 2, 3], [0, 1, 1, 4, 5, 6]], [[7, 8, 9], [7, 8, 9, 10]]]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Remove array \n",
    "artificalData_seqs = [sorted(seq) for seq in artificalData_seqs]\n",
    "artificalData_seqs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/Patients.png\" style=\"height:200px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def padding(seqs, labels, inputDimSize, numClass):\n",
    "    lengths = np.array([len(seq) for seq in seqs]) - 1\n",
    "    maxlen = np.max(lengths)\n",
    "    num_samples = len(seqs)\n",
    "    \n",
    "\n",
    "    x = torch.zeros(maxlen, num_samples, inputDimSize)\n",
    "    y = torch.zeros(maxlen, num_samples, numClass)\n",
    "    mask = torch.zeros(maxlen, num_samples)\n",
    "\n",
    "    for idx, (seq, label) in enumerate(zip(seqs, labels)):\n",
    "        for xvec, subseq in zip(x[:, idx, :], seq[:-1]):\n",
    "            xvec[subseq] = 1.\n",
    "        for yvec, subseq in zip(y[:, idx, :], label[1:]):\n",
    "            yvec[subseq] = 1. \n",
    "\n",
    "        mask[:lengths[idx], idx] = 1.\n",
    "\n",
    "    lengths = torch.LongTensor(lengths)\n",
    "    return x, y, mask, lengths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### So what extactly are we padding with this nested list?\n",
    "Let's break down the padding function:\n",
    "1. `lenghts = np.array([len(seq) for seq in seqs]) - 1` \n",
    "Here were are mysteriously subtracting 1 from the length, in the author's notes he mentioned that both the `visit` and `label` files must match as the algorithm takes care of the time lag for inference time.\n",
    "\n",
    "What does this mean? \n",
    "Given the structure of the data, the last visit in each patient's record will be removed. As illustrated here:\n",
    "<img src=\"img/Lengths.png\" style=\"height:350px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2, 1])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lenghts_artificial = np.array([len(seq) for seq in artificalData_seqs]) - 1\n",
    "lenghts_artificial # here we can see that the final sequences for each patient was reduced by 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_artifical, y_artifical, mask_artifical, lenghts_artifical = padding(artificalData_seqs, artificalData_seqs, 11, 11)\n",
    "# 11 represents the number of tokens in the vocab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Aside: Dealing with variable length sequences in a Character level RNN \n",
    "If this was a character level  problem let's say [`Sparkle`,`Dorian`, `Deep`, `Learning`]. The sequences are first arranged by length, in descending order and padded with zeros (red), where each letter represents a token. As shown here:\n",
    "<img src=\"img/test_seq.png\" style=\"height:250px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### EHR data:\n",
    "However, for EHR data of this form given our current problem, instead of each encoded diagnosis code representing a unique token. In this case, each visit represents a token/sequence. So, using the same approach used with character level RNNs we first arrange each mini-batch by the patient visits in descending order. In this the patient 1 has the longest visit history with a total of two visits, while patient 2's visits will be padded to the max length of 2, since it's the longest sequence. As shown here:\n",
    "<img src=\"img/padding1.png\" style=\"height:175px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, that we have taken care of the variable length problem, we can proceed to multi-one hot encode our sequences. This will result in the desired dimensions of S x B x I ( Sequence length, Batch size, Input dimensions/vocab)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 0.]],\n",
       "\n",
       "        [[1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_artifical"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we can easily see that the sequences will represent the patient with the longest visit history in each minibatch, while all others will be padded to this length (red). Depending on the desired batch size, the batch size will represent how many patients sequences are feed in at each timestep. Finally, the inner list will be encoded to the length of the vocabulary, which in this case the number of unique diagnosis codes in the entire dataset.\n",
    "\n",
    "<img src=\"img/minibatch.png\" style=\"height:275px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 2, 11])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# dimensions: S x B x I\n",
    "x_artifical.shape "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Labels\n",
    "To ensure that the labels are shifted over by one sequence, so that the algorithm can accurately predict the next time step. The author took care of this by ensuring that the training data excluded the last vist within each patient's history, using this logic `for xvec, subseq in zip(x[:, idx, :], seq[:-1]):`, where we took all but the last visit within each patient's visit record `seq[:-1]`. For the labels, this meant that the sequences will start from the patients second visit, or in python's indexing style the first index `for yvec, subseq in zip(y[:, idx, :], label[1:])`, where the label `label[1:]`, is shifted by one.\n",
    "\n",
    "<img src=\"img/labels.png\" style=\"height:475px\" align=\"center\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 0., 0., 1., 1., 1., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_artifical"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What is masking and what does it do?\n",
    "Masking allows the algorithm to know where the true sequences are in one-hot encoded data, simply put ignore/filter out the padding values, which in our case are zeros. This allows us to easily handle variable length sequences in RNNs, which require fixed length inputs. How is it done? Remember the `lengths` variable? This variable stores the effective lengths of each patient's sequences in descending order (**recall**: after removing the last sequence in each record for inference, eg. patient 1 has 3 visits, but length will reflect only 2). The logic in the code `mask[:lengths[idx], idx] = 1.` then fills in our zeroed tensor along the rows with 1's to match the length of each patient sequence from largest to smallest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2, 1])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lenghts_artificial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1.],\n",
       "        [1., 0.]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask_artifical"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Loaders and Sampler\n",
    "The `Dataset` class is an abstract class that represents the data in x and y pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset():\n",
    "    def __init__(self, x, y):\n",
    "        self.x, self.y = x, y\n",
    "    def __len__(self): return len(self.x)\n",
    "    def __getitem__(self, i ): return self.x[i], self.y[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `Sampler` class randomly shuffles the order of the training set (validation set will not be randomized). Additionally, it keeps the exact amount of sequences needed created a full batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sampler():\n",
    "    def __init__(self, ds, bs, shuffle=False):\n",
    "        self.n,self.bs,self.shuffle = (len(ds)//bs)*bs,bs,shuffle  \n",
    "        # Note: self.n = (len(ds)//bs) keeps the exact amount of samples needed for your desired batchSize\n",
    "        \n",
    "    def __iter__(self):\n",
    "        self.idxs = torch.randperm(self.n) if self.shuffle else torch.arange(self.n)\n",
    "        for i in range(0, self.n, self.bs): yield self.idxs[i:i+self.bs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `DataLoader` class combinds the dataset and the data sampler which iterates over the dataset and grabs batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate(batch_pairs):\n",
    "    x,y = zip(*batch_pairs)\n",
    "    return (x,y)\n",
    "\n",
    "class DataLoader():\n",
    "    def __init__(self, ds, sampler, collate_fn=collate):\n",
    "        self.ds,self.sampler,self.collate_fn = ds,sampler,collate_fn\n",
    "    def __len__(self): return len(self.ds)\n",
    "    def __iter__(self):\n",
    "        for s in self.sampler: yield self.collate_fn([self.ds[i] for i in s])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Embedding Layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `Custom_Embedding` class was used to project the high-dimensional multi-hot encoded vectors to a lower dimensional space prior to presenting the input data to the GRU. In this step the auther used two approaches \n",
    "1. Random intialization , then learn the appropriate $W_{(emb)}$ weights during back-prop \n",
    "    - $h_{i}^{(1)} = [tanh(x_{i}^{(T)} W_{(emb)} + b_{emb})]$\n",
    "\n",
    "2. Pre-trained embedding intialized using the Skik-gram algorithm, then refine weights during back-prop \n",
    "    - $h_{i}^{(1)} = [x_{i}^{(T)} W_{(emb)}]$\n",
    "    \n",
    "In this implementation of the paper we used the first approach. Therefore, the `Custom Embedding` class was created to created apply a tanh activation on the embedding layer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Custom_Embedding(nn.Module):\n",
    "    def __init__ (self, inputDimSize, embSize):\n",
    "        super(Custom_Embedding, self).__init__()\n",
    "        self.inputDimSize = inputDimSize\n",
    "        self.embSize = embSize\n",
    "        \n",
    "        self.W_emb = nn.Parameter(torch.randn(self.inputDimSize, self.embSize) * 0.01)\n",
    "        self.b_emb = nn.Parameter(torch.zeros(self.embSize) * 0.01) \n",
    "       \n",
    "    def forward(self, x):\n",
    "        return torch.tanh(x@self.W_emb + self.b_emb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dropout Layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this paper the author used the naive application of dropout that was first introduced by [Srivastava (2014)](http://jmlr.org/papers/volume15/srivastava14a.old/srivastava14a.pdf). While this method works well, it imapcts the RNNs ability to retain long term dependencies, because we are not maintaining the same mask across each timestep. Why is this important? It's simple, if we randonly sample a new mask at each time step, it perturbs our RNNs connections making it difficult for the network to determine what information might be relevant in the long term. In this approach, I tested the a technique proposed by [Gal & Ghahramani (2016)](phttps://arxiv.org/pdf/1512.05287.pdf) and further developed by [Merity (2017)](https://arxiv.org/pdf/1708.02182.pdf) for LSTMs. Here, they proposed overcoming the aforementioned problem with random sampling, by using the same dropout mask across multiple time steps in LSTMs. Here, I will applied the same approach on a GRU between each layer (two layers)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dropout_mask(x, sz, p):\n",
    "    return x.new(*sz).bernoulli_(1-p).div_(1-p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.0000, 0.0000, 1.4286, 0.0000, 1.4286, 0.0000, 1.4286, 1.4286,\n",
       "          1.4286, 1.4286, 1.4286]],\n",
       "\n",
       "        [[1.4286, 1.4286, 1.4286, 1.4286, 1.4286, 1.4286, 1.4286, 1.4286,\n",
       "          0.0000, 1.4286, 1.4286]]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dropout_mask = dropout_mask(x_artifical, (x_artifical.size(0),1,x_artifical.size(2)), 0.3)\n",
    "dropout_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000],\n",
       "         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.4286,\n",
       "          1.4286, 1.4286, 0.0000]],\n",
       "\n",
       "        [[1.4286, 1.4286, 1.4286, 1.4286, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000],\n",
       "         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000]]])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x_artifical * dropout_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.5286), tensor(0.4080))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x_artifical * dropout_mask).std(),x_artifical.std() \n",
    "# here you can see that we are maintaining about the same std as the original dataset without dropout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Doctor AI: Predicting Clinical Events via Recurrent Neural Networks\n",
    "Despite the popularity and preference given to LSTMs. This paper used a GRU architecture, for it's  simpilicty and ability to get similar performance as LSTMs. The dataset used in this paper contained `263, 706 patients`, whereas our dataset (MIMIC III) contained a total of `7537 patients`. However, the author demonstrated transfer learning can be a viable option in cases where one hostipal system lack the large scale datasets need to train deep learning models like Dr. AI. Using the following archiecture, my interest lies in the prediction of the patient's future diagnosis codes. However, one can easily extrapolate the algorithm to predict both diagnoses and duration between visits."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/Model_arch.png\" style=\"height:500px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EHR_GRU(Custom_Embedding):\n",
    "    def __init__(self, inputDimSize, hiddenDimSize, embSize, numClass, numLayers):\n",
    "        super().__init__(inputDimSize, embSize)\n",
    "        \n",
    "        self.numClass = numClass\n",
    "        self.numLayers = numLayers\n",
    "        self.hiddenDimSize = hiddenDimSize\n",
    "        self.emb = Custom_Embedding(inputDimSize, embSize)\n",
    "        \n",
    "        self.W_r = nn.Parameter(torch.randn(embSize, hiddenDimSize)* 0.01)\n",
    "        self.W_z = nn.Parameter(torch.randn(embSize, hiddenDimSize)* 0.01)\n",
    "        self.W_h = nn.Parameter(torch.randn(embSize, hiddenDimSize)* 0.01)\n",
    "        \n",
    "        self.U_r = nn.Parameter(torch.randn(hiddenDimSize, hiddenDimSize)* 0.01)\n",
    "        self.U_z = nn.Parameter(torch.randn(hiddenDimSize, hiddenDimSize)* 0.01)\n",
    "        self.U_h = nn.Parameter(torch.randn(hiddenDimSize, hiddenDimSize)* 0.01)\n",
    "        \n",
    "        self.b_r = nn.Parameter(torch.randn(hiddenDimSize))\n",
    "        self.b_z = nn.Parameter(torch.randn(hiddenDimSize))\n",
    "        self.b_h = nn.Parameter(torch.randn(hiddenDimSize))\n",
    "        \n",
    "        self.W_output = nn.Parameter(torch.randn(embSize, numClass))\n",
    "        self.b_output = nn.Parameter(torch.randn(numClass))\n",
    "        \n",
    "    def forward(self, emb, mask):\n",
    "        h = self.init_hidden(emb.size(1))\n",
    "        \n",
    "        z = torch.sigmoid(emb@self.W_z + h@self.U_z + self.b_z) \n",
    "        r = torch.sigmoid(emb@self.W_r + h@self.U_r + self.b_r)\n",
    "        h_tilde = torch.tanh(emb@self.W_h + (r * h)@self.U_h + self.b_h)\n",
    "        h_new = z * h + ((1. - z) * h_tilde)\n",
    "        h_new = mask[:, :, None] * h_new + (1. - mask)[:, :, None] * h\n",
    "       \n",
    "        return h_new\n",
    "    \n",
    "    def init_hidden(self, batchSize):\n",
    "        return Variable(torch.zeros(1, batchSize, hiddenDimSize))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GRU Layer:\n",
    "This class uses the `EHR_GRU` cell class and allows the iteration over the desired number of layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class build_EHR_GRU(EHR_GRU):\n",
    "    def __init__(self, GRUCell, *kwargs):\n",
    "        super().__init__(inputDimSize, hiddenDimSize, embSize, numClass, numLayers)\n",
    "        self.cell = GRUCell(*kwargs)\n",
    "        self.emb = Custom_Embedding(inputDimSize, embSize)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        inputVector = self.emb(x)\n",
    "        for i in range(numLayers):\n",
    "            memories = self.cell(inputVector, mask)\n",
    "            drop_out = dropout_mask(inputVector, (inputVector.size(0), 1, inputVector.size(2)), 0.5)\n",
    "            inputVector = memories * drop_out\n",
    "        \n",
    "        y_linear = inputVector@self.W_output + self.b_output\n",
    "        output = F.softmax(y_linear, dim=1)\n",
    "        output = output * mask[:,:,None]\n",
    "        return output, inputVector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loss Function:\n",
    "The loss function used to assess model perform, contained a combination of the cross entropy. The prediction loss for each mini-batch was normalized to the sequence length. Finally, L2-norm regularization was applied to all of the weight matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "class cost_function():\n",
    "    def __init__(self, yhat, y, L_2=0.001, logEps=1e-8):\n",
    "        self.yhat = yhat\n",
    "        self.y = y\n",
    "       \n",
    "        self.logEps = logEps\n",
    "        self.L_2 = L_2\n",
    "        \n",
    "        self.W_out = nn.Parameter(torch.randn(hiddenDimSize, numClass)*0.01)\n",
    "        \n",
    "    def cross_entropy(self):\n",
    "        return  -(self.y * torch.log(self.yhat + self.logEps) + (1. - self.y) * torch.log(1. - self.yhat + self.logEps))\n",
    "    \n",
    "    def prediction_loss(self):\n",
    "        return  (torch.sum(torch.sum(self.cross_entropy(), dim=0),dim=1)).float()/  lengths.float()\n",
    "    \n",
    "    def cost(self):\n",
    "        return torch.mean(self.prediction_loss()) + self.L_2 * (self.W_out ** 2).sum() # regularize\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Parameters:\n",
    "The parameters used here were selected from those used in the Dr AI paper. The major difference between this approach and what I present here, was my use of the more updated drop out approach for RNNs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "numClass = 4894\n",
    "inputDimSize = 4894\n",
    "embSize = 200\n",
    "hiddenDimSize = 200\n",
    "batchSize = 100\n",
    "numLayers = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data:\n",
    "It's key to note that you want to pass in the same file for the sequences and lables into the `load_data` function, as the model will take care of the adjusting the time steps for prediction internally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, valid, test = load_data('data/Jan19.seqs', 'data/Jan19.seqs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds= Dataset(train[0], train[1])\n",
    "train_samp = Sampler(train_ds, batchSize, shuffle=True)\n",
    "train_dl = DataLoader(train_ds, sampler=train_samp, collate_fn=collate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_ds= Dataset(valid[0], valid[1])\n",
    "valid_samp = Sampler(valid_ds, batchSize, shuffle=False)\n",
    "valid_dl = DataLoader(valid_ds, sampler=valid_samp, collate_fn=collate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_EHR_GRU(EHR_GRU, inputDimSize, hiddenDimSize, embSize, numClass, numLayers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training and validation loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1/10... Step: 0... Training Loss: 107.3128... Val Loss: 107.5775\n",
      "Epoch: 2/10... Step: 0... Training Loss: 107.1489... Val Loss: 107.6264\n",
      "Epoch: 3/10... Step: 0... Training Loss: 106.9846... Val Loss: 107.7252\n",
      "Epoch: 4/10... Step: 0... Training Loss: 107.2707... Val Loss: 107.6710\n",
      "Epoch: 5/10... Step: 0... Training Loss: 107.2436... Val Loss: 107.5871\n",
      "Epoch: 6/10... Step: 0... Training Loss: 107.4828... Val Loss: 107.6925\n",
      "Epoch: 7/10... Step: 0... Training Loss: 107.0155... Val Loss: 107.4827\n",
      "Epoch: 8/10... Step: 0... Training Loss: 107.0334... Val Loss: 107.5125\n",
      "Epoch: 9/10... Step: 0... Training Loss: 107.2941... Val Loss: 107.5251\n",
      "Epoch: 10/10... Step: 0... Training Loss: 107.0258... Val Loss: 107.4568\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adadelta(model.parameters(), lr = 0.01, rho=0.95)\n",
    "epochs = 10\n",
    "\n",
    "counter = 0\n",
    "for e in range(epochs):\n",
    "    for x, y in train_dl:\n",
    "        x, y , mask, lengths = padding(x, y, inputDimSize, numClass)\n",
    "        \n",
    "        output, h = model(x, mask)\n",
    "        \n",
    "        loss = cost_function(output, y).cost()\n",
    "        loss.backward()\n",
    "        nn.utils.clip_grad_norm_(model.parameters(), 5)\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "            model.eval()\n",
    "            val_loss = []\n",
    "            for x_valid, y_valid in valid_dl:\n",
    "                    x_val, y_val, mask, lengths = padding(x_valid, y_valid, inputDimSize, numClass)\n",
    "                    outputs_val, hidden_val = model(x_val,  mask)\n",
    "                    loss = cost_function(outputs_val, y_val).cost()\n",
    "                    val_loss.append(loss.item())\n",
    "            model.train()\n",
    "\n",
    "            print(\"Epoch: {}/{}...\".format(e+1, epochs),\n",
    "                                  \"Step: {}...\".format(counter),\n",
    "                                  \"Training Loss: {:.4f}...\".format(loss.item()),\n",
    "                                  \"Val Loss: {:.4f}\".format(torch.mean(torch.tensor(val_loss))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of my implementation to the paper's algorithm:\n",
    "I ran the same sequences on the paper's algorithm, which is written in theano and pyton 2.7 and here you can see that the best cross entropy score after 10 epochs is about 86.79 vs. my 107. While, I am not performing better with some more hyperparmeter tuning and optimization the algorithm can definitely perform better.\n",
    "<img src=\"img/Comparing_with_actual_algo.png\" style=\"height:500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Observations:\n",
    "As you can see our training and vaildation losses are about the same, with such a small subset of the data used in the actual paper. It might be difficult to get better performance without overfitting. However, the intent of this tutorial was to provide a detailed walkthrough of how one can use EHR data to drive insights!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next Steps:\n",
    "1. Add Callbacks using Fast.AI's callback approach to track in training stats\n",
    "2. Play around with different intialization approaches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Acknowledgements:\n",
    "1. Fast.ai (Rachel Thomas, Jeremey Howard, and the amazing fast.ai community)\n",
    "2. Dorian Puleri"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}