2_Generate Embeddings from Pickle Files.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ffba6d5f",
   "metadata": {},
   "source": [
    "# Code to Generate HAIM embeddings from HAIM-MIMIC-MM dataset\n",
    "\n",
    "### Project Info\n",
    " ->Copyright 2020 (Last Update: June 07, 2022)\n",
    " \n",
    " -> Authors: \n",
    "        Luis R Soenksen (<soenksen@mit.edu>),\n",
    "        Yu Ma (<midsumer@mit.edu>),\n",
    "        Cynthia Zeng (<czeng12@mit.edu>),\n",
    "        Ignacio Fuentes (<ifuentes@mit.edu>),\n",
    "        Leonard David Jean Boussioux (<leobix@mit.edu>),\n",
    "        Agni Orfanoudaki (<agniorf@mit.edu>),\n",
    "        Holly Mika Wiberg (<hwiberg@mit.edu>),\n",
    "        Michael Lingzhi Li (<mlli@mit.edu>),\n",
    "        Kimberly M Villalobos Carballo (<kimvc@mit.edu>),\n",
    "        Liangyuan Na (<lyna@mit.edu>),\n",
    "        Dimitris J Bertsimas (<dbertsim@mit.edu>),\n",
    "\n",
    "```\n",
    "**Licensed under the Apache License, Version 2.0**\n",
    "You may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
    "https://www.apache.org/licenses/LICENSE-2.0\n",
    "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "449bfc15",
   "metadata": {},
   "source": [
    "### Requires \n",
    "```\n",
    " -> Previously generated pickle files from HAIM-MIMIC-MM Dataset\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2063bb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from MIMIC_IV_HAIM_API import *\n",
    "import gc\n",
    "\n",
    "# Full Core MIMIC-IV database path\n",
    "core_mimiciv_path = 'data/'\n",
    "df_haim_ids = pd.read_csv(core_mimiciv_path + 'pickle/haim_mimiciv_key_ids.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdee1132",
   "metadata": {},
   "outputs": [],
   "source": [
    "# General function that processes all data embeddings\n",
    "def process_cxr_embeddings_haim_id(haim_id, dt_patient, df_init):\n",
    "    # DEMOGRAPHICS EMBEDDINGS EXTRACTION\n",
    "    demo_embeddings = get_demographic_embeddings(dt_patient, verbose=0)\n",
    "    gc.collect() #Clear memory\n",
    "    \n",
    "    # Time Series (TSFRESH-like) CHARTEVENT & LABEVENT EMBEDDINGS EXTRACTION\n",
    "    aggregated_ts_ce_embeddings = get_ts_embeddings(dt_patient, event_type = 'chart')\n",
    "    gc.collect() #Clear memory\n",
    "    \n",
    "    aggregated_ts_le_embeddings = get_ts_embeddings(dt_patient, event_type = 'lab')\n",
    "    gc.collect() #Clear memory\n",
    "    \n",
    "    aggregated_ts_pe_embeddings = get_ts_embeddings(dt_patient, event_type = 'procedure')\n",
    "    gc.collect() #Clear memory\n",
    "    \n",
    "    # CHEST XRAY VISION EMBEDDINGS EXTRACTION\n",
    "    aggregated_densefeature_embeddings, _, aggregated_prediction_embeddings, _, _ = get_chest_xray_embeddings(dt_patient, verbose=0)\n",
    "    gc.collect() #Clear memory\n",
    "    \n",
    "    # NOTES FROM ECGs\n",
    "    aggregated_ecg_embeddings = get_notes_biobert_embeddings(patient, note_type = 'ecgnotes')\n",
    "    gc.collect() #Clear memory\n",
    "    \n",
    "    # NOTES FROM ECOCARDIOGRAMs\n",
    "    aggregated_echo_embeddings = get_notes_biobert_embeddings(patient, note_type = 'echonotes')\n",
    "    gc.collect() #Clear memory\n",
    "    \n",
    "    # NOTES FROM RADIOLOGY\n",
    "    aggregated_rad_embeddings = get_notes_biobert_embeddings(patient, note_type = 'radnotes')\n",
    "    gc.collect() #Clear memory\n",
    "\n",
    "    # CHEST XRAY VISION SINGLE-IMAGE EMBEDDINGS EXTRACTION\n",
    "    print('getting xray')\n",
    "    img = df_imcxr[idx]\n",
    "    densefeature_embeddings, prediction_embeddings = get_single_chest_xray_embeddings(img)\n",
    "    gc.collect() #Clear memory\n",
    "\n",
    "    # Create Dataframes filteed by ordered sample number for Fusion\n",
    "    df_haim_ids_fusion = pd.DataFrame([haim_id],columns=['haim_id'])\n",
    "    df_demographics_embeddings_fusion = pd.DataFrame(demo_embeddings.reshape(1,-1), columns=['de_'+str(i) for i in range(demo_embeddings.shape[0])])\n",
    "    df_ts_ce_embeddings_fusion = pd.DataFrame(aggregated_ts_ce_embeddings.values.reshape(1,-1), columns=['ts_ce_'+str(i) for i in range(aggregated_ts_ce_embeddings.values.shape[0])])\n",
    "    df_ts_le_embeddings_fusion = pd.DataFrame(aggregated_ts_le_embeddings.values.reshape(1,-1), columns=['ts_le_'+str(i) for i in range(aggregated_ts_le_embeddings.values.shape[0])])\n",
    "    df_ts_pe_embeddings_fusion = pd.DataFrame(aggregated_ts_pe_embeddings.values.reshape(1,-1), columns=['ts_pe_'+str(i) for i in range(aggregated_ts_pe_embeddings.values.shape[0])])\n",
    "    \n",
    "    df_vision_dense_embeddings_fusion = pd.DataFrame(densefeature_embeddings.reshape(1,-1), columns=['vd_'+str(i) for i in range(densefeature_embeddings.shape[0])])\n",
    "    df_vision_predictions_embeddings_fusion = pd.DataFrame(prediction_embeddings.reshape(1,-1), columns=['vp_'+str(i) for i in range(prediction_embeddings.shape[0])])\n",
    "    df_vision_multi_dense_embeddings_fusion = pd.DataFrame(aggregated_densefeature_embeddings.reshape(1,-1), columns=['vmd_'+str(i) for i in range(aggregated_densefeature_embeddings.shape[0])])\n",
    "    df_vision_multi_predictions_embeddings_fusion = pd.DataFrame(aggregated_prediction_embeddings.reshape(1,-1), columns=['vmp_'+str(i) for i in range(aggregated_prediction_embeddings.shape[0])])\n",
    "    df_ecgnotes_embeddings_fusion = pd.DataFrame(aggregated_ecg_embeddings.reshape(1,-1), columns=['n_ecg_'+str(i) for i in range(aggregated_ecg_embeddings.shape[0])])\n",
    "    df_echonotes_embeddings_fusion = pd.DataFrame(aggregated_echo_embeddings.reshape(1,-1), columns=['n_ech_'+str(i) for i in range(aggregated_echo_embeddings.shape[0])])\n",
    "    df_radnotes_embeddings_fusion = pd.DataFrame(aggregated_rad_embeddings.reshape(1,-1), columns=['n_rad_'+str(i) for i in range(aggregated_rad_embeddings.shape[0])])\n",
    "    \n",
    "    # Vision targets\n",
    "    cxr_target_columns = ['split','Atelectasis','Cardiomegaly','Consolidation','Edema','Enlarged Cardiomediastinum','Fracture','Lung Lesion','Lung Opacity','No Finding','Pleural Effusion','Pleural Other','Pneumonia','Pneumothorax','Support Devices', 'PerformedProcedureStepDescription','ViewPosition']\n",
    "    df_vision_targets_fusion = df_stay_cxr.loc[idx:idx][cxr_target_columns].reset_index(drop=True)\n",
    "\n",
    "    # Embeddings FUSION\n",
    "    df_fusion = df_haim_ids_fusion\n",
    "    df_fusion = pd.concat([df_fusion, df_init], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_demographics_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_vision_dense_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_vision_predictions_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_vision_multi_dense_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_vision_multi_predictions_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_ts_ce_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_ts_le_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_ts_pe_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_ecgnotes_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_echonotes_embeddings_fusion], axis=1)\n",
    "    df_fusion = pd.concat([df_fusion, df_radnotes_embeddings_fusion], axis=1)\n",
    "    \n",
    "    #Add targets\n",
    "    df_fusion = pd.concat([df_fusion, df_vision_targets_fusion], axis=1)\n",
    "    gc.collect() #Clear memory\n",
    "    \n",
    "    return df_fusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40e58ef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's select a single HAIM Patient from pickle files and check if it fits inclusion criteria\n",
    "haim_patient_idx = 4\n",
    "\n",
    "#Load precomputed file\n",
    "filename = f\"{haim_patient_idx:08d}\" + '.pkl'\n",
    "folder = f\"{haim_patient_idx:05d}\"[:2] + \"/\"\n",
    "patient = load_patient_object(core_mimiciv_path + 'pickle/folder' + folder + filename)\n",
    "\n",
    "# Get information of chest x-rays conducted within this patiewnt stay\n",
    "df_cxr = patient.cxr\n",
    "df_imcxr = patient.imcxr\n",
    "admittime = patient.admissions.admittime.values[0]\n",
    "dischtime = patient.admissions.dischtime.values[0]\n",
    "df_stay_cxr = df_cxr.loc[(df_cxr['charttime'] >= admittime) & (df_cxr['charttime'] <= dischtime)]\n",
    "\n",
    "if not df_stay_cxr.empty:\n",
    "    for idx, df_stay_cxr_row in df_stay_cxr.iterrows():\n",
    "        # Get stay anchor times\n",
    "        img_charttime = df_stay_cxr_row['charttime']\n",
    "        img_deltacharttime = df_stay_cxr_row['deltacharttime']\n",
    "\n",
    "        # Get time to discharge and discharge location/status\n",
    "        img_id = df_stay_cxr_row[\"dicom_id\"]\n",
    "        img_length_of_stay = date_diff_hrs(dischtime, img_charttime)\n",
    "        discharge_location = patient.core['discharge_location'][0]\n",
    "        if discharge_location == \"DIED\": death_status = 1\n",
    "        else: death_status = 0\n",
    "            \n",
    "        # Select allowed timestamp range\n",
    "        start_hr = None\n",
    "        end_hr = img_deltacharttime\n",
    "        \n",
    "        # We need to reload it since the original object has been modified\n",
    "        patient = load_patient_object(core_mimiciv_path + 'pickle/folder' + folder + filename)\n",
    "        dt_patient = get_timebound_patient_icustay(patient, start_hr , end_hr)\n",
    "        is_included = True\n",
    "\n",
    "        if is_included:\n",
    "            df_init = pd.DataFrame([[img_id, img_charttime, img_deltacharttime, discharge_location, img_length_of_stay, death_status]],columns=['img_id', 'img_charttime', 'img_deltacharttime', 'discharge_location', 'img_length_of_stay', 'death_status'])\n",
    "            df_fusion = process_cxr_embeddings_haim_id(haim_id, dt_patient, df_init)\n",
    "            \n",
    "            if os.path.isfile(fname):\n",
    "                df_fusion.to_csv(fname, mode='a', index=False, header=False)\n",
    "            else:\n",
    "                df_fusion.to_csv(fname, mode='w', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}