diff --git a/.gitignore b/.gitignore index 849312887..6e802ed70 100644 --- a/.gitignore +++ b/.gitignore @@ -139,4 +139,7 @@ data/physionet.org/ .vscode/ # Local Data -local_data/ \ No newline at end of file +local_data/ + +# Model weight files (large binaries, distributed separately) +weightfiles/ \ No newline at end of file diff --git a/README.rst b/README.rst index cf12b5874..c3f3e6355 100644 --- a/README.rst +++ b/README.rst @@ -189,8 +189,6 @@ Module 1: root="https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/", # raw CSV table name tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"], - # map all NDC codes to CCS codes in these tables - code_mapping={"NDC": "CCSCM"}, ) .. image:: figure/structured-dataset.png diff --git a/configs/train/base.yaml b/configs/train/base.yaml new file mode 100644 index 000000000..70c7f2c7c --- /dev/null +++ b/configs/train/base.yaml @@ -0,0 +1,80 @@ +# PyHealth unified training config — base defaults +# All values here are the OOM-safe defaults validated on devsplit. +# Override per-condition and per-model in condition/model-specific configs. + +# ── paths ────────────────────────────────────────────────────────────────────── +ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2 +note_root: /shared/rsaas/physionet.org/files/mimic-note +cache_dir: /home/rianatri/pyhealth_cache +output_dir: output/unified + +# ── task ────────────────────────────────────────────────────────────────────── +# One of: notes_labs | labs_only | icd_labs | clinical_notes_icd_labs | stagenet +task: notes_labs +observation_window_hours: 24 + +# Task flags (notes_labs only) +icd_codes: false +include_vitals: false +balanced_sampling: false +balanced_ratio: 1.0 + +# ── model ───────────────────────────────────────────────────────────────────── +# One of: mlp | rnn | transformer | bottleneck_transformer | ehrmamba | jambaehr +model: mlp +freeze_encoder: false + +# Shared embedding dims — safe defaults across all models on 24 GB GPU (frozen) +# or 80 GB A100 (full BERT). +embedding_dim: 128 +hidden_dim: 128 +heads: 4 +num_layers: 2 +dropout: 0.1 + +# ── training ────────────────────────────────────────────────────────────────── +epochs: 50 +batch_size: 16 +lr: null # null = model-specific default (1e-4 for all) +weight_decay: 1.0e-5 +patience: 10 +seed: 42 +num_workers: 2 + +# ── dev mode ────────────────────────────────────────────────────────────────── +# 0 = full dataset; N > 0 = limit to N patients (devsplit) +dev: 0 + +# ── model-specific overrides ────────────────────────────────────────────────── +# These are merged at runtime based on `model` value. +# Keys match CLI args without the leading `--`. +_model_overrides: + transformer: + batch_size: 2 + embedding_dim: 64 + hidden_dim: 64 + heads: 2 + num_layers: 1 + bottleneck_transformer: + 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: + batch_size: 2 + embedding_dim: 96 + hidden_dim: 96 + mamba_state_size: 16 + mamba_conv_kernel: 4 + jambaehr: + 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 diff --git a/configs/train/e2e_balanced.yaml b/configs/train/e2e_balanced.yaml new file mode 100644 index 000000000..74aaedcb9 --- /dev/null +++ b/configs/train/e2e_balanced.yaml @@ -0,0 +1,8 @@ +# PyHealth E2E balanced condition — notes+labs with 1:1 pos:neg undersampling +_inherit: base.yaml + +task: notes_labs +balanced_sampling: true +balanced_ratio: 1.0 +epochs: 50 +output_dir: output/e2e_full/balanced diff --git a/configs/train/e2e_baseline.yaml b/configs/train/e2e_baseline.yaml new file mode 100644 index 000000000..df1a6243e --- /dev/null +++ b/configs/train/e2e_baseline.yaml @@ -0,0 +1,11 @@ +# PyHealth E2E baseline condition — notes+labs, ICD off, no balancing +# Full dataset, 50 epochs, patience=10. + +_inherit: base.yaml + +task: notes_labs +icd_codes: false +include_vitals: false +balanced_sampling: false +epochs: 50 +output_dir: output/e2e_full/baseline diff --git a/configs/train/e2e_icd_on.yaml b/configs/train/e2e_icd_on.yaml new file mode 100644 index 000000000..13f817fd3 --- /dev/null +++ b/configs/train/e2e_icd_on.yaml @@ -0,0 +1,7 @@ +# PyHealth E2E ICD-on condition — notes+labs+ICD (discharge-coded leakage ablation) +_inherit: base.yaml + +task: notes_labs +icd_codes: true +epochs: 50 +output_dir: output/e2e_full/icd_on diff --git a/configs/train/e2e_labs_only.yaml b/configs/train/e2e_labs_only.yaml new file mode 100644 index 000000000..d0ca5ade8 --- /dev/null +++ b/configs/train/e2e_labs_only.yaml @@ -0,0 +1,6 @@ +# PyHealth E2E labs-only condition — EHR-only reference baseline +_inherit: base.yaml + +task: labs_only +epochs: 50 +output_dir: output/e2e_full/labs_only diff --git a/configs/train/smoke.yaml b/configs/train/smoke.yaml new file mode 100644 index 000000000..af0079770 --- /dev/null +++ b/configs/train/smoke.yaml @@ -0,0 +1,14 @@ +# PyHealth smoke test config — devsplit, fast iteration +# Inherits base.yaml defaults; overrides for fast validation. + +_inherit: base.yaml + +task: notes_labs +model: mlp +dev: 1000 +epochs: 5 +patience: 5 +batch_size: 16 +embedding_dim: 64 +hidden_dim: 64 +output_dir: output/smoke diff --git a/docs/api/models/pyhealth.models.AdaCare.rst b/docs/api/models/pyhealth.models.AdaCare.rst index 00aeaf4f0..4d988ba78 100644 --- a/docs/api/models/pyhealth.models.AdaCare.rst +++ b/docs/api/models/pyhealth.models.AdaCare.rst @@ -9,6 +9,11 @@ The separate callable AdaCareLayer and the complete AdaCare model. :show-inheritance: .. autoclass:: pyhealth.models.AdaCare + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.models.MultimodalAdaCare :members: :undoc-members: :show-inheritance: \ No newline at end of file diff --git a/docs/api/models/pyhealth.models.RETAIN.rst b/docs/api/models/pyhealth.models.RETAIN.rst index 88899363e..ac8bbad06 100644 --- a/docs/api/models/pyhealth.models.RETAIN.rst +++ b/docs/api/models/pyhealth.models.RETAIN.rst @@ -9,6 +9,11 @@ The separate callable RETAINLayer and the complete RETAIN model. :show-inheritance: .. autoclass:: pyhealth.models.RETAIN + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.models.MultimodalRETAIN :members: :undoc-members: :show-inheritance: \ No newline at end of file diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index 47c7ce356..347ec4df3 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -212,8 +212,6 @@ Available Tasks COVID-19 CXR Classification DKA Prediction (MIMIC-IV) Drug Recommendation - EEG Abnormal - EEG Events Length of Stay Prediction Medical Transcriptions Classification Mortality Prediction (Next Visit) diff --git a/docs/api/tasks/pyhealth.tasks.EEG_abnormal.rst b/docs/api/tasks/pyhealth.tasks.EEG_abnormal.rst deleted file mode 100644 index c6ca62bd0..000000000 --- a/docs/api/tasks/pyhealth.tasks.EEG_abnormal.rst +++ /dev/null @@ -1,4 +0,0 @@ -pyhealth.tasks.EEG_abnormal -======================================= - -.. autofunction:: pyhealth.tasks.EEG_abnormal.EEG_isAbnormal_fn \ No newline at end of file diff --git a/docs/api/tasks/pyhealth.tasks.EEG_events.rst b/docs/api/tasks/pyhealth.tasks.EEG_events.rst deleted file mode 100644 index 62b15963c..000000000 --- a/docs/api/tasks/pyhealth.tasks.EEG_events.rst +++ /dev/null @@ -1,4 +0,0 @@ -pyhealth.tasks.EEG_events -======================================= - -.. autofunction:: pyhealth.tasks.EEG_events.EEG_events_fn \ No newline at end of file diff --git a/docs/api/tasks/pyhealth.tasks.multimodal_mimic4.rst b/docs/api/tasks/pyhealth.tasks.multimodal_mimic4.rst index 4fb9f660c..0100745c1 100644 --- a/docs/api/tasks/pyhealth.tasks.multimodal_mimic4.rst +++ b/docs/api/tasks/pyhealth.tasks.multimodal_mimic4.rst @@ -10,3 +10,8 @@ pyhealth.tasks.multimodal_mimic4 :members: :undoc-members: :show-inheritance: + +.. autoclass:: pyhealth.tasks.multimodal_mimic4.ClinicalNotesICDLabsCXRMIMIC4 + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/index.rst b/docs/index.rst index 37ec2e705..d85f89d35 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -7,7 +7,14 @@ Welcome to PyHealth **The Python Library for Healthcare AI** -Build, test, and deploy healthcare machine learning models with ease. PyHealth is designed for both **ML researchers and medical practitioners**. We can make your **healthcare AI applications** easier to develop, test and validate. Your development process becomes more flexible and more customizable. `[GitHub] `_ +.. card:: 🌐 Visit the PyHealth Project Website + :link: https://pyhealth.dev + :link-type: url + :class-card: sd-bg-primary sd-text-white sd-text-center + + **pyhealth.dev** — the new home for PyHealth news, updates, and resources → + +Build, test, and deploy healthcare machine learning models with ease. PyHealth is designed for both **ML researchers and medical practitioners**. We can make your **healthcare AI applications** easier to develop, test and validate. Your development process becomes more flexible and more customizable. `[GitHub] `_ **Key Features** diff --git a/docs/research_initiative.rst b/docs/research_initiative.rst index f7ddaba14..f631c0dd2 100644 --- a/docs/research_initiative.rst +++ b/docs/research_initiative.rst @@ -11,6 +11,12 @@ about advancing computational healthcare, regardless of their career stage or in participants work on innovative projects that advance the field of computational healthcare, contributing to publications, open-source software, and the broader healthcare AI community. +Our goals are to build: + +1. **Easily accessible and reproducible research** — Making healthcare AI research transparent and replicable +2. **Solutions to real-world healthcare problems** — Tackling important clinical challenges with practical impact +3. **Connections with healthcare professionals** — Bridging the gap between AI researchers and clinical practitioners + The initiative provides participants with hands-on experience in: - **Healthcare AI Research**: Working on real-world healthcare problems using electronic health records (EHRs) and clinical data @@ -19,17 +25,45 @@ The initiative provides participants with hands-on experience in: - **Academic Publishing**: Co-authoring research papers and presenting findings - **Collaborative Research**: Working alongside researchers and industry partners -Mission & Goals ---------------- +Program Logistics +----------------- -We're an open-source community of researchers with the goal of building: +**Format**: Remote -1. **Easily accessible and reproducible research** — Making healthcare AI research transparent and replicable -2. **Solutions to real-world healthcare problems** — Tackling important clinical challenges with practical impact -3. **Connections with healthcare professionals** — Bridging the gap between AI researchers and clinical practitioners +**Time Commitment**: 10–20 hours per week during active research cycles + +**Eligibility**: Open to anyone! We welcome people from all backgrounds—students, engineers, researchers, +and healthcare professionals. We don't care about your title or institution. We only ask that you have the +ability to write decent-quality code and are self-driven to work hard on healthcare problems. +See `How to Apply`_ to get started. + +**Open Projects**: Browse available research projects and find one that matches your interests: +`Open Projects List `_ -We connect those who want to work on healthcare problems with industry and academic collaborators to create -meaningful impact in healthcare AI. +The program runs on a rolling, year-round basis with recurring terms aligned to major healthcare AI +conference cycles. Each term culminates in a submission to a top-tier venue. + +.. list-table:: Upcoming Research Terms + :widths: 20 25 30 25 + :header-rows: 1 + :class: research-table + + * - Term + - Period + - Target Conference + - Est. Submission Deadline + * - Summer Term + - Apr – Aug 2026 + - `ML4H 2026 `_ + - ~Sep 2026 + * - Fall Term + - Sep – Dec 2026 + - `CHIL 2027 `_ + - ~Feb 2027 + * - Spring Term + - Jan – Apr 2027 + - `MLHC 2027 `_ + - ~May 2027 Research Contributions ---------------------- @@ -48,6 +82,11 @@ drug recommendation, and healthcare AI infrastructure. - Paper Title - Venue - Links + * - 2026 + - Arjun Chatterjee, Sayeed Sajjad Razin + - Making Conformal Predictors Robust in Healthcare Settings: a Case Study on EEG Classification + - Under Review at AIME + - `Paper `_ * - 2025 - Zilal Eiz Al Din - MIMIC-RD: Can LLMs differentially diagnose rare diseases in real-world clinical settings? @@ -71,31 +110,6 @@ drug recommendation, and healthcare AI infrastructure. **Latest from ML4H 2025**: Our first cohort successfully published three papers at ML4H, covering rare diseases, social determinants of health, and prostate cancer genomics—all at the forefront of healthcare and AI research! -Program Information -------------------- - -**Duration**: Ongoing year-round program with targeted submission cycles - -**Format**: Remote - -**Time Commitment**: 10–20 hours per week during active research cycles - -**Eligibility**: Open to anyone! We welcome people from all backgrounds—students, engineers, researchers, -and healthcare professionals. We don't care about your title or institution. We only ask that you have the -ability to write decent-quality code and are self-driven to work hard on healthcare problems. - -**Current Cycle**: Spring 2026, targeting `MLHC 2026 `_ (Machine Learning for Healthcare) - -**Key Dates for Spring 2026**: - -- **Official start date**: January 20, 2026 — Kickoff meeting (details via Discord) -- **Pre-submission intent deadline**: April 10, 2026 -- **Full submission deadline**: April 11, 2026 -- **Conference dates**: August 15–16, 2026 - -**Open Projects**: Browse available research projects and find one that matches your interests: -`Open Projects List `_ - Research Areas -------------- diff --git a/examples/benchmark_perf/loc/minimal_drug_rec.py b/examples/benchmark_perf/loc/minimal_drug_rec.py index d35a23336..4aa00e244 100644 --- a/examples/benchmark_perf/loc/minimal_drug_rec.py +++ b/examples/benchmark_perf/loc/minimal_drug_rec.py @@ -1,7 +1,23 @@ -from pyhealth.datasets import MIMIC4Dataset -from pyhealth.tasks import DrugRecommendationMIMIC4 -base_dataset = MIMIC4Dataset( - ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", - ehr_tables=["patients", "admissions", "diagnoses_icd", "procedures_icd", "prescriptions"], -) -sample_dataset = base_dataset.set_task(DrugRecommendationMIMIC4()) \ No newline at end of file +import polars as pl; from pyhealth.datasets import MIMIC4Dataset; from pyhealth.tasks.base_task import BaseTask +class DrugRecommendationMIMIC4(BaseTask): + task_name="DrugRecommendationMIMIC4" + input_schema={"conditions":"nested_sequence","procedures":"nested_sequence","drugs_hist":"nested_sequence"} + output_schema={"drugs":"multilabel"} + def __call__(self,p): + adms=p.get_events(event_type="admissions") + if len(adms)<2: return [] + S=[] + for adm in adms: + f=[("hadm_id","==",adm.hadm_id)] + c=p.get_events("diagnoses_icd",filters=f,return_df=True).select(pl.concat_str(["diagnoses_icd/icd_version","diagnoses_icd/icd_code"],separator="_")).to_series().to_list() + r=p.get_events("procedures_icd",filters=f,return_df=True).select(pl.concat_str(["procedures_icd/icd_version","procedures_icd/icd_code"],separator="_")).to_series().to_list() + d=[x[:4] for x in p.get_events("prescriptions",filters=f,return_df=True).select(pl.col("prescriptions/ndc")).to_series().to_list() if x] + if not(c and r and d): continue + S.append({"visit_id":adm.hadm_id,"patient_id":p.patient_id,"conditions":c,"procedures":r,"drugs":d,"drugs_hist":d}) + if len(S)<2: return [] + S[0].update({"conditions":[S[0]["conditions"]],"procedures":[S[0]["procedures"]],"drugs_hist":[S[0]["drugs_hist"]]}) + for i in range(1,len(S)): S[i]["conditions"]=S[i-1]["conditions"]+[S[i]["conditions"]]; S[i]["procedures"]=S[i-1]["procedures"]+[S[i]["procedures"]]; S[i]["drugs_hist"]=S[i-1]["drugs_hist"]+[S[i]["drugs_hist"]] + for i in range(len(S)): S[i]["drugs_hist"][i]=[] + return S +base_dataset=MIMIC4Dataset(ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/",ehr_tables=["patients","admissions","diagnoses_icd","procedures_icd","prescriptions"]) +sample_dataset=base_dataset.set_task(DrugRecommendationMIMIC4()) diff --git a/examples/benchmark_perf/loc/minimal_legacy_drug_rec.py b/examples/benchmark_perf/loc/minimal_legacy_drug_rec.py index 48c42bb28..a805ab07b 100644 --- a/examples/benchmark_perf/loc/minimal_legacy_drug_rec.py +++ b/examples/benchmark_perf/loc/minimal_legacy_drug_rec.py @@ -1,7 +1,16 @@ -from pyhealth.datasets import MIMIC4Dataset -from pyhealth.tasks import drug_recommendation_mimic4_fn -base_dataset = MIMIC4Dataset(root="/srv/local/data/physionet.org/files/mimiciv/2.0/hosp", - tables=["diagnoses_icd", "procedures_icd", "prescriptions"], dev=False, - code_mapping={"NDC": "ATC"}, refresh_cache=True) -sample_dataset = base_dataset.set_task(task_fn=drug_recommendation_mimic4_fn) -print(f"Samples: {len(sample_dataset.samples)}") +from pyhealth.data import Patient,Visit; from pyhealth.datasets import MIMIC4Dataset + +def drug_recommendation_mimic4_fn(patient): + S=[] + for v in patient: + c=v.get_code_list(table="diagnoses_icd"); r=v.get_code_list(table="procedures_icd"); d=[x[:4] for x in v.get_code_list(table="prescriptions")] + if not(c and r and d): continue + S.append({"visit_id":v.visit_id,"patient_id":patient.patient_id,"conditions":c,"procedures":r,"drugs":d,"drugs_hist":d}) + if len(S)<2: return [] + S[0].update({"conditions":[S[0]["conditions"]],"procedures":[S[0]["procedures"]],"drugs_hist":[S[0]["drugs_hist"]]}) + for i in range(1,len(S)): S[i]["conditions"]=S[i-1]["conditions"]+[S[i]["conditions"]]; S[i]["procedures"]=S[i-1]["procedures"]+[S[i]["procedures"]]; S[i]["drugs_hist"]=S[i-1]["drugs_hist"]+[S[i]["drugs_hist"]] + for i in range(len(S)): S[i]["drugs_hist"][i]=[] + return S + +base_dataset=MIMIC4Dataset(root="/srv/local/data/physionet.org/files/mimiciv/2.0/hosp",tables=["diagnoses_icd","procedures_icd","prescriptions"],dev=False,code_mapping={"NDC":"ATC"},refresh_cache=True) +sample_dataset=base_dataset.set_task(task_fn=drug_recommendation_mimic4_fn) diff --git a/examples/benchmark_perf/loc/minimal_legacy_los.py b/examples/benchmark_perf/loc/minimal_legacy_los.py index 7d30b4919..2bc8eae8e 100644 --- a/examples/benchmark_perf/loc/minimal_legacy_los.py +++ b/examples/benchmark_perf/loc/minimal_legacy_los.py @@ -1,7 +1,14 @@ from pyhealth.datasets import MIMIC4Dataset -from pyhealth.tasks import length_of_stay_prediction_mimic4_fn -base_dataset = MIMIC4Dataset(root="/srv/local/data/physionet.org/files/mimiciv/2.0/hosp", - tables=["diagnoses_icd", "procedures_icd", "prescriptions"], dev=False, - code_mapping={"ICD10PROC": "CCSPROC", "NDC": "ATC"}, refresh_cache=True) -sample_dataset = base_dataset.set_task(task_fn=length_of_stay_prediction_mimic4_fn) -print(f"Samples: {len(sample_dataset.samples)}") + +def categorize_los(d): return 0 if d<1 else (d if d<=7 else (8 if d<=14 else 9)) + +def length_of_stay_prediction_mimic4_fn(patient): + S=[] + for v in patient: + c=v.get_code_list(table="diagnoses_icd"); r=v.get_code_list(table="procedures_icd"); d=v.get_code_list(table="prescriptions") + if not(c and r and d): continue + S.append({"visit_id":v.visit_id,"patient_id":patient.patient_id,"conditions":[c],"procedures":[r],"drugs":[d],"label":categorize_los((v.discharge_time-v.encounter_time).days)}) + return S + +base_dataset=MIMIC4Dataset(root="/srv/local/data/physionet.org/files/mimiciv/2.0/hosp",tables=["diagnoses_icd","procedures_icd","prescriptions"],dev=False,code_mapping={"ICD10PROC":"CCSPROC","NDC":"ATC"},refresh_cache=True) +sample_dataset=base_dataset.set_task(task_fn=length_of_stay_prediction_mimic4_fn) diff --git a/examples/benchmark_perf/loc/minimal_legacy_mortality.py b/examples/benchmark_perf/loc/minimal_legacy_mortality.py index b285e6789..719959f56 100644 --- a/examples/benchmark_perf/loc/minimal_legacy_mortality.py +++ b/examples/benchmark_perf/loc/minimal_legacy_mortality.py @@ -1,24 +1,27 @@ -from typing import List -from pyhealth.datasets import MIMIC4Dataset -from pyhealth.data import Patient, Visit -LAB_ITEM_IDS = {"50824", "52455", "50983", "52623", "50822", "52452", "50971", "52610", - "50806", "52434", "50902", "52535", "50803", "50804", "50809", "52027", - "50931", "52569", "50808", "51624", "50960", "50868", "52500", "52031", - "50964", "51701", "50970"} -def mortality_task_fn(patient: Patient) -> List[dict]: - samples = [] - for i in range(len(patient) - 1): - visit, next_visit = patient[i], patient[i + 1] - mortality_label = int(next_visit.discharge_status) if next_visit.discharge_status in [0, 1] else 0 - conditions = visit.get_code_list(table="diagnoses_icd") - procedures = visit.get_code_list(table="procedures_icd") - labs = list(dict.fromkeys([e.code for e in visit.get_event_list(table="labevents") if e.code in LAB_ITEM_IDS])) - if conditions and labs: - samples.append({"visit_id": visit.visit_id, "patient_id": patient.patient_id, - "conditions": [conditions], "procedures": [procedures] if procedures else [[]], - "labs": [labs], "label": mortality_label}) - return samples -base_dataset = MIMIC4Dataset(root="/srv/local/data/physionet.org/files/mimiciv/2.0/hosp", - tables=["diagnoses_icd", "procedures_icd", "labevents"], dev=False, refresh_cache=True) -sample_dataset = base_dataset.set_task(task_fn=mortality_task_fn) -print(f"Samples: {len(sample_dataset.samples)}") +from collections import defaultdict; from pyhealth.data import Patient; from pyhealth.datasets import MIMIC4Dataset +from typing import Dict,List + +LAB_CATS: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"]} +LAB_NAMES=list(LAB_CATS); LABITEMS=[x for ids in LAB_CATS.values() for x in ids] + +def mortality_task_fn(patient): + icd_d,icd_t,lab_v,lab_t,mort,prev=[],[],[],[],0,None + for v in patient: + at=v.encounter_time; dt=v.discharge_time + if at is None or dt is None or dt0: + ldf=ldf.with_columns(pl.col("labevents/storetime").str.strptime(pl.Datetime,"%Y-%m-%d %H:%M:%S")).filter(pl.col("labevents/storetime")<=dt).select(pl.col("timestamp"),pl.col("labevents/itemid"),pl.col("labevents/valuenum").cast(pl.Float64)) + for ts in sorted(ldf["timestamp"].unique().to_list()): + r=ldf.filter(pl.col("timestamp")==ts) + vec=[next((r.filter(pl.col("labevents/itemid")==iid)["labevents/valuenum"][0] for iid in self.LAB_CATS[cat] if r.filter(pl.col("labevents/itemid")==iid).height>0),None) for cat in self.LAB_NAMES] + lab_v.append(vec); lab_t.append((ts-at).total_seconds()/3600.0) + if not lab_v or not icd_d: return [] + return [{"patient_id":p.patient_id,"icd_codes":(icd_t,icd_d),"labs":(lab_t,lab_v),"mortality":mort}] +base_dataset=MIMIC4Dataset(ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/",ehr_tables=["patients","admissions","diagnoses_icd","procedures_icd","labevents"]) +sample_dataset=base_dataset.set_task(MortalityPredictionStageNetMIMIC4()) diff --git a/examples/conformal_eeg/tuab_conventional_conformal.py b/examples/conformal_eeg/tuab_conventional_conformal.py index 41fe603a1..0f6b32401 100644 --- a/examples/conformal_eeg/tuab_conventional_conformal.py +++ b/examples/conformal_eeg/tuab_conventional_conformal.py @@ -19,6 +19,7 @@ from __future__ import annotations import argparse +import os import random import sys from pathlib import Path @@ -45,8 +46,8 @@ def flush(self): from pyhealth.calib.predictionset import LABEL -from pyhealth.datasets import TUABDataset, get_dataloader, split_by_sample_conformal_tuh -from pyhealth.models import ContraWR +from pyhealth.datasets import TUABDataset, get_dataloader, split_by_patient_conformal_tuh, split_by_sample_conformal_tuh, split_by_sample_conformal +from pyhealth.models import ContraWR, TFMTokenizer from pyhealth.tasks import EEGAbnormalTUAB from pyhealth.trainer import Trainer, get_metrics_fn @@ -58,7 +59,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--root", type=str, - default="downloads/tuab/v3.0.0/edf", + default="/srv/local/data/TUH/tuh_eeg_abnormal/v3.0.0/edf", help="Path to TUAB edf/ folder.", ) parser.add_argument("--subset", type=str, default="both", choices=["train", "eval", "both"]) @@ -74,6 +75,10 @@ def parse_args() -> argparse.Namespace: "--alpha", type=float, default=0.1, help="Miscoverage rate (e.g., 0.1 => 90% target coverage).", ) + parser.add_argument( + "--alphas", type=str, default=None, + help="Comma-separated miscoverage rates, e.g. '0.2,0.1,0.05,0.01'. Overrides --alpha.", + ) parser.add_argument( "--ratios", type=float, @@ -84,6 +89,10 @@ def parse_args() -> argparse.Namespace: "Must sum to 1.0. Test is fixed as the TUH eval partition.", ) parser.add_argument("--n-fft", type=int, default=128, help="STFT FFT size used by ContraWR.") + parser.add_argument( + "--model", type=str, default="contrawr", choices=["contrawr", "tfm"], + help="Backbone model: 'contrawr' (default) or 'tfm' (TFMTokenizer).", + ) parser.add_argument( "--device", type=str, default=None, help="Device string, e.g. 'cuda:0' or 'cpu'. Defaults to auto-detect.", @@ -111,9 +120,48 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Smoke test: dev=True, max 2000 samples, 2 epochs.", ) + parser.add_argument( + "--weights-dir", + type=str, + default="weightfiles/TFM_Tokenizer_multiple_finetuned_on_TUAB", + help="Root folder of fine-tuned TFM classifier checkpoints (only with --model tfm).", + ) + parser.add_argument( + "--tokenizer-weights", + type=str, + default="weightfiles/tfm_tokenizer_last.pth", + help="Path to the pre-trained TFM tokenizer weights (only with --model tfm).", + ) + parser.add_argument( + "--split-type", + type=str, + default="patient", + choices=["patient", "sample"], + help="Split strategy: 'patient' (default, patient-level, no leakage) or " + "'sample' (original sample-level, for comparison).", + ) return parser.parse_args() +def _do_split(dataset, ratios, seed, split_type): + """Dispatch to the correct TUH split function based on split_type.""" + if split_type == "patient": + return split_by_patient_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + else: + return split_by_sample_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + + +def _load_tfm_weights(model, args, run_idx: int) -> None: + """Load pre-trained tokenizer + fine-tuned classifier for run_idx (0-based).""" + base = os.path.basename(args.weights_dir) + classifier_path = os.path.join(args.weights_dir, f"{base}_{run_idx + 1}", "best_model.pth") + print(f" Loading TFM weights (run {run_idx + 1}): {classifier_path}") + model.load_pretrained_weights( + tokenizer_checkpoint_path=args.tokenizer_weights, + classifier_checkpoint_path=classifier_path, + ) + + def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) @@ -130,19 +178,20 @@ def _run_one_seed( device: str, epochs: int, run_seed: int, + alphas: list, + run_idx: int = 0, ) -> dict: - """Train ContraWR + calibrate LABEL predictor for one seed. + """Train model + calibrate LABEL for one seed across all alphas. - The test set is passed in pre-built (fixed TUH eval partition). - Only train/val/cal (and model weights) vary per seed. + Training and base-model inference are done once; calibration loops over alphas (fast). - Returns a dict with keys: - accuracy, roc_auc, f1, coverage, miscoverage, avg_set_size + Returns {alpha: metrics_dict} where metrics_dict has keys: + accuracy, roc_auc_weighted_ovr, f1_weighted, coverage, miscoverage, avg_set_size """ set_seed(run_seed) - train_ds, val_ds, cal_ds, _ = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=run_seed + train_ds, val_ds, cal_ds, _ = _do_split( + sample_dataset, args.ratios, run_seed, args.split_type ) print(f" Split — Train: {len(train_ds)}, Val: {len(val_ds)}, " f"Cal: {len(cal_ds)}, Test: {len(test_ds)} (fixed)") @@ -153,62 +202,69 @@ def _run_one_seed( if len(val_ds) else None ) - print(" Training ContraWR...") - model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + if args.model == "tfm": + model = TFMTokenizer(dataset=sample_dataset).to(device) + _load_tfm_weights(model, args, run_idx) + else: + model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + print(" Training ContraWR...") + trainer_tmp = Trainer(model=model, device=device, enable_logging=False) + trainer_tmp.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=epochs, + monitor="accuracy" if val_loader is not None else None, + ) trainer = Trainer(model=model, device=device, enable_logging=False) - trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=epochs, - monitor="roc_auc" if val_loader is not None else None, - ) - # Base model metrics on fixed test set + # Base model metrics — computed once, shared across all alphas y_true_base, y_prob_base, _ = trainer.inference(test_loader) - base_metrics = get_metrics_fn("binary")( - y_true_base, y_prob_base, metrics=["accuracy", "roc_auc", "f1"] + base_metrics = get_metrics_fn("multiclass")( + y_true_base, y_prob_base, metrics=["accuracy", "roc_auc_weighted_ovr", "f1_weighted"] ) - # Conformal calibration + evaluation - print(" Calibrating LABEL predictor...") - label_predictor = LABEL(model=model, alpha=float(args.alpha)) - label_predictor.calibrate(cal_dataset=cal_ds) + # Calibration + evaluation — fast; loop over every alpha + results = {} + for alpha in alphas: + print(f" Calibrating LABEL predictor (alpha={alpha})...") + label_predictor = LABEL(model=model, alpha=float(alpha)) + label_predictor.calibrate(cal_dataset=cal_ds) - print(" Evaluating LABEL predictor on test set...") - y_true, y_prob, _, extra = Trainer(model=label_predictor).inference( - test_loader, additional_outputs=["y_predset"] - ) - conf_metrics = get_metrics_fn("binary")( - y_true, y_prob, - metrics=["accuracy", "miscoverage_ps"], - y_predset=extra["y_predset"], - ) + y_true, y_prob, _, extra = Trainer(model=label_predictor).inference( + test_loader, additional_outputs=["y_predset"] + ) + conf_metrics = get_metrics_fn("multiclass")( + y_true, y_prob, + metrics=["accuracy", "miscoverage_ps"], + y_predset=extra["y_predset"], + ) - predset = extra["y_predset"] - predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset - avg_set_size = predset_t.float().sum(dim=1).mean().item() + predset = extra["y_predset"] + predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset + avg_set_size = predset_t.float().sum(dim=1).mean().item() - miscoverage = conf_metrics["miscoverage_ps"] - if isinstance(miscoverage, np.ndarray): - miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) - else: - miscoverage = float(miscoverage) + miscoverage = conf_metrics["miscoverage_ps"] + if isinstance(miscoverage, np.ndarray): + miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) + else: + miscoverage = float(miscoverage) - return { - "accuracy": float(base_metrics["accuracy"]), - "roc_auc": float(base_metrics["roc_auc"]), - "f1": float(base_metrics["f1"]), - "coverage": 1.0 - miscoverage, - "miscoverage": miscoverage, - "avg_set_size": avg_set_size, - } + results[alpha] = { + "accuracy": float(base_metrics["accuracy"]), + "roc_auc_weighted_ovr": float(base_metrics["roc_auc_weighted_ovr"]), + "f1_weighted": float(base_metrics["f1_weighted"]), + "coverage": 1.0 - miscoverage, + "miscoverage": miscoverage, + "avg_set_size": avg_set_size, + } + return results def _print_single_run_results(metrics: dict, alpha: float) -> None: print("\nLABEL Results:") print(f" Accuracy: {metrics['accuracy']:.4f}") - print(f" ROC-AUC: {metrics['roc_auc']:.4f}") - print(f" F1: {metrics['f1']:.4f}") + print(f" ROC-AUC: {metrics['roc_auc_weighted_ovr']:.4f}") + print(f" F1: {metrics['f1_weighted']:.4f}") print(f" Empirical coverage: {metrics['coverage']:.4f}") print(f" Empirical miscoverage: {metrics['miscoverage']:.4f}") print(f" Average set size: {metrics['avg_set_size']:.2f}") @@ -219,8 +275,8 @@ def _print_multi_seed_summary( all_metrics: list, run_seeds: list, alpha: float, n_test: int ) -> None: accs = np.array([m["accuracy"] for m in all_metrics]) - roc_aucs = np.array([m["roc_auc"] for m in all_metrics]) - f1s = np.array([m["f1"] for m in all_metrics]) + roc_aucs = np.array([m["roc_auc_weighted_ovr"] for m in all_metrics]) + f1s = np.array([m["f1_weighted"] for m in all_metrics]) coverages = np.array([m["coverage"] for m in all_metrics]) miscovs = np.array([m["miscoverage"] for m in all_metrics]) set_sizes = np.array([m["avg_set_size"] for m in all_metrics]) @@ -235,7 +291,7 @@ def _print_multi_seed_summary( for i in range(n_runs): m = all_metrics[i] print(f" {i+1:<4} {run_seeds[i]:<6} {m['accuracy']:<10.4f} " - f"{m['roc_auc']:<10.4f} {m['f1']:<8.4f} {m['coverage']:<10.4f} " + f"{m['roc_auc_weighted_ovr']:<10.4f} {m['f1_weighted']:<8.4f} {m['coverage']:<10.4f} " f"{m['miscoverage']:<12.4f} {m['avg_set_size']:<12.2f}") print("\n" + "=" * 80) @@ -294,9 +350,14 @@ def _main(args: argparse.Namespace) -> None: print("\n" + "=" * 80) print("STEP 2: Extract fixed test set (TUH eval partition — same for all seeds)") print("=" * 80) - _, _, _, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=args.seed + _, _, _, test_ds = _do_split( + sample_dataset, args.ratios, args.seed, args.split_type ) + if len(test_ds) == 0 and args.quick_test: + print(" [quick-test] TUH eval partition empty in dev mode — using random 20% as test set.") + _, _, _, test_ds = split_by_sample_conformal( + dataset=sample_dataset, ratios=[0.6, 0.1, 0.1, 0.2], seed=args.seed + ) test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) print(f"Test: {len(test_ds)} (fixed)") @@ -308,36 +369,42 @@ def _main(args: argparse.Namespace) -> None: else: run_seeds = [args.seed + i for i in range(args.n_seeds)] + alphas = [float(a.strip()) for a in args.alphas.split(",")] if args.alphas else [args.alpha] + use_multi_seed = len(run_seeds) > 1 print(f"\nRun config: {'multi-seed (' + str(len(run_seeds)) + ' runs)' if use_multi_seed else 'single run'}") - print(f"Seeds: {run_seeds}, alpha={args.alpha}, target coverage={1 - args.alpha:.0%}") + print(f"Seeds: {run_seeds}, alphas={alphas}") # ------------------------------------------------------------------------- - # STEP 3+: Train + conformal (once per seed) + # STEP 3+: Train once per seed; calibrate for every alpha (fast) # ------------------------------------------------------------------------- - all_metrics = [] + all_metrics = {alpha: [] for alpha in alphas} for run_i, run_seed in enumerate(run_seeds): print("\n" + "=" * 80) if use_multi_seed: print(f"Run {run_i + 1} / {len(run_seeds)} (seed={run_seed})") else: - print(f"STEP 3–4: Train ContraWR + Conformal Calibration (seed={run_seed})") + print(f"STEP 3–4: Train + Conformal Calibration (seed={run_seed})") print("=" * 80) - metrics = _run_one_seed( - args, sample_dataset, test_ds, test_loader, device, epochs, run_seed + seed_results = _run_one_seed( + args, sample_dataset, test_ds, test_loader, device, epochs, run_seed, alphas, + run_idx=run_i, ) - all_metrics.append(metrics) + for alpha in alphas: + all_metrics[alpha].append(seed_results[alpha]) if use_multi_seed: - print(f" [Run {run_i + 1} result] " - f"acc={metrics['accuracy']:.4f}, roc_auc={metrics['roc_auc']:.4f}, " - f"cov={metrics['coverage']:.4f}, set_size={metrics['avg_set_size']:.2f}") - - if not use_multi_seed: - _print_single_run_results(all_metrics[0], args.alpha) - else: - _print_multi_seed_summary(all_metrics, run_seeds, args.alpha, len(test_ds)) + m = seed_results[alphas[0]] + print(f" [Run {run_i + 1} result (alpha={alphas[0]})] " + f"acc={m['accuracy']:.4f}, roc_auc={m['roc_auc_weighted_ovr']:.4f}, " + f"cov={m['coverage']:.4f}, set_size={m['avg_set_size']:.2f}") + + for alpha in alphas: + if not use_multi_seed: + _print_single_run_results(all_metrics[alpha][0], alpha) + else: + _print_multi_seed_summary(all_metrics[alpha], run_seeds, alpha, len(test_ds)) def main() -> None: diff --git a/examples/conformal_eeg/tuab_covariate_shift_conformal.py b/examples/conformal_eeg/tuab_covariate_shift_conformal.py index 4bb200827..33a810ab1 100644 --- a/examples/conformal_eeg/tuab_covariate_shift_conformal.py +++ b/examples/conformal_eeg/tuab_covariate_shift_conformal.py @@ -24,6 +24,7 @@ from __future__ import annotations import argparse +import os import random import sys from pathlib import Path @@ -51,8 +52,8 @@ def flush(self): from pyhealth.calib.predictionset.covariate import CovariateLabel from pyhealth.calib.utils import extract_embeddings -from pyhealth.datasets import TUABDataset, get_dataloader, split_by_sample_conformal_tuh -from pyhealth.models import ContraWR +from pyhealth.datasets import TUABDataset, get_dataloader, split_by_patient_conformal_tuh, split_by_sample_conformal_tuh, split_by_sample_conformal +from pyhealth.models import ContraWR, TFMTokenizer from pyhealth.tasks import EEGAbnormalTUAB from pyhealth.trainer import Trainer, get_metrics_fn @@ -64,7 +65,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--root", type=str, - default="downloads/tuab/v3.0.0/edf", + default="/srv/local/data/TUH/tuh_eeg_abnormal/v3.0.0/edf", help="Path to TUAB edf/ folder.", ) parser.add_argument("--subset", type=str, default="both", choices=["train", "eval", "both"]) @@ -80,6 +81,10 @@ def parse_args() -> argparse.Namespace: "--alpha", type=float, default=0.1, help="Miscoverage rate (e.g., 0.1 => 90% target coverage).", ) + parser.add_argument( + "--alphas", type=str, default=None, + help="Comma-separated miscoverage rates, e.g. '0.2,0.1,0.05,0.01'. Overrides --alpha.", + ) parser.add_argument( "--ratios", type=float, @@ -90,6 +95,10 @@ def parse_args() -> argparse.Namespace: "Must sum to 1.0. Test is fixed as the TUH eval partition.", ) parser.add_argument("--n-fft", type=int, default=128, help="STFT FFT size used by ContraWR.") + parser.add_argument( + "--model", type=str, default="contrawr", choices=["contrawr", "tfm"], + help="Backbone model: 'contrawr' (default) or 'tfm' (TFMTokenizer).", + ) parser.add_argument( "--device", type=str, default=None, help="Device string, e.g. 'cuda:0' or 'cpu'. Defaults to auto-detect.", @@ -117,9 +126,48 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Smoke test: dev=True, max 2000 samples, 2 epochs.", ) + parser.add_argument( + "--weights-dir", + type=str, + default="weightfiles/TFM_Tokenizer_multiple_finetuned_on_TUAB", + help="Root folder of fine-tuned TFM classifier checkpoints (only with --model tfm).", + ) + parser.add_argument( + "--tokenizer-weights", + type=str, + default="weightfiles/tfm_tokenizer_last.pth", + help="Path to the pre-trained TFM tokenizer weights (only with --model tfm).", + ) + parser.add_argument( + "--split-type", + type=str, + default="patient", + choices=["patient", "sample"], + help="Split strategy: 'patient' (default, patient-level, no leakage) or " + "'sample' (original sample-level, for comparison).", + ) return parser.parse_args() +def _do_split(dataset, ratios, seed, split_type): + """Dispatch to the correct TUH split function based on split_type.""" + if split_type == "patient": + return split_by_patient_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + else: + return split_by_sample_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + + +def _load_tfm_weights(model, args, run_idx: int) -> None: + """Load pre-trained tokenizer + fine-tuned classifier for run_idx (0-based).""" + base = os.path.basename(args.weights_dir) + classifier_path = os.path.join(args.weights_dir, f"{base}_{run_idx + 1}", "best_model.pth") + print(f" Loading TFM weights (run {run_idx + 1}): {classifier_path}") + model.load_pretrained_weights( + tokenizer_checkpoint_path=args.tokenizer_weights, + classifier_checkpoint_path=classifier_path, + ) + + def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) @@ -136,20 +184,21 @@ def _run_one_seed( device: str, epochs: int, run_seed: int, + alphas: list, + run_idx: int = 0, ) -> dict: - """Train ContraWR + calibrate CovariateLabel predictor for one seed. + """Train model + calibrate CovariateLabel for one seed across all alphas. - The test set is passed in pre-built (fixed TUH eval partition). - Only train/val/cal (and model weights) vary per seed. - Test embeddings are recomputed each seed because the model changes. + Training, embedding extraction, and base inference are done once; calibration + loops over alphas (fast — only likelihood-ratio weights/threshold recomputed). - Returns a dict with keys: - accuracy, roc_auc, f1, coverage, miscoverage, avg_set_size + Returns {alpha: metrics_dict} where metrics_dict has keys: + accuracy, roc_auc_weighted_ovr, f1_weighted, coverage, miscoverage, avg_set_size """ set_seed(run_seed) - train_ds, val_ds, cal_ds, _ = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=run_seed + train_ds, val_ds, cal_ds, _ = _do_split( + sample_dataset, args.ratios, run_seed, args.split_type ) print(f" Split — Train: {len(train_ds)}, Val: {len(val_ds)}, " f"Cal: {len(cal_ds)}, Test: {len(test_ds)} (fixed)") @@ -160,71 +209,78 @@ def _run_one_seed( if len(val_ds) else None ) - print(" Training ContraWR...") - model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + if args.model == "tfm": + model = TFMTokenizer(dataset=sample_dataset).to(device) + _load_tfm_weights(model, args, run_idx) + else: + model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + print(" Training ContraWR...") + trainer_tmp = Trainer(model=model, device=device, enable_logging=False) + trainer_tmp.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=epochs, + monitor="accuracy" if val_loader is not None else None, + ) trainer = Trainer(model=model, device=device, enable_logging=False) - trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=epochs, - monitor="roc_auc" if val_loader is not None else None, - ) - # Base model metrics on fixed test set + # Base model metrics — computed once, shared across all alphas y_true_base, y_prob_base, _ = trainer.inference(test_loader) - base_metrics = get_metrics_fn("binary")( - y_true_base, y_prob_base, metrics=["accuracy", "roc_auc", "f1"] + base_metrics = get_metrics_fn("multiclass")( + y_true_base, y_prob_base, metrics=["accuracy", "roc_auc_weighted_ovr", "f1_weighted"] ) - # Extract embeddings — both depend on the current model so must redo each seed + # Extract embeddings once — reused for every alpha print(" Extracting embeddings for calibration and test splits...") cal_embeddings = extract_embeddings(model, cal_ds, batch_size=args.batch_size, device=device) test_embeddings = extract_embeddings(model, test_ds, batch_size=args.batch_size, device=device) - # Conformal calibration + evaluation - print(" Calibrating CovariateLabel predictor...") - cov_predictor = CovariateLabel(model=model, alpha=float(args.alpha)) - cov_predictor.calibrate( - cal_dataset=cal_ds, - cal_embeddings=cal_embeddings, - test_embeddings=test_embeddings, - ) + # Calibration + evaluation — fast; loop over every alpha + results = {} + for alpha in alphas: + print(f" Calibrating CovariateLabel predictor (alpha={alpha})...") + cov_predictor = CovariateLabel(model=model, alpha=float(alpha)) + cov_predictor.calibrate( + cal_dataset=cal_ds, + cal_embeddings=cal_embeddings, + test_embeddings=test_embeddings, + ) - print(" Evaluating CovariateLabel predictor on test set...") - y_true, y_prob, _, extra = Trainer(model=cov_predictor).inference( - test_loader, additional_outputs=["y_predset"] - ) - conf_metrics = get_metrics_fn("binary")( - y_true, y_prob, - metrics=["accuracy", "miscoverage_ps"], - y_predset=extra["y_predset"], - ) + y_true, y_prob, _, extra = Trainer(model=cov_predictor).inference( + test_loader, additional_outputs=["y_predset"] + ) + conf_metrics = get_metrics_fn("multiclass")( + y_true, y_prob, + metrics=["accuracy", "miscoverage_ps"], + y_predset=extra["y_predset"], + ) - predset = extra["y_predset"] - predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset - avg_set_size = predset_t.float().sum(dim=1).mean().item() + predset = extra["y_predset"] + predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset + avg_set_size = predset_t.float().sum(dim=1).mean().item() - miscoverage = conf_metrics["miscoverage_ps"] - if isinstance(miscoverage, np.ndarray): - miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) - else: - miscoverage = float(miscoverage) + miscoverage = conf_metrics["miscoverage_ps"] + if isinstance(miscoverage, np.ndarray): + miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) + else: + miscoverage = float(miscoverage) - return { - "accuracy": float(base_metrics["accuracy"]), - "roc_auc": float(base_metrics["roc_auc"]), - "f1": float(base_metrics["f1"]), - "coverage": 1.0 - miscoverage, - "miscoverage": miscoverage, - "avg_set_size": avg_set_size, - } + results[alpha] = { + "accuracy": float(base_metrics["accuracy"]), + "roc_auc_weighted_ovr": float(base_metrics["roc_auc_weighted_ovr"]), + "f1_weighted": float(base_metrics["f1_weighted"]), + "coverage": 1.0 - miscoverage, + "miscoverage": miscoverage, + "avg_set_size": avg_set_size, + } + return results def _print_single_run_results(metrics: dict, alpha: float) -> None: print("\nCovariateLabel Results:") print(f" Accuracy: {metrics['accuracy']:.4f}") - print(f" ROC-AUC: {metrics['roc_auc']:.4f}") - print(f" F1: {metrics['f1']:.4f}") + print(f" ROC-AUC: {metrics['roc_auc_weighted_ovr']:.4f}") + print(f" F1: {metrics['f1_weighted']:.4f}") print(f" Empirical coverage: {metrics['coverage']:.4f}") print(f" Empirical miscoverage: {metrics['miscoverage']:.4f}") print(f" Average set size: {metrics['avg_set_size']:.2f}") @@ -235,8 +291,8 @@ def _print_multi_seed_summary( all_metrics: list, run_seeds: list, alpha: float, n_test: int ) -> None: accs = np.array([m["accuracy"] for m in all_metrics]) - roc_aucs = np.array([m["roc_auc"] for m in all_metrics]) - f1s = np.array([m["f1"] for m in all_metrics]) + roc_aucs = np.array([m["roc_auc_weighted_ovr"] for m in all_metrics]) + f1s = np.array([m["f1_weighted"] for m in all_metrics]) coverages = np.array([m["coverage"] for m in all_metrics]) miscovs = np.array([m["miscoverage"] for m in all_metrics]) set_sizes = np.array([m["avg_set_size"] for m in all_metrics]) @@ -251,7 +307,7 @@ def _print_multi_seed_summary( for i in range(n_runs): m = all_metrics[i] print(f" {i+1:<4} {run_seeds[i]:<6} {m['accuracy']:<10.4f} " - f"{m['roc_auc']:<10.4f} {m['f1']:<8.4f} {m['coverage']:<10.4f} " + f"{m['roc_auc_weighted_ovr']:<10.4f} {m['f1_weighted']:<8.4f} {m['coverage']:<10.4f} " f"{m['miscoverage']:<12.4f} {m['avg_set_size']:<12.2f}") print("\n" + "=" * 80) @@ -310,9 +366,14 @@ def _main(args: argparse.Namespace) -> None: print("\n" + "=" * 80) print("STEP 2: Extract fixed test set (TUH eval partition — same for all seeds)") print("=" * 80) - _, _, _, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=args.seed + _, _, _, test_ds = _do_split( + sample_dataset, args.ratios, args.seed, args.split_type ) + if len(test_ds) == 0 and args.quick_test: + print(" [quick-test] TUH eval partition empty in dev mode — using random 20% as test set.") + _, _, _, test_ds = split_by_sample_conformal( + dataset=sample_dataset, ratios=[0.6, 0.1, 0.1, 0.2], seed=args.seed + ) test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) print(f"Test: {len(test_ds)} (fixed)") @@ -324,36 +385,42 @@ def _main(args: argparse.Namespace) -> None: else: run_seeds = [args.seed + i for i in range(args.n_seeds)] + alphas = [float(a.strip()) for a in args.alphas.split(",")] if args.alphas else [args.alpha] + use_multi_seed = len(run_seeds) > 1 print(f"\nRun config: {'multi-seed (' + str(len(run_seeds)) + ' runs)' if use_multi_seed else 'single run'}") - print(f"Seeds: {run_seeds}, alpha={args.alpha}, target coverage={1 - args.alpha:.0%}") + print(f"Seeds: {run_seeds}, alphas={alphas}") # ------------------------------------------------------------------------- - # STEP 3+: Train + conformal (once per seed) + # STEP 3+: Train once per seed; calibrate for every alpha (fast) # ------------------------------------------------------------------------- - all_metrics = [] + all_metrics = {alpha: [] for alpha in alphas} for run_i, run_seed in enumerate(run_seeds): print("\n" + "=" * 80) if use_multi_seed: print(f"Run {run_i + 1} / {len(run_seeds)} (seed={run_seed})") else: - print(f"STEP 3–4: Train ContraWR + Conformal Calibration (seed={run_seed})") + print(f"STEP 3–4: Train + Conformal Calibration (seed={run_seed})") print("=" * 80) - metrics = _run_one_seed( - args, sample_dataset, test_ds, test_loader, device, epochs, run_seed + seed_results = _run_one_seed( + args, sample_dataset, test_ds, test_loader, device, epochs, run_seed, alphas, + run_idx=run_i, ) - all_metrics.append(metrics) + for alpha in alphas: + all_metrics[alpha].append(seed_results[alpha]) if use_multi_seed: - print(f" [Run {run_i + 1} result] " - f"acc={metrics['accuracy']:.4f}, roc_auc={metrics['roc_auc']:.4f}, " - f"cov={metrics['coverage']:.4f}, set_size={metrics['avg_set_size']:.2f}") - - if not use_multi_seed: - _print_single_run_results(all_metrics[0], args.alpha) - else: - _print_multi_seed_summary(all_metrics, run_seeds, args.alpha, len(test_ds)) + m = seed_results[alphas[0]] + print(f" [Run {run_i + 1} result (alpha={alphas[0]})] " + f"acc={m['accuracy']:.4f}, roc_auc={m['roc_auc_weighted_ovr']:.4f}, " + f"cov={m['coverage']:.4f}, set_size={m['avg_set_size']:.2f}") + + for alpha in alphas: + if not use_multi_seed: + _print_single_run_results(all_metrics[alpha][0], alpha) + else: + _print_multi_seed_summary(all_metrics[alpha], run_seeds, alpha, len(test_ds)) def main() -> None: diff --git a/examples/conformal_eeg/tuab_kmeans_conformal.py b/examples/conformal_eeg/tuab_kmeans_conformal.py index 35669dcf2..67152bfff 100644 --- a/examples/conformal_eeg/tuab_kmeans_conformal.py +++ b/examples/conformal_eeg/tuab_kmeans_conformal.py @@ -23,6 +23,7 @@ from __future__ import annotations import argparse +import os import random import sys from pathlib import Path @@ -50,8 +51,8 @@ def flush(self): from pyhealth.calib.predictionset.cluster import ClusterLabel from pyhealth.calib.utils import extract_embeddings -from pyhealth.datasets import TUABDataset, get_dataloader, split_by_sample_conformal_tuh -from pyhealth.models import ContraWR +from pyhealth.datasets import TUABDataset, get_dataloader, split_by_patient_conformal_tuh, split_by_sample_conformal_tuh, split_by_sample_conformal +from pyhealth.models import ContraWR, TFMTokenizer from pyhealth.tasks import EEGAbnormalTUAB from pyhealth.trainer import Trainer, get_metrics_fn @@ -63,7 +64,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--root", type=str, - default="downloads/tuab/v3.0.0/edf", + default="/srv/local/data/TUH/tuh_eeg_abnormal/v3.0.0/edf", help="Path to TUAB edf/ folder.", ) parser.add_argument("--subset", type=str, default="both", choices=["train", "eval", "both"]) @@ -79,6 +80,10 @@ def parse_args() -> argparse.Namespace: "--alpha", type=float, default=0.1, help="Miscoverage rate (e.g., 0.1 => 90% target coverage).", ) + parser.add_argument( + "--alphas", type=str, default=None, + help="Comma-separated miscoverage rates, e.g. '0.2,0.1,0.05,0.01'. Overrides --alpha.", + ) parser.add_argument( "--ratios", type=float, @@ -93,6 +98,10 @@ def parse_args() -> argparse.Namespace: help="Number of K-means clusters for cluster-specific thresholds.", ) parser.add_argument("--n-fft", type=int, default=128, help="STFT FFT size used by ContraWR.") + parser.add_argument( + "--model", type=str, default="contrawr", choices=["contrawr", "tfm"], + help="Backbone model: 'contrawr' (default) or 'tfm' (TFMTokenizer).", + ) parser.add_argument( "--device", type=str, default=None, help="Device string, e.g. 'cuda:0' or 'cpu'. Defaults to auto-detect.", @@ -120,9 +129,48 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Smoke test: dev=True, max 2000 samples, 2 epochs.", ) + parser.add_argument( + "--weights-dir", + type=str, + default="weightfiles/TFM_Tokenizer_multiple_finetuned_on_TUAB", + help="Root folder of fine-tuned TFM classifier checkpoints (only with --model tfm).", + ) + parser.add_argument( + "--tokenizer-weights", + type=str, + default="weightfiles/tfm_tokenizer_last.pth", + help="Path to the pre-trained TFM tokenizer weights (only with --model tfm).", + ) + parser.add_argument( + "--split-type", + type=str, + default="patient", + choices=["patient", "sample"], + help="Split strategy: 'patient' (default, patient-level, no leakage) or " + "'sample' (original sample-level, for comparison).", + ) return parser.parse_args() +def _do_split(dataset, ratios, seed, split_type): + """Dispatch to the correct TUH split function based on split_type.""" + if split_type == "patient": + return split_by_patient_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + else: + return split_by_sample_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + + +def _load_tfm_weights(model, args, run_idx: int) -> None: + """Load pre-trained tokenizer + fine-tuned classifier for run_idx (0-based).""" + base = os.path.basename(args.weights_dir) + classifier_path = os.path.join(args.weights_dir, f"{base}_{run_idx + 1}", "best_model.pth") + print(f" Loading TFM weights (run {run_idx + 1}): {classifier_path}") + model.load_pretrained_weights( + tokenizer_checkpoint_path=args.tokenizer_weights, + classifier_checkpoint_path=classifier_path, + ) + + def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) @@ -139,19 +187,21 @@ def _run_one_seed( device: str, epochs: int, run_seed: int, + alphas: list, + run_idx: int = 0, ) -> dict: - """Train ContraWR + calibrate ClusterLabel predictor for one seed. + """Train model + calibrate ClusterLabel for one seed across all alphas. - The test set is passed in pre-built (fixed TUH eval partition). - Only train/val/cal (and model weights) vary per seed. + Training, embedding extraction, and base inference are done once; calibration + loops over alphas (fast — only threshold recomputed per alpha). - Returns a dict with keys: - accuracy, roc_auc, f1, coverage, miscoverage, avg_set_size + Returns {alpha: metrics_dict} where metrics_dict has keys: + accuracy, roc_auc_weighted_ovr, f1_weighted, coverage, miscoverage, avg_set_size """ set_seed(run_seed) - train_ds, val_ds, cal_ds, _ = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=run_seed + train_ds, val_ds, cal_ds, _ = _do_split( + sample_dataset, args.ratios, run_seed, args.split_type ) print(f" Split — Train: {len(train_ds)}, Val: {len(val_ds)}, " f"Cal: {len(cal_ds)}, Test: {len(test_ds)} (fixed)") @@ -162,76 +212,83 @@ def _run_one_seed( if len(val_ds) else None ) - print(" Training ContraWR...") - model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + if args.model == "tfm": + model = TFMTokenizer(dataset=sample_dataset).to(device) + _load_tfm_weights(model, args, run_idx) + else: + model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + print(" Training ContraWR...") + trainer_tmp = Trainer(model=model, device=device, enable_logging=False) + trainer_tmp.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=epochs, + monitor="accuracy" if val_loader is not None else None, + ) trainer = Trainer(model=model, device=device, enable_logging=False) - trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=epochs, - monitor="roc_auc" if val_loader is not None else None, - ) - # Base model metrics on fixed test set + # Base model metrics — computed once, shared across all alphas y_true_base, y_prob_base, _ = trainer.inference(test_loader) - base_metrics = get_metrics_fn("binary")( - y_true_base, y_prob_base, metrics=["accuracy", "roc_auc", "f1"] + base_metrics = get_metrics_fn("multiclass")( + y_true_base, y_prob_base, metrics=["accuracy", "roc_auc_weighted_ovr", "f1_weighted"] ) - # Extract embeddings for train and cal splits (model-dependent, must redo each seed) + # Extract embeddings once — reused for every alpha print(" Extracting embeddings for train and calibration splits...") train_embeddings = extract_embeddings(model, train_ds, batch_size=args.batch_size, device=device) cal_embeddings = extract_embeddings(model, cal_ds, batch_size=args.batch_size, device=device) - # Conformal calibration + evaluation - print(" Calibrating ClusterLabel predictor...") - cluster_predictor = ClusterLabel( - model=model, - alpha=float(args.alpha), - n_clusters=args.n_clusters, - random_state=run_seed, - ) - cluster_predictor.calibrate( - cal_dataset=cal_ds, - train_embeddings=train_embeddings, - cal_embeddings=cal_embeddings, - ) + # Calibration + evaluation — fast; loop over every alpha + results = {} + for alpha in alphas: + print(f" Calibrating ClusterLabel predictor (alpha={alpha})...") + cluster_predictor = ClusterLabel( + model=model, + alpha=float(alpha), + n_clusters=args.n_clusters, + random_state=run_seed, + ) + cluster_predictor.calibrate( + cal_dataset=cal_ds, + train_embeddings=train_embeddings, + cal_embeddings=cal_embeddings, + ) - print(" Evaluating ClusterLabel predictor on test set...") - y_true, y_prob, _, extra = Trainer(model=cluster_predictor).inference( - test_loader, additional_outputs=["y_predset"] - ) - conf_metrics = get_metrics_fn("binary")( - y_true, y_prob, - metrics=["accuracy", "miscoverage_ps"], - y_predset=extra["y_predset"], - ) + y_true, y_prob, _, extra = Trainer(model=cluster_predictor).inference( + test_loader, additional_outputs=["y_predset"] + ) + conf_metrics = get_metrics_fn("multiclass")( + y_true, y_prob, + metrics=["accuracy", "miscoverage_ps"], + y_predset=extra["y_predset"], + ) - predset = extra["y_predset"] - predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset - avg_set_size = predset_t.float().sum(dim=1).mean().item() + predset = extra["y_predset"] + predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset + avg_set_size = predset_t.float().sum(dim=1).mean().item() - miscoverage = conf_metrics["miscoverage_ps"] - if isinstance(miscoverage, np.ndarray): - miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) - else: - miscoverage = float(miscoverage) + miscoverage = conf_metrics["miscoverage_ps"] + if isinstance(miscoverage, np.ndarray): + miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) + else: + miscoverage = float(miscoverage) - return { - "accuracy": float(base_metrics["accuracy"]), - "roc_auc": float(base_metrics["roc_auc"]), - "f1": float(base_metrics["f1"]), - "coverage": 1.0 - miscoverage, - "miscoverage": miscoverage, - "avg_set_size": avg_set_size, - } + results[alpha] = { + "accuracy": float(base_metrics["accuracy"]), + "roc_auc_weighted_ovr": float(base_metrics["roc_auc_weighted_ovr"]), + "f1_weighted": float(base_metrics["f1_weighted"]), + "coverage": 1.0 - miscoverage, + "miscoverage": miscoverage, + "avg_set_size": avg_set_size, + } + return results def _print_single_run_results(metrics: dict, alpha: float, n_clusters: int) -> None: print("\nClusterLabel Results:") print(f" Accuracy: {metrics['accuracy']:.4f}") - print(f" ROC-AUC: {metrics['roc_auc']:.4f}") - print(f" F1: {metrics['f1']:.4f}") + print(f" ROC-AUC: {metrics['roc_auc_weighted_ovr']:.4f}") + print(f" F1: {metrics['f1_weighted']:.4f}") print(f" Empirical coverage: {metrics['coverage']:.4f}") print(f" Empirical miscoverage: {metrics['miscoverage']:.4f}") print(f" Average set size: {metrics['avg_set_size']:.2f}") @@ -243,8 +300,8 @@ def _print_multi_seed_summary( all_metrics: list, run_seeds: list, alpha: float, n_test: int, n_clusters: int ) -> None: accs = np.array([m["accuracy"] for m in all_metrics]) - roc_aucs = np.array([m["roc_auc"] for m in all_metrics]) - f1s = np.array([m["f1"] for m in all_metrics]) + roc_aucs = np.array([m["roc_auc_weighted_ovr"] for m in all_metrics]) + f1s = np.array([m["f1_weighted"] for m in all_metrics]) coverages = np.array([m["coverage"] for m in all_metrics]) miscovs = np.array([m["miscoverage"] for m in all_metrics]) set_sizes = np.array([m["avg_set_size"] for m in all_metrics]) @@ -259,7 +316,7 @@ def _print_multi_seed_summary( for i in range(n_runs): m = all_metrics[i] print(f" {i+1:<4} {run_seeds[i]:<6} {m['accuracy']:<10.4f} " - f"{m['roc_auc']:<10.4f} {m['f1']:<8.4f} {m['coverage']:<10.4f} " + f"{m['roc_auc_weighted_ovr']:<10.4f} {m['f1_weighted']:<8.4f} {m['coverage']:<10.4f} " f"{m['miscoverage']:<12.4f} {m['avg_set_size']:<12.2f}") print("\n" + "=" * 80) @@ -319,9 +376,14 @@ def _main(args: argparse.Namespace) -> None: print("\n" + "=" * 80) print("STEP 2: Extract fixed test set (TUH eval partition — same for all seeds)") print("=" * 80) - _, _, _, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=args.seed + _, _, _, test_ds = _do_split( + sample_dataset, args.ratios, args.seed, args.split_type ) + if len(test_ds) == 0 and args.quick_test: + print(" [quick-test] TUH eval partition empty in dev mode — using random 20% as test set.") + _, _, _, test_ds = split_by_sample_conformal( + dataset=sample_dataset, ratios=[0.6, 0.1, 0.1, 0.2], seed=args.seed + ) test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) print(f"Test: {len(test_ds)} (fixed)") @@ -333,37 +395,42 @@ def _main(args: argparse.Namespace) -> None: else: run_seeds = [args.seed + i for i in range(args.n_seeds)] + alphas = [float(a.strip()) for a in args.alphas.split(",")] if args.alphas else [args.alpha] + use_multi_seed = len(run_seeds) > 1 print(f"\nRun config: {'multi-seed (' + str(len(run_seeds)) + ' runs)' if use_multi_seed else 'single run'}") - print(f"Seeds: {run_seeds}, alpha={args.alpha}, n_clusters={args.n_clusters}, " - f"target coverage={1 - args.alpha:.0%}") + print(f"Seeds: {run_seeds}, alphas={alphas}, n_clusters={args.n_clusters}") # ------------------------------------------------------------------------- - # STEP 3+: Train + conformal (once per seed) + # STEP 3+: Train once per seed; calibrate for every alpha (fast) # ------------------------------------------------------------------------- - all_metrics = [] + all_metrics = {alpha: [] for alpha in alphas} for run_i, run_seed in enumerate(run_seeds): print("\n" + "=" * 80) if use_multi_seed: print(f"Run {run_i + 1} / {len(run_seeds)} (seed={run_seed})") else: - print(f"STEP 3–4: Train ContraWR + Conformal Calibration (seed={run_seed})") + print(f"STEP 3–4: Train + Conformal Calibration (seed={run_seed})") print("=" * 80) - metrics = _run_one_seed( - args, sample_dataset, test_ds, test_loader, device, epochs, run_seed + seed_results = _run_one_seed( + args, sample_dataset, test_ds, test_loader, device, epochs, run_seed, alphas, + run_idx=run_i, ) - all_metrics.append(metrics) + for alpha in alphas: + all_metrics[alpha].append(seed_results[alpha]) if use_multi_seed: - print(f" [Run {run_i + 1} result] " - f"acc={metrics['accuracy']:.4f}, roc_auc={metrics['roc_auc']:.4f}, " - f"cov={metrics['coverage']:.4f}, set_size={metrics['avg_set_size']:.2f}") - - if not use_multi_seed: - _print_single_run_results(all_metrics[0], args.alpha, args.n_clusters) - else: - _print_multi_seed_summary(all_metrics, run_seeds, args.alpha, len(test_ds), args.n_clusters) + m = seed_results[alphas[0]] + print(f" [Run {run_i + 1} result (alpha={alphas[0]})] " + f"acc={m['accuracy']:.4f}, roc_auc={m['roc_auc_weighted_ovr']:.4f}, " + f"cov={m['coverage']:.4f}, set_size={m['avg_set_size']:.2f}") + + for alpha in alphas: + if not use_multi_seed: + _print_single_run_results(all_metrics[alpha][0], alpha, args.n_clusters) + else: + _print_multi_seed_summary(all_metrics[alpha], run_seeds, alpha, len(test_ds), args.n_clusters) def main() -> None: diff --git a/examples/conformal_eeg/tuab_ncp_conformal.py b/examples/conformal_eeg/tuab_ncp_conformal.py index ebde5242d..be658d794 100644 --- a/examples/conformal_eeg/tuab_ncp_conformal.py +++ b/examples/conformal_eeg/tuab_ncp_conformal.py @@ -20,6 +20,7 @@ from __future__ import annotations import argparse +import os import random import sys from pathlib import Path @@ -47,8 +48,8 @@ def flush(self): from pyhealth.calib.predictionset.cluster import NeighborhoodLabel from pyhealth.calib.utils import extract_embeddings -from pyhealth.datasets import TUABDataset, get_dataloader, split_by_sample_conformal_tuh -from pyhealth.models import ContraWR +from pyhealth.datasets import TUABDataset, get_dataloader, split_by_patient_conformal_tuh, split_by_sample_conformal_tuh, split_by_sample_conformal +from pyhealth.models import ContraWR, TFMTokenizer from pyhealth.tasks import EEGAbnormalTUAB from pyhealth.trainer import Trainer, get_metrics_fn @@ -86,6 +87,10 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--epochs", type=int, default=20) parser.add_argument("--alpha", type=float, default=0.1, help="Miscoverage rate (e.g., 0.1 => 90% target coverage).") + parser.add_argument( + "--alphas", type=str, default=None, + help="Comma-separated miscoverage rates, e.g. '0.2,0.1,0.05,0.01'. Overrides --alpha.", + ) parser.add_argument( "--ratios", type=float, @@ -107,6 +112,10 @@ def parse_args() -> argparse.Namespace: help="Temperature for NCP exponential weights; smaller => more localization.", ) parser.add_argument("--n-fft", type=int, default=128, help="STFT FFT size used by ContraWR.") + parser.add_argument( + "--model", type=str, default="contrawr", choices=["contrawr", "tfm"], + help="Backbone model: 'contrawr' (default) or 'tfm' (TFMTokenizer).", + ) parser.add_argument( "--device", type=str, @@ -124,9 +133,48 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Smoke test: dev=True, max 2000 samples, 2 epochs, ~5-10 min.", ) + parser.add_argument( + "--weights-dir", + type=str, + default="weightfiles/TFM_Tokenizer_multiple_finetuned_on_TUAB", + help="Root folder of fine-tuned TFM classifier checkpoints (only with --model tfm).", + ) + parser.add_argument( + "--tokenizer-weights", + type=str, + default="weightfiles/tfm_tokenizer_last.pth", + help="Path to the pre-trained TFM tokenizer weights (only with --model tfm).", + ) + parser.add_argument( + "--split-type", + type=str, + default="patient", + choices=["patient", "sample"], + help="Split strategy: 'patient' (default, patient-level, no leakage) or " + "'sample' (original sample-level, for comparison).", + ) return parser.parse_args() +def _do_split(dataset, ratios, seed, split_type): + """Dispatch to the correct TUH split function based on split_type.""" + if split_type == "patient": + return split_by_patient_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + else: + return split_by_sample_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + + +def _load_tfm_weights(model, args, run_idx: int) -> None: + """Load pre-trained tokenizer + fine-tuned classifier for run_idx (0-based).""" + base = os.path.basename(args.weights_dir) + classifier_path = os.path.join(args.weights_dir, f"{base}_{run_idx + 1}", "best_model.pth") + print(f" Loading TFM weights (run {run_idx + 1}): {classifier_path}") + model.load_pretrained_weights( + tokenizer_checkpoint_path=args.tokenizer_weights, + classifier_checkpoint_path=classifier_path, + ) + + def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) @@ -135,15 +183,13 @@ def set_seed(seed: int) -> None: torch.cuda.manual_seed_all(seed) -def _split_train_pool_for_run(sample_dataset, ratios, run_seed): +def _split_train_pool_for_run(sample_dataset, ratios, run_seed, split_type="patient"): """Re-split the TUH train partition into train/val/cal for one run seed. The test set (TUH eval partition) is always fixed regardless of seed, so only train/val/cal change across runs in multi-seed mode. """ - train_ds, val_ds, cal_ds, _ = split_by_sample_conformal_tuh( - sample_dataset, ratios=ratios, seed=run_seed - ) + train_ds, val_ds, cal_ds, _ = _do_split(sample_dataset, ratios, run_seed, split_type) return train_ds, val_ds, cal_ds @@ -156,87 +202,82 @@ def _run_one_ncp( args, device, epochs, - return_metrics=False, + alphas: list, + run_idx: int = 0, ): - """Train ContraWR, calibrate NCP, evaluate on test. Optionally return metrics dict for aggregation.""" + """Train model + calibrate NCP for one seed across all alphas. + + Training, embedding extraction, and base inference are done once; calibration + loops over alphas (fast — only threshold recomputed per alpha). + + Returns {alpha: metrics_dict} where metrics_dict has keys: + accuracy, roc_auc_weighted_ovr, f1_weighted, coverage, miscoverage, avg_set_size + """ 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) if len(val_ds) else None - print("\n" + "=" * 80) - print("STEP 3: Train ContraWR") - print("=" * 80) - model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + if args.model == "tfm": + model = TFMTokenizer(dataset=sample_dataset).to(device) + _load_tfm_weights(model, args, run_idx) + else: + model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + print(" Training ContraWR...") + trainer_tmp = Trainer(model=model, device=device, enable_logging=False) + trainer_tmp.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=epochs, + monitor="accuracy" if val_loader is not None else None, + ) trainer = Trainer(model=model, device=device, enable_logging=False) - trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=epochs, - monitor="roc_auc" if val_loader is not None else None, - ) - print("\nBase model performance on test set:") - y_true_base, y_prob_base, _loss_base = trainer.inference(test_loader) - base_metrics = get_metrics_fn("binary")( - y_true_base, y_prob_base, metrics=["accuracy", "roc_auc", "f1"] + # Base model metrics — computed once, shared across all alphas + y_true_base, y_prob_base, _ = trainer.inference(test_loader) + base_metrics = get_metrics_fn("multiclass")( + y_true_base, y_prob_base, metrics=["accuracy", "roc_auc_weighted_ovr", "f1_weighted"] ) - for metric, value in base_metrics.items(): - print(f" {metric}: {value:.4f}") - - print("\n" + "=" * 80) - print("STEP 4: Neighborhood Conformal Prediction (NCP / NeighborhoodLabel)") - print("=" * 80) - print(f"Target miscoverage alpha: {args.alpha} (target coverage {1 - args.alpha:.0%})") - print(f"k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") + # Extract calibration embeddings once — reused for every alpha + print(" Extracting calibration embeddings...") cal_embeddings = extract_embeddings(model, cal_ds, batch_size=args.batch_size, device=device) - if not return_metrics: - print(f" cal_embeddings shape: {cal_embeddings.shape}") - - ncp_predictor = NeighborhoodLabel( - model=model, - alpha=float(args.alpha), - k_neighbors=args.k_neighbors, - lambda_L=args.lambda_L, - ) - ncp_predictor.calibrate(cal_dataset=cal_ds, cal_embeddings=cal_embeddings) - y_true, y_prob, _loss, extra = Trainer(model=ncp_predictor).inference( - test_loader, additional_outputs=["y_predset"] - ) - ncp_metrics = get_metrics_fn("binary")( - y_true, y_prob, metrics=["accuracy", "miscoverage_ps"], y_predset=extra["y_predset"] - ) - predset = extra["y_predset"] - if isinstance(predset, np.ndarray): - predset_t = torch.tensor(predset) - else: - predset_t = predset - avg_set_size = predset_t.float().sum(dim=1).mean().item() - miscoverage = ncp_metrics["miscoverage_ps"] - if isinstance(miscoverage, np.ndarray): - miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) - else: - miscoverage = float(miscoverage) - coverage = 1.0 - miscoverage - - if return_metrics: - return { - "accuracy": float(base_metrics["accuracy"]), - "roc_auc": float(base_metrics["roc_auc"]), - "f1": float(base_metrics["f1"]), - "coverage": coverage, - "miscoverage": miscoverage, - "avg_set_size": avg_set_size, - } + # Calibration + evaluation — fast; loop over every alpha + results = {} + for alpha in alphas: + print(f" Calibrating NCP predictor (alpha={alpha})...") + ncp_predictor = NeighborhoodLabel( + model=model, + alpha=float(alpha), + k_neighbors=args.k_neighbors, + lambda_L=args.lambda_L, + ) + ncp_predictor.calibrate(cal_dataset=cal_ds, cal_embeddings=cal_embeddings) - print("\nNCP (NeighborhoodLabel) Results:") - print(f" Accuracy: {base_metrics['accuracy']:.4f}") - print(f" ROC-AUC: {base_metrics['roc_auc']:.4f}") - print(f" F1: {base_metrics['f1']:.4f}") - print(f" Empirical miscoverage: {miscoverage:.4f}") - print(f" Empirical coverage: {coverage:.4f}") - print(f" Average set size: {avg_set_size:.2f}") - print(f" k_neighbors: {args.k_neighbors}") + y_true, y_prob, _, extra = Trainer(model=ncp_predictor).inference( + test_loader, additional_outputs=["y_predset"] + ) + ncp_metrics = get_metrics_fn("multiclass")( + y_true, y_prob, metrics=["accuracy", "miscoverage_ps"], y_predset=extra["y_predset"] + ) + predset = extra["y_predset"] + predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset + avg_set_size = predset_t.float().sum(dim=1).mean().item() + + miscoverage = ncp_metrics["miscoverage_ps"] + if isinstance(miscoverage, np.ndarray): + miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) + else: + miscoverage = float(miscoverage) + + results[alpha] = { + "accuracy": float(base_metrics["accuracy"]), + "roc_auc_weighted_ovr": float(base_metrics["roc_auc_weighted_ovr"]), + "f1_weighted": float(base_metrics["f1_weighted"]), + "coverage": 1.0 - miscoverage, + "miscoverage": miscoverage, + "avg_set_size": avg_set_size, + } + return results def main() -> None: @@ -287,79 +328,58 @@ def _run(args: argparse.Namespace) -> None: print(f"Input schema: {sample_dataset.input_schema}") print(f"Output schema: {sample_dataset.output_schema}") - print("\n--- Experiment configuration ---") - print(f" dataset_root: {root}") - print(f" subset: {args.subset}, ratios: train/val/cal = {args.ratios[0]:.2f}/{args.ratios[1]:.2f}/{args.ratios[2]:.2f} (test = TUH eval partition)") - print(f" alpha: {args.alpha} (target coverage {1 - args.alpha:.0%})") - print(f" k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") - print(f" epochs: {epochs}, batch_size: {args.batch_size}, device: {device}, seed: {args.seed}") - if len(sample_dataset) == 0: raise RuntimeError("No samples produced. Verify TUAB root/subset/task.") + # Parse alphas and run seeds + alphas = [float(a.strip()) for a in args.alphas.split(",")] if args.alphas else [args.alpha] ratios = list(args.ratios) use_multi_seed = args.n_seeds > 1 or args.seeds is not None - if use_multi_seed: - run_seeds = ( - [int(s.strip()) for s in args.seeds.split(",")] - if args.seeds - else [args.seed + i for i in range(args.n_seeds)] - ) - n_runs = len(run_seeds) - print(f" multi_seed: n_runs={n_runs}, run_seeds={run_seeds}, split_seed={args.split_seed} (fixed test set)") - print(f"Multi-seed mode: {n_runs} runs (fixed test set), run seeds: {run_seeds}") - - if not use_multi_seed: - print("\n" + "=" * 80) - print("STEP 2: Split train/val/cal/test") - print("=" * 80) - train_ds, val_ds, cal_ds, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=ratios, seed=args.seed - ) - print(f"Train: {len(train_ds)}") - print(f"Val: {len(val_ds)}") - print(f"Cal: {len(cal_ds)}") - print(f"Test: {len(test_ds)}") - - test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) - _run_one_ncp( - sample_dataset=sample_dataset, - train_ds=train_ds, - val_ds=val_ds, - cal_ds=cal_ds, - test_loader=test_loader, - args=args, - device=device, - epochs=epochs, - ) - print("\n--- Split sizes and seed (for reporting) ---") - print(f" train={len(train_ds)}, val={len(val_ds)}, cal={len(cal_ds)}, test={len(test_ds)}, seed={args.seed}") - return + run_seeds = ( + [int(s.strip()) for s in args.seeds.split(",")] + if args.seeds + else [args.seed + i for i in range(args.n_seeds)] + ) + n_runs = len(run_seeds) - # Multi-seed: test set is always fixed as the TUH eval partition (no split-seed needed). - # Each run uses a different seed to re-shuffle the train pool into train/val/cal. + # ------------------------------------------------------------------------- + # STEP 2: Extract the fixed test set ONCE + # ------------------------------------------------------------------------- print("\n" + "=" * 80) - print("STEP 2: Fix test set (TUH eval partition), then run multiple train/cal splits") + print("STEP 2: Extract fixed test set (TUH eval partition — same for all seeds)") print("=" * 80) - _, _, _, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=ratios, seed=args.split_seed + _, _, _, test_ds = _do_split( + sample_dataset, ratios, args.split_seed, args.split_type ) + if len(test_ds) == 0 and args.quick_test: + print(" [quick-test] TUH eval partition empty in dev mode — using random 20% as test set.") + _, _, _, test_ds = split_by_sample_conformal( + dataset=sample_dataset, ratios=[0.6, 0.1, 0.1, 0.2], seed=args.split_seed + ) test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) n_test = len(test_ds) - print(f"Fixed test set size: {n_test} (TUH eval partition)") + print(f"Test: {n_test} (fixed)") - accs, roc_aucs, f1s, coverages, miscoverages, set_sizes = [], [], [], [], [], [] + print(f"\nRun config: {'multi-seed (' + str(n_runs) + ' runs)' if use_multi_seed else 'single run'}") + print(f"Seeds: {run_seeds}, alphas={alphas}, k_neighbors={args.k_neighbors}") + + # ------------------------------------------------------------------------- + # STEP 3+: Train once per seed; calibrate for every alpha (fast) + # ------------------------------------------------------------------------- + all_metrics = {alpha: [] for alpha in alphas} for run_i, run_seed in enumerate(run_seeds): print("\n" + "=" * 80) - print(f"Run {run_i + 1} / {n_runs} (seed={run_seed})") + if use_multi_seed: + print(f"Run {run_i + 1} / {n_runs} (seed={run_seed})") + else: + print(f"STEP 3–4: Train + NCP Calibration (seed={run_seed})") print("=" * 80) set_seed(run_seed) - train_ds, val_ds, cal_ds = _split_train_pool_for_run( - sample_dataset, ratios, run_seed - ) - print(f"Train: {len(train_ds)}, Val: {len(val_ds)}, Cal: {len(cal_ds)}") + train_ds, val_ds, cal_ds = _split_train_pool_for_run(sample_dataset, ratios, run_seed, args.split_type) + print(f" Split — Train: {len(train_ds)}, Val: {len(val_ds)}, " + f"Cal: {len(cal_ds)}, Test: {n_test} (fixed)") - metrics = _run_one_ncp( + seed_results = _run_one_ncp( sample_dataset=sample_dataset, train_ds=train_ds, val_ds=val_ds, @@ -368,50 +388,66 @@ def _run(args: argparse.Namespace) -> None: args=args, device=device, epochs=epochs, - return_metrics=True, + alphas=alphas, + run_idx=run_i, ) - accs.append(metrics["accuracy"]) - roc_aucs.append(metrics["roc_auc"]) - f1s.append(metrics["f1"]) - coverages.append(metrics["coverage"]) - miscoverages.append(metrics["miscoverage"]) - set_sizes.append(metrics["avg_set_size"]) - - accs = np.array(accs) - roc_aucs = np.array(roc_aucs) - f1s = np.array(f1s) - coverages = np.array(coverages) - miscoverages_arr = np.array(miscoverages) - set_sizes = np.array(set_sizes) - - print("\n" + "=" * 80) - print("Per-run NCP results (fixed test set = TUH eval partition)") - print("=" * 80) - print(f" {'Run':<4} {'Seed':<6} {'Accuracy':<10} {'ROC-AUC':<10} {'F1':<8} " - f"{'Coverage':<10} {'Miscoverage':<12} {'Avg set size':<12}") - print(" " + "-" * 76) - for i in range(n_runs): - print(f" {i+1:<4} {run_seeds[i]:<6} {accs[i]:<10.4f} {roc_aucs[i]:<10.4f} " - f"{f1s[i]:<8.4f} {coverages[i]:<10.4f} {miscoverages_arr[i]:<12.4f} {set_sizes[i]:<12.2f}") - - print("\n" + "=" * 80) - print("NCP summary (mean \u00b1 std over {} runs, fixed test set)".format(n_runs)) - print("=" * 80) - print(f" Accuracy: {accs.mean():.4f} \u00b1 {accs.std():.4f}") - print(f" ROC-AUC: {roc_aucs.mean():.4f} \u00b1 {roc_aucs.std():.4f}") - print(f" F1: {f1s.mean():.4f} \u00b1 {f1s.std():.4f}") - print(f" Empirical coverage: {coverages.mean():.4f} \u00b1 {coverages.std():.4f}") - print(f" Empirical miscoverage: {miscoverages_arr.mean():.4f} \u00b1 {miscoverages_arr.std():.4f}") - print(f" Average set size: {set_sizes.mean():.2f} \u00b1 {set_sizes.std():.2f}") - print(f" Target coverage: {1 - args.alpha:.0%} (alpha={args.alpha})") - print(f" k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") - print(f" Test set size: {n_test} (fixed across runs)") - print(f" Run seeds: {run_seeds}") - print("\n--- Min / Max (across runs) ---") - print(f" Coverage: [{coverages.min():.4f}, {coverages.max():.4f}]") - print(f" Set size: [{set_sizes.min():.2f}, {set_sizes.max():.2f}]") - print(f" Accuracy: [{accs.min():.4f}, {accs.max():.4f}]") - print(f" ROC-AUC: [{roc_aucs.min():.4f}, {roc_aucs.max():.4f}]") + for alpha in alphas: + all_metrics[alpha].append(seed_results[alpha]) + + if use_multi_seed: + m = seed_results[alphas[0]] + print(f" [Run {run_i + 1} result (alpha={alphas[0]})] " + f"acc={m['accuracy']:.4f}, roc_auc={m['roc_auc_weighted_ovr']:.4f}, " + f"cov={m['coverage']:.4f}, set_size={m['avg_set_size']:.2f}") + + for alpha in alphas: + mlist = all_metrics[alpha] + accs = np.array([m["accuracy"] for m in mlist]) + roc_aucs = np.array([m["roc_auc_weighted_ovr"] for m in mlist]) + f1s = np.array([m["f1_weighted"] for m in mlist]) + coverages = np.array([m["coverage"] for m in mlist]) + miscovs = np.array([m["miscoverage"] for m in mlist]) + set_sizes = np.array([m["avg_set_size"] for m in mlist]) + + if not use_multi_seed: + print(f"\nNCP Results (alpha={alpha}):") + print(f" Accuracy: {accs[0]:.4f}") + print(f" ROC-AUC: {roc_aucs[0]:.4f}") + print(f" F1: {f1s[0]:.4f}") + print(f" Empirical coverage: {coverages[0]:.4f}") + print(f" Empirical miscoverage: {miscovs[0]:.4f}") + print(f" Average set size: {set_sizes[0]:.2f}") + print(f" Target coverage: {1 - alpha:.0%} (alpha={alpha})") + print(f" k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") + else: + print("\n" + "=" * 80) + print(f"Per-run NCP results — alpha={alpha} (target coverage={1-alpha:.0%})") + print("=" * 80) + print(f" {'Run':<4} {'Seed':<6} {'Accuracy':<10} {'ROC-AUC':<10} {'F1':<8} " + f"{'Coverage':<10} {'Miscoverage':<12} {'Avg set size':<12}") + print(" " + "-" * 76) + for i in range(n_runs): + print(f" {i+1:<4} {run_seeds[i]:<6} {accs[i]:<10.4f} {roc_aucs[i]:<10.4f} " + f"{f1s[i]:<8.4f} {coverages[i]:<10.4f} {miscovs[i]:<12.4f} {set_sizes[i]:<12.2f}") + + print("\n" + "=" * 80) + print(f"NCP summary — alpha={alpha} (mean ± std over {n_runs} runs, fixed test set)") + print("=" * 80) + print(f" Accuracy: {accs.mean():.4f} ± {accs.std():.4f}") + print(f" ROC-AUC: {roc_aucs.mean():.4f} ± {roc_aucs.std():.4f}") + print(f" F1: {f1s.mean():.4f} ± {f1s.std():.4f}") + print(f" Empirical coverage: {coverages.mean():.4f} ± {coverages.std():.4f}") + print(f" Empirical miscoverage: {miscovs.mean():.4f} ± {miscovs.std():.4f}") + print(f" Average set size: {set_sizes.mean():.2f} ± {set_sizes.std():.2f}") + print(f" Target coverage: {1 - alpha:.0%} (alpha={alpha})") + print(f" k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") + print(f" Test set size: {n_test} (fixed across runs)") + print(f" Run seeds: {run_seeds}") + print("\n--- Min / Max (across runs) ---") + print(f" Coverage: [{coverages.min():.4f}, {coverages.max():.4f}]") + print(f" Set size: [{set_sizes.min():.2f}, {set_sizes.max():.2f}]") + print(f" Accuracy: [{accs.min():.4f}, {accs.max():.4f}]") + print(f" ROC-AUC: [{roc_aucs.min():.4f}, {roc_aucs.max():.4f}]") if __name__ == "__main__": diff --git a/examples/conformal_eeg/tuev_conventional_conformal.py b/examples/conformal_eeg/tuev_conventional_conformal.py index f1155545f..fe41aaff4 100644 --- a/examples/conformal_eeg/tuev_conventional_conformal.py +++ b/examples/conformal_eeg/tuev_conventional_conformal.py @@ -1,24 +1,28 @@ -"""Conventional Conformal Prediction (LABEL) on TUEV EEG Events using ContraWR. +"""Conventional Conformal Prediction (LABEL) on TUEV EEG Events. + +Supports both ContraWR (default) and TFMTokenizer via --model contrawr|tfm. This script: 1) Loads the TUEV dataset and applies the EEGEventsTUEV task (once, shared across all seeds). 2) Extracts the fixed test set (TUH eval partition — never changes across seeds). -3) For each seed: splits the TUH train partition into train/val/cal, trains ContraWR, +3) For each seed: splits the TUH train partition into train/val/cal, trains the chosen model, calibrates a LABEL predictor, and evaluates on the fixed test set. 4) Reports per-run results and mean ± std summary across all seeds. Single-seed usage (from repo root): python examples/conformal_eeg/tuev_conventional_conformal.py --root downloads/tuev/v2.0.1/edf + python examples/conformal_eeg/tuev_conventional_conformal.py --model tfm --root downloads/tuev/v2.0.1/edf Multi-seed usage (recommended for papers): python examples/conformal_eeg/tuev_conventional_conformal.py \\ - --root downloads/tuev/v2.0.1/edf --n-seeds 5 --seed 42 --alpha 0.1 \\ - --log-file tuev_conventional_alpha0.1_5seeds.log + --root downloads/tuev/v2.0.1/edf --model tfm --n-seeds 5 --seed 42 --alpha 0.1 \\ + --log-file tuev_conventional_tfm_alpha0.1_5seeds.log """ from __future__ import annotations import argparse +import os import random import sys from pathlib import Path @@ -27,8 +31,8 @@ import torch from pyhealth.calib.predictionset import LABEL -from pyhealth.datasets import TUEVDataset, get_dataloader, split_by_sample_conformal_tuh -from pyhealth.models import ContraWR +from pyhealth.datasets import TUEVDataset, get_dataloader, split_by_patient_conformal_tuh, split_by_sample_conformal_tuh, split_by_sample_conformal +from pyhealth.models import ContraWR, TFMTokenizer from pyhealth.tasks import EEGEventsTUEV from pyhealth.trainer import Trainer, get_metrics_fn @@ -52,12 +56,12 @@ def flush(self): def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Conventional conformal prediction (LABEL) on TUEV EEG events using ContraWR." + description="Conventional conformal prediction (LABEL) on TUEV EEG events." ) parser.add_argument( "--root", type=str, - default="downloads/tuev/v2.0.1/edf", + default="/srv/local/data/TUH/tuh_eeg_events/v2.0.0/edf", help="Path to TUEV edf/ folder.", ) parser.add_argument("--subset", type=str, default="both", choices=["train", "eval", "both"]) @@ -73,6 +77,10 @@ def parse_args() -> argparse.Namespace: "--alpha", type=float, default=0.1, help="Miscoverage rate (e.g., 0.1 => 90% target coverage).", ) + parser.add_argument( + "--alphas", type=str, default=None, + help="Comma-separated miscoverage rates, e.g. '0.2,0.1,0.05,0.01'. Overrides --alpha.", + ) parser.add_argument( "--ratios", type=float, @@ -83,6 +91,10 @@ def parse_args() -> argparse.Namespace: "Must sum to 1.0. Test is fixed as the TUH eval partition.", ) parser.add_argument("--n-fft", type=int, default=128, help="STFT FFT size used by ContraWR.") + parser.add_argument( + "--model", type=str, default="contrawr", choices=["contrawr", "tfm"], + help="Backbone model: 'contrawr' (default) or 'tfm' (TFMTokenizer).", + ) parser.add_argument( "--device", type=str, default=None, help="Device string, e.g. 'cuda:0' or 'cpu'. Defaults to auto-detect.", @@ -101,6 +113,19 @@ def parse_args() -> argparse.Namespace: help="Explicit comma-separated seeds, e.g. '42,43,44,45,46'. " "Overrides --seed and --n-seeds.", ) + parser.add_argument( + "--weights-dir", + type=str, + default="weightfiles/TFM_Tokenizer_multiple_finetuned_on_TUEV", + help="Root folder of fine-tuned TFM classifier checkpoints (used only with --model tfm). " + "Expected sub-folders: /_1/best_model.pth, ..., _5/best_model.pth.", + ) + parser.add_argument( + "--tokenizer-weights", + type=str, + default="weightfiles/tfm_tokenizer_last.pth", + help="Path to the pre-trained TFM tokenizer weights (used only with --model tfm).", + ) parser.add_argument( "--log-file", type=str, default=None, help="Path to log file. Stdout and stderr are teed to this file.", @@ -110,9 +135,36 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Smoke test: dev=True, max 2000 samples, 2 epochs.", ) + parser.add_argument( + "--split-type", + type=str, + default="patient", + choices=["patient", "sample"], + help="Split strategy: 'patient' (default, patient-level, no leakage) or " + "'sample' (original sample-level, for comparison).", + ) return parser.parse_args() +def _do_split(dataset, ratios, seed, split_type): + """Dispatch to the correct TUH split function based on split_type.""" + if split_type == "patient": + return split_by_patient_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + else: + return split_by_sample_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + + +def _load_tfm_weights(model, args, run_idx: int) -> None: + """Load pre-trained tokenizer + fine-tuned classifier for run_idx (0-based).""" + base = os.path.basename(args.weights_dir) + classifier_path = os.path.join(args.weights_dir, f"{base}_{run_idx + 1}", "best_model.pth") + print(f" Loading TFM weights (run {run_idx + 1}): {classifier_path}") + model.load_pretrained_weights( + tokenizer_checkpoint_path=args.tokenizer_weights, + classifier_checkpoint_path=classifier_path, + ) + + def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) @@ -129,19 +181,20 @@ def _run_one_seed( device: str, epochs: int, run_seed: int, + alphas: list, + run_idx: int = 0, ) -> dict: - """Train ContraWR + calibrate LABEL predictor for one seed. + """Train model + calibrate LABEL for one seed across all alphas. - The test set is passed in pre-built (fixed TUH eval partition). - Only train/val/cal vary per seed. + Training and base-model inference are done once; calibration loops over alphas (fast). - Returns a dict with keys: + Returns {alpha: metrics_dict} where metrics_dict has keys: accuracy, f1_weighted, coverage, miscoverage, avg_set_size """ set_seed(run_seed) - train_ds, val_ds, cal_ds, _ = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=run_seed + train_ds, val_ds, cal_ds, _ = _do_split( + sample_dataset, args.ratios, run_seed, args.split_type ) print(f" Split — Train: {len(train_ds)}, Val: {len(val_ds)}, " f"Cal: {len(cal_ds)}, Test: {len(test_ds)} (fixed)") @@ -152,54 +205,61 @@ def _run_one_seed( if len(val_ds) else None ) - print(" Training ContraWR...") - model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + if args.model == "tfm": + model = TFMTokenizer(dataset=sample_dataset).to(device) + _load_tfm_weights(model, args, run_idx) + else: + model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + print(" Training ContraWR...") + trainer_tmp = Trainer(model=model, device=device, enable_logging=False) + trainer_tmp.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=epochs, + monitor="accuracy" if val_loader is not None else None, + ) trainer = Trainer(model=model, device=device, enable_logging=False) - trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=epochs, - monitor="accuracy" if val_loader is not None else None, - ) - # Base model metrics on fixed test set + # Base model metrics — computed once, shared across all alphas y_true_base, y_prob_base, _ = trainer.inference(test_loader) base_metrics = get_metrics_fn("multiclass")( y_true_base, y_prob_base, metrics=["accuracy", "f1_weighted"] ) - # Conformal calibration + evaluation - print(" Calibrating LABEL predictor...") - label_predictor = LABEL(model=model, alpha=float(args.alpha)) - label_predictor.calibrate(cal_dataset=cal_ds) + # Calibration + evaluation — fast; loop over every alpha + results = {} + for alpha in alphas: + print(f" Calibrating LABEL predictor (alpha={alpha})...") + label_predictor = LABEL(model=model, alpha=float(alpha)) + label_predictor.calibrate(cal_dataset=cal_ds) - print(" Evaluating LABEL predictor on test set...") - y_true, y_prob, _, extra = Trainer(model=label_predictor).inference( - test_loader, additional_outputs=["y_predset"] - ) - conf_metrics = get_metrics_fn("multiclass")( - y_true, y_prob, - metrics=["accuracy", "miscoverage_ps"], - y_predset=extra["y_predset"], - ) + y_true, y_prob, _, extra = Trainer(model=label_predictor).inference( + test_loader, additional_outputs=["y_predset"] + ) + conf_metrics = get_metrics_fn("multiclass")( + y_true, y_prob, + metrics=["accuracy", "miscoverage_ps"], + y_predset=extra["y_predset"], + ) - predset = extra["y_predset"] - predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset - avg_set_size = predset_t.float().sum(dim=1).mean().item() + predset = extra["y_predset"] + predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset + avg_set_size = predset_t.float().sum(dim=1).mean().item() - miscoverage = conf_metrics["miscoverage_ps"] - if isinstance(miscoverage, np.ndarray): - miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) - else: - miscoverage = float(miscoverage) + miscoverage = conf_metrics["miscoverage_ps"] + if isinstance(miscoverage, np.ndarray): + miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) + else: + miscoverage = float(miscoverage) - return { - "accuracy": float(base_metrics["accuracy"]), - "f1_weighted": float(base_metrics["f1_weighted"]), - "coverage": 1.0 - miscoverage, - "miscoverage": miscoverage, - "avg_set_size": avg_set_size, - } + results[alpha] = { + "accuracy": float(base_metrics["accuracy"]), + "f1_weighted": float(base_metrics["f1_weighted"]), + "coverage": 1.0 - miscoverage, + "miscoverage": miscoverage, + "avg_set_size": avg_set_size, + } + return results def _print_single_run_results(metrics: dict, alpha: float) -> None: @@ -289,50 +349,63 @@ def _main(args: argparse.Namespace) -> None: print("\n" + "=" * 80) print("STEP 2: Extract fixed test set (TUH eval partition — same for all seeds)") print("=" * 80) - _, _, _, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=args.seed + _, _, _, test_ds = _do_split( + sample_dataset, args.ratios, args.seed, args.split_type ) + if len(test_ds) == 0 and args.quick_test: + # dev mode only loads train-partition patients, so eval partition is empty. + # Fall back to a random 20% hold-out for smoke-testing the pipeline end-to-end. + print(" [quick-test] TUH eval partition empty in dev mode — using random 20% as test set.") + _, _, _, test_ds = split_by_sample_conformal( + dataset=sample_dataset, ratios=[0.6, 0.1, 0.1, 0.2], seed=args.seed + ) test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) print(f"Test: {len(test_ds)} (fixed)") # ------------------------------------------------------------------------- - # Determine run seeds + # Determine run seeds and alphas # ------------------------------------------------------------------------- if args.seeds is not None: run_seeds = [int(s.strip()) for s in args.seeds.split(",")] else: run_seeds = [args.seed + i for i in range(args.n_seeds)] + alphas = [float(a.strip()) for a in args.alphas.split(",")] if args.alphas else [args.alpha] + use_multi_seed = len(run_seeds) > 1 print(f"\nRun config: {'multi-seed (' + str(len(run_seeds)) + ' runs)' if use_multi_seed else 'single run'}") - print(f"Seeds: {run_seeds}, alpha={args.alpha}, target coverage={1 - args.alpha:.0%}") + print(f"Seeds: {run_seeds}, alphas={alphas}") # ------------------------------------------------------------------------- - # STEP 3+: Train + conformal (once per seed) + # STEP 3+: Train once per seed; calibrate for every alpha (fast) # ------------------------------------------------------------------------- - all_metrics = [] + all_metrics = {alpha: [] for alpha in alphas} for run_i, run_seed in enumerate(run_seeds): print("\n" + "=" * 80) if use_multi_seed: print(f"Run {run_i + 1} / {len(run_seeds)} (seed={run_seed})") else: - print(f"STEP 3–4: Train ContraWR + Conformal Calibration (seed={run_seed})") + print(f"STEP 3–4: Train + Conformal Calibration (seed={run_seed})") print("=" * 80) - metrics = _run_one_seed( - args, sample_dataset, test_ds, test_loader, device, epochs, run_seed + seed_results = _run_one_seed( + args, sample_dataset, test_ds, test_loader, device, epochs, run_seed, alphas, + run_idx=run_i, ) - all_metrics.append(metrics) + for alpha in alphas: + all_metrics[alpha].append(seed_results[alpha]) if use_multi_seed: - print(f" [Run {run_i + 1} result] " - f"acc={metrics['accuracy']:.4f}, f1={metrics['f1_weighted']:.4f}, " - f"cov={metrics['coverage']:.4f}, set_size={metrics['avg_set_size']:.2f}") - - if not use_multi_seed: - _print_single_run_results(all_metrics[0], args.alpha) - else: - _print_multi_seed_summary(all_metrics, run_seeds, args.alpha, len(test_ds)) + m = seed_results[alphas[0]] + print(f" [Run {run_i + 1} result (alpha={alphas[0]})] " + f"acc={m['accuracy']:.4f}, f1={m['f1_weighted']:.4f}, " + f"cov={m['coverage']:.4f}, set_size={m['avg_set_size']:.2f}") + + for alpha in alphas: + if not use_multi_seed: + _print_single_run_results(all_metrics[alpha][0], alpha) + else: + _print_multi_seed_summary(all_metrics[alpha], run_seeds, alpha, len(test_ds)) def main() -> None: diff --git a/examples/conformal_eeg/tuev_covariate_shift_conformal.py b/examples/conformal_eeg/tuev_covariate_shift_conformal.py index b32ff9bef..97169a0d7 100644 --- a/examples/conformal_eeg/tuev_covariate_shift_conformal.py +++ b/examples/conformal_eeg/tuev_covariate_shift_conformal.py @@ -24,6 +24,7 @@ from __future__ import annotations import argparse +import os import random import sys from pathlib import Path @@ -51,8 +52,8 @@ def flush(self): from pyhealth.calib.predictionset.covariate import CovariateLabel from pyhealth.calib.utils import extract_embeddings -from pyhealth.datasets import TUEVDataset, get_dataloader, split_by_sample_conformal_tuh -from pyhealth.models import ContraWR +from pyhealth.datasets import TUEVDataset, get_dataloader, split_by_patient_conformal_tuh, split_by_sample_conformal_tuh, split_by_sample_conformal +from pyhealth.models import ContraWR, TFMTokenizer from pyhealth.tasks import EEGEventsTUEV from pyhealth.trainer import Trainer, get_metrics_fn @@ -64,7 +65,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--root", type=str, - default="downloads/tuev/v2.0.1/edf", + default="/srv/local/data/TUH/tuh_eeg_events/v2.0.0/edf", help="Path to TUEV edf/ folder.", ) parser.add_argument("--subset", type=str, default="both", choices=["train", "eval", "both"]) @@ -80,6 +81,10 @@ def parse_args() -> argparse.Namespace: "--alpha", type=float, default=0.1, help="Miscoverage rate (e.g., 0.1 => 90% target coverage).", ) + parser.add_argument( + "--alphas", type=str, default=None, + help="Comma-separated miscoverage rates, e.g. '0.2,0.1,0.05,0.01'. Overrides --alpha.", + ) parser.add_argument( "--ratios", type=float, @@ -90,6 +95,10 @@ def parse_args() -> argparse.Namespace: "Must sum to 1.0. Test is fixed as the TUH eval partition.", ) parser.add_argument("--n-fft", type=int, default=128, help="STFT FFT size used by ContraWR.") + parser.add_argument( + "--model", type=str, default="contrawr", choices=["contrawr", "tfm"], + help="Backbone model: 'contrawr' (default) or 'tfm' (TFMTokenizer).", + ) parser.add_argument( "--device", type=str, default=None, help="Device string, e.g. 'cuda:0' or 'cpu'. Defaults to auto-detect.", @@ -117,9 +126,48 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Smoke test: dev=True, max 2000 samples, 2 epochs.", ) + parser.add_argument( + "--weights-dir", + type=str, + default="weightfiles/TFM_Tokenizer_multiple_finetuned_on_TUEV", + help="Root folder of fine-tuned TFM classifier checkpoints (only with --model tfm).", + ) + parser.add_argument( + "--tokenizer-weights", + type=str, + default="weightfiles/tfm_tokenizer_last.pth", + help="Path to the pre-trained TFM tokenizer weights (only with --model tfm).", + ) + parser.add_argument( + "--split-type", + type=str, + default="patient", + choices=["patient", "sample"], + help="Split strategy: 'patient' (default, patient-level, no leakage) or " + "'sample' (original sample-level, for comparison).", + ) return parser.parse_args() +def _do_split(dataset, ratios, seed, split_type): + """Dispatch to the correct TUH split function based on split_type.""" + if split_type == "patient": + return split_by_patient_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + else: + return split_by_sample_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + + +def _load_tfm_weights(model, args, run_idx: int) -> None: + """Load pre-trained tokenizer + fine-tuned classifier for run_idx (0-based).""" + base = os.path.basename(args.weights_dir) + classifier_path = os.path.join(args.weights_dir, f"{base}_{run_idx + 1}", "best_model.pth") + print(f" Loading TFM weights (run {run_idx + 1}): {classifier_path}") + model.load_pretrained_weights( + tokenizer_checkpoint_path=args.tokenizer_weights, + classifier_checkpoint_path=classifier_path, + ) + + def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) @@ -136,20 +184,21 @@ def _run_one_seed( device: str, epochs: int, run_seed: int, + alphas: list, + run_idx: int = 0, ) -> dict: - """Train ContraWR + calibrate CovariateLabel predictor for one seed. + """Train model + calibrate CovariateLabel for one seed across all alphas. - The test set is passed in pre-built (fixed TUH eval partition). - Only train/val/cal (and model weights) vary per seed. - Test embeddings are recomputed each seed because the model changes. + Training, embedding extraction, and base inference are done once; calibration + loops over alphas (fast — only likelihood-ratio weights/threshold recomputed). - Returns a dict with keys: + Returns {alpha: metrics_dict} where metrics_dict has keys: accuracy, f1_weighted, coverage, miscoverage, avg_set_size """ set_seed(run_seed) - train_ds, val_ds, cal_ds, _ = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=run_seed + train_ds, val_ds, cal_ds, _ = _do_split( + sample_dataset, args.ratios, run_seed, args.split_type ) print(f" Split — Train: {len(train_ds)}, Val: {len(val_ds)}, " f"Cal: {len(cal_ds)}, Test: {len(test_ds)} (fixed)") @@ -160,63 +209,70 @@ def _run_one_seed( if len(val_ds) else None ) - print(" Training ContraWR...") - model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + if args.model == "tfm": + model = TFMTokenizer(dataset=sample_dataset).to(device) + _load_tfm_weights(model, args, run_idx) + else: + model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + print(" Training ContraWR...") + trainer_tmp = Trainer(model=model, device=device, enable_logging=False) + trainer_tmp.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=epochs, + monitor="accuracy" if val_loader is not None else None, + ) trainer = Trainer(model=model, device=device, enable_logging=False) - trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=epochs, - monitor="accuracy" if val_loader is not None else None, - ) - # Base model metrics on fixed test set + # Base model metrics — computed once, shared across all alphas y_true_base, y_prob_base, _ = trainer.inference(test_loader) base_metrics = get_metrics_fn("multiclass")( y_true_base, y_prob_base, metrics=["accuracy", "f1_weighted"] ) - # Extract embeddings — both depend on the current model so must redo each seed + # Extract embeddings once — reused for every alpha print(" Extracting embeddings for calibration and test splits...") cal_embeddings = extract_embeddings(model, cal_ds, batch_size=args.batch_size, device=device) test_embeddings = extract_embeddings(model, test_ds, batch_size=args.batch_size, device=device) - # Conformal calibration + evaluation - print(" Calibrating CovariateLabel predictor...") - cov_predictor = CovariateLabel(model=model, alpha=float(args.alpha)) - cov_predictor.calibrate( - cal_dataset=cal_ds, - cal_embeddings=cal_embeddings, - test_embeddings=test_embeddings, - ) + # Calibration + evaluation — fast; loop over every alpha + results = {} + for alpha in alphas: + print(f" Calibrating CovariateLabel predictor (alpha={alpha})...") + cov_predictor = CovariateLabel(model=model, alpha=float(alpha)) + cov_predictor.calibrate( + cal_dataset=cal_ds, + cal_embeddings=cal_embeddings, + test_embeddings=test_embeddings, + ) - print(" Evaluating CovariateLabel predictor on test set...") - y_true, y_prob, _, extra = Trainer(model=cov_predictor).inference( - test_loader, additional_outputs=["y_predset"] - ) - conf_metrics = get_metrics_fn("multiclass")( - y_true, y_prob, - metrics=["accuracy", "miscoverage_ps"], - y_predset=extra["y_predset"], - ) + y_true, y_prob, _, extra = Trainer(model=cov_predictor).inference( + test_loader, additional_outputs=["y_predset"] + ) + conf_metrics = get_metrics_fn("multiclass")( + y_true, y_prob, + metrics=["accuracy", "miscoverage_ps"], + y_predset=extra["y_predset"], + ) - predset = extra["y_predset"] - predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset - avg_set_size = predset_t.float().sum(dim=1).mean().item() + predset = extra["y_predset"] + predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset + avg_set_size = predset_t.float().sum(dim=1).mean().item() - miscoverage = conf_metrics["miscoverage_ps"] - if isinstance(miscoverage, np.ndarray): - miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) - else: - miscoverage = float(miscoverage) + miscoverage = conf_metrics["miscoverage_ps"] + if isinstance(miscoverage, np.ndarray): + miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) + else: + miscoverage = float(miscoverage) - return { - "accuracy": float(base_metrics["accuracy"]), - "f1_weighted": float(base_metrics["f1_weighted"]), - "coverage": 1.0 - miscoverage, - "miscoverage": miscoverage, - "avg_set_size": avg_set_size, - } + results[alpha] = { + "accuracy": float(base_metrics["accuracy"]), + "f1_weighted": float(base_metrics["f1_weighted"]), + "coverage": 1.0 - miscoverage, + "miscoverage": miscoverage, + "avg_set_size": avg_set_size, + } + return results def _print_single_run_results(metrics: dict, alpha: float) -> None: @@ -305,9 +361,14 @@ def _main(args: argparse.Namespace) -> None: print("\n" + "=" * 80) print("STEP 2: Extract fixed test set (TUH eval partition — same for all seeds)") print("=" * 80) - _, _, _, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=args.seed + _, _, _, test_ds = _do_split( + sample_dataset, args.ratios, args.seed, args.split_type ) + if len(test_ds) == 0 and args.quick_test: + print(" [quick-test] TUH eval partition empty in dev mode — using random 20% as test set.") + _, _, _, test_ds = split_by_sample_conformal( + dataset=sample_dataset, ratios=[0.6, 0.1, 0.1, 0.2], seed=args.seed + ) test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) print(f"Test: {len(test_ds)} (fixed)") @@ -319,36 +380,42 @@ def _main(args: argparse.Namespace) -> None: else: run_seeds = [args.seed + i for i in range(args.n_seeds)] + alphas = [float(a.strip()) for a in args.alphas.split(",")] if args.alphas else [args.alpha] + use_multi_seed = len(run_seeds) > 1 print(f"\nRun config: {'multi-seed (' + str(len(run_seeds)) + ' runs)' if use_multi_seed else 'single run'}") - print(f"Seeds: {run_seeds}, alpha={args.alpha}, target coverage={1 - args.alpha:.0%}") + print(f"Seeds: {run_seeds}, alphas={alphas}") # ------------------------------------------------------------------------- - # STEP 3+: Train + conformal (once per seed) + # STEP 3+: Train once per seed; calibrate for every alpha (fast) # ------------------------------------------------------------------------- - all_metrics = [] + all_metrics = {alpha: [] for alpha in alphas} for run_i, run_seed in enumerate(run_seeds): print("\n" + "=" * 80) if use_multi_seed: print(f"Run {run_i + 1} / {len(run_seeds)} (seed={run_seed})") else: - print(f"STEP 3–4: Train ContraWR + Conformal Calibration (seed={run_seed})") + print(f"STEP 3–4: Train + Conformal Calibration (seed={run_seed})") print("=" * 80) - metrics = _run_one_seed( - args, sample_dataset, test_ds, test_loader, device, epochs, run_seed + seed_results = _run_one_seed( + args, sample_dataset, test_ds, test_loader, device, epochs, run_seed, alphas, + run_idx=run_i, ) - all_metrics.append(metrics) + for alpha in alphas: + all_metrics[alpha].append(seed_results[alpha]) if use_multi_seed: - print(f" [Run {run_i + 1} result] " - f"acc={metrics['accuracy']:.4f}, f1={metrics['f1_weighted']:.4f}, " - f"cov={metrics['coverage']:.4f}, set_size={metrics['avg_set_size']:.2f}") - - if not use_multi_seed: - _print_single_run_results(all_metrics[0], args.alpha) - else: - _print_multi_seed_summary(all_metrics, run_seeds, args.alpha, len(test_ds)) + m = seed_results[alphas[0]] + print(f" [Run {run_i + 1} result (alpha={alphas[0]})] " + f"acc={m['accuracy']:.4f}, f1={m['f1_weighted']:.4f}, " + f"cov={m['coverage']:.4f}, set_size={m['avg_set_size']:.2f}") + + for alpha in alphas: + if not use_multi_seed: + _print_single_run_results(all_metrics[alpha][0], alpha) + else: + _print_multi_seed_summary(all_metrics[alpha], run_seeds, alpha, len(test_ds)) def main() -> None: diff --git a/examples/conformal_eeg/tuev_kmeans_conformal.py b/examples/conformal_eeg/tuev_kmeans_conformal.py index 460f0a8a7..598ad43b1 100644 --- a/examples/conformal_eeg/tuev_kmeans_conformal.py +++ b/examples/conformal_eeg/tuev_kmeans_conformal.py @@ -23,6 +23,7 @@ from __future__ import annotations import argparse +import os import random import sys from pathlib import Path @@ -50,8 +51,8 @@ def flush(self): from pyhealth.calib.predictionset.cluster import ClusterLabel from pyhealth.calib.utils import extract_embeddings -from pyhealth.datasets import TUEVDataset, get_dataloader, split_by_sample_conformal_tuh -from pyhealth.models import ContraWR +from pyhealth.datasets import TUEVDataset, get_dataloader, split_by_patient_conformal_tuh, split_by_sample_conformal_tuh, split_by_sample_conformal +from pyhealth.models import ContraWR, TFMTokenizer from pyhealth.tasks import EEGEventsTUEV from pyhealth.trainer import Trainer, get_metrics_fn @@ -63,7 +64,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--root", type=str, - default="downloads/tuev/v2.0.1/edf", + default="/srv/local/data/TUH/tuh_eeg_events/v2.0.0/edf", help="Path to TUEV edf/ folder.", ) parser.add_argument("--subset", type=str, default="both", choices=["train", "eval", "both"]) @@ -79,6 +80,10 @@ def parse_args() -> argparse.Namespace: "--alpha", type=float, default=0.1, help="Miscoverage rate (e.g., 0.1 => 90% target coverage).", ) + parser.add_argument( + "--alphas", type=str, default=None, + help="Comma-separated miscoverage rates, e.g. '0.2,0.1,0.05,0.01'. Overrides --alpha.", + ) parser.add_argument( "--ratios", type=float, @@ -93,6 +98,10 @@ def parse_args() -> argparse.Namespace: help="Number of K-means clusters for cluster-specific thresholds.", ) parser.add_argument("--n-fft", type=int, default=128, help="STFT FFT size used by ContraWR.") + parser.add_argument( + "--model", type=str, default="contrawr", choices=["contrawr", "tfm"], + help="Backbone model: 'contrawr' (default) or 'tfm' (TFMTokenizer).", + ) parser.add_argument( "--device", type=str, default=None, help="Device string, e.g. 'cuda:0' or 'cpu'. Defaults to auto-detect.", @@ -120,9 +129,48 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Smoke test: dev=True, max 2000 samples, 2 epochs.", ) + parser.add_argument( + "--weights-dir", + type=str, + default="weightfiles/TFM_Tokenizer_multiple_finetuned_on_TUEV", + help="Root folder of fine-tuned TFM classifier checkpoints (only with --model tfm).", + ) + parser.add_argument( + "--tokenizer-weights", + type=str, + default="weightfiles/tfm_tokenizer_last.pth", + help="Path to the pre-trained TFM tokenizer weights (only with --model tfm).", + ) + parser.add_argument( + "--split-type", + type=str, + default="patient", + choices=["patient", "sample"], + help="Split strategy: 'patient' (default, patient-level, no leakage) or " + "'sample' (original sample-level, for comparison).", + ) return parser.parse_args() +def _do_split(dataset, ratios, seed, split_type): + """Dispatch to the correct TUH split function based on split_type.""" + if split_type == "patient": + return split_by_patient_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + else: + return split_by_sample_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + + +def _load_tfm_weights(model, args, run_idx: int) -> None: + """Load pre-trained tokenizer + fine-tuned classifier for run_idx (0-based).""" + base = os.path.basename(args.weights_dir) + classifier_path = os.path.join(args.weights_dir, f"{base}_{run_idx + 1}", "best_model.pth") + print(f" Loading TFM weights (run {run_idx + 1}): {classifier_path}") + model.load_pretrained_weights( + tokenizer_checkpoint_path=args.tokenizer_weights, + classifier_checkpoint_path=classifier_path, + ) + + def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) @@ -139,19 +187,21 @@ def _run_one_seed( device: str, epochs: int, run_seed: int, + alphas: list, + run_idx: int = 0, ) -> dict: - """Train ContraWR + calibrate ClusterLabel predictor for one seed. + """Train model + calibrate ClusterLabel for one seed across all alphas. - The test set is passed in pre-built (fixed TUH eval partition). - Only train/val/cal (and model weights) vary per seed. + Training, embedding extraction, and base inference are done once; calibration + loops over alphas (fast — only threshold recomputed per alpha). - Returns a dict with keys: + Returns {alpha: metrics_dict} where metrics_dict has keys: accuracy, f1_weighted, coverage, miscoverage, avg_set_size """ set_seed(run_seed) - train_ds, val_ds, cal_ds, _ = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=run_seed + train_ds, val_ds, cal_ds, _ = _do_split( + sample_dataset, args.ratios, run_seed, args.split_type ) print(f" Split — Train: {len(train_ds)}, Val: {len(val_ds)}, " f"Cal: {len(cal_ds)}, Test: {len(test_ds)} (fixed)") @@ -162,68 +212,75 @@ def _run_one_seed( if len(val_ds) else None ) - print(" Training ContraWR...") - model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + if args.model == "tfm": + model = TFMTokenizer(dataset=sample_dataset).to(device) + _load_tfm_weights(model, args, run_idx) + else: + model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + print(" Training ContraWR...") + trainer_tmp = Trainer(model=model, device=device, enable_logging=False) + trainer_tmp.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=epochs, + monitor="accuracy" if val_loader is not None else None, + ) trainer = Trainer(model=model, device=device, enable_logging=False) - trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=epochs, - monitor="accuracy" if val_loader is not None else None, - ) - # Base model metrics on fixed test set + # Base model metrics — computed once, shared across all alphas y_true_base, y_prob_base, _ = trainer.inference(test_loader) base_metrics = get_metrics_fn("multiclass")( y_true_base, y_prob_base, metrics=["accuracy", "f1_weighted"] ) - # Extract embeddings for train and cal splits (model-dependent, must redo each seed) + # Extract embeddings once — reused for every alpha print(" Extracting embeddings for train and calibration splits...") train_embeddings = extract_embeddings(model, train_ds, batch_size=args.batch_size, device=device) cal_embeddings = extract_embeddings(model, cal_ds, batch_size=args.batch_size, device=device) - # Conformal calibration + evaluation - print(" Calibrating ClusterLabel predictor...") - cluster_predictor = ClusterLabel( - model=model, - alpha=float(args.alpha), - n_clusters=args.n_clusters, - random_state=run_seed, - ) - cluster_predictor.calibrate( - cal_dataset=cal_ds, - train_embeddings=train_embeddings, - cal_embeddings=cal_embeddings, - ) + # Calibration + evaluation — fast; loop over every alpha + results = {} + for alpha in alphas: + print(f" Calibrating ClusterLabel predictor (alpha={alpha})...") + cluster_predictor = ClusterLabel( + model=model, + alpha=float(alpha), + n_clusters=args.n_clusters, + random_state=run_seed, + ) + cluster_predictor.calibrate( + cal_dataset=cal_ds, + train_embeddings=train_embeddings, + cal_embeddings=cal_embeddings, + ) - print(" Evaluating ClusterLabel predictor on test set...") - y_true, y_prob, _, extra = Trainer(model=cluster_predictor).inference( - test_loader, additional_outputs=["y_predset"] - ) - conf_metrics = get_metrics_fn("multiclass")( - y_true, y_prob, - metrics=["accuracy", "miscoverage_ps"], - y_predset=extra["y_predset"], - ) + y_true, y_prob, _, extra = Trainer(model=cluster_predictor).inference( + test_loader, additional_outputs=["y_predset"] + ) + conf_metrics = get_metrics_fn("multiclass")( + y_true, y_prob, + metrics=["accuracy", "miscoverage_ps"], + y_predset=extra["y_predset"], + ) - predset = extra["y_predset"] - predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset - avg_set_size = predset_t.float().sum(dim=1).mean().item() + predset = extra["y_predset"] + predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset + avg_set_size = predset_t.float().sum(dim=1).mean().item() - miscoverage = conf_metrics["miscoverage_ps"] - if isinstance(miscoverage, np.ndarray): - miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) - else: - miscoverage = float(miscoverage) + miscoverage = conf_metrics["miscoverage_ps"] + if isinstance(miscoverage, np.ndarray): + miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) + else: + miscoverage = float(miscoverage) - return { - "accuracy": float(base_metrics["accuracy"]), - "f1_weighted": float(base_metrics["f1_weighted"]), - "coverage": 1.0 - miscoverage, - "miscoverage": miscoverage, - "avg_set_size": avg_set_size, - } + results[alpha] = { + "accuracy": float(base_metrics["accuracy"]), + "f1_weighted": float(base_metrics["f1_weighted"]), + "coverage": 1.0 - miscoverage, + "miscoverage": miscoverage, + "avg_set_size": avg_set_size, + } + return results def _print_single_run_results(metrics: dict, alpha: float, n_clusters: int) -> None: @@ -314,9 +371,14 @@ def _main(args: argparse.Namespace) -> None: print("\n" + "=" * 80) print("STEP 2: Extract fixed test set (TUH eval partition — same for all seeds)") print("=" * 80) - _, _, _, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=list(args.ratios), seed=args.seed + _, _, _, test_ds = _do_split( + sample_dataset, args.ratios, args.seed, args.split_type ) + if len(test_ds) == 0 and args.quick_test: + print(" [quick-test] TUH eval partition empty in dev mode — using random 20% as test set.") + _, _, _, test_ds = split_by_sample_conformal( + dataset=sample_dataset, ratios=[0.6, 0.1, 0.1, 0.2], seed=args.seed + ) test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) print(f"Test: {len(test_ds)} (fixed)") @@ -328,37 +390,42 @@ def _main(args: argparse.Namespace) -> None: else: run_seeds = [args.seed + i for i in range(args.n_seeds)] + alphas = [float(a.strip()) for a in args.alphas.split(",")] if args.alphas else [args.alpha] + use_multi_seed = len(run_seeds) > 1 print(f"\nRun config: {'multi-seed (' + str(len(run_seeds)) + ' runs)' if use_multi_seed else 'single run'}") - print(f"Seeds: {run_seeds}, alpha={args.alpha}, n_clusters={args.n_clusters}, " - f"target coverage={1 - args.alpha:.0%}") + print(f"Seeds: {run_seeds}, alphas={alphas}, n_clusters={args.n_clusters}") # ------------------------------------------------------------------------- - # STEP 3+: Train + conformal (once per seed) + # STEP 3+: Train once per seed; calibrate for every alpha (fast) # ------------------------------------------------------------------------- - all_metrics = [] + all_metrics = {alpha: [] for alpha in alphas} for run_i, run_seed in enumerate(run_seeds): print("\n" + "=" * 80) if use_multi_seed: print(f"Run {run_i + 1} / {len(run_seeds)} (seed={run_seed})") else: - print(f"STEP 3–4: Train ContraWR + Conformal Calibration (seed={run_seed})") + print(f"STEP 3–4: Train + Conformal Calibration (seed={run_seed})") print("=" * 80) - metrics = _run_one_seed( - args, sample_dataset, test_ds, test_loader, device, epochs, run_seed + seed_results = _run_one_seed( + args, sample_dataset, test_ds, test_loader, device, epochs, run_seed, alphas, + run_idx=run_i, ) - all_metrics.append(metrics) + for alpha in alphas: + all_metrics[alpha].append(seed_results[alpha]) if use_multi_seed: - print(f" [Run {run_i + 1} result] " - f"acc={metrics['accuracy']:.4f}, f1={metrics['f1_weighted']:.4f}, " - f"cov={metrics['coverage']:.4f}, set_size={metrics['avg_set_size']:.2f}") - - if not use_multi_seed: - _print_single_run_results(all_metrics[0], args.alpha, args.n_clusters) - else: - _print_multi_seed_summary(all_metrics, run_seeds, args.alpha, len(test_ds), args.n_clusters) + m = seed_results[alphas[0]] + print(f" [Run {run_i + 1} result (alpha={alphas[0]})] " + f"acc={m['accuracy']:.4f}, f1={m['f1_weighted']:.4f}, " + f"cov={m['coverage']:.4f}, set_size={m['avg_set_size']:.2f}") + + for alpha in alphas: + if not use_multi_seed: + _print_single_run_results(all_metrics[alpha][0], alpha, args.n_clusters) + else: + _print_multi_seed_summary(all_metrics[alpha], run_seeds, alpha, len(test_ds), args.n_clusters) def main() -> None: diff --git a/examples/conformal_eeg/tuev_ncp_conformal.py b/examples/conformal_eeg/tuev_ncp_conformal.py index 0cbe9a48c..2721656b0 100644 --- a/examples/conformal_eeg/tuev_ncp_conformal.py +++ b/examples/conformal_eeg/tuev_ncp_conformal.py @@ -20,6 +20,7 @@ from __future__ import annotations import argparse +import os import random import sys from pathlib import Path @@ -47,8 +48,8 @@ def flush(self): from pyhealth.calib.predictionset.cluster import NeighborhoodLabel from pyhealth.calib.utils import extract_embeddings -from pyhealth.datasets import TUEVDataset, get_dataloader, split_by_sample_conformal_tuh -from pyhealth.models import ContraWR +from pyhealth.datasets import TUEVDataset, get_dataloader, split_by_patient_conformal_tuh, split_by_sample_conformal_tuh, split_by_sample_conformal +from pyhealth.models import ContraWR, TFMTokenizer from pyhealth.tasks import EEGEventsTUEV from pyhealth.trainer import Trainer, get_metrics_fn @@ -86,6 +87,10 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--epochs", type=int, default=20) parser.add_argument("--alpha", type=float, default=0.1, help="Miscoverage rate (e.g., 0.1 => 90% target coverage).") + parser.add_argument( + "--alphas", type=str, default=None, + help="Comma-separated miscoverage rates, e.g. '0.2,0.1,0.05,0.01'. Overrides --alpha.", + ) parser.add_argument( "--ratios", type=float, @@ -107,6 +112,10 @@ def parse_args() -> argparse.Namespace: help="Temperature for NCP exponential weights; smaller => more localization.", ) parser.add_argument("--n-fft", type=int, default=128, help="STFT FFT size used by ContraWR.") + parser.add_argument( + "--model", type=str, default="contrawr", choices=["contrawr", "tfm"], + help="Backbone model: 'contrawr' (default) or 'tfm' (TFMTokenizer).", + ) parser.add_argument( "--device", type=str, @@ -124,9 +133,48 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Smoke test: dev=True, max 2000 samples, 2 epochs, ~5-10 min.", ) + parser.add_argument( + "--weights-dir", + type=str, + default="weightfiles/TFM_Tokenizer_multiple_finetuned_on_TUEV", + help="Root folder of fine-tuned TFM classifier checkpoints (only with --model tfm).", + ) + parser.add_argument( + "--tokenizer-weights", + type=str, + default="weightfiles/tfm_tokenizer_last.pth", + help="Path to the pre-trained TFM tokenizer weights (only with --model tfm).", + ) + parser.add_argument( + "--split-type", + type=str, + default="patient", + choices=["patient", "sample"], + help="Split strategy: 'patient' (default, patient-level, no leakage) or " + "'sample' (original sample-level, for comparison).", + ) return parser.parse_args() +def _do_split(dataset, ratios, seed, split_type): + """Dispatch to the correct TUH split function based on split_type.""" + if split_type == "patient": + return split_by_patient_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + else: + return split_by_sample_conformal_tuh(dataset=dataset, ratios=list(ratios), seed=seed) + + +def _load_tfm_weights(model, args, run_idx: int) -> None: + """Load pre-trained tokenizer + fine-tuned classifier for run_idx (0-based).""" + base = os.path.basename(args.weights_dir) + classifier_path = os.path.join(args.weights_dir, f"{base}_{run_idx + 1}", "best_model.pth") + print(f" Loading TFM weights (run {run_idx + 1}): {classifier_path}") + model.load_pretrained_weights( + tokenizer_checkpoint_path=args.tokenizer_weights, + classifier_checkpoint_path=classifier_path, + ) + + def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) @@ -135,15 +183,13 @@ def set_seed(seed: int) -> None: torch.cuda.manual_seed_all(seed) -def _split_train_pool_for_run(sample_dataset, ratios, run_seed): +def _split_train_pool_for_run(sample_dataset, ratios, run_seed, split_type="patient"): """Re-split the TUH train partition into train/val/cal for one run seed. The test set (TUH eval partition) is always fixed regardless of seed, so only train/val/cal change across runs in multi-seed mode. """ - train_ds, val_ds, cal_ds, _ = split_by_sample_conformal_tuh( - sample_dataset, ratios=ratios, seed=run_seed - ) + train_ds, val_ds, cal_ds, _ = _do_split(sample_dataset, ratios, run_seed, split_type) return train_ds, val_ds, cal_ds @@ -156,85 +202,81 @@ def _run_one_ncp( args, device, epochs, - return_metrics=False, + alphas: list, + run_idx: int = 0, ): - """Train ContraWR, calibrate NCP, evaluate on test. Optionally return metrics dict for aggregation.""" + """Train model + calibrate NCP for one seed across all alphas. + + Training, embedding extraction, and base inference are done once; calibration + loops over alphas (fast — only threshold recomputed per alpha). + + Returns {alpha: metrics_dict} where metrics_dict has keys: + accuracy, f1_weighted, coverage, miscoverage, avg_set_size + """ 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) if len(val_ds) else None - print("\n" + "=" * 80) - print("STEP 3: Train ContraWR") - print("=" * 80) - model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + if args.model == "tfm": + model = TFMTokenizer(dataset=sample_dataset).to(device) + _load_tfm_weights(model, args, run_idx) + else: + model = ContraWR(dataset=sample_dataset, n_fft=args.n_fft).to(device) + print(" Training ContraWR...") + trainer_tmp = Trainer(model=model, device=device, enable_logging=False) + trainer_tmp.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=epochs, + monitor="accuracy" if val_loader is not None else None, + ) trainer = Trainer(model=model, device=device, enable_logging=False) - trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=epochs, - monitor="accuracy" if val_loader is not None else None, - ) - print("\nBase model performance on test set:") - y_true_base, y_prob_base, _loss_base = trainer.inference(test_loader) + # Base model metrics — computed once, shared across all alphas + y_true_base, y_prob_base, _ = trainer.inference(test_loader) base_metrics = get_metrics_fn("multiclass")( y_true_base, y_prob_base, metrics=["accuracy", "f1_weighted"] ) - for metric, value in base_metrics.items(): - print(f" {metric}: {value:.4f}") - - print("\n" + "=" * 80) - print("STEP 4: Neighborhood Conformal Prediction (NCP / NeighborhoodLabel)") - print("=" * 80) - print(f"Target miscoverage alpha: {args.alpha} (target coverage {1 - args.alpha:.0%})") - print(f"k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") + # Extract calibration embeddings once — reused for every alpha + print(" Extracting calibration embeddings...") cal_embeddings = extract_embeddings(model, cal_ds, batch_size=args.batch_size, device=device) - if not return_metrics: - print(f" cal_embeddings shape: {cal_embeddings.shape}") - - ncp_predictor = NeighborhoodLabel( - model=model, - alpha=float(args.alpha), - k_neighbors=args.k_neighbors, - lambda_L=args.lambda_L, - ) - ncp_predictor.calibrate(cal_dataset=cal_ds, cal_embeddings=cal_embeddings) - y_true, y_prob, _loss, extra = Trainer(model=ncp_predictor).inference( - test_loader, additional_outputs=["y_predset"] - ) - ncp_metrics = get_metrics_fn("multiclass")( - y_true, y_prob, metrics=["accuracy", "miscoverage_ps"], y_predset=extra["y_predset"] - ) - predset = extra["y_predset"] - if isinstance(predset, np.ndarray): - predset_t = torch.tensor(predset) - else: - predset_t = predset - avg_set_size = predset_t.float().sum(dim=1).mean().item() - miscoverage = ncp_metrics["miscoverage_ps"] - if isinstance(miscoverage, np.ndarray): - miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) - else: - miscoverage = float(miscoverage) - coverage = 1.0 - miscoverage + # Calibration + evaluation — fast; loop over every alpha + results = {} + for alpha in alphas: + print(f" Calibrating NCP predictor (alpha={alpha})...") + ncp_predictor = NeighborhoodLabel( + model=model, + alpha=float(alpha), + k_neighbors=args.k_neighbors, + lambda_L=args.lambda_L, + ) + ncp_predictor.calibrate(cal_dataset=cal_ds, cal_embeddings=cal_embeddings) - if return_metrics: - return { + y_true, y_prob, _, extra = Trainer(model=ncp_predictor).inference( + test_loader, additional_outputs=["y_predset"] + ) + ncp_metrics = get_metrics_fn("multiclass")( + y_true, y_prob, metrics=["accuracy", "miscoverage_ps"], y_predset=extra["y_predset"] + ) + predset = extra["y_predset"] + predset_t = torch.tensor(predset) if isinstance(predset, np.ndarray) else predset + avg_set_size = predset_t.float().sum(dim=1).mean().item() + + miscoverage = ncp_metrics["miscoverage_ps"] + if isinstance(miscoverage, np.ndarray): + miscoverage = float(miscoverage.item() if miscoverage.size == 1 else miscoverage.mean()) + else: + miscoverage = float(miscoverage) + + results[alpha] = { "accuracy": float(base_metrics["accuracy"]), "f1_weighted": float(base_metrics["f1_weighted"]), - "coverage": coverage, + "coverage": 1.0 - miscoverage, "miscoverage": miscoverage, "avg_set_size": avg_set_size, } - - print("\nNCP (NeighborhoodLabel) Results:") - print(f" Accuracy: {base_metrics['accuracy']:.4f}") - print(f" F1 (weighted): {base_metrics['f1_weighted']:.4f}") - print(f" Empirical miscoverage: {miscoverage:.4f}") - print(f" Empirical coverage: {coverage:.4f}") - print(f" Average set size: {avg_set_size:.2f}") - print(f" k_neighbors: {args.k_neighbors}") + return results def main() -> None: @@ -286,83 +328,58 @@ def _run(args: argparse.Namespace) -> None: print(f"Input schema: {sample_dataset.input_schema}") print(f"Output schema: {sample_dataset.output_schema}") - # Experiment configuration (for PI / reporting) - print("\n--- Experiment configuration ---") - print(f" dataset_root: {root}") - print(f" subset: {args.subset}, ratios: train/val/cal = {args.ratios[0]:.2f}/{args.ratios[1]:.2f}/{args.ratios[2]:.2f} (test = TUH eval partition)") - print(f" alpha: {args.alpha} (target coverage {1 - args.alpha:.0%})") - print(f" k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") - print(f" epochs: {epochs}, batch_size: {args.batch_size}, device: {device}, seed: {args.seed}") - if len(sample_dataset) == 0: raise RuntimeError("No samples produced. Verify TUEV root/subset/task.") + # Parse alphas and run seeds + alphas = [float(a.strip()) for a in args.alphas.split(",")] if args.alphas else [args.alpha] ratios = list(args.ratios) use_multi_seed = args.n_seeds > 1 or args.seeds is not None - if use_multi_seed: - run_seeds = ( - [int(s.strip()) for s in args.seeds.split(",")] - if args.seeds - else [args.seed + i for i in range(args.n_seeds)] - ) - n_runs = len(run_seeds) - print(f" multi_seed: n_runs={n_runs}, run_seeds={run_seeds}, split_seed={args.split_seed} (fixed test set)") - print(f"Multi-seed mode: {n_runs} runs (fixed test set), run seeds: {run_seeds}") - - if not use_multi_seed: - # Single run: original behavior - print("\n" + "=" * 80) - print("STEP 2: Split train/val/cal/test") - print("=" * 80) - train_ds, val_ds, cal_ds, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=ratios, seed=args.seed - ) - print(f"Train: {len(train_ds)}") - print(f"Val: {len(val_ds)}") - print(f"Cal: {len(cal_ds)}") - print(f"Test: {len(test_ds)}") - - test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) - _run_one_ncp( - sample_dataset=sample_dataset, - train_ds=train_ds, - val_ds=val_ds, - cal_ds=cal_ds, - test_loader=test_loader, - args=args, - device=device, - epochs=epochs, - ) - print("\n--- Split sizes and seed (for reporting) ---") - print(f" train={len(train_ds)}, val={len(val_ds)}, cal={len(cal_ds)}, test={len(test_ds)}, seed={args.seed}") - return + run_seeds = ( + [int(s.strip()) for s in args.seeds.split(",")] + if args.seeds + else [args.seed + i for i in range(args.n_seeds)] + ) + n_runs = len(run_seeds) - # Multi-seed: test set is always fixed as the TUH eval partition (no split-seed needed). - # Each run uses a different seed to re-shuffle the train pool into train/val/cal. + # ------------------------------------------------------------------------- + # STEP 2: Extract the fixed test set ONCE + # ------------------------------------------------------------------------- print("\n" + "=" * 80) - print("STEP 2: Fix test set (TUH eval partition), then run multiple train/cal splits") + print("STEP 2: Extract fixed test set (TUH eval partition — same for all seeds)") print("=" * 80) - # Get the fixed test set — seed doesn't affect which samples are in eval, only the - # train-pool shuffle, so any seed works here. - _, _, _, test_ds = split_by_sample_conformal_tuh( - dataset=sample_dataset, ratios=ratios, seed=args.split_seed + _, _, _, test_ds = _do_split( + sample_dataset, ratios, args.split_seed, args.split_type ) + if len(test_ds) == 0 and args.quick_test: + print(" [quick-test] TUH eval partition empty in dev mode — using random 20% as test set.") + _, _, _, test_ds = split_by_sample_conformal( + dataset=sample_dataset, ratios=[0.6, 0.1, 0.1, 0.2], seed=args.split_seed + ) test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) n_test = len(test_ds) - print(f"Fixed test set size: {n_test} (TUH eval partition)") + print(f"Test: {n_test} (fixed)") + + print(f"\nRun config: {'multi-seed (' + str(n_runs) + ' runs)' if use_multi_seed else 'single run'}") + print(f"Seeds: {run_seeds}, alphas={alphas}, k_neighbors={args.k_neighbors}") - accs, f1s, coverages, miscoverages, set_sizes = [], [], [], [], [] + # ------------------------------------------------------------------------- + # STEP 3+: Train once per seed; calibrate for every alpha (fast) + # ------------------------------------------------------------------------- + all_metrics = {alpha: [] for alpha in alphas} for run_i, run_seed in enumerate(run_seeds): print("\n" + "=" * 80) - print(f"Run {run_i + 1} / {n_runs} (seed={run_seed})") + if use_multi_seed: + print(f"Run {run_i + 1} / {n_runs} (seed={run_seed})") + else: + print(f"STEP 3–4: Train + NCP Calibration (seed={run_seed})") print("=" * 80) set_seed(run_seed) - train_ds, val_ds, cal_ds = _split_train_pool_for_run( - sample_dataset, ratios, run_seed - ) - print(f"Train: {len(train_ds)}, Val: {len(val_ds)}, Cal: {len(cal_ds)}") + train_ds, val_ds, cal_ds = _split_train_pool_for_run(sample_dataset, ratios, run_seed, args.split_type) + print(f" Split — Train: {len(train_ds)}, Val: {len(val_ds)}, " + f"Cal: {len(cal_ds)}, Test: {n_test} (fixed)") - metrics = _run_one_ncp( + seed_results = _run_one_ncp( sample_dataset=sample_dataset, train_ds=train_ds, val_ds=val_ds, @@ -371,47 +388,62 @@ def _run(args: argparse.Namespace) -> None: args=args, device=device, epochs=epochs, - return_metrics=True, + alphas=alphas, + run_idx=run_i, ) - accs.append(metrics["accuracy"]) - f1s.append(metrics["f1_weighted"]) - coverages.append(metrics["coverage"]) - miscoverages.append(metrics["miscoverage"]) - set_sizes.append(metrics["avg_set_size"]) - - accs = np.array(accs) - f1s = np.array(f1s) - coverages = np.array(coverages) - miscoverages_arr = np.array(miscoverages) - set_sizes = np.array(set_sizes) - - # Per-run table (for PI / reporting) - print("\n" + "=" * 80) - print("Per-run NCP results (fixed test set = TUH eval partition)") - print("=" * 80) - print(f" {'Run':<4} {'Seed':<6} {'Accuracy':<10} {'F1-Wt':<10} " - f"{'Coverage':<10} {'Miscoverage':<12} {'Avg set size':<12}") - print(" " + "-" * 68) - for i in range(n_runs): - print(f" {i+1:<4} {run_seeds[i]:<6} {accs[i]:<10.4f} {f1s[i]:<10.4f} " - f"{coverages[i]:<10.4f} {miscoverages_arr[i]:<12.4f} {set_sizes[i]:<12.2f}") - - print("\n" + "=" * 80) - print("NCP summary (mean \u00b1 std over {} runs, fixed test set)".format(n_runs)) - print("=" * 80) - print(f" Accuracy: {accs.mean():.4f} \u00b1 {accs.std():.4f}") - print(f" F1 (weighted): {f1s.mean():.4f} \u00b1 {f1s.std():.4f}") - print(f" Empirical coverage: {coverages.mean():.4f} \u00b1 {coverages.std():.4f}") - print(f" Empirical miscoverage: {miscoverages_arr.mean():.4f} \u00b1 {miscoverages_arr.std():.4f}") - print(f" Average set size: {set_sizes.mean():.2f} \u00b1 {set_sizes.std():.2f}") - print(f" Target coverage: {1 - args.alpha:.0%} (alpha={args.alpha})") - print(f" k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") - print(f" Test set size: {n_test} (fixed across runs)") - print(f" Run seeds: {run_seeds}") - print("\n--- Min / Max (across runs) ---") - print(f" Coverage: [{coverages.min():.4f}, {coverages.max():.4f}]") - print(f" Set size: [{set_sizes.min():.2f}, {set_sizes.max():.2f}]") - print(f" Accuracy: [{accs.min():.4f}, {accs.max():.4f}]") + for alpha in alphas: + all_metrics[alpha].append(seed_results[alpha]) + + if use_multi_seed: + m = seed_results[alphas[0]] + print(f" [Run {run_i + 1} result (alpha={alphas[0]})] " + f"acc={m['accuracy']:.4f}, f1={m['f1_weighted']:.4f}, " + f"cov={m['coverage']:.4f}, set_size={m['avg_set_size']:.2f}") + + for alpha in alphas: + mlist = all_metrics[alpha] + accs = np.array([m["accuracy"] for m in mlist]) + f1s = np.array([m["f1_weighted"] for m in mlist]) + coverages = np.array([m["coverage"] for m in mlist]) + miscovs = np.array([m["miscoverage"] for m in mlist]) + set_sizes = np.array([m["avg_set_size"] for m in mlist]) + + if not use_multi_seed: + print(f"\nNCP Results (alpha={alpha}):") + print(f" Accuracy: {accs[0]:.4f}") + print(f" F1 (weighted): {f1s[0]:.4f}") + print(f" Empirical coverage: {coverages[0]:.4f}") + print(f" Empirical miscoverage: {miscovs[0]:.4f}") + print(f" Average set size: {set_sizes[0]:.2f}") + print(f" Target coverage: {1 - alpha:.0%} (alpha={alpha})") + print(f" k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") + else: + print("\n" + "=" * 80) + print(f"Per-run NCP results — alpha={alpha} (target coverage={1-alpha:.0%})") + print("=" * 80) + print(f" {'Run':<4} {'Seed':<6} {'Accuracy':<10} {'F1-Wt':<10} " + f"{'Coverage':<10} {'Miscoverage':<12} {'Avg set size':<12}") + print(" " + "-" * 68) + for i in range(n_runs): + print(f" {i+1:<4} {run_seeds[i]:<6} {accs[i]:<10.4f} {f1s[i]:<10.4f} " + f"{coverages[i]:<10.4f} {miscovs[i]:<12.4f} {set_sizes[i]:<12.2f}") + + print("\n" + "=" * 80) + print(f"NCP summary — alpha={alpha} (mean ± std over {n_runs} runs, fixed test set)") + print("=" * 80) + print(f" Accuracy: {accs.mean():.4f} ± {accs.std():.4f}") + print(f" F1 (weighted): {f1s.mean():.4f} ± {f1s.std():.4f}") + print(f" Empirical coverage: {coverages.mean():.4f} ± {coverages.std():.4f}") + print(f" Empirical miscoverage: {miscovs.mean():.4f} ± {miscovs.std():.4f}") + print(f" Average set size: {set_sizes.mean():.2f} ± {set_sizes.std():.2f}") + print(f" Target coverage: {1 - alpha:.0%} (alpha={alpha})") + print(f" k_neighbors: {args.k_neighbors}, lambda_L: {args.lambda_L}") + print(f" Test set size: {n_test} (fixed across runs)") + print(f" Run seeds: {run_seeds}") + print("\n--- Min / Max (across runs) ---") + print(f" Coverage: [{coverages.min():.4f}, {coverages.max():.4f}]") + print(f" Set size: [{set_sizes.min():.2f}, {set_sizes.max():.2f}]") + print(f" Accuracy: [{accs.min():.4f}, {accs.max():.4f}]") if __name__ == "__main__": diff --git a/examples/cxr/mimic4_cxr_sunlab_stats.py b/examples/cxr/mimic4_cxr_sunlab_stats.py new file mode 100644 index 000000000..eb70f12a9 --- /dev/null +++ b/examples/cxr/mimic4_cxr_sunlab_stats.py @@ -0,0 +1,80 @@ +"""Example script to load Sunlab MIMIC-CXR through MIMIC4Dataset +and print stats. + +Usage: + python examples/cxr/mimic4_cxr_sunlab_stats.py \ + --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \ + --note-root /shared/rsaas/physionet.org/files/mimic-note \ + --cxr-root /shared/rsaas/physionet.org/files/MIMIC-CXR \ + --cache-dir /shared/eng/pyhealth +""" + +import argparse + +from pyhealth.datasets import MIMIC4Dataset + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Load Sunlab MIMIC-CXR variant and run dataset.stats()." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + help="Root directory for MIMIC-IV EHR data.", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + help="Root directory for MIMIC-IV notes data.", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + help="Root directory for the Sunlab MIMIC-CXR mirror.", + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + help="Optional cache directory for PyHealth dataset artifacts.", + ) + parser.add_argument( + "--num-workers", + type=int, + default=4, + help="Number of workers for dataset processing.", + ) + parser.add_argument( + "--dev", + action="store_true", + help="Enable dev mode (small subset).", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant="sunlab", + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + num_workers=args.num_workers, + dev=args.dev, + ) + + # Prints table/patient/event statistics from BaseDataset. + dataset.stats() + + +if __name__ == "__main__": + main() diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_adacare.py b/examples/drug_recommendation/drug_recommendation_mimic4_adacare.py new file mode 100644 index 000000000..2fbad0e38 --- /dev/null +++ b/examples/drug_recommendation/drug_recommendation_mimic4_adacare.py @@ -0,0 +1,121 @@ +""" +Example of using AdaCare for drug recommendation on MIMIC-IV. + +This example demonstrates: +1. Loading MIMIC-IV data +2. Applying the DrugRecommendationMIMIC4 task +3. Creating a SampleDataset with nested sequence processors +4. Training an AdaCare model +""" + +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, +) +from pyhealth.models import AdaCare +from pyhealth.tasks import DrugRecommendationMIMIC4 +from pyhealth.trainer import Trainer + +if __name__ == "__main__": + # STEP 1: Load MIMIC-IV base dataset + base_dataset = MIMIC4Dataset( + ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", + cache_dir="/shared/eng/pyhealth_agent/baselines", # Change this to your desired cache directory + ehr_tables=[ + "patients", + "admissions", + "diagnoses_icd", + "procedures_icd", + "prescriptions", + ], + ) + + # STEP 2: Apply drug recommendation task + sample_dataset = base_dataset.set_task( + DrugRecommendationMIMIC4(), + num_workers=4, + ) + + print(f"Total samples: {len(sample_dataset)}") + print(f"Input schema: {sample_dataset.input_schema}") + print(f"Output schema: {sample_dataset.output_schema}") + + # Inspect a sample + sample = sample_dataset[0] + print("\nSample structure:") + print(f" Patient ID: {sample['patient_id']}") + print(f" Visit ID: {sample['visit_id']}") + print(f" Conditions (history): {len(sample['conditions'])} visits") + print(f" Procedures (history): {len(sample['procedures'])} visits") + print(f" Drugs history: {len(sample['drugs_hist'])} visits") + print(f" Target drugs: {len(sample['drugs'])} drugs") + print(f"\n First visit conditions: {sample['conditions'][0][:5]}...") + print(f" Target drugs sample: {sample['drugs'][:5]}...") + + # STEP 3: Split dataset + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + print("\nDataset split:") + print(f" Train: {len(train_dataset)} samples") + print(f" Validation: {len(val_dataset)} samples") + print(f" Test: {len(test_dataset)} samples") + + # Create dataloaders + train_loader = get_dataloader(train_dataset, batch_size=64, shuffle=True) + val_loader = get_dataloader(val_dataset, batch_size=64, shuffle=False) + test_loader = get_dataloader(test_dataset, batch_size=64, shuffle=False) + + # STEP 4: Initialize AdaCare model + model = AdaCare( + dataset=sample_dataset, + embedding_dim=128, + hidden_dim=128, + ) + + num_params = sum(p.numel() for p in model.parameters()) + print(f"\nModel initialized with {num_params:,} parameters") + print(f"Feature keys: {model.feature_keys}") + print(f"Label key: {model.label_keys[0]}") + + # STEP 5: Train the model + trainer = Trainer( + model=model, + device="cuda:0", # or "cpu" + metrics=["pr_auc_samples", "f1_samples", "jaccard_samples"], + ) + + print("\nStarting training...") + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=50, + monitor="pr_auc_samples", + optimizer_params={"lr": 1e-4}, + optimizer_class=torch.optim.AdamW, + ) + + # STEP 6: Evaluate on test set + print("\nEvaluating on test set...") + results = trainer.evaluate(test_loader) + print("\nTest Results:") + for metric, value in results.items(): + print(f" {metric}: {value:.4f}") + + # STEP 7: Inspect model predictions + print("\nSample predictions:") + sample_batch = next(iter(test_loader)) + + with torch.no_grad(): + output = model(**sample_batch) + + print(f" Batch size: {output['y_prob'].shape[0]}") + print(f" Number of drug classes: {output['y_prob'].shape[1]}") + print(" Predicted probabilities (first 5 drugs of first patient):") + print(f" {output['y_prob'][0, :5].cpu().numpy()}") + print(" True labels (first 5 drugs of first patient):") + print(f" {output['y_true'][0, :5].cpu().numpy()}") diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_adacare_optuna.py b/examples/drug_recommendation/drug_recommendation_mimic4_adacare_optuna.py new file mode 100644 index 000000000..ee618bfb7 --- /dev/null +++ b/examples/drug_recommendation/drug_recommendation_mimic4_adacare_optuna.py @@ -0,0 +1,201 @@ +""" +Optuna hyperparameter tuning for AdaCare on drug recommendation with MIMIC-IV. + +This example demonstrates: +1. Loading MIMIC-IV data and applying the DrugRecommendationMIMIC4 task +2. Defining an Optuna objective that tunes AdaCare-specific hyperparameters +3. Running 10 Optuna trials to find the best configuration +4. Training a final model with the best hyperparameters + +Tuned hyperparameters: + - embedding_dim: embedding size for code tokens + - hidden_dim: GRU hidden state size inside AdaCare + - lr: learning rate for AdamW + - weight_decay: L2 regularization coefficient for AdamW + +Note: + AdaCare.__init__ forwards **kwargs to BaseModel.__init__, which only + accepts `dataset`. Layer-specific parameters (kernel_size, kernel_num, + r_v, r_c, activation, dropout) must not be passed to AdaCare() directly. + The tunable surface here covers the explicit named parameters + (embedding_dim, hidden_dim) and the optimizer settings. +""" + +import torch +import optuna + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, +) +from pyhealth.models import AdaCare +from pyhealth.tasks import DrugRecommendationMIMIC4 +from pyhealth.trainer import Trainer + +if __name__ == "__main__": + # --------------------------------------------------------------------------- + # STEP 1: Load MIMIC-IV base dataset + # --------------------------------------------------------------------------- + base_dataset = MIMIC4Dataset( + ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", + cache_dir="/shared/eng/pyhealth_agent/baselines", + ehr_tables=[ + "patients", + "admissions", + "diagnoses_icd", + "procedures_icd", + "prescriptions", + ], + ) + + # STEP 2: Apply drug recommendation task + sample_dataset = base_dataset.set_task( + DrugRecommendationMIMIC4(), + num_workers=4, + ) + + print(f"Total samples: {len(sample_dataset)}") + print(f"Input schema: {sample_dataset.input_schema}") + print(f"Output schema: {sample_dataset.output_schema}") + + # STEP 3: Split dataset (fixed split so all trials see the same data) + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + print(f"\nDataset split — Train: {len(train_dataset)} " + f"Val: {len(val_dataset)} Test: {len(test_dataset)}") + + # --------------------------------------------------------------------------- + # STEP 4: Define Optuna objective + # --------------------------------------------------------------------------- + DEVICE = "cuda:3" # or "cpu" + TUNE_EPOCHS = 10 # lightweight training per trial + N_TRIALS = 10 + + def objective(trial: optuna.Trial) -> float: + """Suggest hyperparameters and return val pr_auc_samples.""" + + # --- Suggest hyperparameters --------------------------------------- + embedding_dim = trial.suggest_categorical( + "embedding_dim", [64, 128, 256] + ) + hidden_dim = trial.suggest_categorical("hidden_dim", [64, 128, 256]) + lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True) + weight_decay = trial.suggest_float( + "weight_decay", 1e-6, 1e-2, log=True + ) + batch_size = trial.suggest_categorical("batch_size", [32, 64, 128]) + + # --- Build dataloaders --------------------------------------------- + train_loader = get_dataloader( + train_dataset, batch_size=batch_size, shuffle=True + ) + val_loader = get_dataloader( + val_dataset, batch_size=batch_size, shuffle=False + ) + + # --- Build model --------------------------------------------------- + model = AdaCare( + dataset=sample_dataset, + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + ) + + # --- Train --------------------------------------------------------- + trainer = Trainer( + model=model, + device=DEVICE, + metrics=["pr_auc_samples"], + ) + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=TUNE_EPOCHS, + monitor="pr_auc_samples", + optimizer_class=torch.optim.AdamW, + optimizer_params={"lr": lr}, + weight_decay=weight_decay, + ) + + # --- Evaluate on validation set ------------------------------------ + scores = trainer.evaluate(val_loader) + return scores["pr_auc_samples"] + + # --------------------------------------------------------------------------- + # STEP 5: Run Optuna study + # --------------------------------------------------------------------------- + print( + f"\nStarting Optuna search: " + f"{N_TRIALS} trials, {TUNE_EPOCHS} epochs each..." + ) + + study = optuna.create_study(direction="maximize") + study.optimize(objective, n_trials=N_TRIALS) + + best_params = study.best_params + print("\nBest hyperparameters found:") + for k, v in best_params.items(): + print(f" {k}: {v}") + print(f"Best validation pr_auc_samples: {study.best_value:.4f}") + + # --------------------------------------------------------------------------- + # STEP 6: Train final model with best hyperparameters + # --------------------------------------------------------------------------- + print("\nTraining final model with best hyperparameters...") + + train_loader = get_dataloader( + train_dataset, batch_size=best_params["batch_size"], shuffle=True + ) + val_loader = get_dataloader( + val_dataset, batch_size=best_params["batch_size"], shuffle=False + ) + test_loader = get_dataloader( + test_dataset, batch_size=best_params["batch_size"], shuffle=False + ) + + final_model = AdaCare( + dataset=sample_dataset, + embedding_dim=best_params["embedding_dim"], + hidden_dim=best_params["hidden_dim"], + ) + + num_params = sum(p.numel() for p in final_model.parameters()) + print(f"Final model: {num_params:,} parameters") + + final_trainer = Trainer( + model=final_model, + device=DEVICE, + metrics=["pr_auc_samples", "f1_samples", "jaccard_samples"], + ) + final_trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=50, + monitor="pr_auc_samples", + optimizer_class=torch.optim.AdamW, + optimizer_params={"lr": best_params["lr"]}, + weight_decay=best_params["weight_decay"], + ) + + # STEP 7: Evaluate on test set + print("\nEvaluating on test set...") + results = final_trainer.evaluate(test_loader) + print("\nTest Results:") + for metric, value in results.items(): + print(f" {metric}: {value:.4f}") + + # STEP 8: Inspect model predictions + print("\nSample predictions:") + sample_batch = next(iter(test_loader)) + + with torch.no_grad(): + output = final_model(**sample_batch) + + print(f" Batch size: {output['y_prob'].shape[0]}") + print(f" Number of drug classes: {output['y_prob'].shape[1]}") + print(" Predicted probabilities (first 5 drugs of first patient):") + print(f" {output['y_prob'][0, :5].cpu().numpy()}") + print(" True labels (first 5 drugs of first patient):") + print(f" {output['y_true'][0, :5].cpu().numpy()}") diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py b/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py index 6aeb2c8ce..bd5b33cb0 100644 --- a/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py +++ b/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py @@ -25,9 +25,6 @@ def prepare_drug_task_data(): mimicvi = MIMIC4Dataset( root="/srv/local/data/physionet.org/files/mimiciv/2.0/hosp", tables=["diagnoses_icd", "procedures_icd", "prescriptions"], - code_mapping={"NDC": ("ATC", {"target_kwargs": {"level": 3}})}, - dev=_DEV, - refresh_cache=False, ) print("stat") diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_multimodal_retain.py b/examples/drug_recommendation/drug_recommendation_mimic4_multimodal_retain.py new file mode 100644 index 000000000..a47cce1f0 --- /dev/null +++ b/examples/drug_recommendation/drug_recommendation_mimic4_multimodal_retain.py @@ -0,0 +1,210 @@ +""" +Drug Recommendation on MIMIC-IV with MultimodalRETAIN + +This example demonstrates how to use the MultimodalRETAIN model with mixed +input modalities for drug recommendation on MIMIC-IV. + +The MultimodalRETAIN model can handle: +- Sequential features (visit histories with diagnoses, procedures) → RETAIN processing + with reverse time attention mechanism +- Non-sequential features (demographics, static measurements) → Direct embedding + +This example shows: +1. Loading MIMIC-IV data with mixed feature types +2. Applying a drug recommendation task +3. Training a MultimodalRETAIN model with both sequential and non-sequential inputs +4. Evaluating the model performance +5. Comparing to vanilla RETAIN (sequential only) +""" + +from pyhealth.datasets import MIMIC4Dataset +from pyhealth.datasets import split_by_patient, get_dataloader +from pyhealth.models import MultimodalRETAIN +from pyhealth.tasks import DrugRecommendationMIMIC4 +from pyhealth.trainer import Trainer + + +if __name__ == "__main__": + # STEP 1: Load MIMIC-IV base dataset + print("=" * 60) + print("STEP 1: Loading MIMIC-IV Dataset") + print("=" * 60) + + base_dataset = MIMIC4Dataset( + ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", + ehr_tables=["diagnoses_icd", "procedures_icd", "prescriptions"], + dev=True, # Use development mode for faster testing + num_workers=4, + ) + base_dataset.stats() + + # STEP 2: Apply drug recommendation task with multimodal features + print("\n" + "=" * 60) + print("STEP 2: Setting Drug Recommendation Task") + print("=" * 60) + + # Use the DrugRecommendationMIMIC4 task + # This task creates visit-level nested sequences from diagnoses/procedures + # and recommends drugs for the current visit + task = DrugRecommendationMIMIC4() + sample_dataset = base_dataset.set_task( + task, + num_workers=4, + ) + + print(f"\nTotal samples: {len(sample_dataset)}") + print(f"Input schema: {sample_dataset.input_schema}") + print(f"Output schema: {sample_dataset.output_schema}") + + # Inspect a sample + if len(sample_dataset) > 0: + sample = sample_dataset[0] + print("\nSample structure:") + print(f" Patient ID: {sample['patient_id']}") + for key in sample_dataset.input_schema.keys(): + if key in sample: + if isinstance(sample[key], (list, tuple)): + if sample[key] and isinstance(sample[key][0], (list, tuple)): + print(f" {key}: {len(sample[key])} visits") + else: + print(f" {key}: length {len(sample[key])}") + else: + print(f" {key}: {type(sample[key])}") + # Show drugs key from output + if 'drugs' in sample: + print(f" drugs (target): {len(sample['drugs'])} prescriptions") + + # STEP 3: Split dataset + print("\n" + "=" * 60) + print("STEP 3: Splitting Dataset") + print("=" * 60) + + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + print(f"Train samples: {len(train_dataset)}") + print(f"Val samples: {len(val_dataset)}") + print(f"Test samples: {len(test_dataset)}") + + # Create dataloaders + train_loader = get_dataloader(train_dataset, batch_size=64, shuffle=True) + val_loader = get_dataloader(val_dataset, batch_size=64, shuffle=False) + test_loader = get_dataloader(test_dataset, batch_size=64, shuffle=False) + + # STEP 4: Initialize MultimodalRETAIN model + print("\n" + "=" * 60) + print("STEP 4: Initializing MultimodalRETAIN Model") + print("=" * 60) + + model = MultimodalRETAIN( + dataset=sample_dataset, + embedding_dim=128, + dropout=0.5, + ) + + num_params = sum(p.numel() for p in model.parameters()) + print(f"Model initialized with {num_params:,} parameters") + + # Print feature classification + print(f"\nSequential features (RETAIN processing): {model.sequential_features}") + print(f"Non-sequential features (direct embedding): {model.non_sequential_features}") + + # Calculate expected embedding dimensions + total_dim = len(model.feature_keys) * model.embedding_dim + print(f"\nPatient representation dimension: {total_dim}") + + # STEP 5: Train the model + print("\n" + "=" * 60) + print("STEP 5: Training Model") + print("=" * 60) + + trainer = Trainer( + model=model, + device="cuda:0", # Change to "cpu" if no GPU available + metrics=["pr_auc_samples", "roc_auc_samples", "jaccard_samples", "f1_samples"], + ) + + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=10, + monitor="jaccard_samples", + optimizer_params={"lr": 1e-3}, + ) + + # STEP 6: Evaluate on test set + print("\n" + "=" * 60) + print("STEP 6: Evaluating on Test Set") + print("=" * 60) + + results = trainer.evaluate(test_loader) + print("\nTest Results:") + for metric, value in results.items(): + print(f" {metric}: {value:.4f}") + + # STEP 7: Demonstrate model predictions + print("\n" + "=" * 60) + print("STEP 7: Sample Predictions") + print("=" * 60) + + import torch + + sample_batch = next(iter(test_loader)) + with torch.no_grad(): + output = model(**sample_batch) + + print(f"\nBatch size: {output['y_prob'].shape[0]}") + print(f"Output shape: {output['y_prob'].shape}") + print(f"(batch_size, num_drug_types)") + + # Show first patient predictions + print(f"\nFirst patient top-5 drug recommendations:") + first_patient_probs = output['y_prob'][0] + top5_drugs = torch.topk(first_patient_probs, k=min(5, len(first_patient_probs))) + for i, (drug_idx, prob) in enumerate(zip(top5_drugs.indices, top5_drugs.values)): + print(f" {i+1}. Drug index {drug_idx.item()}: probability {prob.item():.4f}") + + # Show ground truth for first patient + print(f"\nFirst patient ground truth drugs:") + first_patient_true = output['y_true'][0] + true_drug_indices = torch.where(first_patient_true > 0)[0] + print(f" Number of prescribed drugs: {len(true_drug_indices)}") + if len(true_drug_indices) > 0: + print(f" Drug indices: {true_drug_indices.tolist()[:10]}...") + + # STEP 8: Compare with vanilla RETAIN (if applicable) + print("\n" + "=" * 60) + print("STEP 8: Model Architecture Comparison") + print("=" * 60) + + print("\nMultimodalRETAIN vs. Vanilla RETAIN:") + print(" Vanilla RETAIN:") + print(" - Only handles sequential (visit-level) features") + print(" - Processes all features through reverse time attention") + print(" ") + print(" MultimodalRETAIN:") + print(" - Handles both sequential and non-sequential features") + print(f" - Sequential features ({len(model.sequential_features)}): " + f"{model.sequential_features}") + print(f" - Non-sequential features ({len(model.non_sequential_features)}): " + f"{model.non_sequential_features}") + print(" - More flexible for heterogeneous EHR data") + + # Summary + print("\n" + "=" * 60) + print("SUMMARY: MultimodalRETAIN Training Complete") + print("=" * 60) + print(f"Model: MultimodalRETAIN") + print(f"Dataset: MIMIC-IV") + print(f"Task: Drug Recommendation") + print(f"Sequential features: {len(model.sequential_features)}") + print(f"Non-sequential features: {len(model.non_sequential_features)}") + print(f"Best validation Jaccard: {results.get('jaccard_samples', 0):.4f}") + print("\nRETAIN advantages:") + print(" - Reverse time attention for interpretability") + print(" - Visit-level attention weights (alpha)") + print(" - Variable-level attention weights (beta)") + print(" - Multimodal extension allows richer feature sets") + print("=" * 60) + diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_retain.py b/examples/drug_recommendation/drug_recommendation_mimic4_retain.py index 496edc7e3..d39c46042 100644 --- a/examples/drug_recommendation/drug_recommendation_mimic4_retain.py +++ b/examples/drug_recommendation/drug_recommendation_mimic4_retain.py @@ -19,100 +19,102 @@ from pyhealth.tasks import DrugRecommendationMIMIC4 from pyhealth.trainer import Trainer -# STEP 1: Load MIMIC-IV base dataset -base_dataset = MIMIC4Dataset( - ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", - ehr_tables=[ - "patients", - "admissions", - "diagnoses_icd", - "procedures_icd", - "prescriptions", - ], -) - -# STEP 2: Apply drug recommendation task -sample_dataset = base_dataset.set_task( - DrugRecommendationMIMIC4(), - num_workers=4, -) - -print(f"Total samples: {len(sample_dataset)}") -print(f"Input schema: {sample_dataset.input_schema}") -print(f"Output schema: {sample_dataset.output_schema}") - -# Inspect a sample -sample = sample_dataset.samples[0] -print("\nSample structure:") -print(f" Patient ID: {sample['patient_id']}") -print(f" Visit ID: {sample['visit_id']}") -print(f" Conditions (history): {len(sample['conditions'])} visits") -print(f" Procedures (history): {len(sample['procedures'])} visits") -print(f" Drugs history: {len(sample['drugs_hist'])} visits") -print(f" Target drugs: {len(sample['drugs'])} drugs") -print(f"\n First visit conditions: {sample['conditions'][0][:5]}...") -print(f" Target drugs sample: {sample['drugs'][:5]}...") - -# STEP 3: Split dataset -train_dataset, val_dataset, test_dataset = split_by_patient( - sample_dataset, [0.8, 0.1, 0.1] -) - -print("\nDataset split:") -print(f" Train: {len(train_dataset)} samples") -print(f" Validation: {len(val_dataset)} samples") -print(f" Test: {len(test_dataset)} samples") - -# Create dataloaders -train_loader = get_dataloader(train_dataset, batch_size=64, shuffle=True) -val_loader = get_dataloader(val_dataset, batch_size=64, shuffle=False) -test_loader = get_dataloader(test_dataset, batch_size=64, shuffle=False) - -# STEP 4: Initialize RETAIN model -model = RETAIN( - dataset=sample_dataset, - embedding_dim=128, - dropout=0.5, -) - -num_params = sum(p.numel() for p in model.parameters()) -print(f"\nModel initialized with {num_params:,} parameters") -print(f"Feature keys: {model.feature_keys}") -print(f"Label key: {model.label_key}") - -# STEP 5: Train the model -trainer = Trainer( - model=model, - device="cuda:4", # or "cpu" - metrics=["pr_auc_samples", "f1_samples", "jaccard_samples"], -) - -print("\nStarting training...") -trainer.train( - train_dataloader=train_loader, - val_dataloader=val_loader, - epochs=50, - monitor="pr_auc_samples", - optimizer_params={"lr": 1e-3}, -) - -# STEP 6: Evaluate on test set -print("\nEvaluating on test set...") -results = trainer.evaluate(test_loader) -print("\nTest Results:") -for metric, value in results.items(): - print(f" {metric}: {value:.4f}") - -# STEP 7: Inspect model predictions -print("\nSample predictions:") -sample_batch = next(iter(test_loader)) - -with torch.no_grad(): - output = model(**sample_batch) - -print(f" Batch size: {output['y_prob'].shape[0]}") -print(f" Number of drug classes: {output['y_prob'].shape[1]}") -print(" Predicted probabilities (first 5 drugs of first patient):") -print(f" {output['y_prob'][0, :5].cpu().numpy()}") -print(" True labels (first 5 drugs of first patient):") -print(f" {output['y_true'][0, :5].cpu().numpy()}") +if __name__ == "__main__": + # STEP 1: Load MIMIC-IV base dataset + base_dataset = MIMIC4Dataset( + ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", + cache_dir="/shared/eng/pyhealth_agent/baselines", + ehr_tables=[ + "patients", + "admissions", + "diagnoses_icd", + "procedures_icd", + "prescriptions", + ], + ) + + # STEP 2: Apply drug recommendation task + sample_dataset = base_dataset.set_task( + DrugRecommendationMIMIC4(), + num_workers=4, + ) + + print(f"Total samples: {len(sample_dataset)}") + print(f"Input schema: {sample_dataset.input_schema}") + print(f"Output schema: {sample_dataset.output_schema}") + + # Inspect a sample + sample = sample_dataset[0] + print("\nSample structure:") + print(f" Patient ID: {sample['patient_id']}") + print(f" Visit ID: {sample['visit_id']}") + print(f" Conditions (history): {len(sample['conditions'])} visits") + print(f" Procedures (history): {len(sample['procedures'])} visits") + print(f" Drugs history: {len(sample['drugs_hist'])} visits") + print(f" Target drugs: {len(sample['drugs'])} drugs") + print(f"\n First visit conditions: {sample['conditions'][0][:5]}...") + print(f" Target drugs sample: {sample['drugs'][:5]}...") + + # STEP 3: Split dataset + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + print("\nDataset split:") + print(f" Train: {len(train_dataset)} samples") + print(f" Validation: {len(val_dataset)} samples") + print(f" Test: {len(test_dataset)} samples") + + # Create dataloaders + train_loader = get_dataloader(train_dataset, batch_size=64, shuffle=True) + val_loader = get_dataloader(val_dataset, batch_size=64, shuffle=False) + test_loader = get_dataloader(test_dataset, batch_size=64, shuffle=False) + + # STEP 4: Initialize RETAIN model + model = RETAIN( + dataset=sample_dataset, + embedding_dim=128, + dropout=0.5, + ) + + num_params = sum(p.numel() for p in model.parameters()) + print(f"\nModel initialized with {num_params:,} parameters") + print(f"Feature keys: {model.feature_keys}") + print(f"Label key: {model.label_key}") + + # STEP 5: Train the model + trainer = Trainer( + model=model, + device="cuda:4", # or "cpu" + metrics=["pr_auc_samples", "f1_samples", "jaccard_samples"], + ) + + print("\nStarting training...") + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=50, + monitor="pr_auc_samples", + optimizer_params={"lr": 1e-3}, + ) + + # STEP 6: Evaluate on test set + print("\nEvaluating on test set...") + results = trainer.evaluate(test_loader) + print("\nTest Results:") + for metric, value in results.items(): + print(f" {metric}: {value:.4f}") + + # STEP 7: Inspect model predictions + print("\nSample predictions:") + sample_batch = next(iter(test_loader)) + + with torch.no_grad(): + output = model(**sample_batch) + + print(f" Batch size: {output['y_prob'].shape[0]}") + print(f" Number of drug classes: {output['y_prob'].shape[1]}") + print(" Predicted probabilities (first 5 drugs of first patient):") + print(f" {output['y_prob'][0, :5].cpu().numpy()}") + print(" True labels (first 5 drugs of first patient):") + print(f" {output['y_true'][0, :5].cpu().numpy()}") diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_rnn.py b/examples/drug_recommendation/drug_recommendation_mimic4_rnn.py new file mode 100644 index 000000000..5f8f38e08 --- /dev/null +++ b/examples/drug_recommendation/drug_recommendation_mimic4_rnn.py @@ -0,0 +1,123 @@ +""" +Example of using RNN for drug recommendation on MIMIC-IV. + +This example demonstrates: +1. Loading MIMIC-IV data +2. Applying the DrugRecommendationMIMIC4 task +3. Creating a SampleDataset with nested sequence processors +4. Training an RNN model +""" + +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, +) +from pyhealth.models import RNN +from pyhealth.tasks import DrugRecommendationMIMIC4 +from pyhealth.trainer import Trainer + +if __name__ == "__main__": + # STEP 1: Load MIMIC-IV base dataset + base_dataset = MIMIC4Dataset( + ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", + cache_dir="/shared/eng/pyhealth_agent/baselines", + ehr_tables=[ + "patients", + "admissions", + "diagnoses_icd", + "procedures_icd", + "prescriptions", + ], + ) + + # STEP 2: Apply drug recommendation task + sample_dataset = base_dataset.set_task( + DrugRecommendationMIMIC4(), + num_workers=4, + ) + + print(f"Total samples: {len(sample_dataset)}") + print(f"Input schema: {sample_dataset.input_schema}") + print(f"Output schema: {sample_dataset.output_schema}") + + # Inspect a sample + sample = sample_dataset[0] + print("\nSample structure:") + print(f" Patient ID: {sample['patient_id']}") + print(f" Visit ID: {sample['visit_id']}") + print(f" Conditions (history): {len(sample['conditions'])} visits") + print(f" Procedures (history): {len(sample['procedures'])} visits") + print(f" Drugs history: {len(sample['drugs_hist'])} visits") + print(f" Target drugs: {len(sample['drugs'])} drugs") + print(f"\n First visit conditions: {sample['conditions'][0][:5]}...") + print(f" Target drugs sample: {sample['drugs'][:5]}...") + + # STEP 3: Split dataset + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + print("\nDataset split:") + print(f" Train: {len(train_dataset)} samples") + print(f" Validation: {len(val_dataset)} samples") + print(f" Test: {len(test_dataset)} samples") + + # Create dataloaders + train_loader = get_dataloader(train_dataset, batch_size=64, shuffle=True) + val_loader = get_dataloader(val_dataset, batch_size=64, shuffle=False) + test_loader = get_dataloader(test_dataset, batch_size=64, shuffle=False) + + # STEP 4: Initialize RNN model + model = RNN( + dataset=sample_dataset, + embedding_dim=128, + hidden_dim=128, + rnn_type="GRU", + dropout=0.5, + ) + + num_params = sum(p.numel() for p in model.parameters()) + print(f"\nModel initialized with {num_params:,} parameters") + print(f"Feature keys: {model.feature_keys}") + print(f"Label key: {model.label_key}") + + # STEP 5: Train the model + trainer = Trainer( + model=model, + device="cuda:0", # or "cpu" + metrics=["pr_auc_samples", "f1_samples", "jaccard_samples"], + ) + + print("\nStarting training...") + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=50, + monitor="pr_auc_samples", + optimizer_params={"lr": 1e-4}, + optimizer_class=torch.optim.AdamW, + ) + + # STEP 6: Evaluate on test set + print("\nEvaluating on test set...") + results = trainer.evaluate(test_loader) + print("\nTest Results:") + for metric, value in results.items(): + print(f" {metric}: {value:.4f}") + + # STEP 7: Inspect model predictions + print("\nSample predictions:") + sample_batch = next(iter(test_loader)) + + with torch.no_grad(): + output = model(**sample_batch) + + print(f" Batch size: {output['y_prob'].shape[0]}") + print(f" Number of drug classes: {output['y_prob'].shape[1]}") + print(" Predicted probabilities (first 5 drugs of first patient):") + print(f" {output['y_prob'][0, :5].cpu().numpy()}") + print(" True labels (first 5 drugs of first patient):") + print(f" {output['y_true'][0, :5].cpu().numpy()}") diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_rnn_optuna.py b/examples/drug_recommendation/drug_recommendation_mimic4_rnn_optuna.py new file mode 100644 index 000000000..4c3fb2a9f --- /dev/null +++ b/examples/drug_recommendation/drug_recommendation_mimic4_rnn_optuna.py @@ -0,0 +1,205 @@ +""" +Optuna hyperparameter tuning for RNN on drug recommendation with MIMIC-IV. + +This example demonstrates: +1. Loading MIMIC-IV data and applying the DrugRecommendationMIMIC4 task +2. Defining an Optuna objective that tunes RNN-specific hyperparameters +3. Running 10 Optuna trials to find the best configuration +4. Training a final model with the best hyperparameters + +Tuned hyperparameters: + - embedding_dim: embedding size for code tokens + - hidden_dim: GRU/LSTM/RNN hidden state size + - rnn_type: recurrent cell type (GRU, LSTM, RNN) + - num_layers: number of stacked recurrent layers + - dropout: dropout rate applied before each recurrent layer + - lr: learning rate for AdamW + - weight_decay: L2 regularization coefficient for AdamW +""" + +import torch +import optuna + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, +) +from pyhealth.models import RNN +from pyhealth.tasks import DrugRecommendationMIMIC4 +from pyhealth.trainer import Trainer + +if __name__ == "__main__": + # --------------------------------------------------------------------------- + # STEP 1: Load MIMIC-IV base dataset + # --------------------------------------------------------------------------- + base_dataset = MIMIC4Dataset( + ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", + cache_dir="/shared/eng/pyhealth_agent/baselines", + ehr_tables=[ + "patients", + "admissions", + "diagnoses_icd", + "procedures_icd", + "prescriptions", + ], + ) + + # STEP 2: Apply drug recommendation task + sample_dataset = base_dataset.set_task( + DrugRecommendationMIMIC4(), + num_workers=4, + ) + + print(f"Total samples: {len(sample_dataset)}") + print(f"Input schema: {sample_dataset.input_schema}") + print(f"Output schema: {sample_dataset.output_schema}") + + # STEP 3: Split dataset (fixed split so all trials see the same data) + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + print(f"\nDataset split — Train: {len(train_dataset)} " + f"Val: {len(val_dataset)} Test: {len(test_dataset)}") + + # --------------------------------------------------------------------------- + # STEP 4: Define Optuna objective + # --------------------------------------------------------------------------- + DEVICE = "cuda:2" # or "cpu" + TUNE_EPOCHS = 10 # lightweight training per trial + N_TRIALS = 10 + + def objective(trial: optuna.Trial) -> float: + """Return validation pr_auc_samples for a sampled RNN configuration.""" + + # --- Suggest hyperparameters ------------------------------------------- + embedding_dim = trial.suggest_categorical( + "embedding_dim", [64, 128, 256] + ) + hidden_dim = trial.suggest_categorical("hidden_dim", [64, 128, 256]) + rnn_type = trial.suggest_categorical( + "rnn_type", ["GRU", "LSTM", "RNN"] + ) + num_layers = trial.suggest_int("num_layers", 1, 3) + dropout = trial.suggest_float("dropout", 0.1, 0.7) + lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True) + weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-2, log=True) + batch_size = trial.suggest_categorical("batch_size", [32, 64, 128]) + + # --- Build dataloaders ------------------------------------------------- + train_loader = get_dataloader( + train_dataset, batch_size=batch_size, shuffle=True + ) + val_loader = get_dataloader( + val_dataset, batch_size=batch_size, shuffle=False + ) + + # --- Build model ------------------------------------------------------- + model = RNN( + dataset=sample_dataset, + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + rnn_type=rnn_type, + num_layers=num_layers, + dropout=dropout, + ) + + # --- Train ------------------------------------------------------------- + trainer = Trainer( + model=model, + device=DEVICE, + metrics=["pr_auc_samples"], + ) + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=TUNE_EPOCHS, + monitor="pr_auc_samples", + optimizer_class=torch.optim.AdamW, + optimizer_params={"lr": lr}, + weight_decay=weight_decay, + ) + + # --- Evaluate on validation set ---------------------------------------- + scores = trainer.evaluate(val_loader) + return scores["pr_auc_samples"] + + # --------------------------------------------------------------------------- + # STEP 5: Run Optuna study + # --------------------------------------------------------------------------- + print( + f"\nStarting Optuna search ({N_TRIALS} trials, {TUNE_EPOCHS} epochs each)..." + ) + + study = optuna.create_study(direction="maximize") + study.optimize(objective, n_trials=N_TRIALS) + + best_params = study.best_params + print("\nBest hyperparameters found:") + for k, v in best_params.items(): + print(f" {k}: {v}") + print(f"Best validation pr_auc_samples: {study.best_value:.4f}") + + # --------------------------------------------------------------------------- + # STEP 6: Train final model with best hyperparameters + # --------------------------------------------------------------------------- + print("\nTraining final model with best hyperparameters...") + + train_loader = get_dataloader( + train_dataset, batch_size=best_params["batch_size"], shuffle=True + ) + val_loader = get_dataloader( + val_dataset, batch_size=best_params["batch_size"], shuffle=False + ) + test_loader = get_dataloader( + test_dataset, batch_size=best_params["batch_size"], shuffle=False + ) + + final_model = RNN( + dataset=sample_dataset, + embedding_dim=best_params["embedding_dim"], + hidden_dim=best_params["hidden_dim"], + rnn_type=best_params["rnn_type"], + num_layers=best_params["num_layers"], + dropout=best_params["dropout"], + ) + + num_params = sum(p.numel() for p in final_model.parameters()) + print(f"Final model: {num_params:,} parameters") + + final_trainer = Trainer( + model=final_model, + device=DEVICE, + metrics=["pr_auc_samples", "f1_samples", "jaccard_samples"], + ) + final_trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=50, + monitor="pr_auc_samples", + optimizer_class=torch.optim.AdamW, + optimizer_params={"lr": best_params["lr"]}, + weight_decay=best_params["weight_decay"], + ) + + # STEP 7: Evaluate on test set + print("\nEvaluating on test set...") + results = final_trainer.evaluate(test_loader) + print("\nTest Results:") + for metric, value in results.items(): + print(f" {metric}: {value:.4f}") + + # STEP 8: Inspect model predictions + print("\nSample predictions:") + sample_batch = next(iter(test_loader)) + + with torch.no_grad(): + output = final_model(**sample_batch) + + print(f" Batch size: {output['y_prob'].shape[0]}") + print(f" Number of drug classes: {output['y_prob'].shape[1]}") + print(" Predicted probabilities (first 5 drugs of first patient):") + print(f" {output['y_prob'][0, :5].cpu().numpy()}") + print(" True labels (first 5 drugs of first patient):") + print(f" {output['y_true'][0, :5].cpu().numpy()}") diff --git a/examples/mortality_prediction/mortality_mimic3_grasp.py b/examples/mortality_prediction/mortality_mimic3_grasp.py index 011bd13fe..e335373d6 100644 --- a/examples/mortality_prediction/mortality_mimic3_grasp.py +++ b/examples/mortality_prediction/mortality_mimic3_grasp.py @@ -9,9 +9,6 @@ base_dataset = MIMIC3Dataset( root="/srv/local/data/physionet.org/files/mimiciii/1.4", tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"], - code_mapping={"ICD9CM": "CCSCM", "ICD9PROC": "CCSPROC", "NDC": "ATC"}, - dev=False, - refresh_cache=False, ) base_dataset.stat() diff --git a/examples/mortality_prediction/mortality_mimic4_multimodal_adacare.py b/examples/mortality_prediction/mortality_mimic4_multimodal_adacare.py new file mode 100644 index 000000000..2fc815f0f --- /dev/null +++ b/examples/mortality_prediction/mortality_mimic4_multimodal_adacare.py @@ -0,0 +1,202 @@ +""" +Mortality Prediction on MIMIC-IV with MultimodalAdaCare + +This example demonstrates how to use the MultimodalAdaCare model with mixed +input modalities for in-hospital mortality prediction on MIMIC-IV. + +The MultimodalAdaCare model can handle: +- Sequential features (diagnoses, procedures) → AdaCare processing with scale-adaptive + feature extraction and recalibration +- Non-sequential features (demographics, static measurements) → Direct embedding + +This example shows: +1. Loading MIMIC-IV data with mixed feature types +2. Applying a mortality prediction task +3. Training a MultimodalAdaCare model with both sequential and non-sequential inputs +4. Evaluating the model performance +5. Analyzing feature importance (from AdaCareLayer for sequential features) +""" + +from pyhealth.datasets import MIMIC4Dataset +from pyhealth.datasets import split_by_patient, get_dataloader +from pyhealth.models import MultimodalAdaCare +from pyhealth.tasks import InHospitalMortalityMIMIC4 +from pyhealth.trainer import Trainer + + +if __name__ == "__main__": + # STEP 1: Load MIMIC-IV base dataset + print("=" * 60) + print("STEP 1: Loading MIMIC-IV Dataset") + print("=" * 60) + + base_dataset = MIMIC4Dataset( + ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/", + ehr_tables=["diagnoses_icd", "procedures_icd", "prescriptions", "labevents"], + dev=True, # Use development mode for faster testing + num_workers=4, + ) + base_dataset.stats() + + # STEP 2: Apply mortality prediction task with multimodal features + print("\n" + "=" * 60) + print("STEP 2: Setting Mortality Prediction Task") + print("=" * 60) + + # Use the InHospitalMortalityMIMIC4 task + # This task will create sequential features from diagnoses and procedures + task = InHospitalMortalityMIMIC4() + sample_dataset = base_dataset.set_task( + task, + num_workers=4, + ) + + print(f"\nTotal samples: {len(sample_dataset)}") + print(f"Input schema: {sample_dataset.input_schema}") + print(f"Output schema: {sample_dataset.output_schema}") + + # Inspect a sample + if len(sample_dataset) > 0: + sample = sample_dataset[0] + print("\nSample structure:") + print(f" Patient ID: {sample['patient_id']}") + for key in sample_dataset.input_schema.keys(): + if key in sample: + if isinstance(sample[key], (list, tuple)): + print(f" {key}: length {len(sample[key])}") + else: + print(f" {key}: {type(sample[key])}") + print(f" Mortality: {sample.get('mortality', 'N/A')}") + + # STEP 3: Split dataset + print("\n" + "=" * 60) + print("STEP 3: Splitting Dataset") + print("=" * 60) + + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + print(f"Train samples: {len(train_dataset)}") + print(f"Val samples: {len(val_dataset)}") + print(f"Test samples: {len(test_dataset)}") + + # Create dataloaders + train_loader = get_dataloader(train_dataset, batch_size=64, shuffle=True) + val_loader = get_dataloader(val_dataset, batch_size=64, shuffle=False) + test_loader = get_dataloader(test_dataset, batch_size=64, shuffle=False) + + # STEP 4: Initialize MultimodalAdaCare model + print("\n" + "=" * 60) + print("STEP 4: Initializing MultimodalAdaCare Model") + print("=" * 60) + + model = MultimodalAdaCare( + dataset=sample_dataset, + embedding_dim=128, + hidden_dim=128, + kernel_size=2, + kernel_num=64, + r_v=4, + r_c=4, + activation="sigmoid", + rnn_type="gru", + dropout=0.3, + ) + + num_params = sum(p.numel() for p in model.parameters()) + print(f"Model initialized with {num_params:,} parameters") + + # Print feature classification + print(f"\nSequential features (AdaCare processing): {model.sequential_features}") + print( + f"Non-sequential features (direct embedding): {model.non_sequential_features}" + ) + + # Calculate expected embedding dimensions + seq_dim = len(model.sequential_features) * model.hidden_dim + non_seq_dim = len(model.non_sequential_features) * model.embedding_dim + total_dim = seq_dim + non_seq_dim + print(f"\nPatient representation dimension:") + print(f" Sequential contribution: {seq_dim}") + print(f" Non-sequential contribution: {non_seq_dim}") + print(f" Total: {total_dim}") + + # STEP 5: Train the model + print("\n" + "=" * 60) + print("STEP 5: Training Model") + print("=" * 60) + + trainer = Trainer( + model=model, + device="cuda:0", # Change to "cpu" if no GPU available + metrics=["pr_auc", "roc_auc", "accuracy", "f1"], + ) + + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=10, + monitor="roc_auc", + optimizer_params={"lr": 1e-3}, + ) + + # STEP 6: Evaluate on test set + print("\n" + "=" * 60) + print("STEP 6: Evaluating on Test Set") + print("=" * 60) + + results = trainer.evaluate(test_loader) + print("\nTest Results:") + for metric, value in results.items(): + print(f" {metric}: {value:.4f}") + + # STEP 7: Demonstrate model predictions and feature importance + print("\n" + "=" * 60) + print("STEP 7: Sample Predictions and Feature Importance") + print("=" * 60) + + import torch + + sample_batch = next(iter(test_loader)) + with torch.no_grad(): + output = model(**sample_batch) + + print(f"\nBatch size: {output['y_prob'].shape[0]}") + print(f"First 10 predicted probabilities:") + for i, (prob, true_label) in enumerate( + zip(output["y_prob"][:10], output["y_true"][:10]) + ): + print(f" Sample {i+1}: prob={prob.item():.4f}, true={int(true_label.item())}") + + # Display feature importance information + print(f"\nFeature Importance outputs:") + print( + f" Number of sequential features with importance: {len(output['feature_importance'])}" + ) + print( + f" Number of sequential features with conv importance: {len(output['conv_feature_importance'])}" + ) + + if len(output["feature_importance"]) > 0: + for i, feat_key in enumerate(model.sequential_features): + feat_imp = output["feature_importance"][i] + conv_imp = output["conv_feature_importance"][i] + print(f"\n Feature '{feat_key}':") + print(f" Input feature importance shape: {feat_imp.shape}") + print(f" Conv feature importance shape: {conv_imp.shape}") + + # Summary + print("\n" + "=" * 60) + print("SUMMARY: MultimodalAdaCare Training Complete") + print("=" * 60) + print(f"Model: MultimodalAdaCare") + print(f"Dataset: MIMIC-IV") + print(f"Task: In-Hospital Mortality Prediction") + print(f"Sequential features: {len(model.sequential_features)}") + print(f"Non-sequential features: {len(model.non_sequential_features)}") + print(f"Best validation ROC-AUC: {max(results.get('roc_auc', 0), 0):.4f}") + print("\nAdaCare provides interpretability through:") + print(" - Input feature importance (original features)") + print(" - Convolutional feature importance (scale-adaptive features)") + print("=" * 60) diff --git a/examples/mortality_prediction/multimodal_dataset_stats.py b/examples/mortality_prediction/multimodal_dataset_stats.py new file mode 100644 index 000000000..ceaaf7765 --- /dev/null +++ b/examples/mortality_prediction/multimodal_dataset_stats.py @@ -0,0 +1,391 @@ +"""Standalone multimodal dataset statistics auditor. + +Loads a MIMIC-IV multimodal task dataset and reports per-modality missingness +rates and token/element counts — no model, no GPU, no trainer required. + +Works directly on the processed SampleDataset. Processed schema (single sample, +no batch dim): + discharge_note_times / radiology_note_times: + (input_ids, attn_mask, token_type_ids, time, type_tag) + input_ids shape: (N_notes, 128) attn_mask shape: (N_notes, 128) + attn_mask.sum() gives real (non-padding) tokens seen by the model. + Note: each note is independently truncated to 128 wordpieces — no chunking. + icd_codes: + (time, value) value shape: (N_visits, vocab_size) [multi-hot] + labs_mask: + (time, value) value shape: (N_timesteps, 10) [bool/float] + cxr_image_times: + (image, time, paths) image shape: (N_images, 3, 224, 224) + +Usage: + python examples/mortality_prediction/multimodal_dataset_stats.py \\ + --dev --quick-test + + python examples/mortality_prediction/multimodal_dataset_stats.py \\ + --task ClinicalNotesICDLabsCXRMIMIC4 --output-csv /tmp/stats.csv +""" + +from __future__ import annotations + +import argparse +import csv +from typing import Any, Dict, List, Tuple + +import numpy as np + +from pyhealth.datasets import MIMIC4Dataset +from pyhealth.tasks.multimodal_mimic4 import ( + ClinicalNotesICDLabsCXRMIMIC4, + ClinicalNotesICDLabsMIMIC4, + ClinicalNotesMIMIC4, + ICDLabsMIMIC4, +) + +TASK_MAP = { + "ClinicalNotesICDLabsCXRMIMIC4": ClinicalNotesICDLabsCXRMIMIC4, + "ClinicalNotesICDLabsMIMIC4": ClinicalNotesICDLabsMIMIC4, + "ClinicalNotesMIMIC4": ClinicalNotesMIMIC4, + "ICDLabsMIMIC4": ICDLabsMIMIC4, +} + +# bert-base-uncased encodes "[MISSING_TEXT]" to ~7 tokens with padding. +# Real notes are always longer. Used to detect the missingness sentinel. +_MISSING_NOTE_TOKEN_THRESHOLD = 15 + + +def _note_stats(tup: tuple, note_max_len: int) -> Tuple[bool, int, int]: + """Stats from a processed note tuple. + + Schema: (input_ids, attn_mask, token_type_ids, time, type_tag) + input_ids shape: (N_notes, note_max_len) + attn_mask shape: (N_notes, note_max_len) + + Returns (is_missing, n_notes, seen_tokens_total). + seen_tokens = attn_mask.sum() — real non-padding tokens, already capped + at note_max_len by the processor's truncation. + """ + input_ids, attn_mask = tup[0], tup[1] + n_notes = input_ids.shape[0] + seen = int(attn_mask.sum().item()) + is_missing = (n_notes == 1) and (seen < _MISSING_NOTE_TOKEN_THRESHOLD) + if is_missing: + return True, 0, 0 + return False, n_notes, seen + + +def _icd_stats(tup: tuple) -> Tuple[bool, int, int]: + """Stats from a processed icd_codes tuple. + + Schema: (time, value) value shape: (N_visits, vocab_size) [multi-hot] + + Returns (is_missing, n_visits, n_code_activations). + """ + value = tup[1] + n_visits = value.shape[0] + n_act = int(value.sum().item()) + return n_act == 0, n_visits, n_act + + +def _labs_stats(tup: tuple) -> Tuple[bool, int]: + """Stats from a processed labs_mask tuple. + + Schema: (time, value) value shape: (N_timesteps, 10) [bool/float] + + Returns (is_missing, n_observed). + """ + value = tup[1] + n_obs = int(value.sum().item()) + return n_obs == 0, n_obs + + +def _cxr_stats(tup: tuple, patch_count: int) -> Tuple[bool, int, int]: + """Stats from a processed cxr_image_times tuple. + + Schema: (image, time, paths) image shape: (N_images, 3, H, W) + + Returns (is_missing, n_images, cxr_tokens). + A zero-valued image tensor indicates the missing-image sentinel. + """ + image = tup[0] + n_images = image.shape[0] + if n_images == 0 or float(image.sum()) < 1e-8: + return True, 0, 0 + return False, n_images, n_images * patch_count + + +def audit_sample( + sample: Dict[str, Any], + note_max_len: int, + cxr_patch_count: int, +) -> Dict[str, Any]: + """Audit one processed SampleDataset sample dict.""" + row: Dict[str, Any] = {} + + for note_key in ("discharge_note_times", "radiology_note_times"): + if note_key not in sample: + continue + miss, n_notes, seen_tok = _note_stats( + sample[note_key], note_max_len + ) + row[f"{note_key}__missing"] = int(miss) + row[f"{note_key}__n_notes"] = n_notes + row[f"{note_key}__seen_tokens"] = seen_tok + + if "icd_codes" in sample: + miss, n_visits, n_act = _icd_stats(sample["icd_codes"]) + row["icd_codes__missing"] = int(miss) + row["icd_codes__n_visits"] = n_visits + row["icd_codes__n_activations"] = n_act + + if "labs_mask" in sample: + miss, n_obs = _labs_stats(sample["labs_mask"]) + row["labs__missing"] = int(miss) + row["labs__n_obs"] = n_obs + + if "cxr_image_times" in sample: + miss, n_img, cxr_tok = _cxr_stats( + sample["cxr_image_times"], cxr_patch_count + ) + row["cxr__missing"] = int(miss) + row["cxr__n_images"] = n_img + row["cxr__tokens"] = cxr_tok + + note_tok = sum( + row.get(f"{k}__seen_tokens", 0) + for k in ("discharge_note_times", "radiology_note_times") + ) + row["total_tokens"] = ( + note_tok + + row.get("icd_codes__n_activations", 0) + + row.get("labs__n_obs", 0) + + row.get("cxr__tokens", 0) + ) + return row + + +def _stats(arr: np.ndarray) -> str: + if len(arr) == 0: + return "N/A" + return ( + f"mean={arr.mean():.1f} median={np.median(arr):.0f}" + f" p90={np.percentile(arr, 90):.0f} max={arr.max():.0f}" + ) + + +def print_report(rows: List[Dict], args: argparse.Namespace) -> None: + n = len(rows) + print() + print(f"Task: {args.task} Samples: {n:,} dev={args.dev}") + print() + + def _arr(key: str) -> np.ndarray: + return np.array([r[key] for r in rows if key in r], dtype=float) + + hdr = ( + f"{'Modality':<28} {'missing%':>10}" + f" {'mean':>8} {'median':>8} {'p90':>8} {'max':>8}" + ) + print(hdr) + print("-" * len(hdr)) + + modality_specs = [] + for note_key, label in [ + ("discharge_note_times", "discharge notes"), + ("radiology_note_times", "radiology notes"), + ]: + miss = _arr(f"{note_key}__missing") + if len(miss): + modality_specs.append( + (label, miss, _arr(f"{note_key}__n_notes")) + ) + + if rows and "icd_codes__missing" in rows[0]: + modality_specs.append(( + "icd codes", + _arr("icd_codes__missing"), + _arr("icd_codes__n_activations"), + )) + if rows and "labs__missing" in rows[0]: + modality_specs.append(( + "labs (observations)", + _arr("labs__missing"), + _arr("labs__n_obs"), + )) + if rows and "cxr__missing" in rows[0]: + modality_specs.append(( + "cxr images", + _arr("cxr__missing"), + _arr("cxr__n_images"), + )) + + for label, miss, counts in modality_specs: + miss_pct = f"{miss.mean() * 100:.1f}%" + print( + f"{label:<28} {miss_pct:>10} " + f"{counts.mean():>8.1f} {np.median(counts):>8.0f} " + f"{np.percentile(counts, 90):>8.0f} {counts.max():>8.0f}" + ) + + print() + print(f"Note seen-tokens (cap@{args.note_max_len}, from attn_mask):") + for note_key, label in [ + ("discharge_note_times", " discharge"), + ("radiology_note_times", " radiology"), + ]: + seen = _arr(f"{note_key}__seen_tokens") + if len(seen) == 0: + continue + pct_full = (seen >= args.note_max_len).mean() * 100 + print( + f"{label}: {_stats(seen)}" + f" % hitting cap={pct_full:.0f}%" + ) + + if rows and "cxr__tokens" in rows[0]: + cxr_tok = _arr("cxr__tokens") + print() + print(f"CXR tokens (@{args.cxr_patch_count} patches/image):") + print(f" {_stats(cxr_tok)}") + + total = _arr("total_tokens") + print() + print("Aggregate tokens/sample:") + print(f" {_stats(total)}") + print() + + +def run(args: argparse.Namespace) -> None: + task_cls = TASK_MAP[args.task] + needs_notes = args.task in ( + "ClinicalNotesMIMIC4", + "ClinicalNotesICDLabsMIMIC4", + "ClinicalNotesICDLabsCXRMIMIC4", + ) + needs_cxr = args.task == "ClinicalNotesICDLabsCXRMIMIC4" + needs_icd = args.task in ( + "ClinicalNotesICDLabsMIMIC4", + "ICDLabsMIMIC4", + "ClinicalNotesICDLabsCXRMIMIC4", + ) + + ehr_tables = ( + ["diagnoses_icd", "procedures_icd", "labevents"] + if needs_icd + else [] + ) + note_tables = ["discharge", "radiology"] if needs_notes else [] + cxr_tables = ( + ["metadata", "negbio", "chexpert", "split"] if needs_cxr else [] + ) + + print("Loading MIMIC4Dataset ...") + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root if needs_notes else None, + cxr_root=args.cxr_root if needs_cxr else None, + cxr_variant=args.cxr_variant, + ehr_tables=ehr_tables, + note_tables=note_tables, + cxr_tables=cxr_tables, + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = task_cls(window_hours=args.observation_window_hours) + print("Running set_task ...") + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + total = len(sample_dataset) + print(f"Total samples: {total:,}") + + limit = ( + min(args.sample_limit, total) + if args.sample_limit and args.sample_limit > 0 + else total + ) + print(f"Auditing {limit:,} samples ...") + + rows: List[Dict] = [] + for i in range(limit): + if i % 5000 == 0 and i > 0: + print(f" {i:,} / {limit:,} ...") + rows.append( + audit_sample( + sample_dataset[i], + args.note_max_len, + args.cxr_patch_count, + ) + ) + + if not rows: + print("No samples. Check roots/tables/task combination.") + return + + print_report(rows, args) + + if args.output_csv: + all_keys = list(rows[0].keys()) + with open(args.output_csv, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=all_keys) + writer.writeheader() + writer.writerows(rows) + print(f"Per-sample CSV written to: {args.output_csv}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Audit modality missingness and token counts " + "for MIMIC-IV multimodal tasks." + ) + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", type=str, default="/shared/eng/pyhealth" + ) + parser.add_argument( + "--task", + type=str, + default="ClinicalNotesICDLabsCXRMIMIC4", + choices=list(TASK_MAP.keys()), + ) + parser.add_argument("--observation-window-hours", type=int, default=24) + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--note-max-len", type=int, default=128) + parser.add_argument("--cxr-patch-count", type=int, default=196) + parser.add_argument("--sample-limit", type=int, default=None) + parser.add_argument("--output-csv", type=str, default=None) + + args = parser.parse_args() + if args.quick_test: + args.dev = True + if args.sample_limit is None: + args.sample_limit = 50 + return args + + +if __name__ == "__main__": + args = parse_args() + run(args) diff --git a/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py new file mode 100644 index 000000000..e807848af --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py @@ -0,0 +1,360 @@ +"""Unified multimodal embedding + BottleneckTransformer runner. + +Runs EHR + notes + X-ray (MIMIC-IV + CXR) with ClinicalNotesICDLabsCXRMIMIC4. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2 +- note_root: /shared/rsaas/physionet.org/files/mimic-note +- cxr_root: /shared/rsaas/physionet.org/files/MIMIC-CXR + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_bottleneck_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_bottleneck_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import UnifiedMultimodalEmbeddingModel +from pyhealth.models.bottleneck_transformer import BottleneckTransformer +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + train_ds, val_ds, test_ds = split_by_patient( + dataset, [0.8, 0.1, 0.1], seed=seed + ) + if len(train_ds) == 0 or len(test_ds) == 0: + train_ds, val_ds, test_ds = split_by_sample( + dataset, [0.8, 0.1, 0.1], seed=seed + ) + return train_ds, val_ds, test_ds + + +def _build_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_layers{args.num_layers}" + f"_heads{args.heads}" + f"_bottleneck{args.bottlenecks_n}" + f"_fuse{args.fusion_startidx}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), + "output", + "multimodal_embedding_bottleneck_mimic4_cxr", + run_tag, + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError( + "Task produced zero samples. Check roots/tables or use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print( + f" - {key}: {type(processor).__name__}, " + f"schema={processor.schema()}" + ) + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = BottleneckTransformer( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + bottlenecks_n=args.bottlenecks_n, + fusion_startidx=args.fusion_startidx, + num_layers=args.num_layers, + heads=args.heads, + dropout=args.dropout, + unified_embedding=unified, + ) + print( + f"BottleneckTransformer bottlenecks_n={args.bottlenecks_n}, " + f"fusion_startidx={args.fusion_startidx}" + ) + + 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) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + if len(test_ds) > 0 + else None + ) + + print( + "Split sizes: " + f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print( + f" tuple[{i}] type={type(elem).__name__} " + f"shape={shape}" + ) + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Run unified multimodal embedding + BottleneckTransformer " + "on MIMIC-IV mortality with CXR." + ) + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--num-layers", type=int, default=2) + parser.add_argument("--heads", type=int, default=1) + parser.add_argument("--dropout", type=float, default=0.1) + parser.add_argument("--bottlenecks-n", type=int, default=4) + parser.add_argument("--fusion-startidx", type=int, default=1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print( + f"Inference completed (patient_ids={num_patient_ids}, " + f"rows={num_rows})." + ) diff --git a/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py new file mode 100644 index 000000000..4746eb342 --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py @@ -0,0 +1,339 @@ +"""Unified multimodal embedding + JambaEHR runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +JambaEHR is a hybrid Transformer + Mamba architecture that interleaves +attention and SSM (state-space model) blocks. + +Default roots are set to shared PhysioNet mounts: +- ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2 +- note_root: /shared/rsaas/physionet.org/files/mimic-note +- cxr_root: /shared/rsaas/physionet.org/files/MIMIC-CXR + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_jamba_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_jamba_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import JambaEHR, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) + if len(train_ds) == 0 or len(test_ds) == 0: + train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) + return train_ds, val_ds, test_ds + + +def _build_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_attn{args.num_transformer_layers}" + f"_mamba{args.num_mamba_layers}" + f"_heads{args.heads}" + f"_state{args.state_size}" + f"_conv{args.conv_kernel}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_jamba_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError( + "Task produced zero samples. Check roots/tables or use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = JambaEHR( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + num_transformer_layers=args.num_transformer_layers, + num_mamba_layers=args.num_mamba_layers, + heads=args.heads, + dropout=args.dropout, + state_size=args.state_size, + conv_kernel=args.conv_kernel, + unified_embedding=unified, + ) + print(f"JambaEHR unified mode: {model._use_unified}") + + 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) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + if len(test_ds) > 0 + else None + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + JambaEHR on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--num-transformer-layers", type=int, default=2) + parser.add_argument("--num-mamba-layers", type=int, default=6) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--state-size", type=int, default=16) + parser.add_argument("--conv-kernel", type=int, default=4) + parser.add_argument("--dropout", type=float, default=0.3) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py new file mode 100644 index 000000000..4896b037d --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py @@ -0,0 +1,330 @@ +"""Unified multimodal embedding + EHRMamba runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2 +- note_root: /shared/rsaas/physionet.org/files/mimic-note +- cxr_root: /shared/rsaas/physionet.org/files/MIMIC-CXR + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_mamba_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_mamba_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import EHRMamba, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) + if len(train_ds) == 0 or len(test_ds) == 0: + train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) + return train_ds, val_ds, test_ds + + +def _build_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_layers{args.num_layers}" + f"_state{args.state_size}" + f"_conv{args.conv_kernel}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_mamba_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError( + "Task produced zero samples. Check roots/tables or use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = EHRMamba( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + num_layers=args.num_layers, + state_size=args.state_size, + conv_kernel=args.conv_kernel, + dropout=args.dropout, + unified_embedding=unified, + ) + print(f"EHRMamba unified mode: {model._use_unified}") + + 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) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + if len(test_ds) > 0 + else None + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + EHRMamba on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--num-layers", type=int, default=2) + parser.add_argument("--state-size", type=int, default=16) + parser.add_argument("--conv-kernel", type=int, default=4) + parser.add_argument("--dropout", type=float, default=0.1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py new file mode 100644 index 000000000..ac09f2c61 --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py @@ -0,0 +1,334 @@ +"""Unified multimodal embedding + MLP runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2 +- note_root: /shared/rsaas/physionet.org/files/mimic-note +- cxr_root: /shared/rsaas/physionet.org/files/MIMIC-CXR + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_mlp_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_mlp_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import MLP, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) + if len(train_ds) == 0 or len(test_ds) == 0: + train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) + return train_ds, val_ds, test_ds + + +def _build_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_hidden{args.hidden_dim}" + f"_layers{args.num_layers}" + f"_act{args.activation}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_mlp_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError( + "Task produced zero samples. Check roots/tables or use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = MLP( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + hidden_dim=args.hidden_dim, + n_layers=args.num_layers, + activation=args.activation, + unified_embedding=unified, + ) + print(f"MLP unified mode: {model._use_unified}") + + 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) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + if len(test_ds) > 0 + else None + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + MLP on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--hidden-dim", type=int, default=128) + parser.add_argument("--num-layers", type=int, default=2) + parser.add_argument( + "--activation", + type=str, + default="relu", + choices=["relu", "tanh", "sigmoid", "leaky_relu", "elu"], + ) + parser.add_argument("--dropout", type=float, default=0.1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py new file mode 100644 index 000000000..4aeeed9ce --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py @@ -0,0 +1,335 @@ +"""Unified multimodal embedding + RNN runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2 +- note_root: /shared/rsaas/physionet.org/files/mimic-note +- cxr_root: /shared/rsaas/physionet.org/files/MIMIC-CXR + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_rnn_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_rnn_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import RNN, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) + if len(train_ds) == 0 or len(test_ds) == 0: + train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) + return train_ds, val_ds, test_ds + + +def _build_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_hidden{args.hidden_dim}" + f"_{args.rnn_type.lower()}" + f"_layers{args.num_layers}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_rnn_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError( + "Task produced zero samples. Check roots/tables or use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = RNN( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + hidden_dim=args.hidden_dim, + unified_embedding=unified, + rnn_type=args.rnn_type, + num_layers=args.num_layers, + dropout=args.dropout, + ) + print(f"RNN unified mode: {model._use_unified}") + + 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) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + if len(test_ds) > 0 + else None + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + RNN on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--hidden-dim", type=int, default=128) + parser.add_argument( + "--rnn-type", + type=str, + default="GRU", + choices=["GRU", "LSTM", "RNN"], + ) + parser.add_argument("--num-layers", type=int, default=1) + parser.add_argument("--dropout", type=float, default=0.1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py new file mode 100644 index 000000000..69942738c --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py @@ -0,0 +1,327 @@ +"""Unified multimodal embedding + Transformer runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2 +- note_root: /shared/rsaas/physionet.org/files/mimic-note +- cxr_root: /shared/rsaas/physionet.org/files/MIMIC-CXR + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_transformer_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_transformer_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import Transformer, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) + if len(train_ds) == 0 or len(test_ds) == 0: + train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) + return train_ds, val_ds, test_ds + + +def _build_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_layers{args.num_layers}" + f"_heads{args.heads}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_transformer_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError( + "Task produced zero samples. Check roots/tables or use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = Transformer( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + heads=args.heads, + dropout=args.dropout, + num_layers=args.num_layers, + unified_embedding=unified, + ) + print(f"Transformer unified mode: {model._use_unified}") + + 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) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + if len(test_ds) > 0 + else None + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + Transformer on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--num-layers", type=int, default=1) + parser.add_argument("--heads", type=int, default=1) + parser.add_argument("--dropout", type=float, default=0.1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/multimodal_mimic4.py b/examples/mortality_prediction/multimodal_mimic4.py index 631e44edb..4261c25b4 100644 --- a/examples/mortality_prediction/multimodal_mimic4.py +++ b/examples/mortality_prediction/multimodal_mimic4.py @@ -4,56 +4,149 @@ # PyHealth Packages from pyhealth.datasets import MIMIC4Dataset -from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4 +from pyhealth.tasks.multimodal_mimic4 import ( + ClinicalNotesMIMIC4, + ICDLabsMIMIC4, + ClinicalNotesICDLabsMIMIC4, + ClinicalNotesICDLabsCXRMIMIC4, +) from pyhealth.tasks.base_task import BaseTask # Load MIMIC4 Files -# There's probably better ways dealing with this on the cluster, but working locally for now +# There's probably better ways dealing with this on the cluster, but working locally for now # (see: https://github.com/sunlabuiuc/PyHealth/blob/master/examples/mortality_prediction/multimodal_mimic4_minimal.py) -TASK = "ClinicalNotesICDLabsMIMIC4" # The idea here is that we want additive tasks so we can evaluate the value in adding more modalities +TASK = "ClinicalNotesICDLabsMIMIC4" # Options: ClinicalNotesMIMIC4, ICDLabsMIMIC4, ClinicalNotesICDLabsMIMIC4, ClinicalNotesICDLabsCXRMIMIC4 # The idea here is that we want additive tasks so we can evaluate the value in adding more modalities +DEV_MODE = True +ENVIRONMENT = "CampusCluster" # Either 'Local' or 'CampusCluster' or "SunLabCluster" +NETID = "wp14" # For personal cache -PYHEALTH_REPO_ROOT = '/Users/wpang/Desktop/PyHealth' +if ENVIRONMENT == "Local": + pyhealth_repo_root = "/Users/wpang/Desktop/PyHealth" + + ehr_root = os.path.join( + pyhealth_repo_root, "local_data/local/data/physionet.org/files/mimiciv/2.2" + ) + note_root = os.path.join( + pyhealth_repo_root, + "local_data/local/data/physionet.org/files/mimic-iv-note/2.2", + ) + cxr_root = os.path.join( + pyhealth_repo_root, + "llocal_data/local/data/physionet.org/files/mimic-cxr-jpg/2.0.0", + ) + cache_dir = os.path.join( + pyhealth_repo_root, "local_data/local/data/wp/pyhealth_cache" + ) +elif ENVIRONMENT == "CampusCluster": + + 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" + cxr_root = None # Please fill this in + cache_dir = f"/u/{NETID}/pyhealth_cache" +elif ENVIRONMENT == "SunLabCluster": + + ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2" + note_root = "/shared/rsaas/physionet.org/files/mimic-note" + cxr_root = None # Please fill this in + cache_dir = f"/home/{NETID}/pyhealth_cache" -EHR_ROOT = os.path.join(PYHEALTH_REPO_ROOT, "local_data/local/data/physionet.org/files/mimiciv/2.2") -NOTE_ROOT = os.path.join(PYHEALTH_REPO_ROOT, "local_data/local/data/physionet.org/files/mimic-iv-note/2.2") -CXR_ROOT = os.path.join(PYHEALTH_REPO_ROOT,"local_data/local/data/physionet.org/files/mimic-cxr-jpg/2.0.0") -CACHE_DIR = os.path.join(PYHEALTH_REPO_ROOT,"local_data/local/data/wp/pyhealth_cache") if __name__ == "__main__": - if TASK == "ClinicalNotesMIMIC4": # A bit janky setup at the moment and open to iteration, but conveys the point for now + if ( + TASK == "ClinicalNotesMIMIC4" + ): # A bit janky setup at the moment and open to iteration, but conveys the point for now + dataset = MIMIC4Dataset( + ehr_root=ehr_root, + note_root=note_root, + ehr_tables=[ + "diagnoses_icd", + "procedures_icd", + "prescriptions", + "labevents", + ], + note_tables=["discharge", "radiology"], + cache_dir=cache_dir, + num_workers=8, + dev=DEV_MODE, + ) + + # Apply multimodal task + task = ClinicalNotesMIMIC4() + samples = dataset.set_task(task) + + # Get and print sample + sample = samples[0] + print(sample) + + elif TASK == "ICDLabsMIMIC4": dataset = MIMIC4Dataset( - ehr_root=EHR_ROOT, - note_root=NOTE_ROOT, - ehr_tables=["diagnoses_icd", "procedures_icd", "prescriptions", "labevents"], - note_tables=["discharge", "radiology"], - cache_dir=CACHE_DIR, - num_workers=8, - dev=True - ) - + ehr_root=ehr_root, + ehr_tables=[ + "diagnoses_icd", + "procedures_icd", + "labevents", + "prescriptions", + ], + cache_dir=cache_dir, + num_workers=8, + dev=DEV_MODE, + ) + + # Apply multimodal task + task = ICDLabsMIMIC4() + samples = dataset.set_task(task) + + # Get and print sample + sample = samples[0] + print(sample) + + elif TASK == "ClinicalNotesICDLabsMIMIC4": + dataset = MIMIC4Dataset( + ehr_root=ehr_root, + note_root=note_root, + ehr_tables=[ + "diagnoses_icd", + "procedures_icd", + "prescriptions", + "labevents", + ], + note_tables=["discharge", "radiology"], + cache_dir=cache_dir, + num_workers=8, + dev=DEV_MODE, + ) + # Apply multimodal task - task = ClinicalNotesMIMIC4() + task = ClinicalNotesICDLabsMIMIC4() samples = dataset.set_task(task) # Get and print sample sample = samples[0] print(sample) - - elif TASK == 'ClinicalNotesICDLabsMIMIC4': + + elif TASK == "ClinicalNotesICDLabsCXRMIMIC4": dataset = MIMIC4Dataset( - ehr_root=EHR_ROOT, - note_root=NOTE_ROOT, - ehr_tables=["diagnoses_icd", "procedures_icd", "prescriptions", "labevents"], - note_tables=["discharge", "radiology"], - cache_dir=CACHE_DIR, - num_workers=8, - dev=True - ) - + ehr_root=ehr_root, + note_root=note_root, + cxr_root=cxr_root, + cxr_variant="sunlab", + ehr_tables=[ + "diagnoses_icd", + "procedures_icd", + "prescriptions", + "labevents", + ], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio"], + cache_dir=cache_dir, + num_workers=8, + dev=DEV_MODE, + ) + # Apply multimodal task - task = ClinicalNotesICDLabsMIMIC4() + task = ClinicalNotesICDLabsCXRMIMIC4() samples = dataset.set_task(task) # Get and print sample diff --git a/examples/mortality_prediction/readme.md b/examples/mortality_prediction/readme.md new file mode 100644 index 000000000..fb82f291b --- /dev/null +++ b/examples/mortality_prediction/readme.md @@ -0,0 +1,18 @@ +# Readme file for mortality prediction run examples here. + + + + +# Multimodality + +## Model Variants + +nohup python examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_mamba_mimic4_cxr_b1.log & + +nohup python examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_mlp_mimic4_cxr_b1.log & + +nohup python examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_rnn_mimic4_cxr_b1.log & + +nohup python examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_transformer_mimic4_cxr_b1.log & + +nohup python examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_jamba_mimic4_cxr_b1.log & \ No newline at end of file diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 0a2dfce91..a572dd1b6 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -10,6 +10,10 @@ MortalityPredictionStageNetMIMIC4: ICD codes + 10-dim lab vectors, patient-level samples aggregated across all admissions. +--task icd_labs + ICDLabsMIMIC4: ICD codes + 10-dim lab vectors via the unified + multimodal pipeline. No notes required. + --task clinical_notes_icd_labs ClinicalNotesICDLabsMIMIC4: discharge/radiology notes + ICD + labs. Requires --note-root. Legacy; ICD codes are discharge-coded (leakage). @@ -72,6 +76,7 @@ from pyhealth.tasks.multimodal_mimic4 import ( ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4, + LabsMIMIC4, NotesLabsMIMIC4, ) from pyhealth.trainer import Trainer @@ -104,6 +109,9 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: 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, @@ -128,6 +136,8 @@ def _build_task(args: argparse.Namespace): 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}") @@ -310,7 +320,7 @@ def run(args: argparse.Namespace) -> Path: trainer = Trainer( model=model, - metrics=["pr_auc", "roc_auc", "f1", "f1_opt", "accuracy"], + metrics=["pr_auc", "roc_auc", "f1", "accuracy"], device=args.device, enable_logging=True, output_path=str(output_dir), @@ -379,7 +389,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--task", type=str, - choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_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. " diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py new file mode 100644 index 000000000..a4e0f7277 --- /dev/null +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py @@ -0,0 +1,346 @@ +"""End-to-end protocol runner for Unified Embedding on MIMIC-IV + CXR. + +Trains and evaluates a unified-embedding model (MLP / RNN / Transformer / +BottleneckTransformer / EHRMamba / JambaEHR) on MIMIC-IV mortality using +clinical notes, ICD codes, lab values, and chest X-rays. + +Example +------- + python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ + --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \\ + --note-root /shared/rsaas/physionet.org/files/mimic-note \\ + --cxr-root /shared/rsaas/physionet.org/files/MIMIC-CXR \\ + --task clinical_notes_icd_labs_cxr \\ + --model ehrmamba --num-layers 2 \\ + --epochs 10 --batch-size 8 --device cuda:0 +""" + +from __future__ import annotations + +import argparse +import csv +from pathlib import Path +from typing import Any, Optional, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +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, + ClinicalNotesICDLabsCXRMIMIC4, +) +from pyhealth.trainer import Trainer + + +def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: + ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] + note_tables = None + + # Notes required for any task that includes clinical text + if args.task in ("clinical_notes_icd_labs", "clinical_notes_icd_labs_cxr"): + if not args.note_root: + raise ValueError(f"--task {args.task} requires --note-root.") + note_tables = ["discharge", "radiology"] + + # CXR metadata tables only loaded for the CXR task + cxr_kwargs = {} + if args.task == "clinical_notes_icd_labs_cxr": + if not args.cxr_root: + raise ValueError("--task clinical_notes_icd_labs_cxr requires --cxr-root.") + cxr_kwargs = dict( + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + cxr_tables=["metadata", "negbio", "chexpert", "split"], + ) + + 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, + num_workers=args.num_workers, + **cxr_kwargs, + ) + + +def _build_task(args: argparse.Namespace): + if args.task == "stagenet": + return MortalityPredictionStageNetMIMIC4() + if args.task == "clinical_notes_icd_labs": + return ClinicalNotesICDLabsMIMIC4(window_hours=args.observation_window_hours) + if args.task == "clinical_notes_icd_labs_cxr": + return ClinicalNotesICDLabsCXRMIMIC4(window_hours=args.observation_window_hours) + raise ValueError(f"Unknown task: {args.task}") + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) + if len(train_ds) == 0 or len(test_ds) == 0: + train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) + return train_ds, val_ds, test_ds + + +def _build_model(args: argparse.Namespace, sample_dataset: Any): + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + + if args.model == "mlp": + return MLP( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + hidden_dim=args.hidden_dim, + unified_embedding=unified, + ) + if args.model == "rnn": + return RNN( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + hidden_dim=args.hidden_dim, + unified_embedding=unified, + rnn_type=args.rnn_type, + num_layers=args.rnn_layers, + dropout=args.dropout, + bidirectional=args.bidirectional, + ) + if args.model == "transformer": + return Transformer( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + heads=args.heads, + num_layers=args.num_layers, + dropout=args.dropout, + unified_embedding=unified, + ) + if args.model == "bottleneck_transformer": + return BottleneckTransformer( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + bottlenecks_n=args.bottlenecks_n, + fusion_startidx=args.fusion_startidx, + num_layers=args.num_layers, + heads=args.heads, + dropout=args.dropout, + unified_embedding=unified, + ) + if args.model == "ehrmamba": + return EHRMamba( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + num_layers=args.num_layers, + state_size=args.mamba_state_size, + conv_kernel=args.mamba_conv_kernel, + dropout=args.dropout, + unified_embedding=unified, + ) + if args.model == "jambaehr": + return JambaEHR( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + num_transformer_layers=args.jamba_transformer_layers, + num_mamba_layers=args.jamba_mamba_layers, + heads=args.heads, + dropout=args.dropout, + state_size=args.mamba_state_size, + conv_kernel=args.mamba_conv_kernel, + unified_embedding=unified, + ) + raise ValueError(f"Unknown model: {args.model}") + + +def _write_predictions( + output_csv: Path, + patient_ids: list[str], + y_true: np.ndarray, + y_prob: np.ndarray, +) -> None: + output_csv.parent.mkdir(parents=True, exist_ok=True) + + y_true_flat = y_true.reshape(-1).tolist() + y_prob_flat = y_prob.reshape(-1).tolist() + + with output_csv.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter( + f, + fieldnames=["patient_id", "y_true", "y_prob", "y_pred_threshold_0_5"], + ) + writer.writeheader() + for idx, prob in enumerate(y_prob_flat): + writer.writerow( + { + "patient_id": patient_ids[idx], + "y_true": int(y_true_flat[idx]), + "y_prob": float(prob), + "y_pred_threshold_0_5": int(float(prob) >= 0.5), + } + ) + + +def run(args: argparse.Namespace) -> Path: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + + base_dataset = _build_base_dataset(args) + task = _build_task(args) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError( + "Task produced zero samples. Check roots/tables or adjust settings." + ) + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + model = _build_model(args, sample_dataset) + + 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) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + if len(test_ds) > 0 + else None + ) + + exp_name = f"{args.model}_seed{args.seed}" + output_dir = Path(args.output_dir) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc", "f1", "accuracy"], + device=args.device, + enable_logging=True, + output_path=str(output_dir), + exp_name=exp_name, + ) + + # BT uses tighter optimizer settings to stabilize training on full MIMIC-IV + effective_lr = args.lr + effective_max_grad_norm = args.max_grad_norm + optimizer_params = {} + + if args.model == "bottleneck_transformer": + if effective_lr is None: + effective_lr = 1e-4 + if effective_max_grad_norm is None: + effective_max_grad_norm = 0.5 + optimizer_params["eps"] = args.adam_eps if args.adam_eps is not None else 1e-6 + else: + if effective_lr is None: + effective_lr = 1e-3 + if args.adam_eps is not None: + optimizer_params["eps"] = args.adam_eps + + optimizer_params["lr"] = effective_lr + + if args.epochs > 0 and len(train_ds) > 0: + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params=optimizer_params, + weight_decay=args.weight_decay, + max_grad_norm=effective_max_grad_norm, + monitor="pr_auc", + load_best_model_at_last=True, + ) + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + + output_csv = output_dir / exp_name / f"predictions_{args.model}.csv" + _write_predictions(output_csv, patient_ids, y_true, y_prob) + return output_csv + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run E2E unified embedding on MIMIC-IV + CXR with any of six backbone models." + ) + parser.add_argument("--ehr-root", type=str, required=True) + parser.add_argument("--note-root", type=str, default=None) + parser.add_argument("--cxr-root", type=str, default=None) + parser.add_argument("--cxr-variant", type=str, default="sunlab", choices=["default", "sunlab"]) + parser.add_argument("--cache-dir", type=str, default=None) + parser.add_argument("--output-dir", type=str, default="./output/unified_e2e_cxr") + + parser.add_argument( + "--task", type=str, default="clinical_notes_icd_labs_cxr", + choices=["stagenet", "clinical_notes_icd_labs", "clinical_notes_icd_labs_cxr"], + ) + parser.add_argument( + "--model", type=str, default="rnn", + choices=["mlp", "rnn", "transformer", "bottleneck_transformer", "ehrmamba", "jambaehr"], + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--hidden-dim", type=int, default=64) + parser.add_argument("--dropout", type=float, default=0.1) + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--lr", type=float, default=None) + parser.add_argument("--adam-eps", type=float, default=None) + parser.add_argument("--weight-decay", type=float, default=0.0) + parser.add_argument("--device", type=str, default=None) + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--dev", action="store_true") + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--rnn-type", type=str, default="GRU") + parser.add_argument("--rnn-layers", type=int, default=1) + parser.add_argument("--bidirectional", action="store_true") + + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--num-layers", type=int, default=2) + + parser.add_argument("--bottlenecks-n", type=int, default=4) + parser.add_argument("--fusion-startidx", type=int, default=1) + + parser.add_argument("--max-grad-norm", type=float, default=None) + + parser.add_argument("--mamba-state-size", type=int, default=16) + parser.add_argument("--mamba-conv-kernel", type=int, default=4) + parser.add_argument("--jamba-transformer-layers", type=int, default=2) + parser.add_argument("--jamba-mamba-layers", type=int, default=6) + + return parser.parse_args() + + +def print_vram_summary(device: Optional[str] = None) -> None: + """Print peak allocated and total VRAM for the given CUDA device.""" + if not torch.cuda.is_available(): + print("VRAM summary: CUDA not available.") + return + idx = torch.device(device).index if device and device.startswith("cuda") else 0 + if idx is None: + idx = 0 + allocated_mb = torch.cuda.max_memory_allocated(idx) / (1024 ** 2) + total_mb = torch.cuda.get_device_properties(idx).total_memory / (1024 ** 2) + print(f"VRAM summary (cuda:{idx}): peak_allocated={allocated_mb:.0f} MB / total={total_mb:.0f} MB") + + +if __name__ == "__main__": + args = parse_args() + output_csv_path = run(args) + print(f"Saved predictions to: {output_csv_path}") + print_vram_summary(args.device) diff --git a/examples/multimodal_mimic4_task_tutorial.ipynb b/examples/multimodal_mimic4_task_tutorial.ipynb index 3fda88002..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,11 +33,16 @@ "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", - "import torch" + "import torch\n", + "import logging" ] }, { @@ -50,13 +55,31 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "201c3a62-2a2e-47ab-9ab8-eea77812f349", "metadata": {}, "outputs": [], "source": [ "from pyhealth.datasets import MIMIC4Dataset\n", - "from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsMIMIC4" + "from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4, ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4, BaseMultimodalMIMIC4Task" + ] + }, + { + "cell_type": "markdown", + "id": "8f92ea87-c73d-4fbe-bae1-649f18922c5f", + "metadata": {}, + "source": [ + "*Disable Logging*" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f9240e90-af5e-4bf3-96e6-821865967db8", + "metadata": {}, + "outputs": [], + "source": [ + "logging.disable(logging.CRITICAL)" ] }, { @@ -69,17 +92,50 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "9a466c4b-2283-443f-94a3-5768babddf50", "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ - "SEED = 42\n", +<<<<<<< HEAD + "SEED = 30\n", +======= + "SEED = 22\n", +>>>>>>> origin/main "random.seed(SEED)\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)" ] }, + { + "cell_type": "markdown", + "id": "a5f48cc8-28a7-41c8-a90d-d8521f337093", + "metadata": {}, + "source": [ + "*Mortality Label Filter `(0, 1, or [0, 1])`*" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "518d0bd0-d320-4173-8f92-b4e1abbd9a33", + "metadata": {}, + "outputs": [], + "source": [ + "MORTALITY_LABEL = 0" + ] + }, { "cell_type": "markdown", "id": "e525de6c-0798-4963-b2ed-0e728037a0e9", @@ -90,11 +146,24 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "05ec2dc9-1a0c-4360-a028-d93173dcba65", "metadata": {}, "outputs": [], "source": [ + "def filter_mortality_samples(samples, label=None):\n", + " \"\"\"Return samples filtered by mortality label.\n", + "\n", + " Args:\n", + " samples: list of sample dicts with a 'mortality' tensor field.\n", + " label: 0, 1, [0, 1], or None to return all samples.\n", + " \"\"\"\n", + " if label is None or label == [0, 1] or label == [1, 0]:\n", + " return list(samples)\n", + " labels = {label} if isinstance(label, int) else set(label)\n", + " return [s for s in samples if s['mortality'].item() in labels]\n", + "\n", + "\n", "def print_patient_admission_info(sample):\n", " patient = dataset.get_patient(sample['patient_id'])\n", " admissions = patient.get_events(\n", @@ -108,6 +177,7 @@ " print(f\" admittime: {adm.timestamp}\")\n", " print(f\" dischtime: {adm.dischtime}\")\n", "\n", + "\n", "def print_patient_note_info(sample, note_type=\"discharge\", char_limit=80):\n", " patient = dataset.get_patient(sample['patient_id'])\n", " notes = patient.get_events(\n", @@ -121,7 +191,88 @@ " print(f\" hadm_id: {note.hadm_id}\")\n", " print(f\" charttime: {note.timestamp}\")\n", " print(f\" storetime: {note.storetime}\")\n", - " print(f\" text: {note.text[:char_limit]}..(Limited to {char_limit} Characters).\")" + " 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", + " for name, itemids in BaseMultimodalMIMIC4Task.LAB_CATEGORIES.items()\n", + " for itemid in itemids\n", + "}\n", + "\n", + "\n", + "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", + " print(f\" itemid: {lab['itemid']} ({lab_name})\")\n", + " print(f\" charttime: {lab.timestamp}\")\n", + " print(f\" storetime: {lab['storetime']}\")\n", + " print(f\" valuenum: {lab['valuenum']}\")\n", + "\n", + "\n", + "def print_patient_icd_info(sample):\n", + " patient = dataset.get_patient(sample['patient_id'])\n", + " diagnoses = patient.get_events(\n", + " event_type=\"diagnoses_icd\",\n", + " start=sample['window_start'],\n", + " end=sample['window_end'],\n", + " )\n", + " procedures = patient.get_events(\n", + " event_type=\"procedures_icd\",\n", + " start=sample['window_start'],\n", + " end=sample['window_end'],\n", + " )\n", + " print(f\"\\ndiagnoses_icd: {len(diagnoses)}\")\n", + " for dx in diagnoses:\n", + " print(f\" hadm_id: {dx.hadm_id}\")\n", + " print(f\" seq_num: {dx.seq_num}\")\n", + " print(f\" icd_code: {dx.icd_code} (ICD-{dx.icd_version})\\n\")\n", + " print(f\"\\nprocedures_icd: {len(procedures)}\")\n", + " for px in procedures:\n", + " print(f\" hadm_id: {px.hadm_id}\")\n", + " print(f\" seq_num: {px.seq_num}\")\n", + " print(f\" icd_code: {px.icd_code} (ICD-{px.icd_version})\\n\")\n", + "\n", + "\n", + "def clear_cache_directory(path):\n", + " for item in os.listdir(path):\n", + " item_path = os.path.join(path, item)\n", + " if os.path.isdir(item_path):\n", + " shutil.rmtree(item_path)\n", + " else:\n", + " os.remove(item_path)\n", + " print(f\"Cache directory cleared: {path}\")" +>>>>>>> origin/main ] }, { @@ -139,60 +290,216 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "tv8ir9mp9t", "metadata": {}, "outputs": [], "source": [ - "PYHEALTH_REPO_ROOT = '/Users/wpang/Desktop/PyHealth'\n", +<<<<<<< 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", - "CACHE_DIR = tempfile.mkdtemp()" + "CACHE_DIR = os.path.join(PYHEALTH_REPO_ROOT, \"local_data/local/data/wp/demo_pyhealth_cache\")" + ] + }, + { + "cell_type": "markdown", + "id": "15ccf548-f1fb-4e51-8c6d-2efcd0e4509c", + "metadata": {}, + "source": [ + "## 3. Multimodal Variants" + ] + }, + { + "cell_type": "markdown", + "id": "9342654c-da12-466a-86e3-ffc010194c89", + "metadata": {}, + "source": [ + "### 3.1 ICDLabsMIMIC4" ] }, { "cell_type": "code", - "execution_count": null, +<<<<<<< HEAD + "execution_count": 6, "id": "2f0f9ac3-1e3e-438b-b463-32e839941ce0", +======= + "execution_count": null, + "id": "7ebcce82-abdc-4b17-b3e5-d536fbb97859", + "metadata": {}, + "outputs": [], + "source": [ + "dataset = MIMIC4Dataset(\n", + " ehr_root=EHR_ROOT,\n", + " ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\", \"prescriptions\"],\n", + " cache_dir=CACHE_DIR,\n", + " num_workers=1,\n", + ")\n", + "\n", + "task = ICDLabsMIMIC4()\n", + "samples = dataset.set_task(task, num_workers=1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9b3d9f90-3257-4703-b68c-2ab286ba1e08", + "metadata": {}, + "outputs": [], + "source": [ + "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", + "\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", + "print(f\"Mortality Flag: {sample['mortality']}\")\n", + "\n", + "# Admissions\n", + "print_patient_admission_info(sample)\n", + "# Discharge Notes\n", + "print_patient_note_info(sample, note_type=\"discharge\")\n", + "# Radiology Notes\n", + "print_patient_note_info(sample, note_type=\"radiology\")\n", + "# Labs\n", + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "95c83b49-1903-4042-a3ef-0bf3f1d81b08", "metadata": {}, "outputs": [], + "source": [ + "clear_cache_directory(CACHE_DIR)" + ] + }, + { + "cell_type": "markdown", + "id": "516dbef8-43a6-4ebe-b0d8-29e423b0eb00", + "metadata": {}, + "source": [ + "### 3.2 ClinicalNotesMIMIC4" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1c777e0-75a0-4096-9e9b-cb3e955fae46", +>>>>>>> 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", " note_root=NOTE_ROOT,\n", - " ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\"],\n", " note_tables=[\"discharge\", \"radiology\"],\n", " cache_dir=CACHE_DIR,\n", " num_workers=1,\n", - ")" + ")\n", + "task = ClinicalNotesMIMIC4()\n", + "samples = dataset.set_task(task, num_workers=1)" ] }, { - "cell_type": "markdown", - "id": "238ac1f5-d26c-4b9b-8e72-8ee05344bdc2", + "cell_type": "code", + "execution_count": null, + "id": "570587e6-d3af-444d-85c6-45295a80b6b3", + "metadata": {}, + "outputs": [], + "source": [ + "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", + "\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", + "print(f\"Mortality Flag: {sample['mortality']}\")\n", + "\n", + "# Admissions\n", + "print_patient_admission_info(sample)\n", + "# Discharge Notes\n", + "print_patient_note_info(sample, note_type=\"discharge\")\n", + "# Radiology Notes\n", + "print_patient_note_info(sample, note_type=\"radiology\")\n", + "# Labs\n", + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cd4abcb3-8085-433f-8df6-a190c3a7c491", "metadata": {}, + "outputs": [], "source": [ - "## 3. Time Filtering" + "clear_cache_directory(CACHE_DIR)" ] }, { "cell_type": "markdown", - "id": "3eeef725-c61c-4320-a3bb-0d26c89d693c", + "id": "2697f301-92f3-43a6-9d7d-8beb15fdef24", "metadata": {}, "source": [ - "### 3.1 Full Patient Window\n", - "By default the task uses all data between the first admission time and the last discharge time across the patient's processed admissions." + "### 3.3 ClinicalNotesICDLabsMIMIC4" ] }, { "cell_type": "code", "execution_count": null, - "id": "d3b027ab-c814-49cd-8542-885d76b2401b", + "id": "546e4811-091b-45c3-97b6-87c79f520f33", "metadata": {}, "outputs": [], "source": [ - "%%capture\n", + "dataset = MIMIC4Dataset(\n", + " ehr_root=EHR_ROOT,\n", + " note_root=NOTE_ROOT,\n", + " ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\", \"prescriptions\"],\n", + " note_tables=[\"discharge\", \"radiology\"],\n", + " cache_dir=CACHE_DIR,\n", + " num_workers=1,\n", + ")\n", + "\n", "task = ClinicalNotesICDLabsMIMIC4()\n", "samples = dataset.set_task(task, num_workers=1)" ] @@ -200,25 +507,234 @@ { "cell_type": "code", "execution_count": null, - "id": "6b70acb2-bdea-4a43-8b9a-7725c34fcfce", + "id": "544f2f04-22e8-49a2-979d-961ec746fae9", "metadata": {}, "outputs": [], "source": [ - "random_patient_record = np.random.randint(0, len(samples)-1)\n", - "sample = samples[random_patient_record]\n", - "sample_patient_id = sample['patient_id']\n", + "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", "\n", - "print(f\"\\nPatient ID: {sample_patient_id}\")\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", "print(f\"Mortality Flag: {sample['mortality']}\")\n", "\n", "# Admissions\n", "print_patient_admission_info(sample)\n", - "\n", "# Discharge Notes\n", "print_patient_note_info(sample, note_type=\"discharge\")\n", + "# Radiology Notes\n", + "print_patient_note_info(sample, note_type=\"radiology\")\n", + "# Labs\n", + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2c9bec37-ff90-48f2-b340-37232a0cd947", + "metadata": {}, + "outputs": [], + "source": [ + "clear_cache_directory(CACHE_DIR)" + ] + }, + { + "cell_type": "markdown", + "id": "9bce1ec8-2108-4c42-bcc0-35916d05fbed", + "metadata": {}, + "source": [ + "## 4. Time Filtering" + ] + }, + { + "cell_type": "markdown", + "id": "3eeef725-c61c-4320-a3bb-0d26c89d693c", + "metadata": {}, + "source": [ + "### 4.1 Full Patient Window\n", + "By default the task uses all data between the first admission time and the last discharge time across the patient's processed admissions." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "d3b027ab-c814-49cd-8542-885d76b2401b", + "metadata": {}, + "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", + " note_root=NOTE_ROOT,\n", + " ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\"],\n", + " note_tables=[\"discharge\", \"radiology\"],\n", + " cache_dir=CACHE_DIR,\n", + " num_workers=1,\n", + ")\n", + "\n", + "task = ClinicalNotesICDLabsMIMIC4()\n", + "samples = dataset.set_task(task, num_workers=1)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "6b70acb2-bdea-4a43-8b9a-7725c34fcfce", + "metadata": {}, + "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", + "sample = random.choice(samples)\n", + "\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", + "print(f\"Mortality Flag: {sample['mortality']}\")\n", "\n", + "# Admissions\n", + "print_patient_admission_info(sample)\n", + "# Discharge Notes\n", + "print_patient_note_info(sample, note_type=\"discharge\")\n", "# Radiology Notes\n", - "print_patient_note_info(sample, note_type=\"radiology\")" + "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", + "print_patient_icd_info(sample)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c6814b1b-f523-4a37-bdef-929440c68952", + "metadata": {}, + "outputs": [], + "source": [ + "clear_cache_directory(CACHE_DIR)" +>>>>>>> origin/main ] }, { @@ -226,43 +742,117 @@ "id": "fd7d6767-ae51-4278-bb55-d12bdf3a4ff2", "metadata": {}, "source": [ - "### 3.2 Random Patient Window" + "### 4.2 Using `window_hours`\n", + "By using `window_hours`, we only look at data from `first admission + window hours` to predict the mortality task." ] }, { "cell_type": "code", "execution_count": null, - "id": "16121916-e52e-4f66-8463-ff0f931cf416", + "id": "28b19943-f5ee-4918-9003-60cf3c2ac623", "metadata": {}, "outputs": [], "source": [ - "%%capture\n", - "task = ClinicalNotesICDLabsMIMIC4(window_hours=12)\n", - "samples = dataset.set_task(task, num_workers=1)\n", - "sample = samples[random_patient_record]\n", - "sample_patient_id = sample['patient_id']" + "WINDOW_HOURS = 24" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, + "id": "16121916-e52e-4f66-8463-ff0f931cf416", + "metadata": {}, + "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", + " note_root=NOTE_ROOT,\n", + " ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\"],\n", + " note_tables=[\"discharge\", \"radiology\"],\n", + " cache_dir=CACHE_DIR,\n", + " num_workers=1,\n", + ")\n", + "\n", + "task = ClinicalNotesICDLabsMIMIC4(window_hours=WINDOW_HOURS)\n", + "samples = dataset.set_task(task, num_workers=1)" + ] + }, + { + "cell_type": "code", + "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": [ - "print(f\"\\nPatient ID: {sample_patient_id}\")\n", + "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", + "\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", "print(f\"Mortality Flag: {sample['mortality']}\")\n", "print(f\"Window start: {sample['window_start']}\")\n", "print(f\"Window end: {sample['window_end']}\")\n", "\n", "# Admissions\n", "print_patient_admission_info(sample)\n", - "\n", "# Discharge Notes\n", "print_patient_note_info(sample, note_type=\"discharge\")\n", - "\n", "# Radiology Notes\n", - "print_patient_note_info(sample, note_type=\"radiology\")" + "print_patient_note_info(sample, note_type=\"radiology\")\n", +<<<<<<< HEAD + "\n", + "# Labs\n", + "print_patient_lab_info(sample)" ] }, { @@ -279,6 +869,12 @@ "metadata": {}, "source": [ "- Remove `CACHE_DIR`" +======= + "# Labs\n", + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" +>>>>>>> origin/main ] }, { @@ -288,15 +884,15 @@ "metadata": {}, "outputs": [], "source": [ - "shutil.rmtree(CACHE_DIR) " + "clear_cache_directory(CACHE_DIR)" ] } ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "pyhealth", "language": "python", - "name": "python3" + "name": "pyhealth" }, "language_info": { "codemirror_mode": { @@ -307,7 +903,12 @@ "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", - "pygments_lexer": "ipython3" + "pygments_lexer": "ipython3", +<<<<<<< HEAD + "version": "3.12.13" +======= + "version": "3.13.3" +>>>>>>> origin/main } }, "nbformat": 4, diff --git a/examples/patient_linkage_mimic3_medlink.py b/examples/patient_linkage_mimic3_medlink.py index 2c38ba184..237d23f14 100644 --- a/examples/patient_linkage_mimic3_medlink.py +++ b/examples/patient_linkage_mimic3_medlink.py @@ -29,9 +29,6 @@ base_dataset = MIMIC3Dataset( root="/srv/local/data/physionet.org/files/mimiciii/1.4", tables=["DIAGNOSES_ICD"], - code_mapping={"ICD9CM": ("CCSCM", {})}, - dev=False, - refresh_cache=False, ) base_dataset.stat() diff --git a/examples/time_image_processor_tutorial.ipynb b/examples/time_image_processor_tutorial.ipynb index 07c50d0d6..3164d0c4b 100644 --- a/examples/time_image_processor_tutorial.ipynb +++ b/examples/time_image_processor_tutorial.ipynb @@ -66,7 +66,7 @@ "source": [ "## 2. Create Synthetic Time-Stamped X-ray Data\n", "\n", - "We simulate a scenario where each patient has 1–5 chest X-rays taken at different times during their hospital stay. Each image gets a timestamp representing days from the patient's first admission." + "We simulate a scenario where each patient has 1–5 chest X-rays taken at different times during their hospital stay, with one patient having no chest X-rays. Each image gets a timestamp representing days from the patient's first admission." ] }, { @@ -82,8 +82,11 @@ "NUM_PATIENTS = 20\n", "MAX_IMAGES_PER_PATIENT = 5\n", "\n", + "TOKEN_REPRESENTING_MISSING_PATH = \"\"\n", + "TOKEN_REPRESENTING_MISSING_FLOAT = 0.0\n", + "\n", "samples = []\n", - "for pid in range(NUM_PATIENTS):\n", + "for pid in range(NUM_PATIENTS-1):\n", " # Each patient has 1-5 X-rays taken at different times\n", " n_images = np.random.randint(1, MAX_IMAGES_PER_PATIENT + 1)\n", "\n", @@ -93,6 +96,11 @@ "\n", " image_paths = []\n", " for j in range(n_images):\n", + " if pid == np.random.randint(1, NUM_PATIENTS - 1) and j == np.random.randint(0, n_images):\n", + " image_paths.append(TOKEN_REPRESENTING_MISSING_PATH)\n", + " time_diffs[j] = TOKEN_REPRESENTING_MISSING_FLOAT\n", + " continue\n", + " \n", " # Synthetic grayscale X-ray with noise\n", " img_array = np.random.normal(80, 25, (224, 224))\n", "\n", @@ -120,6 +128,13 @@ " \"label\": label,\n", " })\n", "\n", + "# Patient with no CXR at all\n", + "samples.append({\n", + " \"patient_id\": f\"p{NUM_PATIENTS-1}\",\n", + " \"visit_id\": f\"v{NUM_PATIENTS-1}\",\n", + " \"chest_xray\": ([TOKEN_REPRESENTING_MISSING_PATH], [TOKEN_REPRESENTING_MISSING_FLOAT]),\n", + " \"label\": label})\n", + "\n", "print(f\"Created {NUM_PATIENTS} patients in {DATA_ROOT}\")\n", "print(f\"Images per patient: 1-{MAX_IMAGES_PER_PATIENT}\")" ] @@ -131,7 +146,9 @@ "outputs": [], "source": [ "# Inspect a sample patient\n", - "sample = samples[0]\n", + "PATIENT_NUMBER = 17\n", + "\n", + "sample = samples[PATIENT_NUMBER]\n", "paths, times = sample[\"chest_xray\"]\n", "\n", "print(f\"Patient {sample['patient_id']}:\")\n", @@ -166,6 +183,7 @@ "proc = TimeImageProcessor(\n", " image_size=224,\n", " mode=\"L\",\n", + " padding=TOKEN_REPRESENTING_MISSING_PATH\n", ")\n", "\n", "images, timestamps, tag = proc.process(sample[\"chest_xray\"])\n", @@ -231,6 +249,7 @@ " image_size=224,\n", " mode=\"L\",\n", " max_images=2,\n", + " padding=TOKEN_REPRESENTING_MISSING_PATH\n", ")\n", "\n", "imgs_trunc, ts_trunc, _ = proc_truncated.process(sample[\"chest_xray\"])\n", @@ -260,6 +279,7 @@ "proc_norm = TimeImageProcessor(\n", " image_size=128,\n", " mode=\"L\",\n", + " padding=TOKEN_REPRESENTING_MISSING_PATH,\n", " normalize=True,\n", " mean=[0.5],\n", " std=[0.5],\n", @@ -302,6 +322,7 @@ " \"chest_xray\": TimeImageProcessor(\n", " image_size=224,\n", " mode=\"L\",\n", + " padding=TOKEN_REPRESENTING_MISSING_PATH,\n", " max_images=4,\n", " ),\n", " },\n", @@ -321,7 +342,7 @@ "outputs": [], "source": [ "# Inspect a processed sample\n", - "processed = dataset[0]\n", + "processed = dataset[PATIENT_NUMBER]\n", "print(f\"Processed sample keys: {list(processed.keys())}\")\n", "print()\n", "\n", @@ -353,12 +374,12 @@ "metadata": {}, "outputs": [], "source": [ - "proc_demo = TimeImageProcessor(image_size=224, mode=\"L\", max_images=4)\n", + "proc_demo = TimeImageProcessor(image_size=224, mode=\"L\", padding=TOKEN_REPRESENTING_MISSING_PATH, max_images=4)\n", "\n", "print(f\"{'Patient':<10} {'N imgs':<8} {'Output Shape':<25} {'Time Range (days)'}\")\n", "print(\"-\" * 65)\n", "\n", - "for i in range(min(10, len(samples))):\n", + "for i in range(len(samples)):\n", " s = samples[i]\n", " paths_i, times_i = s[\"chest_xray\"]\n", " imgs_i, ts_i, _ = proc_demo.process((paths_i, times_i))\n", diff --git a/examples/vision_embedding_tutorial.ipynb b/examples/vision_embedding_tutorial.ipynb index 913422105..d241358c8 100644 --- a/examples/vision_embedding_tutorial.ipynb +++ b/examples/vision_embedding_tutorial.ipynb @@ -9,7 +9,7 @@ "\n", "This notebook demonstrates how to use the `VisionEmbeddingModel` for medical imaging tasks in PyHealth.\n", "\n", - "**Contributors:** Josh Steier \n", + "**Contributors:** Josh Steier, Joshua Chen, William Pang\n", "\n", "\n", "**Overview:**\n", @@ -31,18 +31,10 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "272cc5bb", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Running on device: cpu\n" - ] - } - ], + "outputs": [], "source": [ "import os\n", "import random\n", @@ -84,18 +76,10 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "f07f439c", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Created 200 synthetic images in C:\\Users\\637682\\AppData\\Local\\Temp\\chest_xray_jzp9zbiz\n" - ] - } - ], + "outputs": [], "source": [ "USE_SYNTHETIC = True\n", "MIMIC_CXR_ROOT = \"/path/to/physionet.org/files/mimic-cxr-jpg/2.0.0\"\n", @@ -171,22 +155,10 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "id": "c3f49ff5", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Label label vocab: {0: 0, 1: 1}\n", - "Total task samples: 200\n", - "Input schema: {'image': 'image'}\n", - "Output schema: {'label': 'binary'}\n", - "Train/Val/Test sizes: 140, 28, 32\n" - ] - } - ], + "outputs": [], "source": [ "image_mode = \"L\" if USE_SYNTHETIC else \"RGB\"\n", "\n", @@ -218,19 +190,10 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "4684f6c1", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{'patient_id': 'list(len=16)', 'visit_id': 'list(len=16)', 'image': 'Tensor(shape=(16, 1, 224, 224))', 'label': 'Tensor(shape=(16, 1))'}\n", - "Sample labels: [[0.0], [0.0], [0.0], [1.0], [1.0]]\n" - ] - } - ], + "outputs": [], "source": [ "BATCH_SIZE = 16\n", "\n", @@ -260,21 +223,10 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "id": "a47ebd7e", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "patient_id: list(len=16)\n", - "visit_id: list(len=16)\n", - "image: shape=torch.Size([16, 1, 224, 224]), dtype=torch.float32\n", - "label: shape=torch.Size([16, 1]), dtype=torch.float32\n" - ] - } - ], + "outputs": [], "source": [ "batch = next(iter(train_loader))\n", "\n", @@ -304,52 +256,10 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "c4d35bc6", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "C:\\Users\\637682\\Desktop\\pyhealth\\PyHealth\\pyhealth\\sampler\\sage_sampler.py:3: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n", - " import pkg_resources\n", - "c:\\Users\\637682\\Desktop\\pyhealth\\PyHealth\\pyhealth-env\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "VisionEmbeddingModel Output Shape Verification\n", - "============================================================\n", - "\n", - "Input shape: torch.Size([16, 1, 224, 224]) # (B, C, H, W)\n", - "\n", - "patch:\n", - " Output shape: (16, 197, 128) # (B, num_tokens, E)\n", - " Tokens: 197 = 196 patches + 1 CLS\n", - " Parameters: 58,240\n", - "\n", - "cnn:\n", - " Output shape: (16, 50, 128) # (B, num_tokens, E)\n", - " Tokens: 50 = 49 patches + 1 CLS\n", - " Parameters: 231,808\n", - "\n", - "resnet18:\n", - " Output shape: (16, 50, 128) # (B, num_tokens, E)\n", - " Tokens: 50 = 49 patches + 1 CLS\n", - " Parameters: 11,242,432\n", - "\n", - "resnet50:\n", - " Output shape: (16, 50, 128) # (B, num_tokens, E)\n", - " Tokens: 50 = 49 patches + 1 CLS\n", - " Parameters: 23,770,560\n", - "\n" - ] - } - ], + "outputs": [], "source": [ "from pyhealth.models import VisionEmbeddingModel\n", "\n", @@ -389,6 +299,53 @@ " print()" ] }, + { + "cell_type": "markdown", + "id": "604e21d3", + "metadata": {}, + "source": [ + "## 4.6. Mean Pooling: Compact Single-Vector Embeddings\n", + "\n", + "Pass `pool='mean'` to collapse all patch tokens into a single averaged vector per image:\n", + "\n", + "```\n", + "Input: (B, C, H, W)\n", + "Output with pool='mean': (B, 1, E)\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3be8a484", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Mean Pooling vs Token Sequence Output\")\n", + "print(\"=\" * 60)\n", + "print(f\"\\nInput shape: {batch['image'].shape} # (B, C, H, W)\")\n", + "print()\n", + "\n", + "# With mean pooling: single averaged vector\n", + "model_pooled = VisionEmbeddingModel(\n", + " dataset=sample_dataset,\n", + " embedding_dim=128,\n", + " backbone=\"cnn\",\n", + " pool=\"mean\",\n", + ")\n", + "\n", + "with torch.no_grad():\n", + " out_pooled = model_pooled({\"image\": batch[\"image\"]})\n", + "\n", + "info_seq = model_seq.get_output_info(\"image\")\n", + "info_pooled = model_pooled.get_output_info(\"image\")\n", + "\n", + "print(\"With mean pooling (pool='mean'):\")\n", + "print(f\" Output shape: {tuple(out_pooled['image'].shape)} # (B, 1, E)\")\n", + "print(f\" num_tokens: {info_pooled['num_tokens']}\")\n", + "print()" + ] + }, { "cell_type": "markdown", "id": "4ce223e3", @@ -401,49 +358,10 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": null, "id": "344eac50", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "VisionEmbeddingModel Backbone Comparison\n", - "==================================================\n", - "\n", - "patch:\n", - " Tokens: 197 (196 patches + CLS)\n", - " Parameters: 58,240\n", - "\n", - "cnn:\n", - " Tokens: 50 (49 patches + CLS)\n", - " Parameters: 231,808\n", - "\n", - "resnet18:\n", - " Tokens: 50 (49 patches + CLS)\n", - " Parameters: 11,242,432\n", - "Downloading: \"https://download.pytorch.org/models/resnet50-11ad3fa6.pth\" to C:\\Users\\637682/.cache\\torch\\hub\\checkpoints\\resnet50-11ad3fa6.pth\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "100%|██████████| 97.8M/97.8M [00:22<00:00, 4.63MB/s]\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "resnet50:\n", - " Tokens: 50 (49 patches + CLS)\n", - " Parameters: 23,770,560\n" - ] - } - ], + "outputs": [], "source": [ "from pyhealth.models import VisionEmbeddingModel\n", "\n", @@ -476,18 +394,7 @@ "execution_count": null, "id": "91b5f6a6", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Feature keys: ['image']\n", - "Label key: label\n", - "Mode: binary\n", - "Total parameters: 11,259,329\n" - ] - } - ], + "outputs": [], "source": [ "class VisionClassifier(nn.Module):\n", " \"\"\"End-to-end classifier using VisionEmbeddingModel.\"\"\"\n", @@ -575,30 +482,10 @@ }, { "cell_type": "code", - "execution_count": 25, + "execution_count": null, "id": "699c7e03", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "VisionClassifier(\n", - " (vision_encoder): VisionEmbeddingModel(backbone='resnet18', embedding_dim=128, fields=['image'])\n", - " (classifier): Sequential(\n", - " (0): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n", - " (1): Linear(in_features=128, out_features=128, bias=True)\n", - " (2): GELU(approximate='none')\n", - " (3): Dropout(p=0.1, inplace=False)\n", - " (4): Linear(in_features=128, out_features=1, bias=True)\n", - " )\n", - ")\n", - "Metrics: ['roc_auc']\n", - "Device: cpu\n", - "\n" - ] - } - ], + "outputs": [], "source": [ "from pyhealth.trainer import Trainer\n", "\n", @@ -628,259 +515,10 @@ }, { "cell_type": "code", - "execution_count": 26, + "execution_count": null, "id": "b12a766e", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Training:\n", - "Batch size: 16\n", - "Optimizer: \n", - "Optimizer params: {'lr': 0.0001}\n", - "Weight decay: 0.0\n", - "Max grad norm: 1.0\n", - "Val dataloader: \n", - "Monitor: roc_auc\n", - "Monitor criterion: max\n", - "Epochs: 5\n", - "Patience: None\n", - "\n" - ] - }, - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "3cadbf78b48241199f56d04123c7adb0", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Epoch 0 / 5: 0%| | 0/9 [00:00 bool: """URL detection.""" result = urlparse(path) @@ -124,6 +130,7 @@ def _litdata_merge(cache_dir: Path) -> None: cache_dir (Path): The cache directory containing LitData binary writer files. """ from litdata.streaming.writer import _INDEX_FILENAME + files = os.listdir(cache_dir) # Return if the index already exists @@ -134,13 +141,19 @@ def _litdata_merge(cache_dir: Path) -> None: # Return if there are no index files to merge if len(index_files) == 0: - raise ValueError("There are zero samples in the dataset, please check the task and processors.") + raise ValueError( + "There are zero samples in the dataset, please check the task and processors." + ) - BinaryWriter(cache_dir=str(cache_dir), chunk_bytes="64MB").merge(num_workers=len(index_files)) + BinaryWriter(cache_dir=str(cache_dir), chunk_bytes="64MB").merge( + num_workers=len(index_files) + ) class _ProgressContext: - def __init__(self, queue: multiprocessing.queues.Queue | None, total: int, **kwargs): + def __init__( + self, queue: multiprocessing.queues.Queue | None, total: int, **kwargs + ): """ :param queue: An existing queue (e.g., from multiprocessing). If provided, this class acts as a passthrough. @@ -167,8 +180,10 @@ def __exit__(self, exc_type, exc_val, exc_tb): if self.progress: self.progress.close() + _task_transform_progress: multiprocessing.queues.Queue | None = None + def _task_transform_init(queue: multiprocessing.queues.Queue) -> None: """ Initializer for worker processes to set up a global queue. @@ -179,7 +194,10 @@ def _task_transform_init(queue: multiprocessing.queues.Queue) -> None: global _task_transform_progress _task_transform_progress = queue -def _task_transform_fn(args: tuple[int, BaseTask, Iterable[str], pl.LazyFrame, Path]) -> None: + +def _task_transform_fn( + args: tuple[int, BaseTask, Iterable[str], pl.LazyFrame, Path], +) -> None: """ Worker function to apply task transformation on a chunk of patients. @@ -191,14 +209,16 @@ def _task_transform_fn(args: tuple[int, BaseTask, Iterable[str], pl.LazyFrame, P global_event_df (pl.LazyFrame): The global event dataframe. output_dir (Path): The output directory to save results. """ - BATCH_SIZE = 128 # Use a batch size 128 can reduce runtime by 30%. + BATCH_SIZE = 128 # Use a batch size 128 can reduce runtime by 30%. worker_id, task, patient_ids, global_event_df, output_dir = args total_patients = len(list(patient_ids)) - logger.info(f"Worker {worker_id} started processing {total_patients} patients. (Polars threads: {pl.thread_pool_size()})") + logger.info( + f"Worker {worker_id} started processing {total_patients} patients. (Polars threads: {pl.thread_pool_size()})" + ) with ( set_env(DATA_OPTIMIZER_GLOBAL_RANK=str(worker_id)), - _ProgressContext(_task_transform_progress, total=total_patients) as progress + _ProgressContext(_task_transform_progress, total=total_patients) as progress, ): writer = BinaryWriter(cache_dir=str(output_dir), chunk_bytes="64MB") @@ -208,8 +228,8 @@ def _task_transform_fn(args: tuple[int, BaseTask, Iterable[str], pl.LazyFrame, P complete = 0 patients = ( global_event_df.filter(pl.col("patient_id").is_in(batch)) - .collect(engine="streaming") - .partition_by("patient_id", as_dict=True) + .collect(engine="streaming") + .partition_by("patient_id", as_dict=True) ) for patient_id, patient_df in patients.items(): patient_id = patient_id[0] # Extract string from single-element list @@ -223,8 +243,10 @@ def _task_transform_fn(args: tuple[int, BaseTask, Iterable[str], pl.LazyFrame, P logger.info(f"Worker {worker_id} finished processing patients.") + _proc_transform_progress: multiprocessing.queues.Queue | None = None + def _proc_transform_init(queue: multiprocessing.queues.Queue) -> None: """ Initializer for worker processes to set up a global queue. @@ -235,6 +257,7 @@ def _proc_transform_init(queue: multiprocessing.queues.Queue) -> None: global _proc_transform_progress _proc_transform_progress = queue + def _proc_transform_fn(args: tuple[int, Path, int, int, Path]) -> None: """ Worker function to apply processors on a chunk of samples. @@ -250,11 +273,13 @@ def _proc_transform_fn(args: tuple[int, Path, int, int, Path]) -> None: BATCH_SIZE = 128 worker_id, task_df, start_idx, end_idx, output_dir = args total_samples = end_idx - start_idx - logger.info(f"Worker {worker_id} started processing {total_samples} samples. ({start_idx} to {end_idx})") + logger.info( + f"Worker {worker_id} started processing {total_samples} samples. ({start_idx} to {end_idx})" + ) with ( set_env(DATA_OPTIMIZER_GLOBAL_RANK=str(worker_id)), - _ProgressContext(_proc_transform_progress, total=total_samples) as progress + _ProgressContext(_proc_transform_progress, total=total_samples) as progress, ): writer = BinaryWriter(cache_dir=str(output_dir), chunk_bytes="64MB") @@ -400,9 +425,7 @@ def clean_tmpdir(self) -> None: if tmp_dir.exists(): shutil.rmtree(tmp_dir) - def _scan_csv_tsv_gz( - self, source_path: str - ) -> dd.DataFrame: + def _scan_csv_tsv_gz(self, source_path: str) -> dd.DataFrame: """Scans a CSV/TSV file (possibly gzipped) and returns a Dask DataFrame. If the cached Parquet file does not exist, it converts the source CSV/TSV file @@ -474,6 +497,7 @@ def _scan_csv_tsv_gz( return df.replace("", pd.NA) # Replace empty strings with NaN def _event_transform(self, output_dir: Path) -> None: + compute_ok = False try: df = self.load_data() disable_distributed = os.environ.get( @@ -521,14 +545,29 @@ def _event_transform(self, output_dir: Path) -> None: handle = client.compute(collection) dask_progress(handle) handle.result() # type: ignore + compute_ok = True # Data is fully written to disk + except TimeoutError: + if compute_ok: + # Cluster shutdown timed out after successful compute — data is intact + logger.warning( + "Dask cluster shutdown timed out, but data was written successfully. Continuing." + ) + else: + if output_dir.exists(): + logger.error( + f"Error during caching, removing incomplete file {output_dir}" + ) + shutil.rmtree(output_dir) + raise except Exception as e: if output_dir.exists(): - logger.error(f"Error during caching, removing incomplete file {output_dir}") + logger.error( + f"Error during caching, removing incomplete file {output_dir}" + ) shutil.rmtree(output_dir) raise e finally: self.clean_tmpdir() - pass @property def global_event_df(self) -> pl.LazyFrame: @@ -537,14 +576,16 @@ def global_event_df(self) -> pl.LazyFrame: Returns: Path: The path to the cached event dataframe. """ - self._main_guard(type(self).global_event_df.fget.__name__) # type: ignore + self._main_guard(type(self).global_event_df.fget.__name__) # type: ignore if self._global_event_df is None: ret_path = self.cache_dir / "global_event_df.parquet" cache_valid = ret_path.is_dir() and any(ret_path.glob("*.parquet")) if not cache_valid: if ret_path.exists(): - logger.warning(f"Incomplete parquet cache at {ret_path} (directory exists but contains no parquet files). Removing and rebuilding.") + logger.warning( + f"Incomplete parquet cache at {ret_path} (directory exists but contains no parquet files). Removing and rebuilding." + ) shutil.rmtree(ret_path) logger.info(f"No cached event dataframe found. Creating: {ret_path}") self._event_transform(ret_path) @@ -744,10 +785,14 @@ def default_task(self) -> Optional[BaseTask]: """ return None - def _task_transform(self, task: BaseTask, output_dir: Path, num_workers: int) -> None: + def _task_transform( + self, task: BaseTask, output_dir: Path, num_workers: int + ) -> None: self._main_guard(self._task_transform.__name__) - logger.info(f"Applying task transformations on data with {num_workers} workers...") + logger.info( + f"Applying task transformations on data with {num_workers} workers..." + ) global_event_df = task.pre_filter(self.global_event_df) patient_ids = ( global_event_df.select("patient_id") @@ -759,34 +804,52 @@ def _task_transform(self, task: BaseTask, output_dir: Path, num_workers: int) -> ) if in_notebook(): - logger.info("Detected Jupyter notebook environment, setting num_workers to 1") + logger.info( + "Detected Jupyter notebook environment, setting num_workers to 1" + ) num_workers = 1 - num_workers = min(num_workers, len(patient_ids)) # Avoid spawning empty workers + num_workers = min(num_workers, len(patient_ids)) # Avoid spawning empty workers # This ensures worker's polars threads are limited to avoid oversubscription, # which can lead to additional 75% speedup when num_workers is large. threads_per_worker = max(1, (os.cpu_count() or 1) // num_workers) try: - with set_env(POLARS_MAX_THREADS=str(threads_per_worker), DATA_OPTIMIZER_NUM_WORKERS=str(num_workers)): + with set_env( + POLARS_MAX_THREADS=str(threads_per_worker), + DATA_OPTIMIZER_NUM_WORKERS=str(num_workers), + ): if num_workers == 1: logger.info("Single worker mode, processing sequentially") - _task_transform_fn((0, task, patient_ids, global_event_df, output_dir)) + _task_transform_fn( + (0, task, patient_ids, global_event_df, output_dir) + ) _litdata_merge(output_dir) return # spwan is required for polars in multiprocessing, see https://docs.pola.rs/user-guide/misc/multiprocessing/#summary ctx = multiprocessing.get_context("spawn") queue = ctx.Queue() - args_list = [( - worker_id, - task, - pids, - global_event_df, - output_dir, - ) for worker_id, pids in enumerate(itertools.batched(patient_ids, len(patient_ids) // num_workers + 1))] - with ctx.Pool(processes=num_workers, initializer=_task_transform_init, initargs=(queue,)) as pool: - result = pool.map_async(_task_transform_fn, args_list) # type: ignore + args_list = [ + ( + worker_id, + task, + pids, + global_event_df, + output_dir, + ) + for worker_id, pids in enumerate( + itertools.batched( + patient_ids, len(patient_ids) // num_workers + 1 + ) + ) + ] + with ctx.Pool( + processes=num_workers, + initializer=_task_transform_init, + initargs=(queue,), + ) as pool: + result = pool.map_async(_task_transform_fn, args_list) # type: ignore with tqdm(total=len(patient_ids)) as progress: while not result.ready(): try: @@ -797,26 +860,32 @@ def _task_transform(self, task: BaseTask, output_dir: Path, num_workers: int) -> # remaining items while not queue.empty(): progress.update(queue.get()) - result.get() # ensure exceptions are raised + result.get() # ensure exceptions are raised _litdata_merge(output_dir) logger.info(f"Task transformation completed and saved to {output_dir}") except Exception as e: - logger.error(f"Error during task transformation, cleaning up output directory: {output_dir}") + logger.error( + f"Error during task transformation, cleaning up output directory: {output_dir}" + ) shutil.rmtree(output_dir) raise e - def _proc_transform(self, task_df: Path, output_dir: Path, num_workers: int) -> None: + def _proc_transform( + self, task_df: Path, output_dir: Path, num_workers: int + ) -> None: self._main_guard(self._proc_transform.__name__) logger.info(f"Applying processors on data with {num_workers} workers...") num_samples = len(litdata.StreamingDataset(str(task_df))) if in_notebook(): - logger.info("Detected Jupyter notebook environment, setting num_workers to 1") + logger.info( + "Detected Jupyter notebook environment, setting num_workers to 1" + ) num_workers = 1 - num_workers = min(num_workers, num_samples) # Avoid spawning empty workers + num_workers = min(num_workers, num_samples) # Avoid spawning empty workers try: with set_env(DATA_OPTIMIZER_NUM_WORKERS=str(num_workers)): if num_workers == 1: @@ -827,16 +896,25 @@ def _proc_transform(self, task_df: Path, output_dir: Path, num_workers: int) -> ctx = multiprocessing.get_context("spawn") queue = ctx.Queue() - linspace = more_itertools.sliding_window(np.linspace(0, num_samples, num_workers + 1, dtype=int), 2) - args_list = [( - worker_id, - task_df, - start, - end, - output_dir, - ) for worker_id, (start, end) in enumerate(linspace)] - with ctx.Pool(processes=num_workers, initializer=_proc_transform_init, initargs=(queue,)) as pool: - result = pool.map_async(_proc_transform_fn, args_list) # type: ignore + linspace = more_itertools.sliding_window( + np.linspace(0, num_samples, num_workers + 1, dtype=int), 2 + ) + args_list = [ + ( + worker_id, + task_df, + start, + end, + output_dir, + ) + for worker_id, (start, end) in enumerate(linspace) + ] + with ctx.Pool( + processes=num_workers, + initializer=_proc_transform_init, + initargs=(queue,), + ) as pool: + result = pool.map_async(_proc_transform_fn, args_list) # type: ignore with tqdm(total=num_samples) as progress: while not result.ready(): try: @@ -847,10 +925,12 @@ def _proc_transform(self, task_df: Path, output_dir: Path, num_workers: int) -> # remaining items while not queue.empty(): progress.update(queue.get()) - result.get() # ensure exceptions are raised + result.get() # ensure exceptions are raised _litdata_merge(output_dir) - logger.info(f"Processor transformation completed and saved to {output_dir}") + logger.info( + f"Processor transformation completed and saved to {output_dir}" + ) except Exception as e: logger.error(f"Error during processor transformation.") shutil.rmtree(output_dir) @@ -910,10 +990,14 @@ def set_task( "output_schema": task.output_schema, }, sort_keys=True, - default=str + default=str, ) - cache_dir = self.cache_dir / "tasks" / f"{task.task_name}_{uuid.uuid5(uuid.NAMESPACE_DNS, task_params)}" + cache_dir = ( + self.cache_dir + / "tasks" + / f"{task.task_name}_{uuid.uuid5(uuid.NAMESPACE_DNS, task_params)}" + ) cache_dir.mkdir(parents=True, exist_ok=True) proc_params = json.dumps( @@ -936,54 +1020,79 @@ def set_task( ), }, sort_keys=True, - default=str + default=str, ) task_df_path = Path(cache_dir) / "task_df.ld" - samples_path = Path(cache_dir) / f"samples_{uuid.uuid5(uuid.NAMESPACE_DNS, proc_params)}.ld" + samples_path = ( + Path(cache_dir) + / f"samples_{uuid.uuid5(uuid.NAMESPACE_DNS, proc_params)}.ld" + ) logger.info(f"Task cache paths: task_df={task_df_path}, samples={samples_path}") task_df_path.mkdir(parents=True, exist_ok=True) samples_path.mkdir(parents=True, exist_ok=True) - if not (samples_path / "index.json").exists(): - # Check if index.json exists to verify cache integrity, this - # is the standard file for litdata.StreamingDataset - if not (task_df_path / "index.json").exists(): - self._task_transform( - task, - task_df_path, - num_workers, - ) - else: - logger.info(f"Found cached task dataframe at {task_df_path}, skipping task transformation.") + def _is_valid_litdata_cache(path: Path) -> bool: + """Return True if index.json exists. litdata only writes index.json after + all .bin chunks are flushed, so its presence guarantees a complete cache.""" + return (path / "index.json").exists() + + # Fast path: cache already valid, no lock needed (reads are always safe). + # Slow path: acquire a per-cache-dir file lock so that concurrent processes + # (e.g. parallel hparam jobs) don't race to build the same litdata cache. + # The double-checked pattern inside the lock means the winner builds it + # once; all others wait, re-check, and skip. + if not _is_valid_litdata_cache(samples_path): + lock_path = Path(cache_dir) / "build.lock" + with FileLock(str(lock_path), timeout=7200): + # Re-check inside the lock — another process may have built it + # while we were waiting. + if _is_valid_litdata_cache(samples_path): + logger.info( + f"Found cached processed samples at {samples_path} (built by another process)." + ) + else: + # Check if task_df cache is valid; rebuild if not + if not _is_valid_litdata_cache(task_df_path): + self._task_transform( + task, + task_df_path, + num_workers, + ) + else: + logger.info( + f"Found cached task dataframe at {task_df_path}, skipping task transformation." + ) - # Build processors and fit on the dataset - logger.info(f"Fitting processors on the dataset...") - dataset = litdata.StreamingDataset( - str(task_df_path), - transform=lambda x: pickle.loads(x["sample"]), - ) - builder = SampleBuilder( - input_schema=task.input_schema, # type: ignore - output_schema=task.output_schema, # type: ignore - input_processors=input_processors, - output_processors=output_processors, - ) - builder.fit(dataset) - builder.save(str(samples_path / "schema.pkl")) - - # Apply processors and save final samples to cache_dir - logger.info(f"Processing samples and saving to {samples_path}...") - self._proc_transform( - task_df_path, - samples_path, - num_workers, - ) - logger.info(f"Cached processed samples to {samples_path}") + # Build processors and fit on the dataset + logger.info(f"Fitting processors on the dataset...") + dataset = litdata.StreamingDataset( + str(task_df_path), + transform=lambda x: pickle.loads(x["sample"]), + ) + builder = SampleBuilder( + input_schema=task.input_schema, # type: ignore + output_schema=task.output_schema, # type: ignore + input_processors=input_processors, + output_processors=output_processors, + ) + builder.fit(dataset) + builder.save(str(samples_path / "schema.pkl")) + + # Apply processors and save final samples to cache_dir + logger.info(f"Processing samples and saving to {samples_path}...") + self._proc_transform( + task_df_path, + samples_path, + num_workers, + ) + logger.info(f"Cached processed samples to {samples_path}") else: - logger.info(f"Found cached processed samples at {samples_path}, skipping processing.") + logger.info( + f"Found cached processed samples at {samples_path}, skipping processing." + ) return SampleDataset( path=str(samples_path), diff --git a/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml b/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml new file mode 100644 index 000000000..aa4b1c419 --- /dev/null +++ b/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml @@ -0,0 +1,105 @@ +version: "2.1.0" +tables: + metadata: + file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + patient_id: "subject_id" + timestamp: + - "studydate" + - "studytime" + timestamp_format: "%Y%m%d%H%M%S" + attributes: + - "image_path" + - "dicom_id" + - "study_id" + - "performedprocedurestepdescription" + - "viewposition" + - "rows" + - "columns" + - "procedurecodesequence_codemeaning" + - "viewcodesequence_codemeaning" + - "patientorientationcodesequence_codemeaning" + + chexpert: + file_path: "mimic-cxr-2.0.0-chexpert.csv" + patient_id: "subject_id" + join: + - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + "on": "study_id" + how: "inner" + columns: + - "studydate" + - "studytime" + - "dicom_id" + timestamp: + - "studydate" + - "studytime" + timestamp_format: "%Y%m%d%H%M%S" + attributes: + - "dicom_id" + - "study_id" + - "atelectasis" + - "cardiomegaly" + - "consolidation" + - "edema" + - "enlarged cardiomediastinum" + - "fracture" + - "lung lesion" + - "lung opacity" + - "no finding" + - "pleural effusion" + - "pleural other" + - "pneumonia" + - "pneumothorax" + - "support devices" + + negbio: + file_path: "mimic-cxr-2.0.0-negbio.csv" + patient_id: "subject_id" + join: + - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + "on": "study_id" + how: "inner" + columns: + - "studydate" + - "studytime" + - "dicom_id" + timestamp: + - "studydate" + - "studytime" + timestamp_format: "%Y%m%d%H%M%S" + attributes: + - "dicom_id" + - "study_id" + - "atelectasis" + - "cardiomegaly" + - "consolidation" + - "edema" + - "enlarged cardiomediastinum" + - "fracture" + - "lung lesion" + - "lung opacity" + - "no finding" + - "pleural effusion" + - "pleural other" + - "pneumonia" + - "pneumothorax" + - "support devices" + + split: + file_path: "mimic-cxr-2.0.0-split.csv" + patient_id: "subject_id" + join: + - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + "on": "dicom_id" + how: "inner" + columns: + - "studydate" + - "studytime" + timestamp: + - "studydate" + - "studytime" + timestamp_format: "%Y%m%d%H%M%S" + attributes: + - "dicom_id" + - "study_id" + - "split" diff --git a/pyhealth/datasets/configs/mimic4_ehr.yaml b/pyhealth/datasets/configs/mimic4_ehr.yaml index 5eb74de84..414f31308 100644 --- a/pyhealth/datasets/configs/mimic4_ehr.yaml +++ b/pyhealth/datasets/configs/mimic4_ehr.yaml @@ -135,6 +135,7 @@ tables: patient_id: "subject_id" timestamp: "chartdate" attributes: + - "hadm_id" - "hcpcs_cd" - "seq_num" - "short_description" diff --git a/pyhealth/datasets/medical_transcriptions.py b/pyhealth/datasets/medical_transcriptions.py index 84e1481c5..e52dc20d5 100644 --- a/pyhealth/datasets/medical_transcriptions.py +++ b/pyhealth/datasets/medical_transcriptions.py @@ -39,6 +39,9 @@ def __init__( root: str, dataset_name: Optional[str] = None, config_path: Optional[str] = None, + cache_dir=None, + num_workers: int = 1, + dev: bool = False, ) -> None: if config_path is None: logger.info("No config path provided, using default config") @@ -51,6 +54,9 @@ def __init__( tables=default_tables, dataset_name=dataset_name or "medical_transcriptions", config_path=config_path, + cache_dir=cache_dir, + num_workers=num_workers, + dev=dev, ) return diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py index d4a8e0649..089a6d2f3 100644 --- a/pyhealth/datasets/mimic4.py +++ b/pyhealth/datasets/mimic4.py @@ -223,6 +223,102 @@ def process_image_path(x): return +class MIMIC4CXRSunlabDataset(BaseDataset): + """ + Sunlab variant of the MIMIC-CXR Chest X-ray dataset. + + This variant uses the existing metadata CSV and derives flattened image + paths at ``images/{dicom_id}.jpg``. + """ + + def __init__( + self, + root: str, + tables: List[str], + dataset_name: str = "mimic4_cxr_sunlab", + config_path: Optional[str] = None, + cache_dir: Optional[str] = None, + **kwargs, + ): + if config_path is None: + config_path = os.path.join( + os.path.dirname(__file__), "configs", "mimic4_cxr_sunlab.yaml" + ) + logger.info(f"Using default Sunlab CXR config: {config_path}") + self.prepare_metadata(root) + log_memory_usage(f"Before initializing {dataset_name}") + super().__init__( + root=root, + tables=tables, + dataset_name=dataset_name, + config_path=config_path, + cache_dir=cache_dir, + **kwargs, + ) + log_memory_usage(f"After initializing {dataset_name}") + + @staticmethod + def _resolve_column_name(columns: List[str], target: str) -> str: + lower_to_original = {col.lower(): col for col in columns} + resolved = lower_to_original.get(target.lower()) + if resolved is None: + raise ValueError( + f"Expected column '{target}' in metadata, available columns: {columns}" + ) + return resolved + + def prepare_metadata(self, root: str) -> None: + metadata_path = os.path.join(root, "mimic-cxr-2.0.0-metadata.csv") + if not os.path.exists(metadata_path): + raise FileNotFoundError( + f"Sunlab metadata file not found: {metadata_path}. " + "Expected existing metadata linked by dicom_id/subject_id/study_id." + ) + + images_dir = os.path.join(root, "images") + if not os.path.isdir(images_dir): + raise FileNotFoundError( + f"Sunlab images directory not found: {images_dir}. " + "Expected flattened image files at images/{dicom_id}.jpg." + ) + + metadata = pd.read_csv(metadata_path, dtype=str) + + dicom_col = self._resolve_column_name(metadata.columns.tolist(), "dicom_id") + study_time_col = self._resolve_column_name( + metadata.columns.tolist(), "studytime" + ) + + # Normalize StudyTime so timestamps parse with %Y%m%d%H%M%S in config. + def normalize_studytime(value: Optional[str]) -> str: + if value is None: + return "000000" + value_str = str(value).strip() + if value_str == "" or value_str.lower() == "nan": + return "000000" + try: + return f"{int(float(value_str)):06d}" + except Exception: + digits = "".join(ch for ch in value_str if ch.isdigit()) + if digits == "": + return "000000" + return digits[:6].zfill(6) + + metadata[study_time_col] = metadata[study_time_col].apply(normalize_studytime) + + metadata["image_path"] = metadata[dicom_col].apply( + lambda dicom_id: os.path.join(root, "images", f"{dicom_id}.jpg") + ) + + # Align with existing config conventions by using lowercase headers. + metadata.columns = [col.lower() for col in metadata.columns] + + metadata.to_csv( + os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"), + index=False, + ) + + class MIMIC4Dataset(BaseDataset): """ Unified MIMIC-IV dataset with support for EHR, clinical notes, and X-rays. @@ -242,6 +338,7 @@ class MIMIC4Dataset(BaseDataset): ehr_config_path: Path to the EHR config file note_config_path: Path to the note config file cxr_config_path: Path to the CXR config file + cxr_variant: Which CXR variant to load ("default" or "sunlab") dataset_name: Name of the dataset dev: Whether to enable dev mode (limit to 1000 patients) @@ -279,6 +376,7 @@ def __init__( ehr_config_path: Optional[str] = None, note_config_path: Optional[str] = None, cxr_config_path: Optional[str] = None, + cxr_variant: str = "default", dataset_name: str = "mimic4", dev: Union[bool, int] = False, cache_dir: Optional[str] = None, @@ -340,17 +438,33 @@ def __init__( # Initialize CXR dataset if root is provided if cxr_root is not None: + if cxr_variant not in {"default", "sunlab"}: + raise ValueError( + f"Unknown cxr_variant '{cxr_variant}'. " + "Expected one of {'default', 'sunlab'}." + ) + logger.info( - f"Initializing MIMIC4CXRDataset with tables: {cxr_tables} (dev mode: {dev})" - ) - self.sub_datasets["cxr"] = MIMIC4CXRDataset( - root=cxr_root, - tables=cxr_tables, - config_path=cxr_config_path, - cache_dir=str(self.cache_dir), - dev=dev, - num_workers=num_workers, + f"Initializing MIMIC4 CXR variant '{cxr_variant}' with tables: {cxr_tables} (dev mode: {dev})" ) + if cxr_variant == "sunlab": + self.sub_datasets["cxr"] = MIMIC4CXRSunlabDataset( + root=cxr_root, + tables=cxr_tables, + config_path=cxr_config_path, + cache_dir=str(self.cache_dir), + dev=dev, + num_workers=num_workers, + ) + else: + self.sub_datasets["cxr"] = MIMIC4CXRDataset( + root=cxr_root, + tables=cxr_tables, + config_path=cxr_config_path, + cache_dir=str(self.cache_dir), + dev=dev, + num_workers=num_workers, + ) log_memory_usage("After CXR dataset initialization") log_memory_usage("Completed MIMIC4Dataset init") diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py index 58c306ec0..b2ea98854 100644 --- a/pyhealth/datasets/splitter.py +++ b/pyhealth/datasets/splitter.py @@ -326,6 +326,98 @@ def split_by_patient_conformal( return train_dataset, val_dataset, cal_dataset, test_dataset +def split_by_patient_conformal_tuh( + dataset: SampleDataset, + ratios: Union[Tuple[float, float, float], List[float]], + seed: Optional[int] = None, + get_index: Optional[bool] = False, +): + """Splits a TUH EEG dataset by PATIENT within the TUH train partition. + + Unlike :func:`split_by_sample_conformal_tuh`, which shuffles windows + independently, this function ensures that **all windows from a given patient + stay in the same split** (train, val, or cal). This prevents within-patient + data leakage between training and calibration — a critical requirement for + honest evaluation of conformal prediction methods, especially NCP, which + relies on nearest-neighbor search in embedding space. + + The TUH corpus is designed so that each patient appears exclusively in either + the training partition (265 patients for TUSZ) or the evaluation partition (50 + patients), never both. The test set is therefore always the TUH eval + partition (fixed, fully held-out patients). + + Args: + dataset: a ``SampleDataset`` produced by ``EEGEventsTUEV`` or + ``EEGAbnormalTUAB``. Each sample must carry a ``"split"`` field + (``"train"`` or ``"eval"``) and the dataset must have a + ``patient_to_index`` mapping. + ratios: fraction of *patients* in the TUH train partition to assign to + train / val / cal respectively. Must be a length-3 sequence summing + to 1.0. + seed: random seed used to shuffle the patient list. + get_index: if ``True``, return four :class:`torch.Tensor` index vectors + instead of :class:`~torch.utils.data.Subset` objects. + + Returns: + ``(train_dataset, val_dataset, cal_dataset, test_dataset)`` + """ + assert len(ratios) == 3, ( + "ratios must have exactly 3 elements (train/val/cal). " + "The test set is determined by the TUH eval partition." + ) + assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" + + # Bucket patients by partition using dataset.patient_to_index (fast path). + # TUH guarantees each patient is in exactly one partition, so inspecting the + # first sample per patient is sufficient. + train_patient_to_indices: dict = {} + test_list: List[int] = [] + + for pid, indices in dataset.patient_to_index.items(): + first_sample = dataset[indices[0]] + assert "split" in first_sample, ( + f"Patient {pid}: sample missing 'split' field. " + "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." + ) + if first_sample["split"] == "train": + train_patient_to_indices[pid] = list(indices) + else: + test_list.extend(list(indices)) + + # Shuffle patients deterministically + patient_ids = list(train_patient_to_indices.keys()) + rng = np.random.default_rng(seed) + rng.shuffle(patient_ids) + + n = len(patient_ids) + train_end = int(n * ratios[0]) + val_end = int(n * (ratios[0] + ratios[1])) + + train_pids = patient_ids[:train_end] + val_pids = patient_ids[train_end:val_end] + cal_pids = patient_ids[val_end:] + + train_index = np.array(list(chain(*[train_patient_to_indices[p] for p in train_pids]))) + val_index = np.array(list(chain(*[train_patient_to_indices[p] for p in val_pids]))) + cal_index = np.array(list(chain(*[train_patient_to_indices[p] for p in cal_pids]))) + test_index = np.array(test_list) + + if get_index: + return ( + torch.tensor(train_index), + torch.tensor(val_index), + torch.tensor(cal_index), + torch.tensor(test_index), + ) + else: + return ( + dataset.subset(train_index), # type: ignore + dataset.subset(val_index), # type: ignore + dataset.subset(cal_index), # type: ignore + dataset.subset(test_index), # type: ignore + ) + + def split_by_sample_conformal_tuh( dataset: SampleDataset, ratios: Union[Tuple[float, float, float], List[float]], @@ -396,6 +488,157 @@ def split_by_sample_conformal_tuh( ) +def split_by_patient_tuh( + dataset: SampleDataset, + ratios: Union[Tuple[float, float], List[float]], + seed: Optional[int] = None, + get_index: Optional[bool] = False, +): + """Splits a TUH EEG dataset by PATIENT within the TUH train partition. + + Like :func:`split_by_patient_conformal_tuh` but returns a 3-way split + (train / val / test) instead of 4-way (train / val / cal / test). The + test set is always the TUH eval partition; no calibration set is produced. + + Ensures that **all windows from a given patient stay in the same split** + (train or val), preventing within-patient data leakage. + + Args: + dataset: a ``SampleDataset`` produced by ``EEGEventsTUEV`` or + ``EEGAbnormalTUAB``. Each sample must carry a ``"split"`` field + (``"train"`` or ``"eval"``) and the dataset must have a + ``patient_to_index`` mapping. + ratios: fraction of *patients* in the TUH train partition to assign to + train / val respectively. Must be a length-2 sequence summing + to 1.0. + seed: random seed used to shuffle the patient list. + get_index: if ``True``, return three :class:`torch.Tensor` index + vectors instead of :class:`~torch.utils.data.Subset` objects. + + Returns: + ``(train_dataset, val_dataset, test_dataset)`` + """ + assert len(ratios) == 2, ( + "ratios must have exactly 2 elements (train/val). " + "The test set is determined by the TUH eval partition." + ) + assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" + + train_patient_to_indices: dict = {} + test_list: List[int] = [] + + for pid, indices in dataset.patient_to_index.items(): + first_sample = dataset[indices[0]] + assert "split" in first_sample, ( + f"Patient {pid}: sample missing 'split' field. " + "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." + ) + if first_sample["split"] == "train": + train_patient_to_indices[pid] = list(indices) + else: + test_list.extend(list(indices)) + + patient_ids = list(train_patient_to_indices.keys()) + rng = np.random.default_rng(seed) + rng.shuffle(patient_ids) + + n = len(patient_ids) + train_end = int(n * ratios[0]) + + train_pids = patient_ids[:train_end] + val_pids = patient_ids[train_end:] + + train_index = np.array( + list(chain(*[train_patient_to_indices[p] for p in train_pids])) + ) + val_index = np.array( + list(chain(*[train_patient_to_indices[p] for p in val_pids])) + ) + test_index = np.array(test_list) + + if get_index: + return ( + torch.tensor(train_index), + torch.tensor(val_index), + torch.tensor(test_index), + ) + else: + return ( + dataset.subset(train_index), # type: ignore + dataset.subset(val_index), # type: ignore + dataset.subset(test_index), # type: ignore + ) + + +def split_by_sample_tuh( + dataset: SampleDataset, + ratios: Union[Tuple[float, float], List[float]], + seed: Optional[int] = None, + get_index: Optional[bool] = False, +): + """Splits a TUH EEG dataset (TUEV/TUAB) using its pre-defined split. + + Like :func:`split_by_sample_conformal_tuh` but returns a 3-way split + (train / val / test) instead of 4-way (train / val / cal / test). The + test set is always the TUH eval partition; no calibration set is produced. + + Args: + dataset: a ``SampleDataset`` object produced by ``EEGEventsTUEV`` or + ``EEGAbnormalTUAB`` + ratios: the fraction of the train pool assigned to train / val + respectively. Must be a length-2 sequence summing to 1.0. + seed: random seed for shuffling the train pool + get_index: if True, return three ``torch.Tensor`` index vectors instead + of ``Subset`` objects + + Returns: + train_dataset, val_dataset, test_dataset + """ + assert len(ratios) == 2, ( + "ratios must have exactly 2 elements (train/val). " + "The test set is determined by the dataset's own eval partition." + ) + assert abs(sum(ratios) - 1.0) < 1e-6, "ratios must sum to 1.0" + + for i in range(len(dataset)): + assert "split" in dataset[i], ( + f"Sample {i} is missing the 'split' field. " + "Use EEGEventsTUEV or EEGAbnormalTUAB to build the dataset." + ) + + train_pool: List[int] = [] + test_list: List[int] = [] + for i in range(len(dataset)): + if dataset[i]["split"] == "train": + train_pool.append(i) + else: + test_list.append(i) + + rng = np.random.default_rng(seed) + train_arr = np.array(train_pool) + rng.shuffle(train_arr) + + n = len(train_arr) + train_end = int(n * ratios[0]) + + train_index = train_arr[:train_end] + val_index = train_arr[train_end:] + test_index = np.array(test_list) + + if get_index: + return ( + torch.tensor(train_index), + torch.tensor(val_index), + torch.tensor(test_index), + ) + else: + return ( + dataset.subset(train_index), # type: ignore + dataset.subset(val_index), # type: ignore + dataset.subset(test_index), # type: ignore + ) + + def split_by_sample_conformal( dataset: SampleDataset, ratios: Union[Tuple[float, float, float, float], List[float]], diff --git a/pyhealth/metrics/multiclass.py b/pyhealth/metrics/multiclass.py index ed819b349..5b6db3a01 100644 --- a/pyhealth/metrics/multiclass.py +++ b/pyhealth/metrics/multiclass.py @@ -96,24 +96,28 @@ def multiclass_metrics_fn( output = {} for metric in metrics: if metric == "roc_auc_macro_ovo": + y_prob_roc = y_prob[:, 1] if y_prob.shape[1] == 2 else y_prob roc_auc_macro_ovo = sklearn_metrics.roc_auc_score( - y_true, y_prob, average="macro", multi_class="ovo" - ) + y_true, y_prob_roc, average="macro", multi_class="ovo" + ) if y_prob.shape[1] > 2 else sklearn_metrics.roc_auc_score(y_true, y_prob_roc) output["roc_auc_macro_ovo"] = roc_auc_macro_ovo elif metric == "roc_auc_macro_ovr": + y_prob_roc = y_prob[:, 1] if y_prob.shape[1] == 2 else y_prob roc_auc_macro_ovr = sklearn_metrics.roc_auc_score( - y_true, y_prob, average="macro", multi_class="ovr" - ) + y_true, y_prob_roc, average="macro", multi_class="ovr" + ) if y_prob.shape[1] > 2 else sklearn_metrics.roc_auc_score(y_true, y_prob_roc) output["roc_auc_macro_ovr"] = roc_auc_macro_ovr elif metric == "roc_auc_weighted_ovo": + y_prob_roc = y_prob[:, 1] if y_prob.shape[1] == 2 else y_prob roc_auc_weighted_ovo = sklearn_metrics.roc_auc_score( - y_true, y_prob, average="weighted", multi_class="ovo" - ) + y_true, y_prob_roc, average="weighted", multi_class="ovo" + ) if y_prob.shape[1] > 2 else sklearn_metrics.roc_auc_score(y_true, y_prob_roc) output["roc_auc_weighted_ovo"] = roc_auc_weighted_ovo elif metric == "roc_auc_weighted_ovr": + y_prob_roc = y_prob[:, 1] if y_prob.shape[1] == 2 else y_prob roc_auc_weighted_ovr = sklearn_metrics.roc_auc_score( - y_true, y_prob, average="weighted", multi_class="ovr" - ) + y_true, y_prob_roc, average="weighted", multi_class="ovr" + ) if y_prob.shape[1] > 2 else sklearn_metrics.roc_auc_score(y_true, y_prob_roc) output["roc_auc_weighted_ovr"] = roc_auc_weighted_ovr elif metric == "accuracy": accuracy = sklearn_metrics.accuracy_score(y_true, y_pred) diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index eb6bc0133..3cb4a59af 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -1,4 +1,4 @@ -from .adacare import AdaCare, AdaCareLayer +from .adacare import AdaCare, AdaCareLayer, MultimodalAdaCare from .agent import Agent, AgentLayer from .base_model import BaseModel from .bottleneck_transformer import ( @@ -34,7 +34,7 @@ from .micron import MICRON, MICRONLayer from .mlp import MLP from .molerec import MoleRec, MoleRecLayer -from .retain import RETAIN, RETAINLayer +from .retain import MultimodalRETAIN, RETAIN, RETAINLayer from .rnn import MultimodalRNN, RNN, RNNLayer from .safedrug import SafeDrug, SafeDrugLayer from .sparcnet import DenseBlock, DenseLayer, SparcNet, TransitionLayer diff --git a/pyhealth/models/adacare.py b/pyhealth/models/adacare.py index 5e4ee806b..ab8602b77 100644 --- a/pyhealth/models/adacare.py +++ b/pyhealth/models/adacare.py @@ -5,10 +5,14 @@ from pyhealth.datasets import SampleDataset from pyhealth.processors import ( - SequenceProcessor, - NestedSequenceProcessor, - NestedFloatsProcessor, DeepNestedFloatsProcessor, + DeepNestedSequenceProcessor, + MultiHotProcessor, + NestedFloatsProcessor, + NestedSequenceProcessor, + SequenceProcessor, + TensorProcessor, + TimeseriesProcessor, ) from .base_model import BaseModel from .embedding import EmbeddingModel @@ -362,10 +366,7 @@ def __init__( hidden_dim: int = 128, **kwargs, ): - super().__init__( - dataset=dataset, - **kwargs, - ) + super().__init__(dataset=dataset) self.embedding_dim = embedding_dim self.hidden_dim = hidden_dim @@ -374,7 +375,7 @@ def __init__( raise ValueError("input_dim is automatically determined") assert len(self.label_keys) == 1, "Only one label key is supported" - + # Use EmbeddingModel for unified embedding handling self.embedding_model = EmbeddingModel(dataset, embedding_dim) # AdaCare layers for each feature @@ -390,11 +391,12 @@ def __init__( NestedSequenceProcessor, NestedFloatsProcessor, DeepNestedFloatsProcessor, + TimeseriesProcessor, ), ): raise ValueError( """AdaCare only supports SequenceProcessor, NestedSequenceProcessor, - NestedFloatsProcessor, DeepNestedFloatsProcessor.""" + NestedFloatsProcessor, DeepNestedFloatsProcessor, TimeseriesProcessor.""" ) self.adacare[feature_key] = AdaCareLayer( @@ -429,19 +431,22 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: embedded, masks = self.embedding_model(kwargs, output_mask=True) feature_importance = [] conv_feature_importance = [] - + for _, feature_key in enumerate(self.feature_keys): embeds = embedded[feature_key] mask = masks[feature_key] processor = self.dataset.input_processors[feature_key] - + if embeds.dim() == 3: - if isinstance(processor, NestedFloatsProcessor): + if isinstance(processor, (NestedFloatsProcessor, TimeseriesProcessor)): + # Both produce [batch, seq_len, num_features] masks — reduce to [batch, seq_len] mask = torch.any(mask, dim=2) elif isinstance(processor, SequenceProcessor): - pass + pass # mask already [batch, seq_len] else: - raise ValueError(f"Expected NestedFloatsProcessor or SequenceProcessor for 3D input, got {type(processor)}") + raise ValueError( + f"Expected NestedFloatsProcessor, TimeseriesProcessor, or SequenceProcessor for 3D input, got {type(processor)}" + ) elif embeds.dim() == 4: if isinstance(processor, NestedSequenceProcessor): embeds = torch.sum(embeds, dim=2) @@ -450,10 +455,14 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: embeds = torch.sum(embeds, dim=2) mask = torch.any(mask, dim=(2, 3)) else: - raise ValueError(f"Expected NestedSequenceProcessor or DeepNestedFloatsProcessor for 4D input, got {type(processor)}") + raise ValueError( + f"Expected NestedSequenceProcessor or DeepNestedFloatsProcessor for 4D input, got {type(processor)}" + ) else: - raise NotImplementedError(f"Unsupported input dimension {feature_key}: {embeds.dim()} for AdaCare") - + raise NotImplementedError( + f"Unsupported input dimension {feature_key}: {embeds.dim()} for AdaCare" + ) + embeds, _, inputatt, convatt = self.adacare[feature_key](embeds, mask) feature_importance.append(inputatt) conv_feature_importance.append(convatt) @@ -476,4 +485,259 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: } if kwargs.get("embed", False): results["embed"] = patient_emb - return results \ No newline at end of file + return results + + +class MultimodalAdaCare(BaseModel): + """Multimodal AdaCare model for mixed sequential and non-sequential features. + + This model extends AdaCare to support mixed input modalities: + - Sequential features (sequences, timeseries) go through AdaCareLayer + - Non-sequential features (multi-hot, tensor) bypass AdaCareLayer, use embeddings + + The model automatically classifies input features based on their processor types: + - Sequential processors (apply AdaCareLayer): SequenceProcessor, + NestedSequenceProcessor, DeepNestedSequenceProcessor, NestedFloatsProcessor, + DeepNestedFloatsProcessor, TimeseriesProcessor + - Non-sequential processors (embeddings only): MultiHotProcessor, TensorProcessor + + For sequential features, the model: + 1. Embeds the input using EmbeddingModel + 2. Applies AdaCareLayer with scale-adaptive feature extraction and recalibration + 3. Extracts the patient representation + + For non-sequential features, the model: + 1. Embeds the input using EmbeddingModel + 2. Applies mean pooling if needed to reduce to 2D + 3. Uses the embedding directly + + All feature representations are concatenated and passed through a final + fully connected layer for predictions. Feature importance outputs from + AdaCareLayer are preserved for sequential features. + + Args: + dataset (SampleDataset): the dataset to train the model. It is used to query + certain information such as the set of all tokens and processor types. + embedding_dim (int): the embedding dimension. Default is 128. + hidden_dim (int): the hidden dimension for AdaCare layers. Default is 128. + **kwargs: other parameters for the AdaCareLayer (e.g., kernel_size, kernel_num, + r_v, r_c, activation, rnn_type, dropout). + + Examples: + >>> from pyhealth.datasets import create_sample_dataset + >>> samples = [ + ... { + ... "patient_id": "patient-0", + ... "visit_id": "visit-0", + ... "conditions": ["cond-33", "cond-86"], # sequential + ... "demographics": ["asian", "male"], # multi-hot + ... "vitals": [120.0, 80.0, 98.6], # tensor + ... "label": 1, + ... }, + ... { + ... "patient_id": "patient-1", + ... "visit_id": "visit-1", + ... "conditions": ["cond-12", "cond-52"], # sequential + ... "demographics": ["white", "female"], # multi-hot + ... "vitals": [110.0, 75.0, 98.2], # tensor + ... "label": 0, + ... }, + ... ] + >>> dataset = create_sample_dataset( + ... samples=samples, + ... input_schema={ + ... "conditions": "sequence", + ... "demographics": "multi_hot", + ... "vitals": "tensor", + ... }, + ... output_schema={"label": "binary"}, + ... dataset_name="test" + ... ) + >>> + >>> from pyhealth.datasets import get_dataloader + >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + >>> + >>> model = MultimodalAdaCare(dataset=dataset, hidden_dim=64) + >>> + >>> data_batch = next(iter(train_loader)) + >>> + >>> ret = model(**data_batch) + >>> print(ret) + { + 'loss': tensor(...), + 'y_prob': tensor(...), + 'y_true': tensor(...), + 'logit': tensor(...), + 'feature_importance': [...], + 'conv_feature_importance': [...] + } + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: int = 128, + hidden_dim: int = 128, + **kwargs, + ): + super(MultimodalAdaCare, self).__init__(dataset=dataset) + self.embedding_dim = embedding_dim + self.hidden_dim = hidden_dim + + # validate kwargs + if "input_dim" in kwargs: + raise ValueError("input_dim is determined by embedding_dim") + + assert len(self.label_keys) == 1, "Only one label key is supported" + + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + + # Classify features as sequential or non-sequential + self.sequential_features = [] + self.non_sequential_features = [] + + self.adacare = nn.ModuleDict() + for feature_key in self.feature_keys: + processor = dataset.input_processors[feature_key] + if self._is_sequential_processor(processor): + self.sequential_features.append(feature_key) + self.adacare[feature_key] = AdaCareLayer( + input_dim=embedding_dim, hidden_dim=hidden_dim, **kwargs + ) + else: + self.non_sequential_features.append(feature_key) + + # Calculate final concatenated dimension + final_dim = ( + len(self.sequential_features) * hidden_dim + + len(self.non_sequential_features) * embedding_dim + ) + output_size = self.get_output_size() + self.fc = nn.Linear(final_dim, output_size) + + def _is_sequential_processor(self, processor) -> bool: + """Check if processor represents sequential data. + + Sequential processors are those that benefit from AdaCare processing, + including sequences of codes and timeseries data. + + Note: + StageNetProcessor and StageNetTensorProcessor are excluded as they + are specialized for the StageNet model architecture and should be + treated as non-sequential for standard AdaCare processing. + + Args: + processor: The processor instance to check. + + Returns: + bool: True if processor is sequential, False otherwise. + """ + return isinstance( + processor, + ( + SequenceProcessor, + NestedSequenceProcessor, + DeepNestedSequenceProcessor, + NestedFloatsProcessor, + DeepNestedFloatsProcessor, + TimeseriesProcessor, + ), + ) + + def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + """Forward propagation handling mixed modalities. + + Args: + **kwargs: keyword arguments for the model. The keys must contain + all the feature keys and the label key. + + Returns: + Dict[str, torch.Tensor]: A dictionary with the following keys: + - loss: a scalar tensor representing the loss. + - y_prob: a tensor representing the predicted probabilities. + - y_true: a tensor representing the true labels. + - logit: a tensor representing the logits. + - feature_importance: list of tensors representing input + feature importance for sequential features. + - conv_feature_importance: list of tensors representing + convolutional feature importance for sequential features. + - embed (optional): a tensor representing the patient + embeddings if requested. + """ + patient_emb = [] + embedded, masks = self.embedding_model(kwargs, output_mask=True) + feature_importance = [] + conv_feature_importance = [] + + # Process sequential features through AdaCare + for feature_key in self.sequential_features: + embeds = embedded[feature_key] + mask = masks[feature_key] + processor = self.dataset.input_processors[feature_key] + + # Handle different dimensions + if embeds.dim() == 3: + if isinstance(processor, (NestedFloatsProcessor, TimeseriesProcessor)): + # Both produce [batch, seq_len, num_features] masks — reduce to [batch, seq_len] + mask = torch.any(mask, dim=2) + elif isinstance(processor, SequenceProcessor): + pass # mask already [batch, seq_len] + else: + raise ValueError( + f"Expected NestedFloatsProcessor, TimeseriesProcessor, or " + f"SequenceProcessor for 3D input, got {type(processor)}" + ) + elif embeds.dim() == 4: + if isinstance(processor, NestedSequenceProcessor): + embeds = torch.sum(embeds, dim=2) + mask = torch.any(mask, dim=2) + elif isinstance(processor, DeepNestedFloatsProcessor): + embeds = torch.sum(embeds, dim=2) + mask = torch.any(mask, dim=(2, 3)) + else: + raise ValueError( + f"Expected NestedSequenceProcessor or " + f"DeepNestedFloatsProcessor for 4D input, " + f"got {type(processor)}" + ) + else: + raise NotImplementedError( + f"Unsupported input dimension {feature_key}: " + f"{embeds.dim()} for AdaCare" + ) + + # Apply AdaCare layer + embeds, _, inputatt, convatt = self.adacare[feature_key](embeds, mask) + feature_importance.append(inputatt) + conv_feature_importance.append(convatt) + patient_emb.append(embeds) + + # Process non-sequential features (use embeddings directly) + for feature_key in self.non_sequential_features: + x = embedded[feature_key] + # If multi-dimensional, aggregate (mean pooling) + while x.dim() > 2: + x = x.mean(dim=1) + patient_emb.append(x) + + # Concatenate all representations + patient_emb = torch.cat(patient_emb, dim=1) + # (patient, label_size) + logits = self.fc(patient_emb) + + # Calculate loss and predictions + y_true = kwargs[self.label_keys[0]].to(self.device) + loss = self.get_loss_function()(logits, y_true) + y_prob = self.prepare_y_prob(logits) + + results = { + "loss": loss, + "y_prob": y_prob, + "y_true": y_true, + "logit": logits, + "feature_importance": feature_importance, + "conv_feature_importance": conv_feature_importance, + } + if kwargs.get("embed", False): + results["embed"] = patient_emb + return results diff --git a/pyhealth/models/biot.py b/pyhealth/models/biot.py index cb9a45d8e..e038e092b 100644 --- a/pyhealth/models/biot.py +++ b/pyhealth/models/biot.py @@ -243,8 +243,7 @@ class BIOT(BaseModel): depth: number of transformer encoder layers. Default is 4. n_fft: number of frequency bins used in the STFT transform. Default is 200. hop_length: hop length for the STFT transform. Default is 100. - n_classes: number of output classes for classification tasks (only used for BIOTClassifier). - n_channels: number of channels in the biosignal data. Default is 18. + n_channels: number of channels in the biosignal data. Default is 18. This includes the 16 channels of the TUEV dataset and 2 additional channels for Sleep dataset. Examples: >>> from pyhealth.datasets import TUEVDataset @@ -257,7 +256,6 @@ class BIOT(BaseModel): >>> depth=4, >>> n_fft=200, >>> hop_length=100, - >>> n_classes=6, >>> n_channels=18, >>> ) >>> model.load_pretrained_weights("pretrained-models/EEG-six-datasets-18-channels.ckpt") @@ -266,24 +264,24 @@ class BIOT(BaseModel): >>> output = model(torch.randn(8, 18, TIME_STEPS)) # (batch, channels, time) """ - def __init__(self, + def __init__(self, dataset: SampleDataset, emb_size: int = 256, heads: int = 8, depth: int = 4, n_fft: int = 200, hop_length: int = 100, - n_classes: int = 6, n_channels: int = 18, **kwargs): super().__init__(dataset=dataset) _get_linear_attention_transformer() - self.biot = BIOTClassifier(emb_size=emb_size, - heads=heads, - depth=depth, - n_classes=n_classes, - n_channels=n_channels, - n_fft=n_fft, + output_size = self.get_output_size() + self.biot = BIOTClassifier(emb_size=emb_size, + heads=heads, + depth=depth, + n_classes=output_size, + n_channels=n_channels, + n_fft=n_fft, hop_length=hop_length) def load_pretrained_weights(self, checkpoint_path: str, strict: bool = False, map_location: str = None): diff --git a/pyhealth/models/cnn.py b/pyhealth/models/cnn.py index 9e69de3bd..3d5fe3ce0 100644 --- a/pyhealth/models/cnn.py +++ b/pyhealth/models/cnn.py @@ -261,6 +261,40 @@ def _determine_input_channels(self, feature_key: str, spatial_dim: int) -> int: if spatial_dim == 1: return self.embedding_dim + # For ImageProcessor, derive channel count from the configured mode so + # we don't rely on whichever sample happens to come first (which may be + # grayscale even when most images are RGB). + processor = self.dataset.input_processors.get(feature_key) + if isinstance(processor, ImageProcessor): + _mode_channels = { + "L": 1, "P": 1, "I": 1, "F": 1, + "RGB": 3, "YCbCr": 3, "LAB": 3, "HSV": 3, + "RGBA": 4, "CMYK": 4, + } + if processor.mode is not None: + return _mode_channels.get(processor.mode, 3) + # mode=None: scan dataset and take the maximum channel count seen + # so that a single grayscale sample doesn't under-count. + max_channels = 0 + min_dim = spatial_dim + 1 + for sample in self.dataset: + if feature_key not in sample: + continue + feature = self._extract_feature_tensor(sample[feature_key]) + if feature is None: + continue + tensor = self._ensure_tensor(feature) + if tensor.dim() < min_dim: + continue + max_channels = max(max_channels, tensor.shape[0]) + if max_channels >= 3: + break # RGB is the common maximum; no need to scan further + if max_channels > 0: + return max_channels + raise ValueError( + f"Unable to infer input channels for feature '{feature_key}'." + ) + min_dim = spatial_dim + 1 for sample in self.dataset: if feature_key not in sample: diff --git a/pyhealth/models/concare.py b/pyhealth/models/concare.py index 35301c315..c642ee810 100644 --- a/pyhealth/models/concare.py +++ b/pyhealth/models/concare.py @@ -21,6 +21,7 @@ """ import math +import warnings from typing import Callable, Dict, Optional, Tuple import torch @@ -30,7 +31,7 @@ from pyhealth.models import BaseModel from pyhealth.models.utils import get_last_visit -from .embedding import EmbeddingModel +from pyhealth.processors import NestedSequenceProcessor, SequenceProcessor class FinalAttentionQKV(nn.Module): @@ -73,8 +74,16 @@ def __init__( self.W_out = nn.Linear(attention_hidden_dim, 1) - self.b_in = nn.Parameter(torch.zeros(1,)) - self.b_out = nn.Parameter(torch.zeros(1,)) + self.b_in = nn.Parameter( + torch.zeros( + 1, + ) + ) + self.b_out = nn.Parameter( + torch.zeros( + 1, + ) + ) nn.init.kaiming_uniform_(self.W_q.weight, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.W_k.weight, a=math.sqrt(5)) @@ -85,7 +94,11 @@ def __init__( torch.randn(2 * attention_input_dim, attention_hidden_dim) ) self.Wa = nn.Parameter(torch.randn(attention_hidden_dim, 1)) - self.ba = nn.Parameter(torch.zeros(1,)) + self.ba = nn.Parameter( + torch.zeros( + 1, + ) + ) nn.init.kaiming_uniform_(self.Wh, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.Wa, a=math.sqrt(5)) @@ -95,9 +108,7 @@ def __init__( self.softmax = nn.Softmax(dim=1) self.sigmoid = nn.Sigmoid() - def forward( - self, input: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Forward pass of the final attention layer. Args: @@ -167,9 +178,7 @@ def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1): self.w_2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) - def forward( - self, x: torch.Tensor - ) -> Tuple[torch.Tensor, None]: + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, None]: """Forward pass of the feed-forward network. Args: @@ -293,7 +302,8 @@ def cov(self, m: torch.Tensor, y: Optional[torch.Tensor] = None) -> torch.Tensor m = torch.cat((m, y), dim=0) m_exp = torch.mean(m, dim=1) x = m - m_exp[:, None] - cov = 1 / (x.size(1) - 1) * x.mm(x.t()) + n = max(x.size(1) - 1, 1) + cov = (1 / n) * x.mm(x.t()) return cov def forward( @@ -332,9 +342,7 @@ def forward( query, key, value, mask=mask, dropout=self.dropout ) - x = ( - x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k) - ) + x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k) # Compute DeCov loss decov_contexts = x.transpose(0, 1).transpose(1, 2) @@ -396,10 +404,9 @@ def __init__(self, size: int, dropout: float): def forward( self, - x: torch.Tensor, + x: torch.Tensor, sublayer: Callable[[torch.Tensor], Tuple[torch.Tensor, any]], - ) -> Tuple[torch.Tensor, any]: - + ) -> Tuple[torch.Tensor, any]: """Apply residual connection to sublayer with same size. Args: @@ -447,9 +454,7 @@ def __init__( self.Wx = nn.Parameter( torch.randn(attention_input_dim, attention_hidden_dim) ) - self.Wtime_aware = nn.Parameter( - torch.randn(1, attention_hidden_dim) - ) + self.Wtime_aware = nn.Parameter(torch.randn(1, attention_hidden_dim)) nn.init.kaiming_uniform_(self.Wtime_aware, a=math.sqrt(5)) else: self.Wx = nn.Parameter( @@ -458,9 +463,17 @@ def __init__( self.Wt = nn.Parameter( torch.randn(attention_input_dim, attention_hidden_dim) ) - self.bh = nn.Parameter(torch.zeros(attention_hidden_dim,)) + self.bh = nn.Parameter( + torch.zeros( + attention_hidden_dim, + ) + ) self.Wa = nn.Parameter(torch.randn(attention_hidden_dim, 1)) - self.ba = nn.Parameter(torch.zeros(1,)) + self.ba = nn.Parameter( + torch.zeros( + 1, + ) + ) nn.init.kaiming_uniform_(self.Wx, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.Wt, a=math.sqrt(5)) @@ -470,7 +483,11 @@ def __init__( self.Wa = nn.Parameter( torch.randn(attention_input_dim, attention_input_dim) ) - self.ba = nn.Parameter(torch.zeros(1,)) + self.ba = nn.Parameter( + torch.zeros( + 1, + ) + ) nn.init.kaiming_uniform_(self.Wa, a=math.sqrt(5)) elif attention_type == "concat": @@ -483,7 +500,11 @@ def __init__( torch.randn(2 * attention_input_dim, attention_hidden_dim) ) self.Wa = nn.Parameter(torch.randn(attention_hidden_dim, 1)) - self.ba = nn.Parameter(torch.zeros(1,)) + self.ba = nn.Parameter( + torch.zeros( + 1, + ) + ) nn.init.kaiming_uniform_(self.Wh, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.Wa, a=math.sqrt(5)) @@ -888,13 +909,9 @@ def __init__( if "input_dim" in kwargs: raise ValueError("input_dim is determined by embedding_dim") - assert len(self.label_keys) == 1, ( - "Only one label key is supported for ConCare" - ) + assert len(self.label_keys) == 1, "Only one label key is supported for ConCare" self.label_key = self.label_keys[0] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) - # Determine static dimension self.static_dim = 0 if self.static_key is not None: @@ -912,15 +929,37 @@ def __init__( # Get dynamic feature keys (excluding static key) self.dynamic_feature_keys = [ - k for k in self.feature_keys - if k != self.static_key + k for k in self.feature_keys if k != self.static_key ] # ConCare layers for each dynamic feature self.concare = nn.ModuleDict() + self._embeddings = nn.ModuleDict() + self._proc_types = {} # feature_key -> 'sequence' | 'nested_sequence' | 'other' for feature_key in self.dynamic_feature_keys: + proc = dataset.input_processors[feature_key] + if isinstance(proc, SequenceProcessor): + input_dim = 1 + self._proc_types[feature_key] = "sequence" + elif isinstance(proc, NestedSequenceProcessor): + vocab_size = proc.vocab_size() + self._embeddings[feature_key] = nn.Embedding( + vocab_size, 1, padding_idx=0 + ) + input_dim = 1 + self._proc_types[feature_key] = "nested_sequence" + warnings.warn( + f"ConCare: feature '{feature_key}' uses nested_sequence " + f"(vocab_size={vocab_size}). Codes will be embedded (dim=1) " + "and sum-pooled per visit for memory efficiency. This reduces " + "per-visit code representations to a single scalar.", + UserWarning, + ) + else: + input_dim = proc.size() + self._proc_types[feature_key] = "other" self.concare[feature_key] = ConCareLayer( - input_dim=embedding_dim, + input_dim=input_dim, static_dim=self.static_dim, hidden_dim=self.hidden_dim, **kwargs, @@ -930,7 +969,7 @@ def __init__( self.fc = nn.Linear( len(self.dynamic_feature_keys) * self.hidden_dim, output_size ) - + def forward(self, **kwargs) -> Dict[str, torch.Tensor]: """Forward propagation. @@ -944,12 +983,10 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: - y_prob: a tensor representing the predicted probabilities. - y_true: a tensor representing the true labels. - logit: a tensor representing the logits. - """ + """ patient_emb = [] decov_loss = 0 - embedded, masks = self.embedding_model(kwargs, output_mask=True) - # Get static features if available static = None if self.static_key is not None and self.static_key in kwargs: @@ -962,8 +999,22 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: ) for feature_key in self.dynamic_feature_keys: - x = embedded[feature_key] - mask = masks[feature_key] + raw = kwargs[feature_key] + if self._proc_types[feature_key] == "sequence": + # (batch, T) long → (batch, T, 1) float; PAD=0 + mask = (raw != 0).float().to(self.device) + x = raw.float().unsqueeze(-1).to(self.device) + elif self._proc_types[feature_key] == "nested_sequence": + # raw: (batch, visits, max_codes) long + # embed each code to scalar, sum-pool over codes per visit + emb = self._embeddings[feature_key](raw.to(self.device)) + # emb: (batch, visits, max_codes, 1) + x = emb.sum(dim=2).float() # (batch, visits, 1) + mask = (raw != 0).any(dim=-1).float().to(self.device) + else: + # (batch, T, F) → mask on time axis (all-zero rows are padding) + mask = (raw.sum(-1) != 0).float().to(self.device) + x = raw.float().to(self.device) x, decov = self.concare[feature_key](x, static=static, mask=mask) patient_emb.append(x) diff --git a/pyhealth/models/contrawr.py b/pyhealth/models/contrawr.py index f3b9e58aa..fe359beb5 100644 --- a/pyhealth/models/contrawr.py +++ b/pyhealth/models/contrawr.py @@ -155,6 +155,9 @@ def __init__( assert len(self.label_keys) == 1, ( "Only one label key is supported if ContraWR is initialized" ) + # ContraWR computes its own STFT internally — drop any precomputed + # auxiliary keys (e.g. "stft") that TFMTokenizer adds to input_schema. + self.feature_keys = [k for k in self.feature_keys if k != "stft"] assert len(self.feature_keys) == 1, ( "Only one feature key is supported if ContraWR is initialized" ) diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index ddcee693d..703f77630 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -216,6 +216,8 @@ def __init__( self.encoders: nn.ModuleDict = nn.ModuleDict() self.projections: nn.ModuleDict = nn.ModuleDict() self.modality_types: dict[str, ModalityType] = {} + self._shared_text_field_by_model: dict[str, str] = {} + self._text_canonical: dict[str, str] = {} # field → first field sharing the same tokenizer for field_name, processor in processors.items(): if not isinstance(processor, TemporalFeatureProcessor): @@ -311,6 +313,21 @@ def _build_text_encoder( freeze: bool = False, ) -> None: """Build TEXT encoder: BERT + projection, optionally from TextEmbeddingModel.""" + + def _set_projection( + pre_dim: int, proj_source: Optional[nn.Module] = None + ) -> None: + if pre_dim != embedding_dim: + if proj_source is not None: + self.projections[field_name] = nn.Sequential( + proj_source, + nn.Linear(pre_dim, embedding_dim), + ) + else: + self.projections[field_name] = nn.Linear(pre_dim, embedding_dim) + elif proj_source is not None: + self.projections[field_name] = proj_source + if ( pre_built is not None and hasattr(pre_built, "transformer") @@ -321,26 +338,35 @@ def _build_text_encoder( for p in pre_built.transformer.parameters(): p.requires_grad = False pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) - if pre_dim != embedding_dim: - self.projections[field_name] = nn.Sequential( - pre_built.fc, - nn.Linear(pre_dim, embedding_dim), - ) - else: - self.projections[field_name] = pre_built.fc + _set_projection(pre_dim, pre_built.fc) return if processor.is_token(): from transformers import AutoModel - bert = AutoModel.from_pretrained(processor.tokenizer_model) - if freeze: - for p in bert.parameters(): - p.requires_grad = False - self.encoders[field_name] = bert - hidden = bert.config.hidden_size - if hidden != embedding_dim: - self.projections[field_name] = nn.Linear(hidden, embedding_dim) + tokenizer_model = getattr(processor, "tokenizer_model", None) + if not tokenizer_model: + raise ValueError( + f"TEXT processor '{field_name}' is token-based but does not " + "define tokenizer_model." + ) + + shared_field = self._shared_text_field_by_model.get(tokenizer_model) + if shared_field is not None: + # Second+ field with same tokenizer: reuse existing encoder, do NOT + # register under a new key (avoids duplicate parameter registration). + self._text_canonical[field_name] = shared_field + 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 + + hidden = shared_encoder.config.hidden_size + _set_projection(hidden) else: raise ValueError( f"TEXT processor '{field_name}' must either supply a pre-built " @@ -443,7 +469,8 @@ def forward( time = torch.zeros(value.shape[:2], device=value.device) modality = self.modality_types[field_name] - encoder = self.encoders[field_name] + encoder_key = self._text_canonical.get(field_name, field_name) + encoder = self.encoders[encoder_key] # ── Encode ──────────────────────────────────────────────────── if modality == ModalityType.CODE: diff --git a/pyhealth/models/embedding/vanilla.py b/pyhealth/models/embedding/vanilla.py index d4fbfa9ba..08a90f532 100644 --- a/pyhealth/models/embedding/vanilla.py +++ b/pyhealth/models/embedding/vanilla.py @@ -111,7 +111,7 @@ def init_embedding_with_pretrained( return loaded -class EmbeddingModel(BaseModel, BaseEmbeddingModel): +class EmbeddingModel(BaseModel): """ EmbeddingModel is responsible for creating embedding layers for different types of input data. @@ -164,14 +164,16 @@ def __init__( SequenceProcessor, StageNetProcessor, NestedSequenceProcessor, - DeepNestedSequenceProcessor + DeepNestedSequenceProcessor, ), ): vocab_size = len(processor.code_vocab) # For NestedSequenceProcessor and DeepNestedSequenceProcessor, don't use padding_idx # because empty visits/groups need non-zero embeddings. - if isinstance(processor, (NestedSequenceProcessor, DeepNestedSequenceProcessor)): + if isinstance( + processor, (NestedSequenceProcessor, DeepNestedSequenceProcessor) + ): self.embedding_layers[field_name] = nn.Embedding( num_embeddings=vocab_size, embedding_dim=embedding_dim, @@ -242,16 +244,23 @@ def __init__( try: from transformers import AutoModel except ImportError: - raise ImportError("Please install `transformers` to use token-based processors.") + raise ImportError( + "Please install `transformers` to use token-based processors." + ) # Load the model - self.embedding_layers[field_name] = AutoModel.from_pretrained(processor.tokenizer_model) + self.embedding_layers[field_name] = AutoModel.from_pretrained( + processor.tokenizer_model + ) # Check if we need projection - if self.embedding_layers[field_name].config.hidden_size != self.embedding_dim: + if ( + self.embedding_layers[field_name].config.hidden_size + != self.embedding_dim + ): self.embedding_layers[f"{field_name}_proj"] = nn.Linear( self.embedding_layers[field_name].config.hidden_size, - self.embedding_dim + self.embedding_dim, ) else: @@ -260,15 +269,15 @@ def __init__( field_name, ) - @property - def embedding_dim(self) -> int: - return self._embedding_dim - - def forward(self, - inputs: Dict[str, torch.Tensor], - masks: Dict[str, torch.Tensor] = None, - output_mask: bool = False - ) -> Dict[str, torch.Tensor] | tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: + def forward( + self, + inputs: Dict[str, torch.Tensor], + masks: Dict[str, torch.Tensor] = None, + output_mask: bool = False, + ) -> ( + Dict[str, torch.Tensor] + | tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]] + ): embedded: Dict[str, torch.Tensor] = {} out_masks: Dict[str, torch.Tensor] = {} if output_mask else None @@ -288,14 +297,14 @@ def forward(self, # or check if it has 'config' attribute if hasattr(layer, "config") and hasattr(layer, "forward"): # It's likely a transformer - tensor = tensor.to(self.device).long() # Ensure LongTensor for IDs + tensor = tensor.to(self.device).long() # Ensure LongTensor for IDs mask = None if masks is not None and field_name in masks: mask = masks[field_name].to(self.device) # Handle 3D input (Batch, Num_Notes, Seq_Len) - is_3d = (inputs[field_name].dim() == 3) + is_3d = inputs[field_name].dim() == 3 if is_3d: b, n, l = inputs[field_name].shape @@ -308,16 +317,21 @@ def forward(self, x = output.last_hidden_state # (Batch, Seq, Hidden) if is_3d: - # Pool the sequence dim (L) to one vector per note using CLS token (index 0) - x = x[:, 0, :] # (B*N, H) + # If we had 3D input, we MUST pool the sequence dim (L) to get one vector per note + # Resulting shape: (B, N, H) + + # Pool L dim -> (B*N, H) using CLS token (index 0) + x = x[:, 0, :] + # Check projections if f"{field_name}_proj" in self.embedding_layers: x = self.embedding_layers[f"{field_name}_proj"](x) x = x.view(b, n, -1) else: - # 2D input (Batch, Seq) -> (Batch, Seq, Hidden): token-level embeddings + # 2D input (Batch, Seq) -> (Batch, Seq, Hidden) + # No pooling, treating as sequence of tokens (word embeddings) if f"{field_name}_proj" in self.embedding_layers: x = self.embedding_layers[f"{field_name}_proj"](x) @@ -334,10 +348,14 @@ def forward(self, out_masks[field_name] = masks[field_name].to(self.device) elif hasattr(processor, "code_vocab"): pad_idx = processor.code_vocab.get("", 0) - out_masks[field_name] = (tensor != pad_idx) + out_masks[field_name] = tensor != pad_idx else: + # Default mask generation (e.g. for simple linear layers where 0 might be padding?) + # Be careful changing this behavior. + # Previous code: + # masks[field_name] = (tensor != pad_idx) -> where pad_idx was 0 default pad_idx = 0 - out_masks[field_name] = (tensor != pad_idx) + out_masks[field_name] = tensor != pad_idx if output_mask: return embedded, out_masks diff --git a/pyhealth/models/embedding/vision.py b/pyhealth/models/embedding/vision.py index b1fef865e..57eedda22 100644 --- a/pyhealth/models/embedding/vision.py +++ b/pyhealth/models/embedding/vision.py @@ -116,11 +116,13 @@ def __init__( freeze_backbone: bool = False, dropout: float = 0.0, use_cls_token: bool = False, + pool: Optional[Literal["mean"]] = None, ) -> None: super().__init__(dataset) self._embedding_dim = embedding_dim self.patch_size = patch_size + self.pool = pool self.backbone_type = backbone self.use_cls_token = use_cls_token @@ -286,6 +288,9 @@ def forward( x = x + self.pos_embeddings[field_name] x = self.dropout(x) + if self.pool == "mean": + x = x.mean(dim=1, keepdim=True) + embedded[field_name] = x if output_mask: @@ -303,7 +308,10 @@ def get_output_info(self, field_name: str) -> Dict[str, Any]: info = self._field_info[field_name].copy() info["embedding_dim"] = self._embedding_dim info["has_cls_token"] = self.use_cls_token - info["num_tokens"] = info["num_patches"] + (1 if self.use_cls_token else 0) + if self.pool == "mean": + info["num_tokens"] = 1 + else: + info["num_tokens"] = info["num_patches"] + (1 if self.use_cls_token else 0) return info def __repr__(self) -> str: @@ -350,9 +358,23 @@ def __repr__(self) -> str: use_cls_token=True, ) + model_pooled = VisionEmbeddingModel( + dataset=dataset, + embedding_dim=128, + backbone="cnn", + pool="mean", + ) + + + loader = get_dataloader(dataset, batch_size=4, shuffle=False) batch = next(iter(loader)) + embeddings_pooled = model_pooled({"chest_xray": batch["chest_xray"]}) + print(f"Pooled output shape: {embeddings_pooled['chest_xray'].shape}") # expect (4, 1, 128) + print(f"Pooled output info: {model_pooled.get_output_info('chest_xray')}") # expect num_tokens=1 + + embeddings = model({"chest_xray": batch["chest_xray"]}) print(f"Input shape: {batch['chest_xray'].shape}") print(f"Output shape: {embeddings['chest_xray'].shape}") diff --git a/pyhealth/models/retain.py b/pyhealth/models/retain.py index b272995ec..e871c88d0 100644 --- a/pyhealth/models/retain.py +++ b/pyhealth/models/retain.py @@ -6,6 +6,16 @@ from pyhealth.datasets import SampleDataset from pyhealth.models import BaseModel +from pyhealth.processors import ( + DeepNestedFloatsProcessor, + DeepNestedSequenceProcessor, + MultiHotProcessor, + NestedFloatsProcessor, + NestedSequenceProcessor, + SequenceProcessor, + TensorProcessor, + TimeseriesProcessor, +) from .embedding import EmbeddingModel @@ -56,23 +66,27 @@ def reverse_x(input, lengths): reversed_input[i, :length] = input[i, :length].flip(dims=[0]) return reversed_input - def compute_alpha(self, rx, lengths): + def compute_alpha(self, rx, lengths, total_length: int): """Computes alpha attention.""" rx = rnn_utils.pack_padded_sequence( rx, lengths, batch_first=True, enforce_sorted=False ) g, _ = self.alpha_gru(rx) - g, _ = rnn_utils.pad_packed_sequence(g, batch_first=True) + g, _ = rnn_utils.pad_packed_sequence( + g, batch_first=True, total_length=total_length + ) attn_alpha = torch.softmax(self.alpha_li(g), dim=1) return attn_alpha - def compute_beta(self, rx, lengths): + def compute_beta(self, rx, lengths, total_length: int): """Computes beta attention.""" rx = rnn_utils.pack_padded_sequence( rx, lengths, batch_first=True, enforce_sorted=False ) h, _ = self.beta_gru(rx) - h, _ = rnn_utils.pad_packed_sequence(h, batch_first=True) + h, _ = rnn_utils.pad_packed_sequence( + h, batch_first=True, total_length=total_length + ) attn_beta = torch.tanh(self.beta_li(h)) return attn_beta @@ -95,15 +109,17 @@ def forward( # rnn will only apply dropout between layers x = self.dropout_layer(x) batch_size = x.size(0) + total_length = x.size(1) # capture before packing so pad_packed restores it if mask is None: lengths = torch.full( - size=(batch_size,), fill_value=x.size(1), dtype=torch.int64 + size=(batch_size,), fill_value=total_length, dtype=torch.int64 ) else: lengths = torch.sum(mask.int(), dim=-1).cpu() + lengths = lengths.clamp(min=1) # prevent zero-length crash in GRU rx = self.reverse_x(x, lengths) - attn_alpha = self.compute_alpha(rx, lengths) - attn_beta = self.compute_beta(rx, lengths) + attn_alpha = self.compute_alpha(rx, lengths, total_length) + attn_beta = self.compute_beta(rx, lengths, total_length) c = attn_alpha * attn_beta * x # (patient, sequence len, feature_size) c = torch.sum(c, dim=1) # (patient, feature_size) return c @@ -201,15 +217,11 @@ def __init__( # Create RETAIN layers for each feature self.retain = nn.ModuleDict() for feature_key in self.feature_keys: - self.retain[feature_key] = RETAINLayer( - feature_size=embedding_dim, **kwargs - ) + self.retain[feature_key] = RETAINLayer(feature_size=embedding_dim, **kwargs) output_size = self.get_output_size() num_features = len(self.feature_keys) - self.fc = nn.Linear( - num_features * self.embedding_dim, output_size - ) + self.fc = nn.Linear(num_features * self.embedding_dim, output_size) def forward(self, **kwargs) -> Dict[str, torch.Tensor]: """Forward propagation. @@ -228,30 +240,29 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: """ patient_emb = [] embedded = self.embedding_model(kwargs) - + for feature_key in self.feature_keys: x = embedded[feature_key] - + # Handle different input dimensions # Case 1: 4D tensor from NestedSequenceProcessor # (batch, visits, events, embedding_dim) # Need to sum across events to get (batch, visits, embedding_dim) if len(x.shape) == 4: x = torch.sum(x, dim=2) # Sum across events within visit - + # Case 2: 3D tensor from SequenceProcessor or after summing # (batch, seq_len, embedding_dim) - already correct format elif len(x.shape) == 3: pass # Already correct format - + # Case 3: 2D tensor - shouldn't happen for RETAIN but handle it elif len(x.shape) == 2: x = x.unsqueeze(1) # Add seq dim: (batch, 1, embedding_dim) - + else: raise ValueError( - f"Unexpected tensor shape {x.shape} for feature " - f"{feature_key}" + f"Unexpected tensor shape {x.shape} for feature " f"{feature_key}" ) # Create mask: non-padding entries are valid @@ -331,3 +342,233 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: # try loss backward ret["loss"].backward() + + +class MultimodalRETAIN(BaseModel): + """Multimodal RETAIN model for mixed sequential and non-sequential features. + + This model extends RETAIN to support mixed input modalities: + - Sequential features (sequences, timeseries) go through RETAINLayer + - Non-sequential features (multi-hot, tensor) bypass RETAIN, use embeddings directly + + The model automatically classifies input features based on their processor types: + - Sequential processors (apply RETAINLayer): SequenceProcessor, + NestedSequenceProcessor, DeepNestedSequenceProcessor, NestedFloatsProcessor, + DeepNestedFloatsProcessor, TimeseriesProcessor + - Non-sequential processors (embeddings only): MultiHotProcessor, TensorProcessor + + For sequential features, the model: + 1. Embeds the input using EmbeddingModel + 2. Applies RETAINLayer with reverse time attention mechanism + 3. Extracts the patient representation + + For non-sequential features, the model: + 1. Embeds the input using EmbeddingModel + 2. Applies mean pooling if needed to reduce to 2D + 3. Uses the embedding directly + + All feature representations are concatenated and passed through a final + fully connected layer for predictions. + + Args: + dataset (SampleDataset): the dataset to train the model. It is used to query + certain information such as the set of all tokens and processor types. + embedding_dim (int): the embedding dimension. Default is 128. + **kwargs: other parameters for the RETAIN layer (e.g., dropout). + + Examples: + >>> from pyhealth.datasets import create_sample_dataset + >>> samples = [ + ... { + ... "patient_id": "patient-0", + ... "visit_id": "visit-0", + ... "conditions": [["A", "B"], ["C"]], # nested sequence + ... "demographics": ["asian", "male"], # multi-hot + ... "vitals": [110.0, 75.0, 98.2], # tensor + ... "label": 1, + ... }, + ... { + ... "patient_id": "patient-1", + ... "visit_id": "visit-1", + ... "conditions": [["D"], ["E", "F"]], # nested sequence + ... "demographics": ["white", "female"], # multi-hot + ... "vitals": [120.0, 80.0, 98.6], # tensor + ... "label": 0, + ... }, + ... ] + >>> dataset = create_sample_dataset( + ... samples=samples, + ... input_schema={ + ... "conditions": "nested_sequence", + ... "demographics": "multi_hot", + ... "vitals": "tensor", + ... }, + ... output_schema={"label": "binary"}, + ... dataset_name="test" + ... ) + >>> + >>> from pyhealth.datasets import get_dataloader + >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + >>> + >>> model = MultimodalRETAIN(dataset=dataset) + >>> + >>> data_batch = next(iter(train_loader)) + >>> + >>> ret = model(**data_batch) + >>> print(ret) + { + 'loss': tensor(...), + 'y_prob': tensor(...), + 'y_true': tensor(...), + 'logit': tensor(...) + } + """ + + def __init__(self, dataset: SampleDataset, embedding_dim: int = 128, **kwargs): + super(MultimodalRETAIN, self).__init__(dataset=dataset) + self.embedding_dim = embedding_dim + + # validate kwargs for RETAIN layer + if "feature_size" in kwargs: + raise ValueError("feature_size is determined by embedding_dim") + + assert len(self.label_keys) == 1, "Only one label key is supported" + self.label_key = self.label_keys[0] + self.mode = self.dataset.output_schema[self.label_key] + + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + + # Classify features as sequential or non-sequential + self.sequential_features = [] + self.non_sequential_features = [] + + self.retain = nn.ModuleDict() + for feature_key in self.feature_keys: + processor = dataset.input_processors[feature_key] + if self._is_sequential_processor(processor): + self.sequential_features.append(feature_key) + # Create RETAIN layer for this feature + self.retain[feature_key] = RETAINLayer( + feature_size=embedding_dim, **kwargs + ) + else: + self.non_sequential_features.append(feature_key) + + # Calculate final concatenated dimension + final_dim = ( + len(self.sequential_features) * embedding_dim + + len(self.non_sequential_features) * embedding_dim + ) + output_size = self.get_output_size() + self.fc = nn.Linear(final_dim, output_size) + + def _is_sequential_processor(self, processor) -> bool: + """Check if processor represents sequential data. + + Sequential processors are those that benefit from RETAIN processing, + including sequences of codes and timeseries data. + + Note: + StageNetProcessor and StageNetTensorProcessor are excluded as they + are specialized for the StageNet model architecture and should be + treated as non-sequential for standard RETAIN processing. + + Args: + processor: The processor instance to check. + + Returns: + bool: True if processor is sequential, False otherwise. + """ + return isinstance( + processor, + ( + SequenceProcessor, + NestedSequenceProcessor, + DeepNestedSequenceProcessor, + NestedFloatsProcessor, + DeepNestedFloatsProcessor, + TimeseriesProcessor, + ), + ) + + def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + """Forward propagation handling mixed modalities. + + Args: + **kwargs: keyword arguments for the model. The keys must contain + all the feature keys and the label key. + + Returns: + Dict[str, torch.Tensor]: A dictionary with the following keys: + - loss: a scalar tensor representing the loss. + - y_prob: a tensor representing the predicted probabilities. + - y_true: a tensor representing the true labels. + - logit: a tensor representing the logits. + - embed (optional): a tensor representing the patient + embeddings if requested. + """ + patient_emb = [] + embedded, emb_masks = self.embedding_model(kwargs, output_mask=True) + + # Process sequential features through RETAIN + for feature_key in self.sequential_features: + x = embedded[feature_key] + + # Handle different input dimensions + # Case 1: 4D tensor from NestedSequenceProcessor + # (batch, visits, events, embedding_dim) + # Need to sum across events to get (batch, visits, embedding_dim) + if x.dim() == 4: + x = torch.sum(x, dim=2) # Sum across events within visit + + # Case 2: 3D tensor from SequenceProcessor or after summing + # (batch, seq_len, embedding_dim) - already correct format + elif x.dim() == 3: + pass # Already correct format + + # Case 3: 2D tensor - shouldn't happen for RETAIN but handle it + elif x.dim() == 2: + x = x.unsqueeze(1) # Add seq dim: (batch, 1, embedding_dim) + + else: + raise ValueError( + f"Unexpected tensor shape {x.shape} for feature {feature_key}" + ) + + # Use mask from EmbeddingModel (derived from original unembedded tensor) + mask = emb_masks.get(feature_key) + if mask is not None: + # Ensure 2D (batch, seq_len) — reduce any extra dims + while mask.dim() > 2: + mask = mask.any(dim=-1) + mask = mask.float() + x = self.retain[feature_key](x, mask) + patient_emb.append(x) + + # Process non-sequential features (use embeddings directly) + for feature_key in self.non_sequential_features: + x = embedded[feature_key] + # If multi-dimensional, aggregate (mean pooling) + while x.dim() > 2: + x = x.mean(dim=1) + patient_emb.append(x) + + # Concatenate all representations + patient_emb = torch.cat(patient_emb, dim=1) + # (patient, label_size) + logits = self.fc(patient_emb) + + # Calculate loss and predictions + y_true = kwargs[self.label_key].to(self.device) + loss = self.get_loss_function()(logits, y_true) + y_prob = self.prepare_y_prob(logits) + + results = { + "loss": loss, + "y_prob": y_prob, + "y_true": y_true, + "logit": logits, + } + if kwargs.get("embed", False): + results["embed"] = patient_emb + return results diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py index 1712bed62..b131628a9 100644 --- a/pyhealth/models/rnn.py +++ b/pyhealth/models/rnn.py @@ -110,11 +110,16 @@ def forward( ) else: lengths = torch.sum(mask.int(), dim=-1).cpu() + # Ensure tensor is contiguous for cuDNN compatibility + x = x.contiguous() x = rnn_utils.pack_padded_sequence( x.contiguous(), lengths, batch_first=True, enforce_sorted=False ) outputs, _ = self.rnn(x) outputs, _ = rnn_utils.pad_packed_sequence(outputs, batch_first=True) + # Ensure outputs are contiguous after unpacking + outputs = outputs.contiguous() + if not self.bidirectional: last_outputs = outputs[torch.arange(batch_size), (lengths - 1), :] return outputs, last_outputs @@ -123,7 +128,8 @@ def forward( f_last_outputs = outputs[torch.arange(batch_size), (lengths - 1), 0, :] b_last_outputs = outputs[:, 0, 1, :] last_outputs = torch.cat([f_last_outputs, b_last_outputs], dim=-1) - outputs = outputs.view(batch_size, outputs.shape[1], -1) + # Ensure view result is contiguous for cuDNN + outputs = outputs.view(batch_size, outputs.shape[1], -1).contiguous() last_outputs = self.down_projection(last_outputs) outputs = self.down_projection(outputs) return outputs, last_outputs @@ -327,22 +333,28 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: for feature_key in self.feature_keys: x = embedded[feature_key] - # Use abs() before sum to catch edge cases where embeddings sum to 0 - # @TODO bug with 0 embedding sum can still persist if the embedding is all 0s but the mask is not all 0s. - # despite being valid values (e.g., [1.0, -1.0]) - - # If we have an explicit mask, use it - if feature_key in masks: - mask = masks[feature_key].to(self.device).int() - # Token-level mask (B, N_notes, L): reduce to note-level (B, N_notes) - # by checking whether each note has at least one valid token. - # This is needed when TupleTimeTextProcessor returns 3D token masks that - # EmbeddingModel has already pooled down to (B, N_notes, H). - if mask.dim() == 3: - mask = (mask.sum(dim=-1) > 0).int() # (B, N_notes) + + x_dim_orig = x.dim() + if x_dim_orig == 4: + # nested_sequence: (B, num_visits, num_codes, D) + # @TODO: sum-pooling across codes is a simple baseline. May need to investigate better embeddings for nested codes. + x = x.sum(dim=2) # (B, num_visits, D) + if feature_key in masks: + mask = (masks[feature_key].to(self.device).sum(dim=-1) > 0).int() # (B, V) + else: + mask = (torch.abs(x).sum(dim=-1) != 0).int() + elif x_dim_orig == 2: + x = x.unsqueeze(1) + mask = None else: - mask = (torch.abs(x).sum(dim=-1) != 0).int() - + # 3D: already (B, T, D) + if feature_key in masks: + mask = masks[feature_key].to(self.device).int() + if mask.dim() == 3: + mask = (mask.sum(dim=-1) > 0).int() + else: + mask = (torch.abs(x).sum(dim=-1) != 0).int() + _, x = self.rnn[feature_key](x, mask) patient_emb.append(x) @@ -559,6 +571,27 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: for feature_key in self.sequential_features: x = embedded[feature_key] m = mask[feature_key] + + x_dim_orig = x.dim() + if x_dim_orig == 4: + # nested_sequence: (B, num_visits, num_codes, D) + # Pool codes within each visit, then run RNN over visits. + # Flattening visits*codes would produce length=0 when inner lists are empty. + x = x.sum(dim=2) # (B, num_visits, D) + if feature_key in masks: + m = (masks[feature_key].to(self.device).sum(dim=-1) > 0).int() # (B, V) + else: + m = (torch.abs(x).sum(dim=-1) != 0).int() + elif x_dim_orig == 2: + x = x.unsqueeze(1) + m = None + else: + # 3D: already (B, T, D) + if m is not None and m.dim() == 3: + m = (m.sum(dim=-1) > 0).int() + elif m is not None and m.dim() == 1: + m = m.unsqueeze(1) + _, last_hidden = self.rnn[feature_key](x, m) patient_emb.append(last_hidden) diff --git a/pyhealth/models/tfm_tokenizer.py b/pyhealth/models/tfm_tokenizer.py index 4e9f1559b..38bd6de04 100644 --- a/pyhealth/models/tfm_tokenizer.py +++ b/pyhealth/models/tfm_tokenizer.py @@ -828,11 +828,8 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: y_true = kwargs[label_key].to(self.device) # Reshape tokens to (B, C, T) for multi-channel classifier - # tokens shape: (B, T) -> (B, 1, T) - logits = self.classifier(tokens_reshaped,num_ch=C) + logits = self.classifier(tokens_reshaped, num_ch=C) loss_fn = self.get_loss_function() - print(f"logits shape: {logits.shape}") - print(f"y_true shape: {y_true.shape}") cls_loss = loss_fn(logits, y_true) total_loss = recon_loss + vq_loss + cls_loss y_prob = self.prepare_y_prob(logits) @@ -849,6 +846,11 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: else: results["loss"] = recon_loss + vq_loss + if kwargs.get("embed", False): + # Mean-pool over channels (C) and time steps (T) → (B, emb_size) + # quant_out_reshaped: (B, C, T, E) + results["embed"] = quant_out_reshaped.mean(dim=(1, 2)) + return results def get_embeddings(self, dataloader) -> torch.Tensor: @@ -930,7 +932,25 @@ def load_pretrained_weights( self.tokenizer.load_state_dict(torch.load(tokenizer_checkpoint_path, map_location=map_location, weights_only=True), strict=strict) if classifier_checkpoint_path is not None and not is_masked_training: - self.classifier.load_state_dict(torch.load(classifier_checkpoint_path, map_location=map_location, weights_only=True)) + ckpt = torch.load(classifier_checkpoint_path, map_location=map_location, weights_only=True) + model_n = self.classifier.classification_head.weight.shape[0] + ckpt_n = ckpt["classification_head.weight"].shape[0] + if ckpt_n != model_n: + if ckpt_n == 1 and model_n == 2: + # Checkpoint was trained with 1-class sigmoid BCE; model expects + # 2-class softmax. The conversion is exact: + # softmax([-logit, logit]) = [1-sigmoid(logit), sigmoid(logit)] + w = ckpt["classification_head.weight"] # [1, D] + b = ckpt["classification_head.bias"] # [1] + ckpt["classification_head.weight"] = torch.cat([-w, w], dim=0) + ckpt["classification_head.bias"] = torch.cat([-b, b], dim=0) + print(f" ℹ Adapted classifier head from 1-class sigmoid → 2-class softmax") + else: + raise RuntimeError( + f"Classifier head shape mismatch: checkpoint has {ckpt_n} class(es) " + f"but model expects {model_n}. Cannot auto-adapt." + ) + self.classifier.load_state_dict(ckpt) print(f"✓ Successfully loaded weights from {classifier_checkpoint_path}") elif is_masked_training: load_embedding_weights(self.tokenizer, self.classifier) diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py index 1a4fef6c9..f678403b2 100644 --- a/pyhealth/models/transformer.py +++ b/pyhealth/models/transformer.py @@ -52,10 +52,12 @@ def forward( scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query.size(-1)) if mask is not None: - scores = scores.masked_fill(mask == 0, -1e9) + # Use -inf so softmax produces exact zeros on padded positions, + # avoiding a second masked_fill after softmax (saves one full + # [B, H, S, S] boolean allocation and an extra copy). + pad_mask = mask == 0 + scores = scores.masked_fill(pad_mask, -1e9) p_attn = self.softmax(scores) - if mask is not None: - p_attn = p_attn.masked_fill(mask == 0, 0) if dropout is not None: p_attn = dropout(p_attn) @@ -116,7 +118,7 @@ def get_attn_grad(self) -> Optional[torch.Tensor]: def save_attn_grad(self, attn_grad: torch.Tensor) -> None: """Hook callback that stores attention gradients.""" - self.attn_gradients = attn_grad + self.attn_gradients = attn_grad.detach() def forward( self, @@ -152,12 +154,18 @@ def forward( mask = mask.unsqueeze(1) x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout) - self.attn_map = attn # save the attention map if register_hook: + # Only store attn_map and hook during interpretability passes. + # Using .detach() gives an independent copy whose storage + # is NOT shared with the live graph, so the graph can be freed + # normally after .backward() without leaking GPU memory. + self.attn_map = attn.detach() attn.register_hook(self.save_attn_grad) + else: + self.attn_map = None # 3) "Concat" using a view and apply a final linear. x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) - + return self.output_linear(x) @@ -239,7 +247,7 @@ def set_activation_hooks(self, hooks) -> None: """Deprecated compatibility stub; no-op.""" return None - def forward(self, x, mask=None, register_hook = False): + def forward(self, x, mask=None, register_hook=False): """Forward propagation. Args: @@ -249,7 +257,12 @@ def forward(self, x, mask=None, register_hook = False): Returns: A tensor of shape [batch_size, seq_len, hidden] """ - x = self.input_sublayer(x, lambda _x: self.attention(_x, _x, _x, mask=mask, register_hook=register_hook)) + x = self.input_sublayer( + x, + lambda _x: self.attention( + _x, _x, _x, mask=mask, register_hook=register_hook + ), + ) x = self.output_sublayer(x, lambda _x: self.feed_forward(_x, mask=mask)) return self.dropout(x) @@ -290,7 +303,10 @@ def set_activation_hooks(self, hooks) -> None: return None def forward( - self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, register_hook: bool = False + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + register_hook: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward propagation. @@ -378,6 +394,7 @@ def __init__( heads: int = 1, dropout: float = 0.5, num_layers: int = 1, + max_seq_len: int = 1024, unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, ): super().__init__(dataset=dataset) @@ -385,6 +402,7 @@ def __init__( self.heads = heads self.dropout = dropout self.num_layers = num_layers + self.max_seq_len = max_seq_len self._attention_hooks_enabled = False self._use_unified = unified_embedding is not None @@ -417,19 +435,54 @@ def __init__( ) self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + def _pool_embedding(self, x: torch.Tensor) -> torch.Tensor: + """Pool nested embeddings to ``[batch, seq_len, hidden]`` format. + + Args: + x: Tensor emitted by the embedding model. + + Returns: + torch.Tensor: Sequence-aligned embedding tensor ready for attention. + + Example: + StageNet categorical inputs may have shape + ``[batch, seq_len, inner_len, emb]``. We sum over ``inner_len`` to + obtain per-event representations. + """ + + if x.dim() == 4: + x = x.sum(dim=2) + if x.dim() == 2: + x = x.unsqueeze(1) + # Truncate to max_seq_len to prevent quadratic memory spikes from + # outlier-length sequences (attention is O(S^2)). + if x.size(1) > self.max_seq_len: + x = x[:, : self.max_seq_len, :] + return x + + @staticmethod + def _mask_from_embeddings(x: torch.Tensor) -> torch.Tensor: + """Infer a boolean mask directly from embedded representations.""" + + mask = torch.any(torch.abs(x) > 0, dim=-1) + if mask.dim() == 1: + mask = mask.unsqueeze(1) + invalid_rows = ~mask.any(dim=1) + if invalid_rows.any(): + mask[invalid_rows, 0] = True + return mask.bool() + def _build_unified_inputs( self, kwargs: Dict[str, Any] ) -> Dict[str, Dict[str, torch.Tensor]]: - """Build the inputs dict required by UnifiedMultimodalEmbeddingModel. + """Build inputs expected by UnifiedMultimodalEmbeddingModel.""" - Reads each feature field from *kwargs* using the processor schema to - extract ``value``, ``time``, and optional ``mask`` tensors. - """ inputs: Dict[str, Dict[str, torch.Tensor]] = {} for field_name in self.feature_keys: feature = kwargs[field_name] if isinstance(feature, torch.Tensor): feature = (feature,) + schema = self.dataset.input_processors[field_name].schema() field_dict: Dict[str, torch.Tensor] = {} if "value" in schema: @@ -439,70 +492,41 @@ def _build_unified_inputs( if "mask" in schema: field_dict["mask"] = feature[schema.index("mask")].to(self.device) inputs[field_name] = field_dict + return inputs def _forward_unified( self, **kwargs: torch.Tensor | tuple[torch.Tensor, ...], ) -> Dict[str, torch.Tensor]: - """Forward pass in unified-embedding mode. + """Forward pass in unified-embedding mode.""" - Calls :class:`UnifiedMultimodalEmbeddingModel` to produce a single - temporally-sorted event sequence, then encodes it with one - :class:`TransformerLayer` and projects to label space. - """ - inputs = self._build_unified_inputs(kwargs) + register_hook = self._attention_hooks_enabled + inputs = self._build_unified_inputs(cast(Dict[str, Any], kwargs)) out = self.embedding_model(inputs) - sequence = out["sequence"] # (B, S_total, E) - event_mask = out["mask"].bool() # (B, S_total) + sequence = cast(torch.Tensor, out["sequence"]) + event_mask = cast(torch.Tensor, out["mask"]).bool() - _, cls_emb = self._unified_backbone(sequence, event_mask) - logits = self.fc(cls_emb) + _, patient_emb = self._unified_backbone(sequence, event_mask, register_hook) + + logits = self.fc(patient_emb) y_prob = self.prepare_y_prob(logits) - results: Dict[str, torch.Tensor] = {"logit": logits, "y_prob": y_prob} + results: Dict[str, torch.Tensor] = { + "logit": logits, + "y_prob": y_prob, + } + if self.label_key in kwargs: y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device) - results["loss"] = self.get_loss_function()(logits, y_true) + loss = self.get_loss_function()(logits, y_true) + results["loss"] = loss results["y_true"] = y_true + if kwargs.get("embed", False): - results["embed"] = cls_emb + results["embed"] = patient_emb return results - @staticmethod - def _pool_embedding(x: torch.Tensor) -> torch.Tensor: - """Pool nested embeddings to ``[batch, seq_len, hidden]`` format. - - Args: - x: Tensor emitted by the embedding model. - - Returns: - torch.Tensor: Sequence-aligned embedding tensor ready for attention. - - Example: - StageNet categorical inputs may have shape - ``[batch, seq_len, inner_len, emb]``. We sum over ``inner_len`` to - obtain per-event representations. - """ - - if x.dim() == 4: - x = x.sum(dim=2) - if x.dim() == 2: - x = x.unsqueeze(1) - return x - - @staticmethod - def _mask_from_embeddings(x: torch.Tensor) -> torch.Tensor: - """Infer a boolean mask directly from embedded representations.""" - - mask = torch.any(torch.abs(x) > 0, dim=-1) - if mask.dim() == 1: - mask = mask.unsqueeze(1) - invalid_rows = ~mask.any(dim=1) - if invalid_rows.any(): - mask[invalid_rows, 0] = True - return mask.bool() - def forward_from_embedding( self, **kwargs: torch.Tensor | tuple[torch.Tensor, ...], @@ -560,8 +584,7 @@ def forward_from_embedding( if value is None: raise ValueError( - f"Feature '{feature_key}' must contain 'value' " - f"in the schema." + f"Feature '{feature_key}' must contain 'value' " f"in the schema." ) else: value = value.to(self.device) @@ -575,9 +598,7 @@ def forward_from_embedding( else: mask = self._mask_from_embeddings(value).to(self.device) - _, cls_emb = self.transformer[feature_key]( - value, mask, register_hook - ) + _, cls_emb = self.transformer[feature_key](value, mask, register_hook) patient_emb.append(cls_emb) patient_emb = torch.cat(patient_emb, dim=1) @@ -643,15 +664,16 @@ def forward( if value is None: raise ValueError( - f"Feature '{feature_key}' must contain 'value' " - f"in the schema." + f"Feature '{feature_key}' must contain 'value' " f"in the schema." ) else: value = value.to(self.device) if mask is not None: mask = mask.to(self.device) - value = self.embedding_model({feature_key: value}, masks={feature_key: mask})[feature_key] + value = self.embedding_model( + {feature_key: value}, masks={feature_key: mask} + )[feature_key] else: value = self.embedding_model({feature_key: value})[feature_key] @@ -659,9 +681,9 @@ def forward( # Reconstruct tuple with embedded value # Note: we need to handle list/tuple conversion carefully # feature is a tuple. - + # Simple slice reconstruction - kwargs[feature_key] = feature[:i] + (value,) + feature[i + 1:] + kwargs[feature_key] = feature[:i] + (value,) + feature[i + 1 :] return self.forward_from_embedding(**kwargs) @@ -689,9 +711,7 @@ def get_attention_layers( cast(TransformerBlock, blk).attention.get_attn_map(), cast(TransformerBlock, blk).attention.get_attn_grad(), ) - for blk in cast( - TransformerLayer, self.transformer[key] - ).transformer + for blk in cast(TransformerLayer, self.transformer[key]).transformer ] for key in self.feature_keys } diff --git a/pyhealth/models/transformers_model.py b/pyhealth/models/transformers_model.py index 7116eeee3..5296420cd 100644 --- a/pyhealth/models/transformers_model.py +++ b/pyhealth/models/transformers_model.py @@ -15,12 +15,17 @@ def __init__( self, dataset: SampleDataset, model_name: str, + dropout: float = 0.1, ): super(TransformersModel, self).__init__( dataset=dataset, ) self.model_name = model_name - self.model = AutoModel.from_pretrained(model_name) + self.model = AutoModel.from_pretrained( + model_name, + hidden_dropout_prob=dropout, + attention_probs_dropout_prob=dropout, + ) assert ( len(self.feature_keys) == 1 ), "Only one feature key is supported if Transformers is initialized" @@ -32,6 +37,7 @@ def __init__( self.tokenizer = AutoTokenizer.from_pretrained(model_name) output_size = self.get_output_size() hidden_dim = self.model.config.hidden_size + self.dropout = nn.Dropout(dropout) self.fc = nn.Linear(hidden_dim, output_size) def forward(self, **kwargs) -> Dict[str, torch.Tensor]: @@ -45,7 +51,7 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: x = x.to(self.device) # TODO: should not use pooler_output, but use the last hidden state embeddings = self.model(**x).pooler_output - logits = self.fc(embeddings) + logits = self.fc(self.dropout(embeddings)) y_true = kwargs[self.label_key].to(self.device) loss = self.get_loss_function()(logits, y_true) y_prob = self.prepare_y_prob(logits) diff --git a/pyhealth/processors/stagenet_processor.py b/pyhealth/processors/stagenet_processor.py index b0e430d40..f9b242f0d 100644 --- a/pyhealth/processors/stagenet_processor.py +++ b/pyhealth/processors/stagenet_processor.py @@ -529,4 +529,4 @@ def __repr__(self): return ( f"StageNetTensorProcessor(is_nested={self._is_nested}, " f"feature_dim={self._size})" - ) + ) \ No newline at end of file diff --git a/pyhealth/processors/time_image_processor.py b/pyhealth/processors/time_image_processor.py index 9d313e6bc..421998d07 100644 --- a/pyhealth/processors/time_image_processor.py +++ b/pyhealth/processors/time_image_processor.py @@ -75,6 +75,11 @@ class TimeImageProcessor(TemporalFeatureProcessor): patient has more images, the most recent (by timestamp) are kept. If None, all images are kept. Defaults to None. + padding: Sentinel string that marks a missing image. When + a path equals this value, a zero tensor of shape + (C, H, W) is returned instead of loading from disk. + If None, all paths are treated as real file paths. + Defaults to None. Raises: ValueError: If normalize is True but mean or std is missing. @@ -107,6 +112,7 @@ def __init__( std: Optional[List[float]] = None, mode: Optional[str] = None, max_images: Optional[int] = None, + padding: Optional[str] = None, ) -> None: self.image_size = image_size self.to_tensor = to_tensor @@ -115,18 +121,14 @@ def __init__( self.std = std self.mode = mode self.max_images = max_images + self.padding = padding self.n_channels = None - if self.normalize and ( - self.mean is None or self.std is None - ): + if self.normalize and (self.mean is None or self.std is None): raise ValueError( - "Normalization requires both mean and std to be " - "provided." + "Normalization requires both mean and std to be " "provided." ) - if not self.normalize and ( - self.mean is not None or self.std is not None - ): + if not self.normalize and (self.mean is not None or self.std is not None): raise ValueError( "Mean and std are provided but normalize is set " "to False. Either provide normalize=True, or " @@ -146,36 +148,49 @@ def _build_transform(self) -> transforms.Compose: transform_list = [] if self.mode is not None: transform_list.append( - transforms.Lambda( - partial(_convert_mode, mode=self.mode) - ) + transforms.Lambda(partial(_convert_mode, mode=self.mode)) ) if self.image_size is not None: - transform_list.append( - transforms.Resize( - (self.image_size, self.image_size) - ) - ) + transform_list.append(transforms.Resize((self.image_size, self.image_size))) if self.to_tensor: transform_list.append(transforms.ToTensor()) if self.normalize: - transform_list.append( - transforms.Normalize( - mean=self.mean, std=self.std - ) - ) + transform_list.append(transforms.Normalize(mean=self.mean, std=self.std)) return transforms.Compose(transform_list) - def _load_single_image( - self, path: Union[str, Path] - ) -> torch.Tensor: + def _zero_image_tensor(self) -> torch.Tensor: + """Return a zero tensor matching the expected image shape (C, H, W). + + Used as a placeholder when an image path is an empty string. + Channel count is inferred from self.n_channels if available, + otherwise derived from self.mode ("L"→1, "RGBA"→4, else 3). + + Returns: + Zero tensor of shape (C, image_size, image_size). + """ + if self.n_channels is not None: + c = self.n_channels + elif self.mode == "L": + c = 1 + elif self.mode == "RGBA": + c = 4 + else: + c = 3 + return torch.zeros(c, self.image_size, self.image_size) + + def _load_single_image(self, path: Union[str, Path]) -> torch.Tensor: """Load and transform a single image from disk. + If path equals missing_path_token, returns a zero tensor of + the same shape as a normal image (C, H, W) via _zero_image_tensor. + Called internally by process() for each image path in the input list. Args: - path: Path to the image file. + path: Path to the image file. If this equals + missing_path_token, a zero-filled placeholder tensor + is returned instead. Returns: Transformed image tensor of shape (C, H, W). @@ -183,18 +198,16 @@ def _load_single_image( Raises: FileNotFoundError: If the image file does not exist. """ + if self.padding is not None and str(path) == self.padding: + return self._zero_image_tensor() image_path = Path(path) if not image_path.exists(): - raise FileNotFoundError( - f"Image file not found: {image_path}" - ) + raise FileNotFoundError(f"Image file not found: {image_path}") with Image.open(image_path) as img: img.load() return self.transform(img) - def fit( - self, samples: Iterable[Dict[str, Any]], field: str - ) -> None: + def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: """Fit the processor by inferring n_channels from data. Scans samples to find the first valid entry for the given @@ -214,8 +227,10 @@ def fit( for sample in samples: if field in sample and sample[field] is not None: image_paths, _ = sample[field] - if len(image_paths) > 0: - path = Path(image_paths[0]) + for raw_path in image_paths: + if self.padding is not None and str(raw_path) == self.padding: + continue + path = Path(raw_path) if path.exists(): with Image.open(path) as img: if img.mode == "L": @@ -225,14 +240,14 @@ def fit( else: self.n_channels = 3 break + if self.n_channels is not None: + break if self.n_channels is None: self.n_channels = 3 def process( self, - value: Tuple[ - List[Union[str, Path]], List[float] - ], + value: Tuple[List[Union[str, Path]], List[float]], ) -> Tuple[torch.Tensor, torch.Tensor, str]: """Process paired image paths and timestamps. @@ -278,15 +293,10 @@ def process( if len(image_paths) == 0: raise ValueError("image_paths must be non-empty.") - paired = sorted( - zip(time_diffs, image_paths), key=lambda x: x[0] - ) + paired = sorted(zip(time_diffs, image_paths), key=lambda x: x[0]) - if ( - self.max_images is not None - and len(paired) > self.max_images - ): - paired = paired[-self.max_images:] + if self.max_images is not None and len(paired) > self.max_images: + paired = paired[-self.max_images :] timestamps = [] image_tensors = [] @@ -295,9 +305,7 @@ def process( timestamps.append(t) images = torch.stack(image_tensors, dim=0) - timestamps = torch.tensor( - timestamps, dtype=torch.float32 - ) + timestamps = torch.tensor(timestamps, dtype=torch.float32) if self.n_channels is None: self.n_channels = images.shape[1] @@ -344,5 +352,6 @@ def __repr__(self) -> str: f"mean={self.mean}, " f"std={self.std}, " f"mode={self.mode}, " - f"max_images={self.max_images})" - ) \ No newline at end of file + f"max_images={self.max_images}, " + f"padding={self.padding!r})" + ) diff --git a/pyhealth/tasks/EEG_abnormal.py b/pyhealth/tasks/EEG_abnormal.py deleted file mode 100644 index 5ce4d4b54..000000000 --- a/pyhealth/tasks/EEG_abnormal.py +++ /dev/null @@ -1,165 +0,0 @@ -import os -import pickle -import mne - - -def EEG_isAbnormal_fn(record): - """Processes a single patient for the abnormal EEG detection task on TUAB. - - Abnormal EEG detection aims at determining whether a EEG is abnormal. - - Args: - record: a singleton list of one subject from the TUABDataset. - The (single) record is a dictionary with the following keys: - load_from_path, patient_id, visit_id, signal_file, label_file, save_to_path - - Returns: - samples: a list of samples, each sample is a dict with patient_id, visit_id, record_id, - and epoch_path (the path to the saved epoch {"signal": signal, "label": label} as key. - - Note that we define the task as a binary classification task. - - Examples: - >>> from pyhealth.datasets import TUABDataset - >>> isabnormal = TUABDataset( - ... root="/srv/local/data/TUH/tuh_eeg_abnormal/v3.0.0/edf/", download=True, - ... ) - >>> from pyhealth.tasks import EEG_isabnormal_fn - >>> EEG_abnormal_ds = isabnormal.set_task(EEG_isAbnormal_fn) - >>> EEG_abnormal_ds.samples[0] - { - 'patient_id': 'aaaaamye', - 'visit_id': 's001', - 'record_id': '1', - 'epoch_path': '/home/zhenlin4/.cache/pyhealth/datasets/832afe6e6e8a5c9ea5505b47e7af8125/10-1/1/0.pkl', - 'label': 1 - } - """ - - samples = [] - for visit in record: - root, pid, visit_id, signal, label, save_path = ( - visit["load_from_path"], - visit["patient_id"], - visit["visit_id"], - visit["signal_file"], - visit["label_file"], - visit["save_to_path"], - ) - - raw = mne.io.read_raw_edf(os.path.join(root, signal), preload=True) - raw.resample(200) - ch_name = raw.ch_names - raw_data = raw.get_data() - channeled_data = raw_data.copy()[:16] - try: - channeled_data[0] = ( - raw_data[ch_name.index("EEG FP1-REF")] - - raw_data[ch_name.index("EEG F7-REF")] - ) - channeled_data[1] = ( - raw_data[ch_name.index("EEG F7-REF")] - - raw_data[ch_name.index("EEG T3-REF")] - ) - channeled_data[2] = ( - raw_data[ch_name.index("EEG T3-REF")] - - raw_data[ch_name.index("EEG T5-REF")] - ) - channeled_data[3] = ( - raw_data[ch_name.index("EEG T5-REF")] - - raw_data[ch_name.index("EEG O1-REF")] - ) - channeled_data[4] = ( - raw_data[ch_name.index("EEG FP2-REF")] - - raw_data[ch_name.index("EEG F8-REF")] - ) - channeled_data[5] = ( - raw_data[ch_name.index("EEG F8-REF")] - - raw_data[ch_name.index("EEG T4-REF")] - ) - channeled_data[6] = ( - raw_data[ch_name.index("EEG T4-REF")] - - raw_data[ch_name.index("EEG T6-REF")] - ) - channeled_data[7] = ( - raw_data[ch_name.index("EEG T6-REF")] - - raw_data[ch_name.index("EEG O2-REF")] - ) - channeled_data[8] = ( - raw_data[ch_name.index("EEG FP1-REF")] - - raw_data[ch_name.index("EEG F3-REF")] - ) - channeled_data[9] = ( - raw_data[ch_name.index("EEG F3-REF")] - - raw_data[ch_name.index("EEG C3-REF")] - ) - channeled_data[10] = ( - raw_data[ch_name.index("EEG C3-REF")] - - raw_data[ch_name.index("EEG P3-REF")] - ) - channeled_data[11] = ( - raw_data[ch_name.index("EEG P3-REF")] - - raw_data[ch_name.index("EEG O1-REF")] - ) - channeled_data[12] = ( - raw_data[ch_name.index("EEG FP2-REF")] - - raw_data[ch_name.index("EEG F4-REF")] - ) - channeled_data[13] = ( - raw_data[ch_name.index("EEG F4-REF")] - - raw_data[ch_name.index("EEG C4-REF")] - ) - channeled_data[14] = ( - raw_data[ch_name.index("EEG C4-REF")] - - raw_data[ch_name.index("EEG P4-REF")] - ) - channeled_data[15] = ( - raw_data[ch_name.index("EEG P4-REF")] - - raw_data[ch_name.index("EEG O2-REF")] - ) - except: - with open("tuab-process-error-files.txt", "a") as f: - f.write(os.path.join(root, signal) + "\n") - continue - - # get the label - data_field = pid.split("_")[0] - if data_field == "0" or data_field == "2": - label = 1 - else: - label = 0 - - # load data - for i in range(channeled_data.shape[1] // 2000): - dump_path = os.path.join( - save_path, pid + "_" + visit_id + "_" + str(i) + ".pkl" - ) - pickle.dump( - {"signal": channeled_data[:, i * 2000 : (i + 1) * 2000], "label": label}, - open(dump_path, "wb"), - ) - - samples.append( - { - "patient_id": pid, - "visit_id": visit_id, - "record_id": i, - "epoch_path": dump_path, - "label": label, - } - ) - - return samples - - -if __name__ == "__main__": - from pyhealth.datasets import TUABDataset - - dataset = TUABDataset( - root="/srv/local/data/TUH/tuh_eeg_abnormal/v3.0.0/edf/", - dev=True, - refresh_cache=True, - ) - EEG_abnormal_ds = dataset.set_task(EEG_isAbnormal_fn) - print(EEG_abnormal_ds.samples[0]) - print(EEG_abnormal_ds.input_info) diff --git a/pyhealth/tasks/EEG_events.py b/pyhealth/tasks/EEG_events.py deleted file mode 100644 index 5993413a9..000000000 --- a/pyhealth/tasks/EEG_events.py +++ /dev/null @@ -1,203 +0,0 @@ -import os -import pickle -import mne -import numpy as np - - -def EEG_events_fn(record): - """Processes a single patient for the EEG events task on TUEV. - - This task aims at annotating of EEG segments as one of six classes: (1) spike and sharp wave (SPSW), (2) generalized periodic epileptiform discharges (GPED), (3) periodic lateralized epileptiform discharges (PLED), (4) eye movement (EYEM), (5) artifact (ARTF) and (6) background (BCKG). - - Args: - record: a singleton list of one subject from the TUEVDataset. - The (single) record is a dictionary with the following keys: - load_from_path, patient_id, visit_id, signal_file, label_file, save_to_path - - Returns: - samples: a list of samples, each sample is a dict with patient_id, visit_id, record_id, label, offending_channel, - and epoch_path (the path to the saved epoch {"signal": signal, "label": label} as key. - - Note that we define the task as a multiclass classification task. - - Examples: - >>> from pyhealth.datasets import TUEVDataset - >>> EEGevents = TUEVDataset( - ... root="/srv/local/data/TUH/tuh_eeg_events/v2.0.0/edf/", download=True, - ... ) - >>> from pyhealth.tasks import EEG_events_fn - >>> EEG_events_ds = EEGevents.set_task(EEG_events_fn) - >>> EEG_events_ds.samples[0] - { - 'patient_id': '0_00002265', - 'visit_id': '00000001', - 'record_id': 0, - 'epoch_path': '/Users/liyanjing/.cache/pyhealth/datasets/d8f3cb92cc444d481444d3414fb5240c/0_00002265_00000001_0.pkl', - 'label': 6, - 'offending_channel': array([4.]) - } - """ - - samples = [] - for visit in record: - root, pid, visit_id, signal, label, save_path = ( - visit["load_from_path"], - visit["patient_id"], - visit["visit_id"], - visit["signal_file"], - visit["label_file"], - visit["save_to_path"], - ) - - - # load data - try: - [signals, times, event, Rawdata] = readEDF( - os.path.join(root, signal) - ) # event is the .rec file in the form of an array - signals = convert_signals(signals, Rawdata) - except (ValueError, KeyError): - print("something funky happened in " + os.path.join(root, signal)) - continue - signals, offending_channels, labels = BuildEvents(signals, times, event) - - for idx, (signal, offending_channel, label) in enumerate( - zip(signals, offending_channels, labels) - ): - dump_path = os.path.join( - save_path, pid + "_" + visit_id + "_" + str(idx) + ".pkl" - ) - - pickle.dump( - {"signal": signal, "label": int(label[0])}, - open(dump_path, "wb"), - ) - - samples.append( - { - "patient_id": pid, - "visit_id": visit_id, - "record_id": idx, - "epoch_path": dump_path, - "label": int(label[0]), - "offending_channel": offending_channel, - } - ) - - return samples - -def BuildEvents(signals, times, EventData): - [numEvents, z] = EventData.shape # numEvents is equal to # of rows of the .rec file - fs = 250.0 - [numChan, numPoints] = signals.shape - - features = np.zeros([numEvents, numChan, int(fs) * 5]) - offending_channel = np.zeros([numEvents, 1]) # channel that had the detected thing - labels = np.zeros([numEvents, 1]) - offset = signals.shape[1] - signals = np.concatenate([signals, signals, signals], axis=1) - for i in range(numEvents): # for each event - chan = int(EventData[i, 0]) # chan is channel - start = np.where((times) >= EventData[i, 1])[0][0] - end = np.where((times) >= EventData[i, 2])[0][0] - features[i, :] = signals[ - :, offset + start - 2 * int(fs) : offset + end + 2 * int(fs) - ] - offending_channel[i, :] = int(chan) - labels[i, :] = int(EventData[i, 3]) - return [features, offending_channel, labels] - - -def convert_signals(signals, Rawdata): - signal_names = { - k: v - for (k, v) in zip( - Rawdata.info["ch_names"], list(range(len(Rawdata.info["ch_names"]))) - ) - } - new_signals = np.vstack( - ( - signals[signal_names["EEG FP1-REF"]] - - signals[signal_names["EEG F7-REF"]], # 0 - ( - signals[signal_names["EEG F7-REF"]] - - signals[signal_names["EEG T3-REF"]] - ), # 1 - ( - signals[signal_names["EEG T3-REF"]] - - signals[signal_names["EEG T5-REF"]] - ), # 2 - ( - signals[signal_names["EEG T5-REF"]] - - signals[signal_names["EEG O1-REF"]] - ), # 3 - ( - signals[signal_names["EEG FP2-REF"]] - - signals[signal_names["EEG F8-REF"]] - ), # 4 - ( - signals[signal_names["EEG F8-REF"]] - - signals[signal_names["EEG T4-REF"]] - ), # 5 - ( - signals[signal_names["EEG T4-REF"]] - - signals[signal_names["EEG T6-REF"]] - ), # 6 - ( - signals[signal_names["EEG T6-REF"]] - - signals[signal_names["EEG O2-REF"]] - ), # 7 - ( - signals[signal_names["EEG FP1-REF"]] - - signals[signal_names["EEG F3-REF"]] - ), # 14 - ( - signals[signal_names["EEG F3-REF"]] - - signals[signal_names["EEG C3-REF"]] - ), # 15 - ( - signals[signal_names["EEG C3-REF"]] - - signals[signal_names["EEG P3-REF"]] - ), # 16 - ( - signals[signal_names["EEG P3-REF"]] - - signals[signal_names["EEG O1-REF"]] - ), # 17 - ( - signals[signal_names["EEG FP2-REF"]] - - signals[signal_names["EEG F4-REF"]] - ), # 18 - ( - signals[signal_names["EEG F4-REF"]] - - signals[signal_names["EEG C4-REF"]] - ), # 19 - ( - signals[signal_names["EEG C4-REF"]] - - signals[signal_names["EEG P4-REF"]] - ), # 20 - (signals[signal_names["EEG P4-REF"]] - signals[signal_names["EEG O2-REF"]]), - ) - ) # 21 - return new_signals - - -def readEDF(fileName): - Rawdata = mne.io.read_raw_edf(fileName) - signals, times = Rawdata[:] - RecFile = fileName[0:-3] + "rec" - eventData = np.genfromtxt(RecFile, delimiter=",") - Rawdata.close() - return [signals, times, eventData, Rawdata] - - -if __name__ == "__main__": - from pyhealth.datasets import TUEVDataset - - dataset = TUEVDataset( - root="/srv/local/data/TUH/tuh_eeg_events/v2.0.0/edf/", - dev=True, - refresh_cache=True, - ) - EEG_events_ds = dataset.set_task(EEG_events_fn) - print(EEG_events_ds.samples[0]) - print(EEG_events_ds.input_info) diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index be6ec9fca..e98d98db6 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -21,8 +21,6 @@ drug_recommendation_mimic4_fn, drug_recommendation_omop_fn, ) -from .EEG_abnormal import EEG_isAbnormal_fn -from .EEG_events import EEG_events_fn from .in_hospital_mortality_mimic4 import InHospitalMortalityMIMIC4 from .length_of_stay_prediction import ( LengthOfStayPredictioneICU, @@ -49,6 +47,9 @@ from .multimodal_mimic4 import ( ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4, + ClinicalNotesICDLabsCXRMIMIC4, + ICDLabsMIMIC4, + LabsMIMIC4, NotesLabsMIMIC4, ) from .patient_linkage import patient_linkage_mimic3_fn diff --git a/pyhealth/tasks/covid19_cxr_classification.py b/pyhealth/tasks/covid19_cxr_classification.py index ae6814928..0a7f526f5 100644 --- a/pyhealth/tasks/covid19_cxr_classification.py +++ b/pyhealth/tasks/covid19_cxr_classification.py @@ -27,7 +27,7 @@ class COVID19CXRClassification(BaseTask): """ task_name: str = "COVID19CXRClassification" - input_schema: Dict[str, str] = {"image": "image"} + input_schema: Dict = {"image": ("image", {"mode": "RGB"})} output_schema: Dict[str, str] = {"disease": "multiclass"} def __call__(self, patient: Any) -> List[Dict[str, Any]]: diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 248377322..71f8d3f87 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -13,6 +13,7 @@ class BaseMultimodalMIMIC4Task(BaseTask): """ MISSING_TEXT_TOKEN: ClassVar[str] = "" + MISSING_CODE_TOKEN: ClassVar[str] = "" MISSING_FLOAT_TOKEN: ClassVar[float] = 0.0 LAB_CATEGORIES: ClassVar[Dict[str, List[str]]] = { @@ -769,6 +770,208 @@ 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. + """ + + 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, + ) + radiology_note_times_from_admission = ( + all_radiology_texts, + 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, + "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, + } + + return [single_patient_longitudinal_record] + + class ICDLabsMIMIC4(BaseMultimodalMIMIC4Task): """Task for ICD codes + lab values mortality prediction using MIMIC-IV. @@ -1077,3 +1280,87 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: 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 + ) + if not admissions_to_process: + return [] + + effective_start, effective_end = self._compute_effective_window( + admissions_to_process + ) + + 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, + ) + 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, + "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] diff --git a/pyhealth/tasks/temple_university_EEG_tasks.py b/pyhealth/tasks/temple_university_EEG_tasks.py index 93b69cdf4..fc2ba702a 100644 --- a/pyhealth/tasks/temple_university_EEG_tasks.py +++ b/pyhealth/tasks/temple_university_EEG_tasks.py @@ -129,7 +129,7 @@ def readEDF(fileName: str, Rawdata.filter(l_freq=bandpass_filter[0], h_freq=bandpass_filter[1], verbose="error") Rawdata.notch_filter(notch_filter, verbose="error") - Rawdata.resample(resample_rate, n_jobs=5, verbose="error") + Rawdata.resample(resample_rate, n_jobs=1, verbose="error") _, times = Rawdata[:] signals = Rawdata.get_data(units="uV") @@ -256,7 +256,7 @@ def read_and_process_edf(fileName: str, Rawdata.filter(l_freq=bandpass_filter[0], h_freq=bandpass_filter[1], verbose="error") Rawdata.notch_filter(notch_filter, verbose="error") - Rawdata.resample(resample_rate, n_jobs=5, verbose="error") + Rawdata.resample(resample_rate, n_jobs=1, verbose="error") raw_data = Rawdata.get_data(units="uV") ch_name = Rawdata.ch_names diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index 9894e7285..fc264a3af 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -137,6 +137,9 @@ def train( monitor_criterion: str = "max", load_best_model_at_last: bool = True, patience=None, + accumulation_steps: int = 1, + use_amp: bool = False, + amp_dtype: str = "bf16", ): """Trains the model. @@ -156,10 +159,25 @@ def train( Default is True. patience: Number of epochs to wait for improvement before early stopping. Default is None, which means no early stopping. + accumulation_steps: Gradient accumulation steps to simulate a larger + effective batch size. Default is 1 (no accumulation). + use_amp: Whether to use automatic mixed precision. Default is False. + amp_dtype: AMP dtype — "bf16" (stable, recommended) or "fp16". + Default is "bf16". """ if optimizer_params is None: optimizer_params = {"lr": 1e-3} + _amp_dtype = ( + torch.bfloat16 if amp_dtype == "bf16" else torch.float16 + ) + # GradScaler only needed for fp16; bf16 has fp32 dynamic range + scaler = ( + torch.cuda.amp.GradScaler() + if (use_amp and _amp_dtype == torch.float16) + else None + ) + # logging logger.info("Training:") logger.info(f"Batch size: {train_dataloader.batch_size}") @@ -172,6 +190,8 @@ def train( logger.info(f"Monitor criterion: {monitor_criterion}") logger.info(f"Epochs: {epochs}") logger.info(f"Patience: {patience}") + logger.info(f"Accumulation steps: {accumulation_steps}") + logger.info(f"AMP: {use_amp} (dtype={amp_dtype})") # set optimizer param = list(self.model.named_parameters()) @@ -210,7 +230,7 @@ def train( epoch_start = time.perf_counter() # batch training loop logger.info("") - for _ in trange( + for step_idx in trange( steps_per_epoch, desc=f"Epoch {epoch + 1}/{epochs}", smoothing=0.05, @@ -221,20 +241,39 @@ def train( except StopIteration: data_iterator = iter(train_dataloader) data = next(data_iterator) - # forward - output = self.model(**data) - loss = output["loss"] + # forward (with optional AMP) + if use_amp: + with torch.autocast(device_type="cuda", dtype=_amp_dtype): + output = self.model(**data) + loss = output["loss"] / accumulation_steps + else: + output = self.model(**data) + loss = output["loss"] / accumulation_steps # backward - loss.backward() - if max_grad_norm is not None: - torch.nn.utils.clip_grad_norm_( - self.model.parameters(), max_grad_norm - ) - # update - optimizer.step() - optimizer.zero_grad() - training_loss.append(loss.item()) - global_step += 1 + if scaler is not None: + scaler.scale(loss).backward() + else: + loss.backward() + training_loss.append(loss.item() * accumulation_steps) + # optimizer step every accumulation_steps batches or epoch end + is_update_step = ( + (step_idx + 1) % accumulation_steps == 0 + or (step_idx + 1) == steps_per_epoch + ) + if is_update_step: + if max_grad_norm is not None: + if scaler is not None: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), max_grad_norm + ) + if scaler is not None: + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad() + global_step += 1 epoch_time = time.perf_counter() - epoch_start vram = _vram_stats(self.device) diff --git a/pyproject.toml b/pyproject.toml index c1e0ebcc6..934d4f1bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,16 @@ keywords = [ "deep learning", ] +[project.optional-dependencies] +graph = [ + "torch-geometric>=2.6.0", +] +nlp = [ + "editdistance~=0.8.1", + "rouge_score~=0.1.2", + "nltk~=3.9.1", +] + [project.urls] Homepage = "https://github.com/sunlabuiuc/PyHealth" Documentation = "https://pyhealth.readthedocs.io/en/latest/about.html" @@ -97,7 +107,7 @@ build-env = { features = ["pyverbase", "build-env"], solve-group = "default" } # specify where to download the build backend that builds conda files [tool.pixi.package.build.backend] name = "pixi-build-python" -version = "2.0.0" +version = "==2.0.0" channels = [ "https://prefix.dev/pixi-build-backends", "https://prefix.dev/conda-forge", diff --git a/scripts/slurm/run_table2.sh b/scripts/slurm/run_table2.sh index 8e7356922..97d6b85fb 100755 --- a/scripts/slurm/run_table2.sh +++ b/scripts/slurm/run_table2.sh @@ -77,8 +77,12 @@ 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 diff --git a/scripts/slurm/submit_rnn.sh b/scripts/slurm/submit_rnn.sh index 900dbe51b..847f8210d 100755 --- a/scripts/slurm/submit_rnn.sh +++ b/scripts/slurm/submit_rnn.sh @@ -21,7 +21,7 @@ for seed in "${SEEDS[@]}"; do --mem=32G --gres=gpu:1 --time=6:00:00 \ --output="logs/slurm/table2_rnn_seed${seed}_%j.out" \ --error="logs/slurm/table2_rnn_seed${seed}_%j.err" \ - --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=rnn,SEED="${seed}",TABLE2_BS_RNN=16 \ + --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=rnn,SEED="${seed}",TABLE2_BS_RNN=16,TABLE2_TASK="${TABLE2_TASK:-clinical_notes_icd_labs}" \ scripts/slurm/run_table2.sh | awk '{print $NF}') echo " Submitted rnn seed=${seed} → ${job}" done diff --git a/scripts/sunlab/run_ehrmamba_timewindow.sh b/scripts/sunlab/run_ehrmamba_timewindow.sh new file mode 100644 index 000000000..2ea666be9 --- /dev/null +++ b/scripts/sunlab/run_ehrmamba_timewindow.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# Run ehrmamba jobs for multiple seeds in tmux sessions +# Usage: bash scripts/run_ehrmamba.sh +set -euo pipefail + +SEEDS=(267573289 1872967241 706384748) +TIME_WINDOWS=(24 48 96) +PROJECT_DIR="${PROJECT_DIR:-$(pwd)}" +USER="" + +cd "${PROJECT_DIR}" +mkdir -p logs + +for seed in "${SEEDS[@]}"; do + for tw in "${TIME_WINDOWS[@]}"; do + echo " Launching ehrmamba seed=${seed} time_window=${tw} in tmux" + tmux new -s "ehrmamba_${seed}_observationtimewindow_${tw}" "bash -lc ' + conda activate pyhealth2 && \ + CUDA_VISIBLE_DEVICES=4 python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \ + --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \ + --cache-dir /home/${USER}/pyhealth_cache \ + --task icd_labs \ + --model ehrmamba \ + --embedding-dim 128 \ + --num-layers 2 \ + --observation-window-hours ${tw} \ + --mamba-state-size 16 \ + --mamba-conv-kernel 4 \ + --epochs 20 \ + --batch-size 8 \ + --seed ${seed} \ + --output-dir /home/${USER}/output/table2/ehrmamba_seed${seed}_observationtimewindow_${tw} \ + 2>&1 | tee logs/ehrmamba_seed${seed}_observationtimewindow_${tw}.log; exec bash'" + echo " Launched ehrmamba seed=${seed} time_window=${tw} → tmux session: ehrmamba_${seed}_observationtimewindow_${tw}" + done +done diff --git a/scripts/sunlab/run_rnn.sh b/scripts/sunlab/run_rnn.sh new file mode 100755 index 000000000..f47238805 --- /dev/null +++ b/scripts/sunlab/run_rnn.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Submit Table 2 jobs for rnn — 3 fixed seeds, sunlab cluster (tmux sessions). +# Seeds are shared across all lab members for aligned comparisons. +# Usage: bash scripts/sunlab/run_rnn.sh +set -euo pipefail + +# SEEDS=(267573289 1872967241 706384748) +SEEDS=(42) +PROJECT_DIR="${PROJECT_DIR:-$(pwd)}" +USER="${USER:-wp14}" +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +CUDA_DEVICE="${CUDA_DEVICE:-0}" +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:-/home/${USER}/pyhealth_cache}" +OUTPUT_BASE="${OUTPUT_BASE:-/home/${USER}/output/table2}" + +cd "${PROJECT_DIR}" +mkdir -p logs/sunlab + +for seed in "${SEEDS[@]}"; do + session="rnn_seed${seed}" + echo " Launching rnn seed=${seed} in tmux session: ${session}" + tmux new-session -d -s "${session}" "bash -lc ' + conda activate ${CONDA_ENV} && \ + CUDA_VISIBLE_DEVICES=${CUDA_DEVICE} python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \ + --ehr-root ${EHR_ROOT} \ + --cache-dir ${CACHE_DIR} \ + --task icd_labs \ + --model rnn \ + --embedding-dim 128 \ + --hidden-dim 128 \ + --heads 4 \ + --num-layers 2 \ + --dropout 0.1 \ + --epochs 20 \ + --batch-size 16 \ + --weight-decay 1e-5 \ + --patience 5 \ + --seed ${seed} \ + --output-dir ${OUTPUT_BASE}/rnn_seed${seed} \ + 2>&1 | tee logs/sunlab/rnn_seed${seed}.log; exec bash'" + echo " Launched rnn seed=${seed} → tmux session: ${session}" +done diff --git a/scripts/sunlab/run_transformer.sh b/scripts/sunlab/run_transformer.sh new file mode 100644 index 000000000..0a2aacaff --- /dev/null +++ b/scripts/sunlab/run_transformer.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Submit Table 2 jobs for transformer — 3 fixed seeds, sunlab cluster (tmux sessions). +# Seeds are shared across all lab members for aligned comparisons. +# Usage: bash scripts/sunlab/run_transformer.sh +set -euo pipefail + +# SEEDS=(267573289 1872967241 706384748) +SEEDS=(44) +PROJECT_DIR="${PROJECT_DIR:-$(pwd)}" +USER="${USER:-wp14}" +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +CUDA_DEVICE="${CUDA_DEVICE:-4}" +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:-/home/${USER}/pyhealth_cache}" +OUTPUT_BASE="${OUTPUT_BASE:-/home/${USER}/output/table2}" + +cd "${PROJECT_DIR}" + +for seed in "${SEEDS[@]}"; do + session="transformer_seed${seed}" + log_dir="logs/sunlab/$(date +%Y%m%d)" + mkdir -p "${log_dir}" + tmux new-session -d -s "${session}" "bash -lc ' + conda activate ${CONDA_ENV} && \ + CUDA_VISIBLE_DEVICES=${CUDA_DEVICE} python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \ + --ehr-root ${EHR_ROOT} \ + --cache-dir ${CACHE_DIR} \ + --task icd_labs \ + --model transformer \ + --embedding-dim 64 \ + --hidden-dim 64 \ + --heads 2 \ + --num-layers 1 \ + --dropout 0.1 \ + --epochs 20 \ + --batch-size 4 \ + --weight-decay 1e-5 \ + --patience 5 \ + --seed ${seed} \ + --output-dir ${OUTPUT_BASE}/$(date +%Y%m%d) \ + 2>&1 | tee ${log_dir}/transformer_seed${seed}.log; exec bash'" + echo " Launched transformer seed=${seed} → tmux session: ${session}, log: ${log_dir}/transformer_seed${seed}.log" +done diff --git a/scripts/train_unified.py b/scripts/train_unified.py new file mode 100644 index 000000000..6d908b491 --- /dev/null +++ b/scripts/train_unified.py @@ -0,0 +1,163 @@ +"""Unified training CLI for PyHealth multimodal experiments. + +Reads a YAML config (with optional _inherit chain) then fires the E2E script +with all resolved arguments. CLI flags override YAML values. + +Usage: + python scripts/train_unified.py --config configs/train/e2e_baseline.yaml \ + --model transformer --seed 0 + + # Smoke test + python scripts/train_unified.py --config configs/train/smoke.yaml + + # Full run with frozen encoder + python scripts/train_unified.py --config configs/train/e2e_balanced.yaml \ + --model ehrmamba --freeze-encoder --seed 42 + +The script resolves _model_overrides from the YAML, merges CLI overrides, then +calls unified_embedding_e2e_mimic4.py via subprocess so it can be used from +condor jobs without reimplementing the full arg-build logic. +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict + + +def _load_yaml(path: Path) -> Dict[str, Any]: + try: + import yaml # type: ignore + except ImportError: + raise SystemExit("PyYAML not installed. Run: pip install pyyaml") + with open(path) as f: + return yaml.safe_load(f) or {} + + +def _resolve_config(config_path: Path) -> Dict[str, Any]: + """Load config and recursively merge _inherit chain.""" + cfg = _load_yaml(config_path) + inherit = cfg.pop("_inherit", None) + if inherit: + parent_path = config_path.parent / inherit + parent = _resolve_config(parent_path) + # parent values are defaults; child values win + merged = {**parent, **cfg} + return merged + return cfg + + +def _apply_model_overrides(cfg: Dict[str, Any]) -> Dict[str, Any]: + """Merge _model_overrides[model] into cfg if present.""" + overrides = cfg.pop("_model_overrides", {}) + model = cfg.get("model", "mlp") + if model in overrides: + cfg = {**cfg, **overrides[model]} + return cfg + + +def _to_args(cfg: Dict[str, Any]) -> list[str]: + """Convert resolved config dict to CLI args for unified_embedding_e2e_mimic4.py.""" + # Map config keys to CLI flags (underscores → hyphens, bool flags → store_true) + bool_flags = { + "icd_codes", "freeze_encoder", "include_vitals", + "balanced_sampling", "bidirectional", + } + skip_keys = {"_inherit", "_model_overrides"} + + args = [] + for key, val in cfg.items(): + if key in skip_keys or val is None: + continue + flag = "--" + key.replace("_", "-") + if key in bool_flags: + if val: + args.append(flag) + else: + args += [flag, str(val)] + return args + + +def _parse_cli() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Unified PyHealth training CLI — loads a YAML config then fires the E2E script." + ) + parser.add_argument( + "--config", + required=True, + help="Path to a YAML config file (e.g. configs/train/e2e_baseline.yaml).", + ) + # Pass-through overrides — any key from the YAML can be overridden here. + parser.add_argument("--model", default=None) + parser.add_argument("--seed", type=int, default=None) + parser.add_argument("--task", default=None) + parser.add_argument("--epochs", type=int, default=None) + parser.add_argument("--batch-size", type=int, default=None) + parser.add_argument("--embedding-dim", type=int, default=None) + parser.add_argument("--dev", type=int, default=None) + parser.add_argument("--freeze-encoder", action="store_true", default=None) + parser.add_argument("--balanced-sampling", action="store_true", default=None) + parser.add_argument("--icd-codes", action="store_true", default=None) + parser.add_argument("--output-dir", default=None) + parser.add_argument("--ehr-root", default=None) + parser.add_argument("--note-root", default=None) + parser.add_argument("--cache-dir", default=None) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the resolved command without running it.", + ) + return parser.parse_args() + + +def main() -> None: + cli = _parse_cli() + + config_path = Path(cli.config) + if not config_path.exists(): + raise SystemExit(f"Config not found: {config_path}") + + cfg = _resolve_config(config_path) + cfg = _apply_model_overrides(cfg) + + # Apply CLI overrides (non-None values win over YAML) + cli_overrides = { + "model": cli.model, + "seed": cli.seed, + "task": cli.task, + "epochs": cli.epochs, + "batch_size": cli.batch_size, + "embedding_dim": cli.embedding_dim, + "dev": cli.dev, + "freeze_encoder": cli.freeze_encoder if cli.freeze_encoder else None, + "balanced_sampling": cli.balanced_sampling if cli.balanced_sampling else None, + "icd_codes": cli.icd_codes if cli.icd_codes else None, + "output_dir": cli.output_dir, + "ehr_root": cli.ehr_root, + "note_root": cli.note_root, + "cache_dir": cli.cache_dir, + } + for k, v in cli_overrides.items(): + if v is not None: + cfg[k] = v + + script = Path(__file__).parent.parent / "examples" / "mortality_prediction" / "unified_embedding_e2e_mimic4.py" + cmd = [sys.executable, str(script)] + _to_args(cfg) + + print("Resolved command:") + print(" " + " \\\n ".join(cmd)) + + if cli.dry_run: + print("\n[dry-run] Not executing.") + return + + result = subprocess.run(cmd, env={**os.environ, "PYTHONPATH": str(script.parent.parent.parent)}) + sys.exit(result.returncode) + + +if __name__ == "__main__": + main() 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/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv b/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv index 804e2ae40..21511f015 100644 --- a/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv +++ b/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv @@ -19,6 +19,7 @@ subject_id,hadm_id,seq_num,icd_code,icd_version 10003,20006,2,E1065,10 10003,20006,3,N170,10 10003,20006,4,I509,10 +10003,20006,5,R99,10 10004,20007,1,K219,10 10004,20007,2,I10,10 10005,20008,1,25001,9 diff --git a/test-resources/core/mimic4demo/note/discharge.csv b/test-resources/core/mimic4demo/note/discharge.csv index 116c6d69e..a021ce257 100644 --- a/test-resources/core/mimic4demo/note/discharge.csv +++ b/test-resources/core/mimic4demo/note/discharge.csv @@ -6,7 +6,7 @@ d4,10001,20002,DS,1,2150-06-25 12:00:00,2150-06-25 13:00:00,Patient 10001 admitt d5,10002,20003,DS,1,2151-01-12 12:00:00,2151-01-12 13:00:00,Patient 10002 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. d6,10002,20004,DS,1,2151-04-20 12:00:00,2151-04-20 13:00:00,Patient 10002 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. d7,10003,20005,DS,1,2152-03-05 12:00:00,2152-03-05 13:00:00,Patient 10003 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. -d8,10003,20006,DS,1,2152-08-15 12:00:00,2152-08-15 13:00:00,Patient 10003 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. +d8,10003,20006,DS,1,2152-08-15 12:00:00,2152-08-15 13:00:00,Patient 10003 admitted for evaluation. Vitals unstable on presentation. Clinical status deteriorated despite aggressive management. Patient expired during this admission. Death pronounced on 2152-08-15. d9,10004,20007,DS,1,2150-05-02 12:00:00,2150-05-02 13:00:00,Patient 10004 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. d10,10005,20008,DS,1,2151-07-22 12:00:00,2151-07-22 13:00:00,Patient 10005 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. d11,10006,20009,DS,1,2152-09-08 12:00:00,2152-09-08 13:00:00,Patient 10006 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. diff --git a/tests/core/test_adacare.py b/tests/core/test_adacare.py index daf60ecee..84d2b9e6c 100644 --- a/tests/core/test_adacare.py +++ b/tests/core/test_adacare.py @@ -158,6 +158,5 @@ def test_model_with_embedding(self): expected_embed_dim = len(self.model.feature_keys) * self.model.hidden_dim self.assertEqual(ret["embed"].shape[1], expected_embed_dim) - if __name__ == "__main__": unittest.main() diff --git a/tests/core/test_biot.py b/tests/core/test_biot.py index 2f8a0e0a9..d51eecb41 100644 --- a/tests/core/test_biot.py +++ b/tests/core/test_biot.py @@ -13,7 +13,6 @@ def setUp(self): n_channels = 18 n_time = 10 n_fft = 200 - hop_length = 100 n_samples = n_fft * n_time # 2000 self.samples = [ @@ -62,7 +61,6 @@ def setUp(self): depth=4, n_fft=200, hop_length=100, - n_classes=6, n_channels=18, ) @@ -89,7 +87,8 @@ def test_model_forward(self): self.assertEqual(ret["y_prob"].shape[0], 2) self.assertEqual(ret["y_true"].shape[0], 2) self.assertEqual(ret["logit"].shape[0], 2) - self.assertEqual(ret["logit"].shape[1], 6) # n_classes + expected_size = self.dataset.output_processors["label"].size() + self.assertEqual(ret["logit"].shape[1], expected_size) self.assertEqual(ret["loss"].dim(), 0) def test_model_backward(self): @@ -104,12 +103,16 @@ def test_model_backward(self): param.requires_grad and param.grad is not None for param in self.model.parameters() ) - self.assertTrue(has_gradient, "No parameters have gradients after backward pass") + self.assertTrue( + has_gradient, "No parameters have gradients after backward pass" + ) def test_model_different_batch_sizes(self): """Test BIOT with different batch sizes.""" for batch_size in [1, 2, 4]: - train_loader = get_dataloader(self.dataset, batch_size=batch_size, shuffle=False) + train_loader = get_dataloader( + self.dataset, batch_size=batch_size, shuffle=False + ) data_batch = next(iter(train_loader)) with torch.no_grad(): @@ -139,17 +142,6 @@ def test_missing_signal_raises_error(self): def test_model_different_n_classes(self): """Test BIOT with different number of classes.""" - model_binary = BIOT( - dataset=self.dataset, - emb_size=256, - heads=8, - depth=4, - n_fft=200, - hop_length=100, - n_classes=1, - n_channels=18, - ) - binary_samples = [ { "patient_id": f"patient-{i}", @@ -167,6 +159,16 @@ def test_model_different_n_classes(self): dataset_name="test_cbramod_binary", ) + model_binary = BIOT( + dataset=binary_dataset, + emb_size=256, + heads=8, + depth=4, + n_fft=200, + hop_length=100, + n_channels=18, + ) + train_loader = get_dataloader(binary_dataset, batch_size=2, shuffle=True) data_batch = next(iter(train_loader)) with torch.no_grad(): @@ -183,7 +185,6 @@ def test_model(self): depth=2, n_fft=200, hop_length=100, - n_classes=6, n_channels=18, ) @@ -194,8 +195,9 @@ def test_model(self): ret = model_small(**data_batch) self.assertIn("loss", ret) - self.assertEqual(ret["logit"].shape[1], 6) + expected_size = self.dataset.output_processors["label"].size() + self.assertEqual(ret["logit"].shape[1], expected_size) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/core/test_cluster_label.py b/tests/core/test_cluster_label.py index 91b188090..3b63ba7ed 100644 --- a/tests/core/test_cluster_label.py +++ b/tests/core/test_cluster_label.py @@ -13,8 +13,15 @@ class TestClusterLabel(unittest.TestCase): def setUp(self): """Set up test data and model.""" - # Create samples with 3 classes for multiclass classification + # Stabilize model initialization and downstream calibration behavior. + np.random.seed(42) + torch.manual_seed(42) + + # Create samples with 3 classes for multiclass classification. + # 12 samples (2 per class for train, 2 per class for cal) ensure + # calibration samples are spread across clusters during K-means. self.samples = [ + # --- train set (indices 0–5) --- { "patient_id": "patient-0", "visit_id": "visit-0", @@ -57,6 +64,49 @@ def setUp(self): "procedures": [2.5, 3.5, 4.0, 5.5], "label": 2, }, + # --- calibration set (indices 6–11) --- + { + "patient_id": "patient-6", + "visit_id": "visit-6", + "conditions": ["cond-11", "cond-22", "cond-33"], + "procedures": [6.0, 1.0, 2.5, 3.5], + "label": 0, + }, + { + "patient_id": "patient-7", + "visit_id": "visit-7", + "conditions": ["cond-44", "cond-55", "cond-66"], + "procedures": [4.5, 5.5, 1.5, 2.0], + "label": 1, + }, + { + "patient_id": "patient-8", + "visit_id": "visit-8", + "conditions": ["cond-77", "cond-88"], + "procedures": [3.5, 6.5, 2.0, 1.0], + "label": 2, + }, + { + "patient_id": "patient-9", + "visit_id": "visit-9", + "conditions": ["cond-15", "cond-25", "cond-35", "cond-45"], + "procedures": [7.0, 1.5, 3.0, 4.0], + "label": 0, + }, + { + "patient_id": "patient-10", + "visit_id": "visit-10", + "conditions": ["cond-55", "cond-65"], + "procedures": [2.0, 5.0, 6.5, 1.5], + "label": 1, + }, + { + "patient_id": "patient-11", + "visit_id": "visit-11", + "conditions": ["cond-75", "cond-85", "cond-95"], + "procedures": [5.5, 2.5, 1.0, 3.0], + "label": 2, + }, ] # Define input and output schemas @@ -181,13 +231,13 @@ def test_calibrate_marginal(self): cluster_model = ClusterLabel( model=self.model, alpha=0.3, - n_clusters=3, + n_clusters=2, random_state=42, ) - # Split into train and cal sets - train_indices = [0, 1, 2] - cal_indices = [3, 4, 5] + # Split into train and cal sets (6 train, 6 cal) + train_indices = [0, 1, 2, 3, 4, 5] + cal_indices = [6, 7, 8, 9, 10, 11] train_dataset = self.dataset.subset(train_indices) cal_dataset = self.dataset.subset(cal_indices) @@ -203,15 +253,15 @@ def test_calibrate_marginal(self): # Check that K-means model is fitted self.assertIsNotNone(cluster_model.kmeans_model) - self.assertEqual(cluster_model.kmeans_model.n_clusters, 3) + self.assertEqual(cluster_model.kmeans_model.n_clusters, 2) # Check that cluster thresholds are set self.assertIsNotNone(cluster_model.cluster_thresholds) self.assertIsInstance(cluster_model.cluster_thresholds, dict) - self.assertEqual(len(cluster_model.cluster_thresholds), 3) + self.assertEqual(len(cluster_model.cluster_thresholds), 2) # Check that each cluster has a threshold - for cluster_id in range(3): + for cluster_id in range(2): self.assertIn(cluster_id, cluster_model.cluster_thresholds) threshold = cluster_model.cluster_thresholds[cluster_id] self.assertIsInstance(threshold, (float, np.floating)) @@ -226,9 +276,9 @@ def test_calibrate_class_conditional(self): random_state=42, ) - # Split into train and cal sets - train_indices = [0, 1, 2] - cal_indices = [3, 4, 5] + # Split into train and cal sets (6 train, 6 cal) + train_indices = [0, 1, 2, 3, 4, 5] + cal_indices = [6, 7, 8, 9, 10, 11] train_dataset = self.dataset.subset(train_indices) cal_dataset = self.dataset.subset(cal_indices) @@ -254,13 +304,13 @@ def test_forward_returns_predset(self): cluster_model = ClusterLabel( model=self.model, alpha=0.2, - n_clusters=3, + n_clusters=2, random_state=42, ) # Calibrate - train_indices = [0, 1, 2] - cal_indices = [3, 4, 5] + train_indices = [0, 1, 2, 3, 4, 5] + cal_indices = [6, 7, 8, 9, 10, 11] train_dataset = self.dataset.subset(train_indices) cal_dataset = self.dataset.subset(cal_indices) @@ -301,8 +351,8 @@ def test_prediction_sets_nonempty(self): ) # Calibrate - train_indices = [0, 1, 2] - cal_indices = [3, 4, 5] + train_indices = [0, 1, 2, 3, 4, 5] + cal_indices = [6, 7, 8, 9, 10, 11] train_dataset = self.dataset.subset(train_indices) cal_dataset = self.dataset.subset(cal_indices) @@ -335,7 +385,7 @@ def test_calibrate_requires_train_embeddings(self): n_clusters=3, ) - cal_indices = [3, 4, 5] + cal_indices = [6, 7, 8, 9, 10, 11] cal_dataset = self.dataset.subset(cal_indices) cal_embeddings = self._get_embeddings(cal_dataset) @@ -362,8 +412,13 @@ def test_forward_before_calibration_raises_error(self): cluster_model(**data_batch) def test_different_cluster_counts(self): - """Test that different cluster counts work.""" - for n_clusters in [2, 3, 4]: + """Test that different cluster counts work. + + Cluster counts are capped at 3 (half of the 6 cal samples) so that + K-means can reliably assign at least one calibration sample per cluster + with a small test dataset. + """ + for n_clusters in [2, 3]: cluster_model = ClusterLabel( model=self.model, alpha=0.2, @@ -371,8 +426,8 @@ def test_different_cluster_counts(self): random_state=42, ) - train_indices = [0, 1, 2] - cal_indices = [3, 4, 5] + train_indices = [0, 1, 2, 3, 4, 5] + cal_indices = [6, 7, 8, 9, 10, 11] train_dataset = self.dataset.subset(train_indices) cal_dataset = self.dataset.subset(cal_indices) @@ -395,12 +450,12 @@ def test_model_device_handling(self): cluster_model = ClusterLabel( model=self.model, alpha=0.2, - n_clusters=3, + n_clusters=2, random_state=42, ) - train_indices = [0, 1, 2] - cal_indices = [3, 4, 5] + train_indices = [0, 1, 2, 3, 4, 5] + cal_indices = [6, 7, 8, 9, 10, 11] train_dataset = self.dataset.subset(train_indices) cal_dataset = self.dataset.subset(cal_indices) diff --git a/tests/core/test_multimodal_adacare.py b/tests/core/test_multimodal_adacare.py new file mode 100644 index 000000000..a15546b66 --- /dev/null +++ b/tests/core/test_multimodal_adacare.py @@ -0,0 +1,311 @@ +import unittest +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import MultimodalAdaCare + + +class TestMultimodalAdaCare(unittest.TestCase): + """Test cases for the MultimodalAdaCare model.""" + + def setUp(self): + """Set up test data and model with mixed feature types.""" + # Samples with mixed sequential and non-sequential features + self.samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "conditions": ["cond-33", "cond-86", "cond-80"], # sequential + "procedures": ["proc-12", "proc-45"], # sequential + "demographics": ["asian", "male", "smoker"], # multi-hot + "vitals": [120.0, 80.0, 98.6, 16.0], # tensor + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "conditions": ["cond-12", "cond-52"], # sequential + "procedures": ["proc-23"], # sequential + "demographics": ["white", "female"], # multi-hot + "vitals": [110.0, 75.0, 98.2, 18.0], # tensor + "label": 0, + }, + ] + + # Define input and output schemas with mixed types + self.input_schema = { + "conditions": "sequence", # sequential + "procedures": "sequence", # sequential + "demographics": "multi_hot", # non-sequential + "vitals": "tensor", # non-sequential + } + self.output_schema = {"label": "binary"} + + # Create dataset + self.dataset = create_sample_dataset( + samples=self.samples, + input_schema=self.input_schema, + output_schema=self.output_schema, + dataset_name="test", + ) + + # Create model + self.model = MultimodalAdaCare(dataset=self.dataset) + + def test_model_initialization(self): + """Test that the MultimodalAdaCare model initializes correctly.""" + self.assertIsInstance(self.model, MultimodalAdaCare) + self.assertEqual(self.model.embedding_dim, 128) + self.assertEqual(self.model.hidden_dim, 128) + self.assertEqual(len(self.model.feature_keys), 4) + + # Check that features are correctly classified + self.assertIn("conditions", self.model.sequential_features) + self.assertIn("procedures", self.model.sequential_features) + self.assertIn("demographics", self.model.non_sequential_features) + self.assertIn("vitals", self.model.non_sequential_features) + + # Check that AdaCare layers are only created for sequential features + self.assertIn("conditions", self.model.adacare) + self.assertIn("procedures", self.model.adacare) + self.assertNotIn("demographics", self.model.adacare) + self.assertNotIn("vitals", self.model.adacare) + + def test_model_forward(self): + """Test that the MultimodalAdaCare model forward pass works correctly.""" + # Create data loader + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + # Forward pass + with torch.no_grad(): + ret = self.model(**data_batch) + + # Check output structure + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + self.assertIn("y_true", ret) + self.assertIn("logit", ret) + self.assertIn("feature_importance", ret) + self.assertIn("conv_feature_importance", ret) + + # Check tensor shapes + self.assertEqual(ret["y_prob"].shape[0], 2) # batch size + self.assertEqual(ret["y_true"].shape[0], 2) # batch size + self.assertEqual(ret["logit"].shape[0], 2) # batch size + + # Check that loss is a scalar + self.assertEqual(ret["loss"].dim(), 0) + + # Check feature importance outputs + self.assertEqual(len(ret["feature_importance"]), 2) # 2 sequential features + self.assertEqual(len(ret["conv_feature_importance"]), 2) # 2 sequential + + def test_model_backward(self): + """Test that the MultimodalAdaCare model backward pass works correctly.""" + # Create data loader + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + # Forward pass + ret = self.model(**data_batch) + + # Backward pass + ret["loss"].backward() + + # Check that at least one parameter has gradients + has_gradient = False + for param in self.model.parameters(): + if param.requires_grad and param.grad is not None: + has_gradient = True + break + self.assertTrue( + has_gradient, "No parameters have gradients after backward pass" + ) + + def test_model_with_embedding(self): + """Test that the MultimodalAdaCare model returns embeddings when requested.""" + # Create data loader + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + data_batch["embed"] = True + + # Forward pass + with torch.no_grad(): + ret = self.model(**data_batch) + + # Check that embeddings are returned + self.assertIn("embed", ret) + self.assertEqual(ret["embed"].shape[0], 2) # batch size + + # Check embedding dimension + # 2 sequential features * hidden_dim + 2 non-sequential features * embedding_dim + expected_embed_dim = ( + len(self.model.sequential_features) * self.model.hidden_dim + + len(self.model.non_sequential_features) * self.model.embedding_dim + ) + self.assertEqual(ret["embed"].shape[1], expected_embed_dim) + + def test_custom_hyperparameters(self): + """Test MultimodalAdaCare model with custom hyperparameters.""" + model = MultimodalAdaCare( + dataset=self.dataset, + embedding_dim=64, + hidden_dim=32, + kernel_size=2, + kernel_num=32, + r_v=2, + r_c=2, + activation="sparsemax", + rnn_type="lstm", + dropout=0.3, + ) + + self.assertEqual(model.embedding_dim, 64) + self.assertEqual(model.hidden_dim, 32) + + # Test forward pass + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = model(**data_batch) + + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + self.assertIn("feature_importance", ret) + + def test_only_sequential_features(self): + """Test MultimodalAdaCare with only sequential features.""" + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "conditions": ["cond-33", "cond-86"], + "procedures": ["proc-12", "proc-45"], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "conditions": ["cond-12"], + "procedures": ["proc-23"], + "label": 0, + }, + ] + + dataset = create_sample_dataset( + samples=samples, + input_schema={"conditions": "sequence", "procedures": "sequence"}, + output_schema={"label": "binary"}, + dataset_name="test_seq_only", + ) + + model = MultimodalAdaCare(dataset=dataset, hidden_dim=64) + + # Check that all features are sequential + self.assertEqual(len(model.sequential_features), 2) + self.assertEqual(len(model.non_sequential_features), 0) + + # Test forward pass + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = model(**data_batch) + + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + self.assertIn("feature_importance", ret) + + def test_only_non_sequential_features(self): + """Test MultimodalAdaCare with only non-sequential features.""" + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "demographics": ["asian", "male", "smoker"], + "vitals": [120.0, 80.0, 98.6, 16.0], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "demographics": ["white", "female"], + "vitals": [110.0, 75.0, 98.2, 18.0], + "label": 0, + }, + ] + + dataset = create_sample_dataset( + samples=samples, + input_schema={"demographics": "multi_hot", "vitals": "tensor"}, + output_schema={"label": "binary"}, + dataset_name="test_non_seq_only", + ) + + model = MultimodalAdaCare(dataset=dataset, hidden_dim=64) + + # Check that all features are non-sequential + self.assertEqual(len(model.sequential_features), 0) + self.assertEqual(len(model.non_sequential_features), 2) + + # Test forward pass + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = model(**data_batch) + + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + # No feature importance when no sequential features + self.assertEqual(len(ret["feature_importance"]), 0) + self.assertEqual(len(ret["conv_feature_importance"]), 0) + + def test_output_shapes(self): + """Test that output shapes are correct for multimodal inputs.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = self.model(**data_batch) + + self.assertEqual(ret["y_prob"].shape, (2, 1)) + self.assertEqual(ret["y_true"].shape, (2, 1)) + self.assertEqual(ret["logit"].shape, (2, 1)) + + def test_loss_is_finite(self): + """Test that the loss is finite.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = self.model(**data_batch) + + self.assertTrue(torch.isfinite(ret["loss"]).all()) + + def test_sequential_processor_classification(self): + """Test that _is_sequential_processor correctly identifies processor types.""" + from pyhealth.processors import ( + MultiHotProcessor, + SequenceProcessor, + TensorProcessor, + ) + + # Test with actual processor instances + seq_proc = SequenceProcessor() + self.assertTrue(self.model._is_sequential_processor(seq_proc)) + + # Create simple multi-hot processor + multihot_proc = MultiHotProcessor() + self.assertFalse(self.model._is_sequential_processor(multihot_proc)) + + # Tensor processor + tensor_proc = TensorProcessor() + self.assertFalse(self.model._is_sequential_processor(tensor_proc)) + + +if __name__ == "__main__": + unittest.main() + diff --git a/tests/core/test_multimodal_retain.py b/tests/core/test_multimodal_retain.py new file mode 100644 index 000000000..00ed039be --- /dev/null +++ b/tests/core/test_multimodal_retain.py @@ -0,0 +1,334 @@ +import unittest +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import MultimodalRETAIN + + +class TestMultimodalRETAIN(unittest.TestCase): + """Test cases for the MultimodalRETAIN model.""" + + def setUp(self): + """Set up test data and model with mixed feature types.""" + # Samples with mixed sequential and non-sequential features + # RETAIN typically works with visit-level nested sequences + self.samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "conditions": [["A", "B"], ["C"]], # nested sequence + "procedures": [["P1"], ["P2", "P3"]], # nested sequence + "demographics": ["asian", "male"], # multi-hot + "vitals": [120.0, 80.0, 98.6], # tensor + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "conditions": [["D"], ["E", "F"]], # nested sequence + "procedures": [["P4"]], # nested sequence + "demographics": ["white", "female"], # multi-hot + "vitals": [110.0, 75.0, 98.2], # tensor + "label": 0, + }, + ] + + # Define input and output schemas with mixed types + self.input_schema = { + "conditions": "nested_sequence", # sequential + "procedures": "nested_sequence", # sequential + "demographics": "multi_hot", # non-sequential + "vitals": "tensor", # non-sequential + } + self.output_schema = {"label": "binary"} + + # Create dataset + self.dataset = create_sample_dataset( + samples=self.samples, + input_schema=self.input_schema, + output_schema=self.output_schema, + dataset_name="test", + ) + + # Create model + self.model = MultimodalRETAIN(dataset=self.dataset) + + def test_model_initialization(self): + """Test that the MultimodalRETAIN model initializes correctly.""" + self.assertIsInstance(self.model, MultimodalRETAIN) + self.assertEqual(self.model.embedding_dim, 128) + self.assertEqual(len(self.model.feature_keys), 4) + + # Check that features are correctly classified + self.assertIn("conditions", self.model.sequential_features) + self.assertIn("procedures", self.model.sequential_features) + self.assertIn("demographics", self.model.non_sequential_features) + self.assertIn("vitals", self.model.non_sequential_features) + + # Check that RETAIN layers are only created for sequential features + self.assertIn("conditions", self.model.retain) + self.assertIn("procedures", self.model.retain) + self.assertNotIn("demographics", self.model.retain) + self.assertNotIn("vitals", self.model.retain) + + def test_model_forward(self): + """Test that the MultimodalRETAIN model forward pass works correctly.""" + # Create data loader + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + # Forward pass + with torch.no_grad(): + ret = self.model(**data_batch) + + # Check output structure + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + self.assertIn("y_true", ret) + self.assertIn("logit", ret) + + # Check tensor shapes + self.assertEqual(ret["y_prob"].shape[0], 2) # batch size + self.assertEqual(ret["y_true"].shape[0], 2) # batch size + self.assertEqual(ret["logit"].shape[0], 2) # batch size + + # Check that loss is a scalar + self.assertEqual(ret["loss"].dim(), 0) + + def test_model_backward(self): + """Test that the MultimodalRETAIN model backward pass works correctly.""" + # Create data loader + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + # Forward pass + ret = self.model(**data_batch) + + # Backward pass + ret["loss"].backward() + + # Check that at least one parameter has gradients + has_gradient = False + for param in self.model.parameters(): + if param.requires_grad and param.grad is not None: + has_gradient = True + break + self.assertTrue( + has_gradient, "No parameters have gradients after backward pass" + ) + + def test_model_with_embedding(self): + """Test that the MultimodalRETAIN model returns embeddings when requested.""" + # Create data loader + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + data_batch["embed"] = True + + # Forward pass + with torch.no_grad(): + ret = self.model(**data_batch) + + # Check that embeddings are returned + self.assertIn("embed", ret) + self.assertEqual(ret["embed"].shape[0], 2) # batch size + + # Check embedding dimension + # All features contribute embedding_dim + expected_embed_dim = len(self.model.feature_keys) * self.model.embedding_dim + self.assertEqual(ret["embed"].shape[1], expected_embed_dim) + + def test_custom_hyperparameters(self): + """Test MultimodalRETAIN model with custom hyperparameters.""" + model = MultimodalRETAIN( + dataset=self.dataset, + embedding_dim=64, + dropout=0.3, + ) + + self.assertEqual(model.embedding_dim, 64) + + # Test forward pass + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = model(**data_batch) + + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + + def test_only_sequential_features(self): + """Test MultimodalRETAIN with only sequential features.""" + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "conditions": [["A", "B"], ["C"]], + "procedures": [["P1"], ["P2"]], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "conditions": [["D"], ["E"]], + "procedures": [["P3"]], + "label": 0, + }, + ] + + dataset = create_sample_dataset( + samples=samples, + input_schema={ + "conditions": "nested_sequence", + "procedures": "nested_sequence" + }, + output_schema={"label": "binary"}, + dataset_name="test_seq_only", + ) + + model = MultimodalRETAIN(dataset=dataset) + + # Check that all features are sequential + self.assertEqual(len(model.sequential_features), 2) + self.assertEqual(len(model.non_sequential_features), 0) + + # Test forward pass + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = model(**data_batch) + + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + + def test_only_non_sequential_features(self): + """Test MultimodalRETAIN with only non-sequential features.""" + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "demographics": ["asian", "male"], + "vitals": [120.0, 80.0, 98.6], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "demographics": ["white", "female"], + "vitals": [110.0, 75.0, 98.2], + "label": 0, + }, + ] + + dataset = create_sample_dataset( + samples=samples, + input_schema={"demographics": "multi_hot", "vitals": "tensor"}, + output_schema={"label": "binary"}, + dataset_name="test_non_seq_only", + ) + + model = MultimodalRETAIN(dataset=dataset) + + # Check that all features are non-sequential + self.assertEqual(len(model.sequential_features), 0) + self.assertEqual(len(model.non_sequential_features), 2) + + # Test forward pass + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = model(**data_batch) + + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + + def test_output_shapes(self): + """Test that output shapes are correct for multimodal inputs.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = self.model(**data_batch) + + self.assertEqual(ret["y_prob"].shape, (2, 1)) + self.assertEqual(ret["y_true"].shape, (2, 1)) + self.assertEqual(ret["logit"].shape, (2, 1)) + + def test_loss_is_finite(self): + """Test that the loss is finite.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = self.model(**data_batch) + + self.assertTrue(torch.isfinite(ret["loss"]).all()) + + def test_sequential_processor_classification(self): + """Test that _is_sequential_processor correctly identifies processor types.""" + from pyhealth.processors import ( + MultiHotProcessor, + NestedSequenceProcessor, + TensorProcessor, + ) + + # Test with actual processor instances + nested_seq_proc = NestedSequenceProcessor() + self.assertTrue(self.model._is_sequential_processor(nested_seq_proc)) + + # Create simple multi-hot processor + multihot_proc = MultiHotProcessor() + self.assertFalse(self.model._is_sequential_processor(multihot_proc)) + + # Tensor processor + tensor_proc = TensorProcessor() + self.assertFalse(self.model._is_sequential_processor(tensor_proc)) + + def test_with_simple_sequences(self): + """Test MultimodalRETAIN with simple (non-nested) sequences.""" + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "codes": ["A", "B", "C"], # simple sequence + "demographics": ["asian", "male"], # multi-hot + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "codes": ["D", "E"], # simple sequence + "demographics": ["white", "female"], # multi-hot + "label": 0, + }, + ] + + dataset = create_sample_dataset( + samples=samples, + input_schema={"codes": "sequence", "demographics": "multi_hot"}, + output_schema={"label": "binary"}, + dataset_name="test_simple_seq", + ) + + model = MultimodalRETAIN(dataset=dataset) + + # Check that codes is sequential + self.assertIn("codes", model.sequential_features) + self.assertIn("demographics", model.non_sequential_features) + + # Test forward pass + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = model(**data_batch) + + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + + +if __name__ == "__main__": + unittest.main() + diff --git a/tests/core/test_multimodal_rnn.py b/tests/core/test_multimodal_rnn.py index d2fe5fa0b..560ff6703 100644 --- a/tests/core/test_multimodal_rnn.py +++ b/tests/core/test_multimodal_rnn.py @@ -273,6 +273,118 @@ def test_sequential_processor_classification(self): self.assertFalse(self.model._is_sequential_processor(tensor_proc)) +class TestMultimodalRNNNestedSequence(unittest.TestCase): + """Tests for MultimodalRNN with nested_sequence features. + + Covers the bug where drugs_hist = [[]] (first sample in drug recommendation + tasks) caused: RuntimeError: Length of all samples has to be greater than 0 + because MultimodalRNN was flattening visits*codes into a single sequence + instead of pooling codes within visits first. + """ + + def _make_drug_rec_samples(self): + """Mimics DrugRecommendationMIMIC4 output for a 2-patient dataset. + + samples[0] for each patient always has drugs_hist = [[]] — one visit + with an empty inner list, because the task zeroes out the current + visit's drugs from the history. + """ + return [ + # patient-0, visit-0: first ever visit → drugs_hist has one empty inner list + { + "patient_id": "p0", "visit_id": "v0", + "conditions": [["cond-1", "cond-2"]], + "drugs_hist": [[]], # <-- the problematic case + "label": ["drug-A"], + }, + # patient-0, visit-1: second visit → current slot cleared + { + "patient_id": "p0", "visit_id": "v1", + "conditions": [["cond-1", "cond-2"], ["cond-3"]], + "drugs_hist": [["drug-A"], []], + "label": ["drug-B"], + }, + # patient-1, visit-0: another first-visit case + { + "patient_id": "p1", "visit_id": "v2", + "conditions": [["cond-4"]], + "drugs_hist": [[]], + "label": ["drug-A"], + }, + # patient-1, visit-1 + { + "patient_id": "p1", "visit_id": "v3", + "conditions": [["cond-4"], ["cond-5", "cond-6"]], + "drugs_hist": [["drug-B"], []], + "label": ["drug-B"], + }, + ] + + def test_forward_with_empty_inner_list(self): + """MultimodalRNN must not crash when a nested_sequence has empty inner lists. + + Before the fix, MultimodalRNN flattened (visits * max_codes) into a single + sequence dimension, giving length=0 for patients whose inner lists were all + empty (e.g. drugs_hist=[[]] for first-visit samples). + """ + dataset = create_sample_dataset( + samples=self._make_drug_rec_samples(), + input_schema={ + "conditions": "nested_sequence", + "drugs_hist": "nested_sequence", + }, + output_schema={"label": "multilabel"}, + dataset_name="test_empty_inner", + ) + loader = get_dataloader(dataset, batch_size=4, shuffle=False) + model = MultimodalRNN(dataset=dataset, embedding_dim=32, hidden_dim=32) + model.eval() + + batch = next(iter(loader)) + with torch.no_grad(): + result = model(**batch) + + self.assertIn("loss", result) + self.assertIn("y_prob", result) + self.assertEqual(result["y_prob"].shape[0], 4) + + def test_nested_sequence_visit_level_mask(self): + """Verify the fixed MultimodalRNN uses visit-level masks (not code-level). + + With visit-level masking, a patient with 2 visits where the second is + empty should have sequence length 1, not 0. + """ + samples = [ + # One non-empty visit followed by one empty visit + { + "patient_id": "p0", "visit_id": "v0", + "conditions": [["cond-1"], []], + "label": 1, + }, + { + "patient_id": "p1", "visit_id": "v1", + "conditions": [["cond-2", "cond-3"], ["cond-4"]], + "label": 0, + }, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"conditions": "nested_sequence"}, + output_schema={"label": "binary"}, + dataset_name="test_visit_mask", + ) + loader = get_dataloader(dataset, batch_size=2, shuffle=False) + model = MultimodalRNN(dataset=dataset, embedding_dim=16, hidden_dim=16) + model.eval() + + batch = next(iter(loader)) + with torch.no_grad(): + result = model(**batch) + + self.assertIn("loss", result) + self.assertEqual(result["y_prob"].shape[0], 2) + + if __name__ == "__main__": unittest.main() diff --git a/tests/core/test_retain.py b/tests/core/test_retain.py new file mode 100644 index 000000000..5b21180a6 --- /dev/null +++ b/tests/core/test_retain.py @@ -0,0 +1,136 @@ +import unittest +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import RETAIN + + +class TestRETAIN(unittest.TestCase): + """Test cases for the RETAIN model.""" + + def setUp(self): + """Set up test data and model.""" + self.samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "conditions": [["A", "B"], ["C", "D", "E"]], + "procedures": [["P1"], ["P2", "P3"]], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "conditions": [["F"], ["G", "H"]], + "procedures": [["P4", "P5"], ["P6"]], + "label": 0, + }, + ] + + self.input_schema = { + "conditions": "nested_sequence", + "procedures": "nested_sequence", + } + self.output_schema = {"label": "binary"} + + self.dataset = create_sample_dataset( + samples=self.samples, + input_schema=self.input_schema, + output_schema=self.output_schema, + dataset_name="test", + ) + + self.model = RETAIN(dataset=self.dataset) + + def test_model_initialization(self): + """Test that the RETAIN model initializes correctly.""" + self.assertIsInstance(self.model, RETAIN) + self.assertEqual(self.model.embedding_dim, 128) + self.assertEqual(len(self.model.feature_keys), 2) + self.assertIn("conditions", self.model.feature_keys) + self.assertIn("procedures", self.model.feature_keys) + self.assertEqual(self.model.label_keys[0], "label") + + def test_forward_input_format(self): + """Test that the dataloader provides tensor inputs.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=False) + data_batch = next(iter(train_loader)) + + self.assertIsInstance(data_batch["conditions"], torch.Tensor) + self.assertIsInstance(data_batch["procedures"], torch.Tensor) + self.assertIsInstance(data_batch["label"], torch.Tensor) + + def test_model_forward(self): + """Test that the RETAIN model forward pass works correctly.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = self.model(**data_batch) + + self.assertIn("loss", ret) + self.assertIn("y_prob", ret) + self.assertIn("y_true", ret) + self.assertIn("logit", ret) + + self.assertEqual(ret["y_prob"].shape[0], 2) + self.assertEqual(ret["y_true"].shape[0], 2) + self.assertEqual(ret["loss"].dim(), 0) + + def test_model_backward(self): + """Test that the RETAIN model backward pass works correctly.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + ret = self.model(**data_batch) + ret["loss"].backward() + + has_gradient = any( + param.requires_grad and param.grad is not None + for param in self.model.parameters() + ) + self.assertTrue(has_gradient, "No parameters have gradients after backward pass") + + def test_loss_is_finite(self): + """Test that the loss is finite.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = self.model(**data_batch) + + self.assertTrue(torch.isfinite(ret["loss"]).all()) + + def test_output_shapes(self): + """Test that output shapes are correct.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + with torch.no_grad(): + ret = self.model(**data_batch) + + batch_size = 2 + num_labels = 1 # binary classification + + self.assertEqual(ret["y_prob"].shape, (batch_size, num_labels)) + self.assertEqual(ret["y_true"].shape, (batch_size, num_labels)) + self.assertEqual(ret["logit"].shape, (batch_size, num_labels)) + + def test_model_with_embedding(self): + """Test that the RETAIN model returns embeddings when requested.""" + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True) + data_batch = next(iter(train_loader)) + + data_batch["embed"] = True + + with torch.no_grad(): + ret = self.model(**data_batch) + + self.assertIn("embed", ret) + self.assertEqual(ret["embed"].shape[0], 2) # batch size + expected_embed_dim = len(self.model.feature_keys) * self.model.embedding_dim + self.assertEqual(ret["embed"].shape[1], expected_embed_dim) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_time_image_processor.py b/tests/core/test_time_image_processor.py index 51467aa5c..b1eb46cc2 100644 --- a/tests/core/test_time_image_processor.py +++ b/tests/core/test_time_image_processor.py @@ -281,13 +281,14 @@ def test_size_set_after_process(self): def test_repr(self): """__repr__ contains key parameter values.""" proc = TimeImageProcessor( - image_size=128, max_images=5, mode="L" + image_size=128, max_images=5, mode="L", padding="" ) r = repr(proc) self.assertIn("TimeImageProcessor", r) self.assertIn("image_size=128", r) self.assertIn("max_images=5", r) self.assertIn("mode=L", r) + self.assertIn("padding=", r) # ---- Tensor properties ---- @@ -310,6 +311,18 @@ def test_timestamp_dtype_is_float32(self): self.assertEqual(timestamps.dtype, torch.float32) + # ---- padding ---- + + def test_padding_returns_zero_tensor(self): + """Padding path returns a zero tensor instead of loading.""" + proc = TimeImageProcessor(image_size=32, padding="") + images, timestamps, tag = proc.process( + (["", self.rgb_paths[0]], [0.0, 1.0]) + ) + + self.assertEqual(images.shape, (2, 3, 32, 32)) + self.assertTrue(torch.all(images[0] == 0)) + if __name__ == "__main__": unittest.main() \ No newline at end of file diff --git a/tests/core/test_tuab.py b/tests/core/test_tuab.py index 1104d6f86..bf6e25d45 100644 --- a/tests/core/test_tuab.py +++ b/tests/core/test_tuab.py @@ -89,7 +89,9 @@ def test_prepare_metadata_creates_expected_csvs(self): self.assertEqual(normal_row.iloc[0]["record_id"], "s004_t000") self.assertEqual(normal_row.iloc[0]["label"], "normal") self.assertTrue( - str(normal_row.iloc[0]["signal_file"]).endswith("aaaaaaav_s004_t000.edf") + str(normal_row.iloc[0]["signal_file"]).endswith( + "aaaaaaav_s004_t000.edf" + ) ) # Verify an abnormal train entry @@ -228,10 +230,22 @@ def test_convert_to_bipolar_output_shape_and_values(self): def test_convert_to_bipolar_middle_channels(self): """Spot-check a middle channel (channel 9: F3 - C3).""" ch_names = [ - "EEG FP1-REF", "EEG F7-REF", "EEG T3-REF", "EEG T5-REF", - "EEG O1-REF", "EEG FP2-REF", "EEG F8-REF", "EEG T4-REF", - "EEG T6-REF", "EEG O2-REF", "EEG F3-REF", "EEG C3-REF", - "EEG P3-REF", "EEG F4-REF", "EEG C4-REF", "EEG P4-REF", + "EEG FP1-REF", + "EEG F7-REF", + "EEG T3-REF", + "EEG T5-REF", + "EEG O1-REF", + "EEG FP2-REF", + "EEG F8-REF", + "EEG T4-REF", + "EEG T6-REF", + "EEG O2-REF", + "EEG F3-REF", + "EEG C3-REF", + "EEG P3-REF", + "EEG F4-REF", + "EEG C4-REF", + "EEG P4-REF", ] n = 100 @@ -250,10 +264,22 @@ def test_call_segments_into_10s_windows(self): # 25 seconds of data => 2 full 10-second windows (last 5s discarded) n_samples = fs * 25 ch_names = [ - "EEG FP1-REF", "EEG F7-REF", "EEG T3-REF", "EEG T5-REF", - "EEG O1-REF", "EEG FP2-REF", "EEG F8-REF", "EEG T4-REF", - "EEG T6-REF", "EEG O2-REF", "EEG F3-REF", "EEG C3-REF", - "EEG P3-REF", "EEG F4-REF", "EEG C4-REF", "EEG P4-REF", + "EEG FP1-REF", + "EEG F7-REF", + "EEG T3-REF", + "EEG T5-REF", + "EEG O1-REF", + "EEG FP2-REF", + "EEG F8-REF", + "EEG T4-REF", + "EEG T6-REF", + "EEG O2-REF", + "EEG F3-REF", + "EEG C3-REF", + "EEG P3-REF", + "EEG F4-REF", + "EEG C4-REF", + "EEG P4-REF", ] raw_data = np.random.randn(len(ch_names), n_samples) @@ -285,10 +311,22 @@ def test_call_label_mapping_normal(self): fs = 200 n_samples = fs * 10 # exactly one window ch_names = [ - "EEG FP1-REF", "EEG F7-REF", "EEG T3-REF", "EEG T5-REF", - "EEG O1-REF", "EEG FP2-REF", "EEG F8-REF", "EEG T4-REF", - "EEG T6-REF", "EEG O2-REF", "EEG F3-REF", "EEG C3-REF", - "EEG P3-REF", "EEG F4-REF", "EEG C4-REF", "EEG P4-REF", + "EEG FP1-REF", + "EEG F7-REF", + "EEG T3-REF", + "EEG T5-REF", + "EEG O1-REF", + "EEG FP2-REF", + "EEG F8-REF", + "EEG T4-REF", + "EEG T6-REF", + "EEG O2-REF", + "EEG F3-REF", + "EEG C3-REF", + "EEG P3-REF", + "EEG F4-REF", + "EEG C4-REF", + "EEG P4-REF", ] raw_data = np.random.randn(len(ch_names), n_samples) @@ -314,10 +352,22 @@ def test_call_label_mapping_abnormal(self): fs = 200 n_samples = fs * 10 ch_names = [ - "EEG FP1-REF", "EEG F7-REF", "EEG T3-REF", "EEG T5-REF", - "EEG O1-REF", "EEG FP2-REF", "EEG F8-REF", "EEG T4-REF", - "EEG T6-REF", "EEG O2-REF", "EEG F3-REF", "EEG C3-REF", - "EEG P3-REF", "EEG F4-REF", "EEG C4-REF", "EEG P4-REF", + "EEG FP1-REF", + "EEG F7-REF", + "EEG T3-REF", + "EEG T5-REF", + "EEG O1-REF", + "EEG FP2-REF", + "EEG F8-REF", + "EEG T4-REF", + "EEG T6-REF", + "EEG O2-REF", + "EEG F3-REF", + "EEG C3-REF", + "EEG P3-REF", + "EEG F4-REF", + "EEG C4-REF", + "EEG P4-REF", ] raw_data = np.random.randn(len(ch_names), n_samples) @@ -343,10 +393,22 @@ def test_call_too_short_recording_returns_no_samples(self): fs = 200 n_samples = fs * 5 # only 5 seconds ch_names = [ - "EEG FP1-REF", "EEG F7-REF", "EEG T3-REF", "EEG T5-REF", - "EEG O1-REF", "EEG FP2-REF", "EEG F8-REF", "EEG T4-REF", - "EEG T6-REF", "EEG O2-REF", "EEG F3-REF", "EEG C3-REF", - "EEG P3-REF", "EEG F4-REF", "EEG C4-REF", "EEG P4-REF", + "EEG FP1-REF", + "EEG F7-REF", + "EEG T3-REF", + "EEG T5-REF", + "EEG O1-REF", + "EEG FP2-REF", + "EEG F8-REF", + "EEG T4-REF", + "EEG T6-REF", + "EEG O2-REF", + "EEG F3-REF", + "EEG C3-REF", + "EEG P3-REF", + "EEG F4-REF", + "EEG C4-REF", + "EEG P4-REF", ] raw_data = np.random.randn(len(ch_names), n_samples) @@ -371,10 +433,22 @@ def test_call_multiple_events_per_patient(self): fs = 200 n_samples = fs * 20 # 2 windows each ch_names = [ - "EEG FP1-REF", "EEG F7-REF", "EEG T3-REF", "EEG T5-REF", - "EEG O1-REF", "EEG FP2-REF", "EEG F8-REF", "EEG T4-REF", - "EEG T6-REF", "EEG O2-REF", "EEG F3-REF", "EEG C3-REF", - "EEG P3-REF", "EEG F4-REF", "EEG C4-REF", "EEG P4-REF", + "EEG FP1-REF", + "EEG F7-REF", + "EEG T3-REF", + "EEG T5-REF", + "EEG O1-REF", + "EEG FP2-REF", + "EEG F8-REF", + "EEG T4-REF", + "EEG T6-REF", + "EEG O2-REF", + "EEG F3-REF", + "EEG C3-REF", + "EEG P3-REF", + "EEG F4-REF", + "EEG C4-REF", + "EEG P4-REF", ] raw_data = np.random.randn(len(ch_names), n_samples) @@ -409,10 +483,22 @@ def test_call_segment_ids_are_sequential(self): fs = 200 n_samples = fs * 30 # 3 windows ch_names = [ - "EEG FP1-REF", "EEG F7-REF", "EEG T3-REF", "EEG T5-REF", - "EEG O1-REF", "EEG FP2-REF", "EEG F8-REF", "EEG T4-REF", - "EEG T6-REF", "EEG O2-REF", "EEG F3-REF", "EEG C3-REF", - "EEG P3-REF", "EEG F4-REF", "EEG C4-REF", "EEG P4-REF", + "EEG FP1-REF", + "EEG F7-REF", + "EEG T3-REF", + "EEG T5-REF", + "EEG O1-REF", + "EEG FP2-REF", + "EEG F8-REF", + "EEG T4-REF", + "EEG T6-REF", + "EEG O2-REF", + "EEG F3-REF", + "EEG C3-REF", + "EEG P3-REF", + "EEG F4-REF", + "EEG C4-REF", + "EEG P4-REF", ] raw_data = np.random.randn(len(ch_names), n_samples) @@ -453,4 +539,4 @@ def test_task_schema_no_stft(self): if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/test_unified_multimodal.py b/tests/test_unified_multimodal.py index 22c88785c..f17e0f20e 100644 --- a/tests/test_unified_multimodal.py +++ b/tests/test_unified_multimodal.py @@ -2,49 +2,73 @@ collate_temporal helper, and UnifiedMultimodalEmbeddingModel. Run with: - TOKENIZERS_PARALLELISM=false pytest tests/test_unified_multimodal.py -v + TOKENIZERS_PARALLELISM=false python tests/test_unified_multimodal.py """ + import math +import unittest from datetime import datetime, timedelta +from unittest.mock import patch -import pytest import torch import numpy as np # ── 1. TemporalFeatureProcessor ABC & ModalityType ──────────────────────────── + def test_modality_type_values(): from pyhealth.processors import ModalityType - assert ModalityType.CODE == "code" - assert ModalityType.TEXT == "text" - assert ModalityType.IMAGE == "image" + + assert ModalityType.CODE == "code" + assert ModalityType.TEXT == "text" + assert ModalityType.IMAGE == "image" assert ModalityType.NUMERIC == "numeric" def test_stagenet_is_temporal(): - from pyhealth.processors import StageNetProcessor, TemporalFeatureProcessor, ModalityType + from pyhealth.processors import ( + StageNetProcessor, + TemporalFeatureProcessor, + ModalityType, + ) + p = StageNetProcessor() assert isinstance(p, TemporalFeatureProcessor) assert p.modality() == ModalityType.CODE def test_stagenet_tensor_is_temporal(): - from pyhealth.processors import StageNetTensorProcessor, TemporalFeatureProcessor, ModalityType + from pyhealth.processors import ( + StageNetTensorProcessor, + TemporalFeatureProcessor, + ModalityType, + ) + p = StageNetTensorProcessor() assert isinstance(p, TemporalFeatureProcessor) assert p.modality() == ModalityType.NUMERIC def test_tuple_time_text_is_temporal(): - from pyhealth.processors import TupleTimeTextProcessor, TemporalFeatureProcessor, ModalityType + from pyhealth.processors import ( + TupleTimeTextProcessor, + TemporalFeatureProcessor, + ModalityType, + ) + p = TupleTimeTextProcessor() assert isinstance(p, TemporalFeatureProcessor) assert p.modality() == ModalityType.TEXT def test_time_image_is_temporal(): - from pyhealth.processors import TimeImageProcessor, TemporalFeatureProcessor, ModalityType + from pyhealth.processors import ( + TimeImageProcessor, + TemporalFeatureProcessor, + ModalityType, + ) + p = TimeImageProcessor() assert isinstance(p, TemporalFeatureProcessor) assert p.modality() == ModalityType.IMAGE @@ -52,8 +76,10 @@ def test_time_image_is_temporal(): # ── 2. StageNetProcessor.process_temporal() ─────────────────────────────────── + def test_stagenet_process_temporal(): from pyhealth.processors import StageNetProcessor + samples = [{"codes": (None, ["A", "B", "C"])}] p = StageNetProcessor() p.fit(samples, "codes") @@ -63,13 +89,14 @@ def test_stagenet_process_temporal(): assert set(out.keys()) == {"value", "time"} assert out["value"].dtype == torch.long - assert out["time"].dtype == torch.float32 + assert out["time"].dtype == torch.float32 assert out["value"].shape == (3,) - assert out["time"].shape == (3,) + assert out["time"].shape == (3,) def test_stagenet_tensor_process_temporal(): from pyhealth.processors import StageNetTensorProcessor + samples = [{"vitals": ([0.0, 1.0], [[1.0, 2.0], [3.0, 4.0]])}] p = StageNetTensorProcessor() p.fit(samples, "vitals") @@ -77,35 +104,37 @@ def test_stagenet_tensor_process_temporal(): out = p.process_temporal(([0.0, 1.0], [[1.0, 2.0], [3.0, 4.0]])) assert set(out.keys()) == {"value", "time"} assert out["value"].shape == (2, 2) - assert out["time"].shape == (2,) + assert out["time"].shape == (2,) assert p.value_dim() == 2 assert p.modality().value == "numeric" # ── 3. TemporalTimeseriesProcessor ──────────────────────────────────────────── + def test_temporal_timeseries_basic(): from pyhealth.processors import TemporalTimeseriesProcessor p = TemporalTimeseriesProcessor(sampling_rate=timedelta(hours=2)) ts = [ - datetime(2023, 1, 1, 0), - datetime(2023, 1, 1, 4), - datetime(2023, 1, 1, 8), + datetime(2023, 1, 1, 0), + datetime(2023, 1, 1, 4), + datetime(2023, 1, 1, 8), ] val = np.array([[120.0, 80.0], [115.0, 78.0], [118.0, 82.0]]) out = p.process((ts, val)) # 8 h window / 2 h step + 1 = 5 steps assert out["value"].shape == (5, 2) - assert out["time"].shape == (5,) + assert out["time"].shape == (5,) # Times should be [0, 2, 4, 6, 8] - expected_times = torch.tensor([0., 2., 4., 6., 8.]) + expected_times = torch.tensor([0.0, 2.0, 4.0, 6.0, 8.0]) assert torch.allclose(out["time"], expected_times) def test_temporal_timeseries_fit(): from pyhealth.processors import TemporalTimeseriesProcessor + p = TemporalTimeseriesProcessor() ts = [datetime(2023, 1, 1, 0), datetime(2023, 1, 1, 1)] val = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) @@ -117,29 +146,35 @@ def test_temporal_timeseries_fit(): def test_temporal_timeseries_imputation(): from pyhealth.processors import TemporalTimeseriesProcessor + p = TemporalTimeseriesProcessor(sampling_rate=timedelta(hours=1)) ts = [datetime(2023, 1, 1, 0), datetime(2023, 1, 1, 2)] # gap at h=1 val = np.array([[10.0], [20.0]]) out = p.process((ts, val)) # 3 steps: h=0 → 10, h=1 → forward-filled to 10, h=2 → 20 assert out["value"].shape == (3, 1) - assert float(out["value"][1, 0]) == pytest.approx(10.0) + assert math.isclose(float(out["value"][1, 0]), 10.0, rel_tol=1e-6) # ── 4. collate_temporal ─────────────────────────────────────────────────────── + def test_collate_temporal_basic(): from pyhealth.datasets.collate import collate_temporal batch = [ { - "codes": {"value": torch.tensor([1, 2, 3], dtype=torch.long), - "time": torch.tensor([0., 1., 2.])}, + "codes": { + "value": torch.tensor([1, 2, 3], dtype=torch.long), + "time": torch.tensor([0.0, 1.0, 2.0]), + }, "label": torch.tensor(1), }, { - "codes": {"value": torch.tensor([4, 5, 3], dtype=torch.long), - "time": torch.tensor([0.5, 1.5, 2.5])}, + "codes": { + "value": torch.tensor([4, 5, 3], dtype=torch.long), + "time": torch.tensor([0.5, 1.5, 2.5]), + }, "label": torch.tensor(0), }, ] @@ -147,7 +182,7 @@ def test_collate_temporal_basic(): collated = collate_temporal(batch) assert collated["codes"]["value"].shape == (2, 3) - assert collated["codes"]["time"].shape == (2, 3) + assert collated["codes"]["time"].shape == (2, 3) assert collated["label"].shape == (2,) @@ -156,10 +191,18 @@ def test_collate_temporal_variable_length(): from pyhealth.datasets.collate import collate_temporal batch = [ - {"codes": {"value": torch.tensor([1, 2], dtype=torch.long), - "time": torch.tensor([0., 1.])}}, - {"codes": {"value": torch.tensor([3, 4, 5], dtype=torch.long), - "time": torch.tensor([0., 1., 2.])}}, + { + "codes": { + "value": torch.tensor([1, 2], dtype=torch.long), + "time": torch.tensor([0.0, 1.0]), + } + }, + { + "codes": { + "value": torch.tensor([3, 4, 5], dtype=torch.long), + "time": torch.tensor([0.0, 1.0, 2.0]), + } + }, ] collated = collate_temporal(batch) # Padded to length 3 @@ -168,24 +211,28 @@ def test_collate_temporal_variable_length(): # ── 5. SinusoidalTimeEmbedding ──────────────────────────────────────────────── + def test_sinusoidal_time_embedding_shape(): from pyhealth.models.embedding import SinusoidalTimeEmbedding + emb = SinusoidalTimeEmbedding(dim=64, max_hours=720.0) - t = torch.tensor([[0.0, 12.0, 24.0], [0.0, 6.0, 48.0]]) # (2, 3) + t = torch.tensor([[0.0, 12.0, 24.0], [0.0, 6.0, 48.0]]) # (2, 3) out = emb(t) assert out.shape == (2, 3, 64) def test_sinusoidal_different_times_differ(): from pyhealth.models.embedding import SinusoidalTimeEmbedding + emb = SinusoidalTimeEmbedding(dim=32) - t0 = emb(torch.tensor([0.0])) - t1 = emb(torch.tensor([24.0])) + t0 = emb(torch.tensor([0.0])) + t1 = emb(torch.tensor([24.0])) assert not torch.allclose(t0, t1) # ── 6. UnifiedMultimodalEmbeddingModel, code-only smoke test ───────────────── + def _make_code_processors_and_inputs(batch_size=2, seq_len=5): """Build a minimal dataset mock with a single CODE-modality field.""" from pyhealth.processors import StageNetProcessor @@ -199,7 +246,9 @@ def _make_code_processors_and_inputs(batch_size=2, seq_len=5): # Fake batch dict (as produced by collate_temporal) value = torch.randint(1, vocab_size, (batch_size, seq_len)) - time = torch.arange(seq_len, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1) + time = ( + torch.arange(seq_len, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1) + ) inputs = {"codes": {"value": value, "time": time}} return processors, inputs @@ -214,8 +263,8 @@ def test_unified_model_code_only(): out = model(inputs) assert "sequence" in out - assert "time" in out - assert "mask" in out + assert "time" in out + assert "mask" in out B, S, E = out["sequence"].shape assert B == 2 @@ -228,8 +277,15 @@ def test_unified_model_rejects_non_temporal(): from pyhealth.processors import SequenceProcessor bad_proc = SequenceProcessor() - with pytest.raises(TypeError, match="TemporalFeatureProcessor"): - UnifiedMultimodalEmbeddingModel(processors={"field": bad_proc}, embedding_dim=64) + try: + UnifiedMultimodalEmbeddingModel( + processors={"field": bad_proc}, + embedding_dim=64, + ) + except TypeError as exc: + assert "TemporalFeatureProcessor" in str(exc) + else: + raise AssertionError("Expected TypeError for non-temporal processor") def test_unified_model_gradient_flow(): @@ -239,7 +295,7 @@ def test_unified_model_gradient_flow(): processors, inputs = _make_code_processors_and_inputs() model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=32) - out = model(inputs) + out = model(inputs) loss = out["sequence"].mean() loss.backward() @@ -257,15 +313,13 @@ def test_unified_model_time_sort(): proc = StageNetProcessor() proc.fit(samples, "c") - model = UnifiedMultimodalEmbeddingModel( - processors={"c": proc}, embedding_dim=16 - ) + model = UnifiedMultimodalEmbeddingModel(processors={"c": proc}, embedding_dim=16) # Reverse-order times - value = torch.tensor([[2, 1]]) # (1, 2) - time = torch.tensor([[10.0, 0.0]]) # t=10 then t=0 → should sort to [0, 10] - out = model({"c": {"value": value, "time": time}}) - assert out["time"][0, 0].item() == pytest.approx(0.0) - assert out["time"][0, 1].item() == pytest.approx(10.0) + value = torch.tensor([[2, 1]]) # (1, 2) + time = torch.tensor([[10.0, 0.0]]) # t=10 then t=0 → should sort to [0, 10] + out = model({"c": {"value": value, "time": time}}) + assert math.isclose(out["time"][0, 0].item(), 0.0, rel_tol=1e-6) + assert math.isclose(out["time"][0, 1].item(), 10.0, rel_tol=1e-6) # ── 7. field_embeddings: reuse pre-built unimodal encoder ───────────────────── @@ -322,8 +376,8 @@ class _MockEmbedModel: assert isinstance(model.encoders["codes"], nn.Sequential) # Forward should produce embedding_dim=32 value = torch.randint(1, vocab_size, (2, 3)) - time = torch.arange(3, dtype=torch.float32).unsqueeze(0).expand(2, -1) - out = model({"codes": {"value": value, "time": time}}) + time = torch.arange(3, dtype=torch.float32).unsqueeze(0).expand(2, -1) + out = model({"codes": {"value": value, "time": time}}) assert out["sequence"].shape[-1] == 32 @@ -349,13 +403,90 @@ class _MockEmbedModel: field_embeddings={"codes": _MockEmbedModel()}, ) value = torch.randint(1, vocab_size, (3, 4)) - time = torch.arange(4, dtype=torch.float32).unsqueeze(0).expand(3, -1) - out = model({"codes": {"value": value, "time": time}}) + time = torch.arange(4, dtype=torch.float32).unsqueeze(0).expand(3, -1) + out = model({"codes": {"value": value, "time": time}}) assert out["sequence"].shape == (3, 4, 64) assert out["mask"].shape == (3, 4) +def test_unified_text_encoder_shared_by_tokenizer(): + """Token-based TEXT fields with the same tokenizer share one encoder.""" + from types import SimpleNamespace + import torch.nn as nn + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.processors import ModalityType, TemporalFeatureProcessor + + class _DummyTemporalTextProcessor(TemporalFeatureProcessor): + def __init__(self, tokenizer_model: str): + self.tokenizer_model = tokenizer_model + + def process(self, value): + return value + + def modality(self): + return ModalityType.TEXT + + def value_dim(self): + return 0 + + def is_token(self): + return True + + def schema(self): + return ("value", "mask", "time") + + def dim(self): + return (2, 2, 1) + + def spatial(self): + return (False, False) + + class _DummyBert(nn.Module): + def __init__(self, hidden_size: int): + super().__init__() + self.config = SimpleNamespace(hidden_size=hidden_size) + + def forward(self, input_ids=None, attention_mask=None): + if input_ids is None: + raise ValueError("input_ids is required") + b, l = input_ids.shape + hidden = self.config.hidden_size + out = torch.zeros(b, l, hidden) + return SimpleNamespace(last_hidden_state=out) + + call_count = {"n": 0} + + def _fake_from_pretrained(_name): + call_count["n"] += 1 + return _DummyBert(hidden_size=48) + + with patch("transformers.AutoModel.from_pretrained", _fake_from_pretrained): + processors = { + "discharge_note_times": _DummyTemporalTextProcessor("bert-base-uncased"), + "radiology_note_times": _DummyTemporalTextProcessor("bert-base-uncased"), + } + + model = UnifiedMultimodalEmbeddingModel( + processors=processors, + embedding_dim=32, + ) + + # Both text fields reuse the same encoder instance. + assert ( + model.encoders["discharge_note_times"] is model.encoders["radiology_note_times"] + ) + assert call_count["n"] == 1 + + # Projections remain field-specific. + assert "discharge_note_times" in model.projections + assert "radiology_note_times" in model.projections + assert ( + model.projections["discharge_note_times"] + is not model.projections["radiology_note_times"] + ) + + # ── 8. Downstream models in unified mode ────────────────────────────────────── @@ -408,8 +539,8 @@ def test_transformer_unified_mode(): model = Transformer(dataset=dataset, embedding_dim=32, unified_embedding=unified) loader = get_dataloader(dataset, batch_size=2, shuffle=False) - batch = next(iter(loader)) - out = model(**batch) + batch = next(iter(loader)) + out = model(**batch) assert "loss" in out and "y_prob" in out and "logit" in out out["loss"].backward() @@ -436,8 +567,8 @@ def test_ehrmamba_unified_mode(): ) loader = get_dataloader(dataset, batch_size=2, shuffle=False) - batch = next(iter(loader)) - out = model(**batch) + batch = next(iter(loader)) + out = model(**batch) assert "loss" in out and "y_prob" in out out["loss"].backward() @@ -465,8 +596,8 @@ def test_jamba_ehr_unified_mode(): ) loader = get_dataloader(dataset, batch_size=2, shuffle=False) - batch = next(iter(loader)) - out = model(**batch) + batch = next(iter(loader)) + out = model(**batch) assert "loss" in out and "y_prob" in out out["loss"].backward() @@ -484,7 +615,9 @@ def test_mlp_unified_mode(): processors=dataset.input_processors, embedding_dim=32, ) - model = MLP(dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified) + model = MLP( + dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified + ) loader = get_dataloader(dataset, batch_size=2, shuffle=False) batch = next(iter(loader)) @@ -507,7 +640,9 @@ def test_rnn_unified_mode(): processors=dataset.input_processors, embedding_dim=32, ) - model = RNN(dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified) + model = RNN( + dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified + ) loader = get_dataloader(dataset, batch_size=2, shuffle=False) batch = next(iter(loader)) @@ -560,6 +695,7 @@ def test_unified_per_field_backward_compat(): dataset = _make_stagenet_dataset() # Uses SequenceProcessor-style input_schema for per-field mode from pyhealth.datasets import create_sample_dataset + samples = [ {"patient_id": "p0", "visit_id": "v0", "codes": ["A", "B", "C"], "label": 1}, {"patient_id": "p1", "visit_id": "v1", "codes": ["D", "E"], "label": 0}, @@ -572,7 +708,21 @@ def test_unified_per_field_backward_compat(): ) model = Transformer(dataset=ds, embedding_dim=32) loader = get_dataloader(ds, batch_size=2, shuffle=False) - batch = next(iter(loader)) - out = model(**batch) + batch = next(iter(loader)) + out = model(**batch) assert "loss" in out out["loss"].backward() + + +def load_tests(loader, tests, pattern): + """Expose top-level test_ functions to unittest discovery.""" + suite = unittest.TestSuite() + namespace = globals() + for name in sorted(namespace): + if name.startswith("test_") and callable(namespace[name]): + suite.addTest(unittest.FunctionTestCase(namespace[name])) + return suite + + +if __name__ == "__main__": + unittest.main()