diff --git a/examples/mortality_prediction/medfuse_mimic4.py b/examples/mortality_prediction/medfuse_mimic4.py new file mode 100644 index 000000000..280b6692b --- /dev/null +++ b/examples/mortality_prediction/medfuse_mimic4.py @@ -0,0 +1,290 @@ +"""Run MedFuse (EHR time-series + chest X-ray) in-hospital mortality on real MIMIC-IV. + +Wires the merged MedFuse model to real MIMIC-IV + MIMIC-CXR via MedFuseMortalityMIMIC4 +(no model edits). EHR selectable between ICU vitals (chartevents) and chemistry labs +(labevents); modality between EHR+CXR fusion and EHR-only (same paired cohort, CXR hidden). +Manual train loop; per-epoch curves to wandb project pyhealth-multimodal (team convention). + +Self-contained: the two dataset configs (chartevents EHR + corrected CXR metadata) are +embedded as Python dicts and written to throwaway temp YAMLs at runtime, so this script + +the MedFuseMortalityMIMIC4 task are the only files needed. +""" + +from __future__ import annotations + +import argparse +import os +import platform +import random +import subprocess +import tempfile + +import numpy as np +import pandas as pd +import torch +import yaml + +from pyhealth.datasets import mimic4 as _mimic4_module + +_mimic4_module.MIMIC4CXRDataset.prepare_metadata = lambda self, root: None # no-op + +import wandb + +from pyhealth.datasets import MIMIC4Dataset, get_dataloader, split_by_patient +from pyhealth.metrics import binary_metrics_fn +from pyhealth.models import MedFuse +from pyhealth.tasks import MedFuseMortalityMIMIC4 + +HERE = os.path.dirname(os.path.abspath(__file__)) +CORRECTED_METADATA_NAME = "mimic-cxr-2.0.0-metadata-medfuse.csv" +SOURCE_METADATA_NAME = "mimic-cxr-2.0.0-metadata-pyhealth.csv" +WANDB_PROJECT = "pyhealth-multimodal" +METRICS = ["pr_auc", "roc_auc", "f1", "accuracy"] + +# --- Dataset configs as Python dicts (no committed YAML). Written to temp files at runtime +# --- because MIMIC4Dataset accepts a config *path*. The default mimic4_ehr.yaml lacks +# --- chartevents (needed for vitals); labs uses the default. The CXR config points at a +# --- corrected metadata CSV whose image_path resolves to the real on-disk images. +EHR_CHARTEVENTS_CFG = { + "version": "2.2", + "tables": { + "patients": {"file_path": "hosp/patients.csv.gz", "patient_id": "subject_id", + "timestamp": None, "attributes": ["gender", "anchor_age", "dod"]}, + "admissions": {"file_path": "hosp/admissions.csv.gz", "patient_id": "subject_id", + "timestamp": "admittime", + "attributes": ["hadm_id", "dischtime", "hospital_expire_flag"]}, + "icustays": {"file_path": "icu/icustays.csv.gz", "patient_id": "subject_id", + "timestamp": "intime", "attributes": ["hadm_id", "stay_id", "outtime"]}, + "chartevents": {"file_path": "icu/chartevents.csv.gz", "patient_id": "subject_id", + "timestamp": "charttime", + "attributes": ["hadm_id", "stay_id", "itemid", "valuenum", "valueuom"]}, + }, +} +CXR_METADATA_CFG = { + "version": "2.1.0", + "tables": { + "metadata": {"file_path": CORRECTED_METADATA_NAME, "patient_id": "subject_id", + "timestamp": ["studydate", "studytime"], "timestamp_format": "%Y%m%d%H%M%S", + "attributes": ["image_path", "dicom_id", "study_id"]}, + }, +} + + +def write_temp_yaml(config: dict) -> str: + fd, path = tempfile.mkstemp(suffix=".yaml", prefix="medfuse_cfg_") + with os.fdopen(fd, "w") as f: + yaml.safe_dump(config, f, sort_keys=False) + return path + + +def set_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def git_sha(): + try: + return subprocess.check_output(["git", "rev-parse", "--short", "HEAD"], cwd=HERE).decode().strip() + except Exception: + return "unknown" + + +def ensure_corrected_cxr_metadata(cxr_root, image_subdir="resized_images"): + dst = os.path.join(cxr_root, CORRECTED_METADATA_NAME) + if os.path.exists(dst): + return + src = os.path.join(cxr_root, SOURCE_METADATA_NAME) + if not os.path.exists(src): + raise FileNotFoundError(f"Expected {src} (pyhealth CXR metadata).") + print(f"[cxr] building corrected metadata -> {dst}") + df = pd.read_csv(src, dtype=str) + img_dir = os.path.join(cxr_root, image_subdir) + df["image_path"] = df["dicom_id"].map(lambda d: os.path.join(img_dir, f"{d}.jpg")) + df.to_csv(dst, index=False) + print(f"[cxr] wrote {len(df):,} rows") + + +def build_dataset(args): + if args.ehr_source == "vitals": + ehr_tables = ["chartevents"] + ehr_config_path = write_temp_yaml(EHR_CHARTEVENTS_CFG) + else: # labs — default mimic4_ehr.yaml already defines labevents + ehr_tables = ["labevents"] + ehr_config_path = None + return MIMIC4Dataset( + ehr_root=args.ehr_root, cxr_root=args.cxr_root, + ehr_tables=ehr_tables, cxr_tables=["metadata"], + ehr_config_path=ehr_config_path, cxr_config_path=write_temp_yaml(CXR_METADATA_CFG), + dev=args.dev, num_workers=args.num_workers, + ) + + +def make_inputs(batch, ehr_only, device): + inputs = {"ehr": batch["ehr"].to(device), "mortality": batch["mortality"].to(device)} + if not ehr_only: + inputs["cxr"] = batch["cxr"].to(device) + inputs["cxr_mask"] = batch["cxr_mask"].to(device) + return inputs + + +@torch.no_grad() +def evaluate(model, loader, ehr_only, device): + model.eval() + pids, y_true, y_prob, losses = [], [], [], [] + for batch in loader: + out = model(**make_inputs(batch, ehr_only, device)) + y_prob.append(out["y_prob"].detach().cpu().view(-1)) + y_true.append(out["y_true"].detach().cpu().view(-1)) + losses.append(float(out["loss"])) + pids.extend(batch["patient_id"]) + return (pids, torch.cat(y_true).numpy(), torch.cat(y_prob).numpy(), + float(np.mean(losses)) if losses else float("nan")) + + +def safe_metrics(y_true, y_prob): + try: + return binary_metrics_fn(y_true, y_prob, metrics=METRICS) + except Exception as exc: + print(f"[warn] metric computation failed: {exc}") + return {m: float("nan") for m in METRICS} + + +def main(): + p = argparse.ArgumentParser(description="MedFuse on MIMIC-IV (EHR + CXR).") + p.add_argument("--ehr-root", default="/srv/local/data/physionet.org/files/mimiciv/2.2") + p.add_argument("--cxr-root", default="/srv/local/data/MIMIC-CXR") + p.add_argument("--ehr-source", choices=["vitals", "labs"], default="vitals") + p.add_argument("--modality", choices=["ehr_cxr", "ehr_only"], default="ehr_cxr") + p.add_argument("--window-hours", type=float, default=48.0) + p.add_argument("--fixed-window", action="store_true", default=True) + p.add_argument("--variable-window", dest="fixed_window", action="store_false") + p.add_argument("--bin-hours", type=float, default=2.0) + p.add_argument("--min-los-hours", type=float, default=48.0) + p.add_argument("--cxr-backbone", choices=["resnet18", "resnet34", "resnet50"], default="resnet34") + p.add_argument("--no-cxr-pretrained", dest="cxr_pretrained", action="store_false") + p.add_argument("--epochs", type=int, default=20) + p.add_argument("--batch-size", type=int, default=64) + p.add_argument("--lr", type=float, default=1e-4) + p.add_argument("--dropout", type=float, default=0.5) + p.add_argument("--device", default="cuda") + p.add_argument("--num-workers", type=int, default=8) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--output-dir", default=os.path.join(HERE, "output")) + p.add_argument("--exp-name", default=None) + p.add_argument("--dev", action="store_true") + args = p.parse_args() + + set_seed(args.seed) + ehr_only = args.modality == "ehr_only" + device = args.device if torch.cuda.is_available() else "cpu" + exp_name = args.exp_name or f"medfuse_{args.ehr_source}_{args.modality}_seed{args.seed}" + out_dir = os.path.join(args.output_dir, exp_name) + os.makedirs(out_dir, exist_ok=True) + + ensure_corrected_cxr_metadata(args.cxr_root) + print(f"[data] building dataset (ehr_source={args.ehr_source}, dev={args.dev})") + dataset = build_dataset(args) + + task = MedFuseMortalityMIMIC4( + ehr_source=args.ehr_source, window_hours=args.window_hours, + fixed_window=args.fixed_window, bin_hours=args.bin_hours, + min_los_hours=args.min_los_hours, require_cxr=True, + ) + print(f"[task] set_task {task.task_name} ({args.ehr_source})") + samples = dataset.set_task(task, num_workers=args.num_workers) + print(f"[task] {len(samples)} samples (F={task.n_features}, n_bins={task.n_bins})") + + train_ds, val_ds, test_ds = split_by_patient(samples, [0.7, 0.1, 0.2], seed=args.seed) + 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) + test_loader = get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + + model = MedFuse( + dataset=samples, ehr_feature_key="ehr", + cxr_feature_key="__disabled__" if ehr_only else "cxr", + cxr_mask_key=None if ehr_only else "cxr_mask", + cxr_backbone=args.cxr_backbone, + cxr_pretrained=False if ehr_only else args.cxr_pretrained, + dropout=args.dropout, + ).to(device) + param_count = sum(x.numel() for x in model.parameters()) + + wandb.init( + project=WANDB_PROJECT, name=exp_name, + tags=["medfuse", args.ehr_source, args.modality] + (["dev"] if args.dev else []), + config={ + "task": task.task_name, "model": "medfuse", "exp_name": exp_name, + "modalities": ["ehr"] if ehr_only else ["ehr", "cxr"], + "seed": args.seed, "dev_mode": args.dev, + "hp/batch_size": args.batch_size, "hp/lr": args.lr, "hp/epochs": args.epochs, + "hp/dropout": args.dropout, + "arch/cxr_backbone": args.cxr_backbone, "arch/ehr_hidden_dim": 256, + "arch/ehr_num_layers": 2, "arch/fusion_hidden_dim": 512, + "arch/projection_dim": 256, "arch/param_count": param_count, + "fusion/ehr_source": args.ehr_source, + "fusion/observation_window_hours": args.window_hours, + "fusion/fixed_window": args.fixed_window, "fusion/bin_hours": args.bin_hours, + "data/train_samples": len(train_ds), "data/val_samples": len(val_ds), + "data/test_samples": len(test_ds), "data/cxr_variant": "resized", + "paths/ehr_root": args.ehr_root, "paths/cxr_root": args.cxr_root, + "paths/output_dir": out_dir, + "env/gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu", + "env/cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES", ""), + "env/torch_version": torch.__version__, "env/git_sha": git_sha(), + "env/hostname": platform.node(), + }, + ) + wandb.define_metric("epoch") + for key in ["train_loss", "val_loss"] + [f"val_{m}" for m in METRICS]: + wandb.define_metric(key, step_metric="epoch") + + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) + best_pr, best_path = -1.0, os.path.join(out_dir, "best.ckpt") + + for epoch in range(args.epochs): + model.train() + total, n = 0.0, 0 + for batch in train_loader: + optimizer.zero_grad() + out = model(**make_inputs(batch, ehr_only, device)) + out["loss"].backward() + optimizer.step() + bs = len(batch["patient_id"]) + total += float(out["loss"]) * bs + n += bs + train_loss = total / max(n, 1) + + _, y_true, y_prob, val_loss = evaluate(model, val_loader, ehr_only, device) + scores = safe_metrics(y_true, y_prob) + wandb.log({"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, + **{f"val_{m}": scores[m] for m in METRICS}}) + print(f"[epoch {epoch}] train_loss={train_loss:.4f} val_loss={val_loss:.4f} " + f"val_pr_auc={scores['pr_auc']:.4f} val_roc_auc={scores['roc_auc']:.4f}") + + if scores["pr_auc"] > best_pr: + best_pr = scores["pr_auc"] + torch.save(model.state_dict(), best_path) + print(f"[epoch {epoch}] new best val_pr_auc={best_pr:.4f}") + + if os.path.exists(best_path): + model.load_state_dict(torch.load(best_path, map_location=device)) + pids, y_true, y_prob, test_loss = evaluate(model, test_loader, ehr_only, device) + test_scores = safe_metrics(y_true, y_prob) + print(f"[test] loss={test_loss:.4f} " + " ".join(f"{m}={test_scores[m]:.4f}" for m in METRICS)) + + for m in METRICS: + wandb.summary[f"test_{m}"] = test_scores[m] + wandb.summary["test_loss"] = test_loss + wandb.summary["best_val_pr_auc"] = best_pr + + pred_csv = os.path.join(out_dir, "predictions_medfuse.csv") + pd.DataFrame({"patient_id": pids, "y_true": y_true, "y_prob": y_prob, + "y_pred_threshold_0_5": (y_prob >= 0.5).astype(int)}).to_csv(pred_csv, index=False) + print(f"[test] wrote predictions -> {pred_csv}") + wandb.finish() + + +if __name__ == "__main__": + main() diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index 3cb4a59af..6abf67e1b 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -30,6 +30,7 @@ from .graph_torchvision_model import Graph_TorchvisionModel from .graphcare import GraphCare from .grasp import GRASP, GRASPLayer +from .medfuse import MedFuse, MedFuseLayer from .medlink import MedLink from .micron import MICRON, MICRONLayer from .mlp import MLP diff --git a/pyhealth/models/medfuse.py b/pyhealth/models/medfuse.py new file mode 100644 index 000000000..8e9f8e112 --- /dev/null +++ b/pyhealth/models/medfuse.py @@ -0,0 +1,456 @@ +# Authors: Sean Nian, Tony Hong, Yaqi Qiao +# Description: MedFuse model implementation for PyHealth. + +from __future__ import annotations + +from typing import Dict, Optional, Tuple, cast + +import torch +import torch.nn as nn +import torch.nn.utils.rnn as rnn_utils + +from pyhealth.datasets import SampleDataset +from pyhealth.models.base_model import BaseModel + + +class MedFuseLayer(nn.Module): + """MedFuse fusion layer. + + Fuses EHR time-series and chest X-ray (CXR) representations using an + LSTM-based sequential fusion strategy that supports missing CXR modality. + Follows Paper §3.2, Eq. 4: the fusion LSTM consumes the token sequence + ``v_fusion = [v_ehr, v*_cxr]`` (length 1 when CXR is missing). The EHR + representation is the first input token, not an initial hidden state. + + Paper: + Hayat, N., Geras, K. J., & Shamout, F. E. (2022). + MedFuse: Multi-modal fusion with clinical time-series data and chest + X-ray images. MLHC 2022. + + Args: + ehr_input_dim: Dimension of EHR features at each timestep. + ehr_hidden_dim: Hidden dimension of the EHR LSTM encoder. + Default is 256. + ehr_num_layers: Number of stacked LSTM layers for EHR. + Default is 2. + cxr_backbone: ResNet variant for CXR encoder. + Default is ``"resnet34"``. + cxr_pretrained: Whether to use pretrained CXR encoder weights. + Default is True. + fusion_hidden_dim: Hidden dimension of fusion LSTM. + Default is 512. + projection_dim: Dimension to project CXR features to. + Must match ``ehr_hidden_dim``. + dropout: Dropout rate. + num_labels: Number of output labels. + """ + + SUPPORTED_CXR_BACKBONES = ("resnet18", "resnet34", "resnet50") + + def __init__( + self, + ehr_input_dim: int, + ehr_hidden_dim: int = 256, + ehr_num_layers: int = 2, + cxr_backbone: str = "resnet34", + cxr_pretrained: bool = True, + fusion_hidden_dim: int = 512, + projection_dim: int = 256, + dropout: float = 0.5, + num_labels: int = 1, + ) -> None: + super().__init__() + if projection_dim != ehr_hidden_dim: + raise ValueError( + "projection_dim must match ehr_hidden_dim for fusion sequencing." + ) + if cxr_backbone not in self.SUPPORTED_CXR_BACKBONES: + raise ValueError( + f"Unsupported cxr_backbone: {cxr_backbone}. " + f"Supported values: {self.SUPPORTED_CXR_BACKBONES}." + ) + + self.ehr_encoder = nn.LSTM( + input_size=ehr_input_dim, + hidden_size=ehr_hidden_dim, + num_layers=ehr_num_layers, + batch_first=True, + dropout=dropout if ehr_num_layers > 1 else 0.0, + ) + + self.cxr_encoder, cxr_feature_dim = self._build_cxr_encoder( + cxr_backbone=cxr_backbone, + cxr_pretrained=cxr_pretrained, + ) + self.projection = nn.Linear(cxr_feature_dim, projection_dim) + + self.fusion_lstm = nn.LSTM( + input_size=projection_dim, + hidden_size=fusion_hidden_dim, + num_layers=1, + batch_first=True, + ) + self.dropout_layer = nn.Dropout(dropout) + self.classifier = nn.Linear(fusion_hidden_dim, num_labels) + + def _build_cxr_encoder( + self, + cxr_backbone: str, + cxr_pretrained: bool, + ) -> Tuple[nn.Module, int]: + """Builds the CXR backbone and returns encoder + feature dimension.""" + try: + import torchvision.models as tv_models + except ImportError as exc: + raise ImportError( + "torchvision is required to use MedFuse CXR backbones." + ) from exc + + constructor = getattr(tv_models, cxr_backbone) + + if cxr_pretrained: + try: + weights = tv_models.get_model_weights(cxr_backbone).DEFAULT + backbone = constructor(weights=weights) + except Exception: + backbone = constructor(pretrained=True) + else: + try: + backbone = constructor(weights=None) + except TypeError: + backbone = constructor(pretrained=False) + + if not hasattr(backbone, "fc") or not isinstance(backbone.fc, nn.Linear): + raise ValueError( + f"Backbone {cxr_backbone} must expose a linear `fc` layer." + ) + + feature_dim = backbone.fc.in_features + backbone.fc = nn.Identity() + return backbone, feature_dim + + def _encode_cxr(self, cxr_input: torch.Tensor) -> torch.Tensor: + """Encodes CXR images and projects them to the fusion dimension.""" + if cxr_input.dim() != 4: + raise ValueError( + "cxr_input must have shape [batch, channels, height, width]." + ) + + if cxr_input.size(1) == 1: + cxr_input = cxr_input.repeat(1, 3, 1, 1) + elif cxr_input.size(1) != 3: + raise ValueError("cxr_input must have 1 or 3 channels.") + + cxr_features = self.cxr_encoder(cxr_input) + if cxr_features.dim() > 2: + cxr_features = torch.flatten(cxr_features, start_dim=1) + + projected = self.projection(cxr_features) + return projected + + def forward( + self, + ehr_input: torch.Tensor, + cxr_input: Optional[torch.Tensor] = None, + cxr_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Runs MedFuse fusion. + + Args: + ehr_input: EHR tensor of shape ``[batch, seq_len, ehr_input_dim]``. + cxr_input: Optional CXR tensor of shape + ``[batch, channels, height, width]``. + cxr_mask: Optional tensor of shape ``[batch]`` indicating CXR + availability (1 = present, 0 = missing). + + Returns: + Logits tensor of shape ``[batch, num_labels]``. + """ + if ehr_input.dim() != 3: + raise ValueError( + "ehr_input must have shape [batch, seq_len, ehr_input_dim]." + ) + + _, (ehr_hidden, _) = self.ehr_encoder(ehr_input) + ehr_repr = self.dropout_layer(ehr_hidden[-1]) + + batch_size = ehr_repr.size(0) + sequence_lengths = torch.ones( + batch_size, + dtype=torch.long, + device=ehr_input.device, + ) + fusion_tokens = [ehr_repr.unsqueeze(1)] + + if cxr_input is not None: + if cxr_input.size(0) != batch_size: + raise ValueError( + "cxr_input batch size must match ehr_input batch size." + ) + + if cxr_mask is None: + cxr_present_mask = torch.ones( + batch_size, + dtype=torch.bool, + device=ehr_input.device, + ) + else: + cxr_present_mask = cxr_mask.to(ehr_input.device).view(-1).bool() + if cxr_present_mask.numel() != batch_size: + raise ValueError("cxr_mask must have shape [batch].") + + cxr_projected = torch.zeros( + batch_size, + self.projection.out_features, + dtype=ehr_repr.dtype, + device=ehr_input.device, + ) + + if cxr_present_mask.any(): + present_cxr = cxr_input[cxr_present_mask] + present_projected = self._encode_cxr(present_cxr) + cxr_projected[cxr_present_mask] = present_projected + sequence_lengths[cxr_present_mask] = 2 + + fusion_tokens.append(cxr_projected.unsqueeze(1)) + + fusion_sequence = torch.cat(fusion_tokens, dim=1) + packed_sequence = rnn_utils.pack_padded_sequence( + fusion_sequence, + lengths=sequence_lengths.cpu(), + batch_first=True, + enforce_sorted=False, + ) + _, (fusion_hidden, _) = self.fusion_lstm(packed_sequence) + fused_repr = self.dropout_layer(fusion_hidden[-1]) + logits = self.classifier(fused_repr) + return logits + + +class MedFuse(BaseModel): + """MedFuse model for multi-modal clinical prediction. + + This model fuses clinical time-series (EHR) and chest X-ray (CXR) + representations using an LSTM-based fusion module. It handles missing CXR + via variable sequence lengths, where each sample contributes either: + + - ``[v_ehr]`` (EHR only), or + - ``[v_ehr, v_cxr]`` (EHR + CXR). + + Paper: + Hayat, N., Geras, K. J., & Shamout, F. E. (2022). + MedFuse: Multi-modal fusion with clinical time-series data and chest + X-ray images. MLHC 2022. + + Args: + dataset: Sample dataset used for model configuration. + ehr_feature_key: Input key for EHR tensor. Default is ``"ehr"``. + cxr_feature_key: Input key for CXR tensor. Default is ``"cxr"``. + cxr_mask_key: Optional input key for per-sample CXR availability mask. + Default is ``"cxr_mask"``. To surface the mask through the + standard dataloader, include ``"cxr_mask": "tensor"`` in your + dataset's ``input_schema`` with per-sample availability values + (1 = CXR present, 0 = missing). + ehr_hidden_dim: Hidden dimension for EHR encoder. + ehr_num_layers: Number of EHR LSTM layers. + cxr_backbone: CXR ResNet backbone. Default is ``"resnet34"``. + cxr_pretrained: Whether to use pretrained CXR backbone. + fusion_hidden_dim: Hidden dimension for fusion LSTM. + projection_dim: CXR projection output dimension. + dropout: Dropout rate. + + Examples: + >>> from pyhealth.datasets import create_sample_dataset, get_dataloader + >>> from pyhealth.models import MedFuse + >>> samples = [ + ... { + ... "patient_id": "p0", + ... "visit_id": "v0", + ... "ehr": torch.randn(5, 10).tolist(), + ... "cxr": torch.randn(3, 32, 32).tolist(), + ... "label": 1, + ... }, + ... { + ... "patient_id": "p1", + ... "visit_id": "v1", + ... "ehr": torch.randn(5, 10).tolist(), + ... "cxr": torch.randn(3, 32, 32).tolist(), + ... "label": 0, + ... }, + ... ] + >>> dataset = create_sample_dataset( + ... samples=samples, + ... input_schema={"ehr": "tensor", "cxr": "tensor"}, + ... output_schema={"label": "binary"}, + ... ) + >>> loader = get_dataloader(dataset, batch_size=2) + >>> model = MedFuse(dataset=dataset, cxr_pretrained=False) + >>> batch = next(iter(loader)) + >>> output = model(**batch) + >>> output["y_prob"].shape + torch.Size([2, 1]) + """ + + def __init__( + self, + dataset: SampleDataset, + ehr_feature_key: str = "ehr", + cxr_feature_key: str = "cxr", + cxr_mask_key: Optional[str] = "cxr_mask", + ehr_hidden_dim: int = 256, + ehr_num_layers: int = 2, + cxr_backbone: str = "resnet34", + cxr_pretrained: bool = True, + fusion_hidden_dim: int = 512, + projection_dim: int = 256, + dropout: float = 0.5, + ) -> None: + super().__init__(dataset=dataset) + + if len(self.label_keys) != 1: + raise ValueError("MedFuse supports exactly one label key.") + + if ehr_feature_key not in self.feature_keys: + raise ValueError( + f"ehr_feature_key '{ehr_feature_key}' not found in dataset " + f"features: {self.feature_keys}." + ) + + self.ehr_feature_key = ehr_feature_key + self.cxr_feature_key = cxr_feature_key + self.cxr_mask_key = cxr_mask_key + self.label_key = self.label_keys[0] + self.mode = self._resolve_mode(self.dataset.output_schema[self.label_key]) + + self.has_cxr_feature = cxr_feature_key in self.feature_keys + self.ehr_input_dim = self._infer_ehr_input_dim(ehr_feature_key) + + output_size = self.get_output_size() + self.layer = MedFuseLayer( + ehr_input_dim=self.ehr_input_dim, + ehr_hidden_dim=ehr_hidden_dim, + ehr_num_layers=ehr_num_layers, + cxr_backbone=cxr_backbone, + cxr_pretrained=cxr_pretrained, + fusion_hidden_dim=fusion_hidden_dim, + projection_dim=projection_dim, + dropout=dropout, + num_labels=output_size, + ) + + def _extract_feature_value( + self, + feature_key: str, + feature: torch.Tensor | tuple[torch.Tensor, ...], + ) -> torch.Tensor: + """Extracts the semantic ``value`` tensor from a feature payload.""" + if isinstance(feature, torch.Tensor): + return feature + + if feature_key not in self.dataset.input_processors: + raise ValueError( + f"Feature '{feature_key}' is not defined in dataset processors." + ) + + schema = self.dataset.input_processors[feature_key].schema() + if "value" in schema: + value = feature[schema.index("value")] + if isinstance(value, torch.Tensor): + return value + + raise ValueError( + f"Feature '{feature_key}' must provide a tensor value in its schema." + ) + + def _infer_ehr_input_dim(self, ehr_feature_key: str) -> int: + """Infers EHR input feature dimension from processed dataset samples.""" + for sample in self.dataset: + if ehr_feature_key not in sample: + continue + feature = sample[ehr_feature_key] + if isinstance(feature, tuple): + value = self._extract_feature_value(ehr_feature_key, feature) + elif isinstance(feature, torch.Tensor): + value = feature + else: + value = torch.as_tensor(feature) + + if value.dim() >= 2: + return int(value.shape[-1]) + + raise ValueError( + "Unable to infer EHR input dimension. Ensure EHR feature tensors " + "have shape [seq_len, feature_dim]." + ) + + def forward( + self, + **kwargs: torch.Tensor | tuple[torch.Tensor, ...], + ) -> Dict[str, torch.Tensor]: + """Forward propagation. + + Args: + **kwargs: Model inputs. Expected keys include: + - ``ehr_feature_key`` (required): EHR tensor. + - ``cxr_feature_key`` (optional): CXR tensor. + - ``cxr_mask_key`` (optional): CXR availability mask. + - ``label_key`` (optional): Label tensor for loss computation. + + Returns: + A dictionary containing: + - ``logit``: Raw logits. + - ``y_prob``: Predicted probabilities. + - ``loss``: Loss tensor when labels are provided. + - ``y_true``: Ground-truth labels when provided. + """ + if self.ehr_feature_key not in kwargs: + raise ValueError( + f"Missing required EHR feature key: '{self.ehr_feature_key}'." + ) + + ehr_input = self._extract_feature_value( + self.ehr_feature_key, + cast(torch.Tensor | tuple[torch.Tensor, ...], kwargs[self.ehr_feature_key]), + ) + ehr_input = ehr_input.to(self.device).float() + + cxr_input: Optional[torch.Tensor] = None + if self.cxr_feature_key in kwargs: + cxr_input = self._extract_feature_value( + self.cxr_feature_key, + cast( + torch.Tensor | tuple[torch.Tensor, ...], + kwargs[self.cxr_feature_key], + ), + ) + cxr_input = cxr_input.to(self.device).float() + + cxr_mask: Optional[torch.Tensor] = None + if self.cxr_mask_key is not None and self.cxr_mask_key in kwargs: + raw_mask = kwargs[self.cxr_mask_key] + if isinstance(raw_mask, tuple): + raise ValueError("cxr_mask must be provided as a tensor.") + if isinstance(raw_mask, torch.Tensor): + cxr_mask = raw_mask.to(self.device) + else: + cxr_mask = torch.as_tensor(raw_mask, device=self.device) + + logits = self.layer( + ehr_input=ehr_input, + cxr_input=cxr_input, + cxr_mask=cxr_mask, + ) + y_prob = self.prepare_y_prob(logits) + + 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) + loss = self.get_loss_function()(logits, y_true) + results["loss"] = loss + results["y_true"] = y_true + + return results diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index cadb479ce..3600b3494 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -29,6 +29,7 @@ LengthOfStayPredictionOMOP, ) from .length_of_stay_stagenet_mimic4 import LengthOfStayStageNetMIMIC4 +from .medfuse_mortality import MedFuseMortalityMIMIC4 from .medical_coding import MIMIC3ICD9Coding from .medical_transcriptions_classification import MedicalTranscriptionsClassification from .mortality_prediction import ( diff --git a/pyhealth/tasks/medfuse_mortality.py b/pyhealth/tasks/medfuse_mortality.py new file mode 100644 index 000000000..8d1053562 --- /dev/null +++ b/pyhealth/tasks/medfuse_mortality.py @@ -0,0 +1,340 @@ +"""MedFuse in-hospital mortality task for MIMIC-IV (EHR time series + chest X-ray). + +Produces MedFuse-ready samples directly from a :class:`~pyhealth.datasets.MIMIC4Dataset`: + + {"ehr": <[n_bins, F] time-series>, "cxr": , "cxr_mask": [1.0], + "mortality": <0/1>} + +The ``ehr`` and ``cxr`` keys match :class:`~pyhealth.models.MedFuse`'s default feature +keys, so the resulting ``SampleDataset`` feeds MedFuse with no model changes. The EHR +signal is configurable between ICU **vitals** (``chartevents``) and chemistry **labs** +(``labevents``); the observation window is a fixed post-admission window (paper default: +first 48h, resampled every 2h). + +Paper: + Hayat, N., Geras, K. J., & Shamout, F. E. (2022). MedFuse: Multi-modal fusion with + clinical time-series data and chest X-ray images. MLHC 2022. + +v1 (this class) requires **both** modalities (paired cohort). The missing-CXR (partial) +cohort is a deferred subclass that overrides :meth:`_resolve_cxr`. +""" + +from __future__ import annotations + +import math +from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import polars as pl + +from .base_task import BaseTask +from .mortality_prediction import MultimodalMortalityPredictionMIMIC4 + + +# --------------------------------------------------------------------------- # +# Feature specs +# --------------------------------------------------------------------------- # +# ICU vitals -> MIMIC-IV `chartevents` itemids (d_items, itemids ~220xxx). +# Feature order is fixed so the produced matrix has a constant column count F. +VITAL_CATEGORIES: Dict[str, List[str]] = { + "heart_rate": ["220045"], + "sbp": ["220050", "220179"], # arterial + non-invasive systolic + "dbp": ["220051", "220180"], # arterial + non-invasive diastolic + "mbp": ["220052", "220181", "225312"], # arterial + non-invasive + cuff mean + "resp_rate": ["220210", "224690"], + "spo2": ["220277"], + "temperature": ["223761", "223762"], # 223761 = Fahrenheit (converted), 223762 = Celsius + "glucose": ["220621", "225664", "226537"], + "gcs_eye": ["220739"], + "gcs_verbal": ["223900"], + "gcs_motor": ["223901"], + "weight": ["226512", "224639"], + "height": ["226730"], +} + +# Per-itemid value transforms. Kept module-level (NOT stored on the task instance) so the +# task's ``vars()`` stays JSON-serializable/deterministic for set_task cache keying. +_ITEMID_TRANSFORM = { + "223761": lambda f: (f - 32.0) * 5.0 / 9.0, # Fahrenheit -> Celsius +} + + +class MedFuseMortalityMIMIC4(BaseTask): + """MedFuse in-hospital mortality prediction on MIMIC-IV (paired EHR + CXR). + + One sample per qualifying ICU stay. The EHR feature is a fixed-window time-series + matrix (``chartevents`` vitals or ``labevents`` labs); the CXR feature is the path to + the last chest X-ray taken during the stay (loaded to a tensor by the ``image`` + processor). Mortality is the admission's ``hospital_expire_flag``. + + Args: + ehr_source: ``"vitals"`` (chartevents) or ``"labs"`` (labevents). Default ``"vitals"``. + window_hours: Observation window length after ICU ``intime``. Default 48. + fixed_window: If True, emit a dense ``[n_bins, F]`` matrix on a fixed grid + (``n_bins = window_hours / bin_hours``) via the ``tensor`` processor. If False, + emit ``(timestamps, values)`` for the ``timeseries`` processor (variable length). + Default True. + bin_hours: Bin width for the fixed grid / resample rate for the variable path. Default 2. + min_los_hours: Minimum ICU length of stay (computed from intime/outtime). Default 48. + min_age: Minimum patient anchor_age. Default 18. + require_cxr: If True (v1), skip stays without a CXR. Default True. + image_size: Square resize for the CXR image processor. Default 224. + cxr_event_type: Event type of CXR metadata events. Default ``"metadata"``. + """ + + task_name: str = "medfuse_mortality_mimic4" + + def __init__( + self, + ehr_source: str = "vitals", + window_hours: float = 48.0, + fixed_window: bool = True, + bin_hours: float = 2.0, + min_los_hours: float = 48.0, + min_age: int = 18, + require_cxr: bool = True, + image_size: int = 224, + cxr_event_type: str = "metadata", + ) -> None: + if ehr_source == "vitals": + self.table = "chartevents" + categories = VITAL_CATEGORIES + elif ehr_source == "labs": + self.table = "labevents" + categories = MultimodalMortalityPredictionMIMIC4.LAB_CATEGORIES + else: + raise ValueError(f"ehr_source must be 'vitals' or 'labs', got {ehr_source!r}.") + + self.ehr_source = ehr_source + self.window_hours = float(window_hours) + self.fixed_window = bool(fixed_window) + self.bin_hours = float(bin_hours) + self.min_los_hours = float(min_los_hours) + self.min_age = int(min_age) + self.require_cxr = bool(require_cxr) + self.image_size = int(image_size) + self.cxr_event_type = str(cxr_event_type) + + # Derived, JSON-safe lookups (feature order fixed -> constant F). + self.feature_names: List[str] = list(categories.keys()) + self.n_features: int = len(self.feature_names) + self.n_bins: int = int(round(self.window_hours / self.bin_hours)) + self.all_itemids: List[str] = [i for ids in categories.values() for i in ids] + self.itemid_to_feat: Dict[str, int] = { + itemid: idx + for idx, name in enumerate(self.feature_names) + for itemid in categories[name] + } + + if self.fixed_window: + ehr_schema: Any = "tensor" + else: + ehr_schema = ("timeseries", {"sampling_rate": timedelta(hours=self.bin_hours)}) + self.input_schema: Dict[str, Any] = { + "ehr": ehr_schema, + "cxr": ("image", {"image_size": self.image_size}), + "cxr_mask": "tensor", + } + self.output_schema: Dict[str, str] = {"mortality": "binary"} + super().__init__() + + # ------------------------------------------------------------------ # + # EHR extraction + # ------------------------------------------------------------------ # + def _parse_events( + self, patient: Any, t0: datetime, t1: datetime + ) -> List[Tuple[datetime, int, float]]: + """Returns [(timestamp, feature_idx, value)] for in-window, in-spec measurements.""" + df = patient.get_events( + event_type=self.table, start=t0, end=t1, return_df=True + ) + if df is None or df.height == 0: + return [] + + icol, vcol = f"{self.table}/itemid", f"{self.table}/valuenum" + df = df.select( + pl.col("timestamp"), + pl.col(icol).cast(pl.Utf8).alias("itemid"), + pl.col(vcol).cast(pl.Float64, strict=False).alias("valuenum"), + ).filter(pl.col("itemid").is_in(self.all_itemids)) + if df.height == 0: + return [] + + out: List[Tuple[datetime, int, float]] = [] + for ts, itemid, val in zip( + df["timestamp"].to_list(), + df["itemid"].to_list(), + df["valuenum"].to_list(), + ): + if val is None or (isinstance(val, float) and math.isnan(val)): + continue + feat_idx = self.itemid_to_feat.get(itemid) + if feat_idx is None: + continue + transform = _ITEMID_TRANSFORM.get(itemid) + out.append((ts, feat_idx, transform(val) if transform else float(val))) + return out + + def _build_fixed( + self, events: List[Tuple[datetime, int, float]], t0: datetime + ) -> Optional[List[List[float]]]: + """Bins events onto a fixed [n_bins, F] grid (mean per bin), then forward-fills.""" + sums = np.zeros((self.n_bins, self.n_features), dtype=np.float64) + counts = np.zeros((self.n_bins, self.n_features), dtype=np.float64) + bin_seconds = self.bin_hours * 3600.0 + for ts, feat_idx, value in events: + b = int((ts - t0).total_seconds() // bin_seconds) + if 0 <= b < self.n_bins: + sums[b, feat_idx] += value + counts[b, feat_idx] += 1 + if counts.sum() == 0: + return None + matrix = np.divide( + sums, counts, out=np.full_like(sums, np.nan), where=counts > 0 + ) + matrix = self._forward_fill(matrix) + return matrix.tolist() + + def _build_variable( + self, events: List[Tuple[datetime, int, float]] + ) -> Optional[Tuple[List[datetime], np.ndarray]]: + """One row per distinct timestamp (mean per feature); timeseries processor resamples.""" + per_ts: Dict[datetime, Tuple[np.ndarray, np.ndarray]] = {} + for ts, feat_idx, value in events: + s, c = per_ts.setdefault( + ts, + (np.zeros(self.n_features), np.zeros(self.n_features)), + ) + s[feat_idx] += value + c[feat_idx] += 1 + if not per_ts: + return None + timestamps = sorted(per_ts) + rows = [ + np.divide( + s, c, out=np.full_like(s, np.nan), where=c > 0 + ) + for s, c in (per_ts[ts] for ts in timestamps) + ] + return timestamps, np.vstack(rows) + + @staticmethod + def _forward_fill(matrix: np.ndarray) -> np.ndarray: + """Forward-fills NaNs down the time axis; leading NaNs become 0.0.""" + n_steps, n_features = matrix.shape + for j in range(n_features): + last = 0.0 + for i in range(n_steps): + if np.isnan(matrix[i, j]): + matrix[i, j] = last + else: + last = matrix[i, j] + return matrix + + def _build_ehr(self, patient: Any, t0: datetime, t1: datetime): + events = self._parse_events(patient, t0, t1) + if not events: + return None + if self.fixed_window: + return self._build_fixed(events, t0) + return self._build_variable(events) + + # ------------------------------------------------------------------ # + # CXR extraction (overridden by the v2 partial-cohort subclass) + # ------------------------------------------------------------------ # + def _resolve_cxr( + self, patient: Any, intime: datetime, outtime: datetime + ) -> Optional[Tuple[str, float]]: + """Returns (image_path, mask) for the last CXR in the stay, or None to skip (v1).""" + best_path: Optional[str] = None + best_ts: Optional[datetime] = None + for event in patient.get_events(event_type=self.cxr_event_type): + ts = getattr(event, "timestamp", None) + if ts is None or ts < intime or ts > outtime: + continue + path = getattr(event, "image_path", None) + if not path: + continue + if best_ts is None or ts > best_ts: + best_ts, best_path = ts, path + if best_path is None: + return None + return best_path, 1.0 + + # ------------------------------------------------------------------ # + # Label + # ------------------------------------------------------------------ # + @staticmethod + def _mortality(adm_flag: Dict[Any, Any], hadm_id: Any) -> Optional[int]: + flag = adm_flag.get(hadm_id) + if flag is None: + return None + try: + return int(flag) + except (TypeError, ValueError): + return 1 if str(flag).strip() == "1" else 0 + + @staticmethod + def _parse_dt(value: Any) -> Optional[datetime]: + if value is None: + return None + if isinstance(value, datetime): + return value + try: + return datetime.fromisoformat(str(value)) + except ValueError: + return None + + # ------------------------------------------------------------------ # + # Task entry point + # ------------------------------------------------------------------ # + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + demographics = patient.get_events(event_type="patients") + if not demographics: + return [] + try: + age = int(demographics[0].anchor_age) + except (TypeError, ValueError, AttributeError): + age = None + if age is not None and age < self.min_age: + return [] + + admissions = patient.get_events(event_type="admissions") + adm_flag = {a.hadm_id: getattr(a, "hospital_expire_flag", None) for a in admissions} + + samples: List[Dict[str, Any]] = [] + for stay in patient.get_events(event_type="icustays"): + intime = self._parse_dt(getattr(stay, "timestamp", None)) + outtime = self._parse_dt(getattr(stay, "outtime", None)) + if intime is None or outtime is None: + continue + if (outtime - intime).total_seconds() / 3600.0 < self.min_los_hours: + continue + t0 = intime + t1 = intime + timedelta(hours=self.window_hours) + + ehr = self._build_ehr(patient, t0, t1) + if ehr is None: + continue + + resolved = self._resolve_cxr(patient, intime, outtime) + if resolved is None: + continue + image_path, mask = resolved + + label = self._mortality(adm_flag, getattr(stay, "hadm_id", None)) + if label is None: + continue + + samples.append( + { + "patient_id": patient.patient_id, + "stay_id": getattr(stay, "stay_id", None), + "ehr": ehr, + "cxr": image_path, + "cxr_mask": [float(mask)], + "mortality": label, + } + ) + return samples