diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py new file mode 100644 index 000000000..d0806c8e7 --- /dev/null +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -0,0 +1,698 @@ +"""End-to-end protocol runner for Unified Embedding on MIMIC-IV. + +Trains and evaluates a unified-embedding model (MLP / RNN / Transformer / +BottleneckTransformer / EHRMamba / JambaEHR) on a MIMIC-IV mortality task, +then writes per-sample predictions to CSV. + +Tasks +----- +--task labs (default) + LabsMIMIC4: 10-dim lab vectors only. + +--task notes_labs (recommended for multimodal) + NotesLabsMIMIC4: notes + 10-dim lab vectors. + +--task notes_labs_cxr + NotesLabsCXRMIMIC4: notes + labs + chest-xray. + +Example +------- + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root /path/to/mimiciv/2.2 \\ + --task labs \\ + --model transformer \\ + --heads 4 --num-layers 2 \\ + --dev --device cpu \\ + --epochs 10 --batch-size 32 --lr 1e-3 \\ + --output-dir ./output/unified_e2e + + # EHRMamba on full dataset (no --dev): + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root /data/mimic-iv/2.2 --note-root /data/mimic-iv/note \\ + --task notes_labs --model ehrmamba \\ + --embedding-dim 128 --num-layers 2 --seed 42 + + # JambaEHR: + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root /data/mimic-iv/2.2 --note-root /data/mimic-iv/note \\ + --task notes_labs --model jambaehr \\ + --embedding-dim 128 --jamba-transformer-layers 2 --jamba-mamba-layers 2 +""" + +from __future__ import annotations + +import argparse +import csv +import warnings +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + +import numpy as np + +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.processors import fit_lab_standardizer, lab_standardizer_fit_scope +from pyhealth.tasks.multimodal_mimic4 import ( + LabsMIMIC4, + NotesLabsCXRMIMIC4, + NotesLabsMIMIC4, +) +from pyhealth.trainer import Trainer +from pyhealth.utils import set_seed, write_run_config + + +class WandbLogger: + + def __init__( + self, + enabled: bool, + project: str, + entity: Optional[str], + run_name: str, + tags: list[str], + config: Dict[str, Any], + ) -> None: + self.enabled = enabled + self._run = None + if self.enabled: + import wandb + + self._run = wandb.init( + project=project, + entity=entity, + name=run_name, + tags=tags, + config=config, + ) + + def log(self, data: Dict[str, Any], step: Optional[int] = None) -> None: + if self.enabled: + self._run.log(data, step=step) + + def finish(self) -> None: + if self.enabled: + self._run.finish() + + +def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: + ehr_tables = ["labevents"] + note_tables = None + cxr_kwargs = {} + + if args.task == "notes_labs": + if not args.note_root: + raise ValueError("--task notes_labs requires --note-root.") + note_tables = ["discharge", "radiology"] + + if args.task == "notes_labs_cxr": + if not args.note_root: + raise ValueError("--task notes_labs_cxr requires --note-root.") + if not args.cxr_root: + raise ValueError("--task notes_labs_cxr requires --cxr-root.") + note_tables = ["discharge", "radiology"] + 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 if args.dev else False, + num_workers=args.num_workers, + **cxr_kwargs, + ) + + +def _build_task(args: argparse.Namespace): + if args.task == "notes_labs": + return NotesLabsMIMIC4( + window_hours=args.observation_window_hours, + ) + if args.task == "notes_labs_cxr": + return NotesLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + if args.task == "labs": + return LabsMIMIC4(window_hours=args.observation_window_hours) + raise ValueError(f"Unknown task: {args.task}") + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any, str]: + """Split by patient, falling back to by-sample only if that yields nothing. + + The fallback is leaky: a patient with several admissions can then land in + both train and test, which inflates the metrics. It only triggers on tiny + cohorts, but it must not trigger silently, so the mode is returned and + recorded alongside the run's results. + """ + 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: + warnings.warn( + "split_by_patient produced an empty split, falling back to " + "split_by_sample. The same patient may now appear in train and " + "test, so these metrics are optimistic and not comparable to " + "patient-split runs.", + RuntimeWarning, + stacklevel=2, + ) + 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, "by_sample_fallback_leaky" + return train_ds, val_ds, test_ds, "by_patient" + + +def _build_model( + args: argparse.Namespace, + sample_dataset: Any, + numeric_standardizers: dict[str, Any] | None = None, +): + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + freeze_text_encoder=args.freeze_encoder, + numeric_standardizers=numeric_standardizers, + ) + + 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: + set_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." + ) + + split_seed = args.seed if args.split_seed is None else args.split_seed + train_ds, val_ds, test_ds, split_mode = _split_dataset( + sample_dataset, seed=split_seed + ) + + numeric_standardizers: dict[str, Any] = {} + if "labs" in sample_dataset.input_processors and not args.no_lab_standardization: + if "labs_mask" not in sample_dataset.input_processors: + raise RuntimeError( + "Lab standardisation requires the labs_mask observation field." + ) + lab_standardizer = fit_lab_standardizer( + train_ds, + value_field="labs", + fit_scope=lab_standardizer_fit_scope(train_ds, value_field="labs"), + ) + numeric_standardizers["labs"] = lab_standardizer + print( + "[lab-standardization] fitted on train split only: " + f"counts={lab_standardizer.observed_count.tolist()} " + f"mean={lab_standardizer.mean.tolist()} std={lab_standardizer.std.tolist()}" + ) + elif "labs" in sample_dataset.input_processors: + print("[lab-standardization] disabled; reproducing raw-lab baseline.") + + model = _build_model(args, sample_dataset, numeric_standardizers) + + loader_kwargs = { + "num_workers": args.loader_num_workers, + "pin_memory": args.pin_memory, + "persistent_workers": args.persistent_workers, + "prefetch_factor": ( + args.prefetch_factor if args.loader_num_workers > 0 else None + ), + } + train_loader = get_dataloader( + train_ds, batch_size=args.batch_size, shuffle=True, **loader_kwargs + ) + val_loader = ( + get_dataloader( + val_ds, batch_size=args.batch_size, shuffle=False, **loader_kwargs + ) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader( + test_ds, batch_size=args.batch_size, shuffle=False, **loader_kwargs + ) + if len(test_ds) > 0 + else None + ) + + if test_loader is not None: + inference_loader, eval_split = test_loader, "test" + elif val_loader is not None: + inference_loader, eval_split = val_loader, "val" + warnings.warn( + "No test split available; reporting predictions from the VALIDATION " + "split. These are not test metrics.", + RuntimeWarning, + stacklevel=2, + ) + else: + inference_loader, eval_split = train_loader, "train" + warnings.warn( + "No test or validation split available; reporting predictions from " + "the TRAINING split. These metrics are held-in and meaningless as a " + "generalisation estimate.", + RuntimeWarning, + stacklevel=2, + ) + + # The task MUST be in the name. Without it, two arms of the same comparison + # at the same seed, for example --task labs and --task notes_labs, resolve + # to one directory and the second run overwrites the first. + exp_name = f"{args.task}_{args.model}_seed{args.seed}" + output_dir = Path(args.output_dir) + + wandb_logger = WandbLogger( + enabled=args.wandb, + project=args.wandb_project, + entity=args.wandb_entity, + run_name=args.wandb_run_name or exp_name, + tags=args.wandb_tags.split(",") if args.wandb_tags else [args.task, args.model], + config=vars(args), + ) + + 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, + ) + + # BottleneckTransformer is more fragile on full MIMIC-IV with no warmup. + # Use safer defaults unless explicitly overridden from CLI. + 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-4 + 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 + + optimizer_params["lr"] = effective_lr + + write_run_config( + str(output_dir / exp_name), + { + **vars(args), + "resolved_lr": effective_lr, + "resolved_max_grad_norm": effective_max_grad_norm, + "split_mode": split_mode, + "resolved_split_seed": split_seed, + "eval_split": eval_split, + "n_train": len(train_ds), + "n_val": len(val_ds), + "n_test": len(test_ds), + "lab_standardization": bool(numeric_standardizers), + }, + ) + + if args.epochs > 0 and len(train_ds) > 0: + metrics_history = 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, + patience=args.patience, + use_amp=args.use_amp, + amp_dtype=args.amp_dtype, + ) + for epoch_record in metrics_history: + wandb_logger.log(epoch_record, step=epoch_record["epoch"]) + + if wandb_logger.enabled and test_loader is not None: + test_scores = trainer.evaluate(test_loader) + wandb_logger.log({f"test_{k}": v for k, v in test_scores.items()}) + + 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) + + wandb_logger.finish() + + return output_csv + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run E2E unified embedding on MIMIC-IV with any of six sequence heads." + ) + 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") + + parser.add_argument( + "--task", + type=str, + choices=["labs", "notes_labs", "notes_labs_cxr"], + default="labs", + help=( + "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. " + "Recommended for multimodal. " + "notes_labs_cxr: notes_labs plus in-window chest X-rays; requires " + "--note-root and --cxr-root." + ), + ) + parser.add_argument( + "--model", + type=str, + choices=["mlp", "rnn", "transformer", "bottleneck_transformer", + "ehrmamba", "jambaehr"], + default="rnn", + ) + + parser.add_argument("--embedding-dim", type=int, default=128) + parser.add_argument("--hidden-dim", type=int, default=128) + parser.add_argument("--dropout", type=float, default=0.1) + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument( + "--batch-size", + "--batch_size", + dest="batch_size", + type=int, + default=32, + ) + parser.add_argument( + "--lr", + "--learning-rate", + "--learning_rate", + dest="lr", + type=float, + default=None, + help="Learning rate. Default is 1e-4 for all models.", + ) + parser.add_argument( + "--adam-eps", + type=float, + default=None, + help=( + "Adam epsilon. Default is model-specific: 1e-8 for non-BT models, " + "1e-6 for bottleneck_transformer." + ), + ) + parser.add_argument("--weight-decay", type=float, default=0.0) + parser.add_argument("--device", type=str, default=None) + parser.add_argument( + "--use-amp", + action="store_true", + help="Enable automatic mixed precision training to reduce GPU memory usage.", + ) + parser.add_argument( + "--amp-dtype", + "--amp_dtype", + dest="amp_dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + help="AMP dtype when --use-amp is set. bf16 is more stable (default).", + ) + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument( + "--loader-num-workers", + type=int, + default=0, + help="Worker processes for runtime batch loading/collation.", + ) + parser.add_argument("--pin-memory", action="store_true", default=False) + parser.add_argument("--persistent-workers", action="store_true", default=False) + parser.add_argument("--prefetch-factor", type=int, default=2) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--split-seed", + "--split_seed", + dest="split_seed", + type=int, + default=None, + help=( + "Patient-split RNG seed. Defaults to --seed. Set separately when " + "a tiny --dev cohort puts no positives in val/test." + ), + ) + parser.add_argument("--patience", type=int, default=None) + 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( + "--patients", + type=int, + default=None, + help=( + "Not a supported flag. Do not pass this. Use --dev N for a " + "patient-limited smoke, and omit both for the full table." + ), + ) + parser.add_argument( + "--observation-window-hours", + type=int, + default=None, + help=( + "If set, collect labs/CXR/radiology only this many hours from each " + "admission. Default: full stay (through discharge)." + ), + ) + parser.add_argument( + "--freeze-encoder", + action="store_true", + default=False, + help=( + "Freeze pretrained BERT text encoder weights and train only the " + "downstream backbone (RNN/Transformer head + projection layer). " + ), + ) + 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", + "--num_layers", + dest="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, + help=( + "Gradient clipping max norm. Default is model-specific: None for " + "non-BT models, 0.5 for bottleneck_transformer." + ), + ) + + parser.add_argument( + "--wandb", + action="store_true", + default=False, + help="Log training/eval metrics to Weights & Biases.", + ) + parser.add_argument("--wandb-project", type=str, default="pyhealth-mortality") + parser.add_argument("--wandb-entity", type=str, default=None) + parser.add_argument( + "--wandb-run-name", + type=str, + default=None, + help="Defaults to '{task}_{model}_seed{seed}' if unset.", + ) + parser.add_argument( + "--wandb-tags", + type=str, + default=None, + help="Comma-separated wandb tags, e.g. 'labs,rnn'. Defaults to '{task},{model}' if unset.", + ) + + parser.add_argument( + "--mamba-state-size", + "--mamba_state_size", + dest="mamba_state_size", + type=int, + default=16, + help="SSM state size for EHRMamba and JambaEHR blocks.", + ) + parser.add_argument( + "--mamba-conv-kernel", + "--mamba_conv_kernel", + "--conv-kernel", + "--conv_kernel", + dest="mamba_conv_kernel", + type=int, + default=4, + help="Causal conv kernel size for EHRMamba and JambaEHR blocks.", + ) + parser.add_argument( + "--jamba-transformer-layers", + "--jamba_transformer_layers", + dest="jamba_transformer_layers", + type=int, + default=2, + help="Number of Transformer (attention) layers in JambaEHR.", + ) + parser.add_argument( + "--jamba-mamba-layers", + "--jamba_mamba_layers", + dest="jamba_mamba_layers", + type=int, + default=6, + help="Number of Mamba (SSM) layers in JambaEHR. Library default is 6; " + "pass 2 for a depth-matched comparison against --num-layers 2.", + ) + parser.add_argument( + "--no-lab-standardization", + action="store_true", + default=False, + help="Disable train-split lab z-scoring (raw-lab ablation).", + ) + + args = parser.parse_args() + if args.patients is not None: + parser.error( + "--patients is not a supported flag and is not a 5-patient smoke. " + "Omit it for the full table; use --dev N for a patient-limited run." + ) + return args + + +if __name__ == "__main__": + args = parse_args() + output_csv_path = run(args) + print(f"Saved predictions to: {output_csv_path}") \ No newline at end of file diff --git a/pyhealth/data/data.py b/pyhealth/data/data.py index 14b1b526c..39dd1206f 100644 --- a/pyhealth/data/data.py +++ b/pyhealth/data/data.py @@ -155,7 +155,7 @@ def _filter_by_time_range_fast(self, df: pl.DataFrame, start: Optional[datetime] start_idx = np.searchsorted(ts_col, np.datetime64(start, "ms"), side="left") if end is not None: end_idx = np.searchsorted(ts_col, np.datetime64(end, "ms"), side="right") - return df.slice(start_idx, end_idx - start_idx) + return df.slice(start_idx, max(0, end_idx - start_idx)) def _filter_by_event_type_regular(self, df: pl.DataFrame, event_type: Optional[str]) -> pl.DataFrame: """Regular filtering by event type. Time complexity: O(n).""" diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 3d449d579..b03d9b925 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -3,7 +3,7 @@ import pickle from abc import ABC from pathlib import Path -from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable +from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable, Union import functools import operator from urllib.parse import urlparse, urlunparse @@ -71,6 +71,17 @@ def clean_path(path: str) -> str: return str(Path(path).expanduser().resolve()) +def resolve_table_path(root: str, file_path: str) -> str: + """Resolve a table file_path against dataset root. + + Absolute paths and URLs are kept as-is so a generated file that cannot + live in a read-only data root (PhysioNet) can still be loaded from cache. + """ + if is_url(file_path) or os.path.isabs(file_path): + return clean_path(file_path) + return clean_path(f"{root}/{file_path}") + + def path_exists(path: str) -> bool: """ Check if a path exists. @@ -84,7 +95,13 @@ def path_exists(path: str) -> bool: except requests.RequestException: return False else: - return Path(path).exists() + try: + return Path(path).exists() + except OSError: + # Treat unreadable paths (e.g. stale/corrupted filesystem + # entries that raise I/O errors on stat) as non-existent so + # callers can fall back to an alternate extension. + return False def _csv_tsv_gz_path(path: str) -> str: @@ -314,16 +331,7 @@ class BaseDataset(ABC): dataset_name (str): Name of the dataset. config (dict): Configuration loaded from a YAML file. global_event_df (pl.LazyFrame): The global event data frame. - dev (bool): Whether to enable dev mode (limit to 1000 patients). - - Examples: - >>> from pyhealth.datasets import BaseDataset - >>> dataset = BaseDataset( - ... root="/path/to/source", - ... tables=["patients", "diagnoses"], - ... config_path="/path/to/config.yaml", - ... ) - >>> dataset.stats() + dev (Union[bool, int]): Whether to enable dev mode. If True, limit to 1000 patients. If an int, limit to that many patients. """ def __init__( @@ -334,7 +342,7 @@ def __init__( config_path: Optional[str] = None, cache_dir: str | Path | None = None, num_workers: int = 1, - dev: bool = False, + dev: Union[bool, int] = False, ): """Initializes the BaseDataset. @@ -351,7 +359,7 @@ def __init__( - **str** or **Path**: Used as the root cache directory path. A UUID is appended to the provided path to capture dataset configuration. num_workers (int): Number of worker processes for parallel operations. - dev (bool): Whether to run in dev mode (limits to 1000 patients). + dev (Union[bool, int]): Whether to run in dev mode. If True, limits to 1000 patients. If an int, limits to that many patients. """ if len(set(tables)) != len(tables): logger.warning("Duplicate table names in tables list. Removing duplicates.") @@ -571,30 +579,52 @@ def _event_transform(self, output_dir: Path) -> None: compute_ok = False try: df = self.load_data() - with DaskCluster( - n_workers=self.num_workers, - threads_per_worker=1, - processes=not in_notebook(), - # Use cache_dir for Dask's scratch space to avoid filling up /tmp or home directory - local_directory=str(self.create_tmpdir()), - ) as cluster: - with DaskClient(cluster) as client: - if self.dev: - logger.info("Dev mode enabled: limiting to 1000 patients") - patients = df["patient_id"].unique().head(1000).tolist() - filter = df["patient_id"].isin(patients) - df = df[filter] - - logger.info(f"Caching event dataframe to {output_dir}...") - collection = df.sort_values("patient_id").to_parquet( - output_dir, - write_index=False, - compute=False, - ) - handle = client.compute(collection) - dask_progress(handle) - handle.result() # type: ignore - compute_ok = True # Data is fully written to disk + disable_distributed = os.environ.get( + "PYHEALTH_DISABLE_DASK_DISTRIBUTED", "0" + ) == "1" + + if disable_distributed: + logger.info( + "PYHEALTH_DISABLE_DASK_DISTRIBUTED=1 detected; using local dask scheduler." + ) + if self.dev: + n = 1000 if self.dev is True else int(self.dev) + logger.info(f"Dev mode enabled: limiting to {n} patients") + patients = df["patient_id"].unique().head(n, compute=True).tolist() + patient_filter = df["patient_id"].isin(patients) + df = df[patient_filter] + + logger.info(f"Caching event dataframe to {output_dir}...") + df.sort_values("patient_id").to_parquet( + output_dir, + write_index=False, + compute=True, + ) + else: + with DaskCluster( + n_workers=self.num_workers, + threads_per_worker=1, + processes=not in_notebook(), + # Use cache_dir for Dask's scratch space to avoid filling up /tmp or home directory + local_directory=str(self.create_tmpdir()), + ) as cluster: + with DaskClient(cluster) as client: + if self.dev: + logger.info(f"Dev mode enabled: limiting to {1000 if self.dev is True else int(self.dev)} patients") + patients = df["patient_id"].unique().head(1000 if self.dev is True else int(self.dev)).tolist() + filter = df["patient_id"].isin(patients) + df = df[filter] + + logger.info(f"Caching event dataframe to {output_dir}...") + collection = df.sort_values("patient_id").to_parquet( + output_dir, + write_index=False, + compute=False, + ) + 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 @@ -667,8 +697,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: Raises: ValueError: If the table is not found in the config. - FileNotFoundError: If the source file (CSV/TSV or Parquet) for the - table or join is not found. + FileNotFoundError: If the CSV file for the table or join is not found. """ assert self.config is not None, "Config must be provided to load tables" @@ -676,8 +705,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: raise ValueError(f"Table {table_name} not found in config") table_cfg = self.config.tables[table_name] - csv_path = f"{self.root}/{table_cfg.file_path}" - csv_path = clean_path(csv_path) + csv_path = resolve_table_path(self.root, table_cfg.file_path) logger.info(f"Scanning table: {table_name} from {csv_path}") df = self._scan_table(csv_path) @@ -696,8 +724,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: # Handle joins for join_cfg in table_cfg.join: - other_csv_path = f"{self.root}/{join_cfg.file_path}" - other_csv_path = clean_path(other_csv_path) + other_csv_path = resolve_table_path(self.root, join_cfg.file_path) logger.info(f"Joining with table: {other_csv_path}") join_df = self._scan_table(other_csv_path) join_df = join_df.rename(columns=str.lower) @@ -723,21 +750,14 @@ def load_table(self, table_name: str) -> dd.DataFrame: timestamp_series: dd.Series = functools.reduce( operator.add, (df[col].astype("string") for col in timestamp_col) ) - timestamp_series = dd.to_datetime( - timestamp_series, - format=timestamp_format, - errors="raise", - ) - elif pd.api.types.is_datetime64_any_dtype(df[timestamp_col].dtype): - # Typed sources (e.g. Parquet) already carry native timestamps: - # skip the string round-trip and only normalize the unit below. - timestamp_series: dd.Series = df[timestamp_col] else: - timestamp_series = dd.to_datetime( - df[timestamp_col].astype("string"), - format=timestamp_format, - errors="raise", - ) + timestamp_series: dd.Series = df[timestamp_col].astype("string") + + timestamp_series: dd.Series = dd.to_datetime( + timestamp_series, + format=timestamp_format, + errors="raise", + ) df: dd.DataFrame = df.assign( timestamp=timestamp_series.astype("datetime64[ms]") ) @@ -1165,4 +1185,4 @@ def _main_guard(self, func_name: str): f"{func_name} method accessed from a non-main process. This may lead to unexpected behavior.\n" + "Consider use __name__ == '__main__' guard when using multiprocessing." ) - exit(1) + exit(1) \ No newline at end of file diff --git a/pyhealth/datasets/collate.py b/pyhealth/datasets/collate.py index 9e4c113c0..fe72370ed 100644 --- a/pyhealth/datasets/collate.py +++ b/pyhealth/datasets/collate.py @@ -17,14 +17,51 @@ from typing import Any import torch -from torch.nn.utils.rnn import pad_sequence +import torch.nn.functional as F + + +def _pad_stack(tensors: list[torch.Tensor]) -> torch.Tensor: + """Right-pad same-rank tensors to the per-dimension max, then stack. + + ``pad_sequence`` only pads dimension 0 and requires every trailing dimension + to already match. Tokenized notes are ``(n_notes, seq_len)`` and, once the + text processor pads to the longest note in a sample rather than to a fixed + ``max_length``, BOTH dimensions vary across samples. + """ + if len({t.dim() for t in tensors}) != 1: + raise ValueError("cannot pad tensors of differing rank") + target = [max(t.shape[d] for t in tensors) for d in range(tensors[0].dim())] + padded = [] + for t in tensors: + spec: list[int] = [] + for d in range(t.dim() - 1, -1, -1): + spec.extend([0, target[d] - t.shape[d]]) + padded.append(F.pad(t, spec) if any(spec) else t) + return torch.stack(padded) def _stack_or_pad(tensors: list[torch.Tensor]) -> torch.Tensor: - """Stack if all shapes match; pad along dim-0 otherwise.""" + """Stack if all shapes match; pad every ragged dimension otherwise.""" if all(t.shape == tensors[0].shape for t in tensors): return torch.stack(tensors) - return pad_sequence(tensors, batch_first=True) + return _pad_stack(tensors) + + +def _pad_mask(tensors: list[torch.Tensor]) -> torch.Tensor: + """Event-level validity for the tensor :func:`_stack_or_pad` just built. + + Batch padding is created here and nowhere else, so it has to be recorded + here. A padded slot carries value 0.0 and time 0.0, which is + indistinguishable from a real measurement taken at admission time, so a + model given no mask treats padding as data. + + This is deliberately NOT called ``mask``. A field may carry its own + ``{field}_mask`` meaning "was this value observed", which is a different + question from "is this slot real". + """ + lengths = torch.tensor([t.shape[0] for t in tensors]) + width = int(lengths.max()) + return torch.arange(width)[None, :] < lengths[:, None] def collate_temporal(batch: list[dict[str, Any]]) -> dict[str, Any]: @@ -62,6 +99,8 @@ def collate_temporal(batch: list[dict[str, Any]]) -> dict[str, Any]: sub_result[sub_key] = [None] * len(sub_vals) elif isinstance(sub_vals[0], torch.Tensor): sub_result[sub_key] = _stack_or_pad(sub_vals) + if sub_key == "time": + sub_result["pad_mask"] = _pad_mask(sub_vals) else: sub_result[sub_key] = sub_vals result[key] = sub_result diff --git a/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml b/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml new file mode 100644 index 000000000..e631de4aa --- /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" \ No newline at end of file diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py index 9d1aa55d8..e92470ba8 100644 --- a/pyhealth/datasets/mimic4.py +++ b/pyhealth/datasets/mimic4.py @@ -1,7 +1,7 @@ import logging import os import warnings -from typing import List, Optional +from typing import List, Optional, Union import pandas as pd import dask.dataframe as dd @@ -223,6 +223,158 @@ 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|resized_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}") + metadata_csv = self.prepare_metadata(root, cache_dir=cache_dir) + if os.path.dirname(os.path.abspath(metadata_csv)) != os.path.abspath(root): + config_path = self._rewrite_sunlab_config( + config_path, metadata_csv, cache_dir or os.path.dirname(metadata_csv) + ) + 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 + + @staticmethod + def _rewrite_sunlab_config( + config_path: str, metadata_csv: str, dest_dir: str + ) -> str: + """Point the sunlab YAML at a metadata CSV that is not under root.""" + with open(config_path, encoding="utf-8") as f: + text = f.read() + rewritten = text.replace( + "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv", + metadata_csv, + ) + os.makedirs(dest_dir, exist_ok=True) + out = os.path.join(dest_dir, "mimic4_cxr_sunlab.generated.yaml") + with open(out, "w", encoding="utf-8") as f: + f.write(rewritten) + return out + + def prepare_metadata( + self, root: str, cache_dir: Optional[str] = None + ) -> str: + 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." + ) + + # The flattened layout appears under more than one directory name + # depending on how the set was produced, so accept either rather than + # hardcoding one and failing on a complete, correct dataset. + candidates = ("images", "resized_images") + images_dir = next( + ( + os.path.join(root, name) + for name in candidates + if os.path.isdir(os.path.join(root, name)) + ), + None, + ) + if images_dir is None: + raise FileNotFoundError( + f"No flattened image directory under {root}. Looked for " + f"{', '.join(candidates)}, each expected to hold " + "{dicom_id}.jpg." + ) + images_subdir = os.path.basename(images_dir) + + 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_subdir, f"{dicom_id}.jpg") + ) + + # Align with existing config conventions by using lowercase headers. + metadata.columns = [col.lower() for col in metadata.columns] + + filename = "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + dest_dirs = [] + if cache_dir: + dest_dirs.append(str(cache_dir)) + dest_dirs.append(root) + + for d in dest_dirs: + existing = os.path.join(d, filename) + if os.path.isfile(existing): + return existing + + last_err: Optional[OSError] = None + for d in dest_dirs: + os.makedirs(d, exist_ok=True) + dest = os.path.join(d, filename) + try: + metadata.to_csv(dest, index=False) + return dest + except OSError as exc: + last_err = exc + continue + raise PermissionError( + f"Could not write {filename} under {dest_dirs}: {last_err}" + ) + + class MIMIC4Dataset(BaseDataset): """ Unified MIMIC-IV dataset with support for EHR, clinical notes, and X-rays. @@ -242,6 +394,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,8 +432,9 @@ 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: bool = False, + dev: Union[bool, int] = False, cache_dir: Optional[str] = None, num_workers: int = 1, ): @@ -340,17 +494,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") @@ -374,4 +544,4 @@ def load_data(self) -> dd.DataFrame: if len(frames) == 1: return frames[0] else: - return dd.concat(frames, axis=0, join="outer") + return dd.concat(frames, axis=0, join="outer") \ No newline at end of file diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index 24c87a1d5..b6607b5b8 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -7,17 +7,18 @@ import torch import litdata from dateutil.parser import parse as dateutil_parse -from torch.nn.utils.rnn import pad_sequence from torch.utils.data import DataLoader from pyhealth import BASE_CACHE_PATH from pyhealth.utils import create_directory +from pyhealth.datasets.collate import _pad_stack MODULE_CACHE_PATH = os.path.join(BASE_CACHE_PATH, "datasets") create_directory(MODULE_CACHE_PATH) -#PyG import for graph-based models +# PyG import for graph-based models try: from torch_geometric.data import Data as PyGData, Batch as PyGBatch + HAS_PYG = True except ImportError: HAS_PYG = False @@ -250,6 +251,9 @@ def collate_fn_dict(batch: List[dict]) -> dict: return {key: [d[key] for d in batch] for key in batch[0]} +PAD_MASK_SUFFIX = "__pad_mask" + + def collate_fn_dict_with_padding(batch: List[dict]) -> dict: """Collates a batch of data into a dictionary with padding for tensor values. @@ -267,40 +271,40 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: for key in keys: values = [sample[key] for sample in batch] - # Check if this is a temporal feature tuple (time, values) - if isinstance(values[0], tuple) and len(values[0]) == 2: - # Handle (time, values) tuples from processors - time_tensors = [v[0] for v in values] - value_tensors = [v[1] for v in values] - - # Collate values - if value_tensors[0].dim() == 0: - # Scalars - collated_values = torch.stack(value_tensors) - elif all(v.shape == value_tensors[0].shape for v in value_tensors): - # All same shape - collated_values = torch.stack(value_tensors) - else: - # Variable shapes, use pad_sequence - collated_values = pad_sequence( - value_tensors, batch_first=True, padding_value=0 - ) - - # Collate times (if present) - collated_times = None - # Check if ALL samples have time (not just some) - if all(t is not None for t in time_tensors): - time_tensors_all = [t for t in time_tensors if t is not None] - if all(t.shape == time_tensors_all[0].shape for t in time_tensors_all): - collated_times = torch.stack(time_tensors_all) + if isinstance(values[0], tuple): + # Generic tuple collation for processor outputs, e.g. + # - (time, value) from StageNet processors + # - (value, mask, token_type_ids, time, type_tag) + # from TupleTimeTextProcessor with tokenizer. + transposed = list(zip(*values)) + collated_elems = [] + + event_lengths: Optional[List[int]] = None + for elem_vals in transposed: + first = elem_vals[0] + + if first is None and all(v is None for v in elem_vals): + collated_elems.append(None) + elif isinstance(first, torch.Tensor): + tensor_vals = list(elem_vals) + if all(v.shape == tensor_vals[0].shape for v in tensor_vals): + collated_elems.append(torch.stack(tensor_vals)) + else: + if event_lengths is None: + event_lengths = [v.shape[0] for v in tensor_vals] + collated_elems.append(_pad_stack(tensor_vals)) else: - collated_times = pad_sequence( - time_tensors_all, batch_first=True, padding_value=0 - ) + collated_elems.append(list(elem_vals)) + + collated[key] = tuple(collated_elems) + if event_lengths is not None: + lengths = torch.tensor(event_lengths) + width = int(lengths.max()) + collated[f"{key}{PAD_MASK_SUFFIX}"] = ( + torch.arange(width)[None, :] < lengths[:, None] + ) - # Return as tuple (time, values) - collated[key] = (collated_times, collated_values) - # PyG Data objects (graph processor output) + # PyG Data objects (graph processor output) elif HAS_PYG and isinstance(values[0], PyGData): collated[key] = PyGBatch.from_data_list(values) @@ -316,9 +320,7 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: # Scalars, treat as stackable collated[key] = torch.stack(values) elif values[0].dim() >= 1: - collated[key] = pad_sequence( - values, batch_first=True, padding_value=0 - ) + collated[key] = _pad_stack(values) else: raise ValueError(f"Unsupported tensor shape: {values[0].shape}") else: @@ -329,7 +331,13 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: def get_dataloader( - dataset: litdata.StreamingDataset, batch_size: int, shuffle: bool = False + dataset: litdata.StreamingDataset, + batch_size: int, + shuffle: bool = False, + num_workers: int = 0, + pin_memory: bool = False, + persistent_workers: bool = False, + prefetch_factor: Optional[int] = None, ) -> DataLoader: """Creates a DataLoader for a given dataset. @@ -337,18 +345,47 @@ def get_dataloader( dataset: The dataset to load data from. batch_size: The number of samples per batch. shuffle: Whether to shuffle the data at every epoch. + num_workers: Number of worker processes that load and collate batches. + pin_memory: Copy CPU tensors into page-locked memory before return. + persistent_workers: Keep the loader workers between epochs. This is valid + only when ``num_workers`` is more than 0. + prefetch_factor: Batches that each worker loads in advance. This is valid + only when ``num_workers`` is more than 0. ``None`` keeps the default + of PyTorch. Returns: A DataLoader instance for the dataset. """ dataset.set_shuffle(shuffle) - dataloader = DataLoader( - dataset, - batch_size=batch_size, - collate_fn=collate_fn_dict_with_padding, - ) + if num_workers < 0: + raise ValueError(f"num_workers must be non-negative, got {num_workers}.") + if persistent_workers and num_workers == 0: + raise ValueError("persistent_workers requires num_workers > 0.") + if prefetch_factor is not None and (num_workers == 0 or prefetch_factor <= 0): + raise ValueError( + "prefetch_factor must be positive and requires num_workers > 0." + ) - return dataloader + loader_kwargs = { + "dataset": dataset, + "batch_size": batch_size, + "collate_fn": collate_fn_dict_with_padding, + "num_workers": num_workers, + "pin_memory": pin_memory, + } + if num_workers > 0: + loader_kwargs["persistent_workers"] = persistent_workers + if prefetch_factor is not None: + loader_kwargs["prefetch_factor"] = prefetch_factor + + # StreamingDataLoader coordinates shard reads across workers. With a single + # process it adds no advantage, so keep the plain DataLoader there. + loader_class = ( + litdata.StreamingDataLoader + if isinstance(dataset, litdata.StreamingDataset) and num_workers > 0 + else DataLoader + ) + return loader_class(**loader_kwargs) def save_processors(sample_dataset, output_dir: str) -> Dict[str, str]: @@ -453,4 +490,4 @@ def load_processors(processor_dir: str) -> Tuple[Dict, Dict]: print(list_nested_levels([[1, [2], [[3]]]])) print(is_homo_list([1, 2, 3])) print(is_homo_list([1, 2, [3]])) - print(is_homo_list([1, 2.0])) + print(is_homo_list([1, 2.0])) \ No newline at end of file diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index 2f30ae673..7464f7a90 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -46,7 +46,7 @@ from .text_embedding import TextEmbedding from .sdoh import SdohClassifier from .medlink import MedLink -from .unified_embedding import UnifiedMultimodalEmbeddingModel, SinusoidalTimeEmbedding +from .embedding import UnifiedMultimodalEmbeddingModel, SinusoidalTimeEmbedding from .califorest import CaliForest from .generators.halo import HALO from .generators.gpt2 import GPT2 diff --git a/pyhealth/models/bottleneck_transformer.py b/pyhealth/models/bottleneck_transformer.py new file mode 100644 index 000000000..c3b07d210 --- /dev/null +++ b/pyhealth/models/bottleneck_transformer.py @@ -0,0 +1,524 @@ +from typing import Any, Dict, List, Optional, Tuple, Union, cast + +import torch +import torch.nn as nn + +from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX +from pyhealth.models import BaseModel +from pyhealth.models.embedding import EmbeddingModel +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel + + +class MultimodalBottleneckTransformerEncoder(nn.Module): + """ + Generalized Bottleneck Transformer Encoder for N modalities. + Based on "Attention Bottlenecks for Multimodal Fusion" (Nagrani et al., NeurIPS 2021). + """ + + def __init__( + self, + n_modality: int, + bottlenecks_n: int, + fusion_startidx: int, + n_layers: int, + n_head: int, + d_model: int, + d_ff: int, + dropout: float = 0.1, + ): + super(MultimodalBottleneckTransformerEncoder, self).__init__() + + self.n_modality = n_modality + self.fusion_startidx = fusion_startidx + self.n_layers = n_layers + self.n_fusion_layers = n_layers - fusion_startidx + self.n_prefusion = fusion_startidx + self.d_model = d_model + self.n_bottlenecks = bottlenecks_n + + # Shared Bottleneck Tokens — small init to avoid early gradient explosion + self.bottlenecks = nn.Parameter(torch.randn(1, bottlenecks_n, d_model) * 0.02) + + # Prefusion Stacks: independent layers per modality + self.prefusion_stacks = nn.ModuleList([ + nn.ModuleList([ + nn.TransformerEncoderLayer( + d_model=d_model, + nhead=n_head, + dim_feedforward=d_ff, + dropout=dropout, + batch_first=True + ) for _ in range(n_modality) + ]) for _ in range(self.n_prefusion) + ]) + + # Fusion Stacks: processes [bottleneck_tokens || modality_tokens] + self.fusion_stacks = nn.ModuleList([ + nn.ModuleList([ + nn.TransformerEncoderLayer( + d_model=d_model, + nhead=n_head, + dim_feedforward=d_ff, + dropout=dropout, + batch_first=True + ) for _ in range(n_modality) + ]) for _ in range(self.n_fusion_layers) + ]) + + def forward_prefusion(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor]) -> List[torch.Tensor]: + for enc_layers in self.prefusion_stacks: + enc_outputs = [] + for modal_idx, enc_layer in enumerate(enc_layers): + # Apply mask to padding tokens (src_key_padding_mask requires True for ignoring) + # True in mask = invalid/padding + enc_out = enc_layer(enc_inputs[modal_idx], src_key_padding_mask=~masks[modal_idx] if masks[modal_idx] is not None else None) + enc_outputs.append(enc_out) + enc_inputs = enc_outputs + return enc_inputs + + def forward_fusion(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor], bottleneck_tokens: torch.Tensor, valid_modalities: List[torch.Tensor]) -> List[torch.Tensor]: + # valid_modalities: [B] list of boolean/float tensors indicating if modality is present + batch_size = enc_inputs[0].size(0) + + for modality_encoders in self.fusion_stacks: + enc_outputs = [] + bottleneck_tokens_modality_sum = torch.zeros_like(bottleneck_tokens) + sum_of_modalities = torch.zeros(batch_size, 1, 1, device=bottleneck_tokens.device) + + for idx, enc_layer in enumerate(modality_encoders): + # Concatenate bottleneck tokens with modality tokens + # bottleneck_tokens: [B, num_bottlenecks, d_model] + # enc_inputs[idx]: [B, seq_len, d_model] + fused_input = torch.cat([bottleneck_tokens, enc_inputs[idx]], dim=1) + + # Padding mask for bottleneck tokens is always False (i.e. valid) + # [B, num_bottlenecks] of False + b_mask = torch.zeros(batch_size, self.n_bottlenecks, dtype=torch.bool, device=fused_input.device) + + # Modality padding mask + m_mask = ~masks[idx] if masks[idx] is not None else torch.zeros(batch_size, enc_inputs[idx].size(1), dtype=torch.bool, device=fused_input.device) + + combined_mask = torch.cat([b_mask, m_mask], dim=1) + + # Pass through the layer + enc_out = enc_layer(fused_input, src_key_padding_mask=combined_mask) + + # The output consists of processed bottleneck tokens and modality tokens + # [B, num_bottlenecks, d_model] and [B, seq_len, d_model] + bottleneck_hidden_tokens = enc_out[:, :self.n_bottlenecks, :] + modality_hidden_tokens = enc_out[:, self.n_bottlenecks:, :] + enc_outputs.append(modality_hidden_tokens) + + # Average updated bottlenecks from valid modalities + modality_is_valid = valid_modalities[idx].view(batch_size, 1, 1) + bottleneck_tokens_modality_sum += bottleneck_hidden_tokens * modality_is_valid + sum_of_modalities += modality_is_valid + + # Prevent division by zero if all modalities are missing + # If sum_of_modalities is 0, just pass zeros (or keep previous bottleneck_tokens) + # sum_of_modalities = torch.clamp(sum_of_modalities, min=1.0) + avg_divisor = sum_of_modalities.clone() + avg_divisor[avg_divisor == 0] = 1.0 + + bottleneck_tokens = bottleneck_tokens_modality_sum / avg_divisor + enc_inputs = enc_outputs + + return enc_inputs + + def forward(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor]) -> List[torch.Tensor]: + batch_size = enc_inputs[0].size(0) + + # Determine if a modality is valid for each instance in the batch + # A modality is valid if it has at least one True in its mask + valid_modalities = [] + for mask, inp in zip(masks, enc_inputs): + if mask is not None: + # [B] - True if there's any valid token (1/True) + valid = mask.any(dim=1).float() + else: + valid = torch.ones(batch_size, device=inp.device) + valid_modalities.append(valid) + + bottleneck_tokens = self.bottlenecks.expand(batch_size, -1, -1) + + enc_inputs = self.forward_prefusion(enc_inputs, masks) + enc_inputs = self.forward_fusion(enc_inputs, masks, bottleneck_tokens, valid_modalities) + + return enc_inputs + + +class BottleneckTransformer(BaseModel): + """Bottleneck Transformer model for PyHealth datasets. + + Per-field mode: each feature stream is embedded with :class:`EmbeddingModel`, + prefixed with a learnable per-modality ``[CLS]`` token, processed by + independent prefusion layers, then fused via shared bottleneck tokens. + The per-modality ``[CLS]`` embeddings are averaged and fed to the + classification head. + + Unified mode (``unified_embedding`` supplied): all temporal fields are + jointly embedded and time-sorted by + :class:`~pyhealth.models.embedding.unified.UnifiedMultimodalEmbeddingModel` + into a single interleaved sequence. A single ``[CLS]`` token is prepended + and the encoder runs with ``n_modality=1``, so the bottleneck tokens attend + over the full cross-modal timeline. + + Args: + dataset (SampleDataset): dataset providing processed inputs. + embedding_dim (int): shared embedding dimension. + bottlenecks_n (int): number of shared bottleneck tokens. + fusion_startidx (int): layer index at which bottleneck fusion starts. + Must satisfy ``0 <= fusion_startidx <= num_layers``. + num_layers (int): total transformer layers (prefusion + fusion). + heads (int): number of attention heads per transformer block. + dropout (float): dropout rate inside transformer blocks. + unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when + provided, switches to unified mode. + + Examples: + >>> from pyhealth.datasets import create_sample_dataset, get_dataloader + >>> samples = [ + ... { + ... "patient_id": "patient-0", + ... "visit_id": "visit-0", + ... "conditions": ["A", "B", "C"], + ... "procedures": ["X", "Y"], + ... "label": 1, + ... }, + ... { + ... "patient_id": "patient-1", + ... "visit_id": "visit-0", + ... "conditions": ["D"], + ... "procedures": ["Z", "Y"], + ... "label": 0, + ... }, + ... ] + >>> input_schema = {"conditions": "sequence", "procedures": "sequence"} + >>> output_schema = {"label": "binary"} + >>> dataset = create_sample_dataset( + ... samples, + ... input_schema, + ... output_schema, + ... dataset_name="demo", + ... ) + >>> model = BottleneckTransformer(dataset=dataset, num_layers=3, fusion_startidx=1, bottlenecks_n=4) + >>> loader = get_dataloader(dataset, batch_size=2, shuffle=True) + >>> batch = next(iter(loader)) + >>> output = model(**batch) + >>> sorted(output.keys()) + ['logit', 'loss', 'y_prob', 'y_true'] + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: int = 128, + bottlenecks_n: int = 4, + fusion_startidx: int = 1, + num_layers: int = 3, + heads: int = 4, + dropout: float = 0.5, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, + ): + super().__init__(dataset=dataset) + self.embedding_dim = embedding_dim + self.bottlenecks_n = bottlenecks_n + self.fusion_startidx = fusion_startidx + self.num_layers = num_layers + self.heads = heads + self.dropout = dropout + self._use_unified = unified_embedding is not None + + assert ( + len(self.label_keys) == 1 + ), "Only one label key is supported if BottleneckTransformer is initialized" + self.label_key = self.label_keys[0] + self.mode = self.dataset.output_schema[self.label_key] + + assert 0 <= fusion_startidx <= num_layers, ( + f"fusion_startidx must be in [0, num_layers], got {fusion_startidx}" + ) + + output_size = self.get_output_size() + + if self._use_unified: + self.embedding_model = unified_embedding + # Single CLS token for the unified interleaved sequence + self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim) * 0.02) + self.encoder = MultimodalBottleneckTransformerEncoder( + n_modality=1, + bottlenecks_n=bottlenecks_n, + fusion_startidx=fusion_startidx, + n_layers=num_layers, + n_head=heads, + d_model=embedding_dim, + d_ff=embedding_dim * 4, + dropout=dropout, + ) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.n_modality = len(self.feature_keys) + # Per-modality CLS tokens + self.cls_token_per_modality = nn.ParameterList([ + nn.Parameter(torch.randn(1, 1, embedding_dim) * 0.02) + for _ in range(self.n_modality) + ]) + self.encoder = MultimodalBottleneckTransformerEncoder( + n_modality=self.n_modality, + bottlenecks_n=bottlenecks_n, + fusion_startidx=fusion_startidx, + n_layers=num_layers, + n_head=heads, + d_model=embedding_dim, + d_ff=embedding_dim * 4, + dropout=dropout, + ) + + # fc input is embedding_dim in both modes (CLS token, not concat) + self.fc = nn.Linear(embedding_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Extract value/time/mask tensors for UnifiedMultimodalEmbeddingModel.""" + 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: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified( + self, + **kwargs: Union[torch.Tensor, Tuple[torch.Tensor, ...]], + ) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Calls UnifiedMultimodalEmbeddingModel to produce a single time-sorted + sequence, prepends a CLS token, encodes with the bottleneck encoder + (n_modality=1), and classifies from the CLS output. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + sequence = out["sequence"] # (B, S, E) + event_mask = out["mask"].bool() # (B, S) + + # Prepend CLS token + batch_size = sequence.size(0) + cls = self.cls_token.expand(batch_size, -1, -1) + sequence = torch.cat([cls, sequence], dim=1) + cls_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=sequence.device) + event_mask = torch.cat([cls_mask, event_mask], dim=1) + + enc_outputs = self.encoder([sequence], [event_mask]) + patient_emb = enc_outputs[0][:, 0, :] # CLS token output + + logits = self.fc(patient_emb) + 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) + results["loss"] = self.get_loss_function()(logits, y_true) + results["y_true"] = y_true + return results + + @staticmethod + def _pool_embedding(x: torch.Tensor) -> torch.Tensor: + 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: + 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( + self, + **kwargs: Union[torch.Tensor, Tuple[torch.Tensor, ...]], + ) -> Dict[str, torch.Tensor]: + """Forward propagation. + + In unified mode dispatches to :meth:`_forward_unified`. Otherwise runs + per-field embedding + bottleneck fusion. + + Args: + **kwargs: keyword arguments for the model. + + Returns: + A dictionary with the following keys: + loss: a scalar tensor representing the final loss. + y_prob: a tensor of predicted probabilities. + y_true: a tensor representing the true labels. + logit: the raw logits before activation. + """ + if self._use_unified: + return self._forward_unified(**kwargs) + + enc_inputs = [] + masks = [] + + for idx, feature_key in enumerate(self.feature_keys): + feature = kwargs[feature_key] + + if isinstance(feature, torch.Tensor): + feature = (feature,) + + schema = self.dataset.input_processors[feature_key].schema() + + value = feature[schema.index("value")] if "value" in schema else None + mask = feature[schema.index("mask")] if "mask" in schema else None + + if len(feature) == len(schema) + 1 and mask is None: + mask = feature[-1] + + if value is None: + raise ValueError( + 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] + else: + value = self.embedding_model({feature_key: value})[feature_key] + + value = self._pool_embedding(value) + + if mask is not None: + mask = mask.bool() + if mask.dim() == value.dim(): + mask = mask.any(dim=-1) + else: + mask = self._mask_from_embeddings(value) + + # Prepend Modality CLS token + batch_size = value.size(0) + cls_token = self.cls_token_per_modality[idx].expand(batch_size, -1, -1) + value = torch.cat([cls_token, value], dim=1) + + # Update mask for CLS token (always valid) + cls_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=value.device) + mask = torch.cat([cls_mask, mask], dim=1) + + enc_inputs.append(value) + masks.append(mask) + + # Pass through Bottleneck Transformer Encoder + enc_outputs = self.encoder(enc_inputs, masks) + + # Extract CLS tokens + cls_tokens = [out[:, 0, :].unsqueeze(1) for out in enc_outputs] + cls_tokens = torch.cat(cls_tokens, dim=1) # [B, n_modality, embedding_dim] + + # Average CLS tokens across valid modalities + b_size = cls_tokens.size(0) + valid_modalities = [] + for mask in masks: + # We check if there's any valid token aside from the CLS token (index 0) + if mask.size(1) > 1: + valid = mask[:, 1:].any(dim=1).float() + else: + valid = mask[:, 0].float() # fallback + valid_modalities.append(valid.view(b_size, 1, 1)) + + valid_modality_tensor = torch.cat(valid_modalities, dim=1) # [B, n_modality, 1] + + # Apply valid mask + masked_cls = cls_tokens * valid_modality_tensor + sum_valid = valid_modality_tensor.sum(dim=1) # [B, 1] + + # Avoid division by zero + sum_valid[sum_valid == 0] = 1.0 + patient_emb = masked_cls.sum(dim=1) / sum_valid # [B, embedding_dim] + + logits = self.fc(patient_emb) + y_prob = self.prepare_y_prob(logits) + + results = { + "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 + +if __name__ == "__main__": + from pyhealth.datasets import create_sample_dataset, get_dataloader + + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "conditions": ["A", "B", "C"], + "procedures": ["X", "Y"], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-0", + "conditions": ["D"], + "procedures": ["Z", "Y"], + "label": 0, + }, + ] + + input_schema = { + "conditions": "sequence", + "procedures": "sequence", + } + output_schema = {"label": "binary"} + + dataset = create_sample_dataset( + samples=samples, + input_schema=input_schema, + output_schema=output_schema, + dataset_name="test", + ) + + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + + model = BottleneckTransformer( + dataset=dataset, + embedding_dim=64, + bottlenecks_n=2, + fusion_startidx=1, + num_layers=3, + heads=2 + ) + + data_batch = next(iter(train_loader)) + + result = model(**data_batch) + print(result) + + result["loss"].backward() + print("Test completed successfully.") \ No newline at end of file diff --git a/pyhealth/models/ehrmamba.py b/pyhealth/models/ehrmamba.py index e24c5595f..aa5daedfa 100644 --- a/pyhealth/models/ehrmamba.py +++ b/pyhealth/models/ehrmamba.py @@ -4,8 +4,10 @@ from torch import nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel from pyhealth.models.utils import get_last_visit from pyhealth.processors import ( MultiHotProcessor, @@ -111,6 +113,11 @@ class EHRMamba(BaseModel): Electronic Health Records (arxiv 2405.14567). Uses Mamba (SSM) for linear complexity in sequence length; supports long EHR sequences. + When ``unified_embedding`` is supplied the model switches to **unified + mode**: all temporal fields are jointly embedded and time-sorted by + :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single* + stack of :class:`MambaBlock` layers rather than one stack per field. + Args: dataset: SampleDataset for token/embedding setup. embedding_dim: Embedding and hidden dimension. Default 128. @@ -118,6 +125,8 @@ class EHRMamba(BaseModel): state_size: SSM state size per channel. Default 16. conv_kernel: Causal conv kernel size in block. Default 4. dropout: Dropout before classification head. Default 0.1. + unified_embedding: Optional pre-built UnifiedMultimodalEmbeddingModel. + When provided, enables unified multi-modal mode. """ def __init__( @@ -128,6 +137,7 @@ def __init__( state_size: int = 16, conv_kernel: int = 4, dropout: float = 0.1, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, ): super().__init__(dataset=dataset) self.embedding_dim = embedding_dim @@ -135,19 +145,18 @@ def __init__( self.state_size = state_size self.conv_kernel = conv_kernel self.dropout_rate = dropout + self._use_unified = unified_embedding is not None assert len(self.label_keys) == 1, "EHRMamba supports single label key only" self.label_key = self.label_keys[0] self.mode = self.dataset.output_schema[self.label_key] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) - self.feature_processors = { - k: self.dataset.input_processors[k] for k in self.feature_keys - } + output_size = self.get_output_size() + self.dropout = nn.Dropout(dropout) - self.blocks = nn.ModuleDict() - for feature_key in self.feature_keys: - self.blocks[feature_key] = nn.ModuleList( + if self._use_unified: + self.embedding_model = unified_embedding + self._unified_blocks = nn.ModuleList( [ MambaBlock( d_model=embedding_dim, @@ -157,10 +166,79 @@ def __init__( for _ in range(num_layers) ] ) - - output_size = self.get_output_size() - self.dropout = nn.Dropout(dropout) - self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + self.fc = nn.Linear(embedding_dim, output_size) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.feature_processors = { + k: self.dataset.input_processors[k] for k in self.feature_keys + } + self.blocks = nn.ModuleDict() + for feature_key in self.feature_keys: + self.blocks[feature_key] = nn.ModuleList( + [ + MambaBlock( + d_model=embedding_dim, + state_size=state_size, + conv_kernel=conv_kernel, + ) + for _ in range(num_layers) + ] + ) + self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build the inputs dict required by UnifiedMultimodalEmbeddingModel.""" + 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: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified(self, **kwargs) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Calls UnifiedMultimodalEmbeddingModel to produce a single + temporally-sorted event sequence, then encodes it with one shared + MambaBlock stack and pools to the last valid event. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + x = out["sequence"] # (B, S_total, E) + mask = out["mask"].bool() # (B, S_total) + + for blk in self._unified_blocks: + x = blk(x) + + last_h = get_last_visit(x, mask) + logits = self.fc(self.dropout(last_h)) + y_prob = self.prepare_y_prob(logits) + results: Dict[str, torch.Tensor] = { + "loss": torch.tensor(0.0), # placeholder, overwritten below + "y_prob": y_prob, + "logit": logits, + } + if self.label_key in kwargs: + y_true = kwargs[self.label_key].to(self.device) + results["loss"] = self.get_loss_function()(logits, y_true) + results["y_true"] = y_true + if kwargs.get("embed", False): + results["embed"] = last_h + return results @staticmethod def _split_temporal(feature: Any) -> Tuple[Optional[torch.Tensor], Any]: @@ -211,6 +289,9 @@ def _pool_embedding(x: torch.Tensor) -> torch.Tensor: return x def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + if self._use_unified: + return self._forward_unified(**kwargs) + patient_emb = [] embedding_inputs: Dict[str, torch.Tensor] = {} masks: Dict[str, torch.Tensor] = {} @@ -261,4 +342,4 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: batch = next(iter(loader)) out = model(**batch) print("keys:", sorted(out.keys())) - out["loss"].backward() + out["loss"].backward() \ No newline at end of file diff --git a/pyhealth/models/embedding/__init__.py b/pyhealth/models/embedding/__init__.py new file mode 100644 index 000000000..2bc66d73b --- /dev/null +++ b/pyhealth/models/embedding/__init__.py @@ -0,0 +1,33 @@ +"""Embedding models for PyHealth multimodal pipelines. + +All embedding models share the :class:`BaseEmbeddingModel` interface: +they expose an ``embedding_dim`` property and a ``forward`` method that +transforms processor output tensors into dense vector embeddings. + +Available models: + +- :class:`EmbeddingModel` — generic encoder for codes, sequences, timeseries +- :class:`VisionEmbeddingModel` — ViT-style patch encoder for medical images (Josh) +- :class:`UnifiedMultimodalEmbeddingModel` — temporally-aligned multi-modal encoder + +Helper utilities: + +- :class:`SinusoidalTimeEmbedding` — continuous time positional encoding +- :func:`init_embedding_with_pretrained` — load GloVe-style pretrained vectors +""" + +from .base import BaseEmbeddingModel +from .vanilla import EmbeddingModel, init_embedding_with_pretrained +from .vision import VisionEmbeddingModel, PatchEmbedding, Permute +from .unified import UnifiedMultimodalEmbeddingModel, SinusoidalTimeEmbedding + +__all__ = [ + "BaseEmbeddingModel", + "EmbeddingModel", + "VisionEmbeddingModel", + "PatchEmbedding", + "Permute", + "UnifiedMultimodalEmbeddingModel", + "SinusoidalTimeEmbedding", + "init_embedding_with_pretrained", +] \ No newline at end of file diff --git a/pyhealth/models/embedding/base.py b/pyhealth/models/embedding/base.py new file mode 100644 index 000000000..5ff1ab2a6 --- /dev/null +++ b/pyhealth/models/embedding/base.py @@ -0,0 +1,52 @@ +from abc import ABC, abstractmethod + + +class BaseEmbeddingModel(ABC): + """Abstract base class for all embedding models in PyHealth. + + All embedding models share a common contract: + + - They expose an ``embedding_dim`` property indicating the output vector dimension. + - Their ``forward`` method accepts processor output tensors and returns + vector embeddings. + + Concrete subclasses: + + - :class:`EmbeddingModel` – generic encoder for codes, sequences, timeseries + - :class:`VisionEmbeddingModel` – patch-based encoder for medical images (Josh) + - :class:`TextEmbeddingModel` – BERT-based encoder for clinical text (Rian) + - :class:`UnifiedMultimodalEmbeddingModel` – temporally-aligned multi-modal encoder + """ + + @property + @abstractmethod + def embedding_dim(self) -> int: + """Output embedding dimension shared across all modalities.""" + ... + + @abstractmethod + def forward(self, *args, **kwargs): + """Transform processor outputs into embeddings. + + Subclass return types + --------------------- + EmbeddingModel + ``Dict[str, Tensor]`` mapping each field name to its embedded + tensor. When ``output_mask=True`` is passed, returns a + ``(Dict[str, Tensor], Dict[str, Tensor])`` tuple of + (embeddings, masks). + + VisionEmbeddingModel + ``Tensor`` of shape ``[batch, embedding_dim]``. + + TextEmbeddingModel + ``(Tensor, BoolTensor)`` of shapes ``([B, T, E], [B, T])`` + when ``return_mask=True`` (default), or a plain + ``Tensor [B, T, E]`` when ``return_mask=False``. + + UnifiedMultimodalEmbeddingModel + ``Dict[str, Tensor]`` with keys ``"sequence"`` ``[B, S, E]``, + ``"mask"`` ``[B, S]``, ``"time"`` ``[B, S]``, and + ``"type_ids"`` ``[B, S]``. + """ + ... diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py new file mode 100644 index 000000000..bf8050088 --- /dev/null +++ b/pyhealth/models/embedding/unified.py @@ -0,0 +1,839 @@ +"""UnifiedMultimodalEmbeddingModel, temporally aligned multimodal embedding. + +Takes K temporal features ( dict outputs from ``TemporalFeatureProcessor`` +subclasses ), embeds each event with a modality-specific encoder, then +interleaves all events on a shared timeline by sorting on timestamp and adding +sinusoidal time embeddings + learned modality-type embeddings. + +Output shape: ``(B, S_total, E')``, a single sequence of events usable by +any downstream sequence model (Transformer, Mamba, RNN, …). + +IMAGE encoding delegates to :class:`PatchEmbedding` from +:mod:`pyhealth.models.embedding.vision` (Josh's model), pooling patch tokens +to a single per-image vector via global mean pooling. + +TEXT encoding uses a pretrained BERT tokenizer model directly, extracting the +[CLS] token per note, the same BERT-based approach as +:class:`TextEmbeddingModel` (Rian's model). + +Unimodal model reuse via ``field_embeddings``:: + + vision_model = VisionEmbeddingModel(dataset, embedding_dim=128) + text_model = TextEmbeddingModel(embedding_dim=128) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=128, + field_embeddings={ + "chest_xray": vision_model, # reuses trained backbone + "notes": text_model, # reuses BERT + projection + }, + ) + +Quickstart:: + + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.datasets.collate import collate_temporal + model = UnifiedMultimodalEmbeddingModel(dataset, embedding_dim=128) + # inside forward: + # inputs = {field: {"value": Tensor, "time": Tensor, ...}, ...} + out = model(inputs) + # out["sequence"]: (B, S_total, 128) + # out["mask"]: (B, S_total) , 1 = real event, 0 = padding + # out["time"]: (B, S_total) , hours from first event +""" + +from __future__ import annotations + +import math +import warnings +from contextlib import nullcontext +from typing import Any, Optional + +import torch +import torch.nn.functional as F +import torch.nn as nn + +from ...processors.base_processor import ModalityType, TemporalFeatureProcessor +from .base import BaseEmbeddingModel +from .vision import PatchEmbedding + + +# ── Helpers ─────────────────────────────────────────────────────────────────── + + +class SinusoidalTimeEmbedding(nn.Module): + """Multi-scale sinusoidal embedding for times in hours. + + Wavelengths are spaced geometrically from ``min_hours`` (within-stay + resolution) to ``max_hours`` (longitudinal span). The previous encoding + mapped ``t / 720 * 2π``, so every frequency wrapped every 30 days: a later + stay at +9606h produced the same embedding as +6h even after the task + stopped resetting the clock. + + Args: + dim: Output embedding dimension (must be even). + max_hours: Longest wavelength in hours. Default 87600 (10 years). + min_hours: Shortest wavelength in hours. Default 1.0. + + Shape: + Input: ``(*, )`` float tensor of times in hours + Output: ``(*, dim)`` + """ + + def __init__( + self, + dim: int, + max_hours: float = 87600.0, + min_hours: float = 1.0, + ): + super().__init__() + assert dim % 2 == 0, f"dim must be even, got {dim}" + if min_hours <= 0 or max_hours <= min_hours: + raise ValueError( + f"need 0 < min_hours < max_hours, got min={min_hours}, max={max_hours}" + ) + self.dim = dim + self.max_hours = float(max_hours) + self.min_hours = float(min_hours) + half = dim // 2 + periods = torch.exp( + torch.linspace( + math.log(self.min_hours), + math.log(self.max_hours), + half, + dtype=torch.float32, + ) + ) + self.register_buffer("freqs", 2 * math.pi / periods) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + """:param t: ``(...,)`` float, times in hours.""" + args = t.unsqueeze(-1).to(dtype=self.freqs.dtype) * self.freqs + return torch.cat([args.sin(), args.cos()], dim=-1) + + +class _MeanPool(nn.Module): + """Pool a sequence of patch embeddings to a single vector via global mean.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # (B, num_patches, E) -> (B, E) + return x.mean(dim=1) + + +# ── Main model ─────────────────────────────────────────────────────────────── + + +class UnifiedMultimodalEmbeddingModel(nn.Module, BaseEmbeddingModel): + """Embed heterogeneous temporal features into a single aligned sequence. + + **All** input processors must be ``TemporalFeatureProcessor`` subclasses. + Non-temporal processors (e.g. ``SequenceProcessor``, ``MultiHotProcessor``) + are rejected with a clear error, use :class:`EmbeddingModel` for those fields. + + Modality routing: + + - **CODE**: ``nn.Embedding`` lookup. + - **TEXT**: Pretrained BERT (same approach as :class:`TextEmbeddingModel`), + CLS token extracted per note. + - **IMAGE**: :class:`PatchEmbedding` (from :class:`VisionEmbeddingModel`) + followed by global mean pooling to produce one vector per image event. + - **NUMERIC / SIGNAL**: ``nn.Linear`` projection. + + Unimodal model reuse: + + Pass pre-built :class:`EmbeddingModel`, :class:`VisionEmbeddingModel`, or + :class:`TextEmbeddingModel` instances via ``field_embeddings`` to reuse + their trained encoder weights instead of building new ones from scratch. + The core encoder module is extracted from each pre-built model: + + - ``EmbeddingModel`` → ``embedding_layers[field_name]`` (``nn.Embedding`` / + ``nn.Linear``) + - ``VisionEmbeddingModel`` → ``embedding_layers[field_name]`` backbone + + global mean pooling + - ``TextEmbeddingModel`` → ``transformer`` (BERT) + ``fc`` (projection) + + Algorithm + --------- + For each temporal field: + + 1. Route ``inputs[field]["value"]`` through a modality-specific encoder → + ``(B, N_i, E')`` per-event embeddings. + 2. Retrieve ``inputs[field]["time"]`` → ``(B, N_i)`` timestamps (hours). + 3. (Optional) Retrieve ``inputs[field]["mask"]`` → ``(B, N_i, L)`` or + ``(B, N_i)`` attention mask; reduced to event-level ``(B, N_i)`` if + token-level. + + Then: + + 4. Concatenate across all fields → ``(B, S_total, E')``. + 5. Sort events along dim=1 by timestamp (ascending). + 6. Add ``SinusoidalTimeEmbedding(time)`` + ``type_embedding(modality_idx)``. + 7. Return ``{"sequence", "time", "mask", "type_ids"}``. + + Args: + processors: ``dict[field_name, TemporalFeatureProcessor]``, the + processors for each temporal field in the dataset. Pass + ``dataset.input_processors`` directly. + embedding_dim: Shared embedding dimension ``E'``. + time_embedding: ``"sinusoidal"`` (default) or ``"learned"``. + max_time_hours: Longest wavelength of the time embedding, in hours. + Defaults to 87600 (10 years). Shortest wavelength is 1 hour. + image_size: Image size (H=W) assumed for IMAGE fields when using + PatchEmbedding. Defaults to 224. + image_channels: Number of input channels for IMAGE fields. Defaults to 3. + patch_size: Patch size for IMAGE PatchEmbedding encoder. Defaults to 16. + image_pool: Pooling strategy applied to IMAGE patch tokens to produce + one vector per image event. Only ``"mean"`` (global mean pooling) + is currently implemented. Defaults to ``"mean"``. + field_embeddings: Optional mapping of field names to pre-built unimodal + embedding models. Supported types: + + - :class:`EmbeddingModel` (codes / numeric) — extracts + ``embedding_layers[field_name]``. + - :class:`VisionEmbeddingModel` — extracts the backbone layer and + wraps it with global mean pooling. + - :class:`TextEmbeddingModel` — reuses ``transformer`` and ``fc`` + for BERT-based CLS extraction. + + Fields not present in this dict fall back to the default + internally-built encoders. + + Example:: + + model = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=128, + ) + # inputs: {field: {"value": Tensor, "time": Tensor, "mask": Tensor}} + out = model(inputs) + seq = out["sequence"] # (B, S_total, 128) + mask = out["mask"] # (B, S_total) float, 1=valid 0=pad + + # With pre-built unimodal models: + vision = VisionEmbeddingModel(dataset, embedding_dim=128) + model = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=128, + field_embeddings={"chest_xray": vision}, + ) + """ + + def __init__( + self, + processors: dict[str, Any], + embedding_dim: int = 128, + time_embedding: str = "sinusoidal", + max_time_hours: float = 87600.0, + image_size: int = 224, + image_channels: int = 3, + patch_size: int = 16, + image_pool: str = "mean", + field_embeddings: Optional[dict[str, Any]] = None, + freeze_text_encoder: bool = False, + normalize_content: bool = True, + cache_frozen_text: bool = True, + max_frozen_text_cache: int = 200_000, + numeric_standardizers: Optional[dict[str, Any]] = None, + ): + super().__init__() + if image_pool != "mean": + raise NotImplementedError( + f"Only image_pool='mean' is implemented, got {image_pool!r}." + ) + self._embedding_dim = embedding_dim + self._freeze_text_encoder = freeze_text_encoder + self._frozen_text_fields: set[str] = set() + self.cache_frozen_text = cache_frozen_text + self.max_frozen_text_cache = max_frozen_text_cache + self._frozen_text_cache: dict[str, dict[int, torch.Tensor]] = {} + self.image_pool = image_pool + self.normalize_content = normalize_content + # Statistics live in buffers, so they travel in state_dict. A checkpoint + # therefore applies at inference the same transform it trained under. + self.numeric_standardizers = nn.ModuleDict(numeric_standardizers or {}) + _field_embeddings = field_embeddings or {} + + 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): + raise TypeError( + f"UnifiedMultimodalEmbeddingModel requires every input processor " + f"to be a TemporalFeatureProcessor subclass, but '{field_name}' " + f"uses {type(processor).__name__}. For non-temporal fields use " + f"EmbeddingModel." + ) + + m = processor.modality() + self.modality_types[field_name] = m + pre_built = _field_embeddings.get(field_name) + + if m == ModalityType.CODE: + self.encoders[field_name] = self._build_code_encoder( + field_name, processor, pre_built, embedding_dim + ) + + elif m == ModalityType.TEXT: + self._build_text_encoder( + field_name, processor, pre_built, embedding_dim, + freeze=freeze_text_encoder, + ) + + elif m == ModalityType.IMAGE: + self.encoders[field_name] = self._build_image_encoder( + field_name, + processor, + pre_built, + embedding_dim, + image_size, + image_channels, + patch_size, + image_pool, + ) + + elif m in (ModalityType.NUMERIC, ModalityType.SIGNAL): + self.encoders[field_name] = self._build_numeric_encoder( + field_name, processor, pre_built, embedding_dim + ) + + else: + raise NotImplementedError( + f"No encoder implemented for modality {m!r} (field '{field_name}')." + ) + + # Shared type embedding, one vector per unique modality in this dataset + unique_modalities = sorted(set(self.modality_types.values())) + self._modality_to_idx: dict[ModalityType, int] = { + mod: i for i, mod in enumerate(unique_modalities) + } + self.type_embedding = nn.Embedding(len(unique_modalities), embedding_dim) + self._warned_nested_code_flatten = False + + # Time embedding + if time_embedding == "sinusoidal": + self.time_embed = SinusoidalTimeEmbedding(embedding_dim, max_time_hours) + else: + raise NotImplementedError( + "Only 'sinusoidal' time embedding is implemented." + ) + + # ── Encoder builders ────────────────────────────────────────────────────── + + def _build_code_encoder( + self, + field_name: str, + processor: TemporalFeatureProcessor, + pre_built: Any, + embedding_dim: int, + ) -> nn.Module: + """Build CODE encoder: nn.Embedding, optionally from a pre-built EmbeddingModel.""" + if ( + pre_built is not None + and hasattr(pre_built, "embedding_layers") + and field_name in pre_built.embedding_layers + ): + layer = pre_built.embedding_layers[field_name] + pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) + if pre_dim != embedding_dim: + return nn.Sequential(layer, nn.Linear(pre_dim, embedding_dim)) + return layer + + vocab_size = processor.value_dim() + return nn.Embedding(vocab_size, embedding_dim, padding_idx=0) + + def _build_text_encoder( + self, + field_name: str, + processor: TemporalFeatureProcessor, + pre_built: Any, + embedding_dim: int, + 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") + and hasattr(pre_built, "fc") + ): + self.encoders[field_name] = pre_built.transformer + if freeze: + for p in pre_built.transformer.parameters(): + p.requires_grad = False + self._frozen_text_fields.add(field_name) + pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) + _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._frozen_text_fields.add(field_name) + self.encoders[field_name] = bert + hidden = bert.config.hidden_size + if hidden != embedding_dim: + self.projections[field_name] = nn.Linear(hidden, embedding_dim) + else: + raise ValueError( + f"TEXT processor '{field_name}' must either supply a pre-built " + f"TextEmbeddingModel via field_embeddings or use a tokenizer " + f"(set tokenizer_model=...) to be used with " + f"UnifiedMultimodalEmbeddingModel." + ) + + def _build_image_encoder( + self, + field_name: str, + processor: TemporalFeatureProcessor, + pre_built: Any, + embedding_dim: int, + image_size: int, + image_channels: int, + patch_size: int, + image_pool: str, + ) -> nn.Module: + """Build IMAGE encoder: backbone + pool, optionally from VisionEmbeddingModel.""" + pool_layers: dict[str, nn.Module] = {"mean": _MeanPool()} + pool_layer = pool_layers[image_pool] + + if ( + pre_built is not None + and hasattr(pre_built, "embedding_layers") + and field_name in pre_built.embedding_layers + ): + backbone = pre_built.embedding_layers[field_name] + pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) + if pre_dim != embedding_dim: + return nn.Sequential( + backbone, pool_layer, nn.Linear(pre_dim, embedding_dim) + ) + return nn.Sequential(backbone, pool_layer) + + _image_size = getattr(processor, "image_size", image_size) + _in_channels = getattr(processor, "in_channels", image_channels) + return nn.Sequential( + PatchEmbedding(_image_size, patch_size, _in_channels, embedding_dim), + pool_layer, + ) + + def _build_numeric_encoder( + self, + field_name: str, + processor: TemporalFeatureProcessor, + pre_built: Any, + embedding_dim: int, + ) -> nn.Module: + """Build NUMERIC/SIGNAL encoder: nn.Linear, optionally from EmbeddingModel.""" + if ( + pre_built is not None + and hasattr(pre_built, "embedding_layers") + and field_name in pre_built.embedding_layers + ): + layer = pre_built.embedding_layers[field_name] + pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) + if pre_dim != embedding_dim: + return nn.Sequential(layer, nn.Linear(pre_dim, embedding_dim)) + return layer + + in_features = processor.value_dim() + return nn.Linear(in_features, embedding_dim) + + @property + def embedding_dim(self) -> int: + return self._embedding_dim + + def _encode_text_cls( + self, + field_name: str, + encoder: nn.Module, + flat_ids: torch.Tensor, + flat_mask: Optional[torch.Tensor], + ) -> torch.Tensor: + """Return the ``[CLS]`` vector for each row, from a cache when possible. + + A frozen encoder gives the same output for the same tokens, so a run of + 50 epochs would otherwise repeat the identical 110M-parameter forward + pass 50 times. + + The cache has three conditions. It is used only for a field in + ``_frozen_text_fields``, so a trainable encoder can never read it. The + key is the token identifiers under the attention mask, so a change of + tokenizer or truncation budget gives a different key. The cache has a + maximum size, and it recalculates a row when the cache is full. + + Key on the REAL tokens only. The collator pads each row to the widest + note in its batch, and batch composition changes every epoch because + the loader shuffles, so a key over the padded row gives the same note + a different key each epoch and the cache never hits. Measured on the + full-scale notes run: epoch time did not fall after epoch 1 + (3458s, 3936s, 4048s, 3835s) because every lookup missed. + """ + if flat_ids.shape[0] == 0: + hidden = encoder.config.hidden_size + return flat_ids.new_zeros( + (0, hidden), dtype=next(encoder.parameters()).dtype + ) + + if not (self.cache_frozen_text and field_name in self._frozen_text_fields): + ctx = torch.no_grad() if field_name in self._frozen_text_fields else nullcontext() + with ctx: + out = encoder(input_ids=flat_ids, attention_mask=flat_mask) + return out.last_hidden_state[:, 0, :] + + cache = self._frozen_text_cache.setdefault(field_name, {}) + ids_cpu = flat_ids.detach().cpu() + mask_cpu = ( + flat_mask.detach().cpu().to(torch.int8) + if flat_mask is not None + else torch.ones_like(ids_cpu, dtype=torch.int8) + ) + keys = [ + hash(tuple(i[m.bool()].tolist())) if m.any() else hash(tuple(i.tolist())) + for i, m in zip(ids_cpu, mask_cpu) + ] + + first_row_of_key: dict[int, int] = {} + for k, key in enumerate(keys): + if key not in cache and key not in first_row_of_key: + first_row_of_key[key] = k + missing = list(first_row_of_key.values()) + if missing: + index = torch.tensor(missing, device=flat_ids.device) + with torch.no_grad(): + out = encoder( + input_ids=flat_ids.index_select(0, index), + attention_mask=( + flat_mask.index_select(0, index) + if flat_mask is not None + else None + ), + ) + fresh = out.last_hidden_state[:, 0, :].detach() + for slot, row in zip(missing, fresh): + if len(cache) < self.max_frozen_text_cache: + cache[keys[slot]] = row.cpu() + + rows = [] + for k, key in enumerate(keys): + hit = cache.get(key) + if hit is None: + with torch.no_grad(): + out = encoder( + input_ids=flat_ids[k : k + 1], + attention_mask=( + flat_mask[k : k + 1] if flat_mask is not None else None + ), + ) + rows.append(out.last_hidden_state[0, 0, :].detach()) + else: + rows.append(hit.to(flat_ids.device)) + return torch.stack(rows).to(dtype=self.type_embedding.weight.dtype) + + def train(self, mode: bool = True) -> "UnifiedMultimodalEmbeddingModel": + """Keep a frozen text encoder in eval mode. + + ``nn.Module.train()`` would enable dropout inside the encoder. Its + output would then change between passes even though every weight has + ``requires_grad=False``. That makes the cache incorrect, and it also + makes a frozen encoder give a different answer for the same input. + """ + super().train(mode) + for field_name in self._frozen_text_fields: + self.encoders[field_name].eval() + return self + + # ── Forward ─────────────────────────────────────────────────────────────── + + def forward( + self, + inputs: dict[str, dict[str, torch.Tensor]], + ) -> dict[str, torch.Tensor]: + """Encode and temporally align all temporal features. + + Args: + inputs: ``{field_name: {"value": Tensor, "time": Tensor, + "mask": Tensor (optional)}}`` + , one dict per temporal feature, exactly as produced by + ``collate_temporal``. + + Returns: + A dict with keys: + + * ``"sequence"``, ``(B, S_total, E')`` temporally-sorted events + (content + time + type embeddings) + * ``"time"`` , ``(B, S_total)`` timestamps (hours) + * ``"mask"`` , ``(B, S_total)`` 1=real event, 0=padding + * ``"type_ids"``, ``(B, S_total)`` modality index per event + * ``"token_emb"``, ``(B, S_total, E')`` content-only event embedding + (before time/type are added); the target for masked modeling. + """ + all_embeddings: list[torch.Tensor] = [] + all_times: list[torch.Tensor] = [] + all_masks: list[torch.Tensor] = [] + all_types: list[torch.Tensor] = [] + + for field_name, feat_dict in inputs.items(): + if field_name.endswith("_mask") and field_name[: -len("_mask")] in inputs: + # Observation-flag sibling consumed by the standardiser; not a + # modality of its own. Encoding it would duplicate every lab + # timestamp with a 0/1 vector. + continue + value = feat_dict["value"] # (B, N_i, ...) or (B, S, F) + time = feat_dict["time"] # (B, N_i) + # Three different masks meet here and must not be conflated. + # mask token level, from the processor schema; this is the + # attention mask a text encoder needs. + # pad_mask event level, from the collator; which slots are real + # events rather than batch padding. + # {field}_mask a separate FIELD meaning "was this value + # observed", consumed by the standardiser below. + mask = feat_dict.get("mask") + pad_mask = feat_dict.get("pad_mask") + + if time is None: + # Fallback: treat every event as occurring at t=0 + time = torch.zeros(value.shape[:2], device=value.device) + + modality = self.modality_types[field_name] + encoder_key = self._text_canonical.get(field_name, field_name) + encoder = self.encoders[encoder_key] + + # ── Encode ──────────────────────────────────────────────────── + if modality == ModalityType.CODE: + # CODE values may be either: + # - flat indices: (B, S) + # - nested indices: (B, S, C) where C is codes-per-event + # For nested indices, flatten to (B, S*C, E') so code-level + # detail is preserved, and expand time/mask to match. + if value.dim() == 2: + emb = encoder(value) # (B, S, E') + elif value.dim() == 3: + bsz, seq_len, per_event_codes = value.shape + token_emb = encoder(value.long()) # (B, S, C, E') + emb = token_emb.reshape(bsz, seq_len * per_event_codes, -1) + + if not self._warned_nested_code_flatten: + warnings.warn( + ( + "UnifiedMultimodalEmbeddingModel detected " + f"nested CODE input for '{field_name}' with " + f"shape={tuple(value.shape)}. Flattening to " + f"(B, S*C, E) and repeating time along C." + ), + stacklevel=2, + ) + self._warned_nested_code_flatten = True + + if time is not None: + time = ( + time.unsqueeze(-1) + .expand(-1, -1, per_event_codes) + .reshape(bsz, seq_len * per_event_codes) + ) + + if mask is not None: + if mask.dim() == 2: + mask = ( + mask.unsqueeze(-1) + .expand(-1, -1, per_event_codes) + .reshape(bsz, seq_len * per_event_codes) + ) + elif mask.dim() == 3: + mask = mask.reshape(bsz, seq_len * per_event_codes) + else: + raise ValueError( + f"Unsupported CODE value rank for '{field_name}': " + f"shape={tuple(value.shape)}" + ) + + elif modality == ModalityType.TEXT: + # Collate pads note slots to the longest sample in the batch. + # Running BERT on those empty rows is what OOM'd batch-32 + # notes_labs on a 48 GB GPU (~B*N=full pad width, L=512). + b, n, l = value.shape + flat_ids = value.reshape(b * n, l) + flat_attn = mask.reshape(b * n, l) if mask is not None else None + if pad_mask is not None: + valid = pad_mask.reshape(b * n).bool() + elif flat_attn is not None: + valid = flat_attn.any(dim=-1) + else: + valid = torch.ones( + b * n, dtype=torch.bool, device=value.device + ) + hidden = encoder.config.hidden_size + cls_emb = value.new_zeros( + (b * n, hidden), dtype=next(encoder.parameters()).dtype + ) + if valid.any(): + h = self._encode_text_cls( + field_name, + encoder, + flat_ids[valid], + flat_attn[valid] if flat_attn is not None else None, + ) + cls_emb = cls_emb.to(dtype=h.dtype) + cls_emb[valid] = h + if field_name in self.projections: + cls_emb = self.projections[field_name](cls_emb) + emb = cls_emb.view(b, n, -1) # (B, N, E') + + elif modality == ModalityType.IMAGE: + # encoder = Sequential(PatchEmbedding, _MeanPool) → (B*N, E') + b, n, c, h, w = value.shape + flat_imgs = value.reshape(b * n, c, h, w) + if pad_mask is not None: + valid = pad_mask.reshape(b * n).bool() + else: + valid = flat_imgs.reshape(b * n, -1).abs().sum(dim=-1) > 0 + if valid.any(): + img_valid = encoder(flat_imgs[valid]) + img_emb = img_valid.new_zeros( + (b * n, img_valid.shape[-1]) + ) + img_emb[valid] = img_valid + else: + img_emb = value.new_zeros((b * n, self._embedding_dim)) + emb = img_emb.view(b, n, -1) # (B, N, E') + + else: # NUMERIC / SIGNAL + # Standardise BEFORE the projection. The projection mixes the + # features, so a transform after it cannot correct a feature + # whose physical unit gives it 300 times the magnitude of + # another. + standardizer = ( + self.numeric_standardizers[field_name] + if field_name in self.numeric_standardizers + else None + ) + if standardizer is not None: + # Observation flags live in the sibling ``{field}_mask`` + # FIELD, not in this field's dict. Reading the padding mask + # here would tell the standardiser that every real event + # was measured, which is exactly the distinction the + # standardiser exists to preserve. + sibling = inputs.get(f"{field_name}_mask") + obs = sibling["value"] if isinstance(sibling, dict) else None + if obs is None: + raise ValueError( + f"The standardiser for {field_name!r} needs a paired " + f"{field_name}_mask field in the batch." + ) + if obs.shape != value.shape: + raise ValueError( + f"{field_name}_mask has shape {tuple(obs.shape)}, " + f"which does not match {field_name} " + f"{tuple(value.shape)}." + ) + value = standardizer(value, obs.bool()) + emb = encoder(value) # (B, T, E') + + # ── Build event-level validity mask ─────────────────────────── + if pad_mask is not None: + # The collator is authoritative about batch padding. + event_mask = pad_mask.to(emb.device).float() + if event_mask.shape[1] != emb.shape[1]: + # A nested CODE field was flattened to (B, S*C); repeat the + # event flags along the same axis. + repeat = emb.shape[1] // event_mask.shape[1] + event_mask = ( + event_mask.unsqueeze(-1) + .expand(-1, -1, repeat) + .reshape(emb.shape[0], -1) + ) + elif mask is None: + event_mask = torch.ones(emb.shape[:2], device=emb.device) + else: + if mask.dim() > time.dim(): + # token-level (B, N, L) → event-level (B, N) + event_mask = (mask.sum(dim=-1) > 0).float() + else: + event_mask = mask.float() + + # ── Modality type indices ───────────────────────────────────── + type_idx = self._modality_to_idx[modality] + type_ids = torch.full( + emb.shape[:2], type_idx, dtype=torch.long, device=emb.device + ) + + all_embeddings.append(emb) + all_times.append(time) + all_masks.append(event_mask) + all_types.append(type_ids) + + # ── Concatenate across all fields ───────────────────────────────── + cat_emb = torch.cat(all_embeddings, dim=1) # (B, S_total, E') + cat_time = torch.cat(all_times, dim=1) # (B, S_total) + cat_mask = torch.cat(all_masks, dim=1) # (B, S_total) + cat_types = torch.cat(all_types, dim=1) # (B, S_total) + + # ── Sort by time ────────────────────────────────────────────────── + # Padding carries time 0.0, so a plain ascending sort places it BEFORE + # every real event. Three consumers then read it: RNNLayer packs the + # first ``mask.sum()`` steps, ``get_last_visit`` indexes + # ``mask.sum() - 1``, and TransformerLayer takes position 0 as its CLS + # vector. Push invalid slots past every real one to keep the sequence + # left-aligned, which is what all three assume. + # + # Stable, because the key is heavily tied: all padding shares time 0.0 + # and events from one admission share offsets. An unstable sort makes + # event order differ between torch builds and between CPU and CUDA, + # silently changing RNN and Mamba outputs. + sort_key = cat_time.masked_fill(~cat_mask.bool(), float("inf")) + sort_idx = sort_key.argsort(dim=1, stable=True) + cat_emb = cat_emb.gather(1, sort_idx.unsqueeze(-1).expand_as(cat_emb)) + cat_time = cat_time.gather(1, sort_idx) + cat_mask = cat_mask.gather(1, sort_idx) + cat_types = cat_types.gather(1, sort_idx) + + # ── Add time + type embeddings ──────────────────────────────────── + time_emb = self.time_embed(cat_time) # (B, S_total, E') + type_emb = self.type_embedding(cat_types) # (B, S_total, E') + if self.normalize_content: + # Put the content term on the scale of the additive terms. Without + # this the sum is decided by whichever modality has the larger + # magnitude, which is an accident of feature scaling and not a + # modelling decision. Measured at embedding_dim=128: text content + # norm 3.2, raw laboratory content norm 761.4, time and type + # together 13. F.layer_norm without weight or bias adds NO + # parameters, so an existing checkpoint still loads. + cat_emb = F.layer_norm(cat_emb, (cat_emb.shape[-1],)) + final = cat_emb + time_emb + type_emb + # Zero the padded slots so a consumer that ignores the mask, such as a + # mean pool, still cannot pick them up. + final = final * cat_mask.unsqueeze(-1).to(final.dtype) # (B, S_total, E') + + return { + "sequence": final, # (B, S_total, E') + "time": cat_time, # (B, S_total) + "mask": cat_mask, # (B, S_total) + "type_ids": cat_types, # (B, S_total) + # Per-event content embedding BEFORE time/type are added (same sort + # order as ``sequence``). Masked-modeling pretrainers should + # reconstruct THIS rather than ``sequence``: the time/type + # components are largely recoverable from event position, so + # including them in the target dilutes the content signal. + "token_emb": cat_emb, # (B, S_total, E') + } \ No newline at end of file diff --git a/pyhealth/models/embedding.py b/pyhealth/models/embedding/vanilla.py similarity index 93% rename from pyhealth/models/embedding.py rename to pyhealth/models/embedding/vanilla.py index 4232b2788..2a229bfad 100644 --- a/pyhealth/models/embedding.py +++ b/pyhealth/models/embedding/vanilla.py @@ -6,8 +6,8 @@ import torch import torch.nn as nn -from ..datasets import SampleDataset -from ..processors import ( +from ...datasets import SampleDataset +from ...processors import ( MultiHotProcessor, NestedFloatsProcessor, NestedSequenceProcessor, @@ -19,7 +19,8 @@ DeepNestedSequenceProcessor, DeepNestedFloatsProcessor, ) -from .base_model import BaseModel +from ..base_model import BaseModel +from .base import BaseEmbeddingModel def _iter_text_vectors( @@ -147,7 +148,10 @@ def __init__( normalize_pretrained: bool = False, ): super().__init__(dataset) - self.embedding_dim = embedding_dim + # BaseEmbeddingModel declares `embedding_dim` as an abstract property, + # so we can't set self.embedding_dim directly (no setter). Use a + # private backing attribute and expose it through the property below. + self._embedding_dim = embedding_dim self.embedding_layers = nn.ModuleDict() for field_name, processor in self.dataset.input_processors.items(): @@ -164,21 +168,12 @@ def __init__( ), ): vocab_size = len(processor.code_vocab) - - if isinstance( - processor, (NestedSequenceProcessor, DeepNestedSequenceProcessor) - ): - self.embedding_layers[field_name] = nn.Embedding( - num_embeddings=vocab_size, - embedding_dim=embedding_dim, - padding_idx=0, - ) - else: - self.embedding_layers[field_name] = nn.Embedding( - num_embeddings=vocab_size, - embedding_dim=embedding_dim, - padding_idx=0, - ) + # Keep padding_idx=0 so the pad row receives no gradients. + self.embedding_layers[field_name] = nn.Embedding( + num_embeddings=vocab_size, + embedding_dim=embedding_dim, + padding_idx=0, + ) # Optional pretrained initialization (e.g., GloVe). if pretrained_emb_path is not None: @@ -338,7 +333,6 @@ def forward( if output_mask: # Generate a mask for this field - # For transformers, we might already have a mask, or use pad token if masks is not None and field_name in masks: out_masks[field_name] = masks[field_name].to(self.device) elif hasattr(processor, "code_vocab"): @@ -358,4 +352,4 @@ def forward( return embedded def __repr__(self) -> str: - return f"EmbeddingModel(embedding_layers={self.embedding_layers})" + return f"EmbeddingModel(embedding_layers={self.embedding_layers})" \ No newline at end of file diff --git a/pyhealth/models/embedding/vision.py b/pyhealth/models/embedding/vision.py new file mode 100644 index 000000000..57eedda22 --- /dev/null +++ b/pyhealth/models/embedding/vision.py @@ -0,0 +1,384 @@ +# Author: Josh Steier +# Description: Vision embedding model for medical imaging + +from typing import Any, Dict, Literal, Optional, Tuple, Union + +import torch +import torch.nn as nn +import shutil +from ...datasets import SampleDataset +from ..base_model import BaseModel +from ...processors import ImageProcessor +from .base import BaseEmbeddingModel + + +class Permute(nn.Module): + """Utility module to permute tensor dimensions in nn.Sequential. + + Args: + dims: Variable number of integers specifying the desired ordering + of dimensions. + + Example: + >>> permute = Permute(0, 2, 1) + >>> x = torch.randn(32, 256, 49) # (B, E, spatial) + >>> out = permute(x) # (B, spatial, E) + """ + + def __init__(self, *dims: int) -> None: + super().__init__() + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.permute(*self.dims) + + +class PatchEmbedding(nn.Module): + """Convert images to patch embeddings using ViT-style projection. + + Splits an image into non-overlapping patches and projects each patch + to an embedding vector using a convolutional layer. + + Args: + image_size: Input image size (assumes square images). + patch_size: Size of each square patch. + in_channels: Number of input channels. + embedding_dim: Output embedding dimension for each patch. + + Example: + >>> patch_embed = PatchEmbedding(224, 16, 3, 256) + >>> x = torch.randn(4, 3, 224, 224) + >>> patches = patch_embed(x) # (4, 196, 256) + """ + + def __init__( + self, + image_size: int = 224, + patch_size: int = 16, + in_channels: int = 3, + embedding_dim: int = 128, + ) -> None: + super().__init__() + if image_size % patch_size != 0: + raise ValueError( + f"image_size ({image_size}) must be divisible by " + f"patch_size ({patch_size})" + ) + self.patch_size = patch_size + self.num_patches = (image_size // patch_size) ** 2 + self.proj = nn.Conv2d( + in_channels, embedding_dim, kernel_size=patch_size, stride=patch_size + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # (B, C, H, W) -> (B, E, H/P, W/P) -> (B, num_patches, E) + x = self.proj(x) + x = x.flatten(2).transpose(1, 2) + return x + + +class VisionEmbeddingModel(BaseModel, BaseEmbeddingModel): + """Vision embedding model for medical image inputs. + + Converts medical images to sequences of patch embeddings suitable for + attention-based fusion with other modalities (EHR, text). + + Supports multiple backbone types: + - "patch": ViT-style patch projection (lightweight) + - "cnn": Small CNN encoder (good inductive bias) + - "resnet18"/"resnet50": Pretrained backbones + + Output shape: (batch, num_patches, embedding_dim) + + Args: + dataset: SampleDataset with ImageProcessor fields. + embedding_dim: Output embedding dimension. Default 128. + patch_size: Patch size for "patch" backbone. Default 16. + backbone: One of "patch", "cnn", "resnet18", "resnet50". + pretrained: Use ImageNet weights for ResNet. Default True. + freeze_backbone: Freeze pretrained weights. Default False. + dropout: Dropout rate. Default 0.0. + use_cls_token: Prepend learnable [CLS] token. Default False. + + Example: + >>> from pyhealth.datasets import create_sample_dataset + >>> model = VisionEmbeddingModel(dataset, embedding_dim=256) + >>> embeddings = model({"chest_xray": images}) + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: int = 128, + patch_size: int = 16, + backbone: Literal["patch", "cnn", "resnet18", "resnet50"] = "patch", + pretrained: bool = True, + 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 + + self.embedding_layers = nn.ModuleDict() + self.pos_embeddings = nn.ParameterDict() + self.cls_tokens = nn.ParameterDict() if use_cls_token else None + self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + + self._field_info: Dict[str, Dict[str, Any]] = {} + + for field_name, processor in self.dataset.input_processors.items(): + if not isinstance(processor, ImageProcessor): + continue + + image_size = processor.image_size + in_channels = self._infer_channels(processor) + + num_patches = self._build_embedding_layer( + field_name, image_size, in_channels, backbone, pretrained, freeze_backbone + ) + + num_positions = num_patches + 1 if use_cls_token else num_patches + self.pos_embeddings[field_name] = nn.Parameter( + torch.randn(1, num_positions, embedding_dim) * 0.02 + ) + + if use_cls_token: + self.cls_tokens[field_name] = nn.Parameter( + torch.randn(1, 1, embedding_dim) * 0.02 + ) + + self._field_info[field_name] = { + "num_patches": num_patches, + "image_size": image_size, + "in_channels": in_channels, + } + + @property + def embedding_dim(self) -> int: + return self._embedding_dim + + def _infer_channels(self, processor: ImageProcessor) -> int: + """Infer number of input channels from processor mode.""" + mode = getattr(processor, "mode", None) + if mode == "L": + return 1 + elif mode == "RGBA": + return 4 + return 3 + + def _build_embedding_layer( + self, + field_name: str, + image_size: int, + in_channels: int, + backbone: str, + pretrained: bool, + freeze_backbone: bool, + ) -> int: + """Build embedding layer and return number of output patches.""" + if backbone == "patch": + num_patches = (image_size // self.patch_size) ** 2 + self.embedding_layers[field_name] = PatchEmbedding( + image_size, self.patch_size, in_channels, self._embedding_dim + ) + + elif backbone == "cnn": + num_patches = 7 * 7 + self.embedding_layers[field_name] = nn.Sequential( + nn.Conv2d(in_channels, 64, 7, stride=2, padding=3), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.Conv2d(64, 128, 3, stride=2, padding=1), + nn.BatchNorm2d(128), + nn.ReLU(inplace=True), + nn.Conv2d(128, self._embedding_dim, 3, stride=2, padding=1), + nn.BatchNorm2d(self._embedding_dim), + nn.ReLU(inplace=True), + nn.AdaptiveAvgPool2d((7, 7)), + nn.Flatten(2), + Permute(0, 2, 1), + ) + + elif backbone in ("resnet18", "resnet50"): + num_patches = 7 * 7 + self.embedding_layers[field_name] = self._build_resnet_backbone( + backbone, in_channels, pretrained, freeze_backbone + ) + + else: + raise ValueError(f"Unknown backbone: {backbone}") + + return num_patches + + def _build_resnet_backbone( + self, backbone: str, in_channels: int, pretrained: bool, freeze: bool + ) -> nn.Module: + """Build pretrained ResNet backbone with spatial output.""" + try: + import torchvision.models as models + except ImportError as e: + raise ImportError("torchvision required for ResNet backbones") from e + + if backbone == "resnet18": + weights = models.ResNet18_Weights.DEFAULT if pretrained else None + resnet = models.resnet18(weights=weights) + feature_dim = 512 + else: + weights = models.ResNet50_Weights.DEFAULT if pretrained else None + resnet = models.resnet50(weights=weights) + feature_dim = 2048 + + if in_channels != 3: + resnet.conv1 = nn.Conv2d( + in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False + ) + + layers = list(resnet.children())[:-2] + backbone_net = nn.Sequential(*layers) + + if freeze: + for param in backbone_net.parameters(): + param.requires_grad = False + + return nn.Sequential( + backbone_net, + nn.Conv2d(feature_dim, self._embedding_dim, kernel_size=1), + nn.Flatten(2), + Permute(0, 2, 1), + ) + + def forward( + self, + inputs: Dict[str, torch.Tensor], + output_mask: bool = False, + ) -> Union[Dict[str, torch.Tensor], Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]]: + """Forward pass. + + Args: + inputs: Dict mapping field names to image tensors (B, C, H, W). + output_mask: If True, also return attention masks. + + Returns: + Dict of embeddings (B, num_patches, E), optionally with masks. + """ + embedded: Dict[str, torch.Tensor] = {} + masks: Dict[str, torch.Tensor] = {} if output_mask else None + + for field_name, tensor in inputs.items(): + if field_name not in self.embedding_layers: + embedded[field_name] = tensor + continue + + tensor = tensor.to(self.device) + batch_size = tensor.size(0) + + x = self.embedding_layers[field_name](tensor) + + if self.use_cls_token: + cls = self.cls_tokens[field_name].expand(batch_size, -1, -1) + x = torch.cat([cls, x], dim=1) + + 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: + masks[field_name] = torch.ones( + batch_size, x.size(1), dtype=torch.bool, device=x.device + ) + + return (embedded, masks) if output_mask else embedded + + def get_output_info(self, field_name: str) -> Dict[str, Any]: + """Get metadata about embedding output for a field.""" + if field_name not in self._field_info: + raise KeyError(f"Field '{field_name}' not found") + + info = self._field_info[field_name].copy() + info["embedding_dim"] = self._embedding_dim + info["has_cls_token"] = self.use_cls_token + 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: + fields = list(self.embedding_layers.keys()) + return ( + f"VisionEmbeddingModel(backbone={self.backbone_type!r}, " + f"embedding_dim={self._embedding_dim}, fields={fields})" + ) + + +if __name__ == "__main__": + from pyhealth.datasets import create_sample_dataset + from pyhealth.datasets.utils import get_dataloader + import tempfile + import os + from PIL import Image + import numpy as np + + # Create synthetic images + temp_dir = tempfile.mkdtemp() + samples = [] + for i in range(10): + img_path = os.path.join(temp_dir, f"img_{i}.png") + img = Image.fromarray(np.random.randint(0, 255, (224, 224), dtype=np.uint8), mode="L") + img.save(img_path) + samples.append({ + "patient_id": f"p{i}", + "visit_id": f"v{i}", + "chest_xray": img_path, + "label": i % 2, + }) + + dataset = create_sample_dataset( + samples=samples, + input_schema={"chest_xray": ("image", {"image_size": 224, "mode": "L"})}, + output_schema={"label": "binary"}, + dataset_name="test_vision", + ) + + model = VisionEmbeddingModel( + dataset=dataset, + embedding_dim=128, + backbone="cnn", + 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}") + print(f"Output info: {model.get_output_info('chest_xray')}") + + # Cleanup + shutil.rmtree(temp_dir) diff --git a/pyhealth/models/jamba_ehr.py b/pyhealth/models/jamba_ehr.py index fea902bd1..a08879915 100644 --- a/pyhealth/models/jamba_ehr.py +++ b/pyhealth/models/jamba_ehr.py @@ -13,8 +13,10 @@ import torch.nn as nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel from pyhealth.models.transformer import TransformerBlock from pyhealth.models.ehrmamba import MambaBlock from pyhealth.models.utils import get_last_visit @@ -177,6 +179,11 @@ class JambaEHR(BaseModel): by an independent :class:`JambaLayer`. The resulting patient embeddings are concatenated and projected through a classification head. + When ``unified_embedding`` is supplied the model switches to **unified + mode**: all temporal fields are jointly embedded and time-sorted by + :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single* + :class:`JambaLayer` rather than one layer per field. + Args: dataset (SampleDataset): Dataset providing processed inputs. embedding_dim (int): Embedding and hidden dimension. Default 128. @@ -186,6 +193,8 @@ class JambaEHR(BaseModel): dropout (float): Dropout rate. Default 0.3. state_size (int): SSM state size in Mamba blocks. Default 16. conv_kernel (int): Causal conv kernel in Mamba blocks. Default 4. + unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when + provided, enables unified multi-modal mode with a single JambaLayer. Examples: >>> from pyhealth.datasets import create_sample_dataset, get_dataloader @@ -234,6 +243,7 @@ def __init__( dropout: float = 0.3, state_size: int = 16, conv_kernel: int = 4, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, ): super(JambaEHR, self).__init__(dataset=dataset) self.embedding_dim = embedding_dim @@ -243,6 +253,7 @@ def __init__( self.dropout_rate = dropout self.state_size = state_size self.conv_kernel = conv_kernel + self._use_unified = unified_embedding is not None assert ( len(self.label_keys) == 1 @@ -250,11 +261,12 @@ def __init__( self.label_key = self.label_keys[0] self.mode = self.dataset.output_schema[self.label_key] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) + output_size = self.get_output_size() + self.dropout = nn.Dropout(dropout) - self.jamba: nn.ModuleDict = nn.ModuleDict() - for feature_key in self.feature_keys: - self.jamba[feature_key] = JambaLayer( + if self._use_unified: + self.embedding_model = unified_embedding + self._unified_jamba = JambaLayer( feature_size=embedding_dim, num_transformer_layers=num_transformer_layers, num_mamba_layers=num_mamba_layers, @@ -263,12 +275,69 @@ def __init__( state_size=state_size, conv_kernel=conv_kernel, ) + self.fc = nn.Linear(embedding_dim, output_size) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.jamba: nn.ModuleDict = nn.ModuleDict() + for feature_key in self.feature_keys: + self.jamba[feature_key] = JambaLayer( + feature_size=embedding_dim, + num_transformer_layers=num_transformer_layers, + num_mamba_layers=num_mamba_layers, + heads=heads, + dropout=dropout, + state_size=state_size, + conv_kernel=conv_kernel, + ) + self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build the inputs dict required by UnifiedMultimodalEmbeddingModel.""" + 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: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Calls UnifiedMultimodalEmbeddingModel to produce a single + temporally-sorted event sequence, then encodes it with one shared + JambaLayer and pools to the last valid event. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + sequence = out["sequence"] # (B, S_total, E) + mask = out["mask"] # (B, S_total) float, 1=valid 0=pad - output_size = self.get_output_size() - self.dropout = nn.Dropout(dropout) - self.fc = nn.Linear( - len(self.feature_keys) * embedding_dim, output_size - ) + _, cls_emb = self._unified_jamba(sequence, mask) + logits = self.fc(self.dropout(cls_emb)) + 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) + results["loss"] = self.get_loss_function()(logits, y_true) + results["y_true"] = y_true + if kwargs.get("embed", False): + results["embed"] = cls_emb + return results @staticmethod def _pool_embedding(x: torch.Tensor) -> torch.Tensor: @@ -317,9 +386,10 @@ def forward( ) -> Dict[str, torch.Tensor]: """Forward propagation. - Embeds each feature stream, encodes through the hybrid - Transformer-Mamba stack, concatenates per-stream patient - representations, and projects to label space. + In **unified mode** (when ``unified_embedding`` was supplied at init) + the model jointly embeds all temporal fields and processes them with a + single JambaLayer backbone. Otherwise each field is embedded and + encoded independently. Args: **kwargs: Must include all feature keys (tensors or tuples @@ -330,6 +400,9 @@ def forward( ``y_prob``, ``y_true``, ``logit``, and optionally ``embed`` if ``kwargs["embed"] is True``. """ + if self._use_unified: + return self._forward_unified(**kwargs) + patient_emb = [] for feature_key in self.feature_keys: diff --git a/pyhealth/models/mlp.py b/pyhealth/models/mlp.py index 299dc151e..3d964bab6 100644 --- a/pyhealth/models/mlp.py +++ b/pyhealth/models/mlp.py @@ -1,13 +1,15 @@ -from typing import Dict, cast +from typing import Any, Dict, Optional, cast import torch import torch.nn as nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.interpret.api import Interpretable from .embedding import EmbeddingModel +from .embedding.unified import UnifiedMultimodalEmbeddingModel class MLP(BaseModel, Interpretable): @@ -110,12 +112,14 @@ def __init__( hidden_dim: int = 128, n_layers: int = 2, activation: str = "relu", + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, **kwargs, ): super(MLP, self).__init__(dataset) self.embedding_dim = embedding_dim self.hidden_dim = hidden_dim self.n_layers = n_layers + self._use_unified = unified_embedding is not None # validate kwargs for MLP layer if "input_size" in kwargs: @@ -126,9 +130,6 @@ def __init__( assert len(self.label_keys) == 1, "Only one label key is supported" self.label_key = self.label_keys[0] - # Use the EmbeddingModel to handle embedding logic - self.embedding_model = EmbeddingModel(dataset, embedding_dim) - # Set up activation function if activation == "relu": self.activation = nn.ReLU() @@ -143,18 +144,77 @@ def __init__( else: raise ValueError(f"Unsupported activation function {activation}") - # Create MLP layers for each feature - self.mlp = nn.ModuleDict() - for feature_key in self.feature_keys: - Modules = [] - Modules.append(nn.Linear(self.embedding_dim, self.hidden_dim)) - for _ in range(self.n_layers - 1): - Modules.append(self.activation) - Modules.append(nn.Linear(self.hidden_dim, self.hidden_dim)) - self.mlp[feature_key] = nn.Sequential(*Modules) - output_size = self.get_output_size() - self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size) + + if self._use_unified: + self.embedding_model = unified_embedding + modules = [nn.Linear(embedding_dim, hidden_dim)] + for _ in range(n_layers - 1): + modules.extend([self.activation, nn.Linear(hidden_dim, hidden_dim)]) + self.mlp = nn.ModuleDict({"unified": nn.Sequential(*modules)}) + self.fc = nn.Linear(hidden_dim, output_size) + else: + # Use the EmbeddingModel to handle embedding logic + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + # Create MLP layers for each feature + self.mlp = nn.ModuleDict() + for feature_key in self.feature_keys: + modules = [nn.Linear(self.embedding_dim, self.hidden_dim)] + for _ in range(self.n_layers - 1): + modules.extend([self.activation, nn.Linear(self.hidden_dim, self.hidden_dim)]) + self.mlp[feature_key] = nn.Sequential(*modules) + self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build the inputs dict required by UnifiedMultimodalEmbeddingModel.""" + 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: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Embeds all temporal fields jointly, mean-pools the event sequence, + applies a single MLP, and projects to label space. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + sequence = out["sequence"] # (B, S, E) + mask = out["mask"].float() # (B, S) + + # Masked mean-pool over the event sequence + x = (sequence * mask.unsqueeze(-1)).sum(dim=1) + x = x / mask.sum(dim=1, keepdim=True).clamp(min=1) # (B, E) + + x = self.mlp["unified"](x) # (B, hidden_dim) + logits = self.fc(x) + 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) + results["loss"] = self.get_loss_function()(logits, y_true) + results["y_true"] = y_true + if kwargs.get("embed", False): + results["embed"] = x + return results @staticmethod def mean_pooling(x, mask): @@ -309,6 +369,11 @@ def forward( ) -> Dict[str, torch.Tensor]: """Forward propagation. + In **unified mode** (when ``unified_embedding`` was supplied at init) + the model jointly embeds all temporal fields, mean-pools the event + sequence, and processes it with a single MLP. Otherwise each field + is embedded and encoded independently. + Args: **kwargs: keyword arguments for the model. @@ -326,6 +391,9 @@ def forward( logit: the raw logits before activation. embed: (if embed=True in kwargs) the patient embedding. """ + if self._use_unified: + return self._forward_unified(**kwargs) + for feature_key in self.feature_keys: feature = kwargs[feature_key] @@ -355,10 +423,9 @@ def forward( batch_size, seq_len, inner_len = value.shape value = value.view(batch_size, seq_len * inner_len) if mask is not None: - mask = mask.to(self.device) - # Flatten mask properly if it exists - if mask.dim() == 3: - mask = mask.view(batch_size, seq_len * inner_len) + mask = mask.to(self.device) + if mask.dim() == 3: + mask = mask.view(batch_size, seq_len * inner_len) if mask is not None: mask = mask.to(self.device) diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py index 3393d7287..f68c368a7 100644 --- a/pyhealth/models/rnn.py +++ b/pyhealth/models/rnn.py @@ -1,10 +1,11 @@ -from typing import Dict, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import torch import torch.nn as nn import torch.nn.utils.rnn as rnn_utils from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.processors import ( DeepNestedFloatsProcessor, @@ -20,6 +21,7 @@ ) from .embedding import EmbeddingModel +from .embedding.unified import UnifiedMultimodalEmbeddingModel class RNNLayer(nn.Module): @@ -92,9 +94,7 @@ def forward( Args: x: a tensor of shape [batch size, sequence len, input size]. mask: an optional tensor of shape [batch size, sequence len], where - 1 indicates valid and 0 indicates invalid. Samples with all-zero - masks are clamped to length 1 to prevent pack_padded_sequence - from receiving zero-length sequences. + 1 indicates valid and 0 indicates invalid. Returns: outputs: a tensor of shape [batch size, sequence len, hidden size], @@ -111,13 +111,14 @@ def forward( ) else: lengths = torch.sum(mask.int(), dim=-1).cpu() - # Clamp lengths to at least 1 to handle empty sequences, - # matching TCNLayer (tcn.py:186). - lengths = torch.clamp(lengths, min=1) + # pack_padded_sequence rejects a zero length. Before batch padding + # was masked this was unreachable; a correct mask makes a sample + # with no valid event reachable, so clamp to 1. + lengths = torch.clamp(lengths, min=1) # Ensure tensor is contiguous for cuDNN compatibility x = x.contiguous() x = rnn_utils.pack_padded_sequence( - x, lengths, batch_first=True, enforce_sorted=False + x.contiguous(), lengths, batch_first=True, enforce_sorted=False ) outputs, _ = self.rnn(x) outputs, _ = rnn_utils.pad_packed_sequence(outputs, batch_first=True) @@ -211,6 +212,7 @@ def __init__( dataset: SampleDataset, embedding_dim: int = 128, hidden_dim: int = 128, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, **kwargs ): super(RNN, self).__init__( @@ -218,6 +220,7 @@ def __init__( ) self.embedding_dim = embedding_dim self.hidden_dim = hidden_dim + self._use_unified = unified_embedding is not None # validate kwargs for RNN layer if "input_size" in kwargs: raise ValueError("input_size is determined by embedding_dim") @@ -227,20 +230,74 @@ def __init__( self.label_key = self.label_keys[0] self.mode = self.dataset.output_schema[self.label_key] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) - - self.rnn = nn.ModuleDict() - for feature_key in self.dataset.input_processors.keys(): - self.rnn[feature_key] = RNNLayer( - input_size=embedding_dim, hidden_size=hidden_dim, **kwargs - ) output_size = self.get_output_size() - self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size) + + if self._use_unified: + self.embedding_model = unified_embedding + self.rnn = nn.ModuleDict({ + "unified": RNNLayer(input_size=embedding_dim, hidden_size=hidden_dim, **kwargs) + }) + self.fc = nn.Linear(hidden_dim, output_size) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.rnn = nn.ModuleDict() + for feature_key in self.dataset.input_processors.keys(): + self.rnn[feature_key] = RNNLayer( + input_size=embedding_dim, hidden_size=hidden_dim, **kwargs + ) + self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build the inputs dict required by UnifiedMultimodalEmbeddingModel.""" + 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: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Embeds all temporal fields jointly as a single time-sorted sequence + and processes it with one RNN backbone. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + sequence = out["sequence"] # (B, S, E) + mask = out["mask"].int() # (B, S) + + _, last_hidden = self.rnn["unified"](sequence, mask) # (B, hidden_dim) + logits = self.fc(last_hidden) + 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"] = last_hidden + return results def forward(self, **kwargs) -> Dict[str, torch.Tensor]: """Forward propagation. - The label `kwargs[self.label_key]` is a list of labels for each patient. + In **unified mode** (when ``unified_embedding`` was supplied at init) + the model jointly embeds all temporal fields as a single time-sorted + sequence and processes it with one RNN. Otherwise each field is + embedded and encoded independently. Args: **kwargs: keyword arguments for the model. The keys must contain @@ -254,6 +311,9 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: - logit: a tensor representing the logits. - embed (optional): a tensor representing the patient embeddings if requested. """ + if self._use_unified: + return self._forward_unified(**kwargs) + patient_emb = [] # We need to preprocess kwargs to extract values and masks for EmbeddingModel @@ -569,4 +629,4 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: } if kwargs.get("embed", False): results["embed"] = patient_emb - return results + return results \ No newline at end of file diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py index cc0dfc5ca..9fc835df5 100644 --- a/pyhealth/models/transformer.py +++ b/pyhealth/models/transformer.py @@ -7,11 +7,14 @@ from typing import Any, Dict, Optional, Tuple, Union, cast import torch +import torch.nn.functional as F from torch import nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel from pyhealth.interpret.api import CheferInterpretable # VALID_OPERATION_LEVEL = ["visit", "event"] @@ -54,9 +57,11 @@ def forward( # 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) + pad_mask = mask == 0 + scores = scores.masked_fill(pad_mask, torch.finfo(scores.dtype).min) 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) @@ -149,22 +154,41 @@ def forward( ] # 2) Apply attention on all the projected vectors in batch. - if mask is not None: - mask = mask.unsqueeze(1) - x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout) - - 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. + # Ordinary training uses fused SDPA. The explicit path stays behind + # register_hook=True for interpretability. + if not register_hook: + query_mask = None + attn_mask = None + if mask is not None: + valid = mask.bool() + if mask.dim() == 2: + query_mask = valid[:, None, :, None] + attn_mask = valid[:, None, None, :] + else: + query_mask = valid.any(dim=-1)[:, None, :, None] + attn_mask = valid.unsqueeze(1) + x = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=self.dropout.p if self.training else 0.0, + ) + if query_mask is not None: + x = x * query_mask.to(x.dtype) + self.attn_map = None + self.attn_gradients = None + else: + if mask is not None: + mask = mask.unsqueeze(1) + x, attn = self.attention( + query, key, value, mask=mask, dropout=self.dropout + ) 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) @@ -246,7 +270,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: @@ -256,7 +280,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) @@ -297,7 +326,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. @@ -328,12 +360,21 @@ class Transformer(BaseModel, CheferInterpretable): an independent :class:`TransformerLayer`. The resulting [CLS]-style embeddings are concatenated and passed to a classification head. + When ``unified_embedding`` is supplied the model switches to **unified + mode**: all temporal fields are jointly embedded and time-sorted by + :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single* + :class:`TransformerLayer` rather than one layer per field. This allows + full cross-modal attention over the interleaved event sequence. + Args: dataset (SampleDataset): dataset providing processed inputs. embedding_dim (int): shared embedding dimension. heads (int): number of attention heads per transformer block. dropout (float): dropout rate applied inside transformer blocks. num_layers (int): number of transformer blocks per feature stream. + unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when + provided, the model uses a single backbone over the unified + multi-modal sequence instead of per-field transformers. Examples: >>> from pyhealth.datasets import create_sample_dataset, get_dataloader @@ -377,6 +418,7 @@ def __init__( dropout: float = 0.5, num_layers: int = 1, max_seq_len: int = 1024, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, ): super().__init__(dataset=dataset) self.embedding_dim = embedding_dim @@ -385,6 +427,7 @@ def __init__( 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 assert ( len(self.label_keys) == 1 @@ -392,19 +435,28 @@ def __init__( self.label_key = self.label_keys[0] self.mode = self.dataset.output_schema[self.label_key] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) + output_size = self.get_output_size() - self.transformer: nn.ModuleDict = nn.ModuleDict() - for feature_key in self.feature_keys: - self.transformer[feature_key] = TransformerLayer( + if self._use_unified: + self.embedding_model = unified_embedding + self._unified_backbone = TransformerLayer( feature_size=embedding_dim, heads=heads, dropout=dropout, num_layers=num_layers, ) - - output_size = self.get_output_size() - self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + self.fc = nn.Linear(embedding_dim, output_size) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.transformer: nn.ModuleDict = nn.ModuleDict() + for feature_key in self.feature_keys: + self.transformer[feature_key] = TransformerLayer( + feature_size=embedding_dim, + heads=heads, + dropout=dropout, + num_layers=num_layers, + ) + 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. @@ -443,6 +495,64 @@ def _mask_from_embeddings(x: torch.Tensor) -> torch.Tensor: mask[invalid_rows, 0] = True return mask.bool() + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build inputs expected by UnifiedMultimodalEmbeddingModel.""" + + 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: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_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.""" + + register_hook = self._attention_hooks_enabled + inputs = self._build_unified_inputs(cast(Dict[str, Any], kwargs)) + out = self.embedding_model(inputs) + sequence = cast(torch.Tensor, out["sequence"]) + event_mask = cast(torch.Tensor, out["mask"]).bool() + + _, 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, + } + + 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 + + if kwargs.get("embed", False): + results["embed"] = patient_emb + return results + def forward_from_embedding( self, **kwargs: torch.Tensor | tuple[torch.Tensor, ...], @@ -500,8 +610,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) @@ -515,9 +624,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) @@ -545,6 +652,11 @@ def forward( ) -> Dict[str, torch.Tensor]: """Forward propagation. + In **unified mode** (when ``unified_embedding`` was supplied at init) + the model jointly embeds all temporal fields and processes them with a + single transformer backbone. Otherwise each field is embedded and + encoded independently. + Args: **kwargs: keyword arguments for the model. @@ -562,6 +674,9 @@ def forward( logit: the raw logits before activation. embed: (if embed=True in kwargs) the patient embedding. """ + if self._use_unified: + return self._forward_unified(**kwargs) + for feature_key in self.feature_keys: feature = kwargs[feature_key] @@ -575,15 +690,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] @@ -591,9 +707,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) @@ -621,9 +737,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 } @@ -683,4 +797,4 @@ def get_relevance_tensor( result = model(**data_batch) print(result) - result["loss"].backward() + result["loss"].backward() \ No newline at end of file diff --git a/pyhealth/models/unified_embedding.py b/pyhealth/models/unified_embedding.py index 014326b41..0b2771989 100644 --- a/pyhealth/models/unified_embedding.py +++ b/pyhealth/models/unified_embedding.py @@ -1,327 +1,16 @@ -"""UnifiedMultimodalEmbeddingModel — temporally aligned multimodal embedding. +"""Deprecated import path for the unified multimodal encoder. -Takes K temporal features ( dict outputs from ``TemporalFeatureProcessor`` -subclasses ), embeds each event with a modality-specific encoder, then -interleaves all events on a shared timeline by sorting on timestamp and adding -sinusoidal time embeddings + learned modality-type embeddings. - -Output shape: ``(B, S_total, E')`` — a single sequence of events usable by -any downstream sequence model (Transformer, Mamba, RNN, …). - -Quickstart:: - - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel - from pyhealth.datasets.collate import collate_temporal - model = UnifiedMultimodalEmbeddingModel(dataset, embedding_dim=128) - # inside forward: - # inputs = {field: {"value": Tensor, "time": Tensor, ...}, ...} - out = model(inputs) - # out["sequence"]: (B, S_total, 128) - # out["mask"]: (B, S_total) — 1 = real event, 0 = padding - # out["time"]: (B, S_total) — hours from first event +Use ``pyhealth.models.embedding`` instead. This module re-exports the live +classes so older ``from pyhealth.models.unified_embedding import ...`` call +sites keep working against one implementation. """ -from __future__ import annotations - -import math -from typing import Any - -import torch -import torch.nn as nn - -from pyhealth.processors.base_processor import ModalityType, TemporalFeatureProcessor - - -# ── Helpers ─────────────────────────────────────────────────────────────────── - - -class SinusoidalTimeEmbedding(nn.Module): - """Continuous sinusoidal embedding for scalar time values (in hours). - - Identical in spirit to the positional encoding in "Attention is All You - Need" but operating on real-valued timestamps rather than integer positions. - - Args: - dim: Output embedding dimension (must be even). - max_hours: Maximum expected time value in hours. Values are normalised - to ``[0, 2π]`` before the sin/cos projection. Default 720 (30 days). - - Shape: - Input: ``(*, )`` float tensor of times in hours - Output: ``(*, dim)`` - """ - - def __init__(self, dim: int, max_hours: float = 720.0): - super().__init__() - assert dim % 2 == 0, f"dim must be even, got {dim}" - self.dim = dim - self.max_hours = max_hours - half = dim // 2 - freqs = torch.exp( - -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / (half - 1) - ) - self.register_buffer("freqs", freqs) # (dim//2,) - - def forward(self, t: torch.Tensor) -> torch.Tensor: - """:param t: ``(...,)`` float, times in hours.""" - t_norm = t / self.max_hours * 2 * math.pi # (...,) - args = t_norm.unsqueeze(-1) * self.freqs # (..., dim//2) - return torch.cat([args.sin(), args.cos()], dim=-1) # (..., dim) - - -def _build_image_encoder(embedding_dim: int) -> nn.Module: - """Lightweight 5-layer CNN encoder: C × H × W → embedding_dim. - - Uses ``torchvision.models.resnet18`` pre-trained backbone, strips the - final FC layer, and adds a projection to ``embedding_dim``. Falls back to - a toy Conv-pool-flatten network if torchvision is not installed. - """ - try: - import torchvision.models as tv - - backbone = tv.resnet18(weights=None) - in_features = backbone.fc.in_features - backbone.fc = nn.Linear(in_features, embedding_dim) - return backbone - except ImportError: - # Minimal fallback: single conv → global avg pool → linear - return nn.Sequential( - nn.Conv2d(3, 32, 3, padding=1), - nn.ReLU(), - nn.AdaptiveAvgPool2d(1), - nn.Flatten(), - nn.Linear(32, embedding_dim), - ) - - -# ── Main model ─────────────────────────────────────────────────────────────── - - -class UnifiedMultimodalEmbeddingModel(nn.Module): - """Embed heterogeneous temporal features into a single aligned sequence. - - **All** input processors must be ``TemporalFeatureProcessor`` subclasses. - Non-temporal processors (e.g. ``SequenceProcessor``, ``MultiHotProcessor``) - are rejected with a clear error — use the existing ``EmbeddingModel`` for - those fields. - - Algorithm - --------- - For each temporal field: - - 1. Route ``inputs[field]["value"]`` through a modality-specific encoder → - ``(B, N_i, E')`` per-event embeddings. - 2. Retrieve ``inputs[field]["time"]`` → ``(B, N_i)`` timestamps (hours). - 3. (Optional) Retrieve ``inputs[field]["mask"]`` → ``(B, N_i, L)`` or - ``(B, N_i)`` attention mask; reduced to event-level ``(B, N_i)`` if - token-level. - - Then: - - 4. Concatenate across all fields → ``(B, S_total, E')``. - 5. Sort events along dim=1 by timestamp (ascending). - 6. Add ``SinusoidalTimeEmbedding(time)`` + ``type_embedding(modality_idx)``. - 7. Return ``{"sequence", "time", "mask", "type_ids"}``. - - Args: - processors: ``dict[field_name, TemporalFeatureProcessor]`` — the - processors for each temporal field in the dataset. Pass - ``dataset.input_processors`` directly. - embedding_dim: Shared embedding dimension ``E'``. - time_embedding: ``"sinusoidal"`` (default) or ``"learned"``. - max_time_hours: Normalisation constant for the time embedding. - Defaults to 720 h (30 days). - - Example:: - - model = UnifiedMultimodalEmbeddingModel( - processors=dataset.input_processors, - embedding_dim=128, - ) - # inputs: {field: {"value": Tensor, "time": Tensor, "mask": Tensor}} - out = model(inputs) - seq = out["sequence"] # (B, S_total, 128) - mask = out["mask"] # (B, S_total) float, 1=valid 0=pad - """ - - def __init__( - self, - processors: dict[str, Any], - embedding_dim: int = 128, - time_embedding: str = "sinusoidal", - max_time_hours: float = 720.0, - ): - super().__init__() - self.embedding_dim = embedding_dim - - self.encoders: nn.ModuleDict = nn.ModuleDict() - self.projections: nn.ModuleDict = nn.ModuleDict() - self.modality_types: dict[str, ModalityType] = {} - - for field_name, processor in processors.items(): - if not isinstance(processor, TemporalFeatureProcessor): - raise TypeError( - f"UnifiedMultimodalEmbeddingModel requires every input processor " - f"to be a TemporalFeatureProcessor subclass, but '{field_name}' " - f"uses {type(processor).__name__}. For non-temporal fields use " - f"the existing EmbeddingModel." - ) - - m = processor.modality() - self.modality_types[field_name] = m - - if m == ModalityType.CODE: - vocab_size = processor.value_dim() - self.encoders[field_name] = nn.Embedding( - vocab_size, embedding_dim, padding_idx=0 - ) - - elif m == ModalityType.TEXT: - if processor.is_token(): - from transformers import AutoModel - - bert = AutoModel.from_pretrained(processor.tokenizer_model) - self.encoders[field_name] = bert - hidden = bert.config.hidden_size - if hidden != embedding_dim: - self.projections[field_name] = nn.Linear(hidden, embedding_dim) - else: - raise ValueError( - f"TEXT processor '{field_name}' must use a tokenizer " - f"(set tokenizer_model=...) to be used with " - f"UnifiedMultimodalEmbeddingModel." - ) - - elif m == ModalityType.IMAGE: - self.encoders[field_name] = _build_image_encoder(embedding_dim) - - elif m in (ModalityType.NUMERIC, ModalityType.SIGNAL): - in_features = processor.value_dim() - self.encoders[field_name] = nn.Linear(in_features, embedding_dim) - - else: - raise NotImplementedError( - f"No encoder implemented for modality {m!r} (field '{field_name}')." - ) - - # Shared type embedding — one vector per unique modality in this dataset - unique_modalities = sorted(set(self.modality_types.values())) - self._modality_to_idx: dict[ModalityType, int] = { - mod: i for i, mod in enumerate(unique_modalities) - } - self.type_embedding = nn.Embedding(len(unique_modalities), embedding_dim) - - # Time embedding - if time_embedding == "sinusoidal": - self.time_embed = SinusoidalTimeEmbedding(embedding_dim, max_time_hours) - else: - raise NotImplementedError("Only 'sinusoidal' time embedding is implemented.") - - # ── Forward ─────────────────────────────────────────────────────────────── - - def forward( - self, - inputs: dict[str, dict[str, torch.Tensor]], - ) -> dict[str, torch.Tensor]: - """Encode and temporally align all temporal features. - - Args: - inputs: ``{field_name: {"value": Tensor, "time": Tensor, - "mask": Tensor (optional)}}`` - — one dict per temporal feature, exactly as produced by - ``collate_temporal``. - - Returns: - A dict with keys: - - * ``"sequence"`` — ``(B, S_total, E')`` temporally-sorted events - * ``"time"`` — ``(B, S_total)`` timestamps (hours) - * ``"mask"`` — ``(B, S_total)`` 1=real event, 0=padding - * ``"type_ids"`` — ``(B, S_total)`` modality index per event - """ - all_embeddings: list[torch.Tensor] = [] - all_times: list[torch.Tensor] = [] - all_masks: list[torch.Tensor] = [] - all_types: list[torch.Tensor] = [] - - for field_name, feat_dict in inputs.items(): - value = feat_dict["value"] # (B, N_i, ...) or (B, S, F) - time = feat_dict["time"] # (B, N_i) - mask = feat_dict.get("mask") - - if time is None: - # Fallback: treat every event as occurring at t=0 - time = torch.zeros(value.shape[:2], device=value.device) - - modality = self.modality_types[field_name] - encoder = self.encoders[field_name] - - # ── Encode ──────────────────────────────────────────────────── - if modality == ModalityType.CODE: - emb = encoder(value) # (B, S, E') - - elif modality == ModalityType.TEXT: - b, n, l = value.shape - flat_ids = value.view(b * n, l) - flat_mask = mask.view(b * n, l) if mask is not None else None - out = encoder(input_ids=flat_ids, attention_mask=flat_mask) - cls_emb = out.last_hidden_state[:, 0, :] # (B*N, H) - if field_name in self.projections: - cls_emb = self.projections[field_name](cls_emb) - emb = cls_emb.view(b, n, -1) # (B, N, E') - - elif modality == ModalityType.IMAGE: - b, n, c, h, w = value.shape - flat_imgs = value.view(b * n, c, h, w) - img_emb = encoder(flat_imgs) # (B*N, E') - emb = img_emb.view(b, n, -1) - - else: # NUMERIC / SIGNAL - emb = encoder(value) # (B, T, E') - - # ── Build event-level validity mask ─────────────────────────── - if mask is None: - event_mask = torch.ones(emb.shape[:2], device=emb.device) - else: - if mask.dim() > time.dim(): - # token-level (B, N, L) → event-level (B, N) - event_mask = (mask.sum(dim=-1) > 0).float() - else: - event_mask = mask.float() - - # ── Modality type indices ───────────────────────────────────── - type_idx = self._modality_to_idx[modality] - type_ids = torch.full( - emb.shape[:2], type_idx, dtype=torch.long, device=emb.device - ) - - all_embeddings.append(emb) - all_times.append(time) - all_masks.append(event_mask) - all_types.append(type_ids) - - # ── Concatenate across all fields ───────────────────────────────── - cat_emb = torch.cat(all_embeddings, dim=1) # (B, S_total, E') - cat_time = torch.cat(all_times, dim=1) # (B, S_total) - cat_mask = torch.cat(all_masks, dim=1) # (B, S_total) - cat_types = torch.cat(all_types, dim=1) # (B, S_total) - - # ── Sort by time ────────────────────────────────────────────────── - sort_idx = cat_time.argsort(dim=1) - cat_emb = cat_emb.gather( - 1, sort_idx.unsqueeze(-1).expand_as(cat_emb) - ) - cat_time = cat_time.gather(1, sort_idx) - cat_mask = cat_mask.gather(1, sort_idx) - cat_types = cat_types.gather(1, sort_idx) - # ── Add time + type embeddings ──────────────────────────────────── - time_emb = self.time_embed(cat_time) # (B, S_total, E') - type_emb = self.type_embedding(cat_types) # (B, S_total, E') - final = cat_emb + time_emb + type_emb # (B, S_total, E') +from pyhealth.models.embedding.unified import ( + SinusoidalTimeEmbedding, + UnifiedMultimodalEmbeddingModel, +) - return { - "sequence": final, # (B, S_total, E') - "time": cat_time, # (B, S_total) - "mask": cat_mask, # (B, S_total) - "type_ids": cat_types, # (B, S_total) - } +__all__ = [ + "SinusoidalTimeEmbedding", + "UnifiedMultimodalEmbeddingModel", +] diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index 4568a5ece..7ac3d9651 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -82,4 +82,12 @@ def get_processor(name: str): "TupleTimeTextProcessor", "CehrProcessor", "ConceptVocab", + "LabStandardizer", + "fit_lab_standardizer", + "lab_standardizer_fit_scope", ] +from .lab_standardizer import ( + LabStandardizer, + fit_lab_standardizer, + lab_standardizer_fit_scope, +) diff --git a/pyhealth/processors/lab_standardizer.py b/pyhealth/processors/lab_standardizer.py new file mode 100644 index 000000000..3a9aa1559 --- /dev/null +++ b/pyhealth/processors/lab_standardizer.py @@ -0,0 +1,289 @@ +"""Train-split-only standardisation for masked temporal laboratory values. + +The task records keep labs and their observation mask as separate temporal +fields. This module deliberately fits only rows whose corresponding mask is +true: zero-filled / forward-filled missing values must never affect a lab's +mean or variance. +""" + +from __future__ import annotations + +from collections.abc import Iterable +import hashlib +import json +from typing import Any, Optional + +import torch +from torch import nn + + +def _provenance_indices(dataset: Any) -> Optional[list[int]]: + """Every sample index of this split, or ``None`` for a plain iterable. + + ``SampleDataset`` subclasses ``litdata.StreamingDataset``, whose ``__iter__`` + and ``__len__`` are both sharded by ``WORLD_SIZE``. Under ``torchrun`` that + silently reduces a fit to 1/WORLD_SIZE of the train split, and to the *same* + shard on every rank, because ``torch.distributed`` is not yet initialised + when the dataset is built. Indexing is not sharded, so the fit is driven by + explicit indices. + + The count comes from ``region_of_interest``, which is the only unsharded + description of what this dataset holds. Measured against real litdata with + 20 samples: ``len()`` reports 5 under ``WORLD_SIZE=4`` while the region of + interest still sums to 20. + + ``patient_to_index`` is NOT usable here. ``SampleDataset.subset`` copies it + unchanged, so after ``split_by_patient`` it still holds indices into the + PARENT dataset while ``__getitem__`` is restricted to the subset's own + region. Driving the fit from it made a real training split raise + ``ValueError: The provided index 237 didn't find a match within the chunk + intervals``. + """ + roi = getattr(dataset, "region_of_interest", None) + if not roi: + return None + return list(range(sum(end - start for start, end in roi))) + + +def _is_shardable_dataset(obj: Any) -> bool: + """Whether iterating ``obj`` risks silently yielding only one shard. + + ``litdata.StreamingDataset`` shards ``__iter__`` by ``WORLD_SIZE``. Anything + that subclasses it must be driven by explicit indices instead. + """ + try: + from litdata.streaming.dataset import StreamingDataset + except Exception: # litdata absent: nothing can be sharded + return False + return isinstance(obj, StreamingDataset) + + +class LabStandardizer(nn.Module): + """Per-feature z-score transform with persistent train-only statistics. + + ``mean``, ``std`` and ``observed_count`` are buffers rather than ordinary + attributes. Consequently they are included in every model ``state_dict`` + and a checkpoint always transforms raw serving inputs exactly as it did at + training time. + + Constant train features use a unit denominator. Features with no observed + train values are emitted as zero, because a scale for them cannot be learnt + without looking at validation/test data. + """ + + def __init__( + self, + mean: torch.Tensor, + std: torch.Tensor, + observed_count: torch.Tensor, + *, + version: int = 2, + fit_scope: str | bytes | None = None, + ) -> None: + super().__init__() + if mean.ndim != 1 or std.shape != mean.shape or observed_count.shape != mean.shape: + raise ValueError("Lab standardisation statistics must be one vector per feature.") + if not torch.isfinite(mean).all() or not torch.isfinite(std).all(): + raise ValueError("Lab standardisation statistics must be finite.") + if (std <= 0).any() or (observed_count < 0).any(): + raise ValueError("Lab standardisation std must be positive and counts non-negative.") + self.register_buffer("mean", mean.detach().to(torch.float32).clone()) + self.register_buffer("std", std.detach().to(torch.float32).clone()) + self.register_buffer( + "observed_count", observed_count.detach().to(torch.long).clone() + ) + self.register_buffer("version", torch.tensor([version], dtype=torch.long)) + # This is a privacy-preserving SHA-256 digest of the exact train-split + # identity. It lets transfer loading reject a checkpoint whose fitted + # transform came from a different cohort/split before that transform can + # reach a downstream test patient. + self.register_buffer( + "fit_scope_digest", self._scope_digest(fit_scope), persistent=True + ) + + @property + def feature_dim(self) -> int: + return int(self.mean.numel()) + + @classmethod + def fit( + cls, + samples: Iterable[dict[str, Any]], + *, + value_field: str = "labs", + observation_mask_field: Optional[str] = None, + fit_scope: str | bytes | None = None, + ) -> "LabStandardizer": + """Fit only observed, finite values from the supplied samples. + + The caller supplies the already split training dataset. No reference + to a parent/full dataset is used here, which makes train-only fitting + auditable at the call site. + """ + observation_mask_field = observation_mask_field or f"{value_field}_mask" + indices = _provenance_indices(samples) + if indices is None and _is_shardable_dataset(samples): + # A plain list of dicts is safe to iterate; a StreamingDataset is not, + # because __iter__ is sharded by WORLD_SIZE. Falling back to iteration + # for one while silently doing it for the other is how a fit ends up + # on 1/N of the training data with nothing reporting it. + raise ValueError( + "Refusing to fit lab standardisation by iterating a " + f"{type(samples).__name__}: its __iter__ is sharded by WORLD_SIZE, " + "so under torchrun this would silently fit on a fraction of the " + "training split. The dataset exposes no region_of_interest " + "to drive an unsharded fit." + ) + stream = samples if indices is None else (samples[index] for index in indices) + consumed = 0 + count: Optional[torch.Tensor] = None + total: Optional[torch.Tensor] = None + total_sq: Optional[torch.Tensor] = None + + for sample in stream: + consumed += 1 + if value_field not in sample or observation_mask_field not in sample: + continue + values = cls._value_tensor(sample[value_field]).to(torch.float64) + observed = cls._value_tensor(sample[observation_mask_field]).bool() + if values.shape != observed.shape: + raise ValueError( + f"{value_field!r} and {observation_mask_field!r} must have the " + f"same shape, got {tuple(values.shape)} and {tuple(observed.shape)}." + ) + if values.ndim == 1: + values = values.unsqueeze(-1) + observed = observed.unsqueeze(-1) + if values.ndim != 2: + raise ValueError( + f"Expected temporal values shaped (time, features), got {tuple(values.shape)}." + ) + valid = observed & torch.isfinite(values) + if count is None: + feature_dim = values.shape[-1] + count = torch.zeros(feature_dim, dtype=torch.long) + total = torch.zeros(feature_dim, dtype=torch.float64) + total_sq = torch.zeros(feature_dim, dtype=torch.float64) + elif values.shape[-1] != count.numel(): + raise ValueError("All fitted samples must have the same laboratory feature dimension.") + + assert total is not None and total_sq is not None + count += valid.sum(dim=0).to(torch.long) + total += torch.where(valid, values, torch.zeros_like(values)).sum(dim=0) + total_sq += torch.where(valid, values.square(), torch.zeros_like(values)).sum(dim=0) + + if indices is not None and consumed != len(indices): + raise RuntimeError( + f"Lab standardisation consumed {consumed} of the {len(indices)} " + "samples its split declares; the statistics would be fitted on a " + "fraction of the training data." + ) + + if count is None or total is None or total_sq is None: + raise ValueError( + f"No samples with both {value_field!r} and {observation_mask_field!r} were available." + ) + + has_observed = count > 0 + denominator = count.clamp(min=1).to(torch.float64) + mean = torch.where(has_observed, total / denominator, torch.zeros_like(total)) + variance = torch.where( + has_observed, + (total_sq / denominator - mean.square()).clamp_min(0), + torch.ones_like(total), + ) + # A constant train feature is well-defined: it maps to zero. Unit std + # avoids division by zero and gives a finite, explicit OOD behaviour. + std = torch.where(variance > 0, variance.sqrt(), torch.ones_like(variance)) + return cls( + mean.to(torch.float32), std.to(torch.float32), count, + fit_scope=fit_scope, + ) + + @staticmethod + def _scope_digest(fit_scope: str | bytes | None) -> torch.Tensor: + """Return a stable, non-reversible identifier for the fitting split.""" + if fit_scope is None: + return torch.zeros(32, dtype=torch.uint8) + payload = fit_scope.encode("utf-8") if isinstance(fit_scope, str) else fit_scope + return torch.tensor(list(hashlib.sha256(payload).digest()), dtype=torch.uint8) + + @staticmethod + def _value_tensor(value: Any) -> torch.Tensor: + """Extract the ``value`` component from processor output or raw tuples.""" + if isinstance(value, dict): + value = value["value"] + elif isinstance(value, tuple): + # StageNet temporal processors return ``(time, value)``. + value = value[1] + return torch.as_tensor(value) + + def forward( + self, values: torch.Tensor, observed_mask: torch.Tensor + ) -> torch.Tensor: + """Standardise observed values and map missing/unfittable values to zero. + + We intentionally do not clip. The pipeline has no universally valid + physiological range for these MIMIC category aggregates; clipping would + silently overwrite potentially meaningful values. Values far outside + train support are therefore represented by a large (but finite) z-score + and remain auditable. + """ + if values.shape[-1] != self.feature_dim: + raise ValueError( + f"Expected {self.feature_dim} lab features, got {values.shape[-1]}." + ) + if observed_mask.shape != values.shape: + raise ValueError( + "Laboratory observation mask must have exactly the values shape; " + f"got {tuple(observed_mask.shape)} for {tuple(values.shape)}." + ) + values = values.to(dtype=self.mean.dtype) + observed = observed_mask.bool() & torch.isfinite(values) + fitted = self.observed_count > 0 + z = (values - self.mean) / self.std + return torch.where(observed & fitted, z, torch.zeros_like(z)) + + +def fit_lab_standardizer( + train_dataset: Iterable[dict[str, Any]], + *, + value_field: str = "labs", + observation_mask_field: Optional[str] = None, + fit_scope: str | bytes | None = None, +) -> LabStandardizer: + """Explicit train-dataset entry point used by downstream experiment scripts.""" + return LabStandardizer.fit( + train_dataset, + value_field=value_field, + observation_mask_field=observation_mask_field, + fit_scope=fit_scope, + ) + + +def lab_standardizer_fit_scope(dataset: Any, *, value_field: str = "labs") -> str: + """Fingerprint the patient/record membership of a split for safe transfer. + + ``SampleDataset.subset`` preserves these maps, so this checks the actual + split membership without iterating protected test data or serialising any + patient identifier into a checkpoint. A dataset without this provenance is + refused by the experiment scripts rather than being treated as safe by + default. + """ + patients = getattr(dataset, "patient_to_index", None) or {} + records = getattr(dataset, "record_to_index", None) or {} + indices = _provenance_indices(dataset) + if indices is None: + raise ValueError( + "Cannot bind lab standardisation to a train split: dataset has no " + "region_of_interest to drive an unsharded fit." + ) + payload = { + "value_field": value_field, + # ``len(dataset)`` is sharded by WORLD_SIZE on a StreamingDataset, so a + # digest taken under torchrun could never match the single-process one. + "n_samples": len(indices), + "patients": sorted(str(patient_id) for patient_id in patients), + "records": sorted(str(record_id) for record_id in records), + } + return json.dumps(payload, sort_keys=True, separators=(",", ":")) diff --git a/pyhealth/processors/stagenet_processor.py b/pyhealth/processors/stagenet_processor.py index 604376ec1..557fc4a7b 100644 --- a/pyhealth/processors/stagenet_processor.py +++ b/pyhealth/processors/stagenet_processor.py @@ -86,7 +86,7 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: if len(first_elem) > 0 and isinstance(first_elem[0], str): # Case 2: [["A", "B"], ["C"], ...] self._is_nested = True - break + break # Build vocabulary for codes and find max nested length max_inner_len = 0 @@ -178,9 +178,9 @@ def process( def _encode_codes(self, codes: List[str]) -> torch.Tensor: """Encode flat code list to indices.""" - # Handle empty code list - return single padding token + # Handle empty code list — zero events, not a fake pad token. if len(codes) == 0: - return torch.tensor([self.code_vocab[""]], dtype=torch.long) + return torch.zeros((0,), dtype=torch.long) indices = [] for code in codes: @@ -198,10 +198,9 @@ def _encode_nested_codes(self, nested_codes: List[List[str]]) -> torch.Tensor: assert self._max_nested_len is not None, "Max nested length must be set during fit()" # Handle empty nested codes (no visits/events) - # Return single padding token with shape (1, max_len) if len(nested_codes) == 0: - pad_token = self.code_vocab[""] - return torch.tensor([[pad_token] * self._max_nested_len], dtype=torch.long) + max_len = self._max_nested_len if self._max_nested_len is not None else 1 + return torch.zeros((0, max_len), dtype=torch.long) encoded_sequences = [] # Use global max length determined during fit @@ -345,9 +344,10 @@ class StageNetTensorProcessor(TemporalFeatureProcessor): >>> time.shape # (3,) """ - def __init__(self): + def __init__(self, forward_fill: bool = True): self._size = None # Feature dimension (set during fit) self._is_nested = None + self.forward_fill = forward_fill def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: """Determine input structure. @@ -370,13 +370,14 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: # Flat numeric: [1.5, 2.0, ...] self._is_nested = False self._size = 1 + break elif isinstance(first_elem, list): if len(first_elem) > 0: if isinstance(first_elem[0], (int, float)): # Nested numerics: [[1.0, 2.0], [3.0, 4.0]] self._is_nested = True self._size = len(first_elem) - break + break def process( self, value: Tuple[Optional[List], List] @@ -395,14 +396,31 @@ def process( """ # Unpack tuple: (time, values) time_data, value_data = value + value_data = list(value_data or []) - # Convert to numpy for easier imputation handling import numpy as np + if len(value_data) == 0: + n_feat = self._size if self._size is not None else 1 + nested = True if self._is_nested is None else self._is_nested + if nested: + value_tensor = torch.zeros((0, n_feat), dtype=torch.float) + else: + value_tensor = torch.zeros((0,), dtype=torch.float) + time_tensor = ( + torch.zeros((0,), dtype=torch.float) if time_data is not None else None + ) + return time_tensor, value_tensor + + # Convert to numpy for easier imputation handling value_array = np.array(value_data, dtype=float) - # Apply forward-fill imputation - if value_array.ndim == 1: + # Observation-mask fields must preserve false after a true observation; + # forward-filling a 0/1 mask would turn later missing labs into observed + # values. Ordinary numeric time-series retain the historical behaviour. + if not self.forward_fill: + value_array = np.nan_to_num(value_array, nan=0.0, posinf=0.0, neginf=0.0) + elif value_array.ndim == 1: # Flat numeric: [1.5, 2.0, nan, 3.0, ...] last_value = 0.0 for i in range(len(value_array)): diff --git a/pyhealth/processors/time_image_processor.py b/pyhealth/processors/time_image_processor.py index 9d313e6bc..449fff729 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. @@ -264,7 +279,6 @@ def process( Raises: ValueError: If image_paths and time_diffs have different lengths. - ValueError: If image_paths is empty. FileNotFoundError: If any image file does not exist. """ image_paths, time_diffs = value @@ -276,17 +290,24 @@ def process( f"match." ) if len(image_paths) == 0: - raise ValueError("image_paths must be non-empty.") + 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 + images = torch.zeros( + (0, c, self.image_size, self.image_size), dtype=torch.float32 + ) + timestamps = torch.zeros((0,), dtype=torch.float32) + return images, timestamps, "image" - 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 +316,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 +363,6 @@ def __repr__(self) -> str: f"mean={self.mean}, " f"std={self.std}, " f"mode={self.mode}, " - f"max_images={self.max_images})" + f"max_images={self.max_images}, " + f"padding={self.padding!r})" ) \ No newline at end of file diff --git a/pyhealth/processors/tuple_time_text_processor.py b/pyhealth/processors/tuple_time_text_processor.py index bbe74c4e6..7f1fe6a6f 100644 --- a/pyhealth/processors/tuple_time_text_processor.py +++ b/pyhealth/processors/tuple_time_text_processor.py @@ -21,7 +21,7 @@ def __init__( self, type_tag: str = "note", tokenizer_model: Optional[str] = None, - max_length: int = 128, + max_length: int = 512, padding: bool = True, truncation: bool = True, ): @@ -31,8 +31,9 @@ def __init__( type_tag: Modality identifier for automatic routing. Default: "note" tokenizer_model: Name or path of the HuggingFace tokenizer to use. If None, texts are returned as raw strings. Default: None - max_length: Maximum sequence length for tokenization. Default: 128 - padding: Whether to pad sequences to max_length. Default: True + max_length: Maximum sequence length for tokenization. Default: 512 + padding: Whether to pad sequences to the longest note in the sample. + Default: True truncation: Whether to truncate sequences to max_length. Default: True """ super().__init__() @@ -81,28 +82,65 @@ def process(self, value: Tuple[List[str], List[float]]) -> Union[Tuple[List[str] - str: Type tag """ texts, time_diffs = value + texts = list(texts or []) + time_diffs = list(time_diffs or []) + + # Keep text/time aligned and filter malformed text entries. + pair_count = min(len(texts), len(time_diffs)) + cleaned_texts: List[str] = [] + cleaned_times: List[float] = [] + for i in range(pair_count): + raw_text = texts[i] + raw_time = time_diffs[i] + + # Normalize text; skip null/whitespace-only entries. + if raw_text is None: + continue + text = str(raw_text).strip() + if text == "": + continue + + # Best-effort float normalization; skip unparseable timestamps. + try: + t = float(raw_time) + except (TypeError, ValueError): + continue + + cleaned_texts.append(text) + cleaned_times.append(t) + + texts = cleaned_texts + time_diffs = cleaned_times time_tensor = torch.tensor(time_diffs, dtype=torch.float32) if self.tokenizer is not None: - # Tokenize the list of texts + # Fast tokenizers crash on tokenizer([]). Build empty tensors + # ourselves so a patient with no notes is zero events, not a + # fake "[MISSING_TEXT]" row whose BERT embedding is a constant + # the classifier can use as a mortality feature. + if len(texts) == 0: + empty = torch.zeros((0, 1), dtype=torch.long) + return empty, empty.clone(), empty.clone(), time_tensor, self.type_tag encoded = self.tokenizer( texts, - padding="max_length" if self.padding else False, + padding="longest" if self.padding else False, truncation=self.truncation, max_length=self.max_length, return_tensors="pt" ) - input_ids = encoded["input_ids"] - attention_mask = encoded["attention_mask"] + input_ids = encoded.get("input_ids") + if input_ids is None: + raise ValueError("Tokenizer output is missing required `input_ids`.") + + attention_mask = encoded.get("attention_mask") + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) # Not all tokenizers return token_type_ids (e.g. RoBERTa might not, BERT does) - if "token_type_ids" in encoded: - token_type_ids = encoded["token_type_ids"] - else: - # meaningful text usually 0, padding 0? BERT uses 0 for sent A. - # If not provided, we can just use zeros or omit. - # For consistency with schema, let's provide zeros if expected. + token_type_ids = encoded.get("token_type_ids") + if token_type_ids is None: + # Some tokenizers do not return token_type_ids. token_type_ids = torch.zeros_like(input_ids) return input_ids, attention_mask, token_type_ids, time_tensor, self.type_tag @@ -172,4 +210,4 @@ def process_temporal(self, value) -> dict: def __repr__(self): if self.tokenizer_model: return f"TupleTimeTextProcessor(type_tag='{self.type_tag}', tokenizer='{self.tokenizer_model}')" - return f"TupleTimeTextProcessor(type_tag='{self.type_tag}')" + return f"TupleTimeTextProcessor(type_tag='{self.type_tag}')" \ No newline at end of file diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py new file mode 100644 index 000000000..ff9022cbe --- /dev/null +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -0,0 +1,922 @@ +import logging +import re +from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional, Union, Tuple, ClassVar + +from pyhealth.tasks.base_task import BaseTask + +logger = logging.getLogger(__name__) + + +class BaseMultimodalMIMIC4Task(BaseTask): + """Base class for multimodal MIMIC-IV tasks. + + Provides shared constants and utility methods used across all multimodal + task variants (notes, ICD codes, lab values). + """ + + MISSING_TEXT_TOKEN: ClassVar[str] = "" + MISSING_CODE_TOKEN: ClassVar[str] = "" + MISSING_FLOAT_TOKEN: ClassVar[float] = 0.0 + + LAB_CATEGORIES: ClassVar[Dict[str, List[str]]] = { + "Sodium": ["50824", "52455", "50983", "52623"], + "Potassium": ["50822", "52452", "50971", "52610"], + "Chloride": ["50806", "52434", "50902", "52535"], + "Bicarbonate": ["50803", "50804"], + "Glucose": ["50809", "52027", "50931", "52569"], + "Calcium": ["50808", "51624"], + "Magnesium": ["50960"], + "Anion Gap": ["50868", "52500"], + "Osmolality": ["52031", "50964", "51701"], + "Phosphate": ["50970"], + } + + LAB_CATEGORY_NAMES: ClassVar[List[str]] = [ + "Sodium", + "Potassium", + "Chloride", + "Bicarbonate", + "Glucose", + "Calcium", + "Magnesium", + "Anion Gap", + "Osmolality", + "Phosphate", + ] + + LABITEMS: ClassVar[List[str]] = [ + item for itemids in LAB_CATEGORIES.values() for item in itemids + ] + + RADIOLOGY_CLINICAL_HEADERS: ClassVar[List[str]] = [ + "indication", + "impression", + # "findings", + # "clinical history", + # "history", + # "comparison", + # "technique", + # "conclusion", + # "summary" + ] + + DISCHARGE_CLINICAL_HEADERS: ClassVar[List[str]] = [ + "chief complaint", + # "history of present illness", + # "hpi", + # "past medical history", + # "past medical and surgical history", + # "past medical/surgical history", + # "past surgical history", + # "medications on admission", + # "admission medications", + # "home medications", + # "social history", + # "family history", + # "allergies", + # "review of systems", + ] + + def __init__( + self, + window_hours: Optional[float] = None, + ): + self.window_hours = window_hours + # Task cache key is uuid5 over {**vars(task), schemas}. Bump when + # emitted data changes so leaky caches cannot be reused. + # 1: empty events instead of placeholders; per-admission collection span. + # 2: CXR/notes_labs_cxr no longer drop later stays against the first + # admission's clock (admission_time >= first_admit + window_hours). + # 3: event times are hours from the first stay in the sample, not + # reset per admission (reset times made stay 2 at +6h sort with + # stay 1 at +6h). + # 4: class/runner default is full stay (window_hours=None); + # admission-context discharge sections stamped at admit, not + # charttime. v3 caches used a 24h class default and discharge + # charttime on those notes. + self.emitted_data_version = 4 + + @staticmethod + def _clean_text(text: Optional[str]) -> Optional[str]: + """Return text if non-empty, otherwise None.""" + return text if text else None + + @staticmethod + def _parse_note_sections(text: str, note_type: str) -> Dict[str, str]: + """Split a note into {lowercased_header: content_text} pairs.""" + ext_text = text + '\n\n' + if note_type == "radiology": + section_re = re.compile(r'([a-zA-Z ]+):[ \t\n]+(.+?)\n{2,}', re.DOTALL) + elif note_type == "discharge": + section_re = re.compile(r'([a-zA-Z ]+):\n+(.+?)\n{2,}', re.DOTALL) + else: + raise ValueError(f"Note Type '{note_type}' not supported.") + return { + m.group(1).strip().lower(): m.group(2).strip() + for m in section_re.finditer(ext_text) + if m.end() - m.start() > 0 + } + + @staticmethod + def _parse_datetime(value: Any) -> Optional[datetime]: + if isinstance(value, datetime): + return value + if isinstance(value, str): + for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d"): + try: + return datetime.strptime(value, fmt) + except ValueError: + continue + return None + + @staticmethod + def _to_hours(delta_seconds: float) -> float: + return delta_seconds / 3600.0 + + @classmethod + def _hours_since(cls, timestamp: datetime, origin: datetime) -> float: + """Hours from ``origin`` to ``timestamp``. + + Collection windows stay per admission (full stay, or admit+window + when ``window_hours`` is set). The value written onto the unified + timeline is hours from the first stay in this sample, so a later + stay at +6h does not sort with the first stay at +6h. + """ + return cls._to_hours((timestamp - origin).total_seconds()) + + def _compute_effective_window( + self, + admissions_to_process: List[Any], + ) -> Tuple[datetime, Optional[datetime]]: + """Compute effective start/end from the global span of processed admissions. + + Returns: + Tuple of (effective_start, effective_end). + """ + global_start = admissions_to_process[0].timestamp + global_end: Optional[datetime] = None + + for a in admissions_to_process: + dt = self._parse_datetime(getattr(a, "dischtime", None)) + if dt is not None and (global_end is None or dt > global_end): + global_end = dt + + if self.window_hours is not None: + effective_start = global_start + effective_end = effective_start + timedelta(hours=self.window_hours) + return effective_start, effective_end + + effective_start = global_start + effective_end = global_end + + return effective_start, effective_end + + def _admission_window_end( + self, + admission_time: datetime, + admission_dischtime: datetime, + ) -> datetime: + """End of the observation window for one admission. + + Callers previously passed ``admission_dischtime`` directly, so + ``window_hours`` was inert and labs were collected through discharge. + For a mortality label that reads the outcome. Re-anchor per admission + and clamp to discharge so a later stay cannot inherit the first + admission's window. + """ + if self.window_hours is None: + return admission_dischtime + end = admission_time + timedelta(hours=self.window_hours) + return min(end, admission_dischtime) if admission_dischtime else end + + def _build_admissions_to_process(self, patient: Any) -> Tuple[List[Any], int]: + """Build admissions to process and derive mortality label. + + Includes all admissions up to and including the first death admission. + Patients who die in their first (and only) admission are included as + positives — previously they were dropped, which silently discarded most + ICU mortality positives and collapsed positive rate from ~10% to ~2.7%. + This now matches stagenet's semantics: use all available admission data + and label as positive if any admission has hospital_expire_flag=1. + """ + admissions = patient.get_events(event_type="admissions") + if len(admissions) == 0: + return [], 0 + + admissions_to_process: List[Any] = [] + mortality_label = 0 + + for admission in admissions: + admissions_to_process.append(admission) + if admission.hospital_expire_flag in [1, "1"]: + mortality_label = 1 + break + + return admissions_to_process, mortality_label + + def _collect_icd_codes(self, patient: Any, hadm_id: Any) -> List[str]: + """Collect ICD diagnosis and procedure codes for one admission. + + Returns: + List of ICD code strings, or an empty list if none found. + """ + diagnoses_icd = patient.get_events( + event_type="diagnoses_icd", filters=[("hadm_id", "==", hadm_id)] + ) + procedures_icd = patient.get_events( + event_type="procedures_icd", filters=[("hadm_id", "==", hadm_id)] + ) + return [ + e.icd_code for e in diagnoses_icd if hasattr(e, "icd_code") and e.icd_code + ] + [ + e.icd_code for e in procedures_icd if hasattr(e, "icd_code") and e.icd_code + ] + + def _collect_labs( + self, + patient: Any, + admission_time: datetime, + end_time: datetime, + time_origin: datetime, + ) -> Tuple[List[float], List[List[float]], List[List[bool]]]: + """Collect lab values and observation masks for one admission. + + Args: + patient: Patient object. + admission_time: Start of this stay's collection window. + end_time: End of the window (inclusive). + time_origin: First stay in this sample. Event times are hours + from here, not from ``admission_time``. + + Returns: + Tuple of (lab_times, lab_values, lab_masks). ``lab_masks`` is a + parallel boolean tensor where ``True`` means observed and ``False`` + means imputed with 0.0. Returns empty lists when no valid lab + events are found; do not invent a placeholder row. + """ + try: + import polars as pl + except ImportError as exc: + raise ImportError("Polars is required for lab collection.") from exc + + labevents_df = patient.get_events( + event_type="labevents", + start=admission_time, + end=end_time, + return_df=True, + ) + + lab_times: List[float] = [] + lab_values: List[List[float]] = [] + lab_masks: List[List[bool]] = [] + + labevents_df = labevents_df.filter( + pl.col("labevents/itemid").is_in(self.LABITEMS) + ) + if labevents_df.height > 0: + labevents_df = labevents_df.with_columns( + pl.col("labevents/storetime").str.strptime( + pl.Datetime, "%Y-%m-%d %H:%M:%S" + ) + ) + labevents_df = labevents_df.filter( + pl.col("labevents/storetime") <= end_time + ) + if labevents_df.height > 0: + labevents_df = labevents_df.select( + pl.col("timestamp"), + pl.col("labevents/itemid"), + pl.col("labevents/valuenum").cast(pl.Float64), + ) + for lab_ts in sorted(labevents_df["timestamp"].unique().to_list()): + ts_labs = labevents_df.filter(pl.col("timestamp") == lab_ts) + lab_vector: List[float] = [] + lab_mask: List[bool] = [] + for category_name in self.LAB_CATEGORY_NAMES: + category_value = self.MISSING_FLOAT_TOKEN + observed = False + for itemid in self.LAB_CATEGORIES[category_name]: + matching = ts_labs.filter( + pl.col("labevents/itemid") == itemid + ) + if matching.height > 0: + category_value = matching["labevents/valuenum"][0] + observed = True + break + lab_vector.append(category_value) + lab_mask.append(observed) + lab_times.append(self._hours_since(lab_ts, time_origin)) + lab_values.append(lab_vector) + lab_masks.append(lab_mask) + return lab_times, lab_values, lab_masks + + def _collect_notes( + self, + patient: Any, + note_event_type: str, + hadm_id: Any, + admission_time: datetime, + start_time: Optional[datetime] = None, + end_time: Optional[datetime] = None, + section_headers: Optional[List[str]] = None, + fallback_to_full_note: bool = False, + time_origin: Optional[datetime] = None, + event_time: Optional[datetime] = None, + ) -> Tuple[List[str], List[float]]: + """Collect notes of a given type for one admission. + + Args: + patient: Patient object. + note_event_type: Event type string (e.g. "discharge", "radiology"). + hadm_id: Admission ID to filter by. + admission_time: This stay's admit time. Used as the timeline + origin when ``time_origin`` is omitted. + start_time: Optional start of the time window. + end_time: Optional end of the time window. + section_headers: When provided, extract only these named sections + from each note (lowercased match against parsed headers). + fallback_to_full_note: When True, falls back to the full note text + if no matching sections are found. When False, notes with no + matching sections are dropped entirely. + time_origin: Timeline origin. Hours are measured from here. + event_time: If set, stamp every collected note at this instant. + Discharge-section text is admission-context and should pass + this stay's admit. Radiology omits it and uses exam + ``charttime``. + + Returns: + Tuple of (texts, hours from the sample's first stay). Empty lists + when the events list is empty; do not invent a placeholder note. + """ + notes = patient.get_events( + event_type=note_event_type, + start=start_time, + end=end_time, + filters=[("hadm_id", "==", hadm_id)], + ) + + texts: List[str] = [] + note_times: List[float] = [] + for note in notes: + try: + note_text = self._clean_text(note.text) + if note_text: + if section_headers is not None: + parsed = self._parse_note_sections(note_text, note_type=note_event_type) + extracted = [f"{k}: {v}" for k, v in parsed.items() if k in section_headers and v] + if extracted: + note_text = " [SEP] ".join(extracted) + elif not fallback_to_full_note: + continue + + origin = time_origin if time_origin is not None else admission_time + stamp = event_time if event_time is not None else note.timestamp + texts.append(note_text) + note_times.append(self._hours_since(stamp, origin)) + except ( + AttributeError + ): # note object is missing .text or .timestamp attribute (e.g. malformed note) + pass + + return texts, note_times + + +class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): + """Mortality prediction from admission-context notes and lab values. + + Follows the approach of Lee et al. (2023): use text that is clinically + available *at admission* rather than discharge notes. ICD codes are + excluded by default but can be re-enabled for ablation experiments. + + Text is extracted from the MIMIC-IV discharge note by parsing the Chief + Complaint, History of Present Illness, Past Medical History, and Medications + on Admission sections — all of which describe the patient's state at the + start of the stay. The extracted text is stamped at that stay's admit. + + Radiology reports are also included, parsed for their Indication and + Impression sections and bounded to the same observation window as labs + (rather than that stay's admit), since — unlike the discharge summary — they + are written at exam time and describe findings from later in the stay. + + Fields: + admission_note_times: Admission-context discharge-note text stamped + at that stay's admit (hours from the first stay), plus in-window + radiology note text at exam time. + labs: 10-dim lab vectors at each measurement timestamp. + labs_mask: Boolean observation mask parallel to ``labs``. + icd_codes: (only when ``include_icd=True``) Diagnosis + procedure codes + per admission with inter-admission time offsets. + + Args: + window_hours: Hours from admission for lab collection. ``None`` + collects for the full admission span. Default: ``None``. + include_icd: When ``True``, collect discharge-coded ICD codes and add + ``icd_codes`` to the sample dict / input schema. Default: ``False``. + MIMIC-IV timestamps those codes at ``dischtime``, so this leaks + the in-hospital mortality label. Keep it off except as an + explicit leaky ablation. + """ + + PADDING: int = 0 + + task_name: str = "NotesLabsMIMIC4" + + _BASE_INPUT_SCHEMA: ClassVar[Dict] = { + "admission_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", + "type_tag": "note", + "max_length": 512, + }, + ), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {"forward_fill": False}), + } + + input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = _BASE_INPUT_SCHEMA + output_schema: Dict[str, str] = {"mortality": "binary"} + + def __init__( + self, + window_hours: Optional[float] = None, + include_icd: bool = False, + ) -> None: + super().__init__(window_hours=window_hours) + self.include_icd = include_icd + schema = dict(self._BASE_INPUT_SCHEMA) + if include_icd: + schema["icd_codes"] = ("stagenet", {"padding": self.PADDING}) + self.input_schema = schema + logger.info( + "NotesLabsMIMIC4: filtering discharge notes to sections: %s; " + "radiology notes to sections: %s", + self.DISCHARGE_CLINICAL_HEADERS, + self.RADIOLOGY_CLINICAL_HEADERS, + ) + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + if not patient.get_events(event_type="patients"): + return [] + + 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 + ) + time_origin = admissions_to_process[0].timestamp + + all_note_texts: List[str] = [] + all_note_times: List[float] = [] + all_lab_values: List[List[float]] = [] + all_lab_masks: List[List[bool]] = [] + all_lab_times: List[float] = [] + all_icd_codes: List[List[str]] = [] + all_icd_times: List[float] = [] + + 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 + + note_texts, note_times = self._collect_notes( + patient, + "discharge", + admission.hadm_id, + admission_time, + section_headers=self.DISCHARGE_CLINICAL_HEADERS, + time_origin=time_origin, + event_time=admission_time, + ) + all_note_texts.extend(note_texts) + all_note_times.extend(note_times) + + # Labs within the observation window of THIS admission. + lab_end = self._admission_window_end(admission_time, admission_dischtime) + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=lab_end, + time_origin=time_origin, + ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + # Radiology notes within the observation window. Unlike the + # discharge note (a retrospective summary parsed for its + # admission-context sections), radiology reports are written at + # exam time, so they're bounded to the same window as labs + # to avoid pulling in findings from later in the stay. + radiology_texts, radiology_times = self._collect_notes( + patient, + "radiology", + admission.hadm_id, + admission_time, + start_time=admission_time, + end_time=lab_end, + section_headers=self.RADIOLOGY_CLINICAL_HEADERS, + time_origin=time_origin, + ) + all_note_texts.extend(radiology_texts) + all_note_times.extend(radiology_times) + + if self.include_icd: + visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) + if visit_icd_codes: + all_icd_codes.append(visit_icd_codes) + all_icd_times.append( + self._hours_since(admission_time, time_origin) + ) + + record: Dict[str, Any] = { + "patient_id": patient.patient_id, + "admission_note_times": (all_note_texts, all_note_times), + "labs": (all_lab_times, all_lab_values), + "labs_mask": (all_lab_times, all_lab_masks), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } + + if self.include_icd: + record["icd_codes"] = (all_icd_times, all_icd_codes) + + return [record] + + +class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): + """Mortality prediction from admission-context notes, labs, and CXR. + + Extends ``NotesLabsMIMIC4`` with chest X-ray images: the same + admission-context discharge-note sections, in-window radiology reports, + and labs, plus CXR studies (StudyDate+StudyTime from the ``metadata`` + event table) bounded to the same observation window as labs/radiology. + ICD codes are excluded by default but can be re-enabled for ablation + experiments, same as ``NotesLabsMIMIC4``. + + Fields: + admission_note_times: Admission-context discharge-note text stamped + at that stay's admit (hours from the first stay), plus in-window + radiology note text at exam time. + labs: 10-dim lab vectors at each measurement timestamp. + labs_mask: Boolean observation mask parallel to ``labs``. + cxr_image_times: In-window CXR image paths at their exam-relative + timestamp. + icd_codes: (only when ``include_icd=True``) Diagnosis + procedure codes + per admission with inter-admission time offsets. + + Args: + window_hours: Hours from admission for lab/CXR collection. + ``None`` collects for the full admission span. Default: ``None``. + include_icd: When ``True``, collect discharge-coded ICD codes and add + ``icd_codes`` to the sample dict / input schema. Default: ``False``. + MIMIC-IV timestamps those codes at ``dischtime``, so this leaks + the in-hospital mortality label. Keep it off except as an + explicit leaky ablation. + """ + + PADDING: int = 0 + + task_name: str = "NotesLabsCXRMIMIC4" + + _BASE_INPUT_SCHEMA: ClassVar[Dict] = { + "admission_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", + "type_tag": "note", + "max_length": 512, + }, + ), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {"forward_fill": False}), + "cxr_image_times": ( + "time_image", + { + "image_size": 224, + "mode": "RGB", + "padding": "", + }, + ), + } + + input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = _BASE_INPUT_SCHEMA + output_schema: Dict[str, str] = {"mortality": "binary"} + + def __init__( + self, + window_hours: Optional[float] = None, + include_icd: bool = False, + ) -> None: + super().__init__(window_hours=window_hours) + self.include_icd = include_icd + schema = dict(self._BASE_INPUT_SCHEMA) + if include_icd: + schema["icd_codes"] = ("stagenet", {"padding": self.PADDING}) + self.input_schema = schema + logger.info( + "NotesLabsCXRMIMIC4: filtering discharge notes to sections: %s; " + "radiology notes to sections: %s", + self.DISCHARGE_CLINICAL_HEADERS, + self.RADIOLOGY_CLINICAL_HEADERS, + ) + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + if not patient.get_events(event_type="patients"): + return [] + + 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 + ) + time_origin = admissions_to_process[0].timestamp + + all_note_texts: List[str] = [] + all_note_times: List[float] = [] + all_lab_values: List[List[float]] = [] + all_lab_masks: List[List[bool]] = [] + all_lab_times: List[float] = [] + all_icd_codes: List[List[str]] = [] + all_icd_times: List[float] = [] + all_cxr_paths: List[str] = [] + all_cxr_times: List[float] = [] + + 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 + + note_texts, note_times = self._collect_notes( + patient, + "discharge", + admission.hadm_id, + admission_time, + section_headers=self.DISCHARGE_CLINICAL_HEADERS, + time_origin=time_origin, + event_time=admission_time, + ) + all_note_texts.extend(note_texts) + all_note_times.extend(note_times) + + # Labs within the observation window of THIS admission. + lab_end = self._admission_window_end(admission_time, admission_dischtime) + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=lab_end, + time_origin=time_origin, + ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + # Radiology notes within the observation window. Unlike the + # discharge note (a retrospective summary parsed for its + # admission-context sections), radiology reports are written at + # exam time, so they're bounded to the same window as labs/CXR + # to avoid pulling in findings from later in the stay. + radiology_texts, radiology_times = self._collect_notes( + patient, + "radiology", + admission.hadm_id, + admission_time, + start_time=admission_time, + end_time=lab_end, + section_headers=self.RADIOLOGY_CLINICAL_HEADERS, + time_origin=time_origin, + ) + all_note_texts.extend(radiology_texts) + all_note_times.extend(radiology_times) + + # CXR studies within the same observation window as labs/radiology. + # CXR metadata is filtered by timestamp; this includes StudyTime. + metadata_events = patient.get_events( + event_type="metadata", + start=admission_time, + end=lab_end, + ) + for event in metadata_events: + try: + if event.image_path: + all_cxr_paths.append(event.image_path) + all_cxr_times.append( + self._hours_since(event.timestamp, time_origin) + ) + except AttributeError: + continue + + if self.include_icd: + visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) + if visit_icd_codes: + all_icd_codes.append(visit_icd_codes) + all_icd_times.append( + self._hours_since(admission_time, time_origin) + ) + + record: Dict[str, Any] = { + "patient_id": patient.patient_id, + "admission_note_times": (all_note_texts, all_note_times), + "labs": (all_lab_times, all_lab_values), + "labs_mask": (all_lab_times, all_lab_masks), + "cxr_image_times": (all_cxr_paths, all_cxr_times), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } + + if self.include_icd: + 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: ``None``. + """ + + PADDING: int = 0 + + task_name: str = "LabsMIMIC4" + + input_schema: ClassVar[Dict] = { + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {"forward_fill": False}), + } + output_schema: ClassVar[Dict] = {"mortality": "binary"} + + def __init__(self, window_hours: Optional[float] = None) -> None: + super().__init__(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 + ) + time_origin = admissions_to_process[0].timestamp + + 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=self._admission_window_end( + admission_time, admission_dischtime + ), + time_origin=time_origin, + ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + 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] + + +class CXRMIMIC4(BaseMultimodalMIMIC4Task): + """CXR-only mortality prediction using chest X-ray images. + + Serves as the imaging-only reference baseline for multimodal ablations — + no notes, no ICD codes, no labs — isolating the chest X-ray + modality the same way ``LabsMIMIC4`` isolates labs. + + CXR studies are filtered by timestamp (StudyDate+StudyTime, from the + ``metadata`` event table) within each admission's observation window. + + Args: + window_hours: Hours from admission to collect CXR studies. ``None`` + collects for the full admission span. Default: None. + """ + + task_name: str = "CXRMIMIC4" + + input_schema: ClassVar[Dict] = { + "cxr_image_times": ( + "time_image", + { + "image_size": 224, + "mode": "RGB", + "padding": "", + }, + ), + } + output_schema: ClassVar[Dict] = {"mortality": "binary"} + + 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 + ) + time_origin = admissions_to_process[0].timestamp + + all_cxr_paths: List[str] = [] + all_cxr_times: List[float] = [] + + 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 + + admission_end = self._admission_window_end( + admission_time, admission_dischtime + ) + + # CXR metadata is filtered by timestamp; this includes StudyTime. + metadata_events = patient.get_events( + event_type="metadata", + start=admission_time, + end=admission_end, + ) + for event in metadata_events: + try: + if event.image_path: + all_cxr_paths.append(event.image_path) + all_cxr_times.append( + self._hours_since(event.timestamp, time_origin) + ) + except AttributeError: + continue + + single_patient_longitudinal_record = { + "patient_id": patient.patient_id, + "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] \ No newline at end of file diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index bc6a28677..2085221b3 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -1,5 +1,7 @@ +import json import logging import os +import time from datetime import datetime from typing import Callable, Dict, List, Optional, Type @@ -18,6 +20,48 @@ logger = logging.getLogger(__name__) +_AMP_DTYPES = { + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp16": torch.float16, + "float16": torch.float16, +} + + +def resolve_amp_dtype(amp_dtype: str, use_amp: bool = False) -> torch.dtype: + """Validate the mixed-precision dtype and return it. + + The previous expression was ``bfloat16 if amp_dtype == "bf16" else float16``, + so every other spelling silently selected fp16 and, through GradScaler, + changed gradient behaviour too. + """ + name = str(amp_dtype).lower() + if name not in _AMP_DTYPES: + raise ValueError( + f"amp_dtype must be one of {sorted(_AMP_DTYPES)}, got {amp_dtype!r}." + ) + dtype = _AMP_DTYPES[name] + if ( + use_amp + and dtype == torch.bfloat16 + and torch.cuda.is_available() + and not torch.cuda.is_bf16_supported() + ): + raise RuntimeError( + "bf16 mixed precision was requested, but this CUDA device does not " + "support bf16." + ) + return dtype + + +def autocast_device_type(device) -> str: + """Device type string for ``torch.autocast``, from the trainer device.""" + name = str(device).split(":", 1)[0].lower() + if name in ("cuda", "cpu", "mps"): + return name + return "cpu" + + def is_best(best_score: float, score: float, monitor_criterion: str) -> bool: if monitor_criterion == "max": return score > best_score @@ -37,6 +81,15 @@ def set_logger(log_path: str) -> None: return +def _vram_stats(device: str) -> Dict[str, float]: + """Returns current and peak VRAM usage in MB for a CUDA device.""" + if not torch.cuda.is_available() or not str(device).startswith("cuda"): + return {} + allocated = torch.cuda.memory_allocated(device) / 1024**2 + peak = torch.cuda.max_memory_allocated(device) / 1024**2 + return {"vram_allocated_mb": allocated, "vram_peak_mb": peak} + + def get_metrics_fn(mode: str) -> Callable: if mode == "binary": return binary_metrics_fn @@ -126,6 +179,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. @@ -145,10 +201,23 @@ 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 = resolve_amp_dtype(amp_dtype, use_amp=use_amp) + # 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}") @@ -161,6 +230,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()) @@ -184,50 +255,132 @@ def train( steps_per_epoch = len(train_dataloader) global_step = 0 patience_counter = 0 + metrics_history: List[Dict] = [] + train_start = time.perf_counter() + total_skipped_steps = 0 # epoch training loop - for epoch in range(epochs): + epoch_iterator = tqdm(range(epochs), desc="Epochs", unit="epoch") + for epoch in epoch_iterator: + epoch_iterator.set_postfix_str(f"{epoch + 1}/{epochs}", refresh=False) training_loss = [] + epoch_skipped_steps = 0 self.model.zero_grad() self.model.train() + if torch.cuda.is_available() and str(self.device).startswith("cuda"): + torch.cuda.reset_peak_memory_stats(self.device) + 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} / {epochs}", + desc=f"Epoch {epoch + 1}/{epochs}", smoothing=0.05, + leave=False, ): try: data = next(data_iterator) 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=autocast_device_type(self.device), + 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 + 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 scaler is not None: + scaler.unscale_(optimizer) + # Always compute the grad norm (even with no clipping + # configured) so non-finite gradients can be detected and + # skipped before they permanently poison the model with + # NaN weights. + grad_norm = torch.nn.utils.clip_grad_norm_( + self.model.parameters(), + max_grad_norm if max_grad_norm is not None else float("inf"), ) - # update - optimizer.step() - optimizer.zero_grad() - training_loss.append(loss.item()) - global_step += 1 + step_ok = bool(torch.isfinite(grad_norm)) + if not step_ok: + epoch_skipped_steps += 1 + total_skipped_steps += 1 + logger.warning( + f"epoch-{epoch} step-{global_step}: non-finite " + f"gradient norm ({grad_norm}); skipping optimizer " + f"step." + ) + if scaler is not None: + if step_ok: + scaler.step(optimizer) + scaler.update() + elif step_ok: + optimizer.step() + optimizer.zero_grad() + global_step += 1 + + epoch_time = time.perf_counter() - epoch_start + vram = _vram_stats(self.device) + + epochs_done = epoch + 1 + epochs_left = epochs - epochs_done + elapsed_total = time.perf_counter() - train_start + avg_epoch_time = elapsed_total / epochs_done + eta_s = avg_epoch_time * epochs_left + eta_h, eta_rem = divmod(int(eta_s), 3600) + eta_m = eta_rem // 60 + eta_str = f"{eta_h}h{eta_m:02d}m" + # log and save logger.info(f"--- Train epoch-{epoch}, step-{global_step} ---") logger.info(f"loss: {sum(training_loss) / len(training_loss):.4f}") + logger.info(f"epoch_time: {epoch_time:.2f}s elapsed: {elapsed_total:.0f}s ETA: {eta_str} ({epochs_done}/{epochs} epochs)") + print(f"[ETA] epoch {epochs_done}/{epochs} done in {epoch_time:.0f}s — ETA to finish: {eta_str}", flush=True) + if vram: + logger.info( + f"vram_peak: {vram['vram_peak_mb']:.1f} MB " + f"vram_current: {vram['vram_allocated_mb']:.1f} MB" + ) + if epoch_skipped_steps > 0: + logger.warning( + f"Skipped {epoch_skipped_steps} optimizer step(s) this " + f"epoch due to non-finite gradients (total so far: " + f"{total_skipped_steps})." + ) if self.exp_path is not None: self.save_ckpt(os.path.join(self.exp_path, "last.ckpt")) + epoch_record: Dict = { + "epoch": epoch, + "global_step": global_step, + "train_loss": sum(training_loss) / len(training_loss), + "epoch_time_s": round(epoch_time, 3), + "skipped_steps": epoch_skipped_steps, + **{f"train_{k}": v for k, v in vram.items()}, + } + # validation if val_dataloader is not None: scores = self.evaluate(val_dataloader) logger.info(f"--- Eval epoch-{epoch}, step-{global_step} ---") for key in scores.keys(): logger.info("{}: {:.4f}".format(key, scores[key])) + epoch_record.update({f"val_{k}": v for k, v in scores.items()}) # save best model if monitor is not None: score = scores[monitor] @@ -247,8 +400,21 @@ def train( logger.info( f"Early stopping at epoch-{epoch}, step-{global_step}" ) + metrics_history.append(epoch_record) break + metrics_history.append(epoch_record) + + total_time = time.perf_counter() - train_start + logger.info(f"--- Training complete: {total_time:.2f}s total ---") + + # persist metrics history + if self.exp_path is not None: + history_path = os.path.join(self.exp_path, "metrics_history.json") + with open(history_path, "w") as f: + json.dump(metrics_history, f, indent=2) + logger.info(f"Metrics history saved to {history_path}") + # load best model if load_best_model_at_last and self.exp_path is not None and os.path.isfile( os.path.join(self.exp_path, "best.ckpt")): @@ -262,7 +428,7 @@ def train( for key in scores.keys(): logger.info("{}: {:.4f}".format(key, scores[key])) - return + return metrics_history def inference(self, dataloader, additional_outputs=None, return_patient_ids=False) -> Dict[str, float]: @@ -430,4 +596,4 @@ def forward(self, x, y, **kwargs): monitor="accuracy", epochs=5, test_dataloader=val_dataloader, - ) + ) \ No newline at end of file diff --git a/pyhealth/utils.py b/pyhealth/utils.py index b4af8980a..eb14acfe1 100644 --- a/pyhealth/utils.py +++ b/pyhealth/utils.py @@ -1,7 +1,9 @@ +import hashlib import json import os import pickle import random +import subprocess import contextlib import numpy as np @@ -44,6 +46,82 @@ def save_json(data, filename): with open(filename, "w") as f: json.dump(data, f) + +def _git_revision(): + repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + try: + commit = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=repo, stderr=subprocess.DEVNULL + ).decode().strip() + dirty = bool(subprocess.check_output( + ["git", "status", "--porcelain"], cwd=repo, stderr=subprocess.DEVNULL + ).decode().strip()) + return {"commit": commit, "dirty": dirty} + except Exception: + return {"commit": None, "dirty": None} + + +def _source_digest(): + """Hash the package source so code identity survives a non-git deploy. + + Cluster runs typically execute from an unpacked tarball rather than a + clone, so the git lookup returns nothing exactly where provenance matters + most. Hashing the sources keeps "which code produced this result" + answerable either way. + """ + package = os.path.dirname(os.path.abspath(__file__)) + digest = hashlib.sha256() + try: + for root, dirs, files in os.walk(package): + dirs[:] = sorted(d for d in dirs if d != "__pycache__") + for name in sorted(files): + if not name.endswith(".py"): + continue + path = os.path.join(root, name) + digest.update(os.path.relpath(path, package).encode()) + with open(path, "rb") as f: + digest.update(f.read()) + return digest.hexdigest() + except OSError: + return None + + +def _jsonable(value): + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, (list, tuple)): + return [_jsonable(v) for v in value] + if isinstance(value, dict): + return {str(k): _jsonable(v) for k, v in value.items()} + return str(value) + + +def write_run_config(exp_path, config): + """Persist the resolved run configuration next to the run's metrics. + + ``metrics_history.json`` records what a run scored but not the conditions + that produced it. Record the resolved settings, not the raw flags, so + derived conditions (lr, split mode, eval split) are recoverable. + """ + record = { + "config": {str(k): _jsonable(v) for k, v in config.items()}, + "git": _git_revision(), + "source_sha256": _source_digest(), + "torch": torch.__version__, + } + os.makedirs(exp_path, exist_ok=True) + path = os.path.join(exp_path, "run_config.json") + tmp = f"{path}.tmp.{os.getpid()}" + try: + with open(tmp, "w") as f: + json.dump(record, f, indent=2, sort_keys=True) + os.replace(tmp, path) + finally: + if os.path.exists(tmp): + os.remove(tmp) + return path + + @contextlib.contextmanager def set_env(**environ): """ diff --git a/scripts/will/condor/labs_notes/lab_notes_rnn.sub b/scripts/will/condor/labs_notes/lab_notes_rnn.sub new file mode 100644 index 000000000..a5ede861e --- /dev/null +++ b/scripts/will/condor/labs_notes/lab_notes_rnn.sub @@ -0,0 +1,58 @@ +# HTCondor submission — labs+notes RNN mortality run +# +# Condor equivalent of scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py: +# same task (notes_labs), same model (rnn), same hyperparameters. Runs unattended +# instead of in a tmux session; Condor assigns the GPU (no manual +# nvidia-smi / CUDA_VISIBLE_DEVICES step needed). +# +# To submit (from the project root): +# mkdir -p /home/wp14/logs/condor +# condor_submit scripts/will/condor/labs_notes/labs_notes_rnn.sub +# +# Monitor: +# condor_q +# tail -f /home/wp14/logs/condor/labs_notes_rnn__0.out + +initialdir = /home/wp14/PyHealth +executable = /home/wp14/PyHealth/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh +transfer_executable = False +arguments = $(seed) +getenv = True + +environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \ + NOTE_ROOT=/shared/rsaas/physionet.org/files/mimic-note \ + CACHE_DIR=/shared/rsaas/wp14/pyhealth_cache_labs_notes \ + OUTPUT_DIR=/home/wp14/output \ + USE_WANDB=1 \ + WANDB_PROJECT=pyhealth-multimodal-labs-notes \ + WANDB_RUN_NAME=labs_notes_rnn_seed$(seed)" + +output = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).out +error = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).err +log = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).log + +stream_output = False +stream_error = False + +request_gpus = 1 +request_cpus = 4 +# Was 163840MB (160GB) — cgroup-killed job 10821 at 162747MB during the +# full-scale patient_id sort/shuffle (see run_labs_notes_rnn.sh comment). +# Bumped for headroom now that the distributed cluster (with disk-spilling) +# is back in play; c02 has ~1TB total and is otherwise idle. +request_memory = 400000MB +request_disk = 20GB + +# Previously hardcoded to sunlab-c01 (A100 80GB) because the previous run +# OOM'd on a 47GB card with ~46.8GB resident. sunlab-c01's condor_startd is +# currently down (master alive, STARTD_StartTime=0), so it never matches — +# the only machine in the pool is sunlab-c02 (8x RTX 6000 Ada, 48509MB each). +# Match any GPU with enough headroom and prefer the biggest; FREEZE_ENCODER=1 +# below (frozen Bio_ClinicalBERT text encoder, ~50% less VRAM for the text +# branch) is what actually keeps this under 48GB instead of the hostname pin. +Requirements = (TARGET.GPUs_GlobalMemoryMb >= 40000) +Rank = TARGET.GPUs_GlobalMemoryMb + +queue seed from ( + 12 +) \ No newline at end of file diff --git a/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh b/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh new file mode 100755 index 000000000..33661ff11 --- /dev/null +++ b/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +# HTCondor executable — labs+notes RNN mortality run. +# +# Condor equivalent of scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py: +# same task (notes_labs), same model (rnn), same hyperparameter defaults. GPU +# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor +# assigns the GPU via request_gpus / cgroups. +# +# usage: run_labs_notes_rnn.sh +# to remove logs: rm -rf logs/condor/* +# to remove cache: rm -rf /shared/rsaas/wp14/pyhealth_cache_labs_notes/* +set -euo pipefail + +SEED="${1:?usage: run_labs_notes_rnn.sh }" + +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}" +EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}" +NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}" +CACHE_DIR="${CACHE_DIR:-/shared/eng/wp14/pyhealth_cache_labs_notes}" +OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}" +CONDA_SH="${CONDA_SH:-}" + +DEV_MODE="${DEV_MODE:-1}" +EMBEDDING_DIM="${EMBEDDING_DIM:-128}" +HIDDEN_DIM="${HIDDEN_DIM:-128}" +RNN_TYPE="${RNN_TYPE:-GRU}" +RNN_LAYERS="${RNN_LAYERS:-2}" +DROPOUT="${DROPOUT:-0.1}" +EPOCHS="${EPOCHS:-15}" +BATCH_SIZE="${BATCH_SIZE:-32}" +LR="${LR:-1e-3}" +WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}" +PATIENCE="${PATIENCE:-5}" +NUM_WORKERS="${NUM_WORKERS:-4}" +FREEZE_ENCODER="${FREEZE_ENCODER:-1}" +INCLUDE_VITALS="${INCLUDE_VITALS:-0}" +USE_AMP="${USE_AMP:-0}" +AMP_DTYPE="${AMP_DTYPE:-bf16}" + +# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that +# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup. +# Disabling the whole distributed cluster (PYHEALTH_DISABLE_DASK_DISTRIBUTED=1) +# also throws away its disk-spilling/memory limits, which OOM'd the full-scale +# patient_id sort during event-dataframe caching (job 10821, 160GB cgroup +# limit hit). Scope the fix to just the NVML probe instead, keeping the real +# distributed cluster (with spilling) for the sort. +export DASK_DISTRIBUTED__DIAGNOSTICS__NVML="${DASK_DISTRIBUTED__DIAGNOSTICS__NVML:-0}" + +USE_WANDB="${USE_WANDB:-0}" +WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-notes}" +WANDB_RUN_NAME="${WANDB_RUN_NAME:-}" + +resolve_conda_sh() { + if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then + echo "${CONDA_SH}" + return 0 + fi + if command -v conda >/dev/null 2>&1; then + local base + base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then + echo "${base}/etc/profile.d/conda.sh" + return 0 + fi + fi + if [ -f /etc/profile.d/modules.sh ]; then + source /etc/profile.d/modules.sh >/dev/null 2>&1 || true + if command -v module >/dev/null 2>&1; then + module load miniconda3 >/dev/null 2>&1 || true + module load anaconda3 >/dev/null 2>&1 || true + if command -v conda >/dev/null 2>&1; then + local mod_base + mod_base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then + echo "${mod_base}/etc/profile.d/conda.sh" + return 0 + fi + fi + fi + fi + local user_name home_dir + user_name="${USER:-$(id -un 2>/dev/null || true)}" + home_dir="${HOME:-/home/${user_name}}" + local candidates=( + "${home_dir}/miniconda3/etc/profile.d/conda.sh" + "/home/${user_name}/miniconda3/etc/profile.d/conda.sh" + "${home_dir}/anaconda3/etc/profile.d/conda.sh" + "/home/${user_name}/anaconda3/etc/profile.d/conda.sh" + "/opt/miniconda3/etc/profile.d/conda.sh" + "/opt/anaconda3/etc/profile.d/conda.sh" + "/opt/conda/etc/profile.d/conda.sh" + ) + local c + for c in "${candidates[@]}"; do + if [[ -f "${c}" ]]; then + echo "${c}" + return 0 + fi + done + local found="" + found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)" + if [[ -n "${found}" && -f "${found}" ]]; then + echo "${found}" + return 0 + fi + return 1 +} + +CONDA_SH="$(resolve_conda_sh || true)" +if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then + echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2 + exit 1 +fi +source "${CONDA_SH}" +eval "$(conda shell.bash hook)" +conda activate "${CONDA_ENV}" + +cd "${PROJECT_DIR}" +export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}" + +JOB_TAG="rnn_labs_notes_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}" + +echo "========================================================" +echo " Labs+notes RNN run | ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo " Conda env : ${CONDA_ENV}" +echo " EHR root : ${EHR_ROOT}" +echo " Note root : ${NOTE_ROOT}" +echo " Cache dir : ${CACHE_DIR}" +echo " Output dir: ${OUTPUT_DIR}" +echo " Seed : ${SEED}" +echo " Dev mode : ${DEV_MODE}" +echo " Use AMP : ${USE_AMP} (dtype=${AMP_DTYPE})" +echo "========================================================" + +if ! python -c "import pyhealth" >/dev/null 2>&1; then + echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2 + exit 1 +fi + +COMMON=( + --ehr-root "${EHR_ROOT}" + --note-root "${NOTE_ROOT}" + --cache-dir "${CACHE_DIR}" + --task notes_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}" +) + +if [[ "${DEV_MODE}" == "1" ]]; then + COMMON+=(--dev) +fi + +if [[ "${FREEZE_ENCODER}" == "1" ]]; then + COMMON+=(--freeze-encoder) +fi + +if [[ "${INCLUDE_VITALS}" == "1" ]]; then + COMMON+=(--include-vitals) +fi + +if [[ "${USE_AMP}" == "1" ]]; then + COMMON+=(--use-amp --amp-dtype "${AMP_DTYPE}") +fi + +if [[ "${USE_WANDB}" == "1" ]]; then + COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}") + if [[ -n "${WANDB_RUN_NAME}" ]]; then + COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}") + fi +fi + +python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" + +echo "========================================================" +echo " Completed ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo "========================================================" diff --git a/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub b/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub new file mode 100644 index 000000000..118fa8149 --- /dev/null +++ b/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub @@ -0,0 +1,61 @@ +# HTCondor submission — labs+notes+CXR RNN mortality run +# +# Condor equivalent of +# scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py: +# same task (notes_labs_cxr), same model (rnn), same hyperparameters. Runs +# unattended instead of in a tmux session; Condor assigns the GPU (no manual +# nvidia-smi / CUDA_VISIBLE_DEVICES step needed). +# +# To submit (from the project root): +# mkdir -p /home/wp14/logs/condor +# condor_submit scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub +# +# Monitor: +# condor_q +# tail -f /home/wp14/logs/condor/labs_notes_cxr_rnn__0.out + +initialdir = /home/wp14/PyHealth +executable = /home/wp14/PyHealth/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh +transfer_executable = False +arguments = $(seed) +getenv = True + +environment = "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 \ + CXR_VARIANT=sunlab \ + CACHE_DIR=/shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr \ + OUTPUT_DIR=/home/wp14/output \ + USE_WANDB=1 \ + WANDB_PROJECT=pyhealth-multimodal-labs-notes-cxr \ + WANDB_RUN_NAME=labs_notes_cxr_rnn_seed$(seed) \ + FREEZE_ENCODER=1" + +output = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).out +error = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).err +log = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).log + +stream_output = False +stream_error = False + +request_gpus = 1 +request_cpus = 8 +# Starting point copied from labs_notes_rnn.sub (400000MB was sized for the +# full-scale patient_id sort/shuffle there). Untested for this variant: the +# added CXR branch means more dataloader workers decoding JPEGs plus an +# image-metadata join during caching, so this may need to be raised further +# if the job gets cgroup-killed. +request_memory = 400000MB +request_disk = 20GB + +# Same reasoning as labs_notes_rnn.sub: FREEZE_ENCODER=1 (frozen +# Bio_ClinicalBERT text encoder) keeps VRAM down for the text branch, but the +# CXR image encoder is unfrozen here and adds its own footprint on top, so +# the GPU memory floor is a conservative starting guess pending a real OOM +# data point on this variant. +Requirements = (TARGET.GPUs_GlobalMemoryMb >= 40000) +Rank = TARGET.GPUs_GlobalMemoryMb + +queue seed from ( + 12 +) diff --git a/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh b/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh new file mode 100755 index 000000000..aa9037ae7 --- /dev/null +++ b/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh @@ -0,0 +1,197 @@ +#!/usr/bin/env bash +# HTCondor executable — labs+notes+CXR RNN mortality run. +# +# Condor equivalent of +# scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py: +# same task (notes_labs_cxr), same model (rnn), same hyperparameter defaults. GPU +# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor +# assigns the GPU via request_gpus / cgroups. +# +# usage: run_labs_notes_cxr_rnn.sh +# to remove logs: rm -rf logs/condor/* +# to remove cache: rm -rf /shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr/* +set -euo pipefail + +SEED="${1:?usage: run_labs_notes_cxr_rnn.sh }" + +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}" +EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}" +NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}" +CXR_ROOT="${CXR_ROOT:-/shared/rsaas/physionet.org/files/MIMIC-CXR}" +CXR_VARIANT="${CXR_VARIANT:-default}" +CACHE_DIR="${CACHE_DIR:-/shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr}" +OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}" +CONDA_SH="${CONDA_SH:-}" + +DEV_MODE="${DEV_MODE:-1}" +EMBEDDING_DIM="${EMBEDDING_DIM:-128}" +HIDDEN_DIM="${HIDDEN_DIM:-128}" +RNN_TYPE="${RNN_TYPE:-GRU}" +RNN_LAYERS="${RNN_LAYERS:-2}" +DROPOUT="${DROPOUT:-0.1}" +EPOCHS="${EPOCHS:-15}" +BATCH_SIZE="${BATCH_SIZE:-32}" +LR="${LR:-1e-3}" +WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}" +PATIENCE="${PATIENCE:-5}" +NUM_WORKERS="${NUM_WORKERS:-4}" +FREEZE_ENCODER="${FREEZE_ENCODER:-1}" +INCLUDE_VITALS="${INCLUDE_VITALS:-0}" +USE_AMP="${USE_AMP:-0}" +AMP_DTYPE="${AMP_DTYPE:-bf16}" + +# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that +# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup. +# Disabling the whole distributed cluster (PYHEALTH_DISABLE_DASK_DISTRIBUTED=1) +# also throws away its disk-spilling/memory limits, which OOM'd the full-scale +# patient_id sort during event-dataframe caching (job 10821, 160GB cgroup +# limit hit). Scope the fix to just the NVML probe instead, keeping the real +# distributed cluster (with spilling) for the sort. +export DASK_DISTRIBUTED__DIAGNOSTICS__NVML="${DASK_DISTRIBUTED__DIAGNOSTICS__NVML:-0}" + +USE_WANDB="${USE_WANDB:-0}" +WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-notes-cxr}" +WANDB_RUN_NAME="${WANDB_RUN_NAME:-}" + +resolve_conda_sh() { + if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then + echo "${CONDA_SH}" + return 0 + fi + if command -v conda >/dev/null 2>&1; then + local base + base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then + echo "${base}/etc/profile.d/conda.sh" + return 0 + fi + fi + if [ -f /etc/profile.d/modules.sh ]; then + source /etc/profile.d/modules.sh >/dev/null 2>&1 || true + if command -v module >/dev/null 2>&1; then + module load miniconda3 >/dev/null 2>&1 || true + module load anaconda3 >/dev/null 2>&1 || true + if command -v conda >/dev/null 2>&1; then + local mod_base + mod_base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then + echo "${mod_base}/etc/profile.d/conda.sh" + return 0 + fi + fi + fi + fi + local user_name home_dir + user_name="${USER:-$(id -un 2>/dev/null || true)}" + home_dir="${HOME:-/home/${user_name}}" + local candidates=( + "${home_dir}/miniconda3/etc/profile.d/conda.sh" + "/home/${user_name}/miniconda3/etc/profile.d/conda.sh" + "${home_dir}/anaconda3/etc/profile.d/conda.sh" + "/home/${user_name}/anaconda3/etc/profile.d/conda.sh" + "/opt/miniconda3/etc/profile.d/conda.sh" + "/opt/anaconda3/etc/profile.d/conda.sh" + "/opt/conda/etc/profile.d/conda.sh" + ) + local c + for c in "${candidates[@]}"; do + if [[ -f "${c}" ]]; then + echo "${c}" + return 0 + fi + done + local found="" + found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)" + if [[ -n "${found}" && -f "${found}" ]]; then + echo "${found}" + return 0 + fi + return 1 +} + +CONDA_SH="$(resolve_conda_sh || true)" +if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then + echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2 + exit 1 +fi +source "${CONDA_SH}" +eval "$(conda shell.bash hook)" +conda activate "${CONDA_ENV}" + +cd "${PROJECT_DIR}" +export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}" + +JOB_TAG="rnn_labs_notes_cxr_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}" + +echo "========================================================" +echo " Labs+notes+CXR RNN run | ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo " Conda env : ${CONDA_ENV}" +echo " EHR root : ${EHR_ROOT}" +echo " Note root : ${NOTE_ROOT}" +echo " CXR root : ${CXR_ROOT} (variant=${CXR_VARIANT})" +echo " Cache dir : ${CACHE_DIR}" +echo " Output dir : ${OUTPUT_DIR}" +echo " Seed : ${SEED}" +echo " Dev mode : ${DEV_MODE}" +echo " Use AMP : ${USE_AMP} (dtype=${AMP_DTYPE})" +echo "========================================================" + +if ! python -c "import pyhealth" >/dev/null 2>&1; then + echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2 + exit 1 +fi + +COMMON=( + --ehr-root "${EHR_ROOT}" + --note-root "${NOTE_ROOT}" + --cxr-root "${CXR_ROOT}" + --cxr-variant "${CXR_VARIANT}" + --cache-dir "${CACHE_DIR}" + --task notes_labs_cxr + --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}" +) + +if [[ "${DEV_MODE}" == "1" ]]; then + COMMON+=(--dev) +fi + +if [[ "${FREEZE_ENCODER}" == "1" ]]; then + COMMON+=(--freeze-encoder) +fi + +if [[ "${INCLUDE_VITALS}" == "1" ]]; then + COMMON+=(--include-vitals) +fi + +if [[ "${USE_AMP}" == "1" ]]; then + COMMON+=(--use-amp --amp-dtype "${AMP_DTYPE}") +fi + +if [[ "${USE_WANDB}" == "1" ]]; then + COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}") + if [[ -n "${WANDB_RUN_NAME}" ]]; then + COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}") + fi +fi + +python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" + +echo "========================================================" +echo " Completed ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo "========================================================" diff --git a/scripts/will/condor/labs_only/labs_only_rnn.sub b/scripts/will/condor/labs_only/labs_only_rnn.sub new file mode 100644 index 000000000..e7d66a1b1 --- /dev/null +++ b/scripts/will/condor/labs_only/labs_only_rnn.sub @@ -0,0 +1,45 @@ +# HTCondor submission — labs-only RNN mortality run +# +# Condor equivalent of scripts/will/sunlab/tmux_run_labs_only_rnn_variant.py: +# same task (labs), same model (rnn), same hyperparameters. Runs unattended +# instead of in a tmux session; Condor assigns the GPU (no manual +# nvidia-smi / CUDA_VISIBLE_DEVICES step needed). +# +# To submit (from the project root): +# mkdir -p /home/wp14/logs/condor +# condor_submit scripts/will/condor/labs_only/labs_only_rnn.sub +# +# Monitor: +# condor_q +# tail -f /home/wp14/logs/condor/labs_only_rnn__0.out + +initialdir = /home/wp14/PyHealth +executable = /home/wp14/PyHealth/scripts/will/condor/labs_only/run_labs_only_rnn.sh +transfer_executable = False +arguments = $(seed) +getenv = True + +environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \ + CACHE_DIR=/shared/eng/wp14/pyhealth_cache_labs \ + OUTPUT_DIR=/home/wp14/output \ + USE_WANDB=1 \ + WANDB_PROJECT=pyhealth-multimodal-labs-only \ + WANDB_RUN_NAME=labs_rnn_seed$(seed)" + +output = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).out +error = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).err +log = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).log + +stream_output = False +stream_error = False + +request_gpus = 1 +request_cpus = 4 +request_memory = 163840MB +request_disk = 20GB + +Rank = TARGET.CUDAGlobalMemoryMb + +queue seed from ( + 12 +) \ No newline at end of file diff --git a/scripts/will/condor/labs_only/run_labs_only_rnn.sh b/scripts/will/condor/labs_only/run_labs_only_rnn.sh new file mode 100755 index 000000000..e05a8895f --- /dev/null +++ b/scripts/will/condor/labs_only/run_labs_only_rnn.sh @@ -0,0 +1,168 @@ +#!/usr/bin/env bash +# HTCondor executable — labs-only RNN mortality run. +# +# Condor equivalent of scripts/will/sunlab/tmux_run_labs_only_rnn_variant.py: +# same task (labs), same model (rnn), same hyperparameter defaults. GPU +# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor +# assigns the GPU via request_gpus / cgroups. +# +# usage: run_labs_only_rnn.sh +# to remove logs: rm -rf logs/condor/* +# to remove cache: rm -rf /shared/eng/wp14/pyhealth_cache_labs/* +set -euo pipefail + +SEED="${1:?usage: run_labs_only_rnn.sh }" + +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}" +EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}" +CACHE_DIR="${CACHE_DIR:-/shared/eng/wp14/pyhealth_cache_labs}" +OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}" +CONDA_SH="${CONDA_SH:-}" + +DEV_MODE="${DEV_MODE:-1}" +EMBEDDING_DIM="${EMBEDDING_DIM:-64}" +HIDDEN_DIM="${HIDDEN_DIM:-64}" +RNN_TYPE="${RNN_TYPE:-GRU}" +RNN_LAYERS="${RNN_LAYERS:-1}" +DROPOUT="${DROPOUT:-0.1}" +EPOCHS="${EPOCHS:-15}" +BATCH_SIZE="${BATCH_SIZE:-32}" +LR="${LR:-1e-3}" +WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}" +PATIENCE="${PATIENCE:-5}" +NUM_WORKERS="${NUM_WORKERS:-4}" + +# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that +# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup. +# Skip the distributed Dask cluster (falls back to the plain local scheduler) +# to avoid touching NVML during event-dataframe preprocessing. +export PYHEALTH_DISABLE_DASK_DISTRIBUTED="${PYHEALTH_DISABLE_DASK_DISTRIBUTED:-1}" + +USE_WANDB="${USE_WANDB:-0}" +WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-only}" +WANDB_RUN_NAME="${WANDB_RUN_NAME:-}" + +resolve_conda_sh() { + if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then + echo "${CONDA_SH}" + return 0 + fi + if command -v conda >/dev/null 2>&1; then + local base + base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then + echo "${base}/etc/profile.d/conda.sh" + return 0 + fi + fi + if [ -f /etc/profile.d/modules.sh ]; then + source /etc/profile.d/modules.sh >/dev/null 2>&1 || true + if command -v module >/dev/null 2>&1; then + module load miniconda3 >/dev/null 2>&1 || true + module load anaconda3 >/dev/null 2>&1 || true + if command -v conda >/dev/null 2>&1; then + local mod_base + mod_base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then + echo "${mod_base}/etc/profile.d/conda.sh" + return 0 + fi + fi + fi + fi + local user_name home_dir + user_name="${USER:-$(id -un 2>/dev/null || true)}" + home_dir="${HOME:-/home/${user_name}}" + local candidates=( + "${home_dir}/miniconda3/etc/profile.d/conda.sh" + "/home/${user_name}/miniconda3/etc/profile.d/conda.sh" + "${home_dir}/anaconda3/etc/profile.d/conda.sh" + "/home/${user_name}/anaconda3/etc/profile.d/conda.sh" + "/opt/miniconda3/etc/profile.d/conda.sh" + "/opt/anaconda3/etc/profile.d/conda.sh" + "/opt/conda/etc/profile.d/conda.sh" + ) + local c + for c in "${candidates[@]}"; do + if [[ -f "${c}" ]]; then + echo "${c}" + return 0 + fi + done + local found="" + found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)" + if [[ -n "${found}" && -f "${found}" ]]; then + echo "${found}" + return 0 + fi + return 1 +} + +CONDA_SH="$(resolve_conda_sh || true)" +if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then + echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2 + exit 1 +fi +source "${CONDA_SH}" +eval "$(conda shell.bash hook)" +conda activate "${CONDA_ENV}" + +cd "${PROJECT_DIR}" +export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}" + +JOB_TAG="rnn_labs_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}" + +echo "========================================================" +echo " Labs-only RNN run | ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo " Conda env : ${CONDA_ENV}" +echo " EHR root : ${EHR_ROOT}" +echo " Cache dir : ${CACHE_DIR}" +echo " Output dir: ${OUTPUT_DIR}" +echo " Seed : ${SEED}" +echo " Dev mode : ${DEV_MODE}" +echo "========================================================" + +if ! python -c "import pyhealth" >/dev/null 2>&1; then + echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2 + exit 1 +fi + +COMMON=( + --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}" +) + +if [[ "${DEV_MODE}" == "1" ]]; then + COMMON+=(--dev) +fi + +if [[ "${USE_WANDB}" == "1" ]]; then + COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}") + if [[ -n "${WANDB_RUN_NAME}" ]]; then + COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}") + fi +fi + +python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" + +echo "========================================================" +echo " Completed ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo "========================================================" \ No newline at end of file diff --git a/tests/core/test_stagenet_processor.py b/tests/core/test_stagenet_processor.py index 6e217dc7d..1cde77f6a 100644 --- a/tests/core/test_stagenet_processor.py +++ b/tests/core/test_stagenet_processor.py @@ -199,9 +199,8 @@ def test_empty_codes_flat(self): time, values = processor.process((None, [])) - # Should return single padding token - self.assertEqual(values.shape, (1,)) - self.assertEqual(values[0].item(), processor.code_vocab[""]) + # Should return zero events, not a fake pad token + self.assertEqual(values.shape, (0,)) def test_empty_codes_nested(self): """Test processing empty nested codes.""" @@ -211,10 +210,8 @@ def test_empty_codes_nested(self): time, values = processor.process((None, [])) - # Should return single row of padding tokens - self.assertEqual(values.shape, (1, 2)) - self.assertEqual(values[0, 0].item(), processor.code_vocab[""]) - self.assertEqual(values[0, 1].item(), processor.code_vocab[""]) + # Should return zero events, not a fake pad row + self.assertEqual(values.shape, (0, 2)) def test_vocab_size_method(self): """Test vocab_size() returns correct size.""" diff --git a/tests/test_frozen_text.py b/tests/test_frozen_text.py new file mode 100644 index 000000000..f39c8eae3 --- /dev/null +++ b/tests/test_frozen_text.py @@ -0,0 +1,88 @@ +"""Proofs for frozen-text encoder behaviour in UnifiedMultimodalEmbeddingModel.""" + +from __future__ import annotations + +import unittest + +import torch +import torch.nn as nn + +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel +from pyhealth.processors.stagenet_processor import StageNetTensorProcessor + + +def _numeric_model(**kwargs) -> UnifiedMultimodalEmbeddingModel: + proc = StageNetTensorProcessor() + proc.fit([{"labs": ([0.0], [[1.0] * 10])}], "labs") + return UnifiedMultimodalEmbeddingModel( + {"labs": proc}, embedding_dim=8, freeze_text_encoder=True, **kwargs + ) + + +class TinyEnc(nn.Module): + def __init__(self): + super().__init__() + self.drop = nn.Dropout(p=0.9) + self.lin = nn.Linear(1, 8) + self.config = type("C", (), {"hidden_size": 8})() + + def forward(self, input_ids, attention_mask=None): + b, l = input_ids.shape + h = self.drop(torch.ones(b, l, 8)) + return type("O", (), {"last_hidden_state": h})() + + +class TestFrozenEncoderEval(unittest.TestCase): + def test_train_keeps_frozen_text_encoder_in_eval(self): + model = _numeric_model() + enc = TinyEnc() + model.encoders["notes"] = enc + model._frozen_text_fields.add("notes") + model.train() + self.assertTrue(model.training) + self.assertFalse(enc.training) + model.eval() + self.assertFalse(enc.training) + model.train() + self.assertFalse(enc.training) + + +class CountingEnc(nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + self.lin = nn.Linear(1, 8) + self.config = type("C", (), {"hidden_size": 8})() + + def forward(self, input_ids, attention_mask=None): + self.calls += 1 + b, l = input_ids.shape + scale = input_ids[:, :1].float() + h = torch.ones(b, l, 8) * scale + return type("O", (), {"last_hidden_state": h})() + + +class TestFrozenTextCache(unittest.TestCase): + def test_cache_keys_ignore_padding_tokens(self): + model = _numeric_model(cache_frozen_text=True) + enc = CountingEnc() + model.encoders["notes"] = enc + model._frozen_text_fields.add("notes") + + ids_a = torch.tensor([[1, 2, 3, 0, 0], [1, 2, 3, 9, 9]]) + mask_a = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 0, 0]]) + h1 = model._encode_text_cls("notes", enc, ids_a, mask_a) + self.assertEqual(enc.calls, 1) + h2 = model._encode_text_cls("notes", enc, ids_a, mask_a) + self.assertEqual(enc.calls, 1) + ids_b = torch.tensor([[1, 2, 3, 7, 7, 7]]) + mask_b = torch.tensor([[1, 1, 1, 0, 0, 0]]) + h3 = model._encode_text_cls("notes", enc, ids_b, mask_b) + self.assertEqual(enc.calls, 1) + self.assertEqual(h1.shape[0], 2) + self.assertEqual(h3.shape[0], 1) + + +if __name__ == "__main__": + unittest.main() + diff --git a/tests/test_p0_parquet_scan.py b/tests/test_p0_parquet_scan.py new file mode 100644 index 000000000..cd9b58153 --- /dev/null +++ b/tests/test_p0_parquet_scan.py @@ -0,0 +1,62 @@ +"""Proof that BaseDataset still scans parquet files. + +Will's a0f1422 deleted _scan_table/_scan_parquet while MEDS still calls +_scan_parquet. These tests load a real parquet file, not just inspect source. +""" + +from __future__ import annotations + +import os +import tempfile +import unittest +from pathlib import Path + +import pandas as pd + + +class TestP0ParquetScan(unittest.TestCase): + def test_scan_parquet_reads_a_real_file(self): + from pyhealth.datasets.base_dataset import BaseDataset + + self.assertTrue(callable(getattr(BaseDataset, "_scan_table"))) + self.assertTrue(callable(getattr(BaseDataset, "_scan_parquet"))) + + class _Host: + pass + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "events.parquet" + pd.DataFrame({"subject_id": ["a", "b"], "n": [1, 2]}).to_parquet(path) + df = BaseDataset._scan_parquet(_Host(), str(path)) + out = df.compute() + self.assertEqual(len(out), 2) + self.assertIn("subject_id", out.columns) + + def test_scan_table_routes_a_parquet_file_to_the_parquet_scanner(self): + from pyhealth.datasets.base_dataset import BaseDataset + + class _Host: + def _scan_parquet(self, source_path): + return f"parquet:{source_path}" + + def _scan_csv_tsv_gz(self, source_path): + return f"csv:{source_path}" + + host = _Host() + pq = BaseDataset._scan_table(host, "/tmp/events.parquet") + csv = BaseDataset._scan_table(host, "/tmp/events.csv.gz") + self.assertTrue(pq.startswith("parquet:")) + self.assertTrue(csv.startswith("csv:")) + + +class TestP0AbsoluteTablePath(unittest.TestCase): + def test_resolve_table_path_keeps_absolute(self): + from pyhealth.datasets.base_dataset import resolve_table_path + + abs_csv = "/tmp/generated/mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + resolved = resolve_table_path("/data/root", abs_csv) + self.assertTrue(os.path.isabs(resolved)) + self.assertTrue(resolved.endswith("mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv")) + self.assertEqual(resolved, str(Path(abs_csv).expanduser().resolve())) + rel = resolve_table_path("/data/root", "hosp/patients.csv.gz") + self.assertTrue(rel.endswith("hosp/patients.csv.gz")) diff --git a/tests/test_p1_amp_dtype.py b/tests/test_p1_amp_dtype.py new file mode 100644 index 000000000..57fb1b959 --- /dev/null +++ b/tests/test_p1_amp_dtype.py @@ -0,0 +1,43 @@ +"""Proof that amp_dtype is validated instead of silently coerced to fp16. + +The previous expression was ``bfloat16 if amp_dtype == "bf16" else float16``. +Any other spelling, including ``"bfloat16"``, selected fp16 with no message. +fp16 also constructs a GradScaler, so the silent path changed gradients. + +Measured: ``resolve_amp_dtype("bfloat16")`` is bf16; ``"f16"`` raises. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_amp_dtype.py -q +""" + +from __future__ import annotations + +import unittest + +import torch + + +class TestP1AmpDtype(unittest.TestCase): + def test_known_spellings_map_to_the_named_dtype(self): + from pyhealth.trainer import resolve_amp_dtype + + self.assertIs(resolve_amp_dtype("bf16"), torch.bfloat16) + self.assertIs(resolve_amp_dtype("bfloat16"), torch.bfloat16) + self.assertIs(resolve_amp_dtype("fp16"), torch.float16) + self.assertIs(resolve_amp_dtype("float16"), torch.float16) + + def test_autocast_device_follows_the_trainer_device(self): + from pyhealth.trainer import autocast_device_type + + self.assertEqual(autocast_device_type("cuda:0"), "cuda") + self.assertEqual(autocast_device_type("cpu"), "cpu") + self.assertEqual(autocast_device_type("mps"), "mps") + + def test_unknown_spelling_raises(self): + from pyhealth.trainer import resolve_amp_dtype + + for bad in ("bfloat_16", "f16", "int8", ""): + with self.assertRaises(ValueError): + resolve_amp_dtype(bad) diff --git a/tests/test_p1_fp16_attention.py b/tests/test_p1_fp16_attention.py new file mode 100644 index 000000000..512943967 --- /dev/null +++ b/tests/test_p1_fp16_attention.py @@ -0,0 +1,77 @@ +"""Proof that attention mask fill is fp16-safe and ordinary forwards use SDPA. + +``-1e9`` is outside the fp16 range, so AMP overflowed on padded positions +(``value cannot be converted to type at``). Ordinary training uses fused +SDPA; the explicit path stays for interpretability and fills with +``finfo(dtype).min``. + +Measured (CPU, this checkout, ``TransformerLayer`` 128/4 heads/2 layers, +seed 0, ``B=4 S=32``): + + fused vs explicit max abs diff: 7.153e-07 (no padding), 4.768e-07 (with padding) + fp16 padded forward: finite; pad weight exactly 0.0 + raw ``masked_fill(..., -1e9)`` on fp16: raises + +A10 GPU, ``notes_labs``, transformer 128/2/4, full scale, 2 epochs, batch 8 +(characterises mixed precision already on ``main``; this commit makes the +fp16 path numerically valid): + + bf16 5,275 s (1471.3 / 1007.2 s/epoch) 1,814 MB loss 1.2345 -> 1.1412 + fp32 10,198 s (4176.1 / 1808.8 s/epoch) 2,402 MB loss 1.2348 -> 1.1341 + epoch-1 2.84x, mean 2.41x, VRAM -24.5%, final train loss within 0.63% + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_fp16_attention.py -q +""" + +from __future__ import annotations + +import inspect +import unittest + +import torch + + +class TestP1Fp16Mask(unittest.TestCase): + def test_fp16_attention_mask_fill_does_not_overflow(self): + from pyhealth.models.transformer import Attention + + attn = Attention() + q = torch.zeros(1, 1, 2, 4, dtype=torch.float16) + k = torch.zeros(1, 1, 2, 4, dtype=torch.float16) + v = torch.ones(1, 1, 2, 4, dtype=torch.float16) + mask = torch.tensor([[[[1, 0], [1, 0]]]], dtype=torch.float16) + out, weights = attn(q, k, v, mask=mask) + self.assertTrue(torch.isfinite(out).all()) + self.assertTrue(torch.isfinite(weights).all()) + self.assertEqual(float(weights[0, 0, 0, 1]), 0.0) + + def test_raw_minus_1e9_still_overflows_fp16(self): + scores = torch.zeros(2, 2, dtype=torch.float16) + with self.assertRaises(RuntimeError): + scores.masked_fill(torch.tensor([[True, False], [False, True]]), -1e9) + + def test_ordinary_forward_uses_fused_sdpa(self): + from pyhealth.models.transformer import MultiHeadedAttention + + src = inspect.getsource(MultiHeadedAttention.forward) + self.assertIn("scaled_dot_product_attention", src) + self.assertIn("register_hook", src) + + def test_fused_and_explicit_paths_agree(self): + from pyhealth.models.transformer import TransformerLayer + + torch.manual_seed(0) + layer = TransformerLayer(feature_size=128, heads=4, num_layers=2).eval() + x = torch.randn(4, 32, 128, requires_grad=True) + mask = torch.ones(4, 32) + fused, _ = layer(x, mask, register_hook=False) + explicit, _ = layer(x, mask, register_hook=True) + self.assertLessEqual(float((fused - explicit).abs().max()), 1e-5) + + mask_pad = torch.cat([torch.ones(4, 20), torch.zeros(4, 12)], dim=1) + fused_pad, _ = layer(x, mask_pad, register_hook=False) + explicit_pad, _ = layer(x, mask_pad, register_hook=True) + self.assertLessEqual(float((fused_pad - explicit_pad).abs().max()), 1e-5) diff --git a/tests/test_p1_observation_window.py b/tests/test_p1_observation_window.py new file mode 100644 index 000000000..6cb694ec4 --- /dev/null +++ b/tests/test_p1_observation_window.py @@ -0,0 +1,143 @@ +"""Proof that every lab/CXR task honours a per-admission observation window. + +Four task bodies computed a window from ``window_hours`` and then collected +labs through discharge. For a mortality label that reads the outcome. + +Measured on MIMIC-IV ``labs_only``: + + collection through discharge PR-AUC 0.6204 ROC 0.90 + window honoured PR-AUC 0.2137 ROC 0.7136 + +A sweep over 24, 48, and 96 hours produced three identical datasets, because +``window_hours`` was inert. The window also anchored on the patient's first +admission globally, so a later stay received a span that had already closed. + +CXR / ``notes_labs_cxr`` still skipped those later stays with +``admission_time >= first_admit + window_hours``. That skip is gone. +``emitted_data_version`` is 4 so caches from version 1-3 cannot be reused. +The table protocol default is full stay (``window_hours=None``). Passing +``window_hours=24`` still caps collection per admission; CXR arms must not +skip later stays against the first admission's clock. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_observation_window.py -q +""" + +from __future__ import annotations + +import inspect +import json +import unittest +import uuid +from datetime import datetime, timedelta + + +LAB_TASKS = [ + "LabsMIMIC4", + "NotesLabsMIMIC4", + "NotesLabsCXRMIMIC4", + "CXRMIMIC4", +] + + +class TestP1ObservationWindow(unittest.TestCase): + def test_every_task_honours_a_24h_window(self): + from pyhealth.tasks import multimodal_mimic4 as m + + admit = datetime(2180, 5, 6, 8, 0, 0) + discharge = admit + timedelta(days=9) + for name in LAB_TASKS: + task = getattr(m, name)(window_hours=24) + end = task._admission_window_end(admit, discharge) + horizon = (end - admit).total_seconds() / 3600.0 + self.assertAlmostEqual( + horizon, + 24.0, + places=2, + msg=f"{name} collects {horizon:.0f}h past admission", + ) + self.assertLess(end, discharge) + + def test_window_is_anchored_per_admission(self): + from pyhealth.tasks import multimodal_mimic4 as m + + first = datetime(2180, 5, 6, 8, 0, 0) + later = first + timedelta(days=400) + for name in LAB_TASKS: + task = getattr(m, name)(window_hours=24) + self.assertEqual( + task._admission_window_end(first, first + timedelta(days=9)), + first + timedelta(hours=24), + ) + self.assertEqual( + task._admission_window_end(first, first + timedelta(hours=6)), + first + timedelta(hours=6), + ) + end = task._admission_window_end(later, later + timedelta(days=5)) + self.assertEqual(end, later + timedelta(hours=24)) + self.assertGreater(end, later, msg=f"{name} expired before later stay") + + def test_cxr_arms_do_not_skip_later_stays_on_the_first_admit_clock(self): + from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4, NotesLabsCXRMIMIC4 + + for cls in (NotesLabsCXRMIMIC4, CXRMIMIC4): + src = inspect.getsource(cls.__call__) + self.assertNotIn( + "admission_time >= effective_end", + src, + msg=f"{cls.__name__} still drops later stays against first admit + window", + ) + + def test_window_change_invalidates_the_cache(self): + from pyhealth.tasks import multimodal_mimic4 as m + + task = m.LabsMIMIC4(window_hours=24) + self.assertIsNotNone(vars(task).get("emitted_data_version")) + self.assertGreaterEqual(task.emitted_data_version, 4) + + def cache_key(t, drop_version=False): + v = dict(vars(t)) + if drop_version: + v.pop("emitted_data_version", None) + params = json.dumps( + { + **v, + "input_schema": t.input_schema, + "output_schema": t.output_schema, + }, + sort_keys=True, + default=str, + ) + return str(uuid.uuid5(uuid.NAMESPACE_DNS, params)) + + self.assertNotEqual(cache_key(task), cache_key(task, drop_version=True)) + + def test_window_none_still_collects_through_discharge(self): + from pyhealth.tasks.multimodal_mimic4 import LabsMIMIC4 + + task = LabsMIMIC4(window_hours=None) + admit = datetime(2180, 5, 6, 8, 0, 0) + discharge = admit + timedelta(days=9) + self.assertEqual(task._admission_window_end(admit, discharge), discharge) + + def test_protocol_default_is_full_stay(self): + from pyhealth.tasks.multimodal_mimic4 import ( + CXRMIMIC4, + LabsMIMIC4, + NotesLabsCXRMIMIC4, + NotesLabsMIMIC4, + ) + + self.assertIsNone(NotesLabsMIMIC4().window_hours) + self.assertIsNone(NotesLabsCXRMIMIC4().window_hours) + self.assertIsNone(LabsMIMIC4().window_hours) + self.assertIsNone(CXRMIMIC4().window_hours) + + def test_discharge_coded_icd_is_not_a_mortality_task(self): + from pyhealth.tasks import multimodal_mimic4 as m + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + self.assertFalse(hasattr(m, "ICDLabsMIMIC4")) + self.assertFalse(NotesLabsMIMIC4().include_icd) diff --git a/tests/test_p1_pad_masks.py b/tests/test_p1_pad_masks.py new file mode 100644 index 000000000..4f210b27b --- /dev/null +++ b/tests/test_p1_pad_masks.py @@ -0,0 +1,95 @@ +"""Proof that batch padding is recorded and skipped by the unified path. + +The collator padded short samples with 0.0 and nothing recorded that padding, +so padded slots looked like real measurements at admission time. RNN packed +lengths of 0 also crash once a correct mask exists. + +Measured: + + ``collate_temporal`` had zero callers; the dataloader uses + ``collate_fn_dict_with_padding``. First ``pad_mask`` on the unused collator + never reached a model. + Token-budget notes then crashed: + ``RuntimeError: The size of tensor a (7) must match the size of tensor b (14)`` + (``pad_sequence`` only pads dim 0). + Reusing event ``pad_mask`` as the BERT token mask then crashed: + ``RuntimeError: shape '[96, 512]' is invalid for input of size 96``. + After the collator emits ``{field}__pad_mask``, a 3-event + 1-event note + batch is ``(2, 3, 4)`` with mask ``[[True, True, True], [True, False, False]]``. + An all-pad RNN step stays finite (lengths clamped at 1). + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_pad_masks.py -q +""" + +from __future__ import annotations + +import unittest + +import torch + + +class TestP1RnnClamp(unittest.TestCase): + def test_all_pad_mask_forward_is_finite(self): + from pyhealth.models.rnn import RNNLayer + + layer = RNNLayer(input_size=4, hidden_size=8, dropout=0.0).eval() + x = torch.zeros(2, 5, 4) + mask = torch.zeros(2, 5) + with torch.no_grad(): + outputs, last = layer(x, mask) + self.assertEqual(outputs.shape[0], 2) + self.assertEqual(tuple(last.shape), (2, 8)) + self.assertTrue(torch.isfinite(last).all()) + + +class TestP1BertPadSkip(unittest.TestCase): + def test_collate_emits_false_pad_mask_on_empty_note_slots(self): + from pyhealth.datasets.utils import PAD_MASK_SUFFIX, collate_fn_dict_with_padding + + batch = [ + {"notes": (torch.ones(3, 4, dtype=torch.long), torch.ones(3, 4))}, + {"notes": (torch.ones(1, 4, dtype=torch.long), torch.ones(1, 4))}, + ] + collated = collate_fn_dict_with_padding(batch) + ids = collated["notes"][0] + pad = collated[f"notes{PAD_MASK_SUFFIX}"] + self.assertEqual(tuple(ids.shape), (2, 3, 4)) + self.assertEqual(pad.tolist(), [[True, True, True], [True, False, False]]) + + def test_padded_events_sort_last_and_are_zeroed(self): + from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel + from pyhealth.processors.stagenet_processor import StageNetTensorProcessor + + proc = StageNetTensorProcessor() + proc.fit([{"labs": ([0.0], [[1.0, 2.0]])}], "labs") + model = UnifiedMultimodalEmbeddingModel( + processors={"labs": proc}, + embedding_dim=8, + normalize_content=True, + ) + model.eval() + value = torch.tensor([[[3.0, 4.0], [0.0, 0.0]]]) + time = torch.tensor([[6.0, 0.0]]) + pad_mask = torch.tensor([[True, False]]) + with torch.no_grad(): + out = model({"labs": {"value": value, "time": time, "pad_mask": pad_mask}}) + self.assertEqual(out["mask"].tolist(), [[1.0, 0.0]]) + self.assertAlmostEqual(out["time"][0, 0].item(), 6.0) + self.assertTrue(torch.allclose(out["sequence"][0, 1], torch.zeros(8), atol=1e-6)) + + def test_unified_heads_thread_pad_mask(self): + import inspect + + from pyhealth.models.bottleneck_transformer import BottleneckTransformer + from pyhealth.models.ehrmamba import EHRMamba + from pyhealth.models.jamba_ehr import JambaEHR + from pyhealth.models.rnn import RNN + from pyhealth.models.transformer import Transformer + + for cls in (RNN, Transformer, BottleneckTransformer, EHRMamba, JambaEHR): + src = inspect.getsource(cls._build_unified_inputs) + self.assertIn("PAD_MASK_SUFFIX", src, msg=cls.__name__) + self.assertIn("pad_mask", src, msg=cls.__name__) diff --git a/tests/test_p1_time_axis.py b/tests/test_p1_time_axis.py new file mode 100644 index 000000000..04f21f137 --- /dev/null +++ b/tests/test_p1_time_axis.py @@ -0,0 +1,132 @@ +"""Proof that concatenated stays share one time origin. + +The sample is patient-level: every admission up to the first death is one +sequence. Event times were hours from *that stay's* admit, then concatenated, +so stay 2 at +6h sorted with stay 1 at +6h. Collection is still per stay +(admit through discharge, or admit+window if ``window_hours`` is set). Times +are hours from the first stay in the sample. + +Admission-context discharge sections are stamped at that stay's admit, not +at the discharge note's ``charttime``. Radiology stays at exam time. + +Single-stay patients are unchanged. Time embeddings use geometric +wavelengths from 1 hour to 10 years so a later stay at +9606h does not +alias with +6h. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_time_axis.py -q +""" + +from __future__ import annotations + +import inspect +import unittest +from datetime import datetime, timedelta + +import torch + + +class TestP1TimeAxis(unittest.TestCase): + def test_hours_since_does_not_reset_per_stay(self): + from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task + + first = datetime(2180, 5, 6, 8, 0, 0) + later = first + timedelta(days=400) + stay1_event = first + timedelta(hours=6) + stay2_event = later + timedelta(hours=6) + + self.assertEqual(BaseMultimodalMIMIC4Task._hours_since(stay1_event, first), 6.0) + # Old convention: both events were 6.0 and collided after concat. + self.assertEqual(BaseMultimodalMIMIC4Task._hours_since(stay2_event, later), 6.0) + self.assertAlmostEqual( + BaseMultimodalMIMIC4Task._hours_since(stay2_event, first), + 400 * 24 + 6.0, + ) + self.assertGreater( + BaseMultimodalMIMIC4Task._hours_since(stay2_event, first), + BaseMultimodalMIMIC4Task._hours_since(stay1_event, first), + ) + + def test_collectors_write_hours_from_the_sample_origin(self): + from pyhealth.tasks.multimodal_mimic4 import ( + BaseMultimodalMIMIC4Task, + CXRMIMIC4, + LabsMIMIC4, + NotesLabsCXRMIMIC4, + NotesLabsMIMIC4, + ) + + for cls in (LabsMIMIC4, NotesLabsMIMIC4, NotesLabsCXRMIMIC4, CXRMIMIC4): + src = inspect.getsource(cls.__call__) + self.assertIn( + "time_origin = admissions_to_process[0].timestamp", + src, + msg=cls.__name__, + ) + self.assertNotIn( + "event.timestamp - admission_time", + src, + msg=f"{cls.__name__} still resets times per stay", + ) + + labs_src = inspect.getsource(BaseMultimodalMIMIC4Task._collect_labs) + self.assertIn("time_origin", labs_src) + self.assertIn("_hours_since", labs_src) + self.assertNotIn("lab_ts - admission_time", labs_src) + + def test_sinusoid_does_not_wrap_every_thirty_days(self): + from pyhealth.models.embedding.unified import SinusoidalTimeEmbedding + + emb = SinusoidalTimeEmbedding(dim=32) + t6 = emb(torch.tensor([6.0])) + t726 = emb(torch.tensor([6.0 + 720.0])) + t9606 = emb(torch.tensor([9606.0])) + self.assertFalse(torch.allclose(t6, t726, atol=1e-5)) + self.assertFalse(torch.allclose(t6, t9606, atol=1e-5)) + + def test_discharge_sections_are_stamped_at_admit_not_charttime(self): + from types import SimpleNamespace + + from pyhealth.tasks.multimodal_mimic4 import NotesLabsCXRMIMIC4, NotesLabsMIMIC4 + + first = datetime(2180, 5, 6, 8, 0, 0) + stay2_admit = first + timedelta(days=400) + charttime = stay2_admit + timedelta(days=9) + note = SimpleNamespace( + text="Chief Complaint:\n\nshortness of breath\n\n", + timestamp=charttime, + ) + + class FakePatient: + def get_events(self, event_type, start=None, end=None, filters=None): + return [note] + + task = NotesLabsMIMIC4() + _texts, times = task._collect_notes( + FakePatient(), + "discharge", + hadm_id=1, + admission_time=stay2_admit, + section_headers=task.DISCHARGE_CLINICAL_HEADERS, + time_origin=first, + event_time=stay2_admit, + ) + self.assertEqual(times, [400 * 24.0]) + + _texts, rad_times = task._collect_notes( + FakePatient(), + "radiology", + hadm_id=1, + admission_time=stay2_admit, + time_origin=first, + ) + self.assertAlmostEqual(rad_times[0], 400 * 24.0 + 9 * 24.0) + + for cls in (NotesLabsMIMIC4, NotesLabsCXRMIMIC4): + self.assertIn( + "event_time=admission_time", + inspect.getsource(cls.__call__), + msg=cls.__name__, + ) diff --git a/tests/test_p2_lab_standardizer.py b/tests/test_p2_lab_standardizer.py new file mode 100644 index 000000000..2b3e574a4 --- /dev/null +++ b/tests/test_p2_lab_standardizer.py @@ -0,0 +1,109 @@ +"""Proof that lab z-scores fit on observed train rows, not a WORLD_SIZE shard. + +``SampleDataset`` subclasses ``litdata.StreamingDataset``. Under ``torchrun``, +``WORLD_SIZE`` is set before ``torch.distributed`` is initialised, so +``__len__`` / ``__iter__`` silently yield 1/N of the train split (the same +shard on every rank). Measured on real litdata with 20 samples: ``len()`` +reports 5 under ``WORLD_SIZE=4`` while ``region_of_interest`` still sums to 20. + +Fitting padded 0.0 as if it were a measurement also moves sodium's mean from +140 to 105. ``patient_to_index`` is unusable after ``subset()``: it still +holds parent indices and raised ``ValueError: index 237 didn't find a match +within the chunk intervals``. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_lab_standardizer.py -q +""" + +from __future__ import annotations + +import os +import unittest +from unittest import mock + +import torch + + +def _lab_samples(n: int = 40): + torch.manual_seed(0) + return [ + { + "labs": torch.stack( + [140.0 + torch.randn(1) * 4, 1.0 + torch.randn(1) * 0.2] + ).view(1, 2), + "labs_mask": torch.ones(1, 2, dtype=torch.bool), + } + for _ in range(n) + ] + + +class TestP2LabStandardizer(unittest.TestCase): + def test_fit_ignores_padded_zeros(self): + from pyhealth.processors import fit_lab_standardizer + + samples = [ + { + "labs": torch.tensor([[140.0, 1.0], [0.0, 0.0]]), + "labs_mask": torch.tensor([[True, True], [False, False]]), + }, + { + "labs": torch.tensor([[142.0, 1.2], [138.0, 0.8]]), + "labs_mask": torch.tensor([[True, True], [True, True]]), + }, + ] + standardizer = fit_lab_standardizer(samples) + # Observed sodium 140, 142, 138. Mean 140, not 105 from the padded 0.0. + self.assertAlmostEqual(standardizer.mean[0].item(), 140.0, places=4) + + def test_unobserved_slot_maps_to_zero(self): + from pyhealth.processors import fit_lab_standardizer + + standardizer = fit_lab_standardizer(_lab_samples()) + values = torch.tensor([[[140.0, 1.0], [0.0, 0.0]]]) + observed = torch.tensor([[[True, True], [False, False]]]) + out = standardizer(values, observed) + self.assertEqual(out[0, 1].abs().sum().item(), 0.0) + self.assertTrue(torch.isfinite(out).all()) + + def test_world_size_does_not_shrink_the_fit(self): + from pyhealth.processors import fit_lab_standardizer + + samples = _lab_samples(40) + + class _ShardedByWorldSize: + def __init__(self, records): + self._records = records + self.region_of_interest = [(0, len(records))] + + def _visible(self): + world = int(os.environ.get("WORLD_SIZE", "1")) + return self._records[: len(self._records) // world] + + def __len__(self): + return len(self._visible()) + + def __iter__(self): + return iter(self._visible()) + + def __getitem__(self, index): + return self._records[index] + + dataset = _ShardedByWorldSize(samples) + single = fit_lab_standardizer(dataset) + with mock.patch.dict(os.environ, {"WORLD_SIZE": "4"}): + sharded = fit_lab_standardizer(dataset) + self.assertTrue(torch.allclose(single.mean, sharded.mean)) + self.assertTrue(torch.allclose(single.std, sharded.std)) + self.assertEqual( + int(single.observed_count.sum()), int(sharded.observed_count.sum()) + ) + + def test_statistics_travel_in_the_state_dict(self): + from pyhealth.processors import fit_lab_standardizer + + standardizer = fit_lab_standardizer(_lab_samples()) + keys = set(standardizer.state_dict()) + self.assertTrue({"mean", "std"} <= keys) + self.assertEqual(tuple(standardizer.state_dict()["mean"].shape), (2,)) diff --git a/tests/test_p2_mlp_pad_mask.py b/tests/test_p2_mlp_pad_mask.py new file mode 100644 index 000000000..56aaae247 --- /dev/null +++ b/tests/test_p2_mlp_pad_mask.py @@ -0,0 +1,48 @@ +"""Proof that unified MLP threads collate pad_mask like the other heads. + +Without this, ``--model mlp`` is missing from the six-backbone table, and a +unified MLP would take the ``mask is None`` branch and score padded slots. + +Measured: a 3-event + 1-event lab batch collates to pad_mask +``[[True, True, True], [True, False, False]]``, and MLP copies it into +``inputs["labs"]["pad_mask"]``. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_mlp_pad_mask.py -q +""" + +from __future__ import annotations + +import unittest +from types import SimpleNamespace + +import torch + + +class TestP2MlpPadMask(unittest.TestCase): + def test_mlp_threads_pad_mask_into_unified_inputs(self): + from pyhealth.datasets.utils import PAD_MASK_SUFFIX, collate_fn_dict_with_padding + from pyhealth.models.mlp import MLP + from pyhealth.processors.stagenet_processor import StageNetTensorProcessor + + batch = [ + {"labs": (torch.tensor([6.0, 12.0, 24.0]), torch.ones(3, 2))}, + {"labs": (torch.tensor([6.0]), torch.ones(1, 2))}, + ] + collated = collate_fn_dict_with_padding(batch) + self.assertIn(f"labs{PAD_MASK_SUFFIX}", collated) + + host = SimpleNamespace( + feature_keys=["labs"], + device="cpu", + dataset=SimpleNamespace( + input_processors={"labs": StageNetTensorProcessor()} + ), + ) + inputs = MLP._build_unified_inputs(host, collated) + self.assertEqual( + inputs["labs"]["pad_mask"].tolist(), + [[True, True, True], [True, False, False]], + ) diff --git a/tests/test_p2_nested_padding_idx.py b/tests/test_p2_nested_padding_idx.py new file mode 100644 index 000000000..5c8ae7f53 --- /dev/null +++ b/tests/test_p2_nested_padding_idx.py @@ -0,0 +1,63 @@ +"""Proof that nested code embeddings freeze the pad row. + +``NestedSequenceProcessor`` used ``padding_idx=None``, so a fake empty visit +could have a non-zero vector and index 0 received gradients. Empty visits are +now zero events, so the pad row stays frozen zeros. + +Measured (CPU, this checkout, ``embedding_dim=8``, backward on pad indices): + + ``emb.padding_idx == 0`` + ``||weight[0]|| == 0`` + ``||grad[0]|| == 0`` + +No cluster table for this row: the defect is that pad embeddings were trainable. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_nested_padding_idx.py -q +""" + +from __future__ import annotations + +import unittest + +import torch + + +class TestP2NestedPaddingIdx(unittest.TestCase): + def test_nested_embedding_pad_row_is_zeros_and_frozen(self): + from pyhealth.datasets import create_sample_dataset + from pyhealth.models.embedding.vanilla import EmbeddingModel + + samples = [ + { + "patient_id": "p0", + "visit_id": "v0", + "conditions": [["A", "B"], ["C"]], + "label": 0, + }, + { + "patient_id": "p1", + "visit_id": "v0", + "conditions": [["A"]], + "label": 1, + }, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"conditions": "nested_sequence"}, + output_schema={"label": "binary"}, + in_memory=True, + ) + model = EmbeddingModel(dataset, embedding_dim=8) + emb = model.embedding_layers["conditions"] + self.assertEqual(emb.padding_idx, 0) + self.assertTrue(torch.equal(emb.weight[0], torch.zeros_like(emb.weight[0]))) + x = torch.zeros(2, 3, dtype=torch.long) + y = emb(x).sum() + y.backward() + self.assertIsNotNone(emb.weight.grad) + self.assertTrue( + torch.equal(emb.weight.grad[0], torch.zeros_like(emb.weight.grad[0])) + ) diff --git a/tests/test_p2_run_config.py b/tests/test_p2_run_config.py new file mode 100644 index 000000000..06c5168ff --- /dev/null +++ b/tests/test_p2_run_config.py @@ -0,0 +1,35 @@ +"""Proof that a finished run records the conditions that produced the score. + +``metrics_history.json`` stores PR-AUC but not whether the encoder was frozen, +which split ran, or which code produced it. Cluster jobs start from an unpacked +archive, so git is often empty; the SHA-256 of the package source still +identifies the code. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_run_config.py -q +""" + +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + + +class TestP2RunConfig(unittest.TestCase): + def test_write_run_config_records_resolved_fields(self): + from pyhealth.utils import write_run_config + + with tempfile.TemporaryDirectory() as tmp: + path = write_run_config( + tmp, {"resolved_lr": 1e-4, "split_mode": "by_patient"} + ) + data = json.loads(Path(path).read_text()) + self.assertEqual(data["config"]["resolved_lr"], 0.0001) + self.assertIn("git", data) + self.assertIn("source_sha256", data) + self.assertEqual(len(data["source_sha256"]), 64) + self.assertTrue(path.endswith("run_config.json")) diff --git a/tests/test_p2_runner_measurement.py b/tests/test_p2_runner_measurement.py new file mode 100644 index 000000000..84658451d --- /dev/null +++ b/tests/test_p2_runner_measurement.py @@ -0,0 +1,131 @@ +"""Proof that paired runs cannot overwrite each other or report train as test. + +The run directory was ``{model}_seed{seed}``. ``--task labs`` and +``--task notes_labs`` at seed 42 resolved to one path and the second run +destroyed the first. ``split_by_patient`` fell back to ``split_by_sample`` +with no warning (patient leak). Predictions came from +``test_loader or val_loader or train_loader``. + +The sixth backbone was also missing: ``--model`` had no ``mlp``. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_runner_measurement.py -q +""" + +from __future__ import annotations + +import importlib.util +import sys +import unittest +import warnings +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + + +RUNNER = ( + Path(__file__).resolve().parents[1] + / "examples" + / "mortality_prediction" + / "unified_embedding_e2e_mimic4.py" +) + + +def _load_runner(): + spec = importlib.util.spec_from_file_location("e2e_runner_measurement", RUNNER) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _parse(mod, *argv): + with mock.patch.object(sys, "argv", ["e2e.py", "--ehr-root", "/tmp", *argv]): + return mod.parse_args() + + +def _exp_name_assignment() -> str: + for line in RUNNER.read_text().splitlines(): + stripped = line.strip() + if stripped.startswith("exp_name = "): + return stripped + raise AssertionError("exp_name is no longer assigned in the runner") + + +class TestP2RunnerMeasurement(unittest.TestCase): + def test_run_directory_includes_the_task(self): + assignment = _exp_name_assignment() + self.assertIn("args.task", assignment) + template = assignment.split("=", 1)[1].strip() + args = SimpleNamespace(task="labs", model="transformer", seed=42) + labs = eval(template, {"args": args}) # noqa: S307 + args.task = "notes_labs" + notes = eval(template, {"args": args}) # noqa: S307 + self.assertEqual(labs, "labs_transformer_seed42") + self.assertEqual(notes, "notes_labs_transformer_seed42") + self.assertNotEqual(labs, notes) + + def test_cli_accepts_mlp(self): + from pyhealth.models import MLP + + self.assertTrue(hasattr(MLP, "_forward_unified")) + mod = _load_runner() + args = _parse(mod, "--model", "mlp") + self.assertEqual(args.model, "mlp") + + def test_split_warns_and_labels_leaky_fallback(self): + from pyhealth.datasets import create_sample_dataset + + mod = _load_runner() + samples = [ + { + "patient_id": "only", + "visit_id": "v0", + "labs": [1.0, 2.0], + "label": 0, + }, + { + "patient_id": "only", + "visit_id": "v1", + "labs": [3.0, 4.0], + "label": 1, + }, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"labs": "tensor"}, + output_schema={"label": "binary"}, + in_memory=True, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _train, _val, _test, mode = mod._split_dataset(dataset, seed=1) + self.assertEqual(mode, "by_sample_fallback_leaky") + self.assertTrue(any("split_by_sample" in str(w.message) for w in caught)) + + def test_eval_split_is_named_and_recorded(self): + src = RUNNER.read_text() + self.assertIn('inference_loader, eval_split = test_loader, "test"', src) + self.assertIn("write_run_config", src) + self.assertIn("These are not test metrics", src) + self.assertNotIn( + "inference_loader = test_loader or val_loader or train_loader", + src, + ) + + def test_jamba_cli_default_matches_the_library(self): + from pyhealth.models.jamba_ehr import JambaLayer + + self.assertEqual(JambaLayer.__init__.__defaults__[1], 6) + mod = _load_runner() + args = _parse(mod, "--model", "jambaehr") + self.assertEqual(args.jamba_transformer_layers, 2) + self.assertEqual(args.jamba_mamba_layers, 6) + + def test_observation_window_defaults_to_full_stay(self): + mod = _load_runner() + args = _parse(mod) + self.assertIsNone(args.observation_window_hours) + args_24 = _parse(mod, "--observation-window-hours", "24") + self.assertEqual(args_24.observation_window_hours, 24) diff --git a/tests/test_p2_sunlab_cache.py b/tests/test_p2_sunlab_cache.py new file mode 100644 index 000000000..8ce5b9a8e --- /dev/null +++ b/tests/test_p2_sunlab_cache.py @@ -0,0 +1,92 @@ +"""Proof that sunlab CXR metadata can be written to cache, not the data root. + +The loader required a directory named ``images`` and wrote the derived CSV +into the PhysioNet root. The complete resized set lives under +``resized_images``, and the cluster root is read-only. + +Measured on the cluster: + + Complete resized cohort: 377,110 / 377,110 images, 0 dropped. + 256x256 greyscale, 3.3 GB. PhysioNet ``images/`` was incomplete (p13-p19 + absent). Hardcoded ``images/`` raised ``FileNotFoundError`` on the complete + set. Default CXR config raised ``KeyError: 'studytime_normalized'``. + ``chmod 0o555`` on the root: CSV lands under ``cache_dir``. + +After the layout worked (one seed, 6 epochs, 18,542 train / 2,285 test, +prevalence 0.0565, split seed 42). These compare with each other, not with +the primary notes_labs table: + + cxr_only PR-AUC 0.0602 ROC 0.5267 + cxr + labs PR-AUC 0.3082 ROC 0.8047 + cxr + notes+labs PR-AUC 0.4096 ROC 0.8625 + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_sunlab_cache.py -q +""" + +from __future__ import annotations + +import os +import tempfile +import unittest +from pathlib import Path + +import pandas as pd + + +class TestP2SunlabLayout(unittest.TestCase): + def _fake_cxr_root(self, tmp: str, image_dirname: str) -> str: + root = Path(tmp) / "cxr" + (root / image_dirname).mkdir(parents=True) + pd.DataFrame( + {"dicom_id": ["abc"], "StudyTime": ["93000"], "subject_id": ["1"]} + ).to_csv(root / "mimic-cxr-2.0.0-metadata.csv", index=False) + return str(root) + + def test_resized_images_and_cache_write(self): + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + host = MIMIC4CXRSunlabDataset.__new__(MIMIC4CXRSunlabDataset) + with tempfile.TemporaryDirectory() as tmp: + root = self._fake_cxr_root(tmp, "resized_images") + cache = Path(tmp) / "cache" + cache.mkdir() + dest = host.prepare_metadata(root, cache_dir=str(cache)) + self.assertTrue(dest.startswith(str(cache))) + self.assertTrue(Path(dest).is_file()) + root_csv = Path(root) / "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + self.assertFalse(root_csv.exists()) + written = pd.read_csv(dest) + self.assertTrue( + str(written.loc[0, "image_path"]).endswith( + os.path.join("resized_images", "abc.jpg") + ) + ) + + def test_unwritable_root_falls_back_to_cache(self): + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + host = MIMIC4CXRSunlabDataset.__new__(MIMIC4CXRSunlabDataset) + with tempfile.TemporaryDirectory() as tmp: + root = self._fake_cxr_root(tmp, "images") + cache = Path(tmp) / "cache" + cache.mkdir() + os.chmod(root, 0o555) + try: + dest = host.prepare_metadata(root, cache_dir=str(cache)) + finally: + os.chmod(root, 0o755) + self.assertTrue(Path(dest).is_file()) + self.assertTrue(dest.startswith(str(cache))) + + def test_writes_to_root_when_no_cache_dir(self): + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + host = MIMIC4CXRSunlabDataset.__new__(MIMIC4CXRSunlabDataset) + with tempfile.TemporaryDirectory() as tmp: + root = self._fake_cxr_root(tmp, "images") + dest = host.prepare_metadata(root, cache_dir=None) + self.assertTrue(dest.startswith(root)) + self.assertTrue(Path(dest).is_file()) diff --git a/tests/test_tuple_time_text_processor.py b/tests/test_tuple_time_text_processor.py index 7a8dfd5b5..1f51ee481 100644 --- a/tests/test_tuple_time_text_processor.py +++ b/tests/test_tuple_time_text_processor.py @@ -25,6 +25,12 @@ def test_tuple_time_text_processor(): assert torch.equal(time_tensor, torch.tensor([0.0, 24.0, 72.0])) assert tag == "clinical_note" + # Empty input is zero events, not a fake "[MISSING_TEXT]" token. + empty_texts, empty_time, empty_tag = processor.process(([], [])) + assert empty_texts == [] + assert empty_time.shape == (0,) + assert empty_tag == "clinical_note" + # Test registration from pyhealth.processors import get_processor ProcessorClass = get_processor("tuple_time_text") diff --git a/tests/test_unified_multimodal.py b/tests/test_unified_multimodal.py index 01bc194e6..f2e3f5fac 100644 --- a/tests/test_unified_multimodal.py +++ b/tests/test_unified_multimodal.py @@ -168,22 +168,44 @@ def test_collate_temporal_variable_length(): # ── 5. SinusoidalTimeEmbedding ──────────────────────────────────────────────── +def test_legacy_unified_import_is_the_package(): + from pyhealth.models.embedding.unified import ( + SinusoidalTimeEmbedding as PackageTime, + UnifiedMultimodalEmbeddingModel as PackageModel, + ) + from pyhealth.models.unified_embedding import ( + SinusoidalTimeEmbedding as LegacyTime, + UnifiedMultimodalEmbeddingModel as LegacyModel, + ) + + assert PackageTime is LegacyTime + assert PackageModel is LegacyModel + + def test_sinusoidal_time_embedding_shape(): - from pyhealth.models.unified_embedding import SinusoidalTimeEmbedding - emb = SinusoidalTimeEmbedding(dim=64, max_hours=720.0) + from pyhealth.models.embedding import SinusoidalTimeEmbedding + emb = SinusoidalTimeEmbedding(dim=64) 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.unified_embedding import SinusoidalTimeEmbedding + from pyhealth.models.embedding import SinusoidalTimeEmbedding emb = SinusoidalTimeEmbedding(dim=32) t0 = emb(torch.tensor([0.0])) t1 = emb(torch.tensor([24.0])) assert not torch.allclose(t0, t1) +def test_sinusoidal_does_not_alias_every_720h(): + from pyhealth.models.embedding import SinusoidalTimeEmbedding + emb = SinusoidalTimeEmbedding(dim=32) + t6 = emb(torch.tensor([6.0])) + t726 = emb(torch.tensor([726.0])) + assert not torch.allclose(t6, t726, atol=1e-5) + + # ── 6. UnifiedMultimodalEmbeddingModel — code-only smoke test ───────────────── def _make_code_processors_and_inputs(batch_size=2, seq_len=5): @@ -206,7 +228,7 @@ def _make_code_processors_and_inputs(batch_size=2, seq_len=5): def test_unified_model_code_only(): - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel processors, inputs = _make_code_processors_and_inputs() model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=64) @@ -224,7 +246,7 @@ def test_unified_model_code_only(): def test_unified_model_rejects_non_temporal(): - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel from pyhealth.processors import SequenceProcessor bad_proc = SequenceProcessor() @@ -234,7 +256,7 @@ def test_unified_model_rejects_non_temporal(): def test_unified_model_gradient_flow(): """Loss.backward() should propagate through time + type embeddings.""" - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel processors, inputs = _make_code_processors_and_inputs() model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=32) @@ -250,7 +272,7 @@ def test_unified_model_gradient_flow(): def test_unified_model_time_sort(): """Events should be sorted by time ascending in the output.""" - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel from pyhealth.processors import StageNetProcessor samples = [{"c": (None, ["a", "b"])}]