{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extract irregular time series features from eICU dataset\n",
    "## We consider those features\n",
    "* hemoglobin\n",
    "* creatinine\n",
    "* sodium\n",
    "* BUN\n",
    "\n",
    "## Overview of the procedures\n",
    "1. Draw $N$ patients randomly from `pivoted_lab`, which is the preprocessed lab measurements provided by official eICU repository `https://github.com/MIT-LCP/eicu-code/blob/master/concepts/pivoted/pivoted-lab.sql`, do this for both train and test splits.\n",
    "2. Extract above time series features between minute $t_0$ and $T$, we use $t_0 = 0$ and $T = 1440$ in this study (first 24 hours of unit stay)\n",
    "3. Round time to nearest $r$ minutes, we use $r = 5$ (every 5 minutes), we expand the time dimension for every patient, so that the values over the entire first 24 hours is considered as part of the data, missingness is included. This will introduce a lot of missingness, but this is an inherent nature of sparse, irregular time series data like lab measurements.\n",
    "4. Add flags indicating whether a value is present or not, then do a cumulative sum on this flag (this is useful for Neural CDE since it indicates the presence of values overtime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers.loader import DataBaseLoader\n",
    "loader = DataBaseLoader(user='mt361', password=\"tian01050417\", dbname=\"eicu_gen\", schema=\"eicu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import timeit\n",
    "\n",
    "import joblib\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "TIMER = timeit.default_timer\n",
    "\n",
    "PATIENT_DF = loader[\"patient_id\"]\n",
    "\n",
    "EXTRACT_SQL = \"\"\"\n",
    "    SELECT\n",
    "        patientunitstayid,\n",
    "        chartoffset,\n",
    "        hemoglobin,\n",
    "        creatinine,\n",
    "        sodium,\n",
    "        bun\n",
    "    FROM\n",
    "        pivoted_lab\n",
    "    WHERE\n",
    "        chartoffset >= {start_time} AND chartoffset <= {end_time}\n",
    "        AND patientunitstayid = {id}\n",
    "\"\"\"\n",
    "\n",
    "def extract_patient(id, start_time, end_time):\n",
    "    extracted_df = loader.query(EXTRACT_SQL.format(start_time = start_time, end_time = end_time, id = id))\n",
    "    if len(extracted_df) == 0:\n",
    "        return None\n",
    "    else:\n",
    "        return extracted_df\n",
    "\n",
    "def drop_patientunitstayid(df):\n",
    "    if isinstance(df, type(None)):\n",
    "        return None\n",
    "    else:\n",
    "        return df.drop(columns = [\"patientunitstayid\"])\n",
    "\n",
    "def get_all_ids():\n",
    "    return loader.query(\"SELECT DISTINCT(patientunitstayid) FROM patient_id\").values.flatten()\n",
    "\n",
    "def get_all_time(start_time, end_time, every_minute=5):\n",
    "    return np.sort(np.unique((loader.query(f\"SELECT DISTINCT(chartoffset) FROM pivoted_lab WHERE chartoffset >= {start_time} AND chartoffset <= {end_time}\").values.flatten() // every_minute) * every_minute))\n",
    "\n",
    "def round_time(df, every_minute=5):\n",
    "    if isinstance(df, type(None)):\n",
    "        return None\n",
    "    else:\n",
    "        df[\"chartoffset\"] = np.round(df[\"chartoffset\"].values / every_minute) * every_minute\n",
    "        return df\n",
    "\n",
    "def expand_time_dim(df, all_times, fill_forward = False):\n",
    "    if isinstance(df, type(None)):\n",
    "        return None\n",
    "    else:\n",
    "        # Custom aggregation function\n",
    "        def custom_aggregate(series):\n",
    "            non_nans = series.dropna()\n",
    "            if not non_nans.empty:\n",
    "                return non_nans.median()\n",
    "            return np.nan\n",
    "    \n",
    "        # Group by chartoffset and aggregate\n",
    "        df = df.groupby('chartoffset').agg(custom_aggregate).reset_index()\n",
    "\n",
    "        # Reindex with all_times\n",
    "        df = df.set_index(\"chartoffset\")\n",
    "        df = df.reindex(all_times, method=\"ffill\" if fill_forward else None)\n",
    "        df = df.reset_index()\n",
    "        df = df.rename(columns={\"index\": \"chartoffset\"})\n",
    "        return df\n",
    "    \n",
    "def add_nonnan_cum_flag(df):\n",
    "    if isinstance(df, type(None)):\n",
    "        return None\n",
    "    else:\n",
    "        for col in df.columns:\n",
    "            if col != \"chartoffset\":\n",
    "                df[f\"{col}_nonnan\"] = np.cumsum(~np.isnan(df[col].values))\n",
    "        return df\n",
    "\n",
    "def reverse_to_nonnan_flag(df):\n",
    "    if isinstance(df, type(None)):\n",
    "        return None\n",
    "    else:\n",
    "        arr = []\n",
    "        for col in df.columns:\n",
    "            if \"nonnan\" in col:\n",
    "                flag =  np.concatenate([df[col].to_numpy()[:1], df[col].to_numpy()[1:] - df[col].to_numpy()[:-1]])\n",
    "                arr.append(flag)\n",
    "        arr = np.stack(arr, axis=1)\n",
    "        return arr\n",
    "\n",
    "def add_mortality_flag(df, id):\n",
    "    mortality = PATIENT_DF[PATIENT_DF[\"patientunitstayid\"] == id][\"hospital_expire_flag\"].values[0]\n",
    "    dummy = np.full((df.shape[0], ), mortality)\n",
    "    df[\"mortality\"] = dummy                             # mortality flag is at the last column\n",
    "    return df\n",
    "    \n",
    "def random_selection(df, num):\n",
    "    return df.sample(n=num, replace=False, random_state=2023)\n",
    "\n",
    "def collect_all_patients(all_ids, start_time, end_time, every_minute=5, prefix=\"\"):\n",
    "    all_times = get_all_time(start_time, end_time, every_minute=5)          # get all rounded times\n",
    "    all_patients = []\n",
    "    \n",
    "    with tqdm(total=len(all_ids), desc=f\"Collecting {len(all_ids)} patients...\") as pbar:\n",
    "        \n",
    "        last_time = TIMER()\n",
    "        for id in all_ids:\n",
    "            patient_df = drop_patientunitstayid(extract_patient(id=id, start_time=start_time, end_time=end_time))       # extract patient and drop id\n",
    "            if patient_df is None:\n",
    "                continue            # skip if does not exist\n",
    "            patient_df = round_time(patient_df, every_minute=every_minute)              # round time of charttime to nearest 5 minutes\n",
    "            patient_df = expand_time_dim(patient_df, all_times)                         # expand time dimension to include all times\n",
    "            patient_df = add_nonnan_cum_flag(patient_df)                                # add nonnan cumulative flag to indicate missingness\n",
    "            patient_df = add_mortality_flag(patient_df, id)                             # add mortality flag\n",
    "            \n",
    "            nonnan_flag = reverse_to_nonnan_flag(patient_df)                            # reverse process, this is simply for checking, not used as part of the training data\n",
    "            assert np.all(nonnan_flag == (~patient_df[[\"hemoglobin\", \"creatinine\", \"sodium\", \"bun\"]].isna()).to_numpy()), \"FAILED!\"\n",
    "            \n",
    "            all_patients.append(patient_df.to_numpy())\n",
    "            pbar.update(1)\n",
    "            \n",
    "            if TIMER() - last_time >= 120:\n",
    "                pbar.set_description(f\"Collected {np.asarray(all_patients).shape}... (Saving)\")\n",
    "                last_time = TIMER()\n",
    "                joblib.dump(all_patients, f\"data/eicu-extract/{prefix}irregular_all_patients_{start_time}_{end_time}_{every_minute}.joblib\")\n",
    "            \n",
    "    joblib.dump(all_patients, f\"data/eicu-extract/{prefix}irregular_all_patients_{start_time}_{end_time}_{every_minute}.joblib\")\n",
    "    print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Seed set to 2023\n",
      "\n",
      "Mortality rate in train set: 0.088, mortality rate in test set: 0.088\n",
      "Number of patients in train set: 40000, number of patients in test set: 10000\n"
     ]
    }
   ],
   "source": [
    "from helpers.utils import seed_everything\n",
    "from sklearn.model_selection import train_test_split\n",
    "seed_everything(2023)\n",
    "sub_patient_df = random_selection(PATIENT_DF, 50000)\n",
    "X, y = sub_patient_df.drop('hospital_expire_flag', axis=1), sub_patient_df['hospital_expire_flag']\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=2023, stratify=y, test_size=0.2)\n",
    "print(f\"Mortality rate in train set: {y_train.mean():.3f}, mortality rate in test set: {y_test.mean():.3f}\\nNumber of patients in train set: {len(y_train)}, number of patients in test set: {len(y_test)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ids = X_train['patientunitstayid'].values\n",
    "test_ids = X_test['patientunitstayid'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collected (36778, 289, 10)... (Saving):  94%|█████████▍| 37527/40000 [1:09:37<04:35,  8.98it/s]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done!\n"
     ]
    }
   ],
   "source": [
    "collect_all_patients(\n",
    "    all_ids = train_ids, start_time = 0, end_time = 1440, \n",
    "    every_minute = 5, prefix=\"TRAIN_\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collected (8963, 289, 10)... (Saving):  94%|█████████▍| 9387/10000 [16:46<01:05,  9.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done!\n"
     ]
    }
   ],
   "source": [
    "collect_all_patients(\n",
    "    all_ids = test_ids, start_time = 0, end_time = 1440, \n",
    "    every_minute = 5, prefix=\"TEST_\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-process into tensor and some checking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(8963, 289, 10)\n",
      "(36778, 289, 10)\n"
     ]
    }
   ],
   "source": [
    "import joblib\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "test = np.asarray(joblib.load(\"data/eicu-extract/TEST_irregular_all_patients_0_1440_5.joblib\"))\n",
    "train = np.asarray(joblib.load(\"data/eicu-extract/TRAIN_irregular_all_patients_0_1440_5.joblib\"))\n",
    "print(test.shape)\n",
    "print(train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8963, 10, 289])\n",
      "torch.Size([36778, 10, 289])\n"
     ]
    }
   ],
   "source": [
    "test = torch.tensor(test).permute(0, 2, 1)\n",
    "train = torch.tensor(train).permute(0, 2, 1)\n",
    "print(test.shape)\n",
    "print(train.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "check class imbalance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Imbalance for test: 8.368%\n",
      "Imbalance for train: 8.554%\n"
     ]
    }
   ],
   "source": [
    "def imbalance(data_tensor):\n",
    "    data_tensor = data_tensor[:, 9, :]          # last channel is label\n",
    "    return data_tensor.sum() / data_tensor.numel()\n",
    "\n",
    "print(f\"Imbalance for test: {imbalance(test)*100:.3f}%\")\n",
    "print(f\"Imbalance for train: {imbalance(train)*100:.3f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(test, \"data/eicu-extract/TEST_irregular_all_patients_0_1440_5.pt\")\n",
    "torch.save(train, \"data/eicu-extract/TRAIN_irregular_all_patients_0_1440_5.pt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gen",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
