{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Preprocessing for MIMIC III"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from helpers.loader import DataBaseLoader\n",
    "\n",
    "loader = DataBaseLoader(user=\"mt361\", password=\"tian01050417\", dbname=\"mimic\", schema=\"mimiciii\")\n",
    "SQL_SCRIPT = \"\"\"\n",
    "with first_table as (\n",
    "SELECT\n",
    "    ce.subject_id\n",
    "    , ce.icustay_id\n",
    "    , extract(epoch from (ce.charttime - FIRST_VALUE(ce.charttime) OVER(PARTITION BY ce.subject_id, ce.icustay_id ORDER BY ce.charttime)))/3600 AS time\n",
    "\t, ce.charttime\n",
    "    , AVG(CASE WHEN itemid IN (220045)\n",
    "            AND valuenum > 0\n",
    "            AND valuenum < 300\n",
    "            THEN valuenum END\n",
    "    ) AS heart_rate\n",
    "    , AVG(CASE WHEN itemid IN (220179, 220050, 225309)\n",
    "            AND valuenum > 0\n",
    "            AND valuenum < 400\n",
    "            THEN valuenum END\n",
    "    ) AS sbp\n",
    "    , AVG(CASE WHEN itemid IN (220180, 220051, 225310)\n",
    "                AND valuenum > 0\n",
    "                AND valuenum < 300\n",
    "                THEN valuenum END\n",
    "    ) AS dbp\n",
    "    , AVG(CASE WHEN itemid IN (220052, 220181, 225312)\n",
    "                AND valuenum > 0\n",
    "                AND valuenum < 300\n",
    "                THEN valuenum END\n",
    "    ) AS mbp\n",
    "    , AVG(CASE WHEN itemid = 220179\n",
    "                AND valuenum > 0\n",
    "                AND valuenum < 400\n",
    "                THEN valuenum END\n",
    "    ) AS sbp_ni\n",
    "    , AVG(CASE WHEN itemid = 220180\n",
    "                AND valuenum > 0\n",
    "                AND valuenum < 300\n",
    "                THEN valuenum END\n",
    "    ) AS dbp_ni\n",
    "    , AVG(CASE WHEN itemid = 220181\n",
    "                AND valuenum > 0\n",
    "                AND valuenum < 300\n",
    "                THEN valuenum END\n",
    "    ) AS mbp_ni\n",
    "    , AVG(CASE WHEN itemid IN (220210, 224690)\n",
    "                AND valuenum > 0\n",
    "                AND valuenum < 70\n",
    "                THEN valuenum END\n",
    "    ) AS resp_rate\n",
    "    , ROUND(CAST(\n",
    "            AVG(CASE\n",
    "                -- converted to degC in valuenum call\n",
    "                WHEN itemid IN (223761)\n",
    "                    AND valuenum > 70\n",
    "                    AND valuenum < 120\n",
    "                    THEN (valuenum - 32) / 1.8\n",
    "                -- already in degC, no conversion necessary\n",
    "                WHEN itemid IN (223762)\n",
    "                    AND valuenum > 10\n",
    "                    AND valuenum < 50\n",
    "                    THEN valuenum END)\n",
    "            AS NUMERIC), 2) AS temperature\n",
    "    , MAX(CASE WHEN itemid = 224642 THEN value END\n",
    "    ) AS temperature_site\n",
    "    , AVG(CASE WHEN itemid IN (220277)\n",
    "                AND valuenum > 0\n",
    "                AND valuenum <= 100\n",
    "                THEN valuenum END\n",
    "    ) AS spo2\n",
    "    , AVG(CASE WHEN itemid IN (225664, 220621, 226537)\n",
    "                AND valuenum > 0\n",
    "                THEN valuenum END\n",
    "    ) AS glucose\n",
    "FROM chartevents ce\n",
    "WHERE ce.icustay_id IS NOT NULL\n",
    "    AND ce.itemid IN\n",
    "    (\n",
    "        220045 -- Heart Rate\n",
    "        , 225309 -- ART BP Systolic\n",
    "        , 225310 -- ART BP Diastolic\n",
    "        , 225312 -- ART BP Mean\n",
    "        , 220050 -- Arterial Blood Pressure systolic\n",
    "        , 220051 -- Arterial Blood Pressure diastolic\n",
    "        , 220052 -- Arterial Blood Pressure mean\n",
    "        , 220179 -- Non Invasive Blood Pressure systolic\n",
    "        , 220180 -- Non Invasive Blood Pressure diastolic\n",
    "        , 220181 -- Non Invasive Blood Pressure mean\n",
    "        , 220210 -- Respiratory Rate\n",
    "        , 224690 -- Respiratory Rate (Total)\n",
    "        , 220277 -- SPO2, peripheral\n",
    "        -- GLUCOSE, both lab and fingerstick\n",
    "        , 225664 -- Glucose finger stick\n",
    "        , 220621 -- Glucose (serum)\n",
    "        , 226537 -- Glucose (whole blood)\n",
    "        -- TEMPERATURE\n",
    "        -- 226329 -- Blood Temperature CCO (C)\n",
    "        , 223762 -- \"Temperature Celsius\"\n",
    "        , 223761  -- \"Temperature Fahrenheit\"\n",
    "        , 224642 -- Temperature Site\n",
    "    )\n",
    "GROUP BY ce.subject_id, ce.icustay_id, ce.charttime\n",
    "),\n",
    "time_diffs AS (\n",
    "    SELECT \n",
    "        *,\n",
    "        ROW_NUMBER() OVER(PARTITION BY subject_id, icustay_id, ROUND(time) ORDER BY ABS(time - ROUND(time))) AS rn\t\t-- find the measurement to the nearest hour\n",
    "    FROM first_table\n",
    "),\n",
    "hours AS (\n",
    "    SELECT generate_series(0,24) AS hr\n",
    "),\n",
    "all_hours AS (\n",
    "    SELECT distinct icustay_id, hr\n",
    "    FROM time_diffs, hours\n",
    ")\n",
    "SELECT \n",
    "    a.icustay_id,\n",
    "    a.hr as time,\n",
    "    t.heart_rate,\n",
    "    t.sbp,\n",
    "    t.dbp,\n",
    "    t.mbp,\n",
    "    t.resp_rate as resprate,\n",
    "    t.temperature as temp,\n",
    "    t.spo2\n",
    "FROM all_hours a\n",
    "LEFT JOIN time_diffs t ON a.icustay_id = t.icustay_id AND a.hr = ROUND(t.time)\n",
    "WHERE t.rn = 1 OR t.rn IS NULL\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>icustay_id</th>\n",
       "      <th>time</th>\n",
       "      <th>heart_rate</th>\n",
       "      <th>sbp</th>\n",
       "      <th>dbp</th>\n",
       "      <th>mbp</th>\n",
       "      <th>resprate</th>\n",
       "      <th>temp</th>\n",
       "      <th>spo2</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>200001</td>\n",
       "      <td>0</td>\n",
       "      <td>115.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>200001</td>\n",
       "      <td>1</td>\n",
       "      <td>113.0</td>\n",
       "      <td>110.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>76.0</td>\n",
       "      <td>20.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>97.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>200001</td>\n",
       "      <td>2</td>\n",
       "      <td>108.0</td>\n",
       "      <td>113.0</td>\n",
       "      <td>68.0</td>\n",
       "      <td>79.0</td>\n",
       "      <td>18.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>98.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>200001</td>\n",
       "      <td>3</td>\n",
       "      <td>110.0</td>\n",
       "      <td>116.0</td>\n",
       "      <td>68.0</td>\n",
       "      <td>79.0</td>\n",
       "      <td>27.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>98.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>200001</td>\n",
       "      <td>4</td>\n",
       "      <td>102.0</td>\n",
       "      <td>102.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>71.0</td>\n",
       "      <td>21.0</td>\n",
       "      <td>37.67</td>\n",
       "      <td>96.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   icustay_id  time  heart_rate    sbp   dbp   mbp  resprate   temp  spo2\n",
       "0      200001     0       115.0    NaN   NaN   NaN       NaN    NaN   NaN\n",
       "1      200001     1       113.0  110.0  65.0  76.0      20.0    NaN  97.0\n",
       "2      200001     2       108.0  113.0  68.0  79.0      18.0    NaN  98.0\n",
       "3      200001     3       110.0  116.0  68.0  79.0      27.0    NaN  98.0\n",
       "4      200001     4       102.0  102.0  61.0  71.0      21.0  37.67  96.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vitals_table = loader.query(SQL_SCRIPT)\n",
    "vitals_table.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>row_id</th>\n",
       "      <th>subject_id</th>\n",
       "      <th>hadm_id</th>\n",
       "      <th>admittime</th>\n",
       "      <th>dischtime</th>\n",
       "      <th>deathtime</th>\n",
       "      <th>admission_type</th>\n",
       "      <th>admission_location</th>\n",
       "      <th>discharge_location</th>\n",
       "      <th>insurance</th>\n",
       "      <th>language</th>\n",
       "      <th>religion</th>\n",
       "      <th>marital_status</th>\n",
       "      <th>ethnicity</th>\n",
       "      <th>edregtime</th>\n",
       "      <th>edouttime</th>\n",
       "      <th>diagnosis</th>\n",
       "      <th>hospital_expire_flag</th>\n",
       "      <th>has_chartevents_data</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>21</td>\n",
       "      <td>22</td>\n",
       "      <td>165315</td>\n",
       "      <td>2196-04-09 12:26:00</td>\n",
       "      <td>2196-04-10 15:54:00</td>\n",
       "      <td>NaT</td>\n",
       "      <td>EMERGENCY</td>\n",
       "      <td>EMERGENCY ROOM ADMIT</td>\n",
       "      <td>DISC-TRAN CANCER/CHLDRN H</td>\n",
       "      <td>Private</td>\n",
       "      <td>None</td>\n",
       "      <td>UNOBTAINABLE</td>\n",
       "      <td>MARRIED</td>\n",
       "      <td>WHITE</td>\n",
       "      <td>2196-04-09 10:06:00</td>\n",
       "      <td>2196-04-09 13:24:00</td>\n",
       "      <td>BENZODIAZEPINE OVERDOSE</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>22</td>\n",
       "      <td>23</td>\n",
       "      <td>152223</td>\n",
       "      <td>2153-09-03 07:15:00</td>\n",
       "      <td>2153-09-08 19:10:00</td>\n",
       "      <td>NaT</td>\n",
       "      <td>ELECTIVE</td>\n",
       "      <td>PHYS REFERRAL/NORMAL DELI</td>\n",
       "      <td>HOME HEALTH CARE</td>\n",
       "      <td>Medicare</td>\n",
       "      <td>None</td>\n",
       "      <td>CATHOLIC</td>\n",
       "      <td>MARRIED</td>\n",
       "      <td>WHITE</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaT</td>\n",
       "      <td>CORONARY ARTERY DISEASE\\CORONARY ARTERY BYPASS...</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>23</td>\n",
       "      <td>23</td>\n",
       "      <td>124321</td>\n",
       "      <td>2157-10-18 19:34:00</td>\n",
       "      <td>2157-10-25 14:00:00</td>\n",
       "      <td>NaT</td>\n",
       "      <td>EMERGENCY</td>\n",
       "      <td>TRANSFER FROM HOSP/EXTRAM</td>\n",
       "      <td>HOME HEALTH CARE</td>\n",
       "      <td>Medicare</td>\n",
       "      <td>ENGL</td>\n",
       "      <td>CATHOLIC</td>\n",
       "      <td>MARRIED</td>\n",
       "      <td>WHITE</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaT</td>\n",
       "      <td>BRAIN MASS</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>24</td>\n",
       "      <td>24</td>\n",
       "      <td>161859</td>\n",
       "      <td>2139-06-06 16:14:00</td>\n",
       "      <td>2139-06-09 12:48:00</td>\n",
       "      <td>NaT</td>\n",
       "      <td>EMERGENCY</td>\n",
       "      <td>TRANSFER FROM HOSP/EXTRAM</td>\n",
       "      <td>HOME</td>\n",
       "      <td>Private</td>\n",
       "      <td>None</td>\n",
       "      <td>PROTESTANT QUAKER</td>\n",
       "      <td>SINGLE</td>\n",
       "      <td>WHITE</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaT</td>\n",
       "      <td>INTERIOR MYOCARDIAL INFARCTION</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>25</td>\n",
       "      <td>25</td>\n",
       "      <td>129635</td>\n",
       "      <td>2160-11-02 02:06:00</td>\n",
       "      <td>2160-11-05 14:55:00</td>\n",
       "      <td>NaT</td>\n",
       "      <td>EMERGENCY</td>\n",
       "      <td>EMERGENCY ROOM ADMIT</td>\n",
       "      <td>HOME</td>\n",
       "      <td>Private</td>\n",
       "      <td>None</td>\n",
       "      <td>UNOBTAINABLE</td>\n",
       "      <td>MARRIED</td>\n",
       "      <td>WHITE</td>\n",
       "      <td>2160-11-02 01:01:00</td>\n",
       "      <td>2160-11-02 04:27:00</td>\n",
       "      <td>ACUTE CORONARY SYNDROME</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   row_id  subject_id  hadm_id           admittime           dischtime  \\\n",
       "0      21          22   165315 2196-04-09 12:26:00 2196-04-10 15:54:00   \n",
       "1      22          23   152223 2153-09-03 07:15:00 2153-09-08 19:10:00   \n",
       "2      23          23   124321 2157-10-18 19:34:00 2157-10-25 14:00:00   \n",
       "3      24          24   161859 2139-06-06 16:14:00 2139-06-09 12:48:00   \n",
       "4      25          25   129635 2160-11-02 02:06:00 2160-11-05 14:55:00   \n",
       "\n",
       "  deathtime admission_type         admission_location  \\\n",
       "0       NaT      EMERGENCY       EMERGENCY ROOM ADMIT   \n",
       "1       NaT       ELECTIVE  PHYS REFERRAL/NORMAL DELI   \n",
       "2       NaT      EMERGENCY  TRANSFER FROM HOSP/EXTRAM   \n",
       "3       NaT      EMERGENCY  TRANSFER FROM HOSP/EXTRAM   \n",
       "4       NaT      EMERGENCY       EMERGENCY ROOM ADMIT   \n",
       "\n",
       "          discharge_location insurance language           religion  \\\n",
       "0  DISC-TRAN CANCER/CHLDRN H   Private     None       UNOBTAINABLE   \n",
       "1           HOME HEALTH CARE  Medicare     None           CATHOLIC   \n",
       "2           HOME HEALTH CARE  Medicare     ENGL           CATHOLIC   \n",
       "3                       HOME   Private     None  PROTESTANT QUAKER   \n",
       "4                       HOME   Private     None       UNOBTAINABLE   \n",
       "\n",
       "  marital_status ethnicity           edregtime           edouttime  \\\n",
       "0        MARRIED     WHITE 2196-04-09 10:06:00 2196-04-09 13:24:00   \n",
       "1        MARRIED     WHITE                 NaT                 NaT   \n",
       "2        MARRIED     WHITE                 NaT                 NaT   \n",
       "3         SINGLE     WHITE                 NaT                 NaT   \n",
       "4        MARRIED     WHITE 2160-11-02 01:01:00 2160-11-02 04:27:00   \n",
       "\n",
       "                                           diagnosis  hospital_expire_flag  \\\n",
       "0                            BENZODIAZEPINE OVERDOSE                     0   \n",
       "1  CORONARY ARTERY DISEASE\\CORONARY ARTERY BYPASS...                     0   \n",
       "2                                         BRAIN MASS                     0   \n",
       "3                     INTERIOR MYOCARDIAL INFARCTION                     0   \n",
       "4                            ACUTE CORONARY SYNDROME                     0   \n",
       "\n",
       "   has_chartevents_data  \n",
       "0                     1  \n",
       "1                     1  \n",
       "2                     1  \n",
       "3                     1  \n",
       "4                     1  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adm = loader[\"admissions\"]\n",
    "adm.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "26150"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unique_stay_id = vitals_table[\"icustay_id\"].unique()\n",
    "len(unique_stay_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "for id in unique_stay_id:\n",
    "    if len(vitals_table[vitals_table[\"icustay_id\"] == id]) != 25:\n",
    "        print(f\"actual length is {len(vitals_table[vitals_table['icustay_id'] == id])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Extracting...: 100%|██████████| 26150/26150 [00:48<00:00, 536.99it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(26150, 15, 25)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "def turn_to_ndarray(values_df, ids, adm_df, icu_df):\n",
    "    \n",
    "    def add_mortality_flag(tmp_df, adm_df, id):\n",
    "        tmp_df = tmp_df.copy()\n",
    "        subject_id = icu_df[icu_df[\"icustay_id\"] == id][\"subject_id\"].values[0]         # given icustay_id, find subject_id\n",
    "        flag = adm_df[adm_df[\"subject_id\"] == subject_id][\"hospital_expire_flag\"].values[0]\n",
    "        flag_vec = np.ones((tmp_df.shape[0], 1)) * flag\n",
    "        tmp_df[\"mortality\"] = flag_vec\n",
    "        return tmp_df\n",
    "    \n",
    "    def add_cumulative_nan_flag(tmp_array):\n",
    "        result = []\n",
    "        for idx in range(tmp_array.shape[1]):\n",
    "            if idx <= tmp_array.shape[1] - 2:\n",
    "                result.append(tmp_array[:, idx])\n",
    "                result.append(np.cumsum(~np.isnan(tmp_array[:, idx]), axis=0))\n",
    "            else:\n",
    "                result.append(tmp_array[:, idx])\n",
    "        return np.asarray(result)\n",
    "        \n",
    "    result = []\n",
    "    for id in tqdm(ids, desc=\"Extracting...\"):\n",
    "        tmp_df = values_df[values_df[\"icustay_id\"] == id]\n",
    "        tmp_df = add_mortality_flag(tmp_df, adm_df, id)\n",
    "        tmp_array = add_cumulative_nan_flag(tmp_df.drop([\"icustay_id\", \"time\"], axis=1).values)\n",
    "        result.append(tmp_array)\n",
    "    \n",
    "    return np.asarray(result)\n",
    "\n",
    "icu_df = loader[\"icustays\"]\n",
    "vitals_ndarray = turn_to_ndarray(vitals_table, unique_stay_id, adm, icu_df)\n",
    "vitals_ndarray.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "train test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train shape: (20920, 15, 25), test shape: (5230, 15, 25)\n"
     ]
    }
   ],
   "source": [
    "def train_test_split(array, train_percent=0.8):\n",
    "    indices = np.arange(len(array))\n",
    "    np.random.shuffle(indices)\n",
    "    train, test = array[indices[:int(len(array) * train_percent)]], array[indices[int(len(array) * train_percent):]]\n",
    "    return train, test\n",
    "\n",
    "def check_class_imbalance(array):\n",
    "    return np.sum(array[:,-1,:]) / (array.shape[0] * array.shape[1])\n",
    "\n",
    "train, test = train_test_split(vitals_ndarray)\n",
    "print(f\"train shape: {train.shape}, test shape: {test.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.13304652644996814, 0.12938177182919056)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "check_class_imbalance(train), check_class_imbalance(test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "save data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "np.save(\"data_pack/mimic-iii/train.npy\", train)\n",
    "train = torch.tensor(train)\n",
    "torch.save(train, \"data_pack/mimic-iii/train.pt\")\n",
    "np.save(\"data_pack/mimic-iii/test.npy\", test)\n",
    "test = torch.tensor(test)\n",
    "torch.save(test, \"data_pack/mimic-iii/test.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
}
