Ra/labs only variant

examples/mortality_prediction/unified_embedding_e2e_mimic4.py

@@ -16,7 +16,16 @@
 
 --task clinical_notes_icd_labs
     ClinicalNotesICDLabsMIMIC4: discharge/radiology notes + ICD + labs.
-    Requires --note-root.  Used for Table 2 (EHR + clinical text).
+    Requires --note-root.  Legacy; ICD codes are discharge-coded (leakage).
+
+--task notes_labs (recommended for multimodal)
+    NotesLabsMIMIC4: admission-context note sections + labs, no ICD codes.
+    Extracts Chief Complaint, HPI, PMH, Medications on Admission from the
+    discharge note — text available at admission time, ~90%+ coverage.
+    Requires --note-root.
+    Add --freeze-encoder to freeze Bio_ClinicalBERT and train only the
+    backbone; cuts BERT VRAM by ~50%, useful on smaller GPUs (≤24 GB).
+    Add --icd-codes to include discharge-coded ICD codes (ablation only).
 
 Example
 -------
@@ -50,19 +59,26 @@
 from typing import Any, Tuple
 
 import numpy as np
+import torch
 
 from pyhealth.datasets import (
     MIMIC4Dataset,
     get_dataloader,
     split_by_patient,
     split_by_sample,
+    sample_balanced,
 )
 from pyhealth.models import MLP, RNN, Transformer, UnifiedMultimodalEmbeddingModel
 from pyhealth.models.bottleneck_transformer import BottleneckTransformer
 from pyhealth.models.ehrmamba import EHRMamba
 from pyhealth.models.jamba_ehr import JambaEHR
 from pyhealth.tasks import MortalityPredictionStageNetMIMIC4
-from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4
+from pyhealth.tasks.multimodal_mimic4 import (
+    ClinicalNotesICDLabsMIMIC4,
+    ICDLabsMIMIC4,
+    LabsMIMIC4,
+    NotesLabsMIMIC4,
+)
 from pyhealth.trainer import Trainer
 from pyhealth.utils import set_seed
 
@@ -79,13 +95,30 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset:
     if args.task == "icd_labs":
         ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"]
 
+    if args.task == "notes_labs":
+        if not args.note_root:
+            raise ValueError("--task notes_labs requires --note-root.")
+        note_tables = ["discharge"]
+        # Load ICD tables only when explicitly requested (they are discharge-coded).
+        ehr_tables = (
+            ["diagnoses_icd", "procedures_icd", "labevents"]
+            if args.icd_codes
+            else ["labevents"]
+        )
+        if args.include_vitals:
+            if "chartevents" not in ehr_tables:
+                ehr_tables.append("chartevents")
+
+    if args.task == "labs":
+        ehr_tables = ["labevents"]
+
     return MIMIC4Dataset(
         ehr_root=args.ehr_root,
         ehr_tables=ehr_tables,
         note_root=args.note_root if note_tables else None,
         note_tables=note_tables,
         cache_dir=args.cache_dir,
-        dev=args.dev,
+        dev=args.dev if args.dev else False,
         num_workers=args.num_workers,
     )
 
@@ -97,6 +130,14 @@ def _build_task(args: argparse.Namespace):
         return ICDLabsMIMIC4(window_hours=args.observation_window_hours)
     if args.task == "clinical_notes_icd_labs":
         return ClinicalNotesICDLabsMIMIC4(window_hours=args.observation_window_hours)
+    if args.task == "notes_labs":
+        return NotesLabsMIMIC4(
+            window_hours=args.observation_window_hours,
+            include_icd=args.icd_codes,
+            include_vitals=args.include_vitals,
+        )
+    if args.task == "labs":
+        return LabsMIMIC4(window_hours=args.observation_window_hours)
     raise ValueError(f"Unknown task: {args.task}")
 
 
@@ -111,6 +152,7 @@ def _build_model(args: argparse.Namespace, sample_dataset: Any):
     unified = UnifiedMultimodalEmbeddingModel(
         processors=sample_dataset.input_processors,
         embedding_dim=args.embedding_dim,
+        freeze_text_encoder=args.freeze_encoder,
     )
 
     if args.model == "mlp":
@@ -204,6 +246,26 @@ def _write_predictions(
             )
 
 
+def _compute_pos_weight(train_ds, label_key: str = "mortality") -> float:
+    """Count pos/neg in train_ds and return n_neg/n_pos for BCE pos_weight."""
+    n_pos = n_neg = 0
+    for i in range(len(train_ds)):
+        sample = train_ds[i]
+        label = sample.get(label_key, 0)
+        if hasattr(label, "__iter__"):
+            label = next(iter(label))
+        if float(label) > 0.5:
+            n_pos += 1
+        else:
+            n_neg += 1
+    if n_pos == 0:
+        return 1.0
+    # Cap at 10: n_neg/n_pos ≈ 37 on MIMIC-IV mortality is too extreme with
+    # typical LRs and causes training oscillation. 10 still strongly corrects
+    # for imbalance while keeping gradient magnitudes tractable.
+    return min(10.0, n_neg / n_pos)
+
+
 def run(args: argparse.Namespace) -> Path:
     set_seed(args.seed)
 
@@ -217,8 +279,29 @@ def run(args: argparse.Namespace) -> Path:
         )
 
     train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+
+    label_key = list(sample_dataset.output_schema.keys())[0]
+
+    # Balanced sampling: undersample negatives to achieve a target pos:neg ratio.
+    if args.balanced_sampling:
+        ratio = args.balanced_ratio
+        print(f"[balanced_sampling] Undersampling training set to pos:neg ratio 1:{ratio}")
+        train_ds = sample_balanced(train_ds, ratio=ratio, seed=args.seed, label_key=label_key)
+        print(f"[balanced_sampling] Training set size after sampling: {len(train_ds)}")
+
     model = _build_model(args, sample_dataset)
 
