Skip to content

Add/task caching #546

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jhnwu3 merged 2 commits into from
Sep 23, 2025
+650 −27
Merged

Add/task caching #546

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
73ebc05
added task caching with parquet and pickle formats
Sep 14, 2025
2f4e11b
added notebook with more realistic testing
Sep 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
291 changes: 291 additions & 0 deletions examples/mimic3_mortality_prediction_cached.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,291 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "cc001dde",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Adding PyHealth to sys.path: /home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth\n"
]
}
],
"source": [
"import os\n",
"import sys\n",
"\n",
"pyhealth_path = os.path.dirname(os.getcwd())\n",
"if pyhealth_path not in sys.path:\n",
" print(f\"Adding PyHealth to sys.path: {pyhealth_path}\")\n",
" sys.path.insert(0, pyhealth_path)"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "9887d8db",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'/home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth'"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"pyhealth_path"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "3e05a68b",
"metadata": {},
"outputs": [],
"source": [
"from pyhealth.tasks import BaseTask\n",
"from typing import Any, Dict, List, Optional\n",
"from datetime import datetime\n",
"\n",
"class MortalityPredictionMIMIC3Heterogeneous(BaseTask):\n",
" \"\"\"Task for predicting mortality using MIMIC-III dataset with text data.\n",
"\n",
" This task aims to predict whether the patient will decease in the next hospital\n",
" visit based on clinical information from the current visit.\n",
" \"\"\"\n",
"\n",
" task_name: str = \"MortalityPredictionMIMIC3\"\n",
" input_schema: Dict[str, str] = {\n",
" \"conditions\": \"sequence\",\n",
" \"procedures\": \"sequence\",\n",
" \"drugs\": \"sequence\",\n",
" }\n",
" output_schema: Dict[str, str] = {\"mortality\": \"binary\"}\n",
"\n",
" def __call__(self, patient: Any) -> List[Dict[str, Any]]:\n",
" \"\"\"Processes a single patient for the mortality prediction task.\"\"\"\n",
" samples = []\n",
"\n",
" # We will drop the last visit\n",
" visits = patient.get_events(event_type=\"admissions\")\n",
"\n",
" if len(visits) <= 1:\n",
" return []\n",
"\n",
" for i in range(len(visits) - 1):\n",
" visit = visits[i]\n",
" next_visit = visits[i + 1]\n",
"\n",
" # Check discharge status for mortality label - more robust handling\n",
" if next_visit.hospital_expire_flag not in [0, 1, \"0\", \"1\"]:\n",
" mortality_label = 0\n",
" else:\n",
" mortality_label = int(next_visit.hospital_expire_flag)\n",
"\n",
" # Convert string timestamps to datetime objects\n",
" try:\n",
" # Check the type and convert if necessary\n",
" if isinstance(visit.dischtime, str):\n",
" discharge_time = datetime.strptime(\n",
" visit.dischtime, \"%Y-%m-%d %H:%M:%S\"\n",
" )\n",
" else:\n",
" discharge_time = visit.dischtime\n",
" except (ValueError, AttributeError):\n",
" # If conversion fails, skip this visit\n",
" print(\"Error parsing discharge time:\", visit.dischtime)\n",
" continue\n",
"\n",
" # Get clinical codes\n",
" diagnoses = patient.get_events(\n",
" event_type=\"diagnoses_icd\",\n",
" start=visit.timestamp,\n",
" end=discharge_time, # Now using a datetime object\n",
" )\n",
" procedures = patient.get_events(\n",
" event_type=\"procedures_icd\",\n",
" start=visit.timestamp,\n",
" end=discharge_time, # Now using a datetime object\n",
" )\n",
" prescriptions = patient.get_events(\n",
" event_type=\"prescriptions\",\n",
" start=visit.timestamp,\n",
" end=discharge_time, # Now using a datetime object\n",
" )\n",
"\n",
" conditions = [event.icd9_code for event in diagnoses]\n",
" procedures_list = [event.icd9_code for event in procedures]\n",
" drugs = [event.drug for event in prescriptions]\n",
"\n",
" # Exclude visits without condition, procedure, or drug code\n",
" samples.append(\n",
" {\n",
" \"hadm_id\": visit.hadm_id,\n",
" \"patient_id\": patient.patient_id,\n",
" \"conditions\": conditions,\n",
" \"procedures\": procedures_list,\n",
" \"drugs\": drugs,\n",
" \"mortality\": mortality_label,\n",
" }\n",
" )\n",
"\n",
" return samples\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "da8360f8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"No config path provided, using default config\n",
"Initializing mimic3 dataset from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III (dev mode: False)\n",
"Scanning table: patients from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/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",
" warnings.warn(\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Some column names were converted to lowercase\n",
"Scanning table: admissions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
"Some column names were converted to lowercase\n",
"Scanning table: icustays from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv\n",
"Some column names were converted to lowercase\n",
"Scanning table: diagnoses_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv\n",
"Some column names were converted to lowercase\n",
"Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
"Scanning table: procedures_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv\n",
"Some column names were converted to lowercase\n",
"Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
"Scanning table: prescriptions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv\n",
"Some column names were converted to lowercase\n",
"Scanning table: noteevents from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/NOTEEVENTS.csv.gz\n",
"Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/NOTEEVENTS.csv\n",
"Some column names were converted to lowercase\n",
"Preprocessing table: noteevents with preprocess_noteevents\n"
]
}
],
"source": [
"from pyhealth.datasets import MIMIC3Dataset\n",
"dataset = MIMIC3Dataset(\n",
" root=\"https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III\",\n",
" tables=[\"diagnoses_icd\", \"procedures_icd\", \"prescriptions\", \"noteevents\"]\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "dd5c282f",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Setting task MortalityPredictionMIMIC3 for mimic3 base dataset...\n",
"Loading cached samples from cache/MortalityPredictionMIMIC3.parquet\n",
"Loaded 2776 cached samples\n",
"Label mortality vocab: {0: 0, 1: 1}\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Processing samples: 100%|██████████| 2776/2776 [00:00<00:00, 63155.03it/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Generated 2776 samples for task MortalityPredictionMIMIC3\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n"
]
}
],
"source": [
"from pyhealth.tasks.mortality_prediction import MortalityPredictionMIMIC3\n",
"from pyhealth.datasets import split_by_patient, get_dataloader\n",
"mimic3_mortality_prediction = MortalityPredictionMIMIC3Heterogeneous()\n",
"samples = dataset.set_task(mimic3_mortality_prediction, num_workers=1, cache_dir=\"cache/\") # use default task"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "06b5a9a3",
"metadata": {},
"outputs": [],
"source": [
"from pyhealth.datasets import split_by_sample\n",
"\n",
"\n",
"train_dataset, val_dataset, test_dataset = split_by_sample(\n",
" dataset=samples,\n",
" ratios=[0.7, 0.1, 0.2]\n",
")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "medical_coding_demo",
"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.10.16"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Loading