Skip to content

Commit 077e798

Browse files
jhnwu3John Wu
jhnwu3
and
John Wu
authored
Add/task caching (sunlabuiuc#546)
* added task caching with parquet and pickle formats * added notebook with more realistic testing --------- Co-authored-by: John Wu <johnwu3@sunlab-serv-03.cs.illinois.edu>
1 parent e861318 commit 077e798

3 files changed

Lines changed: 650 additions & 27 deletions

File tree

‎examples/mimic3_mortality_prediction_cached.ipynb‎

Lines changed: 291 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,291 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": 1,
6+
"id": "cc001dde",
7+
"metadata": {},
8+
"outputs": [
9+
{
10+
"name": "stdout",
11+
"output_type": "stream",
12+
"text": [
13+
"Adding PyHealth to sys.path: /home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth\n"
14+
]
15+
}
16+
],
17+
"source": [
18+
"import os\n",
19+
"import sys\n",
20+
"\n",
21+
"pyhealth_path = os.path.dirname(os.getcwd())\n",
22+
"if pyhealth_path not in sys.path:\n",
23+
" print(f\"Adding PyHealth to sys.path: {pyhealth_path}\")\n",
24+
" sys.path.insert(0, pyhealth_path)"
25+
]
26+
},
27+
{
28+
"cell_type": "code",
29+
"execution_count": 2,
30+
"id": "9887d8db",
31+
"metadata": {},
32+
"outputs": [
33+
{
34+
"data": {
35+
"text/plain": [
36+
"'/home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth'"
37+
]
38+
},
39+
"execution_count": 2,
40+
"metadata": {},
41+
"output_type": "execute_result"
42+
}
43+
],
44+
"source": [
45+
"pyhealth_path"
46+
]
47+
},
48+
{
49+
"cell_type": "code",
50+
"execution_count": 3,
51+
"id": "3e05a68b",
52+
"metadata": {},
53+
"outputs": [],
54+
"source": [
55+
"from pyhealth.tasks import BaseTask\n",
56+
"from typing import Any, Dict, List, Optional\n",
57+
"from datetime import datetime\n",
58+
"\n",
59+
"class MortalityPredictionMIMIC3Heterogeneous(BaseTask):\n",
60+
" \"\"\"Task for predicting mortality using MIMIC-III dataset with text data.\n",
61+
"\n",
62+
" This task aims to predict whether the patient will decease in the next hospital\n",
63+
" visit based on clinical information from the current visit.\n",
64+
" \"\"\"\n",
65+
"\n",
66+
" task_name: str = \"MortalityPredictionMIMIC3\"\n",
67+
" input_schema: Dict[str, str] = {\n",
68+
" \"conditions\": \"sequence\",\n",
69+
" \"procedures\": \"sequence\",\n",
70+
" \"drugs\": \"sequence\",\n",
71+
" }\n",
72+
" output_schema: Dict[str, str] = {\"mortality\": \"binary\"}\n",
73+
"\n",
74+
" def __call__(self, patient: Any) -> List[Dict[str, Any]]:\n",
75+
" \"\"\"Processes a single patient for the mortality prediction task.\"\"\"\n",
76+
" samples = []\n",
77+
"\n",
78+
" # We will drop the last visit\n",
79+
" visits = patient.get_events(event_type=\"admissions\")\n",
80+
"\n",
81+
" if len(visits) <= 1:\n",
82+
" return []\n",
83+
"\n",
84+
" for i in range(len(visits) - 1):\n",
85+
" visit = visits[i]\n",
86+
" next_visit = visits[i + 1]\n",
87+
"\n",
88+
" # Check discharge status for mortality label - more robust handling\n",
89+
" if next_visit.hospital_expire_flag not in [0, 1, \"0\", \"1\"]:\n",
90+
" mortality_label = 0\n",
91+
" else:\n",
92+
" mortality_label = int(next_visit.hospital_expire_flag)\n",
93+
"\n",
94+
" # Convert string timestamps to datetime objects\n",
95+
" try:\n",
96+
" # Check the type and convert if necessary\n",
97+
" if isinstance(visit.dischtime, str):\n",
98+
" discharge_time = datetime.strptime(\n",
99+
" visit.dischtime, \"%Y-%m-%d %H:%M:%S\"\n",
100+
" )\n",
101+
" else:\n",
102+
" discharge_time = visit.dischtime\n",
103+
" except (ValueError, AttributeError):\n",
104+
" # If conversion fails, skip this visit\n",
105+
" print(\"Error parsing discharge time:\", visit.dischtime)\n",
106+
" continue\n",
107+
"\n",
108+
" # Get clinical codes\n",
109+
" diagnoses = patient.get_events(\n",
110+
" event_type=\"diagnoses_icd\",\n",
111+
" start=visit.timestamp,\n",
112+
" end=discharge_time, # Now using a datetime object\n",
113+
" )\n",
114+
" procedures = patient.get_events(\n",
115+
" event_type=\"procedures_icd\",\n",
116+
" start=visit.timestamp,\n",
117+
" end=discharge_time, # Now using a datetime object\n",
118+
" )\n",
119+
" prescriptions = patient.get_events(\n",
120+
" event_type=\"prescriptions\",\n",
121+
" start=visit.timestamp,\n",
122+
" end=discharge_time, # Now using a datetime object\n",
123+
" )\n",
124+
"\n",
125+
" conditions = [event.icd9_code for event in diagnoses]\n",
126+
" procedures_list = [event.icd9_code for event in procedures]\n",
127+
" drugs = [event.drug for event in prescriptions]\n",
128+
"\n",
129+
" # Exclude visits without condition, procedure, or drug code\n",
130+
" samples.append(\n",
131+
" {\n",
132+
" \"hadm_id\": visit.hadm_id,\n",
133+
" \"patient_id\": patient.patient_id,\n",
134+
" \"conditions\": conditions,\n",
135+
" \"procedures\": procedures_list,\n",
136+
" \"drugs\": drugs,\n",
137+
" \"mortality\": mortality_label,\n",
138+
" }\n",
139+
" )\n",
140+
"\n",
141+
" return samples\n",
142+
"\n"
143+
]
144+
},
145+
{
146+
"cell_type": "code",
147+
"execution_count": 4,
148+
"id": "da8360f8",
149+
"metadata": {},
150+
"outputs": [
151+
{
152+
"name": "stdout",
153+
"output_type": "stream",
154+
"text": [
155+
"No config path provided, using default config\n",
156+
"Initializing mimic3 dataset from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III (dev mode: False)\n",
157+
"Scanning table: patients from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv.gz\n",
158+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv\n"
159+
]
160+
},
161+
{
162+
"name": "stderr",
163+
"output_type": "stream",
164+
"text": [
165+
"/home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth/pyhealth/datasets/mimic3.py:50: UserWarning: Events from prescriptions table only have date timestamp (no specific time). This may affect temporal ordering of events.\n",
166+
" warnings.warn(\n"
167+
]
168+
},
169+
{
170+
"name": "stdout",
171+
"output_type": "stream",
172+
"text": [
173+
"Some column names were converted to lowercase\n",
174+
"Scanning table: admissions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
175+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
176+
"Some column names were converted to lowercase\n",
177+
"Scanning table: icustays from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv.gz\n",
178+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv\n",
179+
"Some column names were converted to lowercase\n",
180+
"Scanning table: diagnoses_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv.gz\n",
181+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv\n",
182+
"Some column names were converted to lowercase\n",
183+
"Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
184+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
185+
"Scanning table: procedures_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv.gz\n",
186+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv\n",
187+
"Some column names were converted to lowercase\n",
188+
"Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
189+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
190+
"Scanning table: prescriptions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv.gz\n",
191+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv\n",
192+
"Some column names were converted to lowercase\n",
193+
"Scanning table: noteevents from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/NOTEEVENTS.csv.gz\n",
194+
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/NOTEEVENTS.csv\n",
195+
"Some column names were converted to lowercase\n",
196+
"Preprocessing table: noteevents with preprocess_noteevents\n"
197+
]
198+
}
199+
],
200+
"source": [
201+
"from pyhealth.datasets import MIMIC3Dataset\n",
202+
"dataset = MIMIC3Dataset(\n",
203+
" root=\"https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III\",\n",
204+
" tables=[\"diagnoses_icd\", \"procedures_icd\", \"prescriptions\", \"noteevents\"]\n",
205+
")"
206+
]
207+
},
208+
{
209+
"cell_type": "code",
210+
"execution_count": 8,
211+
"id": "dd5c282f",
212+
"metadata": {},
213+
"outputs": [
214+
{
215+
"name": "stdout",
216+
"output_type": "stream",
217+
"text": [
218+
"Setting task MortalityPredictionMIMIC3 for mimic3 base dataset...\n",
219+
"Loading cached samples from cache/MortalityPredictionMIMIC3.parquet\n",
220+
"Loaded 2776 cached samples\n",
221+
"Label mortality vocab: {0: 0, 1: 1}\n"
222+
]
223+
},
224+
{
225+
"name": "stderr",
226+
"output_type": "stream",
227+
"text": [
228+
"Processing samples: 100%|██████████| 2776/2776 [00:00<00:00, 63155.03it/s]"
229+
]
230+
},
231+
{
232+
"name": "stdout",
233+
"output_type": "stream",
234+
"text": [
235+
"Generated 2776 samples for task MortalityPredictionMIMIC3\n"
236+
]
237+
},
238+
{
239+
"name": "stderr",
240+
"output_type": "stream",
241+
"text": [
242+
"\n"
243+
]
244+
}
245+
],
246+
"source": [
247+
"from pyhealth.tasks.mortality_prediction import MortalityPredictionMIMIC3\n",
248+
"from pyhealth.datasets import split_by_patient, get_dataloader\n",
249+
"mimic3_mortality_prediction = MortalityPredictionMIMIC3Heterogeneous()\n",
250+
"samples = dataset.set_task(mimic3_mortality_prediction, num_workers=1, cache_dir=\"cache/\") # use default task"
251+
]
252+
},
253+
{
254+
"cell_type": "code",
255+
"execution_count": 9,
256+
"id": "06b5a9a3",
257+
"metadata": {},
258+
"outputs": [],
259+
"source": [
260+
"from pyhealth.datasets import split_by_sample\n",
261+
"\n",
262+
"\n",
263+
"train_dataset, val_dataset, test_dataset = split_by_sample(\n",
264+
" dataset=samples,\n",
265+
" ratios=[0.7, 0.1, 0.2]\n",
266+
")"
267+
]
268+
}
269+
],
270+
"metadata": {
271+
"kernelspec": {
272+
"display_name": "medical_coding_demo",
273+
"language": "python",
274+
"name": "python3"
275+
},
276+
"language_info": {
277+
"codemirror_mode": {
278+
"name": "ipython",
279+
"version": 3
280+
},
281+
"file_extension": ".py",
282+
"mimetype": "text/x-python",
283+
"name": "python",
284+
"nbconvert_exporter": "python",
285+
"pygments_lexer": "ipython3",
286+
"version": "3.10.16"
287+
}
288+
},
289+
"nbformat": 4,
290+
"nbformat_minor": 5
291+
}

0 commit comments

Comments
 (0)