diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index f72e41f8e..a572dd1b6 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -16,7 +16,16 @@ --task clinical_notes_icd_labs ClinicalNotesICDLabsMIMIC4: discharge/radiology notes + ICD + labs. - Requires --note-root. Used for Table 2 (EHR + clinical text). + Requires --note-root. Legacy; ICD codes are discharge-coded (leakage). + +--task notes_labs (recommended for multimodal) + NotesLabsMIMIC4: admission-context note sections + labs, no ICD codes. + Extracts Chief Complaint, HPI, PMH, Medications on Admission from the + discharge note — text available at admission time, ~90%+ coverage. + Requires --note-root. + Add --freeze-encoder to freeze Bio_ClinicalBERT and train only the + backbone; cuts BERT VRAM by ~50%, useful on smaller GPUs (≤24 GB). + Add --icd-codes to include discharge-coded ICD codes (ablation only). Example ------- @@ -50,19 +59,26 @@ from typing import Any, Tuple import numpy as np +import torch from pyhealth.datasets import ( MIMIC4Dataset, get_dataloader, split_by_patient, split_by_sample, + sample_balanced, ) from pyhealth.models import MLP, RNN, Transformer, UnifiedMultimodalEmbeddingModel from pyhealth.models.bottleneck_transformer import BottleneckTransformer from pyhealth.models.ehrmamba import EHRMamba from pyhealth.models.jamba_ehr import JambaEHR from pyhealth.tasks import MortalityPredictionStageNetMIMIC4 -from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4 +from pyhealth.tasks.multimodal_mimic4 import ( + ClinicalNotesICDLabsMIMIC4, + ICDLabsMIMIC4, + LabsMIMIC4, + NotesLabsMIMIC4, +) from pyhealth.trainer import Trainer from pyhealth.utils import set_seed @@ -79,13 +95,30 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: if args.task == "icd_labs": ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] + if args.task == "notes_labs": + if not args.note_root: + raise ValueError("--task notes_labs requires --note-root.") + note_tables = ["discharge"] + # Load ICD tables only when explicitly requested (they are discharge-coded). + ehr_tables = ( + ["diagnoses_icd", "procedures_icd", "labevents"] + if args.icd_codes + else ["labevents"] + ) + if args.include_vitals: + if "chartevents" not in ehr_tables: + ehr_tables.append("chartevents") + + if args.task == "labs": + ehr_tables = ["labevents"] + return MIMIC4Dataset( ehr_root=args.ehr_root, ehr_tables=ehr_tables, note_root=args.note_root if note_tables else None, note_tables=note_tables, cache_dir=args.cache_dir, - dev=args.dev, + dev=args.dev if args.dev else False, num_workers=args.num_workers, ) @@ -97,6 +130,14 @@ def _build_task(args: argparse.Namespace): return ICDLabsMIMIC4(window_hours=args.observation_window_hours) if args.task == "clinical_notes_icd_labs": return ClinicalNotesICDLabsMIMIC4(window_hours=args.observation_window_hours) + if args.task == "notes_labs": + return NotesLabsMIMIC4( + window_hours=args.observation_window_hours, + include_icd=args.icd_codes, + include_vitals=args.include_vitals, + ) + if args.task == "labs": + return LabsMIMIC4(window_hours=args.observation_window_hours) raise ValueError(f"Unknown task: {args.task}") @@ -111,6 +152,7 @@ def _build_model(args: argparse.Namespace, sample_dataset: Any): unified = UnifiedMultimodalEmbeddingModel( processors=sample_dataset.input_processors, embedding_dim=args.embedding_dim, + freeze_text_encoder=args.freeze_encoder, ) if args.model == "mlp": @@ -204,6 +246,26 @@ def _write_predictions( ) +def _compute_pos_weight(train_ds, label_key: str = "mortality") -> float: + """Count pos/neg in train_ds and return n_neg/n_pos for BCE pos_weight.""" + n_pos = n_neg = 0 + for i in range(len(train_ds)): + sample = train_ds[i] + label = sample.get(label_key, 0) + if hasattr(label, "__iter__"): + label = next(iter(label)) + if float(label) > 0.5: + n_pos += 1 + else: + n_neg += 1 + if n_pos == 0: + return 1.0 + # Cap at 10: n_neg/n_pos ≈ 37 on MIMIC-IV mortality is too extreme with + # typical LRs and causes training oscillation. 10 still strongly corrects + # for imbalance while keeping gradient magnitudes tractable. + return min(10.0, n_neg / n_pos) + + def run(args: argparse.Namespace) -> Path: set_seed(args.seed) @@ -217,8 +279,29 @@ def run(args: argparse.Namespace) -> Path: ) train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + label_key = list(sample_dataset.output_schema.keys())[0] + + # Balanced sampling: undersample negatives to achieve a target pos:neg ratio. + if args.balanced_sampling: + ratio = args.balanced_ratio + print(f"[balanced_sampling] Undersampling training set to pos:neg ratio 1:{ratio}") + train_ds = sample_balanced(train_ds, ratio=ratio, seed=args.seed, label_key=label_key) + print(f"[balanced_sampling] Training set size after sampling: {len(train_ds)}") + model = _build_model(args, sample_dataset) + # Apply class-imbalance correction via BCE pos_weight. + # pos_weight = n_neg / n_pos so the rare positive class gets proportionally + # higher gradient signal, preventing all-negative collapse (F1=0). + if args.pos_weight is not None: + pw_value = args.pos_weight + else: + print(f"[pos_weight] Computing class balance from {len(train_ds)} training samples...") + pw_value = _compute_pos_weight(train_ds, label_key=label_key) + print(f"[pos_weight] Using pos_weight={pw_value:.2f} for binary BCE loss.") + model._pos_weight = torch.tensor([pw_value], dtype=torch.float32) + train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True) val_loader = ( get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False) @@ -257,8 +340,15 @@ def run(args: argparse.Namespace) -> Path: effective_max_grad_norm = 0.5 optimizer_params["eps"] = args.adam_eps if args.adam_eps is not None else 1e-6 else: + # All non-BT models: 1e-4 (was 1e-3). With pos_weight correction, + # effective gradient magnitude for positives is ~10x higher, so a + # smaller LR is needed to avoid training oscillation. if effective_lr is None: - effective_lr = 1e-3 + effective_lr = 1e-4 + # Universal grad clipping: prevents runaway updates from the weighted + # positive-class loss (pos_weight ≈ 10 scales positive gradients 10x). + if effective_max_grad_norm is None: + effective_max_grad_norm = 1.0 if args.adam_eps is not None: optimizer_params["eps"] = args.adam_eps @@ -299,8 +389,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--task", type=str, - choices=["icd_labs", "clinical_notes_icd_labs"], + choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "labs"], default="stagenet", + help=( + "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. " + "No ICD codes (discharge-coded = leakage). Recommended for multimodal." + ), ) parser.add_argument( "--model", @@ -321,9 +415,9 @@ def parse_args() -> argparse.Namespace: type=float, default=None, help=( - "Learning rate. Default is model-specific: 1e-3 for " - "mlp/rnn/transformer/ehrmamba/jambaehr, 1e-4 for " - "bottleneck_transformer." + "Learning rate. Default is 1e-4 for all models. " + "(Previously 1e-3 for mlp/rnn/transformer/ehrmamba/jambaehr — " + "reduced after pos_weight correction caused oscillation at 1e-3.)" ), ) parser.add_argument( @@ -340,10 +434,82 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--num-workers", type=int, default=1) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--patience", type=int, default=None) - parser.add_argument("--dev", action="store_true") + parser.add_argument( + "--dev", + nargs="?", + type=int, + const=1000, + default=0, + help=( + "Dev mode: limit dataset to N patients for fast iteration. " + "--dev (no value) defaults to 1000 patients. " + "--dev 5000 limits to 5000. Omit for full dataset." + ), + ) + parser.add_argument( + "--pos-weight", + type=float, + default=None, + help=( + "BCE pos_weight for the positive class (float). " + "Default: auto-computed as n_neg/n_pos from training split. " + "Set to 1.0 to disable class-imbalance correction." + ), + ) # Task-specific parser.add_argument("--observation-window-hours", type=int, default=24) + parser.add_argument( + "--icd-codes", + action="store_true", + default=False, + help=( + "Include discharge-coded ICD codes in notes_labs task. " + "Default: off (ICD codes are coded at discharge and constitute " + "data leakage for in-hospital mortality prediction). " + "Enable only for ablation / legacy comparison experiments." + ), + ) + parser.add_argument( + "--freeze-encoder", + action="store_true", + default=False, + help=( + "Freeze pretrained BERT text encoder weights and train only the " + "downstream backbone (MLP/RNN/Transformer head + projection layer). " + "Reduces VRAM by ~50% for the text branch; useful when GPU memory " + "is limited or for faster iteration on backbone architectures." + ), + ) + parser.add_argument( + "--include-vitals", + action="store_true", + default=False, + help=( + "Include ICU vital signs (HeartRate, SysBP, DiasBP, MeanBP, " + "RespRate, SpO2, Temperature) from chartevents as an additional " + "modality alongside labs and notes. Adds chartevents to EHR tables." + ), + ) + parser.add_argument( + "--balanced-sampling", + action="store_true", + default=False, + help=( + "Undersample the majority (negative) class in training to improve " + "PR-AUC on imbalanced datasets. Uses sample_balanced() to create a " + "1:--balanced-ratio pos:neg training set." + ), + ) + parser.add_argument( + "--balanced-ratio", + type=float, + default=1.0, + help=( + "Negatives per positive in the balanced training set. " + "Default: 1.0 (equal pos/neg). Only used with --balanced-sampling." + ), + ) # RNN-specific parser.add_argument("--rnn-type", type=str, default="GRU") diff --git a/examples/multimodal_mimic4_task_tutorial.ipynb b/examples/multimodal_mimic4_task_tutorial.ipynb index 87b12d4db..26cd34b20 100644 --- a/examples/multimodal_mimic4_task_tutorial.ipynb +++ b/examples/multimodal_mimic4_task_tutorial.ipynb @@ -25,7 +25,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "8477e2ac-322f-49e9-819f-ce7621b214c6", "metadata": {}, "outputs": [], @@ -33,6 +33,11 @@ "from datetime import datetime\n", "from typing import Any, Dict, List, Optional\n", "import os\n", +<<<<<<< HEAD + "from pathlib import Path\n", + "import tempfile\n", +======= +>>>>>>> origin/main "import shutil \n", "import random\n", "import numpy as np\n", @@ -50,7 +55,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "201c3a62-2a2e-47ab-9ab8-eea77812f349", "metadata": {}, "outputs": [], @@ -87,12 +92,27 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "9a466c4b-2283-443f-94a3-5768babddf50", "metadata": {}, - "outputs": [], - "source": [ + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ +<<<<<<< HEAD + "SEED = 30\n", +======= "SEED = 22\n", +>>>>>>> origin/main "random.seed(SEED)\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)" @@ -126,7 +146,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "05ec2dc9-1a0c-4360-a028-d93173dcba65", "metadata": {}, "outputs": [], @@ -173,6 +193,11 @@ " print(f\" storetime: {note.storetime}\")\n", " print(f\" text: {note.text[:char_limit]}..(Limited to {char_limit} Characters).\")\n", "\n", +<<<<<<< HEAD + "def print_patient_lab_info(sample, event_limit=5):\n", + " patient = dataset.get_patient(sample['patient_id'])\n", + " lab_events = patient.get_events(\n", +======= "\n", "LABCATEGORIES = {\n", " itemid: name\n", @@ -184,10 +209,28 @@ "def print_patient_lab_info(sample):\n", " patient = dataset.get_patient(sample['patient_id'])\n", " labs = patient.get_events(\n", +>>>>>>> origin/main " event_type=\"labevents\",\n", " start=sample['window_start'],\n", " end=sample['window_end'],\n", " )\n", +<<<<<<< HEAD + " print(f\"\\nlab events: {len(lab_events)}\")\n", + " for event in lab_events[:event_limit]:\n", + " print(f\" itemid: {event.itemid}\")\n", + " print(f\" charttime: {event.timestamp}\")\n", + " print(f\" storetime: {event.storetime}\")\n", + " print(f\" valuenum: {event.valuenum}\")\n", + " if len(lab_events) > event_limit:\n", + " print(f\" ... ({len(lab_events) - event_limit} more events)\")\n", + "\n", + "def get_project_root(marker=\"PyHealth\"):\n", + " path = Path(os.getcwd())\n", + " for parent in [path, *path.parents]:\n", + " if parent.name == marker:\n", + " return str(parent)\n", + " raise ValueError(f\"'{marker}' not found in path\")" +======= " print(f\"\\nlabevents: {len(labs)}\")\n", " for lab in labs:\n", " lab_name = LABCATEGORIES.get(lab['itemid'], \"Unknown\")\n", @@ -229,6 +272,7 @@ " else:\n", " os.remove(item_path)\n", " print(f\"Cache directory cleared: {path}\")" +>>>>>>> origin/main ] }, { @@ -246,12 +290,16 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "tv8ir9mp9t", "metadata": {}, "outputs": [], "source": [ +<<<<<<< HEAD + "PYHEALTH_REPO_ROOT = get_project_root()\n", +======= "PYHEALTH_REPO_ROOT = '/Users/williampang/Desktop/PyHealth'\n", +>>>>>>> origin/main "\n", "EHR_ROOT = os.path.join(PYHEALTH_REPO_ROOT, \"test-resources/core/mimic4demo\")\n", "NOTE_ROOT = os.path.join(PYHEALTH_REPO_ROOT, \"test-resources/core/mimic4demo\")\n", @@ -276,6 +324,10 @@ }, { "cell_type": "code", +<<<<<<< HEAD + "execution_count": 6, + "id": "2f0f9ac3-1e3e-438b-b463-32e839941ce0", +======= "execution_count": null, "id": "7ebcce82-abdc-4b17-b3e5-d536fbb97859", "metadata": {}, @@ -340,8 +392,42 @@ "cell_type": "code", "execution_count": null, "id": "a1c777e0-75a0-4096-9e9b-cb3e955fae46", - "metadata": {}, - "outputs": [], +>>>>>>> origin/main + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Memory usage Starting MIMIC4Dataset init: 576.3 MB\n", + "Initializing mimic4 dataset from /u/wp14/PyHealth/test-resources/core/mimic4demo|/u/wp14/PyHealth/test-resources/core/mimic4demo|None (dev mode: False)\n", + "Using provided cache_dir: /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132\n", + "Initializing MIMIC4EHRDataset with tables: ['diagnoses_icd', 'procedures_icd', 'labevents'] (dev mode: False)\n", + "Using default EHR config: /u/wp14/PyHealth/pyhealth/datasets/configs/mimic4_ehr.yaml\n", + "Memory usage Before initializing mimic4_ehr: 576.3 MB\n", + "Initializing mimic4_ehr dataset from /u/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n", + "Using provided cache_dir: /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/1756cd94-34cf-553f-8bbf-732cb4253291\n", + "Memory usage After initializing mimic4_ehr: 576.3 MB\n", + "Memory usage After EHR dataset initialization: 576.3 MB\n", + "Initializing MIMIC4NoteDataset with tables: ['discharge', 'radiology'] (dev mode: False)\n", + "Using default note config: /u/wp14/PyHealth/pyhealth/datasets/configs/mimic4_note.yaml\n", + "Memory usage Before initializing mimic4_note: 576.3 MB\n", + "Initializing mimic4_note dataset from /u/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n", + "Using provided cache_dir: /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/7ddd05e8-68d1-57f8-a246-264b9bfb6f49\n", + "Memory usage After initializing mimic4_note: 576.3 MB\n", + "Memory usage After Note dataset initialization: 576.3 MB\n", + "Memory usage Completed MIMIC4Dataset init: 576.3 MB\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/u/wp14/PyHealth/pyhealth/datasets/mimic4.py:121: UserWarning: Events from discharge table only have date timestamp (no specific time). This may affect temporal ordering of events.\n", + " warnings.warn(\n" + ] + } + ], "source": [ "dataset = MIMIC4Dataset(\n", " ehr_root=EHR_ROOT,\n", @@ -473,10 +559,49 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "id": "d3b027ab-c814-49cd-8542-885d76b2401b", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Setting task ClinicalNotesICDLabsMIMIC4 for mimic4 base dataset...\n", + "Task cache paths: task_df=/tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_493ab707-d120-581c-a795-319ffd5c468e/task_df.ld, samples=/tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_493ab707-d120-581c-a795-319ffd5c468e/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n", + "Applying task transformations on data with 1 workers...\n", + "No cached event dataframe found. Creating: /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/global_event_df.parquet\n", + "Combining data from ehr dataset\n", + "Scanning table: diagnoses_icd from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv.gz\n", + "Joining with table: /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n", + "Scanning table: procedures_icd from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/procedures_icd.csv.gz\n", + "Joining with table: /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n", + "Scanning table: labevents from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/labevents.csv.gz\n", + "Joining with table: /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/d_labitems.csv.gz\n", + "Scanning table: patients from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/patients.csv.gz\n", + "Scanning table: admissions from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n", + "Scanning table: icustays from /u/wp14/PyHealth/test-resources/core/mimic4demo/icu/icustays.csv.gz\n", + "Combining data from note dataset\n", + "Scanning table: discharge from /u/wp14/PyHealth/test-resources/core/mimic4demo/note/discharge.csv.gz\n", + "Scanning table: radiology from /u/wp14/PyHealth/test-resources/core/mimic4demo/note/radiology.csv.gz\n", + "Creating combined dataframe\n", + "Caching event dataframe to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/global_event_df.parquet...\n", + "Detected Jupyter notebook environment, setting num_workers to 1\n", + "Single worker mode, processing sequentially\n", + "Worker 0 started processing 13 patients. (Polars threads: 6)\n", + "Worker 0 finished processing patients.\n", + "Fitting processors on the dataset...\n", + "Label mortality vocab: {0: 0, 1: 1}\n", + "Processing samples and saving to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_493ab707-d120-581c-a795-319ffd5c468e/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n", + "Applying processors on data with 1 workers...\n", + "Detected Jupyter notebook environment, setting num_workers to 1\n", + "Single worker mode, processing sequentially\n", + "Worker 0 started processing 13 samples. (0 to 13)\n", + "Worker 0 finished processing samples.\n", + "Cached processed samples to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_493ab707-d120-581c-a795-319ffd5c468e/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n" + ] + } + ], "source": [ "dataset = MIMIC4Dataset(\n", " ehr_root=EHR_ROOT,\n", @@ -493,10 +618,89 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 8, "id": "6b70acb2-bdea-4a43-8b9a-7725c34fcfce", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Patient ID: 10001\n", + "Mortality Flag: tensor([0.])\n", + "Found 13 unique patient IDs\n", + "\n", + "num_admissions: 3\n", + " hadm_id: 19999\n", + " admittime: 2150-02-15 08:00:00\n", + " dischtime: 2150-02-18 14:00:00\n", + " hadm_id: 20001\n", + " admittime: 2150-03-15 08:00:00\n", + " dischtime: 2150-03-18 14:00:00\n", + " hadm_id: 20002\n", + " admittime: 2150-06-20 10:30:00\n", + " dischtime: 2150-06-25 16:00:00\n", + "\n", + "discharge notes: 3\n", + " note_id: d1\n", + " hadm_id: 19999\n", + " charttime: 2150-02-18 12:00:00\n", + " storetime: 2150-02-18 13:00:00\n", + " text: Patient 10001 admitted for evaluation. Vitals stable. Labs reviewed. Discharged ..(Limited to 80 Characters).\n", + " note_id: d3\n", + " hadm_id: 20001\n", + " charttime: 2150-03-18 12:00:00\n", + " storetime: 2150-03-18 13:00:00\n", + " text: Patient 10001 admitted for evaluation. Vitals stable. Labs reviewed. Discharged ..(Limited to 80 Characters).\n", + " note_id: d4\n", + " hadm_id: 20002\n", + " charttime: 2150-06-25 12:00:00\n", + " storetime: 2150-06-25 13:00:00\n", + " text: Patient 10001 admitted for evaluation. Vitals stable. Labs reviewed. Discharged ..(Limited to 80 Characters).\n", + "\n", + "radiology notes: 3\n", + " note_id: r1\n", + " hadm_id: 19999\n", + " charttime: 2150-02-15 14:00:00\n", + " storetime: 2150-02-15 15:00:00\n", + " text: Chest X-ray for patient 10001. No acute cardiopulmonary process identified. Hear..(Limited to 80 Characters).\n", + " note_id: r3\n", + " hadm_id: 20001\n", + " charttime: 2150-03-15 14:00:00\n", + " storetime: 2150-03-15 15:00:00\n", + " text: Chest X-ray for patient 10001. No acute cardiopulmonary process identified. Hear..(Limited to 80 Characters).\n", + " note_id: r4\n", + " hadm_id: 20002\n", + " charttime: 2150-06-20 14:00:00\n", + " storetime: 2150-06-20 15:00:00\n", + " text: Chest X-ray for patient 10001. No acute cardiopulmonary process identified. Hear..(Limited to 80 Characters).\n", + "\n", + "lab events: 12\n", + " itemid: 50806\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 102.0\n", + " itemid: 50809\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 285.0\n", + " itemid: 50803\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 18.0\n", + " itemid: 50868\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 16.0\n", + " itemid: 50822\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 4.8\n", + " ... (7 more events)\n" + ] + } + ], "source": [ "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", "random.seed(SEED)\n", @@ -511,6 +715,11 @@ "print_patient_note_info(sample, note_type=\"discharge\")\n", "# Radiology Notes\n", "print_patient_note_info(sample, note_type=\"radiology\")\n", +<<<<<<< HEAD + "\n", + "# Labs\n", + "print_patient_lab_info(sample)" +======= "# Labs\n", "print_patient_lab_info(sample)\n", "# ICD Codes\n", @@ -525,6 +734,7 @@ "outputs": [], "source": [ "clear_cache_directory(CACHE_DIR)" +>>>>>>> origin/main ] }, { @@ -548,10 +758,33 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, "id": "16121916-e52e-4f66-8463-ff0f931cf416", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Setting task ClinicalNotesICDLabsMIMIC4 for mimic4 base dataset...\n", + "Task cache paths: task_df=/tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_6c674c02-3635-50c0-bdde-3a45cb1e6f31/task_df.ld, samples=/tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_6c674c02-3635-50c0-bdde-3a45cb1e6f31/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n", + "Applying task transformations on data with 1 workers...\n", + "Detected Jupyter notebook environment, setting num_workers to 1\n", + "Single worker mode, processing sequentially\n", + "Worker 0 started processing 13 patients. (Polars threads: 6)\n", + "Worker 0 finished processing patients.\n", + "Fitting processors on the dataset...\n", + "Label mortality vocab: {0: 0, 1: 1}\n", + "Processing samples and saving to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_6c674c02-3635-50c0-bdde-3a45cb1e6f31/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n", + "Applying processors on data with 1 workers...\n", + "Detected Jupyter notebook environment, setting num_workers to 1\n", + "Single worker mode, processing sequentially\n", + "Worker 0 started processing 13 samples. (0 to 13)\n", + "Worker 0 finished processing samples.\n", + "Cached processed samples to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_6c674c02-3635-50c0-bdde-3a45cb1e6f31/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n" + ] + } + ], "source": [ "dataset = MIMIC4Dataset(\n", " ehr_root=EHR_ROOT,\n", @@ -568,10 +801,38 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 10, "id": "42df85ab-2649-497b-92a2-415dc169798e", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Patient ID: 10001\n", + "Mortality Flag: tensor([0.])\n", + "Window start: 2150-02-15 08:00:00\n", + "Window end: 2150-02-15 20:00:00\n", + "\n", + "num_admissions: 1\n", + " hadm_id: 19999\n", + " admittime: 2150-02-15 08:00:00\n", + " dischtime: 2150-02-18 14:00:00\n", + "\n", + "discharge notes: 0\n", + "\n", + "radiology notes: 1\n", + " note_id: r1\n", + " hadm_id: 19999\n", + " charttime: 2150-02-15 14:00:00\n", + " storetime: 2150-02-15 15:00:00\n", + " text: Chest X-ray for patient 10001. No acute cardiopulmonary process identified. Hear..(Limited to 80 Characters).\n", + "\n", + "lab events: 0\n" + ] + } + ], "source": [ "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", "random.seed(SEED)\n", @@ -588,10 +849,32 @@ "print_patient_note_info(sample, note_type=\"discharge\")\n", "# Radiology Notes\n", "print_patient_note_info(sample, note_type=\"radiology\")\n", +<<<<<<< HEAD + "\n", + "# Labs\n", + "print_patient_lab_info(sample)" + ] + }, + { + "cell_type": "markdown", + "id": "a447c4e2-cfae-4e15-a2ac-0fac5754c1f2", + "metadata": {}, + "source": [ + "## 4. Cleanup" + ] + }, + { + "cell_type": "markdown", + "id": "32dc97fd-6d9a-4212-9045-f65a092dd9f5", + "metadata": {}, + "source": [ + "- Remove `CACHE_DIR`" +======= "# Labs\n", "print_patient_lab_info(sample)\n", "# ICD Codes\n", "print_patient_icd_info(sample)" +>>>>>>> origin/main ] }, { @@ -621,7 +904,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", +<<<<<<< HEAD + "version": "3.12.13" +======= "version": "3.13.3" +>>>>>>> origin/main } }, "nbformat": 4, diff --git a/pyhealth/data/data.py b/pyhealth/data/data.py index 14b1b526c..39dd1206f 100644 --- a/pyhealth/data/data.py +++ b/pyhealth/data/data.py @@ -155,7 +155,7 @@ def _filter_by_time_range_fast(self, df: pl.DataFrame, start: Optional[datetime] start_idx = np.searchsorted(ts_col, np.datetime64(start, "ms"), side="left") if end is not None: end_idx = np.searchsorted(ts_col, np.datetime64(end, "ms"), side="right") - return df.slice(start_idx, end_idx - start_idx) + return df.slice(start_idx, max(0, end_idx - start_idx)) def _filter_by_event_type_regular(self, df: pl.DataFrame, event_type: Optional[str]) -> pl.DataFrame: """Regular filtering by event type. Time complexity: O(n).""" diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index b21a37177..fd80f958a 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -3,7 +3,7 @@ import pickle from abc import ABC from pathlib import Path -from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable +from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable, Union import functools import operator from urllib.parse import urlparse, urlunparse @@ -314,7 +314,7 @@ class BaseDataset(ABC): dataset_name (str): Name of the dataset. config (dict): Configuration loaded from a YAML file. global_event_df (pl.LazyFrame): The global event data frame. - dev (bool): Whether to enable dev mode (limit to 1000 patients). + dev (Union[bool, int]): Whether to enable dev mode. If True, limit to 1000 patients. If an int, limit to that many patients. """ def __init__( @@ -325,7 +325,7 @@ def __init__( config_path: Optional[str] = None, cache_dir: str | Path | None = None, num_workers: int = 1, - dev: bool = False, + dev: Union[bool, int] = False, ): """Initializes the BaseDataset. @@ -342,7 +342,7 @@ def __init__( - **str** or **Path**: Used as the root cache directory path. A UUID is appended to the provided path to capture dataset configuration. num_workers (int): Number of worker processes for parallel operations. - dev (bool): Whether to run in dev mode (limits to 1000 patients). + dev (Union[bool, int]): Whether to run in dev mode. If True, limits to 1000 patients. If an int, limits to that many patients. """ if len(set(tables)) != len(tables): logger.warning("Duplicate table names in tables list. Removing duplicates.") @@ -509,8 +509,9 @@ def _event_transform(self, output_dir: Path) -> None: "PYHEALTH_DISABLE_DASK_DISTRIBUTED=1 detected; using local dask scheduler." ) if self.dev: - logger.info("Dev mode enabled: limiting to 1000 patients") - patients = df["patient_id"].unique().head(1000, compute=True).tolist() + n = 1000 if self.dev is True else int(self.dev) + logger.info(f"Dev mode enabled: limiting to {n} patients") + patients = df["patient_id"].unique().head(n, compute=True).tolist() patient_filter = df["patient_id"].isin(patients) df = df[patient_filter] @@ -530,8 +531,8 @@ def _event_transform(self, output_dir: Path) -> None: ) as cluster: with DaskClient(cluster) as client: if self.dev: - logger.info("Dev mode enabled: limiting to 1000 patients") - patients = df["patient_id"].unique().head(1000).tolist() + logger.info(f"Dev mode enabled: limiting to {1000 if self.dev is True else int(self.dev)} patients") + patients = df["patient_id"].unique().head(1000 if self.dev is True else int(self.dev)).tolist() filter = df["patient_id"].isin(patients) df = df[filter] diff --git a/pyhealth/datasets/configs/mimic4_ehr.yaml b/pyhealth/datasets/configs/mimic4_ehr.yaml index 84c570bb9..414f31308 100644 --- a/pyhealth/datasets/configs/mimic4_ehr.yaml +++ b/pyhealth/datasets/configs/mimic4_ehr.yaml @@ -108,6 +108,28 @@ tables: - "flag" - "storetime" + chartevents: + file_path: "icu/chartevents.csv.gz" + patient_id: "subject_id" + join: + - file_path: "icu/d_items.csv.gz" + "on": "itemid" + how: "inner" + columns: + - "label" + - "category" + timestamp: "charttime" + attributes: + - "hadm_id" + - "stay_id" + - "itemid" + - "label" + - "category" + - "value" + - "valuenum" + - "valueuom" + - "storetime" + hcpcsevents: file_path: "hosp/hcpcsevents.csv.gz" patient_id: "subject_id" diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py index 3f7623e67..089a6d2f3 100644 --- a/pyhealth/datasets/mimic4.py +++ b/pyhealth/datasets/mimic4.py @@ -1,7 +1,7 @@ import logging import os import warnings -from typing import List, Optional +from typing import List, Optional, Union import pandas as pd import dask.dataframe as dd @@ -378,7 +378,7 @@ def __init__( cxr_config_path: Optional[str] = None, cxr_variant: str = "default", dataset_name: str = "mimic4", - dev: bool = False, + dev: Union[bool, int] = False, cache_dir: Optional[str] = None, num_workers: int = 1, ): diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 2dbc94186..b2ea98854 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -22,6 +22,7 @@ def sample_balanced( ratio: float = 1.0, subsample: float = 1.0, seed: Optional[int] = None, + label_key: str = "label", ) -> SampleDataset: """Keep positives and negatives at a target ratio, then cap total size. @@ -32,6 +33,7 @@ def sample_balanced( exceeds ``len(dataset) * subsample``, both positives and negatives are downsampled proportionally while preserving the ratio as closely as possible. seed: Optional RNG seed for reproducible negative sampling. + label_key: Key to use for accessing the label field in each sample. Default ``"label"``. Returns: A new ``SampleDataset`` containing all positives plus sampled negatives, @@ -49,7 +51,7 @@ def sample_balanced( neg_indices: List[int] = [] for idx in range(len(dataset)): - label = _label_to_int(dataset[idx]["label"]) + label = _label_to_int(dataset[idx][label_key]) if label == 1: pos_indices.append(idx) else: diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index d70992e03..703f77630 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -206,9 +206,11 @@ def __init__( image_channels: int = 3, patch_size: int = 16, field_embeddings: Optional[dict[str, Any]] = None, + freeze_text_encoder: bool = False, ): super().__init__() self._embedding_dim = embedding_dim + self._freeze_text_encoder = freeze_text_encoder _field_embeddings = field_embeddings or {} self.encoders: nn.ModuleDict = nn.ModuleDict() @@ -237,7 +239,8 @@ def __init__( elif m == ModalityType.TEXT: self._build_text_encoder( - field_name, processor, pre_built, embedding_dim + field_name, processor, pre_built, embedding_dim, + freeze=freeze_text_encoder, ) elif m == ModalityType.IMAGE: @@ -307,6 +310,7 @@ def _build_text_encoder( processor: TemporalFeatureProcessor, pre_built: Any, embedding_dim: int, + freeze: bool = False, ) -> None: """Build TEXT encoder: BERT + projection, optionally from TextEmbeddingModel.""" @@ -330,6 +334,9 @@ def _set_projection( and hasattr(pre_built, "fc") ): self.encoders[field_name] = pre_built.transformer + if freeze: + for p in pre_built.transformer.parameters(): + p.requires_grad = False pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) _set_projection(pre_dim, pre_built.fc) return @@ -352,6 +359,9 @@ def _set_projection( shared_encoder = self.encoders[shared_field] else: shared_encoder = AutoModel.from_pretrained(tokenizer_model) + if freeze: + for p in shared_encoder.parameters(): + p.requires_grad = False self._shared_text_field_by_model[tokenizer_model] = field_name self.encoders[field_name] = shared_encoder diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index 29307d785..e98d98db6 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -47,6 +47,10 @@ from .multimodal_mimic4 import ( ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4, + ClinicalNotesICDLabsCXRMIMIC4, + ICDLabsMIMIC4, + LabsMIMIC4, + NotesLabsMIMIC4, ) from .patient_linkage import patient_linkage_mimic3_fn from .readmission_prediction import ( @@ -69,9 +73,4 @@ MutationPathogenicityPrediction, VariantClassificationClinVar, ) -from .multimodal_mimic4 import ( - ClinicalNotesMIMIC4, - ClinicalNotesICDLabsMIMIC4, - ClinicalNotesICDLabsCXRMIMIC4, -) from .patient_linkage_mimic3 import PatientLinkageMIMIC3Task diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index f8ef11ab1..71f8d3f87 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -1,3 +1,4 @@ +import re from datetime import datetime, timedelta from typing import Any, Dict, List, Optional, Union, Tuple, ClassVar @@ -45,17 +46,124 @@ class BaseMultimodalMIMIC4Task(BaseTask): item for itemids in LAB_CATEGORIES.values() for item in itemids ] + VITAL_CATEGORIES: ClassVar[Dict[str, List[str]]] = { + "HeartRate": ["220045", "220046", "220047"], + "SysBP": ["220050", "220179"], + "DiasBP": ["220051", "220180"], + "MeanBP": ["220052", "220181"], + "RespRate": ["220210", "220227"], + "SpO2": ["220277"], + "Temperature": ["223761", "223762"], + } + + VITAL_CATEGORY_NAMES: ClassVar[List[str]] = [ + "HeartRate", + "SysBP", + "DiasBP", + "MeanBP", + "RespRate", + "SpO2", + "Temperature", + ] + + VITALITEMS: ClassVar[List[str]] = [ + item for itemids in VITAL_CATEGORIES.values() for item in itemids + ] + def __init__( self, window_hours: Optional[float] = None, ): self.window_hours = window_hours + _ADMISSION_SECTION_TARGETS: ClassVar[frozenset] = frozenset({ + "chief complaint", + "history of present illness", + "hpi", + "past medical history", + "past medical and surgical history", + "past medical/surgical history", + "past surgical history", + "medications on admission", + "admission medications", + "home medications", + }) + + # Matches a line that is purely a section header: words + colon + nothing else. + # E.g. "Chief Complaint:", "Past Medical History:", "HPI:". + _SECTION_HEADER_RE: ClassVar[re.Pattern] = re.compile( + r"^\s*([A-Za-z][A-Za-z\s,/\-\.]{1,60}?)\s*:\s*$" + ) + @staticmethod def _clean_text(text: Optional[str]) -> Optional[str]: """Return text if non-empty, otherwise None.""" return text if text else None + @staticmethod + def _extract_admission_sections(text: str) -> str: + """Extract admission-context sections from a MIMIC-IV discharge note. + + Parses Chief Complaint, HPI, Past Medical History, and Medications on + Admission sections — information clinically available at admission time. + Falls back to the first 1 024 characters if no target sections are found. + """ + sections: Dict[str, str] = {} + current_key: Optional[str] = None + current_lines: List[str] = [] + + for line in text.split("\n"): + m = BaseMultimodalMIMIC4Task._SECTION_HEADER_RE.match(line) + if m: + if current_key is not None: + sections[current_key] = "\n".join(current_lines).strip() + current_key = m.group(1).strip().lower() + current_lines = [] + elif current_key is not None: + current_lines.append(line) + + if current_key is not None: + sections[current_key] = "\n".join(current_lines).strip() + + targets = BaseMultimodalMIMIC4Task._ADMISSION_SECTION_TARGETS + extracted = [v for k, v in sections.items() if k in targets and v] + return " [SEP] ".join(extracted) if extracted else text[:1024] + + def _collect_admission_note_sections( + self, + patient: Any, + hadm_id: Any, + admission_time: datetime, + ) -> Tuple[List[str], List[float]]: + """Collect admission-context text from the discharge note. + + No time filter is applied because the target sections (CC, HPI, PMH, + Medications on Admission) describe the patient's state *at admission* + regardless of when the note was finalised. Returned timestamps are 0.0 + so downstream models treat the text as admission-time context. + """ + notes = patient.get_events( + event_type="discharge", + filters=[("hadm_id", "==", hadm_id)], + ) + + texts: List[str] = [] + for note in notes: + try: + raw = note.text + if not raw: + continue + extracted = self._extract_admission_sections(raw) + if extracted: + texts.append(extracted) + except AttributeError: + pass + + if not texts: + return [self.MISSING_TEXT_TOKEN], [self.MISSING_FLOAT_TOKEN] + + return texts, [0.0] * len(texts) + @staticmethod def _parse_datetime(value: Any) -> Optional[datetime]: if isinstance(value, datetime): @@ -102,8 +210,12 @@ def _compute_effective_window( def _build_admissions_to_process(self, patient: Any) -> Tuple[List[Any], int]: """Build admissions to process and derive mortality label. - The task includes admissions until the first mortality event and labels - the sample as positive if death occurs in the current or next admission. + Includes all admissions up to and including the first death admission. + Patients who die in their first (and only) admission are included as + positives — previously they were dropped, which silently discarded most + ICU mortality positives and collapsed positive rate from ~10% to ~2.7%. + This now matches stagenet's semantics: use all available admission data + and label as positive if any admission has hospital_expire_flag=1. """ admissions = patient.get_events(event_type="admissions") if len(admissions) == 0: @@ -112,20 +224,12 @@ def _build_admissions_to_process(self, patient: Any) -> Tuple[List[Any], int]: admissions_to_process: List[Any] = [] mortality_label = 0 - for i, admission in enumerate(admissions): + for admission in admissions: + admissions_to_process.append(admission) if admission.hospital_expire_flag in [1, "1"]: mortality_label = 1 break - if i + 1 < len(admissions): - next_admission = admissions[i + 1] - if next_admission.hospital_expire_flag in [1, "1"]: - admissions_to_process.append(admission) - mortality_label = 1 - break - - admissions_to_process.append(admission) - return admissions_to_process, mortality_label def _collect_icd_codes(self, patient: Any, hadm_id: Any) -> List[str]: @@ -234,6 +338,85 @@ def _collect_labs( lab_times.append(self.MISSING_FLOAT_TOKEN) return lab_times, lab_values, lab_masks + def _collect_vitals( + self, + patient: Any, + admission_time: datetime, + end_time: datetime, + ) -> Tuple[List[float], List[List[float]], List[List[bool]]]: + """Collect vital sign values and observation masks for one admission. + + Returns: + Tuple of (vital_times, vital_values, vital_masks). + vital_masks is parallel boolean: True = observed, False = imputed 0.0. + Falls back to a single missing-placeholder row when no valid vitals are found. + """ + try: + import polars as pl + except ImportError as exc: + raise ImportError("Polars is required for vitals collection.") from exc + + chartevents_df = patient.get_events( + event_type="chartevents", + start=admission_time, + end=end_time, + return_df=True, + ) + + vital_times: List[float] = [] + vital_values: List[List[float]] = [] + vital_masks: List[List[bool]] = [] + + chartevents_df = chartevents_df.filter( + pl.col("chartevents/itemid").is_in(self.VITALITEMS) + ) + if chartevents_df.height > 0: + chartevents_df = chartevents_df.with_columns( + pl.col("chartevents/storetime").str.strptime( + pl.Datetime, "%Y-%m-%d %H:%M:%S" + ) + ) + chartevents_df = chartevents_df.filter( + pl.col("chartevents/storetime") <= end_time + ) + if chartevents_df.height > 0: + chartevents_df = chartevents_df.select( + pl.col("timestamp"), + pl.col("chartevents/itemid"), + pl.col("chartevents/valuenum").cast(pl.Float64), + ) + for vital_ts in sorted(chartevents_df["timestamp"].unique().to_list()): + ts_vitals = chartevents_df.filter(pl.col("timestamp") == vital_ts) + vital_vector: List[float] = [] + vital_mask: List[bool] = [] + for category_name in self.VITAL_CATEGORY_NAMES: + category_value = self.MISSING_FLOAT_TOKEN + observed = False + for itemid in self.VITAL_CATEGORIES[category_name]: + matching = ts_vitals.filter( + pl.col("chartevents/itemid") == itemid + ) + if matching.height > 0: + category_value = matching["chartevents/valuenum"][0] + observed = True + break + vital_vector.append(category_value) + vital_mask.append(observed) + vital_times.append( + self._to_hours((vital_ts - admission_time).total_seconds()) + ) + vital_values.append(vital_vector) + vital_masks.append(vital_mask) + + if len(vital_values) == 0: + vital_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.VITAL_CATEGORY_NAMES) + ) + vital_masks.append([False] * len(self.VITAL_CATEGORY_NAMES)) + vital_times.append(self.MISSING_FLOAT_TOKEN) + + return vital_times, vital_values, vital_masks + def _collect_notes( self, patient: Any, @@ -314,21 +497,21 @@ class ClinicalNotesMIMIC4(BaseMultimodalMIMIC4Task): task_name: str = "ClinicalNotesMIMIC4" input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { - "discharge_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", - "type_tag": "note", - }, - ), - "radiology_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", - "type_tag": "note", - }, - ) - } + "discharge_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "radiology_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + } output_schema: Dict[str, str] = {"mortality": "binary"} def __call__(self, patient: Any) -> List[Dict[str, Any]]: @@ -391,17 +574,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_radiology_times_from_admission, ) - if len(all_discharge_texts) == 0: - discharge_note_times_from_admission = ( - [self.MISSING_TEXT_TOKEN], - [self.MISSING_FLOAT_TOKEN], - ) - if len(all_radiology_texts) == 0: - radiology_note_times_from_admission = ( - [self.MISSING_TEXT_TOKEN], - [self.MISSING_FLOAT_TOKEN], - ) - single_patient_longitudinal_record = { "patient_id": patient.patient_id, "discharge_note_times": discharge_note_times_from_admission, @@ -451,37 +623,23 @@ class ClinicalNotesICDLabsMIMIC4(BaseMultimodalMIMIC4Task): task_name: str = "ClinicalNotesICDLabsMIMIC4" input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { - "discharge_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", - "type_tag": "note", - }, - ), - "radiology_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", - "type_tag": "note", - }, - ), - "icd_codes": ("stagenet", {"padding": PADDING}), - "labs": ("stagenet_tensor", {}), - "labs_mask": ("stagenet_tensor", {}), - } - output_schema: Dict[str, str] = {"mortality": "binary"} - - LAB_CATEGORIES: ClassVar[Dict[str, List[str]]] = { - "Sodium": ["50824", "52455", "50983", "52623"], - "Potassium": ["50822", "52452", "50971", "52610"], - "Chloride": ["50806", "52434", "50902", "52535"], - "Bicarbonate": ["50803", "50804"], - "Glucose": ["50809", "52027", "50931", "52569"], - "Calcium": ["50808", "51624"], - "Magnesium": ["50960"], - "Anion Gap": ["50868", "52500"], - "Osmolality": ["52031", "50964", "51701"], - "Phosphate": ["50970"], + "discharge_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "radiology_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "icd_codes": ("stagenet", {"padding": PADDING}), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), } output_schema: Dict[str, str] = {"mortality": "binary"} @@ -528,12 +686,10 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: admission.dischtime, "%Y-%m-%d %H:%M:%S" ) except (ValueError, AttributeError): - # Skip malformed admissions; note collectors are not called for these. - continue + admission_dischtime = admission_time if admission_dischtime < admission_time: - # Guard against invalid chronology in source records. - continue + admission_dischtime = admission_time discharge_texts, discharge_times = self._collect_notes( patient, @@ -569,7 +725,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_icd_codes.append(visit_icd_codes) all_icd_times.append(time_from_previous) else: # Add missingness token if there are no ICD diagnosis/inpatient procedure codes - all_icd_codes.append([self.MISSING_CODE_TOKEN]) + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) all_icd_times.append(self.MISSING_FLOAT_TOKEN) previous_admission_time = admission_time @@ -591,12 +747,195 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) all_lab_times.append(self.MISSING_FLOAT_TOKEN) - # If all admissions were skipped before ICD collection, ensure a - # single placeholder step so StageNetProcessor does not emit None time. + discharge_note_times_from_admission = ( + all_discharge_texts, + all_discharge_times_from_admission, + ) + radiology_note_times_from_admission = ( + all_radiology_texts, + all_radiology_times_from_admission, + ) + + single_patient_longitudinal_record = { + "patient_id": patient.patient_id, + "discharge_note_times": discharge_note_times_from_admission, + "radiology_note_times": radiology_note_times_from_admission, + "icd_codes": (all_icd_times, all_icd_codes), + "labs": (all_lab_times, all_lab_values), + "labs_mask": (all_lab_times, all_lab_masks), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } + + return [single_patient_longitudinal_record] + +class ClinicalNotesICDLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): + """Task combining notes, ICD, labs, and CXR for MIMIC-IV mortality. + + Adds temporally filtered CXR signals on top of ``ClinicalNotesICDLabsMIMIC4``: + - ``cxr_image_times``: ``(image_paths, hours_from_admission)`` processed by + ``TimeImageProcessor``. + + CXR filtering uses event timestamps from the metadata table. Since + timestamps are built from StudyDate+StudyTime in dataset configs, studytime + is naturally respected in temporal windows. + """ + + PADDING: int = 0 + + task_name: str = "ClinicalNotesICDLabsCXRMIMIC4" + input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { + "discharge_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "radiology_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "icd_codes": ("stagenet", {"padding": PADDING}), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), + "cxr_image_times": ( + "time_image", + { + "image_size": 224, + "mode": "RGB", + "padding": "", + }, + ), + } + output_schema: Dict[str, str] = {"mortality": "binary"} + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + demographics = patient.get_events(event_type="patients") + if not demographics: + return [] + + admissions_to_process, mortality_label = self._build_admissions_to_process( + patient + ) + + if len(admissions_to_process) == 0: + return [] + + effective_start, effective_end = self._compute_effective_window( + admissions_to_process + ) + + all_discharge_texts: List[str] = [] + all_discharge_times_from_admission: List[float] = [] + all_radiology_texts: List[str] = [] + all_radiology_times_from_admission: List[float] = [] + all_icd_codes: List[List[str]] = [] + all_icd_times: List[float] = [] + all_lab_values: List[List[float]] = [] + all_lab_masks: List[List[bool]] = [] + all_lab_times: List[float] = [] + all_cxr_paths: List[str] = [] + all_cxr_times: List[float] = [] + previous_admission_time = None + + for admission in admissions_to_process: + admission_time = admission.timestamp + + try: + admission_dischtime = datetime.strptime( + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + admission_dischtime = admission_time + + if admission_dischtime < admission_time: + admission_dischtime = admission_time + + discharge_texts, discharge_times = self._collect_notes( + patient, + "discharge", + admission.hadm_id, + admission_time, + start_time=effective_start, + end_time=effective_end, + ) + all_discharge_texts.extend(discharge_texts) + all_discharge_times_from_admission.extend(discharge_times) + + radiology_texts, radiology_times = self._collect_notes( + patient, + "radiology", + admission.hadm_id, + admission_time, + start_time=effective_start, + end_time=effective_end, + ) + all_radiology_texts.extend(radiology_texts) + all_radiology_times_from_admission.extend(radiology_times) + + visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) + if visit_icd_codes: + if previous_admission_time is None: + time_from_previous = 0.0 + else: + time_from_previous = self._to_hours( + (admission_time - previous_admission_time).total_seconds() + ) + all_icd_codes.append(visit_icd_codes) + all_icd_times.append(time_from_previous) + else: + all_icd_codes.append([self.MISSING_CODE_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + + previous_admission_time = admission_time + + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=admission_dischtime, + ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + # CXR metadata is filtered by timestamp; this includes StudyTime. + metadata_events = patient.get_events( + event_type="metadata", + start=admission_time, + end=admission_dischtime, + ) + for event in metadata_events: + try: + if event.image_path: + all_cxr_paths.append(event.image_path) + all_cxr_times.append( + self._to_hours( + (event.timestamp - admission_time).total_seconds() + ) + ) + except AttributeError: + continue + + if len(all_lab_values) == 0: + all_lab_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) + ) + all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) + all_lab_times.append(self.MISSING_FLOAT_TOKEN) + if len(all_icd_codes) == 0: all_icd_codes.append([self.MISSING_CODE_TOKEN]) all_icd_times.append(self.MISSING_FLOAT_TOKEN) + if len(all_cxr_paths) == 0: + all_cxr_paths = [self.MISSING_TEXT_TOKEN] + all_cxr_times = [self.MISSING_FLOAT_TOKEN] + discharge_note_times_from_admission = ( all_discharge_texts, all_discharge_times_from_admission, @@ -606,9 +945,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_radiology_times_from_admission, ) - # Per-admission note fallback happens inside _collect_notes(). - # This final guard handles the edge case where every admission was - # skipped before _collect_notes() was reached. if len(all_discharge_texts) == 0: discharge_note_times_from_admission = ( [self.MISSING_TEXT_TOKEN], @@ -627,6 +963,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: "icd_codes": (all_icd_times, all_icd_codes), "labs": (all_lab_times, all_lab_values), "labs_mask": (all_lab_times, all_lab_masks), + "cxr_image_times": (all_cxr_paths, all_cxr_times), "mortality": mortality_label, "window_start": effective_start, "window_end": effective_end, @@ -635,6 +972,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: return [single_patient_longitudinal_record] + class ICDLabsMIMIC4(BaseMultimodalMIMIC4Task): """Task for ICD codes + lab values mortality prediction using MIMIC-IV. @@ -697,10 +1035,9 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: admission.dischtime, "%Y-%m-%d %H:%M:%S" ) except (ValueError, AttributeError): - continue - + admission_dischtime = admission_time if admission_dischtime < admission_time: - continue + admission_dischtime = admission_time visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) if visit_icd_codes: @@ -713,7 +1050,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_icd_codes.append(visit_icd_codes) all_icd_times.append(time_from_previous) else: - all_icd_codes.append([self.MISSING_CODE_TOKEN]) + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) all_icd_times.append(self.MISSING_FLOAT_TOKEN) previous_admission_time = admission_time @@ -735,7 +1072,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_lab_times.append(self.MISSING_FLOAT_TOKEN) if len(all_icd_codes) == 0: - all_icd_codes.append([self.MISSING_CODE_TOKEN]) + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) all_icd_times.append(self.MISSING_FLOAT_TOKEN) single_patient_longitudinal_record = { @@ -751,215 +1088,276 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: return [single_patient_longitudinal_record] -class ClinicalNotesICDLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): - """Task combining notes, ICD, labs, and CXR for MIMIC-IV mortality. - - Adds temporally filtered CXR signals on top of ``ClinicalNotesICDLabsMIMIC4``: - - ``cxr_image_times``: ``(image_paths, hours_from_admission)`` processed by - ``TimeImageProcessor``. - - CXR filtering uses event timestamps from the metadata table. Since - timestamps are built from StudyDate+StudyTime in dataset configs, studytime - is naturally respected in temporal windows. +class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): + """Mortality prediction from admission-context notes and lab values. + + Follows the approach of Lee et al. (2023): use text that is clinically + available *at admission* rather than discharge notes. ICD codes and vitals + are excluded by default but can be re-enabled for ablation experiments. + + Text is extracted from the MIMIC-IV discharge note by parsing the Chief + Complaint, History of Present Illness, Past Medical History, and Medications + on Admission sections — all of which describe the patient's state at the + start of the stay. The extracted text is assigned timestamp 0.0. + + Fields: + admission_note_times: Admission-context note text at time 0.0. + labs: 10-dim lab vectors at each measurement timestamp. + labs_mask: Boolean observation mask parallel to ``labs``. + vitals: (only when ``include_vitals=True``) 7-dim vital-sign vectors at + each measurement timestamp. + vitals_mask: (only when ``include_vitals=True``) Boolean observation mask + parallel to ``vitals``. + icd_codes: (only when ``include_icd=True``) Diagnosis + procedure codes + per admission with inter-admission time offsets. + + Args: + window_hours: Hours from admission for lab/vital collection. ``None`` + collects for the full admission span. Default: 24. + include_icd: When ``True``, collect discharge-coded ICD codes and add + ``icd_codes`` to the sample dict / input schema. Default: ``False``. + include_vitals: When ``True``, collect ICU vital signs from chartevents + and add ``vitals`` and ``vitals_mask`` to the sample dict / input + schema. Default: ``False``. """ PADDING: int = 0 - task_name: str = "ClinicalNotesICDLabsCXRMIMIC4" - input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { - "discharge_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "bert-base-uncased", - "type_tag": "note", - }, - ), - "radiology_note_times": ( + task_name: str = "NotesLabsMIMIC4" + + _BASE_INPUT_SCHEMA: ClassVar[Dict] = { + "admission_note_times": ( "tuple_time_text", { - "tokenizer_model": "bert-base-uncased", + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", "type_tag": "note", }, ), - "icd_codes": ("stagenet", {"padding": PADDING}), "labs": ("stagenet_tensor", {}), "labs_mask": ("stagenet_tensor", {}), - "cxr_image_times": ( - "time_image", - { - "image_size": 224, - "mode": "RGB", - "padding": "", - }, - ), } + + input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = _BASE_INPUT_SCHEMA output_schema: Dict[str, str] = {"mortality": "binary"} + def __init__( + self, + window_hours: Optional[float] = None, + include_icd: bool = False, + include_vitals: bool = False, + ) -> None: + super().__init__(window_hours=window_hours) + self.include_icd = include_icd + self.include_vitals = include_vitals + schema = dict(self._BASE_INPUT_SCHEMA) + if include_vitals: + schema["vitals"] = ("stagenet_tensor", {}) + schema["vitals_mask"] = ("stagenet_tensor", {}) + if include_icd: + schema["icd_codes"] = ("stagenet", {"padding": self.PADDING}) + self.input_schema = schema + def __call__(self, patient: Any) -> List[Dict[str, Any]]: - demographics = patient.get_events(event_type="patients") - if not demographics: + if not patient.get_events(event_type="patients"): return [] admissions_to_process, mortality_label = self._build_admissions_to_process( patient ) - if len(admissions_to_process) == 0: + if not admissions_to_process: return [] effective_start, effective_end = self._compute_effective_window( admissions_to_process ) - all_discharge_texts: List[str] = [] - all_discharge_times_from_admission: List[float] = [] - all_radiology_texts: List[str] = [] - all_radiology_times_from_admission: List[float] = [] - all_icd_codes: List[List[str]] = [] - all_icd_times: List[float] = [] + all_note_texts: List[str] = [] + all_note_times: List[float] = [] all_lab_values: List[List[float]] = [] all_lab_masks: List[List[bool]] = [] all_lab_times: List[float] = [] - all_cxr_paths: List[str] = [] - all_cxr_times: List[float] = [] + all_vital_values: List[List[float]] = [] + all_vital_masks: List[List[bool]] = [] + all_vital_times: List[float] = [] + all_icd_codes: List[List[str]] = [] + all_icd_times: List[float] = [] previous_admission_time = None for admission in admissions_to_process: admission_time = admission.timestamp - # Skip admissions that start at or after the observation window closes, prevents Polars searchsorted OverflowError. - if effective_end is not None and admission_time >= effective_end: - continue - try: admission_dischtime = datetime.strptime( admission.dischtime, "%Y-%m-%d %H:%M:%S" ) except (ValueError, AttributeError): - # Skip malformed admissions; note collectors are not called for these. - continue - + admission_dischtime = admission_time if admission_dischtime < admission_time: - # Guard against invalid chronology in source records. - continue - - admission_end = admission_dischtime - if effective_end is not None and effective_end < admission_end: - admission_end = effective_end + admission_dischtime = admission_time - discharge_texts, discharge_times = self._collect_notes( - patient, - "discharge", - admission.hadm_id, - admission_time, - start_time=effective_start, - end_time=admission_end, + # Admission-context sections — no time window applied to notes + note_texts, note_times = self._collect_admission_note_sections( + patient, admission.hadm_id, admission_time ) - all_discharge_texts.extend(discharge_texts) - all_discharge_times_from_admission.extend(discharge_times) - - radiology_texts, radiology_times = self._collect_notes( - patient, - "radiology", - admission.hadm_id, - admission_time, - start_time=effective_start, - end_time=admission_end, + all_note_texts.extend(note_texts) + all_note_times.extend(note_times) + + # Labs within the observation window + lab_end = ( + effective_end + if self.window_hours is not None + else admission_dischtime ) - all_radiology_texts.extend(radiology_texts) - all_radiology_times_from_admission.extend(radiology_times) - - visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) - if visit_icd_codes: - if previous_admission_time is None: - time_from_previous = 0.0 - else: - time_from_previous = self._to_hours( - (admission_time - previous_admission_time).total_seconds() - ) - all_icd_codes.append(visit_icd_codes) - all_icd_times.append(time_from_previous) - else: - all_icd_codes.append([self.MISSING_CODE_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) - - previous_admission_time = admission_time - lab_times, lab_values, lab_masks = self._collect_labs( patient=patient, admission_time=admission_time, - end_time=admission_dischtime, + end_time=lab_end, ) all_lab_times.extend(lab_times) all_lab_values.extend(lab_values) all_lab_masks.extend(lab_masks) - # CXR metadata is filtered by timestamp; this includes StudyTime. - metadata_events = patient.get_events( - event_type="metadata", - start=admission_time, - end=admission_end, - ) - for event in metadata_events: - try: - if event.image_path: - all_cxr_paths.append(event.image_path) - all_cxr_times.append( - self._to_hours( - (event.timestamp - admission_time).total_seconds() - ) - ) - except AttributeError: - continue + # Vitals within the observation window + if self.include_vitals: + vital_times, vital_values, vital_masks = self._collect_vitals( + patient=patient, + admission_time=admission_time, + end_time=lab_end, + ) + all_vital_times.extend(vital_times) + all_vital_values.extend(vital_values) + all_vital_masks.extend(vital_masks) + + if self.include_icd: + visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) + time_from_previous = ( + 0.0 + if previous_admission_time is None + else self._to_hours( + (admission_time - previous_admission_time).total_seconds() + ) + ) + if visit_icd_codes: + all_icd_codes.append(visit_icd_codes) + all_icd_times.append(time_from_previous) + else: + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + previous_admission_time = admission_time - if len(all_lab_values) == 0: + if not all_lab_values: all_lab_values.append( [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) ) all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) all_lab_times.append(self.MISSING_FLOAT_TOKEN) - # If all admissions were skipped before ICD collection, ensure a - # single placeholder step so StageNetProcessor does not emit None time. - if len(all_icd_codes) == 0: - all_icd_codes.append([self.MISSING_CODE_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) + if self.include_vitals and not all_vital_values: + all_vital_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.VITAL_CATEGORY_NAMES) + ) + all_vital_masks.append([False] * len(self.VITAL_CATEGORY_NAMES)) + all_vital_times.append(self.MISSING_FLOAT_TOKEN) - # time_image processor expects at least one path/time pair. - if len(all_cxr_paths) == 0: - all_cxr_paths = [self.MISSING_TEXT_TOKEN] - all_cxr_times = [self.MISSING_FLOAT_TOKEN] + record: Dict[str, Any] = { + "patient_id": patient.patient_id, + "admission_note_times": (all_note_texts, all_note_times), + "labs": (all_lab_times, all_lab_values), + "labs_mask": (all_lab_times, all_lab_masks), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } - discharge_note_times_from_admission = ( - all_discharge_texts, - all_discharge_times_from_admission, + if self.include_vitals: + record["vitals"] = (all_vital_times, all_vital_values) + record["vitals_mask"] = (all_vital_times, all_vital_masks) + + if self.include_icd: + if not all_icd_codes: + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + record["icd_codes"] = (all_icd_times, all_icd_codes) + + return [record] + + +class LabsMIMIC4(BaseMultimodalMIMIC4Task): + """EHR-only mortality prediction using lab values — no notes, no ICD codes. + + Serves as the structured-EHR reference baseline for multimodal ablations. + Collecting only ``labevents`` keeps the dataset loader fast and avoids any + leakage from discharge-coded ICD tables. + + Schema mirrors the ``labs`` / ``labs_mask`` fields from ``NotesLabsMIMIC4`` + so the same backbone models (MLP, RNN, Transformer, etc.) work unchanged. + + Args: + window_hours: Hours from admission to collect lab measurements. + ``None`` collects for the full admission span. Default: 24. + """ + + PADDING: int = 0 + + task_name: str = "LabsMIMIC4" + + input_schema: ClassVar[Dict] = { + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), + } + output_schema: ClassVar[Dict] = {"mortality": "binary"} + + def __init__(self, window_hours: Optional[float] = 24) -> None: + super().__init__() + self.window_hours = window_hours + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[override] + admissions_to_process, mortality_label = self._build_admissions_to_process( + patient ) - radiology_note_times_from_admission = ( - all_radiology_texts, - all_radiology_times_from_admission, + if not admissions_to_process: + return [] + + effective_start, effective_end = self._compute_effective_window( + admissions_to_process ) - # Per-admission note fallback happens inside _collect_notes(). - # This final guard handles the edge case where every admission was - # skipped before _collect_notes() was reached. - if len(all_discharge_texts) == 0: - discharge_note_times_from_admission = ( - [self.MISSING_TEXT_TOKEN], - [self.MISSING_FLOAT_TOKEN], + all_lab_times: List[float] = [] + all_lab_values: List[List[float]] = [] + all_lab_masks: List[List[bool]] = [] + + for admission in admissions_to_process: + admission_time = admission.timestamp + + try: + admission_dischtime = datetime.strptime( + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + admission_dischtime = admission_time + if admission_dischtime < admission_time: + admission_dischtime = admission_time + + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=admission_dischtime, ) - if len(all_radiology_texts) == 0: - radiology_note_times_from_admission = ( - [self.MISSING_TEXT_TOKEN], - [self.MISSING_FLOAT_TOKEN], + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + if len(all_lab_values) == 0: + all_lab_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) ) + all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) + all_lab_times.append(self.MISSING_FLOAT_TOKEN) single_patient_longitudinal_record = { "patient_id": patient.patient_id, - "discharge_note_times": discharge_note_times_from_admission, - "radiology_note_times": radiology_note_times_from_admission, - "icd_codes": (all_icd_times, all_icd_codes), "labs": (all_lab_times, all_lab_values), "labs_mask": (all_lab_times, all_lab_masks), - "cxr_image_times": (all_cxr_paths, all_cxr_times), "mortality": mortality_label, "window_start": effective_start, "window_end": effective_end, diff --git a/scripts/compute_token_stats.py b/scripts/compute_token_stats.py new file mode 100644 index 000000000..4430a874a --- /dev/null +++ b/scripts/compute_token_stats.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +"""Compute per-modality token counts and missing-token rates for Table 2. + +Iterates patients directly via dataset.iter_patients() and calls the task +function on each one — never materializes all samples in memory, no litdata +write, no GPU needed. + +Usage (on CC): + python scripts/compute_token_stats.py \ + --ehr-root /projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2 \ + --note-root /projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-note \ + --cache-dir /u/rianatri/pyhealth_cache \ + --task notes_labs +""" +import argparse +import os + +import numpy as np +from tqdm import tqdm + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--ehr-root", required=True) + p.add_argument("--note-root", required=True) + p.add_argument("--cache-dir", default=None) + p.add_argument("--dev", action="store_true", help="Limit to 1000 patients") + p.add_argument( + "--task", + type=str, + choices=["clinical_notes_icd_labs", "notes_labs"], + default="notes_labs", + help=( + "Task to profile. 'notes_labs' (default) uses admission-context " + "text sections; 'clinical_notes_icd_labs' profiles the legacy task." + ), + ) + p.add_argument( + "--icd-codes", + action="store_true", + default=False, + help="Include discharge-coded ICD codes when using --task notes_labs.", + ) + return p.parse_args() + + +def main(): + args = parse_args() + + os.environ.setdefault("PYHEALTH_DISABLE_DASK_DISTRIBUTED", "1") + + from pyhealth.datasets import MIMIC4Dataset + from pyhealth.tasks.multimodal_mimic4 import ( + ClinicalNotesICDLabsMIMIC4, + NotesLabsMIMIC4, + ) + + print("Building dataset (uses cache if available)...") + + if args.task == "clinical_notes_icd_labs": + ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] + note_tables = ["discharge", "radiology"] + task = ClinicalNotesICDLabsMIMIC4() + else: # notes_labs + ehr_tables = ( + ["diagnoses_icd", "procedures_icd", "labevents"] + if args.icd_codes + else ["labevents"] + ) + note_tables = ["discharge"] + task = NotesLabsMIMIC4(window_hours=24, include_icd=args.icd_codes) + + kwargs = dict( + ehr_root=args.ehr_root, + note_root=args.note_root, + ehr_tables=ehr_tables, + note_tables=note_tables, + dev=args.dev, + ) + if args.cache_dir: + kwargs["cache_dir"] = args.cache_dir + + dataset = MIMIC4Dataset(**kwargs) + + MISSING_TEXT = "" + LAB_CATEGORIES = task.LAB_CATEGORY_NAMES # 10 names + + # Counters + n_patients = n_samples = 0 + note_total = note_missing = 0 + note_empty_extracted = 0 # notes present but section extraction returned nothing + icd_total_visits = icd_missing_visits = 0 + icd_total_codes = 0 + lab_total_timesteps = lab_missing_timesteps = 0 + lab_per_cat_missing = np.zeros(len(LAB_CATEGORIES), dtype=np.int64) + + print("Iterating patients and accumulating stats (no litdata write)...") + for patient in tqdm(dataset.iter_patients(), total=len(dataset.unique_patient_ids)): + n_patients += 1 + samples = task(patient) + if not samples: + continue + for s in samples: + n_samples += 1 + + # ── notes ──────────────────────────────────────────── + if args.task == "clinical_notes_icd_labs": + disc_texts, _ = s["discharge_note_times"] + note_total += len(disc_texts) + note_missing += sum(1 for t in disc_texts if t == MISSING_TEXT) + + rad_texts, _ = s["radiology_note_times"] + note_total += len(rad_texts) + note_missing += sum(1 for t in rad_texts if t == MISSING_TEXT) + else: + note_texts, _ = s["admission_note_times"] + note_total += len(note_texts) + for t in note_texts: + if t == MISSING_TEXT: + note_missing += 1 + elif len(t) <= 1024 and t == t[:1024]: + # Heuristic: if the note is exactly 1024 chars, it likely + # came from the fallback raw-note path (no sections found). + note_empty_extracted += 1 + + # ── ICD codes ──────────────────────────────────────── + if "icd_codes" in s: + _, icd_visits = s["icd_codes"] + icd_total_visits += len(icd_visits) + for visit_codes in icd_visits: + if visit_codes == [MISSING_TEXT]: + icd_missing_visits += 1 + else: + icd_total_codes += len(visit_codes) + + # ── labs ───────────────────────────────────────────── + _, lab_masks = s["labs_mask"] + for mask_row in lab_masks: + lab_total_timesteps += 1 + if not any(mask_row): + lab_missing_timesteps += 1 + for i, observed in enumerate(mask_row): + if not observed: + lab_per_cat_missing[i] += 1 + + def pct(n, d): + return 100.0 * n / d if d else float("nan") + + print("\n" + "=" * 60) + print(f"TOKEN STATS — {task.task_name}") + print("=" * 60) + print(f" Patients processed : {n_patients:>8,}") + print(f" Samples (patients) : {n_samples:>8,}") + + if args.task == "clinical_notes_icd_labs": + print("\n── Notes (discharge + radiology) ───────────────────────────") + else: + print("\n── Admission-context notes ─────────────────────────────────") + print(f" Total note tokens : {note_total:>8,}") + print(f" Missing (empty str) : {note_missing:>8,} ({pct(note_missing, note_total):.1f}%)") + if args.task == "notes_labs": + print( + f" Fallback (raw prefix) : {note_empty_extracted:>8,} " + f"({pct(note_empty_extracted, note_total):.1f}%)" + ) + print( + f" Effective coverage : {note_total - note_missing:>8,} " + f"({pct(note_total - note_missing, note_total):.1f}%)" + ) + + if "icd_codes" in s: + print("\n── ICD Codes ────────────────────────────────────────────────") + print(f" Total visit tokens : {icd_total_visits:>8,}") + print( + f' Missing visits [""] : {icd_missing_visits:>8,} ' + f"({pct(icd_missing_visits, icd_total_visits):.1f}%)" + ) + print(f" Total ICD codes : {icd_total_codes:>8,} (in non-missing visits)") + + print("\n── Labs ─────────────────────────────────────────────────────") + print(f" Total timestep tokens : {lab_total_timesteps:>8,}") + print( + f" Missing timesteps : {lab_missing_timesteps:>8,} " + f"({pct(lab_missing_timesteps, lab_total_timesteps):.1f}%)" + ) + print(f"\n Per-category missingness (across all timesteps):") + for i, cat in enumerate(LAB_CATEGORIES): + n_miss = int(lab_per_cat_missing[i]) + print( + f" {cat:<15}: {n_miss:>8,} / {lab_total_timesteps:>8,} " + f"({pct(n_miss, lab_total_timesteps):.1f}% missing)" + ) + + print("\n── Summary ──────────────────────────────────────────────────") + total_tokens = note_total + icd_total_visits + lab_total_timesteps + total_missing = note_missing + icd_missing_visits + lab_missing_timesteps + print(f" All modalities total tokens : {total_tokens:>8,}") + print( + f" All modalities missing tokens: {total_missing:>8,} " + f"({pct(total_missing, total_tokens):.1f}%)" + ) + print() + + +if __name__ == "__main__": + main() diff --git a/scripts/condor/batch_ablation.sub b/scripts/condor/batch_ablation.sub new file mode 100644 index 000000000..04c650b64 --- /dev/null +++ b/scripts/condor/batch_ablation.sub @@ -0,0 +1,100 @@ +# HTCondor submission — PyHealth Ablation Runs +# +# Four conditions on devsplit (1000 patients), MLP, 100 epochs: +# 1. baseline: notes+labs only (ICD off, vitals off, balanced off) +# 2. icd_on: notes+labs+ICD codes (leakage ablation) +# 3. vitals_on: notes+labs+vitals (chartevents ablation) +# 4. balanced: notes+labs with balanced sampling (1:1 pos:neg) +# +# All use MLP backbone, seed=42, embedding_dim=128, patience=10. +# Jobs start with "PyHealth" for easy queue filtering. +# +# To submit from the project root on the submit host: +# cd /home/rianatri/PyHealth && condor_submit scripts/condor/batch_ablation.sub + +initialdir = /home/rianatri/PyHealth +executable = /home/rianatri/PyHealth/scripts/condor/run_table2.sh +transfer_executable = False +arguments = $(model) $(seed) +getenv = True + ++JobBatchName = "PyHealth" + +output = logs/condor/ablation_$(ClusterId)_$(Process)_$(model)_seed$(seed).out +error = logs/condor/ablation_$(ClusterId)_$(Process)_$(model)_seed$(seed).err +log = logs/condor/ablation_$(ClusterId)_$(Process)_$(model)_seed$(seed).log + +request_gpus = 1 +request_cpus = 4 +request_memory = 48000 + +Rank = TARGET.CUDAGlobalMemoryMb + +# ── Condition 1: baseline (notes+labs, no ICD, no vitals, no balancing) ──────── +queue model, seed, environment from ( + mlp, 42, "TABLE2_TASK=notes_labs \ + TABLE2_RUN_LABEL=PyHealth_baseline \ + TABLE2_DEV_MODE=1 \ + TABLE2_DEV_COUNT=1000 \ + TABLE2_EPOCHS=100 \ + TABLE2_NUM_WORKERS=2 \ + TABLE2_FREEZE_ENCODER=0 \ + TABLE2_ICD_CODES=0 \ + TABLE2_INCLUDE_VITALS=0 \ + TABLE2_BALANCED_SAMPLING=0 \ + TABLE2_OUTPUT_DIR=output/ablation_baseline \ + TABLE2_WINDOW_HOURS=24 \ + TABLE2_PATIENCE=10" +) + +# ── Condition 2: icd_on (notes+labs+ICD codes — discharge-coded leakage) ────── +queue model, seed, environment from ( + mlp, 42, "TABLE2_TASK=notes_labs \ + TABLE2_RUN_LABEL=PyHealth_icd_on \ + TABLE2_DEV_MODE=1 \ + TABLE2_DEV_COUNT=1000 \ + TABLE2_EPOCHS=100 \ + TABLE2_NUM_WORKERS=2 \ + TABLE2_FREEZE_ENCODER=0 \ + TABLE2_ICD_CODES=1 \ + TABLE2_INCLUDE_VITALS=0 \ + TABLE2_BALANCED_SAMPLING=0 \ + TABLE2_OUTPUT_DIR=output/ablation_icd_on \ + TABLE2_WINDOW_HOURS=24 \ + TABLE2_PATIENCE=10" +) + +# ── Condition 3: vitals_on (notes+labs+vitals from chartevents) ─────────────── +queue model, seed, environment from ( + mlp, 42, "TABLE2_TASK=notes_labs \ + TABLE2_RUN_LABEL=PyHealth_vitals_on \ + TABLE2_DEV_MODE=1 \ + TABLE2_DEV_COUNT=1000 \ + TABLE2_EPOCHS=100 \ + TABLE2_NUM_WORKERS=2 \ + TABLE2_FREEZE_ENCODER=0 \ + TABLE2_ICD_CODES=0 \ + TABLE2_INCLUDE_VITALS=1 \ + TABLE2_BALANCED_SAMPLING=0 \ + TABLE2_OUTPUT_DIR=output/ablation_vitals_on \ + TABLE2_WINDOW_HOURS=24 \ + TABLE2_PATIENCE=10" +) + +# ── Condition 4: balanced (notes+labs with 1:1 pos:neg training sampling) ───── +queue model, seed, environment from ( + mlp, 42, "TABLE2_TASK=notes_labs \ + TABLE2_RUN_LABEL=PyHealth_balanced \ + TABLE2_DEV_MODE=1 \ + TABLE2_DEV_COUNT=1000 \ + TABLE2_EPOCHS=100 \ + TABLE2_NUM_WORKERS=2 \ + TABLE2_FREEZE_ENCODER=0 \ + TABLE2_ICD_CODES=0 \ + TABLE2_INCLUDE_VITALS=0 \ + TABLE2_BALANCED_SAMPLING=1 \ + TABLE2_BALANCED_RATIO=1.0 \ + TABLE2_OUTPUT_DIR=output/ablation_balanced \ + TABLE2_WINDOW_HOURS=24 \ + TABLE2_PATIENCE=10" +) \ No newline at end of file diff --git a/scripts/condor/batch_noteslabs_poc.sub b/scripts/condor/batch_noteslabs_poc.sub new file mode 100644 index 000000000..8ae75df70 --- /dev/null +++ b/scripts/condor/batch_noteslabs_poc.sub @@ -0,0 +1,45 @@ +# HTCondor submission — notes_labs POC (single seed, lightweight params) +# +# Purpose: proof-of-concept run for the NotesLabsMIMIC4 task (admission-context +# note sections + StageNet labs, no ICD leakage). Single seed; no multi-seed +# coverage needed for a POC. Models chosen for low VRAM: mlp, rnn, +# bottleneck_transformer. BERT encoder is frozen (--freeze-encoder) to halve +# BERT VRAM and keep each job well under 24 GB (A10 / A40 compatible). +# +# Ablation variants: +# noteslabs_poc — recommended: notes + labs, no ICD codes +# noteslabs_icd_poc — ablation: notes + labs + discharge-coded ICD codes +# +# To submit: +# condor_submit scripts/condor/batch_noteslabs_poc.sub + +executable = scripts/condor/run_table2.sh +arguments = $(model) $(seed) +getenv = True + +# notes_labs task, frozen encoder, ablation flag controlled per queue block +environment = "TABLE2_TASK=notes_labs \ + TABLE2_OUTPUT_DIR=/home/rianatri/noteslabs_poc \ + TABLE2_RUN_LABEL=noteslabs_poc \ + TABLE2_EPOCHS=10 \ + TABLE2_FREEZE_ENCODER=1 \ + TABLE2_WINDOW_HOURS=24" + +output = logs/condor/noteslabs_poc_$(ClusterId)_$(Process)_$(model)_seed$(seed).out +error = logs/condor/noteslabs_poc_$(ClusterId)_$(Process)_$(model)_seed$(seed).err +log = logs/condor/noteslabs_poc_$(ClusterId)_$(Process)_$(model)_seed$(seed).log + +# 24 GB covers A10/A40 with frozen BERT + embedding_dim=64. +# For bottleneck_transformer bump to 40 GB if needed on smaller slots. +request_gpus = 1 +request_cpus = 4 +request_memory = 48000 + +Rank = TARGET.CUDAGlobalMemoryMb + +# ── noteslabs_poc: notes + labs, no ICD codes ─────────────────────────────── +queue model, seed from ( + mlp, 42 + rnn, 42 + bottleneck_transformer, 42 +) diff --git a/scripts/condor/run_table2.sh b/scripts/condor/run_table2.sh index d211fa2ba..09d2a5d57 100755 --- a/scripts/condor/run_table2.sh +++ b/scripts/condor/run_table2.sh @@ -1,14 +1,28 @@ #!/usr/bin/env bash set -euo pipefail -MODEL="${1:?usage: run_table2.sh }" +MODEL="${1:?usage: run_table2.sh }" +SEED="${2:?usage: run_table2.sh }" CONDA_ENV="${CONDA_ENV:-pyhealth2}" PROJECT_DIR="${PROJECT_DIR:-/home/rianatri/PyHealth}" EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}" NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}" -CACHE_DIR="${CACHE_DIR:-/shared/eng/pyhealth_cache}" +CACHE_DIR="${CACHE_DIR:-/home/rianatri/pyhealth_cache}" CONDA_SH="${CONDA_SH:-}" +PYHEALTH_DISABLE_DASK_DISTRIBUTED="${PYHEALTH_DISABLE_DASK_DISTRIBUTED:-1}" +TABLE2_EPOCHS="${TABLE2_EPOCHS:-20}" +TABLE2_NUM_WORKERS="${TABLE2_NUM_WORKERS:-2}" +TABLE2_DEV_MODE="${TABLE2_DEV_MODE:-0}" +TABLE2_TASK="${TABLE2_TASK:-clinical_notes_icd_labs}" +TABLE2_WINDOW_HOURS="${TABLE2_WINDOW_HOURS:-24}" +TABLE2_OUTPUT_DIR="${TABLE2_OUTPUT_DIR:-output/table2}" +TABLE2_RUN_LABEL="${TABLE2_RUN_LABEL:-full}" +TABLE2_FREEZE_ENCODER="${TABLE2_FREEZE_ENCODER:-0}" +TABLE2_ICD_CODES="${TABLE2_ICD_CODES:-0}" +TABLE2_INCLUDE_VITALS="${TABLE2_INCLUDE_VITALS:-0}" +TABLE2_BALANCED_SAMPLING="${TABLE2_BALANCED_SAMPLING:-0}" +TABLE2_BALANCED_RATIO="${TABLE2_BALANCED_RATIO:-1.0}" resolve_conda_sh() { if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then @@ -79,9 +93,7 @@ conda activate "${CONDA_ENV}" cd "${PROJECT_DIR}" -# Isolate per-job cache/temp paths so concurrent Condor jobs do not race -# on the same tmp directory tree. -JOB_TAG="${MODEL}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}" +JOB_TAG="${MODEL}_seed${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}" JOB_CACHE_DIR="${CACHE_DIR}/${JOB_TAG}" mkdir -p "${JOB_CACHE_DIR}" @@ -93,6 +105,8 @@ fi mkdir -p "${DASK_TEMPORARY_DIRECTORY}" # Ensure local repo package is importable even if not installed into env. export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}" +export PYHEALTH_DISABLE_DASK_DISTRIBUTED +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True,max_split_size_mb:256}" if ! python -c "import pyhealth" >/dev/null 2>&1; then echo "ERROR: pyhealth is not importable. Run: bash setup.sh" >&2 @@ -100,76 +114,136 @@ if ! python -c "import pyhealth" >/dev/null 2>&1; then fi echo "========================================================" -echo " Table 2 run | model=${MODEL}" +echo " Table 2 run | label=${TABLE2_RUN_LABEL}" +echo " Model : ${MODEL}" +echo " Seed : ${SEED}" echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" echo " Conda env: ${CONDA_ENV}" echo " Conda sh : ${CONDA_SH}" echo " EHR root : ${EHR_ROOT}" echo " Note root: ${NOTE_ROOT}" echo " Cache dir: ${CACHE_DIR}" -echo " Job cache: ${JOB_CACHE_DIR}" +echo " Job cache : ${JOB_CACHE_DIR}" echo " Dask temp: ${DASK_TEMPORARY_DIRECTORY}" +echo " Dask dist: ${PYHEALTH_DISABLE_DASK_DISTRIBUTED} (1=local scheduler)" +echo " Epochs : ${TABLE2_EPOCHS}" +echo " Workers : ${TABLE2_NUM_WORKERS}" +echo " Dev mode : ${TABLE2_DEV_MODE}" +echo " Task : ${TABLE2_TASK}" +echo " Window : ${TABLE2_WINDOW_HOURS}h" +echo " Output dir: ${TABLE2_OUTPUT_DIR}" +echo " Patience : 5 (early stopping)" +echo " Freeze enc: ${TABLE2_FREEZE_ENCODER} (1=freeze BERT)" +echo " ICD codes : ${TABLE2_ICD_CODES} (1=include, ablation only)" +echo " Vitals : ${TABLE2_INCLUDE_VITALS} (1=include chartevents)" +echo " Balanced : ${TABLE2_BALANCED_SAMPLING} (1=undersample negatives)" +echo " Bal ratio : ${TABLE2_BALANCED_RATIO}" echo "========================================================" COMMON=( --ehr-root "${EHR_ROOT}" --note-root "${NOTE_ROOT}" - --cache-dir "${JOB_CACHE_DIR}" - --task clinical_notes_icd_labs + --cache-dir "${CACHE_DIR}" + --task "${TABLE2_TASK}" + --observation-window-hours "${TABLE2_WINDOW_HOURS}" --model "${MODEL}" --embedding-dim 128 --hidden-dim 128 --heads 4 --num-layers 2 --dropout 0.1 - --epochs 30 - --batch-size 32 + --epochs "${TABLE2_EPOCHS}" + --batch-size 16 --weight-decay 1e-5 - --num-workers 4 - --output-dir output/table2 + --num-workers "${TABLE2_NUM_WORKERS}" + --output-dir "${TABLE2_OUTPUT_DIR}" + --patience "${TABLE2_PATIENCE:-5}" ) +if [[ "${TABLE2_DEV_MODE}" == "1" ]]; then + COMMON+=(--dev "${TABLE2_DEV_COUNT:-1000}") +fi + +if [[ "${TABLE2_FREEZE_ENCODER}" == "1" ]]; then + COMMON+=(--freeze-encoder) +fi + +if [[ "${TABLE2_ICD_CODES}" == "1" ]]; then + COMMON+=(--icd-codes) +fi + +if [[ "${TABLE2_INCLUDE_VITALS}" == "1" ]]; then + COMMON+=(--include-vitals) +fi + +if [[ "${TABLE2_BALANCED_SAMPLING}" == "1" ]]; then + COMMON+=(--balanced-sampling --balanced-ratio "${TABLE2_BALANCED_RATIO}") +fi + if [[ "${TABLE2_DRY_RUN:-0}" == "1" ]]; then echo "Dry-run complete: conda activation and argument assembly succeeded." exit 0 fi case "${MODEL}" in + mlp) + # No dim overrides needed; A6000/A100 handles bs=16 at embedding-dim=128. + COMMON+=(--batch-size "${TABLE2_BS_MLP:-16}") + ;; + rnn) + COMMON+=(--batch-size "${TABLE2_BS_RNN:-16}") + ;; + transformer) + # Bumped batch-size 1→2: at bs=1 ~9000s/epoch × 20 = 50h, exceeds deadline. + # embedding-dim=64 is tiny; bs=2 is safe on A6000 (48GB) and A100 (80GB). + COMMON+=( + --batch-size 2 + --embedding-dim 64 + --hidden-dim 64 + --heads 2 + --num-layers 1 + ) + ;; bottleneck_transformer) - COMMON+=(--max-grad-norm 0.5 --bottlenecks-n 4 --fusion-startidx 1) + # Same timing issue as transformer; bs=2 is safe given embedding-dim=96. + COMMON+=( + --batch-size 2 + --embedding-dim 96 + --hidden-dim 96 + --heads 2 + --num-layers 1 + --max-grad-norm 0.5 + --bottlenecks-n 4 + --fusion-startidx 1 + ) ;; ehrmamba) - COMMON+=(--mamba-state-size 16 --mamba-conv-kernel 4) + COMMON+=( + --batch-size 2 + --embedding-dim 96 + --hidden-dim 96 + --mamba-state-size 16 + --mamba-conv-kernel 4 + ) ;; jambaehr) - COMMON+=(--jamba-transformer-layers 2 --jamba-mamba-layers 6 --mamba-state-size 16 --mamba-conv-kernel 4) + # Bumped batch-size 1→2 for same timing reason as transformer. + # Reduced jamba-mamba-layers 4→2: 3 total layers still a valid Jamba model, + # saves ~30-40% per-step cost. + COMMON+=( + --batch-size 2 + --embedding-dim 64 + --hidden-dim 64 + --jamba-transformer-layers 1 + --jamba-mamba-layers 2 + --mamba-state-size 16 + --mamba-conv-kernel 4 + ) ;; esac - -# Sample 3 distinct random seeds each run. -mapfile -t SEEDS < <( - python - <<'PY' -import random - -for seed in random.sample(range(1, 2_147_483_647), 3): - print(seed) -PY -) - -if [[ "${#SEEDS[@]}" -ne 3 ]]; then - echo "ERROR: failed to generate 3 random seeds." >&2 - exit 1 -fi - -echo " Seeds : ${SEEDS[*]}" -echo "========================================================" - -for SEED in "${SEEDS[@]}"; do - python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" --seed "${SEED}" -done +python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" --seed "${SEED}" echo "========================================================" -echo " Completed model=${MODEL}" -echo " Seeds used: ${SEEDS[*]}" +echo " Completed label=${TABLE2_RUN_LABEL} model=${MODEL} seed=${SEED}" echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" echo "========================================================" diff --git a/scripts/slurm/run_table2.sh b/scripts/slurm/run_table2.sh index dcd66769b..97d6b85fb 100755 --- a/scripts/slurm/run_table2.sh +++ b/scripts/slurm/run_table2.sh @@ -16,8 +16,12 @@ CACHE_DIR="${CACHE_DIR:-/u/${USER}/pyhealth_cache}" TABLE2_EPOCHS="${TABLE2_EPOCHS:-20}" TABLE2_NUM_WORKERS="${TABLE2_NUM_WORKERS:-2}" TABLE2_DEV_MODE="${TABLE2_DEV_MODE:-0}" +TABLE2_TASK="${TABLE2_TASK:-clinical_notes_icd_labs}" +TABLE2_WINDOW_HOURS="${TABLE2_WINDOW_HOURS:-24}" TABLE2_OUTPUT_DIR="${TABLE2_OUTPUT_DIR:-output/table2}" TABLE2_RUN_LABEL="${TABLE2_RUN_LABEL:-full}" +TABLE2_FREEZE_ENCODER="${TABLE2_FREEZE_ENCODER:-0}" +TABLE2_ICD_CODES="${TABLE2_ICD_CODES:-0}" # ── Activate conda ──────────────────────────────────────────────── module load miniconda3/24.9.2 @@ -62,14 +66,23 @@ echo " Epochs : ${TABLE2_EPOCHS}" echo " Workers : ${TABLE2_NUM_WORKERS}" echo " Dev mode : ${TABLE2_DEV_MODE}" echo " Output dir: ${TABLE2_OUTPUT_DIR}" +echo " Task : ${TABLE2_TASK}" +echo " Window : ${TABLE2_WINDOW_HOURS}h" echo " Patience : 5 (early stopping)" +echo " Freeze enc: ${TABLE2_FREEZE_ENCODER} (1=freeze BERT)" +echo " ICD codes : ${TABLE2_ICD_CODES} (1=include, ablation only)" echo "========================================================" COMMON=( --ehr-root "${EHR_ROOT}" --note-root "${NOTE_ROOT}" --cache-dir "${CACHE_DIR}" +<<<<<<< HEAD + --task "${TABLE2_TASK}" + --observation-window-hours "${TABLE2_WINDOW_HOURS}" +======= --task "${TABLE2_TASK:-clinical_notes_icd_labs}" +>>>>>>> origin/main --model "${MODEL}" --embedding-dim 128 --hidden-dim 128 @@ -88,6 +101,14 @@ if [[ "${TABLE2_DEV_MODE}" == "1" ]]; then COMMON+=(--dev) fi +if [[ "${TABLE2_FREEZE_ENCODER}" == "1" ]]; then + COMMON+=(--freeze-encoder) +fi + +if [[ "${TABLE2_ICD_CODES}" == "1" ]]; then + COMMON+=(--icd-codes) +fi + if [[ "${TABLE2_DRY_RUN:-0}" == "1" ]]; then echo "Dry-run complete." exit 0 diff --git a/scripts/will/slurm_run_labs_only_variant.py b/scripts/will/slurm_run_labs_only_variant.py new file mode 100644 index 000000000..6e13a38ea --- /dev/null +++ b/scripts/will/slurm_run_labs_only_variant.py @@ -0,0 +1,116 @@ +project_dir = "/u/wp14/PyHealth" +seed = 12 +account = "jimeng-ic" +partition = "eng-research-gpu" +time = "24:00:00" +mem = "64G" +gres = "gpu:A10:1" +conda_env = "pyhealth2" +ehr_root = "/projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2" +cache_dir = "/u/wp14/pyhealth_cache" +output_dir = "output/rnn_labs" +embedding_dim = 128 +hidden_dim = 128 +rnn_type = "GRU" +rnn_layers = 2 +dropout = 0.1 +epochs = 5 +batch_size = 32 +lr = 1e-3 +weight_decay = 1e-5 +patience = 5 +num_workers = 4 +dev = False + +# ── Step 0: Clean logs and cache ──────────────────────────────────────────── +print(f""" +### STEP 0 (optional): Clean logs and cache + +rm -rf {project_dir}/logs/* +rm -rf {cache_dir}/* +""") + +# ── Step 1: Reserve resources ───────────────────────────────────────────────── +if dev: + print(f""" +### STEP 1: Reserve an interactive node + +srun \\ + --account={account} \\ + --partition={partition} \\ + --nodes=1 --ntasks=1 --cpus-per-task={num_workers} \\ + --mem={mem} --gres={gres} --time={time} \\ + --pty bash +""") +else: + print(f""" +### STEP 1: Submit batch job (includes step 2 automatically) + +mkdir -p {project_dir}/logs && sbatch \\ + --job-name=rnn_labs_s{seed} \\ + --account={account} \\ + --partition={partition} \\ + --nodes=1 --ntasks=1 --cpus-per-task={num_workers} \\ + --mem={mem} --gres={gres} --time={time} \\ + --output={project_dir}/logs/rnn_labs_s{seed}_%j.out \\ + --error={project_dir}/logs/rnn_labs_s{seed}_%j.err \\ + --wrap=' + module load miniconda3/24.9.2 && + eval "$(conda shell.bash hook)" && + conda activate {conda_env} && + cd {project_dir} && + export PYTHONPATH={project_dir}:$PYTHONPATH && + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root {ehr_root} \\ + --cache-dir {cache_dir} \\ + --task labs \\ + --model rnn \\ + --embedding-dim {embedding_dim} \\ + --hidden-dim {hidden_dim} \\ + --rnn-type {rnn_type} \\ + --rnn-layers {rnn_layers} \\ + --dropout {dropout} \\ + --epochs {epochs} \\ + --batch-size {batch_size} \\ + --lr {lr} \\ + --weight-decay {weight_decay} \\ + --patience {patience} \\ + --num-workers {num_workers} \\ + --seed {seed} \\ + --output-dir {output_dir} + ' +""") + +# ── Step 2: Run (only needed for dev/interactive) ───────────────────────────── +print("\n" + "=" * 60 + "\n") +if dev: + print(f""" +### STEP 2: Once on the node, run + +module load miniconda3/24.9.2 && +eval "$(conda shell.bash hook)" && +conda activate {conda_env} && +cd {project_dir} && +export PYTHONPATH={project_dir}:$PYTHONPATH && +mkdir -p {project_dir}/logs/dev && +python {project_dir}/examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root {ehr_root} \\ + --cache-dir {cache_dir} \\ + --task labs \\ + --model rnn \\ + --embedding-dim {embedding_dim} \\ + --hidden-dim {hidden_dim} \\ + --rnn-type {rnn_type} \\ + --rnn-layers {rnn_layers} \\ + --dropout {dropout} \\ + --epochs {epochs} \\ + --batch-size {batch_size} \\ + --lr {lr} \\ + --weight-decay {weight_decay} \\ + --patience {patience} \\ + --num-workers {num_workers} \\ + --seed {seed} \\ + --output-dir {output_dir} --dev \\ + > >(tee {project_dir}/logs/dev/rnn_labs_s{seed}.out) \\ + 2> >(tee {project_dir}/logs/dev/rnn_labs_s{seed}.err >&2) +""") diff --git a/scripts/will/tmux_run_labs_only_variant.py b/scripts/will/tmux_run_labs_only_variant.py new file mode 100644 index 000000000..3ba7301f3 --- /dev/null +++ b/scripts/will/tmux_run_labs_only_variant.py @@ -0,0 +1,89 @@ +project_dir = "/home/wp14/PyHealth" +seed = 12 +conda_env = "pyhealth2" +ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2" +cache_dir = "/home/wp14/pyhealth_cache" +logs_dir = "/home/wp14/logs" +output_dir = "/home/wp14/output" +embedding_dim = 128 +hidden_dim = 128 +rnn_type = "GRU" +rnn_layers = 2 +dropout = 0.1 +epochs = 50 +batch_size = 32 +lr = 1e-3 +weight_decay = 1e-5 +patience = 5 +num_workers = 4 +dev = False +cuda_visible_devices = "0" +session_name = f"rnn_labs_s{seed}" + +# ── Step 0: Clean logs and cache ──────────────────────────────────────────── +print(f""" +### STEP 0 (optional): Clean logs and cache + +rm -rf {logs_dir}/* +rm -rf {cache_dir}/* +rm -rf {output_dir}/* +""") + +# ── Step 1: Start a tmux session ──────────────────────────────────────────── +print(f""" +### STEP 1: Start a tmux session (attached) + +tmux new-session -s {session_name} +""") + +# ── Step 2: Paste this into the session to run training ───────────────────── +print("\n" + "=" * 60 + "\n") + +log_dir = logs_dir +log_tag = f"rnn_labs_s{seed}" + +print(f""" +### STEP 2: Paste this into the tmux session + +mkdir -p {log_dir} && +eval "$(conda shell.bash hook)" && +conda activate {conda_env} && +cd {project_dir} && +export PYTHONPATH={project_dir}:$PYTHONPATH && +export CUDA_VISIBLE_DEVICES={cuda_visible_devices} && +python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root {ehr_root} \\ + --cache-dir {cache_dir} \\ + --task labs{' --dev' if dev else ''} \\ + --model rnn \\ + --embedding-dim {embedding_dim} \\ + --hidden-dim {hidden_dim} \\ + --rnn-type {rnn_type} \\ + --rnn-layers {rnn_layers} \\ + --dropout {dropout} \\ + --epochs {epochs} \\ + --batch-size {batch_size} \\ + --lr {lr} \\ + --weight-decay {weight_decay} \\ + --patience {patience} \\ + --num-workers {num_workers} \\ + --seed {seed} \\ + --output-dir {output_dir} \\ + > >(tee {log_dir}/{log_tag}.out) \\ + 2> >(tee {log_dir}/{log_tag}.err >&2) +""") + +# ── Step 3: Detach / reattach / monitor ────────────────────────────────────── +print(f""" +### STEP 3: Detach without killing (Ctrl+b d), then reattach later with + +tmux attach -t {session_name} + +### To check on it later without attaching: + +tail -f {log_dir}/{log_tag}.out + +### To kill the session when done: + +tmux kill-session -t {session_name} +""") diff --git a/tests/core/test_notes_labs_mimic4.py b/tests/core/test_notes_labs_mimic4.py new file mode 100644 index 000000000..4f6c5cb6b --- /dev/null +++ b/tests/core/test_notes_labs_mimic4.py @@ -0,0 +1,409 @@ +"""Unit tests for NotesLabsMIMIC4, admission-section extraction, and ICDLabsMIMIC4 fixes.""" + +from datetime import datetime +import unittest + + +class _DummyEvent: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + +class _DummyPatientWithNotes: + def __init__(self, note_texts=None, icd_codes=None, lab_df=None, vital_df=None) -> None: + self.patient_id = "p-1" + self._admissions = [ + _DummyEvent( + timestamp=datetime(2020, 1, 1, 0, 0, 0), + dischtime="2020-01-03 12:00:00", + hadm_id=101, + hospital_expire_flag=0, + ) + ] + self._patients = [_DummyEvent(anchor_age=55)] + self._note_texts = note_texts or [] + self._icd_codes = icd_codes or [] + self._lab_df = lab_df + self._vital_df = vital_df + + def get_events(self, event_type, start=None, end=None, filters=None, return_df=False): + if event_type == "patients": + return self._patients + if event_type == "admissions": + return self._admissions + if event_type == "discharge": + out = [] + for text in self._note_texts: + out.append( + _DummyEvent( + timestamp=datetime(2020, 1, 3, 12, 0, 0), + text=text, + ) + ) + return out + if event_type in {"diagnoses_icd", "procedures_icd"}: + return self._icd_codes + if event_type == "labevents" and return_df: + return self._lab_df + if event_type == "chartevents" and return_df: + return self._vital_df if self._vital_df is not None else pl.DataFrame( + { + "timestamp": [], + "chartevents/itemid": [], + "chartevents/storetime": [], + "chartevents/valuenum": [], + } + ) + if event_type == "chartevents" and not return_df: + return [] + return [] + + +class _DummyPatientMalformedDischtime: + def __init__(self) -> None: + self.patient_id = "p-2" + self._admissions = [ + _DummyEvent( + timestamp=datetime(2020, 1, 1, 0, 0, 0), + dischtime="malformed-dischtime", + hadm_id=102, + hospital_expire_flag=0, + ) + ] + self._patients = [_DummyEvent(anchor_age=60)] + + def get_events(self, event_type, start=None, end=None, filters=None, return_df=False): + if event_type == "patients": + return self._patients + if event_type == "admissions": + return self._admissions + if event_type in {"diagnoses_icd", "procedures_icd", "discharge", "radiology"}: + return [] + if event_type == "labevents" and return_df: + import polars as pl + + return pl.DataFrame( + { + "timestamp": [], + "labevents/itemid": [], + "labevents/storetime": [], + "labevents/valuenum": [], + } + ) + if event_type == "chartevents" and return_df: + import polars as pl + + return pl.DataFrame( + { + "timestamp": [], + "chartevents/itemid": [], + "chartevents/storetime": [], + "chartevents/valuenum": [], + } + ) + return [] + + +class TestExtractAdmissionSections(unittest.TestCase): + def test_extracts_target_sections(self): + from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task + + text = """Chief Complaint: +Shortness of breath + +Past Medical History: +1. Hypertension + +Medications on Admission: +1. Metoprolol 25 mg PO BID + +Discharge Diagnosis: +Acute MI +""" + result = BaseMultimodalMIMIC4Task._extract_admission_sections(text) + self.assertIn("Shortness of breath", result) + self.assertIn("Hypertension", result) + self.assertIn("Metoprolol", result) + self.assertNotIn("Acute MI", result) + self.assertIn("[SEP]", result) + + def test_fallback_to_first_1024_when_no_sections(self): + from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task + + text = "This is a note with no section headers at all. " * 50 + result = BaseMultimodalMIMIC4Task._extract_admission_sections(text) + self.assertEqual(result, text[:1024]) + + def test_case_insensitive_headers(self): + from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task + + text = """CHIEF COMPLAINT: +Chest pain + +Past Medical/Surgical History: +Appendectomy +""" + result = BaseMultimodalMIMIC4Task._extract_admission_sections(text) + self.assertIn("Chest pain", result) + self.assertIn("Appendectomy", result) + + def test_empty_string_fallback(self): + from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task + + result = BaseMultimodalMIMIC4Task._extract_admission_sections("") + self.assertEqual(result, "") + + +class TestCollectAdmissionNoteSections(unittest.TestCase): + def test_collects_sections_and_returns_time_zero(self): + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4() + patient = _DummyPatientWithNotes( + note_texts=["Chief Complaint:\nFever\n\nPast Medical History:\nDiabetes"] + ) + texts, times = task._collect_admission_note_sections( + patient, 101, datetime(2020, 1, 1, 0, 0, 0) + ) + self.assertEqual(len(texts), 1) + self.assertIn("Fever", texts[0]) + self.assertEqual(times, [0.0]) + + def test_missing_note_fallback(self): + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4() + patient = _DummyPatientWithNotes(note_texts=[]) + texts, times = task._collect_admission_note_sections( + patient, 101, datetime(2020, 1, 1, 0, 0, 0) + ) + self.assertEqual(texts, [""]) + self.assertEqual(times, [0.0]) + + def test_no_time_filter_applied(self): + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4() + # Discharge note timestamp is 2020-01-03, well outside any 24h window + patient = _DummyPatientWithNotes( + note_texts=["Chief Complaint:\nFever"] + ) + texts, times = task._collect_admission_note_sections( + patient, 101, datetime(2020, 1, 1, 0, 0, 0) + ) + self.assertEqual(len(texts), 1) + self.assertIn("Fever", texts[0]) + + +class TestNotesLabsMIMIC4(unittest.TestCase): + def test_default_schema_excludes_icd(self): + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4() + self.assertNotIn("icd_codes", task.input_schema) + self.assertNotIn("vitals", task.input_schema) + self.assertNotIn("vitals_mask", task.input_schema) + self.assertIn("admission_note_times", task.input_schema) + self.assertIn("labs", task.input_schema) + self.assertIn("labs_mask", task.input_schema) + + def test_include_icd_adds_icd_schema(self): + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4(include_icd=True) + self.assertIn("icd_codes", task.input_schema) + + def test_include_vitals_adds_vitals_schema(self): + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4(include_vitals=True) + self.assertIn("vitals", task.input_schema) + self.assertIn("vitals_mask", task.input_schema) + self.assertNotIn("icd_codes", task.input_schema) + + def test_include_vitals_and_icd_adds_both(self): + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4(include_icd=True, include_vitals=True) + self.assertIn("vitals", task.input_schema) + self.assertIn("vitals_mask", task.input_schema) + self.assertIn("icd_codes", task.input_schema) + + def test_output_structure_with_vitals(self): + import polars as pl + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4(window_hours=24, include_vitals=True) + lab_df = pl.DataFrame( + { + "timestamp": [datetime(2020, 1, 1, 2, 0, 0)], + "labevents/itemid": ["50824"], + "labevents/storetime": ["2020-01-01 02:00:00"], + "labevents/valuenum": [138.0], + } + ) + vital_df = pl.DataFrame( + { + "timestamp": [datetime(2020, 1, 1, 1, 0, 0)], + "chartevents/itemid": ["220045"], # HeartRate + "chartevents/storetime": ["2020-01-01 01:00:00"], + "chartevents/valuenum": [80.0], + } + ) + patient = _DummyPatientWithNotes( + note_texts=["Chief Complaint:\nFever"], + lab_df=lab_df, + vital_df=vital_df, + ) + samples = task(patient) + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertIn("vitals", sample) + self.assertIn("vitals_mask", sample) + self.assertNotIn("icd_codes", sample) + vital_times, vital_values = sample["vitals"] + self.assertGreater(len(vital_times), 0) + + def test_vitals_fallback_when_empty(self): + import polars as pl + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4(window_hours=24, include_vitals=True) + lab_df = pl.DataFrame( + { + "timestamp": [], + "labevents/itemid": [], + "labevents/storetime": [], + "labevents/valuenum": [], + } + ) + vital_df = pl.DataFrame( + { + "timestamp": [], + "chartevents/itemid": [], + "chartevents/storetime": [], + "chartevents/valuenum": [], + } + ) + patient = _DummyPatientWithNotes( + note_texts=["Chief Complaint:\nFever"], + lab_df=lab_df, + vital_df=vital_df, + ) + samples = task(patient) + self.assertEqual(len(samples), 1) + vital_times, vital_values = samples[0]["vitals"] + self.assertEqual(len(vital_times), 1) + self.assertEqual(len(vital_values[0]), len(task.VITAL_CATEGORY_NAMES)) + + def test_output_structure_no_icd(self): + import polars as pl + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4(window_hours=24) + lab_df = pl.DataFrame( + { + "timestamp": [datetime(2020, 1, 1, 2, 0, 0)], + "labevents/itemid": ["50824"], # Sodium + "labevents/storetime": ["2020-01-01 02:00:00"], + "labevents/valuenum": [138.0], + } + ) + patient = _DummyPatientWithNotes( + note_texts=["Chief Complaint:\nFever"], + lab_df=lab_df, + ) + samples = task(patient) + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertIn("admission_note_times", sample) + self.assertIn("labs", sample) + self.assertIn("labs_mask", sample) + self.assertNotIn("icd_codes", sample) + self.assertEqual(sample["mortality"], 0) + + def test_output_structure_with_icd(self): + import polars as pl + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4(window_hours=24, include_icd=True) + lab_df = pl.DataFrame( + { + "timestamp": [], + "labevents/itemid": [], + "labevents/storetime": [], + "labevents/valuenum": [], + } + ) + patient = _DummyPatientWithNotes( + note_texts=["Chief Complaint:\nFever"], + icd_codes=[_DummyEvent(icd_code="I21")], + lab_df=lab_df, + ) + samples = task(patient) + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertIn("icd_codes", sample) + + def test_mortality_label_positive(self): + import polars as pl + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + task = NotesLabsMIMIC4(window_hours=24) + lab_df = pl.DataFrame( + { + "timestamp": [], + "labevents/itemid": [], + "labevents/storetime": [], + "labevents/valuenum": [], + } + ) + patient = _DummyPatientWithNotes( + note_texts=["Chief Complaint:\nFever"], + lab_df=lab_df, + ) + patient._admissions[0].hospital_expire_flag = 1 + samples = task(patient) + self.assertEqual(samples[0]["mortality"], 1) + + +class TestICDLabsMIMIC4Fixes(unittest.TestCase): + def test_malformed_dischtime_does_not_drop_admission(self): + from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4 + + task = ICDLabsMIMIC4(window_hours=24) + patient = _DummyPatientMalformedDischtime() + samples = task(patient) + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertGreater(len(sample["icd_codes"][0]), 0) + self.assertGreater(len(sample["labs"][0]), 0) + self.assertGreater(len(sample["labs_mask"][0]), 0) + + def test_missing_icd_code_uses_text_token(self): + import polars as pl + from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4 + + task = ICDLabsMIMIC4(window_hours=24) + lab_df = pl.DataFrame( + { + "timestamp": [], + "labevents/itemid": [], + "labevents/storetime": [], + "labevents/valuenum": [], + } + ) + patient = _DummyPatientWithNotes( + note_texts=[], + icd_codes=[], # no ICD codes + lab_df=lab_df, + ) + samples = task(patient) + self.assertEqual(len(samples), 1) + _, icd_visits = samples[0]["icd_codes"] + self.assertEqual(icd_visits, [[""]]) + + +if __name__ == "__main__": + unittest.main()