notebooks/mnist_survival_analysis_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Survival Analysis with ML4H"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Survival analysis_ is a core modeling paradigm of clinical ML that data scientists from other domains are unlikely to be familiar with. In survival analysis, the aim is to predict when an event might occur, such as a heart attack, stroke, or the onset of heart feailure. Survival analysis can also require incorporating other time-dependent phenomenon, such as patients dropping out of a study, or the random times at which patients see a provider and generate new data.  \n",
    "\n",
    "In this notebook, we will use the classic MNIST dataset to develop a toy model of survival analysis using ML4H. First, we will train a classifier on MNIST using ML4H abstractions. Next, we will use the digit label to determine the threshold for randomly generating synthetic events. Finally, we will incorporate a random lag to model the process of our \"patients\" seeing their physicians in a health system.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "<div class=\"alert alert-block alert-warning\">\n",
    "    If you're running this notebook on Terra, use custom Docker image: <kbd>ghcr.io/broadinstitute/ml4h/ml4h_terra:20211101_143643</kbd>.</br>\n",
    "    This notebook will run fine on a CPU, but you can also run it on a <a href=https://support.terra.bio/hc/en-us/articles/4403006001947>GPU-enabled Cloud Environment</a> if you like.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import os\n",
    "import sys\n",
    "import pickle\n",
    "from typing import List, Dict, Callable\n",
    "from collections import defaultdict, Counter\n",
    "\n",
    "import gzip\n",
    "import h5py\n",
    "import shutil\n",
    "import zipfile\n",
    "import numpy as np\n",
    "\n",
    "from ml4h.arguments import parse_args\n",
    "from ml4h.recipes import train_multimodal_multitask\n",
    "from ml4h.TensorMap import TensorMap, Interpretation\n",
    "from ml4h.models.train import train_model_from_generators\n",
    "from ml4h.models.legacy_models import make_multimodal_multitask_model\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Paths\n",
    "HD5_FOLDER = './mnist_hd5s/'\n",
    "OUTPUT_FOLDER = './runs/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(dataset):\n",
    "    ''' Loads the dataset\n",
    "    :param dataset: the path to the dataset (here MNIST)'''\n",
    "    data_dir, data_file = os.path.split(dataset)\n",
    "    if data_dir == \"\" and not os.path.isfile(dataset):\n",
    "        # Check if dataset is in the data directory.\n",
    "        new_path = os.path.join(\"data\", dataset)\n",
    "        if os.path.isfile(new_path) or data_file == 'mnist.pkl.gz':\n",
    "            dataset = new_path\n",
    "\n",
    "    if (not os.path.isfile(dataset)) and data_file == 'mnist.pkl.gz':\n",
    "        from urllib.request import urlretrieve\n",
    "        origin = ('http://www.iro.umontreal.ca/~lisa/deep/data/mnist/mnist.pkl.gz')\n",
    "        print('Downloading data from %s' % origin)\n",
    "        if not os.path.exists(os.path.dirname(dataset)):\n",
    "            os.makedirs(os.path.dirname(dataset))\t\n",
    "        urlretrieve(origin, dataset)\n",
    "\n",
    "    print('loading data...')\n",
    "    f = gzip.open(dataset, 'rb')\n",
    "    if sys.version_info[0] == 3:\n",
    "        u = pickle._Unpickler(f)\n",
    "        u.encoding = 'latin1'\n",
    "        train_set, valid_set, test_set = u.load()\n",
    "    else:\n",
    "        train_set, valid_set, test_set = pickle.load(f)\n",
    "    f.close()\n",
    "\n",
    "    return train_set, valid_set, test_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_mnist(sides):\n",
    "    train, _, _ = load_data('mnist.pkl.gz')\n",
    "    print(train[0].shape)\n",
    "    mnist_images = train[0].reshape((-1, 28, 28, 1))\n",
    "    sides = int(np.ceil(np.sqrt(min(sides, mnist_images.shape[0]))))\n",
    "    _, axes = plt.subplots(sides, sides, figsize=(16, 16))\n",
    "    for i in range(sides*sides):\n",
    "        axes[i // sides, i % sides].imshow(mnist_images[i, ..., 0], cmap='gray')\n",
    "        axes[i // sides, i % sides].set_xticks(())\n",
    "        axes[i // sides, i % sides].set_yticks(())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mnist(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mnist(25)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensorization\n",
    "It is often helpful to separate data preparation from model training.  In ml4h we call the final data preparation process tensorization.  Tensorization involves gathering all input files (XMLS, CSVs, DICOMs, PNGs, etc) and consolidating them into compressed HD5 files.  We tend to make one HD5 file per individual in the cohort we are studying.  The files contain the raw data and labels (inputs and outputs) we will use to train models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnist_as_hd5(hd5_folder):\n",
    "    train, _, _ = load_data('mnist.pkl.gz')\n",
    "    mnist_images = train[0].reshape((-1, 28, 28, 1))\n",
    "    if not os.path.exists(hd5_folder):\n",
    "        os.makedirs(hd5_folder)\n",
    "    for i, mnist_image in enumerate(mnist_images):\n",
    "        with h5py.File(os.path.join(hd5_folder, f'{i}.hd5'), 'w') as hd5:\n",
    "            hd5.create_dataset('mnist_image', data=mnist_image)\n",
    "            hd5.create_dataset('mnist_label', data=[train[1][i]])\n",
    "        if (i+1) % 5000 == 0:\n",
    "            print(f'Wrote {i+1} MNIST images and labels as HD5 files')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_as_hd5(HD5_FOLDER)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TensorMaps\n",
    "The central data structure in the ml4h codebase is the TensorMap.\n",
    "This abstraction provides a way to translate ***any*** kind of input data, into structured numeric tensors with clear semantics for interpretation and modeling.  TensorMaps guarantee a shape, a way to construct tensors of that shape from the HD5 files created during tensorization and a meaning to the values in the tensor that the TensorMap yields.  The most important method of each TensorMap is their ***tensor_from_file*** function.  This callback function takes the TensorMap, an HD5 file handle, and an optional dictionary as input arguments and it returns a numpy array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnist_image_from_hd5(tm, hd5, dependents={}):\n",
    "     return np.array(hd5['mnist_image'])\n",
    "\n",
    "def mnist_label_from_hd5(tm, hd5, dependents={}):\n",
    "    one_hot = np.zeros(tm.shape, dtype=np.float32)\n",
    "    one_hot[int(hd5['mnist_label'][0])] = 1.0\n",
    "    return one_hot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Command Line Friendly Jupyter Notebooks\n",
    "By directly setting the `sys.argv` array in our jupyter notebooks we make the process of translating from notebook to command line straightforward.  For example, the cell below can be replicated on the command line by running:\n",
    "```\n",
    "./scripts/tf.sh $HOME/ml/ml4h/recipes.py --mode train --tensors ./mnist_hd5s/ \\\n",
    "    --input_tensors mnist_image --output_tensors mnist_label \\\n",
    "    --batch_size 64 --test_steps 64 --epochs 24 \\\n",
    "    --output_folder ./runs/ --id learn_mnist\n",
    "```\n",
    "The script `tf.sh` starts the appropriate docker container and then calls python on the provided arguments.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.argv = ['train', \n",
    "            '--tensors', HD5_FOLDER, \n",
    "            '--input_tensors', 'mnist.mnist_image',\n",
    "            '--output_tensors', 'mnist.mnist_label',\n",
    "            '--batch_size', '64',\n",
    "            '--test_steps', '64',\n",
    "            '--epochs', '6',\n",
    "            '--output_folder', OUTPUT_FOLDER,\n",
    "            '--id', 'learn_mnist'\n",
    "           ]\n",
    "args = parse_args()\n",
    "metrics = train_multimodal_multitask(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating ML Models in ml4h\n",
    "Each ml4h training run creates several plots to give insight into model performance and learning dynamics.  The plots created will depend on the TensorMaps used but in general will include a metric history showing learning curves of each metric tracked during training, performance plots like ROC and Precision Recall curves for classifiers or scatter plots for regressors, calibration plots, and a t-SNE plot showing a 2D representation of the learned embedding of the trained model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Survival Analysis\n",
    "In addition to categorical classification tasks shown above, the ML4Health code base supports survival analysis.  In survival analysis, we consider the relative time before an outcome was observed, and we allow for samples to leave our study at anytime, before having an event.  Survival analysis is implemented via the TensorMap interpretations TIME_TO_EVENT and SURVIVAL_CURVE. \n",
    "\n",
    "The TIME_TO_EVENT interpretation is the most straight forward.  TIME_TO_EVENT TensorMaps *tensor_from_file* function returns an array with two values.  The first value indicates whether or not an event occurred for this sample.  The second value indicates the total days of follow up from when the sample enrolled in the study until either an event ocurred or they left the study.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defined in ml4h/tensormap/mnist.py \n",
    "def mnist_label_as_time_to_event(tm, hd5, dependents={}):\n",
    "    tensor = np.zeros(tm.shape, dtype=np.float32)\n",
    "    label = float(hd5['mnist_label'][0])\n",
    "    tensor[0] = 1.0 if np.random.rand() > (label / 10) else 0.0\n",
    "    tensor[1] = np.random.randint(1, 3650)\n",
    "    return tensor\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.argv = ['train', \n",
    "            '--tensors', HD5_FOLDER, \n",
    "            '--input_tensors', 'mnist.mnist_image',\n",
    "            '--output_tensors', 'mnist.mnist_time_to_event',\n",
    "            '--training_steps', '64',\n",
    "            '--validation_steps', '24',\n",
    "            '--test_steps', '30',\n",
    "            '--batch_size', '32',\n",
    "            '--epochs', '6',\n",
    "            '--eager',\n",
    "            '--output_folder', OUTPUT_FOLDER,\n",
    "            '--id', 'mnist_time_to_event'\n",
    "           ]\n",
    "args = parse_args()\n",
    "train_multimodal_multitask(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model the Entire Survival Curve\n",
    "SURVIVAL_CURVE TensorMaps model the survival curve with quantized steps.  These TensorMaps **tensor_from_file** functions return an array of 1s and 0s. The first half of the array contains 1s for each time step that the sample survived and 0s elsewhere. The second half of the array contains a 1 at the time step at which an event occurred if there was one, otherwise it contains only zeros.  Here we use the MNIST dataset to create a synthetic survival analysis cohort.  We use the MNIST label to determine the likelihood of having events and we model several different shapes of survival curve and distributions of follow up.  The **days_window** and **shape** fields of these TensorMaps controls how quantized the predicted survival curves will be.  Specifically each time bin will cover **days_window** / (**shape[-1]**/2) days. **shape** should be an even number because half the array is used for indicating survival up until an event or censorship and the other half indicates events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defined in ml4h/tensormap/mnist.py\n",
    "def _mnist_label_as_survival_curve(tm, hd5, dependents={}):\n",
    "    label = float(hd5['mnist_label'][0])\n",
    "    has_disease = 1.0 if np.random.rand() > (label / 10) else 0.0\n",
    "    days_follow_up = np.random.randint(1, 3650)\n",
    "        \n",
    "    intervals = int(tm.shape[0] / 2)\n",
    "    days_per_interval = tm.days_window / intervals\n",
    "    survival_then_censor = np.zeros(tm.shape, dtype=np.float32)\n",
    "    for i, day_delta in enumerate(np.arange(0, tm.days_window, days_per_interval)):\n",
    "        survival_then_censor[i] = float(day_delta < days_follow_up)\n",
    "        if day_delta <= days_follow_up < day_delta + days_per_interval:\n",
    "            survival_then_censor[intervals+i] = has_disease\n",
    "    return survival_then_censor\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.argv = ['train', \n",
    "            '--tensors', HD5_FOLDER, \n",
    "            '--input_tensors', 'mnist.mnist_image',\n",
    "            '--output_tensors', 'mnist.mnist_survival_curve',\n",
    "            '--training_steps', '128',\n",
    "            '--validation_steps', '32',\n",
    "            '--test_steps', '32',\n",
    "            '--batch_size', '64',\n",
    "            '--epochs', '8',\n",
    "            '--output_folder', OUTPUT_FOLDER,\n",
    "            '--id', 'mnist_label_as_survival_curve'\n",
    "           ]\n",
    "args = parse_args()\n",
    "train_multimodal_multitask(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Provenance\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "print(datetime.datetime.now())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "pip3 freeze"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Questions about these particular notebooks?  Join the discussion: https://github.com/broadinstitute/ml4h/discussions."
   ]
  }
 ],
 "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.6.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "325px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}