DL4H_Team_24_Project_Draft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/abarrie2/cs598-dlh-project/blob/main/DL4H_Team_24.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LICENSING NOTICE\n",
    "\n",
    "Note that all users who use Vital DB, an open biosignal dataset, must agree to the Data Use Agreement below. If you do not agree, please close this window. The Data Use Agreement is available here:\n",
    "https://vitaldb.net/dataset/#h.vcpgs1yemdb5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Project Draft\n",
    "\n",
    "This is the Project Draft submission. For the development version see the DL4H_Team_24.ipynb notebook in the project repository."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Project repository\n",
    "\n",
    "The project repository can be found at: https://github.com/abarrie2/cs598-dlh-project"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MQ0sNuMePBXx"
   },
   "source": [
    "# Introduction\n",
    "\n",
    "This project aims to reproduce findings from the paper titled \"Predicting intraoperative hypotension using deep learning with waveforms of arterial blood pressure, electroencephalogram, and electrocardiogram: Retrospective study\" by Jo Y-Y et al. (2022) [1]. This study introduces a deep learning model that predicts intraoperative hypotension (IOH) events before they occur, utilizing a combination of arterial blood pressure (ABP), electroencephalogram (EEG), and electrocardiogram (ECG) signals.\n",
    "\n",
    "\n",
    "## Background of the Problem\n",
    "\n",
    "Intraoperative hypotension (IOH) is a common and significant surgical complication defined by a mean arterial pressure drop below 65 mmHg. It is associated with increased risks of myocardial infarction, acute kidney injury, and heightened postoperative mortality. Effective prediction and timely intervention can substantially enhance patient outcomes.\n",
    "\n",
    "### Evolution of IOH Prediction\n",
    "\n",
    "Initial attempts to predict IOH primarily used arterial blood pressure (ABP) waveforms. A foundational study by Hatib F et al. (2018) titled \"Machine-learning Algorithm to Predict Hypotension Based on High-fidelity Arterial Pressure Waveform Analysis\" [2] showed that machine learning could forecast IOH events using ABP with reasonable accuracy. This finding spurred further research into utilizing various physiological signals for IOH prediction.\n",
    "\n",
    "Subsequent advancements included the development of the Acumen™ hypotension prediction index, which was studied in \"AcumenTM hypotension prediction index guidance for prevention and treatment of hypotension in noncardiac surgery: a prospective, single-arm, multicenter trial\" by Bao X et al. (2024) [3]. This trial integrated a hypotension prediction index into blood pressure monitoring equipment, demonstrating its effectiveness in reducing the number and duration of IOH events during surgeries. Further study is needed to determine whether this resultant reduction in IOH events transalates into improved postoperative patient outcomes.\n",
    "\n",
    "\n",
    "### Current Study\n",
    "\n",
    "Building on these advancements, the paper by Jo Y-Y et al. (2022) proposes a deep learning approach that enhances prediction accuracy by incorporating EEG and ECG signals along with ABP. This multi-modal method, evaluated over prediction windows of 3, 5, 10, and 15 minutes, aims to provide a comprehensive physiological profile that could predict IOH more accurately and earlier. Their results indicate that the combination of ABP and EEG significantly improves performance metrics such as AUROC and AUPRC, outperforming models that use fewer signals or different combinations.\n",
    "\n",
    "Our project seeks to reproduce and verify Jo Y-Y et al.'s results to assess whether this integrated approach can indeed improve IOH prediction accuracy, thereby potentially enhancing surgical safety and patient outcomes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uygL9tTPSVHB"
   },
   "source": [
    "# Scope of Reproducibility:\n",
    "\n",
    "The original paper investigated the following hypotheses:\n",
    "\n",
    "1.   Hypothesis 1: A model using ABP and ECG will outperform a model using ABP alone in predicting IOH.\n",
    "2.   Hypothesis 2: A model using ABP and EEG will outperform a model using ABP alone in predicting IOH.\n",
    "3.   Hypothesis 3: A model using ABP, EEG, and ECG will outperform a model using ABP alone in predicting IOH.\n",
    "\n",
    "Results were compared using AUROC and AUPRC scores. Based on the results described in the original paper, we expect that Hypothesis 2 will be confirmed, and that Hypotheses 1 and 3 will not be confirmed.\n",
    "\n",
    "In order to perform the corresponding experiments, we will implement a CNN-based model that can be configured to train and infer using the following four model variations:\n",
    "\n",
    "1.   ABP data alone\n",
    "2.   ABP and ECG data\n",
    "3.   ABP and EEG data\n",
    "4.   ABP, ECG, and EEG data\n",
    "\n",
    "We will measure the performance of these configurations using the same AUROC and AUPRC metrics as used in the original paper. To test hypothesis 1 we will compare the AUROC and AUPRC measures between model variation 1 and model variation 2. To test hypothesis 2 we will compare the AUROC and AUPRC measures between model variation 1 and model variation 3. To test hypothesis 3 we will compare the AUROC and AUPRC measures between model variation 1 and model variation 4. For all of the above measures and experiment combinations, we will operate multiple experiments where the time-to-IOH event prediction will use the following prediction windows:\n",
    "\n",
    "1. 3 minutes before event\n",
    "2. 5 minutes before event\n",
    "3. 10 minutes before event\n",
    "4. 15 minutes before event\n",
    "\n",
    "In the event that we are compute-bound, we will prioritize the 3-minute prediction window experiments as they are the most relevant to the original paper's findings.\n",
    "\n",
    "The predictive power of ABP, ECG and ABP + ECG models at 3-, 5-, 10- and 15-minute prediction windows:\n",
    "![Predictive power of ABP, ECG and ABP + ECG models at 3-, 5-, 10- and 15-minute prediction windows](https://journals.plos.org/plosone/article/figure/image?download&size=large&id=10.1371/journal.pone.0272055.g004)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modifications made for demo mode\n",
    "\n",
    "In order to demonstrate the functioning of the code in a short (ie, <8 minute limit) the following options and modifications were used:\n",
    "\n",
    "1. `MAX_CASES` was set to 20. The total number of cases to be used in the full training set is 3296, but the smaller numbers allows demonstration of each section of the pipeline.\n",
    "2. `vitaldb_cache` is prepopulated in Google Colab. The cache file is approx. 800MB and contains the raw and mini-fied copies of the source dataset and is downloaded from Google Drive. This is much faster than using the `vitaldb` API, but is again only a fraction of the data. The full dataset can be downloaded with the API or prepopulated by following the instructions in the \"Bulk Data Download\" section below.\n",
    "3. `max_epochs` is set to 6. With the small dataset, training is fast and shows the decreasing training and validation losses. In the full model run, `max_epochs` will be set to 100. In both cases early stopping is enabled and will stop training if the validation losses stop decreasing for five consecutive epochs.\n",
    "4. Only the \"ABP + EEG\" combination will be run. In the final report, additional combinations will be run, as discussed later.\n",
    "5. Only the 3-minute prediction window will be run. In the final report, additional prediction windows (5, 10 and 15 minutes) will be run, as discussed later.\n",
    "6. No ablations are run in the demo. These will be completed for the final report."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xWAHJ_1CdtaA"
   },
   "source": [
    "# Methodology\n",
    "\n",
    "The methodology section is composed of the following subsections: Environment, Data and Model.\n",
    "\n",
    "- **Environment**: This section describes the setup of the environment, including the installation of necessary libraries and the configuration of the runtime environment.\n",
    "- **Data**: This section describes the dataset used in the study, including its collection and preprocessing.\n",
    "    - **Data Collection**: This section describes the process of downloading the dataset from VitalDB and populating the local data cache.\n",
    "    - **Data Preprocessing**: This section describes the preprocessing steps applied to the dataset, including data selection, data cleaning, and feature extraction.\n",
    "- **Model**: This section describes the deep learning model used in the study, including its implementation, training, and evaluation.\n",
    "    - **Model Implementation**: This section describes the implementation of the deep learning model, including the architecture, loss function, and optimization algorithm.\n",
    "    - **Model Training**: This section describes the training process, including the training loop, hyperparameters, and training strategy.\n",
    "    - **Model Evaluation**: This section describes the evaluation process, including the metrics used, the evaluation strategy, and the results obtained."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment\n",
    "\n",
    "### Create environment\n",
    "\n",
    "The environment setup differs based on whether you are running the code on a local machine or on Google Colab. The following sections provide instructions for setting up the environment in each case.\n",
    "\n",
    "#### Local machine\n",
    "\n",
    "Create `conda` environment for the project using the `environment.yml` file:\n",
    "\n",
    "```bash\n",
    "conda env create --prefix .envs/dlh-team24 -f environment.yml\n",
    "```\n",
    "\n",
    "Activate the environment with:\n",
    "```bash\n",
    "conda activate .envs/dlh-team24\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Google Colab\n",
    "\n",
    "The following code snippet installs the required packages and downloads the necessary files in a Google Colab environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yu61Jp1xrnKk"
   },
   "outputs": [],
   "source": [
    "#install vitaldb\n",
    "%pip install vitaldb\n",
    "\n",
    "# Download cached preprocessed data\n",
    "import gdown\n",
    "gdown.download(id=\"15b5Nfhgj3McSO2GmkVUKkhSSxQXX14hJ\", output=\"vitaldb_cache.tgz\")\n",
    "!tar -zxf vitaldb_cache.tgz\n",
    "\n",
    "# Download sqi_filter.csv from github repo\n",
    "!wget https://raw.githubusercontent.com/abarrie2/cs598-dlh-project/main/sqi_filter.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All other required packages are already installed in the Google Colab environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import packages\n",
    "import os\n",
    "import random\n",
    "\n",
    "from timeit import default_timer as timer\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import roc_auc_score, precision_recall_curve, auc, confusion_matrix\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "import vitaldb\n",
    "import pickle\n",
    "import _pickle as cPickle\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from tqdm import tqdm\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set random seeds to generate consistent results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RANDOM_SEED = 42\n",
    "random.seed(RANDOM_SEED)\n",
    "np.random.seed(RANDOM_SEED)\n",
    "torch.manual_seed(RANDOM_SEED)\n",
    "#torch.use_deterministic_algorithms(False)\n",
    "os.environ['PYTHONASHSEED'] = str(RANDOM_SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2NbPHUTMbkD3"
   },
   "source": [
    "##  Data\n",
    "\n",
    "### Data Description\n",
    "\n",
    "#### Source\n",
    "\n",
    "Data for this project is sourced from the open biosignal VitalDB dataset as described in \"VitalDB, a high-fidelity multi-parameter vital signs database in surgical patients\" by Lee H-C et al. (2022) [4], which contains perioperative vital signs and numerical data from 6,388 cases of non-cardiac (general, thoracic, urological, and gynecological) surgery patients who underwent routine or emergency surgery at Seoul National University Hospital between 2016 and 2017. The dataset includes ABP, ECG, and EEG signals, as well as other physiological data. The dataset is available through an [API](https://vitaldb.net/dataset/?query=api) and [Python library](https://vitaldb.net/dataset/?query=lib), and at PhysioNet: https://physionet.org/content/vitaldb/1.0.0/\n",
    "\n",
    "#### Statistics\n",
    "\n",
    "Characteristics of the dataset:\n",
    "| Characteristic        | Value                       | Details                |\n",
    "|-----------------------|-----------------------------|------------------------|\n",
    "| Total number of cases | 6,388                       |                        |\n",
    "| Sex (male)            | 3,243 (50.8%)               |                        |\n",
    "| Age (years)           | 59                          | Range: 48-68           |\n",
    "| Height (cm)           | 162                         | Range: 156-169         |\n",
    "| Weight (kg)           | 61                          | Range: 53-69           |\n",
    "| Tram-Rac 4A tracks    | 6,355 (99.5%)               | Sampling rate: 500Hz   |\n",
    "| BIS Vista tracks      | 5,566 (87.1%)               | Sampling rate: 128Hz   |\n",
    "| Case duration (min)   | 189                         | Range: 27-1041         |\n",
    "\n",
    "Labels are only known after processing the data. In the original paper, there were an average of 1.6 IOH events per case and 5.7 non-events per case so we expect approximately 10,221 IOH events and 364,116 non-events in the dataset.\n",
    "\n",
    "#### Data Processing\n",
    "\n",
    "Data will be processed as follows:\n",
    "1. Load the dataset from VitalDB, or from a local cache if previously downloaded.\n",
    "2. Apply the inclusion and exclusion selection criteria to filter the dataset according to surgery metadata.\n",
    "3. Generate a minified dataset by discarding all tracks except ABP, ECG, and EEG.\n",
    "4. Preprocess the data by applying band-pass and z-score normalization to the ECG and EEG signals, and filtering out ABP signals below a Signal Quality Index (SQI) threshold.\n",
    "5. Generate event and non-event samples by extracting 60-second segments around IOH events and non-events.\n",
    "6. Split the dataset into training, validation, and test sets with a 6:1:3 ratio, ensuring that samples from a single case are not split across different sets to avoid data leakage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set Up Local Data Caches\n",
    "\n",
    "VitalDB data is static, so local copies can be stored and reused to avoid expensive downloads and to speed up data processing.\n",
    "\n",
    "The default directory defined below is in the project `.gitignore` file. If this is modified, the new directory should also be added to the project `.gitignore`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VITALDB_CACHE = './vitaldb_cache'\n",
    "VITAL_ALL = f\"{VITALDB_CACHE}/vital_all\"\n",
    "VITAL_MINI = f\"{VITALDB_CACHE}/vital_mini\"\n",
    "VITAL_METADATA = f\"{VITALDB_CACHE}/metadata\"\n",
    "VITAL_MODELS = f\"{VITALDB_CACHE}/models\"\n",
    "VITAL_PREPROCESS_SCRATCH = f\"{VITALDB_CACHE}/data_scratch\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRACK_CACHE = None\n",
    "# when USE_DISK_CACHING is enabled, track and segment data will be flushed to disk\n",
    "USE_DISK_CACHING = False\n",
    "\n",
    "# When RESET_CACHE is set to True, it will ensure the TRACK_CACHE is disposed and recreated when we do dataset initialization. Use as a shortcut to wiping cache rather than restarting kernel\n",
    "RESET_CACHE = False\n",
    "PREDICTION_WINDOW = 3\n",
    "\n",
    "# Maximum number of cases of interest for which to download data.\n",
    "# Set to a small value (ex: 20) for demo purposes, else set to None to disable and download and process all.\n",
    "#MAX_CASES = None\n",
    "MAX_CASES = 20\n",
    "\n",
    "# Preloading Cases: when true, all matched cases will have the _mini tracks extracted and put into in-mem dict\n",
    "PRELOADING_CASES = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(VITALDB_CACHE):\n",
    "  os.mkdir(VITALDB_CACHE)\n",
    "if not os.path.exists(VITAL_ALL):\n",
    "  os.mkdir(VITAL_ALL)\n",
    "if not os.path.exists(VITAL_MINI):\n",
    "  os.mkdir(VITAL_MINI)\n",
    "if not os.path.exists(VITAL_METADATA):\n",
    "  os.mkdir(VITAL_METADATA)\n",
    "if not os.path.exists(VITAL_MODELS):\n",
    "  os.mkdir(VITAL_MODELS)\n",
    "if not os.path.exists(VITAL_PREPROCESS_SCRATCH):\n",
    "  os.mkdir(VITAL_PREPROCESS_SCRATCH)\n",
    "\n",
    "print(os.listdir(VITALDB_CACHE))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bulk Data Download\n",
    "\n",
    "**This step is not required, but will significantly speed up downstream processing and avoid a high volume of API requests to the VitalDB web site.**\n",
    "\n",
    "The cache population code checks if the `.vital` files are locally available, and can be populated by calling the vitaldb API or by manually prepopulating the cache (recommended)\n",
    "\n",
    "- Manually downloaded the dataset from the following site: https://physionet.org/content/vitaldb/1.0.0/\n",
    "    - Download the [zip file](https://physionet.org/static/published-projects/vitaldb/vitaldb-a-high-fidelity-multi-parameter-vital-signs-database-in-surgical-patients-1.0.0.zip) in a browser, or\n",
    "    - Use `wget -r -N -c -np https://physionet.org/files/vitaldb/1.0.0/` to download the files in a terminal\n",
    "- Move the contents of `vital_files` into the `${VITAL_ALL}` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Returns the Pandas DataFrame for the specified dataset.\n",
    "#   One of 'cases', 'labs', or 'trks'\n",
    "# If the file exists locally, create and return the DataFrame.\n",
    "# Else, download and cache the csv first, then return the DataFrame.\n",
    "def vitaldb_dataframe_loader(dataset_name):\n",
    "    if dataset_name not in ['cases', 'labs', 'trks']:\n",
    "        raise ValueError(f'Invalid dataset name: {dataset_name}')\n",
    "    file_path = f'{VITAL_METADATA}/{dataset_name}.csv'\n",
    "    if os.path.isfile(file_path):\n",
    "        print(f'{dataset_name}.csv exists locally.')\n",
    "        df = pd.read_csv(file_path)\n",
    "        return df\n",
    "    else:\n",
    "        print(f'downloading {dataset_name} and storing in the local cache for future reuse.')\n",
    "        df = pd.read_csv(f'https://api.vitaldb.net/{dataset_name}')\n",
    "        df.to_csv(file_path, index=False)\n",
    "        return df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploratory Data Analysis\n",
    "\n",
    "#### Cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cases = vitaldb_dataframe_loader('cases')\n",
    "cases = cases.set_index('caseid')\n",
    "cases.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cases.index.nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cases.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cases['sex'].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tracks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks = vitaldb_dataframe_loader('trks')\n",
    "trks = trks.set_index('caseid')\n",
    "trks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks.index.nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks.groupby('caseid')[['tid']].count().plot();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks.groupby('caseid')[['tid']].count().hist();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks.groupby('tname').count().sort_values(by='tid', ascending=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters of Interest"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hemodynamic Parameters Reference\n",
    "https://vitaldb.net/dataset/?query=overview#h.f7d712ycdpk2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SNUADC/ART**\n",
    "\n",
    "arterial blood pressure waveform\n",
    "\n",
    "Parameter, Description, Type/Hz, Unit\n",
    "\n",
    "SNUADC/ART, Arterial pressure wave, W/500, mmHg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks[trks['tname'].str.contains('SNUADC/ART')].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SNUADC/ECG_II**\n",
    "\n",
    "electrocardiogram waveform\n",
    "\n",
    "Parameter, Description, Type/Hz, Unit\n",
    "\n",
    "SNUADC/ECG_II, ECG lead II wave, W/500, mV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks[trks['tname'].str.contains('SNUADC/ECG_II')].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**BIS/EEG1_WAV**\n",
    "\n",
    "electroencephalogram waveform\n",
    "\n",
    "Parameter, Description, Type/Hz, Unit\n",
    "\n",
    "BIS/EEG1_WAV, EEG wave from channel 1, W/128, uV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks[trks['tname'].str.contains('BIS/EEG1_WAV')].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cases of Interest\n",
    "\n",
    "These are the subset of case ids for which modelling and analysis will be performed based upon inclusion criteria and waveform data availability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRACK_NAMES = ['SNUADC/ART', 'SNUADC/ECG_II', 'BIS/EEG1_WAV']\n",
    "TRACK_SRATES = [500, 500, 128]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# As in the paper, select cases which meet the following criteria:\n",
    "#\n",
    "# For patients, the inclusion criteria were as follows:\n",
    "# (1) adults (age >= 18)\n",
    "# (2) administered general anaesthesia\n",
    "# (3) undergone non-cardiac surgery. \n",
    "#\n",
    "# For waveform data, the inclusion criteria were as follows:\n",
    "# (1) no missing monitoring for ABP, ECG, and EEG waveforms\n",
    "# (2) no cases containing false events or non-events due to poor signal quality\n",
    "#     (checked in second stage of data preprocessing)\n",
    "\n",
    "# Adult\n",
    "inclusion_1 = cases.loc[cases['age'] >= 18].index\n",
    "print(f'{len(cases)-len(inclusion_1)} cases excluded, {len(inclusion_1)} remaining due to age criteria')\n",
    "\n",
    "# General Anesthesia\n",
    "inclusion_2 = cases.loc[cases['ane_type'] == 'General'].index\n",
    "print(f'{len(cases)-len(inclusion_2)} cases excluded, {len(inclusion_2)} remaining due to anesthesia criteria')\n",
    "\n",
    "# Non-cardiac surgery\n",
    "inclusion_3 = cases.loc[\n",
    "    ~cases['opname'].str.contains(\"cardiac\", case=False)\n",
    "    & ~cases['opname'].str.contains(\"aneurysmal\", case=False)\n",
    "].index\n",
    "print(f'{len(cases)-len(inclusion_3)} cases excluded, {len(inclusion_3)} remaining due to non-cardiac surgery criteria')\n",
    "\n",
    "# ABP, ECG, EEG waveforms\n",
    "inclusion_4 = trks.loc[trks['tname'].isin(TRACK_NAMES)].index.value_counts()\n",
    "inclusion_4 = inclusion_4[inclusion_4 == len(TRACK_NAMES)].index\n",
    "print(f'{len(cases)-len(inclusion_4)} cases excluded, {len(inclusion_4)} remaining due to missing waveform data')\n",
    "\n",
    "# SQI filter\n",
    "# NOTE: this depends on a sqi_filter.csv generated by external processing\n",
    "inclusion_5 = pd.read_csv('sqi_filter.csv', header=None, names=['caseid','sqi']).set_index('caseid').index\n",
    "print(f'{len(cases)-len(inclusion_5)} cases excluded, {len(inclusion_5)} remaining due to SQI threshold not being met')\n",
    "\n",
    "cases_of_interest_idx = inclusion_1 \\\n",
    "    .intersection(inclusion_2) \\\n",
    "    .intersection(inclusion_3) \\\n",
    "    .intersection(inclusion_4) \\\n",
    "    .intersection(inclusion_5)\n",
    "\n",
    "cases_of_interest = cases.loc[cases_of_interest_idx]\n",
    "\n",
    "print()\n",
    "print(f'{cases_of_interest_idx.shape[0]} out of {cases.shape[0]} total cases remaining after exclusions applied')\n",
    "\n",
    "# Trim cases of interest to MAX_CASES\n",
    "if MAX_CASES:\n",
    "    cases_of_interest_idx = cases_of_interest_idx[:MAX_CASES]\n",
    "print(f'{cases_of_interest_idx.shape[0]} cases of interest selected')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cases_of_interest.head(n=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tracks of Interest\n",
    "\n",
    "These are the subset of tracks (waveforms) for the cases of interest identified above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A single case maps to one or more waveform tracks. Select only the tracks required for analysis.\n",
    "trks_of_interest = trks.loc[cases_of_interest_idx][trks.loc[cases_of_interest_idx]['tname'].isin(TRACK_NAMES)]\n",
    "trks_of_interest.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks_of_interest.head(n=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trks_of_interest_idx = trks_of_interest.set_index('tid').index\n",
    "trks_of_interest_idx.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build Tracks Cache for Local Processing\n",
    "\n",
    "Tracks data are large and therefore expensive to download every time used.\n",
    "By default, the `.vital` file format stores all tracks for each case internally. Since only select tracks per case are required, each `.vital` file can be further reduced by discarding the unused tracks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ensure the full vital file dataset is available for cases of interest.\n",
    "count_downloaded = 0\n",
    "count_present = 0\n",
    "\n",
    "#for i, idx in enumerate(cases.index):\n",
    "for i, idx in enumerate(cases_of_interest_idx):\n",
    "    full_path = f'{VITAL_ALL}/{idx:04d}.vital'\n",
    "    if not os.path.isfile(full_path):\n",
    "        print(f'Missing vital file: {full_path}')\n",
    "        # Download and save the file.\n",
    "        vf = vitaldb.VitalFile(idx)\n",
    "        vf.to_vital(full_path)\n",
    "        count_downloaded += 1\n",
    "    else:\n",
    "        count_present += 1\n",
    "\n",
    "print()\n",
    "print(f'Count of cases of interest:           {cases_of_interest_idx.shape[0]}')\n",
    "print(f'Count of vital files downloaded:      {count_downloaded}')\n",
    "print(f'Count of vital files already present: {count_present}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert vital files to \"mini\" versions including only the subset of tracks defined in TRACK_NAMES above.\n",
    "# Only perform conversion for the cases of interest.\n",
    "# NOTE: If this cell is interrupted, it can be restarted and will continue where it left off.\n",
    "count_minified = 0\n",
    "count_present = 0\n",
    "\n",
    "for i, idx in enumerate(cases_of_interest_idx):\n",
    "    full_path = f'{VITAL_ALL}/{idx:04d}.vital'\n",
    "    mini_path = f'{VITAL_MINI}/{idx:04d}_mini.vital'\n",
    "    if not os.path.isfile(mini_path):\n",
    "        print(f'Creating mini vital file: {idx}')\n",
    "        vf = vitaldb.VitalFile(full_path, TRACK_NAMES)\n",
    "        vf.to_vital(mini_path)\n",
    "        count_minified += 1\n",
    "    else:\n",
    "        count_present += 1\n",
    "\n",
    "print()\n",
    "print(f'Count of cases of interest:           {cases_of_interest_idx.shape[0]}')\n",
    "print(f'Count of vital files minified:        {count_minified}')\n",
    "print(f'Count of vital files already present: {count_present}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Filtering\n",
    "\n",
    "Preprocessing characteristics are different for each of the three signal categories:\n",
    " * ABP: no preprocessing, use as-is\n",
    " * ECG: apply a 1-40Hz bandpass filter, then perform Z-score normalization\n",
    " * EEG: apply a 0.5-50Hz bandpass filter\n",
    "\n",
    "`apply_bandpass_filter()` implements the bandpass filter using scipy.signal\n",
    "\n",
    "`apply_zscore_normalization()` implements the Z-score normalization using numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import butter, lfilter, spectrogram\n",
    "\n",
    "# define two methods for data preprocessing\n",
    "\n",
    "def apply_bandpass_filter(data, lowcut, highcut, fs, order=5):\n",
    "    b, a = butter(order, [lowcut, highcut], fs=fs, btype='band')\n",
    "    y = lfilter(b, a, np.nan_to_num(data))\n",
    "    return y\n",
    "\n",
    "def apply_zscore_normalization(signal):\n",
    "    mean = np.nanmean(signal)\n",
    "    std = np.nanstd(signal)\n",
    "    return (signal - mean) / std\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filtering Demonstration\n",
    "\n",
    "# temp experimental, code to be incorporated into overall preloader process\n",
    "# for now it's just dumping example plots of the before/after filtered signal data\n",
    "caseidx = 1\n",
    "file_path = f\"{VITAL_MINI}/{caseidx:04d}_mini.vital\"\n",
    "vf = vitaldb.VitalFile(file_path, TRACK_NAMES)\n",
    "\n",
    "originalAbp = None\n",
    "filteredAbp = None\n",
    "originalEcg = None\n",
    "filteredEcg = None\n",
    "originalEeg = None\n",
    "filteredEeg = None\n",
    "\n",
    "ABP_TRACK_NAME = \"SNUADC/ART\"\n",
    "ECG_TRACK_NAME = \"SNUADC/ECG_II\"\n",
    "EEG_TRACK_NAME = \"BIS/EEG1_WAV\"\n",
    "\n",
    "for i, (track_name, rate) in enumerate(zip(TRACK_NAMES, TRACK_SRATES)):\n",
    "    # Get samples for this track\n",
    "    track_samples = vf.get_track_samples(track_name, 1/rate)\n",
    "    #track_samples, _ = vf.get_samples(track_name, 1/rate)\n",
    "    print(f\"Track {track_name} @ {rate}Hz shape {len(track_samples)}\")\n",
    "\n",
    "    if track_name == ABP_TRACK_NAME:\n",
    "        # ABP waveforms are used without further pre-processing\n",
    "        originalAbp = track_samples\n",
    "        filteredAbp = track_samples\n",
    "    elif track_name == ECG_TRACK_NAME:\n",
    "        originalEcg = track_samples\n",
    "        # ECG waveforms are band-pass filtered between 1 and 40 Hz, and Z-score normalized\n",
    "        # first apply bandpass filter\n",
    "        filteredEcg = apply_bandpass_filter(track_samples, 1, 40, rate)\n",
    "        # then do z-score normalization\n",
    "        filteredEcg = apply_zscore_normalization(filteredEcg)\n",
    "    elif track_name == EEG_TRACK_NAME:\n",
    "        # EEG waveforms are band-pass filtered between 0.5 and 50 Hz\n",
    "        originalEeg = track_samples\n",
    "        filteredEeg = apply_bandpass_filter(track_samples, 0.5, 50, rate, 2)\n",
    "\n",
    "def plotSignal(data, title):\n",
    "    plt.figure(figsize=(20, 5))\n",
    "    plt.plot(data)\n",
    "    plt.title(title)\n",
    "    plt.show()\n",
    "\n",
    "plotSignal(originalAbp, \"Original ABP\")\n",
    "plotSignal(originalAbp, \"Unfiltered ABP\")\n",
    "plotSignal(originalEcg, \"Original ECG\")\n",
    "plotSignal(filteredEcg, \"Filtered ECG\")\n",
    "plotSignal(originalEeg, \"Original EEG\")\n",
    "plotSignal(filteredEeg, \"Filtered EEG\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_pickled_data(path):\n",
    "    if USE_DISK_CACHING and os.path.isfile(path):\n",
    "        with open(path, 'rb') as disk_cache_file:\n",
    "            result = np.load(disk_cache_file)\n",
    "            return result\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "def save_pickled_data(path, data):\n",
    "    if USE_DISK_CACHING:\n",
    "        with open(path, 'wb') as disk_cache_file:\n",
    "            np.save(data, disk_cache_file)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocess data tracks\n",
    "ABP_TRACK_NAME = \"SNUADC/ART\"\n",
    "ECG_TRACK_NAME = \"SNUADC/ECG_II\"\n",
    "EEG_TRACK_NAME = \"BIS/EEG1_WAV\"\n",
    "MINI_FILE_FOLDER = VITAL_MINI\n",
    "CACHE_FILE_FOLDER = VITAL_PREPROCESS_SCRATCH\n",
    "\n",
    "if RESET_CACHE:\n",
    "    TRACK_CACHE = None\n",
    "\n",
    "if TRACK_CACHE is None:\n",
    "    TRACK_CACHE = {}\n",
    "\n",
    "def get_track_data(case, print_when_file_loaded = False):\n",
    "    parsedFile = None\n",
    "    abp = None\n",
    "    eeg = None\n",
    "    ecg = None\n",
    "    for i, (track_name, rate) in enumerate(zip(TRACK_NAMES, TRACK_SRATES)):\n",
    "        # use integer case id and track name, delimited by pipe, as cache key\n",
    "        cache_label = f\"{case}|{track_name}\"\n",
    "        if cache_label not in TRACK_CACHE:\n",
    "            if parsedFile is None:\n",
    "                file_path = f\"{MINI_FILE_FOLDER}/{case:04d}_mini.vital\"\n",
    "                if print_when_file_loaded:\n",
    "                    print(f\"[{datetime.now()}] Loading vital file {file_path}\")\n",
    "                parsedFile = vitaldb.VitalFile(file_path, TRACK_NAMES)\n",
    "            dataset = np.array(vf.get_track_samples(track_name, 1/rate))\n",
    "            if track_name == ABP_TRACK_NAME:\n",
    "                # no filtering for ABP\n",
    "                abp = dataset\n",
    "                TRACK_CACHE[cache_label] = abp\n",
    "            elif track_name == ECG_TRACK_NAME:\n",
    "                ecg = dataset\n",
    "                # apply ECG filtering: first bandpass then do z-score normalization\n",
    "                ecg = apply_bandpass_filter(ecg, 1, 40, rate, 2)\n",
    "                ecg = apply_zscore_normalization(ecg)\n",
    "                TRACK_CACHE[cache_label] = ecg\n",
    "            elif track_name == EEG_TRACK_NAME:\n",
    "                eeg = dataset\n",
    "                # apply EEG filtering: bandpass only\n",
    "                eeg = apply_bandpass_filter(eeg, 0.5, 50, rate, 2)\n",
    "                TRACK_CACHE[cache_label] = eeg\n",
    "        else:\n",
    "            # cache hit, pull from cache\n",
    "            if track_name == ABP_TRACK_NAME:\n",
    "                abp = TRACK_CACHE[cache_label]\n",
    "            elif track_name == ECG_TRACK_NAME:\n",
    "                ecg = TRACK_CACHE[cache_label]\n",
    "            elif track_name == EEG_TRACK_NAME:\n",
    "                eeg = TRACK_CACHE[cache_label]\n",
    "\n",
    "    return (abp, ecg, eeg)\n",
    "\n",
    "# ABP waveforms are used without further pre-processing\n",
    "# ECG waveforms are band-pass filtered between 1 and 40 Hz, and Z-score normalized\n",
    "# EEG waveforms are band-pass filtered between 0.5 and 50 Hz\n",
    "if PRELOADING_CASES:\n",
    "    # determine disk cache file label\n",
    "    maxlabel = \"ALL\"\n",
    "    if MAX_CASES is not None:\n",
    "        maxlabel = str(MAX_CASES)\n",
    "    picklefile = f\"{CACHE_FILE_FOLDER}/{PREDICTION_WINDOW}_minutes_MAX{maxlabel}.trackcache\"\n",
    "\n",
    "    loaded = False\n",
    "    TRACK_CACHE = load_pickled_data(picklefile)\n",
    "    if TRACK_CACHE is None:\n",
    "        TRACK_CACHE = {}\n",
    "    if TRACK_CACHE is not None and (MAX_CASES is None or len(TRACK_CACHE) >= MAX_CASES):\n",
    "        loaded = True\n",
    "        print(f\"Loaded track cache from {picklefile}, {len(TRACK_CACHE)} records loaded\")\n",
    "\n",
    "    if not loaded:\n",
    "        print(f\"At beginning of process, the track cache has {len(TRACK_CACHE)} entries present.\")\n",
    "        for track in tqdm(cases_of_interest_idx):\n",
    "            # getting track data will cause a cache-check and fill when missing\n",
    "            # will also apply appropriate filtering per track\n",
    "            get_track_data(track, False)\n",
    "        \n",
    "        save_pickled_data(picklefile, TRACK_CACHE)\n",
    "        print(f\"Saved track cache to {picklefile}, {len(TRACK_CACHE)} records saved\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following method is adapted from the preprocessing block of reference [6] (https://github.com/vitaldb/examples/blob/master/hypotension_art.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate hypotensive events\n",
    "# Hypotensive events are defined as a 1-minute interval with sustained ABP of less than 65 mmHg\n",
    "# Note: Hypotensive events should be at least 20 minutes apart to minimize potential residual effects from previous events\n",
    "# Generate hypotension non-events\n",
    "# To sample non-events, 30-minute segments where the ABP was above 75 mmHG were selected, and then\n",
    "# three one-minute samples of each waveform were obtained from the middle of the segment\n",
    "# both occur in extract_segments\n",
    "\n",
    "def extract_segments(cases_of_interest_idx, min_before_event=3, debug=False):\n",
    "    # Sampling rate for ABP and ECG, Hz. These rates should be the same. Default = 500\n",
    "    ABP_ECG_SRATE_HZ = 500\n",
    "\n",
    "    # Sampling rate for EEG. Default = 128\n",
    "    EEG_SRATE_HZ = 128\n",
    "\n",
    "    # Length of feature segment, seconds.\n",
    "    FEATURE_LENGTH_SEC = 60\n",
    "    # Look ahead to predict hypotension, seconds.\n",
    "    MIDDLE_LENGTH_SEC  = 60 * min_before_event\n",
    "    # Length of label segment, seconds.\n",
    "    LABEL_LENGTH_SEC   = 60\n",
    "\n",
    "    # Length to move down the ABP track for starting a new analysis segment, seconds.\n",
    "    NEW_SEGMENT_OFFSET_SEC = 10\n",
    "\n",
    "    # Final dataset for training and testing the model.\n",
    "    # inputs with shape of (segments, timepoints)\n",
    "    samples = []\n",
    "    invalid_samples = []\n",
    "\n",
    "    # Process each case and extract segments. For each segment identify presence of an event in the label zone.\n",
    "    time_start = timer()\n",
    "    \n",
    "    count_cases = len(cases_of_interest_idx)\n",
    "\n",
    "    for case_count, caseid in tqdm(enumerate(cases_of_interest_idx), total=count_cases):\n",
    "        if debug:\n",
    "            print(f'Loading case: {caseid:04d}, ({case_count + 1} of {count_cases})')\n",
    "        \n",
    "        segment_key = []\n",
    "        segment_abp = []\n",
    "        segment_ecg = []\n",
    "        segment_eeg = []\n",
    "        segment_label = []\n",
    "        segment_validity = []\n",
    "        segment_caseid = caseid\n",
    "\n",
    "        # read the arterial waveform\n",
    "        (abp, ecg, eeg) = get_track_data(caseid)\n",
    "\n",
    "\n",
    "        # EEG - Different sample rate, process alone\n",
    "        if debug:\n",
    "            print(f'Length of {TRACK_NAMES[2]}:     {eeg.shape[0]}')\n",
    "\n",
    "        print_first_segment = True\n",
    "\n",
    "        last_sample_start_index = len(eeg) - EEG_SRATE_HZ * (FEATURE_LENGTH_SEC + MIDDLE_LENGTH_SEC + LABEL_LENGTH_SEC)\n",
    "        sample_index_offset = NEW_SEGMENT_OFFSET_SEC * EEG_SRATE_HZ\n",
    "\n",
    "        for i in range(0, last_sample_start_index, sample_index_offset):\n",
    "            segx_start = i\n",
    "            segx_end   = i + EEG_SRATE_HZ * FEATURE_LENGTH_SEC\n",
    "            segx = eeg[segx_start:segx_end]\n",
    "\n",
    "            if debug and print_first_segment:\n",
    "                print(f'  Feature Segment Length:   {segx.shape[0]} pts, {segx.shape[0] / EEG_SRATE_HZ} sec')\n",
    "                print_first_segment = False\n",
    "\n",
    "            # handle eeg, only care about extracting data from the same time interval used for abp\n",
    "            segment_eeg.append(eeg[segx_start:segx_end])\n",
    "\n",
    "        # ABP and ECG - Shared sample rate, process together    \n",
    "        if debug:\n",
    "            print(f'Length of {TRACK_NAMES[0]}:       {abp.shape[0]}')\n",
    "            print(f'Length of {TRACK_NAMES[1]}:    {ecg.shape[0]}')\n",
    "\n",
    "        segment_count = 0\n",
    "        segment_valid = 0\n",
    "        segment_event = 0\n",
    "        print_first_segment = True\n",
    "\n",
    "        last_sample_start_index = len(abp) - ABP_ECG_SRATE_HZ * (FEATURE_LENGTH_SEC + MIDDLE_LENGTH_SEC + LABEL_LENGTH_SEC)\n",
    "        sample_index_offset = NEW_SEGMENT_OFFSET_SEC * ABP_ECG_SRATE_HZ\n",
    "\n",
    "        for i in range(0, last_sample_start_index, sample_index_offset):\n",
    "            segment_count += 1\n",
    "\n",
    "            segx_start = i\n",
    "            segx_end   = i + ABP_ECG_SRATE_HZ * FEATURE_LENGTH_SEC\n",
    "            segx = abp[segx_start:segx_end]\n",
    "\n",
    "            segy_start = i + ABP_ECG_SRATE_HZ * (FEATURE_LENGTH_SEC + MIDDLE_LENGTH_SEC)\n",
    "            segy_end   = i + ABP_ECG_SRATE_HZ * (FEATURE_LENGTH_SEC + MIDDLE_LENGTH_SEC + LABEL_LENGTH_SEC)\n",
    "            segy = abp[segy_start:segy_end]\n",
    "\n",
    "            if debug and print_first_segment:\n",
    "                print(f'  Feature Segment Length:   {segx.shape[0]} pts, {segx.shape[0] / ABP_ECG_SRATE_HZ} sec')\n",
    "                print(f'  Middle Segment Length:    {segy_start - segx_end} pts, {(segy_start - segx_end) / ABP_ECG_SRATE_HZ} sec')\n",
    "                print(f'  Label Segment Length:     {segy.shape[0]} pts, {segy.shape[0] / ABP_ECG_SRATE_HZ} sec')\n",
    "                print_first_segment = False\n",
    "\n",
    "            # check the validity of this segment\n",
    "            valid = True\n",
    "            if np.isnan(segx).mean() > 0.1:\n",
    "                valid = False\n",
    "            elif np.isnan(segy).mean() > 0.1:\n",
    "                valid = False\n",
    "            elif (segx > 200).any():\n",
    "                valid = False\n",
    "            elif (segy > 200).any():\n",
    "                valid = False\n",
    "            elif (segx < 30).any():\n",
    "                valid = False\n",
    "            elif (segy < 30).any():\n",
    "                valid = False\n",
    "            elif np.max(segx) - np.min(segx) < 30:\n",
    "                valid = False\n",
    "            elif np.max(segy) - np.min(segy) < 30:\n",
    "                valid = False\n",
    "            elif (np.abs(np.diff(segx)) > 30).any():  # abrupt change -> noise\n",
    "                valid = False\n",
    "            elif (np.abs(np.diff(segy)) > 30).any():  # abrupt change -> noise\n",
    "                valid = False\n",
    "\n",
    "            # 2 sec moving avg\n",
    "            n = 2 * ABP_ECG_SRATE_HZ  \n",
    "            segy = np.nancumsum(segy, dtype=np.float32)\n",
    "            segy[n:] = segy[n:] - segy[:-n]\n",
    "            segy = segy[n - 1:] / n\n",
    "\n",
    "            # forward filling - do this per case to avoid massive resource utilization at the end.\n",
    "            segx = pd.DataFrame(segx).ffill(axis=0).bfill(axis=0)[0].values\n",
    "\n",
    "            # Identify IOH event as < 65mm HG\n",
    "            evt = np.nanmax(segy) < 65\n",
    "\n",
    "            segment_abp.append(segx)\n",
    "            segment_label.append(evt)\n",
    "            segment_validity.append(valid)\n",
    "\n",
    "            # handle ecg, only care about extracting the same segment used for abp.\n",
    "            # data is already time aligned and has same sample rate.\n",
    "            segment_ecg.append(ecg[segx_start:segx_end])\n",
    "            segment_key.append(f\"{caseid}|{segx_start}\")\n",
    "\n",
    "            \n",
    "\n",
    "            if valid:\n",
    "                segment_valid += 1\n",
    "                if evt:\n",
    "                    segment_event += 1\n",
    "\n",
    "        for i in range(0, len(segment_abp)):\n",
    "            if (segment_validity[i]):\n",
    "                samples.append((segment_abp[i], segment_ecg[i], segment_eeg[i], segment_label[i], segment_validity[i], segment_caseid, segment_key[i]))\n",
    "            else:\n",
    "                invalid_samples.append((segment_abp[i], segment_ecg[i], segment_eeg[i], segment_label[i], segment_validity[i], segment_caseid, segment_key[i]))\n",
    "\n",
    "        if debug:\n",
    "            print(f'Total Segments Evaluated:   {segment_count}')\n",
    "            segment_valid_percent = 0 if segment_count == 0 else 100 * segment_valid / segment_count \n",
    "            print(f'  Segments Valid:           {segment_valid}, {segment_valid_percent:.1f}%')\n",
    "            segment_event_percent = 0 if segment_valid == 0 else 100 * segment_event / segment_valid\n",
    "            print(f'  Segments with Event:      {segment_event}, {segment_event_percent:.1f}%')\n",
    "            time_delta = np.round(timer() - time_start, 3)\n",
    "            print(f'Total Processing Time:      {time_delta:.4f} sec')\n",
    "            print()\n",
    "\n",
    "\n",
    "    # total processing time\n",
    "    time_end = timer()\n",
    "    time_delta = np.round(time_end - time_start, 3)\n",
    "\n",
    "    # if debug:\n",
    "    #     print('OVERALL SUMMARY')\n",
    "    #     print(f'Total Processing Time:      {time_delta:.4f} sec')\n",
    "    #     print(f'Total Cases Processed:      {caseids.shape[0]}')\n",
    "    #     print(f'Total Segments Evaluated:   {x_abp.shape[0]}')\n",
    "\n",
    "    #     segment_valid_count = np.sum(valid_mask)\n",
    "    #     segment_valid_percent = 0 if x_abp.shape[0] == 0 else 100 * segment_valid_count / x_abp.shape[0] \n",
    "    #     print(f'  Segments Valid:           {segment_valid_count}, {segment_valid_percent:.1f}%')\n",
    "    #     segment_event_count = np.sum(y & valid_mask)\n",
    "    #     segment_event_percent = 0 if y.shape[0] == 0 else 100 * segment_event_count / y.shape[0]\n",
    "    #     print(f'  Segments with Event:      {segment_event_count}, {segment_event_percent:.1f}%')\n",
    "\n",
    "\n",
    "    #     print(f'Valid Samples Generated:    {(100 * np.mean(valid_mask)):.1f}%')\n",
    "    #     print()\n",
    "    #     print(f'Valid Mask Shape:           {valid_mask.shape}')\n",
    "    #     print(f'X_ABP Shape:                {x_abp.shape}')\n",
    "    #     print(f'X_ECG Shape:                {x_ecg.shape}')\n",
    "    #     print(f'X_EEG Shape:                {x_eeg.shape}')\n",
    "    #     print(f'Y Shape:                    {y.shape}')\n",
    "    #     print(f'CIPS Shape:                 {case_id_per_segment.shape}')\n",
    "    \n",
    "    return pd.DataFrame(samples, columns=['segment_abp', 'segment_ecg', 'segment_eeg', 'segment_label', 'segment_valid', 'caseidx', 'segment_key'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x_abp, x_ecg, x_eeg, y, valid_mask, case_id_per_segment = \\\n",
    "#     extract_segments(cases_of_interest_idx, min_before_event=3, debug=True)\n",
    "\n",
    "\n",
    "\n",
    "if PRELOADING_CASES:\n",
    "    # determine disk cache file label\n",
    "    maxlabel = \"ALL\"\n",
    "    if MAX_CASES is not None:\n",
    "        maxlabel = str(MAX_CASES)\n",
    "    picklefile = f\"{CACHE_FILE_FOLDER}/{PREDICTION_WINDOW}_minutes_MAX{maxlabel}.segmentcache\"\n",
    "\n",
    "    loaded = False\n",
    "    samples = load_pickled_data(picklefile)\n",
    "    if samples is not None:\n",
    "        loaded = True\n",
    "        print(f\"Loaded segment cache from {picklefile}, {len(samples)} samples loaded\")\n",
    "\n",
    "    if not loaded:\n",
    "        samples = \\\n",
    "            extract_segments(cases_of_interest_idx, min_before_event=PREDICTION_WINDOW, debug=False)\n",
    "        \n",
    "        save_pickled_data(picklefile, samples)\n",
    "        print(f\"Saved segment cache to {picklefile}, {len(TRACK_CACHE)} samples saved\")\n",
    "else:\n",
    "    samples = extract_segments(cases_of_interest_idx, min_before_event=PREDICTION_WINDOW, debug=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split data into training, validation, and test sets\n",
    "# Use 6:1:3 ratio and prevent samples from a single case from being split across different sets\n",
    "# Note: number of samples at each time point is not the same, because the first event can occur before the 3/5/10/15 minute mark\n",
    "\n",
    "# Set target sizes\n",
    "train_ratio = 0.6\n",
    "val_ratio = 0.1\n",
    "test_ratio = 1 - train_ratio - val_ratio # ensure ratios sum to 1\n",
    "\n",
    "# Split samples into train and other\n",
    "samples_train, samples_other = train_test_split(samples, test_size=(1 - train_ratio), random_state=RANDOM_SEED)\n",
    "# Split other into val and test\n",
    "samples_val, samples_test = train_test_split(samples_other, test_size=(test_ratio / (1 - train_ratio)), random_state=RANDOM_SEED)\n",
    "\n",
    "# Check how many samples are in each set\n",
    "print(f\"Train samples: {len(samples_train)}, ({len(samples_train) / len(samples):.2%})\")\n",
    "print(f\"Val samples: {len(samples_val)}, ({len(samples_val) / len(samples):.2%})\")\n",
    "print(f\"Test samples: {len(samples_test)}, ({len(samples_test) / len(samples):.2%})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create vitalDataset class\n",
    "class vitalDataset(Dataset):\n",
    "    def __init__(self, samples):\n",
    "        self.samples = samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # Get metadata for this event\n",
    "        segment = self.samples.iloc[idx]\n",
    "\n",
    "        abp = segment['segment_abp']\n",
    "        ecg = segment['segment_ecg']\n",
    "        eeg = segment['segment_eeg']\n",
    "        label = segment['segment_label']\n",
    "\n",
    "        return abp, ecg, eeg, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = samples.iloc[0]\n",
    "(fullAbp, fullEcg, fullEeg) = get_track_data(sample['caseidx'])\n",
    "\n",
    "plt.plot(fullAbp.ravel())\n",
    "plt.plot(fullEcg.ravel())\n",
    "plt.plot(fullEeg.ravel())\n",
    "plt.show()\n",
    "plt.plot(sample['segment_abp'])\n",
    "plt.plot(sample['segment_ecg'])\n",
    "plt.plot(sample['segment_eeg'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = vitalDataset(samples_train)\n",
    "val_dataset = vitalDataset(samples_val)\n",
    "test_dataset = vitalDataset(samples_test)\n",
    "\n",
    "BATCH_SIZE = 32\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##   Model\n",
    "\n",
    "The model implementation is based on the CNN architecture described in Jo Y-Y et al. (2022). It is designed to handle 1, 2, or 3 signal categories simultaneously, allowing for flexible model configurations based on different combinations of physiological signals:\n",
    " * ABP alone\n",
    " * EEG alone\n",
    " * ECG alone\n",
    " * ABP + EEG\n",
    " * ABP + ECG\n",
    " * EEG + ECG\n",
    " * ABP + EEG + ECG\n",
    "\n",
    "### Model Architecture\n",
    "\n",
    "The architecture, as depicted in Figure 2 from the original paper, utilizes a ResNet-based approach tailored for time-series data from different physiological signals. The model architecture is adapted to handle varying input signal frequencies, with specific hyperparameters for each signal type, particularly EEG, due to its distinct characteristics compared to ABP and ECG. A diagram of the model architecture is shown below:\n",
    "\n",
    "![Architecture of the hypotension risk prediction model using multiple waveforms](https://journals.plos.org/plosone/article/figure/image?download&size=large&id=10.1371/journal.pone.0272055.g002)\n",
    "\n",
    "Each input signal is processed through a sequence of 12 7-layer residual blocks, followed by a flattening process and a linear transformation to produce a 32-dimensional feature vector per signal type. These vectors are then concatenated (if multiple signals are used) and passed through two additional linear layers to produce a single output vector, representing the IOH index. A threshold is determined experimentally in order to minimize the differene between the sensitivity and specificity and is applied to this index to perform binary classification for predicting IOH events.\n",
    "\n",
    "The hyperparameters for the residual blocks are specified in Supplemental Table 1 from the original paper and vary for different signal type.\n",
    "\n",
    "A forward pass through the model passes through 85 layers before concatenation, followed by two more linear layers and finally a sigmoid activation layer to produce the prediction measure.\n",
    "\n",
    "### Residual Block Definition\n",
    "\n",
    "Each residual block consists of the following seven layers:\n",
    " \n",
    " * Batch normalization\n",
    " * ReLU\n",
    " * Dropout (0.5)\n",
    " * 1D convolution\n",
    " * Batch normalization\n",
    " * ReLU\n",
    " * 1D convolution\n",
    "\n",
    "Skip connections are included to aid in gradient flow during training, with optional 1D convolution in the skip connection to align dimensions.\n",
    "\n",
    "#### Residual Block Hyperparameters\n",
    "\n",
    "The hyperparameters are detailed in Supplemental Table 1 of the original paper. A screenshot of these hyperparameters is provided for reference below:\n",
    "\n",
    "![Supplemental Table 1 from original paper](<https://github.com/abarrie2/cs598-dlh-project/blob/main/img/table_1_hyperparameters.png?raw=true>)\n",
    "\n",
    "**Note**: Please be aware of a transcription error in the original paper's Supplemental Table 1 for the ECG+ABP configuration in Residual Blocks 11 and 12, where the output size should be 469 * 6 instead of the reported 496 * 6.\n",
    "\n",
    "### Training Objectives\n",
    "\n",
    "Our model uses binary cross entropy as the loss function and Adam as the optimizer, consistent with the original study. The learning rate is set at 0.0001, and training is configured to run for up to 100 epochs, with early stopping implemented if no improvement in loss is observed over five consecutive epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gBdVZoTvsSFV"
   },
   "outputs": [],
   "source": [
    "# First define the residual block which is reused 12x for each data track for each sample.\n",
    "# Second define the primary model.\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_features: int, out_features: int, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, size_down: bool = False) -> None:\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        \n",
    "        # calculate the appropriate padding required to ensure expected sequence lengths out of each residual block\n",
    "        padding = int((((stride-1)*in_features)-stride+kernel_size)/2)\n",
    "\n",
    "        self.size_down = size_down\n",
    "        self.bn1 = nn.BatchNorm1d(in_channels)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=False)\n",
    "        self.bn2 = nn.BatchNorm1d(out_channels)\n",
    "        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=False)\n",
    "        \n",
    "        self.residualConv = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=False)\n",
    "\n",
    "        # unclear where in sequence this should take place. Size down expressed in Supplemental table S1\n",
    "        if self.size_down:\n",
    "            pool_padding = (1 if (in_features % 2 > 0) else 0)\n",
    "            self.downsample = nn.MaxPool1d(kernel_size=2, stride=2, padding = pool_padding)\n",
    "        \n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        identity = x\n",
    "        \n",
    "        out = self.bn1(x)\n",
    "        out = self.relu(out)\n",
    "        out = self.dropout(out)\n",
    "        out = self.conv1(out)\n",
    "\n",
    "        if self.size_down:\n",
    "            out = self.downsample(out)\n",
    "\n",
    "        out = self.bn2(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "        \n",
    "        if out.shape != identity.shape:\n",
    "            # run the residual through a convolution when necessary\n",
    "            identity = self.residualConv(identity)\n",
    "            \n",
    "            outlen = np.prod(out.shape)\n",
    "            idlen = np.prod(identity.shape)\n",
    "            # downsample when required\n",
    "            if idlen > outlen:\n",
    "                identity = self.downsample(identity)\n",
    "            # match dimensions\n",
    "            identity = identity.reshape(out.shape)\n",
    "       \n",
    "        # add the residual       \n",
    "        out += identity\n",
    "\n",
    "        return  out\n",
    "\n",
    "class HypotensionCNN(nn.Module):\n",
    "    def __init__(self, useAbp: bool = True, useEeg: bool = False, useEcg: bool = False) -> None:\n",
    "        super(HypotensionCNN, self).__init__()\n",
    "\n",
    "        self.useAbp = useAbp\n",
    "        self.useEeg = useEeg\n",
    "        self.useEcg = useEcg\n",
    "\n",
    "        if useAbp:\n",
    "            self.abpBlock1 = ResidualBlock(30000, 15000, 1, 2, 15, 1, True)\n",
    "            self.abpBlock2 = ResidualBlock(15000, 15000, 2, 2, 15, 1, False)\n",
    "            self.abpBlock3 = ResidualBlock(15000, 7500, 2, 2, 15, 1, True)\n",
    "            self.abpBlock4 = ResidualBlock(7500, 7500, 2, 2, 15, 1, False)\n",
    "            self.abpBlock5 = ResidualBlock(7500, 3750, 2, 2, 15, 1, True)\n",
    "            self.abpBlock6 = ResidualBlock(3750, 3750, 2, 4, 15, 1, False)\n",
    "            self.abpBlock7 = ResidualBlock(3750, 1875, 4, 4, 7, 1, True)\n",
    "            self.abpBlock8 = ResidualBlock(1875, 1875, 4, 4, 7, 1, False)\n",
    "            self.abpBlock9 = ResidualBlock(1875, 938, 4, 4, 7, 1, True)\n",
    "            self.abpBlock10 = ResidualBlock(938, 938, 4, 4, 7, 1, False)\n",
    "            self.abpBlock11 = ResidualBlock(938, 469, 4, 6, 7, 1, True)\n",
    "            self.abpBlock12 = ResidualBlock(469, 469, 6, 6, 7, 1, False)\n",
    "            self.abpFc = nn.Linear(6*469, 32)\n",
    "        \n",
    "        if useEcg:\n",
    "            self.ecgBlock1 = ResidualBlock(30000, 15000, 1, 2, 15, 1, True)\n",
    "            self.ecgBlock2 = ResidualBlock(15000, 15000, 2, 2, 15, 1, False)\n",
    "            self.ecgBlock3 = ResidualBlock(15000, 7500, 2, 2, 15, 1, True)\n",
    "            self.ecgBlock4 = ResidualBlock(7500, 7500, 2, 2, 15, 1, False)\n",
    "            self.ecgBlock5 = ResidualBlock(7500, 3750, 2, 2, 15, 1, True)\n",
    "            self.ecgBlock6 = ResidualBlock(3750, 3750, 2, 4, 15, 1, False)\n",
    "            self.ecgBlock7 = ResidualBlock(3750, 1875, 4, 4, 7, 1, True)\n",
    "            self.ecgBlock8 = ResidualBlock(1875, 1875, 4, 4, 7, 1, False)\n",
    "            self.ecgBlock9 = ResidualBlock(1875, 938, 4, 4, 7, 1, True)\n",
    "            self.ecgBlock10 = ResidualBlock(938, 938, 4, 4, 7, 1, False)\n",
    "            self.ecgBlock11 = ResidualBlock(938, 469, 4, 6, 7, 1, True)\n",
    "            self.ecgBlock12 = ResidualBlock(469, 469, 6, 6, 7, 1, False)\n",
    "            self.ecgFc = nn.Linear(6 * 469, 32)\n",
    "        \n",
    "        if useEeg:\n",
    "            self.eegBlock1 = ResidualBlock(7680, 3840, 1, 2, 7, 1, True)\n",
    "            self.eegBlock2 = ResidualBlock(3840, 3840, 2, 2, 7, 1, False)\n",
    "            self.eegBlock3 = ResidualBlock(3840, 1920, 2, 2, 7, 1, True)\n",
    "            self.eegBlock4 = ResidualBlock(1920, 1920, 2, 2, 7, 1, False)\n",
    "            self.eegBlock5 = ResidualBlock(1920, 960, 2, 2, 7, 1, True)\n",
    "            self.eegBlock6 = ResidualBlock(960, 960, 2, 4, 7, 1, False)\n",
    "            self.eegBlock7 = ResidualBlock(960, 480, 4, 4, 3, 1, True)\n",
    "            self.eegBlock8 = ResidualBlock(480, 480, 4, 4, 3, 1, False)\n",
    "            self.eegBlock9 = ResidualBlock(480, 240, 4, 4, 3, 1, True)\n",
    "            self.eegBlock10 = ResidualBlock(240, 240, 4, 4, 3, 1, False)\n",
    "            self.eegBlock11 = ResidualBlock(240, 120, 4, 6, 3, 1, True)\n",
    "            self.eegBlock12 = ResidualBlock(120, 120, 6, 6, 3, 1, False)\n",
    "            self.eegFc = nn.Linear(6 * 120, 32)\n",
    "\n",
    "        concatSize = 0\n",
    "        if useAbp:\n",
    "            concatSize += 32\n",
    "        if useEeg:\n",
    "            concatSize += 32\n",
    "        if useEcg:\n",
    "            concatSize += 32\n",
    "\n",
    "        self.fullLinear1 = nn.Linear(concatSize, 16)\n",
    "        self.fullLinear2 = nn.Linear(16, 1)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, abp: torch.Tensor, eeg: torch.Tensor, ecg: torch.Tensor) -> torch.Tensor:\n",
    "\n",
    "        batchSize = len(abp)\n",
    "\n",
    "        # conditionally operate ABP, EEG, and ECG networks\n",
    "        if self.useAbp:\n",
    "            abp = self.abpBlock1(abp)\n",
    "            abp = self.abpBlock2(abp)\n",
    "            abp = self.abpBlock3(abp)\n",
    "            abp = self.abpBlock4(abp)\n",
    "            abp = self.abpBlock5(abp)\n",
    "            abp = self.abpBlock6(abp)\n",
    "            abp = self.abpBlock7(abp)\n",
    "            abp = self.abpBlock8(abp)\n",
    "            abp = self.abpBlock9(abp)\n",
    "            abp = self.abpBlock10(abp)\n",
    "            abp = self.abpBlock11(abp)\n",
    "            abp = self.abpBlock12(abp)\n",
    "            totalLen = np.prod(abp.shape)\n",
    "            abp = torch.reshape(abp, (batchSize, int(totalLen / batchSize)))\n",
    "            abp = self.abpFc(abp)\n",
    "\n",
    "        if self.useEeg:\n",
    "            eeg = self.eegBlock1(eeg)\n",
    "            eeg = self.eegBlock2(eeg)\n",
    "            eeg = self.eegBlock3(eeg)\n",
    "            eeg = self.eegBlock4(eeg)\n",
    "            eeg = self.eegBlock5(eeg)\n",
    "            eeg = self.eegBlock6(eeg)\n",
    "            eeg = self.eegBlock7(eeg)\n",
    "            eeg = self.eegBlock8(eeg)\n",
    "            eeg = self.eegBlock9(eeg)\n",
    "            eeg = self.eegBlock10(eeg)\n",
    "            eeg = self.eegBlock11(eeg)\n",
    "            eeg = self.eegBlock12(eeg)\n",
    "            totalLen = np.prod(eeg.shape)\n",
    "            eeg = torch.reshape(eeg, (batchSize, int(totalLen / batchSize)))\n",
    "            eeg = self.eegFc(eeg)\n",
    "        \n",
    "        if self.useEcg:\n",
    "            ecg = self.ecgBlock1(ecg)\n",
    "            ecg = self.ecgBlock2(ecg)\n",
    "            ecg = self.ecgBlock3(ecg)\n",
    "            ecg = self.ecgBlock4(ecg)\n",
    "            ecg = self.ecgBlock5(ecg)\n",
    "            ecg = self.ecgBlock6(ecg)\n",
    "            ecg = self.ecgBlock7(ecg)\n",
    "            ecg = self.ecgBlock8(ecg)\n",
    "            ecg = self.ecgBlock9(ecg)\n",
    "            ecg = self.ecgBlock10(ecg)\n",
    "            ecg = self.ecgBlock11(ecg)\n",
    "            ecg = self.ecgBlock12(ecg)\n",
    "            #ecg = torch.flatten(ecg)\n",
    "            totalLen = np.prod(ecg.shape)\n",
    "            ecg = torch.reshape(ecg, (batchSize, int(totalLen / batchSize)))\n",
    "            ecg = self.ecgFc(ecg)\n",
    "        \n",
    "        # concatenation\n",
    "        merged = None\n",
    "        if self.useAbp and self.useEeg and self.useEcg:\n",
    "            merged = torch.cat((abp, eeg, ecg), dim=1)\n",
    "        elif self.useAbp and self.useEeg:\n",
    "            merged = torch.cat((abp, eeg), dim=1)\n",
    "        elif self.useAbp and self.useEcg:\n",
    "            merged = torch.cat((abp, ecg), dim=1)\n",
    "        elif self.useEeg and self.useEcg:\n",
    "            merged = torch.cat((eeg, ecg), dim=1)\n",
    "        elif self.useAbp:\n",
    "            merged = abp\n",
    "        elif self.useEeg:\n",
    "            merged = eeg\n",
    "        elif self.useEcg:\n",
    "            merged = ecg\n",
    "\n",
    "        totalLen = np.prod(merged.shape)\n",
    "        merged = torch.reshape(merged, (batchSize, int(totalLen / batchSize)))\n",
    "        out = self.fullLinear1(merged)\n",
    "        out = self.fullLinear2(out)\n",
    "        out = self.sigmoid(out)\n",
    "\n",
    "        out = torch.nan_to_num(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "As discussed earlier, our model uses binary cross entropy as the loss function and Adam as the optimizer, consistent with the original study. The learning rate is set at 0.0001, and training is configured to run for up to 100 epochs, with early stopping implemented if no improvement in loss is observed over five consecutive epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experimentName = \"ABP_EEG_ECG_DEFAULT\"\n",
    "useAbp = True\n",
    "useEeg = True\n",
    "useEcg = False\n",
    "\n",
    "model = HypotensionCNN(useAbp, useEeg, useEcg)\n",
    "loss_func = nn.BCELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"mps\" if (torch.backends.mps.is_available() and torch.backends.mps.is_built()) else \"cpu\")\n",
    "print(f\"Using device: {device}\")\n",
    "model = model.to(device)\n",
    "\n",
    "def train_model_one_iter(model, loss_func, optimizer, train_loader):\n",
    "  model.train()\n",
    "  train_losses = []\n",
    "  for abp, ecg, eeg, label in tqdm(train_loader):\n",
    "    batch = len(abp)\n",
    "    \n",
    "    abp = torch.nan_to_num(abp.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "    ecg = torch.nan_to_num(ecg.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "    eeg = torch.nan_to_num(eeg.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "    label = label.type(torch.float).reshape(batch, 1).to(device)\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    mdl = model(abp, eeg, ecg)\n",
    "    loss = loss_func(torch.nan_to_num(mdl), label)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    train_losses.append(loss.cpu().data.numpy())\n",
    "  return np.mean(train_losses)\n",
    "\n",
    "def evaluate_model(model, loss_func, val_loader):\n",
    "  model.eval()\n",
    "  val_losses = []\n",
    "  for abp, ecg, eeg, label in tqdm(val_loader):\n",
    "    batch = len(abp)\n",
    "    \n",
    "    abp = torch.nan_to_num(abp.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "    ecg = torch.nan_to_num(ecg.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "    eeg = torch.nan_to_num(eeg.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "    label = label.type(torch.float).reshape(batch, 1).to(device)\n",
    "   \n",
    "    mdl = model(abp, eeg, ecg)\n",
    "    loss = loss_func(torch.nan_to_num(mdl), label)\n",
    "    val_losses.append(loss.cpu().data.numpy())\n",
    "  return np.mean(val_losses)\n",
    "\n",
    "\n",
    "# Training loop\n",
    "max_epochs = 6\n",
    "best_epoch = 0\n",
    "train_losses = []\n",
    "val_losses = []\n",
    "best_loss = float('inf')\n",
    "patience = 5\n",
    "no_improve_epochs = 0\n",
    "model_path = os.path.join(VITAL_MODELS, f\"{experimentName}.model\")\n",
    "\n",
    "for i in range(max_epochs):\n",
    "  # Train the model and get the training loss\n",
    "  train_loss = train_model_one_iter(model, loss_func, optimizer, train_loader)\n",
    "  train_losses.append(train_loss)\n",
    "  # Calculate validate loss\n",
    "  val_loss = evaluate_model(model, loss_func, val_loader)\n",
    "  val_losses.append(val_loss)\n",
    "  print(f\"[{datetime.now()}] Completed epoch {i} with training loss {train_loss:.8f}, validation loss {val_loss:.8f}\")\n",
    "\n",
    "  # Check if validation loss has improved\n",
    "  if val_loss < best_loss:\n",
    "    best_epoch = i\n",
    "    best_loss = val_loss\n",
    "    no_improve_epochs = 0\n",
    "    torch.save(model.state_dict(), model_path)\n",
    "    print(f\"Validation loss improved to {val_loss:.8f}. Model saved.\")\n",
    "  else:\n",
    "    no_improve_epochs += 1\n",
    "    print(f\"No improvement in validation loss. {no_improve_epochs} epochs without improvement.\")\n",
    "\n",
    "  # exit early if no improvement in loss over last 'patience' epochs\n",
    "  if no_improve_epochs >= patience:\n",
    "    print(\"Early stopping due to no improvement in validation loss.\")\n",
    "    break\n",
    "\n",
    "# Load best model from disk\n",
    "if os.path.exists(model_path):\n",
    "  model.load_state_dict(torch.load(model_path))\n",
    "  print(f\"Loaded best model from disk from epoch {best_epoch}.\")\n",
    "else:\n",
    "  print(\"No saved model found for f{experimentName}.\")\n",
    "\n",
    "model.train(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the training and validation losses after each epoch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create x-axis values for epochs\n",
    "epochs = range(0, len(train_losses))\n",
    "\n",
    "# Plot the training and validation losses\n",
    "plt.plot(epochs, train_losses, 'b', label='Training Loss')\n",
    "plt.plot(epochs, val_losses, 'r', label='Validation Loss')\n",
    "\n",
    "# Add a vertical bar at the best_epoch\n",
    "plt.axvline(x=best_epoch, color='g', linestyle='--', label='Best Epoch')\n",
    "\n",
    "# Shade everything to the right of the best_epoch a light red\n",
    "plt.axvspan(best_epoch, max(epochs), facecolor='r', alpha=0.1)\n",
    "\n",
    "# Add labels and title\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training and Validation Losses')\n",
    "\n",
    "# Add legend\n",
    "plt.legend(loc='upper right')\n",
    "\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_model(model, dataloader):\n",
    "    model.eval()\n",
    "    model = model.to(device)\n",
    "    total_loss = 0\n",
    "    all_predictions = []\n",
    "    all_labels = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for abp, ecg, eeg, label in tqdm(dataloader):\n",
    "            batch = len(abp)\n",
    "    \n",
    "            abp = torch.nan_to_num(abp.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "            ecg = torch.nan_to_num(ecg.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "            eeg = torch.nan_to_num(eeg.reshape(batch, 1, -1)).type(torch.FloatTensor).to(device)\n",
    "            label = label.type(torch.float).reshape(batch, 1).to(device)\n",
    "   \n",
    "            pred = model(abp, eeg, ecg)\n",
    "            loss = loss_func(pred, label)\n",
    "            total_loss += loss.item()\n",
    "\n",
    "            all_predictions.append(pred.detach().cpu().numpy())\n",
    "            all_labels.append(label.detach().cpu().numpy())\n",
    "\n",
    "    # Flatten the lists\n",
    "    all_predictions = np.concatenate(all_predictions).flatten()\n",
    "    all_labels = np.concatenate(all_labels).flatten()\n",
    "\n",
    "    # Calculate AUROC and AUPRC\n",
    "    auroc = roc_auc_score(all_labels, all_predictions)\n",
    "    precision, recall, _ = precision_recall_curve(all_labels, all_predictions)\n",
    "    auprc = auc(recall, precision)\n",
    "\n",
    "    # Determine the optimal threshold, which is argmin(abs(sensitivity - specificity)) per the paper\n",
    "    thresholds = np.linspace(0, 1, 101) # 0 to 1 in 0.01 steps\n",
    "    min_diff = float('inf')\n",
    "    optimal_sensitivity = None\n",
    "    optimal_specificity = None\n",
    "    optimal_threshold = None\n",
    "\n",
    "    for threshold in thresholds:\n",
    "        all_predictions_binary = (all_predictions > threshold).astype(int)\n",
    "\n",
    "        tn, fp, fn, tp = confusion_matrix(all_labels, all_predictions_binary).ravel()\n",
    "        sensitivity = tp / (tp + fn)\n",
    "        specificity = tn / (tn + fp)\n",
    "        diff = abs(sensitivity - specificity)\n",
    "\n",
    "        if diff < min_diff:\n",
    "            min_diff = diff\n",
    "            optimal_threshold = threshold\n",
    "            optimal_sensitivity = sensitivity\n",
    "            optimal_specificity = specificity\n",
    "\n",
    "    avg_loss = total_loss / len(dataloader)\n",
    "    return all_predictions, all_labels, avg_loss, auroc, auprc, optimal_sensitivity, optimal_specificity, optimal_threshold\n",
    "\n",
    "# validation loop\n",
    "valid_predictions, valid_labels, valid_loss, valid_auroc, valid_auprc, valid_sensitivity, valid_specificity, valid_threshold = eval_model(model, val_loader)\n",
    "\n",
    "# test loop\n",
    "test_predictions, test_labels, test_loss, test_auroc, test_auprc, test_sensitivity, test_specificity, test_threshold = eval_model(model, test_loader)\n",
    "\n",
    "print(f\"Validation predictions: {valid_predictions}\")\n",
    "print(f\"Validation labels: {valid_labels}\")\n",
    "print(f\"Validation loss: {valid_loss}\")\n",
    "print(f\"Validation AUROC: {valid_auroc}\")\n",
    "print(f\"Validation AUPRC: {valid_auprc}\")\n",
    "print(f\"Validation Sensitivity: {valid_sensitivity}\")\n",
    "print(f\"Validation Specificity: {valid_specificity}\")\n",
    "print(f\"Validation Threshold: {valid_threshold}\")\n",
    "print()\n",
    "print(f\"Test predictions: {test_predictions}\")\n",
    "print(f\"Test labels: {test_labels}\")\n",
    "print(f\"Test loss: {test_loss}\")\n",
    "print(f\"Test AUROC: {test_auroc}\")\n",
    "print(f\"Test AUPRC: {test_auprc}\")\n",
    "print(f\"Test Sensitivity: {test_sensitivity}\")\n",
    "print(f\"Test Specificity: {test_specificity}\")\n",
    "print(f\"Test Threshold: {test_threshold}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gX6bCcZNuxmz"
   },
   "source": [
    "# Results (Planned results for Draft submission)\n",
    "\n",
    "When we complete our experiments, we will build comparison tables that compare a set of measures for each experiment performed. The full set of experiments and measures are listed below.\n",
    "\n",
    "## Experiments\n",
    "\n",
    " * ABP only\n",
    " * ECG only\n",
    " * EEG only\n",
    " * ABP + ECG\n",
    " * ABP + EEG\n",
    " * ECG + EEG\n",
    " * ABP + ECG + EEG\n",
    "\n",
    "Note: each experiment will be repeated with the following time-to-IOH-event durations:\n",
    " * 3 minutes\n",
    " * 5 minutes\n",
    " * 10 minutes\n",
    " * 15 minutes\n",
    "\n",
    "Note: the above list of experiments will be performed if there is sufficient time and gpu capability to complete that before the submission deadline. Should we experience any constraints on this front, we will reduce our experimental coverage to the following 4 core experiments that are necessary to measure the hypotheses included at the head of this report:\n",
    " * ABP only @ 3 minutes\n",
    " * ABP + ECG @ 3 minutes\n",
    " * ABP + EEG @ 3 minutes\n",
    " * ABP + ECG + EEG @ 3 minutes\n",
    "\n",
    "For additional details please review the \"Planned Actions\" in the Discussion section of this report.\n",
    "\n",
    "## Measures\n",
    "\n",
    " * AUROC\n",
    " * AUPRC\n",
    " * Sensitivity\n",
    " * Specificity\n",
    " * Threshold\n",
    " * Loss Shrinkage\n",
    "\n",
    "[ TODO for final report - collect data for all measures listed above. ]\n",
    "\n",
    "[ TODO for final report - generate ROC and PRC plots for each experiment ]\n",
    "\n",
    "We are collecting a broad set of measures across each experiment in order to perform a comprehensive comparison of all measures listed across all comparable experiments executed in the original paper. However, our key experimental results will be focused on a subset of these results that address the main experiments defined at the beginning of this notebook.\n",
    "\n",
    "The key experimental result measures will be as follows:\n",
    "\n",
    "* For 3 minutes ahead of the predicted IOH event:\n",
    "  * compare AUROC and AUPRC for ABP only vs ABP+ECG\n",
    "  * compare AUROC and AUPRC for ABP only vs ABP+EEG\n",
    "  * compare AUROC and AUPRC for ABP only vs ABP+ECG+EEG\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8EAWAy_LwHlV"
   },
   "source": [
    "## Model comparison\n",
    "\n",
    "The following table is Table 3 from the original paper which presents the measured values for each signal combination across each of the four temporal predictive categories:\n",
    "\n",
    "![Area under the Receiver-operating Characteristic Curve, Area under the Precision-Recall Curve, Sensitivity, and Specificity of the model in predicting intraoperative hypotension](https://journals.plos.org/plosone/article/figure/image?download&size=large&id=10.1371/journal.pone.0272055.t003)\n",
    "\n",
    "We have not yet completed the execution of the experiments necessary to determine our reproduced model performance in order determine whether our results are accurately representing those of the original paper. These details are expected to be included in the final report.\n",
    "\n",
    "As of the draft submission, the reported evaluation measures of our model are too good to be true (all measures are 1.0). We suspect that there is data leakage in the dataset splitting process and will address this in time for the final report."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qH75TNU71eRH"
   },
   "source": [
    "# Discussion\n",
    "\n",
    "### Feasibility of reproduction\n",
    "Our assessment is that this paper will be reproducible. The outstanding risk is that each experiment can take up to 7 hours to run on hardware within the team (i.e., 7h to run ~70 epochs on a desktop with AMD Ryzen 7 3800X 8-core CPU w/ RTX 2070 SUPER GPU and 32GB RAM). There are a total of 28 experiments (7 different combinations of signal inputs, 4 different time horizons for each combination). Should our team find it not possible to complete the necessary experiments across all of the experiments represented in Table 3 of our selected paper, we will reduce the number of experiments to focus solely on the ones directly related to our hypotheses described in the beginning of this notebook (i.e., reduce the number of combinations of interest to 4: ABP alone, ABP+EEG, ABP+ECG, ABP+ECG+EEG). This will result in a new total of 16 experiments to run.\n",
    "\n",
    "### Planned ablations\n",
    "Our proposal included a collection of potential ablations to be investigated:\n",
    "\n",
    "* Remove ResNet skip connection\n",
    "* Reduce # of residual blocks from 12 to 6\n",
    "* Reduce # of residual blocks from 12 to 1\n",
    "* Eliminate dropout from residual block\n",
    "* Max pooling configuration\n",
    "  * smaller size/stride\n",
    "  * eliminate max pooling\n",
    "\n",
    "Given the amount of time required to conduct each experiment, our team intends to choose only a small number of ablations from this set. Further, we only intend to perform ablation analysis against the best performing signal combination and time horizon from the reproduction experiments. In order words, we intend to perform ablation analysis against the following training combinations, and only against the models trained with data measured 3 minutes prior to an IOH event:\n",
    "  * ABP alone\n",
    "  * ABP + ECG\n",
    "  * ABP + EEG\n",
    "  * ABP + ECG + EEG\n",
    "\n",
    "Time and GPU resource permitting, we will complete a broader range of experiments. For additional details, please see the section below titled \"Plans for next phase\".\n",
    "\n",
    "### Nature of reproduced results\n",
    "Our team intends to address the manner in which the experimental results align with the published results in the paper in the final submission of this report. The amount of time required to complete model training and result analysis during the preparation of the Draft notebook was not sufficient to complete a large number of experiments.\n",
    "\n",
    "### What was easy? What was difficult?\n",
    "The difficult aspect of the preparation of this draft involved the data preprocessing.\n",
    " * First, the source data is unlabelled, so our team was responsible for implementing analysis methods for identifying positive (IOH event occurred) and negative (IOH event did not occur) by running a lookahead analysis of our input training set.\n",
    " * Second, the volume of raw data is in excess of 90GB. A non-trivial amount of compute was required to minify the input data to only include the data tracks of interest to our experiments (i.e., ABP, ECG, and EEG tracks).\n",
    " * Third, our team found it difficult to trace back to the definition of the jSQI signal quality index referenced in the paper. Multiple references through multiple papers needed to be traversed to understand which variant of the quality index \n",
    "   * The only available source code related to the signal quality index as referenced by our paper in [5]. Source code was not directly linked from the paper, but the GitHub repository for the corresponding author for reference [5] did result in the identification of MATLAB source code for the signal quality index as described in the referenced paper. That code is available here: https://github.com/cliffordlab/PhysioNet-Cardiovascular-Signal-Toolbox/tree/master/Tools/BP_Tools\n",
    "   * Our team had insufficient time to port this signal quality index to Python for use in our investigation, or to setup a MATLAB environment in which to assess our source data using the above MATLAB functions, but we expect to complete this as part of our final report.\n",
    "\n",
    "### Suggestions to paper author\n",
    "The most notable suggestion would be to correct the hyperparameters published in Supplemental Table 1. Specifically, the output size for residual blocks 11 and 12 for the ECG and ABP data sets was 496x6. This is a typo, and should read 469x6. This typo became apparent when operating the size down operation within Residual Block 11 and recognizing the tensor dimensions were misaligned.\n",
    "\n",
    "Additionally, more explicit references to the signal quality index assessment tools should be added. Our team could not find a reference to the MATLAB source code as described in reference [3], and had to manually discover the GitHub profile for the lab of the corresponding author of reference [3] in order to find MATLAB source that corresponded to the metrics described therein.\n",
    "\n",
    "### Plans for next phase\n",
    "Our team plans to accomplish the following goals in service of preparing the Final Report:\n",
    " * Implement the jSQI filter to remove any training data with aberrent signal quality per the threshold defined in our original paper.\n",
    " * Execute the following experiments:\n",
    "   * Measure predictive quality of the model trained solely with ABP data at 3 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained with ABP+ECG data at 3 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained with ABP+EEG data at 3 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained with ABP+ECG+EEG data at 3 minutes prior to IOH events.\n",
    " * Gather our measures for these experiments and perform a comparison against the published results from our selected paper and determine whether or not we are succesfully reproducing the results outlined in the paper.\n",
    " * Ablation analysis:\n",
    "   * Execute the following ablation experiments:\n",
    "     * Repeat the four experiments described above while reducing the numnber of residual blocks in the model from 12 to 6.\n",
    " * Time- and/or GPU-resource permitting, we will complete the remaining 24 experiments as described in the paper:\n",
    "   * Measure predictive quality of the model trained solely with ABP data at 5, 10, and 15 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained with ABP+ECG data at 5, 10, and 15 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained with ABP+EEG data at 5, 10, and 15 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained with ABP+ECG+EEG data at 5, 10, and 15 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained solely with ECG data at 3, 5, 10, and 15 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained solely with EEG data at 3, 5, 10, and 15 minutes prior to IOH events.\n",
    "   * Measure predictive quality of the model trained with ECG+EEG data at 3, 5, 10, and 15 minutes prior to IOH events.\n",
    "   * Additional ablation experiments:\n",
    "     * For the four core experiments (ABP, ABP+ECG, ABP+EEG, ABP+ECG+EEG each trained on event data occurring 3 minutes prior to IOH events), perform the following ablations:\n",
    "       * Repeat experiment while eliminating dropout from every residual block\n",
    "       * Repeat experiment while removing the skip connection from every residual block\n",
    "       * Repeat the four experiments described above while reducing the numnber of residual blocks in the model from 12 to 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SHMI2chl9omn"
   },
   "source": [
    "# References\n",
    "\n",
    "1. Jo Y-Y, Jang J-H, Kwon J-m, Lee H-C, Jung C-W, Byun S, et al. “Predicting intraoperative hypotension using deep learning with waveforms of arterial blood pressure, electroencephalogram, and electrocardiogram: Retrospective study.” PLoS ONE, (2022) 17(8): e0272055 https://doi.org/10.1371/journal.pone.0272055\n",
    "2. Hatib, Feras, Zhongping J, Buddi S, Lee C, Settels J, Sibert K, Rhinehart J, Cannesson M “Machine-learning Algorithm to Predict Hypotension Based on High-fidelity Arterial Pressure Waveform Analysis” Anesthesiology (2018) 129:4 https://doi.org/10.1097/ALN.0000000000002300\n",
    "3. Bao, X., Kumar, S.S., Shah, N.J. et al. \"AcumenTM hypotension prediction index guidance for prevention and treatment of hypotension in noncardiac surgery: a prospective, single-arm, multicenter trial.\" Perioperative Medicine (2024) 13:13 https://doi.org/10.1186/s13741-024-00369-9\n",
    "4. Lee, HC., Park, Y., Yoon, S.B. et al. VitalDB, a high-fidelity multi-parameter vital signs database in surgical patients. Sci Data 9, 279 (2022). https://doi.org/10.1038/s41597-022-01411-5\n",
    "5. Li Q., Mark R.G. & Clifford G.D. \"Artificial arterial blood pressure artifact models and an evaluation of a robust blood pressure and heart rate estimator.\" BioMed Eng OnLine. (2009) 8:13. pmid:19586547 https://doi.org/10.1186/1475-925X-8-13\n",
    "6. Park H-J, \"VitalDB Python Example Notebooks\" GitHub Repository https://github.com/vitaldb/examples/blob/master/hypotension_art.ipynb"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}