+    # Apply class-imbalance correction via BCE pos_weight.
+    # pos_weight = n_neg / n_pos so the rare positive class gets proportionally
+    # higher gradient signal, preventing all-negative collapse (F1=0).
+    if args.pos_weight is not None:
+        pw_value = args.pos_weight
+    else:
+        print(f"[pos_weight] Computing class balance from {len(train_ds)} training samples...")
+        pw_value = _compute_pos_weight(train_ds, label_key=label_key)
+    print(f"[pos_weight] Using pos_weight={pw_value:.2f} for binary BCE loss.")
+    model._pos_weight = torch.tensor([pw_value], dtype=torch.float32)
+
     train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
     val_loader = (
         get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
@@ -257,8 +340,15 @@ def run(args: argparse.Namespace) -> Path:
             effective_max_grad_norm = 0.5
         optimizer_params["eps"] = args.adam_eps if args.adam_eps is not None else 1e-6
     else:
+        # All non-BT models: 1e-4 (was 1e-3). With pos_weight correction,
+        # effective gradient magnitude for positives is ~10x higher, so a
+        # smaller LR is needed to avoid training oscillation.
         if effective_lr is None:
-            effective_lr = 1e-3
+            effective_lr = 1e-4
+        # Universal grad clipping: prevents runaway updates from the weighted
+        # positive-class loss (pos_weight ≈ 10 scales positive gradients 10x).
+        if effective_max_grad_norm is None:
+            effective_max_grad_norm = 1.0
         if args.adam_eps is not None:
             optimizer_params["eps"] = args.adam_eps
 
@@ -299,8 +389,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--task",
         type=str,
-        choices=["icd_labs", "clinical_notes_icd_labs"],
+        choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "labs"],
         default="stagenet",
+        help=(
+            "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. "
+            "No ICD codes (discharge-coded = leakage). Recommended for multimodal."
+        ),
     )
     parser.add_argument(
         "--model",
@@ -321,9 +415,9 @@ def parse_args() -> argparse.Namespace:
         type=float,
         default=None,
         help=(
-            "Learning rate. Default is model-specific: 1e-3 for "
-            "mlp/rnn/transformer/ehrmamba/jambaehr, 1e-4 for "
-            "bottleneck_transformer."
+            "Learning rate. Default is 1e-4 for all models. "
+            "(Previously 1e-3 for mlp/rnn/transformer/ehrmamba/jambaehr — "
+            "reduced after pos_weight correction caused oscillation at 1e-3.)"
         ),
     )
     parser.add_argument(
@@ -340,10 +434,82 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--num-workers", type=int, default=1)
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--patience", type=int, default=None)
-    parser.add_argument("--dev", action="store_true")
+    parser.add_argument(
+        "--dev",
+        nargs="?",
+        type=int,
+        const=1000,
+        default=0,
+        help=(
+            "Dev mode: limit dataset to N patients for fast iteration. "
+            "--dev (no value) defaults to 1000 patients. "
+            "--dev 5000 limits to 5000. Omit for full dataset."
+        ),
+    )
+    parser.add_argument(
+        "--pos-weight",
+        type=float,
+        default=None,
+        help=(
+            "BCE pos_weight for the positive class (float). "
+            "Default: auto-computed as n_neg/n_pos from training split. "
+            "Set to 1.0 to disable class-imbalance correction."
+        ),
+    )
 
     # Task-specific
     parser.add_argument("--observation-window-hours", type=int, default=24)
+    parser.add_argument(
+        "--icd-codes",
+        action="store_true",
+        default=False,
+        help=(
+            "Include discharge-coded ICD codes in notes_labs task. "
+            "Default: off (ICD codes are coded at discharge and constitute "
+            "data leakage for in-hospital mortality prediction). "
+            "Enable only for ablation / legacy comparison experiments."
+        ),
+    )
+    parser.add_argument(
+        "--freeze-encoder",
+        action="store_true",
+        default=False,
+        help=(
+            "Freeze pretrained BERT text encoder weights and train only the "
+            "downstream backbone (MLP/RNN/Transformer head + projection layer). "
+            "Reduces VRAM by ~50% for the text branch; useful when GPU memory "
+            "is limited or for faster iteration on backbone architectures."
+        ),
+    )
+    parser.add_argument(
+        "--include-vitals",
+        action="store_true",
+        default=False,
+        help=(
+            "Include ICU vital signs (HeartRate, SysBP, DiasBP, MeanBP, "
+            "RespRate, SpO2, Temperature) from chartevents as an additional "
+            "modality alongside labs and notes. Adds chartevents to EHR tables."
+        ),
+    )
+    parser.add_argument(
+        "--balanced-sampling",
+        action="store_true",
+        default=False,
+        help=(
+            "Undersample the majority (negative) class in training to improve "
+            "PR-AUC on imbalanced datasets. Uses sample_balanced() to create a "
+            "1:--balanced-ratio pos:neg training set."
+        ),
+    )
+    parser.add_argument(
+        "--balanced-ratio",
+        type=float,
+        default=1.0,
+        help=(
+            "Negatives per positive in the balanced training set. "
+            "Default: 1.0 (equal pos/neg). Only used with --balanced-sampling."
+        ),
+    )
 
     # RNN-specific
     parser.add_argument("--rnn-type", type=str, default="GRU")