diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 75968787c..0736cb0e0 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -65,8 +65,9 @@ NotesLabsCXRMIMIC4, NotesLabsMIMIC4, ) +from pyhealth.processors import fit_lab_standardizer from pyhealth.trainer import Trainer -from pyhealth.utils import set_seed +from pyhealth.utils import set_seed, write_run_config logger = logging.getLogger(__name__) @@ -238,7 +239,11 @@ def _inert_arch_flags(model: str) -> list[str]: return ["--" + f.replace("_", "-") for f in _ARCH_FLAGS_ALL if f not in used] -def _build_model(args: argparse.Namespace, sample_dataset: Any): +def _build_model( + args: argparse.Namespace, + sample_dataset: Any, + numeric_standardizers: Optional[dict[str, Any]] = None, +): inert = _inert_arch_flags(args.model) if inert: logger.warning( @@ -252,6 +257,7 @@ def _build_model(args: argparse.Namespace, sample_dataset: Any): processors=sample_dataset.input_processors, embedding_dim=args.embedding_dim, freeze_text_encoder=args.freeze_encoder, + numeric_standardizers=numeric_standardizers, max_frozen_text_cache=args.max_frozen_text_cache, text_grad_checkpoint_rows=args.text_grad_checkpoint_rows, ) @@ -365,7 +371,13 @@ def run(args: argparse.Namespace) -> Path: sample_dataset, seed=args.seed, allow_leaky_split=args.allow_leaky_split ) - model = _build_model(args, sample_dataset) + # Lab z-scores, fit on the training split only. Missing values stay + # missing; --no-lab-standardization runs raw labs as an ablation. + numeric_standardizers: dict[str, Any] = {} + if "labs" in sample_dataset.input_processors and not args.no_lab_standardization: + numeric_standardizers["labs"] = fit_lab_standardizer(train_ds) + + model = _build_model(args, sample_dataset, numeric_standardizers) train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True) val_loader = ( @@ -379,7 +391,7 @@ def run(args: argparse.Namespace) -> Path: else None ) - exp_name = f"{args.model}_seed{args.seed}" + exp_name = f"{args.task}_{args.model}_seed{args.seed}" output_dir = Path(args.output_dir) wandb_logger = WandbLogger( @@ -430,11 +442,33 @@ def run(args: argparse.Namespace) -> Path: test_scores = trainer.evaluate(test_loader) wandb_logger.log({f"test_{k}": v for k, v in test_scores.items()}) - inference_loader = test_loader or val_loader or train_loader + 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; predictions come from VAL.", RuntimeWarning) + else: + inference_loader, eval_split = train_loader, "train" + warnings.warn( + "No test or val split; predictions come from TRAIN and are held-in.", + RuntimeWarning, + ) y_true, y_prob, _, patient_ids = trainer.inference( inference_loader, return_patient_ids=True ) + write_run_config( + str(output_dir / exp_name), + { + **vars(args), + "eval_split": eval_split, + "lab_standardization": bool(numeric_standardizers), + "n_train": len(train_ds), + "n_val": len(val_ds), + "n_test": len(test_ds), + }, + ) + output_csv = output_dir / exp_name / f"predictions_{args.model}.csv" _write_predictions(output_csv, patient_ids, y_true, y_prob) @@ -501,6 +535,12 @@ def parse_args() -> argparse.Namespace: "tests only — the metrics are not usable." ), ) + parser.add_argument( + "--no-lab-standardization", + action="store_true", + default=False, + help="Disable train-split lab z-scoring (raw-lab ablation).", + ) parser.add_argument("--weight-decay", type=float, default=0.0) parser.add_argument("--device", type=str, default=None) parser.add_argument( diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index 4568a5ece..d7df193ae 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -82,4 +82,11 @@ def get_processor(name: str): "TupleTimeTextProcessor", "CehrProcessor", "ConceptVocab", + "LabStandardizer", + "fit_lab_standardizer", + "", ] +from .lab_standardizer import ( + LabStandardizer, + fit_lab_standardizer, +) diff --git a/pyhealth/processors/lab_standardizer.py b/pyhealth/processors/lab_standardizer.py new file mode 100644 index 000000000..0e193c453 --- /dev/null +++ b/pyhealth/processors/lab_standardizer.py @@ -0,0 +1,130 @@ +"""Train-split-only z-scoring for masked temporal laboratory values. + +Labs and their observation mask are separate temporal fields. Only rows whose +mask is true are fitted: zero-filled / forward-filled missing values must never +affect a lab's mean or variance. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any, Optional + +import torch +from torch import nn + + +def _indices(dataset: Any) -> Optional[list[int]]: + """Explicit sample indices for this split, or None for a plain iterable. + + ``SampleDataset`` subclasses ``litdata.StreamingDataset``, whose ``__iter__`` + and ``__len__`` are sharded by ``WORLD_SIZE``: under torchrun, iterating + would silently fit on 1/WORLD_SIZE of the train split. Indexing is not + sharded, and ``region_of_interest`` is the only unsharded description of + what the split holds. (``patient_to_index`` is not usable — ``subset()`` + copies it unchanged, so after ``split_by_patient`` it still indexes the + parent and raises "index ... didn't find a match within the chunk + intervals".) + """ + roi = getattr(dataset, "region_of_interest", None) + return list(range(sum(end - start for start, end in roi))) if roi else None + + +class LabStandardizer(nn.Module): + """Per-feature z-score with persistent train-only statistics. + + ``mean``/``std``/``observed_count`` are buffers, so they travel in the + model ``state_dict`` and a checkpoint transforms serving inputs exactly as + it did at training time. + """ + + def __init__( + self, + mean: torch.Tensor, + std: torch.Tensor, + observed_count: torch.Tensor, + ) -> None: + super().__init__() + if (std <= 0).any() or not torch.isfinite(mean).all(): + raise ValueError("Lab statistics must be finite with positive std.") + 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().clone()) + + @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, + ) -> "LabStandardizer": + """Fit on observed, finite values of the supplied (already split) data.""" + mask_field = observation_mask_field or f"{value_field}_mask" + idx = _indices(samples) + stream = samples if idx is None else (samples[i] for i in idx) + + count = total = total_sq = None + for sample in stream: + if value_field not in sample or mask_field not in sample: + continue + v = sample[value_field] + m = sample[mask_field] + v = v[1] if isinstance(v, (tuple, list)) else v + m = m[1] if isinstance(m, (tuple, list)) else m + v = torch.as_tensor(v, dtype=torch.float64) + m = torch.as_tensor(m).bool() & torch.isfinite(v) + if v.ndim == 1: + v, m = v.unsqueeze(0), m.unsqueeze(0) + if count is None: + z = torch.zeros(v.shape[-1], dtype=torch.float64) + count, total, total_sq = z.clone(), z.clone(), z.clone() + obs = torch.where(m, v, torch.zeros_like(v)) + count += m.sum(0).to(torch.float64) + total += obs.sum(0) + total_sq += (obs * obs).sum(0) + + if count is None: + raise ValueError( + f"No samples carried both {value_field!r} and {mask_field!r}." + ) + + seen = count > 0 + mean = torch.where(seen, total / count.clamp(min=1), torch.zeros_like(total)) + var = torch.where( + seen, + (total_sq / count.clamp(min=1) - mean.square()).clamp_min(0), + torch.ones_like(total), + ) + # A constant train feature maps to zero; unit std keeps that finite. + std = torch.where(var > 0, var.sqrt(), torch.ones_like(var)) + return cls(mean.to(torch.float32), std.to(torch.float32), count) + + def forward( + self, values: torch.Tensor, observed_mask: torch.Tensor + ) -> torch.Tensor: + """Z-score observed values; missing or unfittable features map to zero. + + Deliberately not clipped: there is no universally valid physiological + range for these MIMIC category aggregates, so values outside train + support stay as large finite z-scores and remain auditable. + """ + if values.shape[-1] != self.feature_dim: + raise ValueError( + f"Expected {self.feature_dim} lab features, got {values.shape[-1]}." + ) + values = values.to(dtype=self.mean.dtype) + observed = observed_mask.bool() & torch.isfinite(values) + z = (values - self.mean) / self.std + return torch.where(observed & (self.observed_count > 0), z, torch.zeros_like(z)) + + +def fit_lab_standardizer( + train_dataset: Iterable[dict[str, Any]], **kwargs: Any +) -> LabStandardizer: + """Fit on the training split only. Train-only is enforced by the caller.""" + return LabStandardizer.fit(train_dataset, **kwargs) diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index aa1d588b5..20dbef61c 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -83,6 +83,12 @@ def __init__( window_hours: Optional[float] = None, ): self.window_hours = window_hours + # Part of vars(task), which is what the task-cache uuid5 key is built + # from. Without it, a cache built before an emitted-data change is + # silently reused. Bump whenever the emitted data changes. + # 2: event times are hours from the first stay in the sample, not + # reset per admission. + self.emitted_data_version = 2 @staticmethod def _clean_text(text: Optional[str]) -> Optional[str]: diff --git a/pyhealth/utils.py b/pyhealth/utils.py index b4af8980a..b46c66ac4 100644 --- a/pyhealth/utils.py +++ b/pyhealth/utils.py @@ -1,4 +1,6 @@ import json +import hashlib +import subprocess import os import pickle import random @@ -65,4 +67,76 @@ def set_env(**environ): yield finally: os.environ.clear() - os.environ.update(old_environ) \ No newline at end of file + os.environ.update(old_environ) + + +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 _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 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 diff --git a/scripts/paper/common.sh b/scripts/paper/common.sh new file mode 100644 index 000000000..ac6491011 --- /dev/null +++ b/scripts/paper/common.sh @@ -0,0 +1,47 @@ +# Shared Tranche 1 protocol. Sourced by rian.sh / will.sh, which add only the +# data roots and the CPU tuning appropriate to their machine. +# +# TASK=labs_notes MODEL=ehrmamba SEED=3 bash scripts/paper/rian.sh +# +# Env: TASK MODEL SEED TREE EHR_ROOT NOTE_ROOT CXR_ROOT CACHE_DIR OUT GPU + +TASK="${TASK:?set TASK=labs|labs_notes|labs_notes_cxr}" +MODEL="${MODEL:?set MODEL=mlp|rnn|transformer|bottleneck_transformer|ehrmamba|jambaehr}" +SEED="${SEED:?set SEED}" +TREE="${TREE:-$HOME/PyHealth}" +CACHE_DIR="${CACHE_DIR:-$HOME/pyhealth_cache/$TASK}" +OUT="${OUT:-$TREE/output}" + +case "$TASK" in + labs) TASK_FLAG=labs; ROOTS=() ;; + labs_notes) TASK_FLAG=notes_labs; ROOTS=(--note-root "$NOTE_ROOT") ;; + labs_notes_cxr) TASK_FLAG=notes_labs_cxr; ROOTS=(--note-root "$NOTE_ROOT" --cxr-root "$CXR_ROOT" --cxr-variant sunlab) ;; + *) echo "unknown TASK=$TASK" >&2; exit 2 ;; +esac + +case "$MODEL" in + mlp) ARCH=(--mlp-layers 2 --mlp-activation relu) ;; + rnn) ARCH=(--rnn-type GRU --rnn-layers 1) ;; + transformer) ARCH=(--heads 4 --num-layers 2) ;; + bottleneck_transformer) ARCH=(--heads 4 --num-layers 2 --bottlenecks-n 4 --fusion-startidx 1) ;; + ehrmamba) ARCH=(--num-layers 2 --mamba-state-size 16 --mamba-conv-kernel 4) ;; + jambaehr) ARCH=(--heads 4 --jamba-transformer-layers 2 --jamba-mamba-layers 6) ;; + *) echo "unknown MODEL=$MODEL" >&2; exit 2 ;; +esac + +# Identical for every cell: the Tranche 1 protocol. +PROTOCOL=(--embedding-dim 128 --hidden-dim 128 --dropout 0.1 + --batch-size 32 --lr 1e-4 --epochs 50 --patience 5 + --use-amp --amp-dtype bf16 --freeze-encoder) + +launch () { # any extra flags are passed through + cd "$TREE"; export PYTHONPATH="$TREE" + mkdir -p logs "$OUT" + [[ -n "${GPU:-}" ]] && export CUDA_VISIBLE_DEVICES="$GPU" + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \ + --task "$TASK_FLAG" --model "$MODEL" --seed "$SEED" \ + --ehr-root "$EHR_ROOT" "${ROOTS[@]}" \ + --cache-dir "$CACHE_DIR" --output-dir "$OUT" \ + "${PROTOCOL[@]}" "${ARCH[@]}" "$@" \ + 2>&1 | tee "logs/${TASK}_${MODEL}_seed${SEED}.out" +} diff --git a/scripts/paper/rian.sh b/scripts/paper/rian.sh new file mode 100755 index 000000000..4d42736df --- /dev/null +++ b/scripts/paper/rian.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# Rian's cells: seeds 1,3,5 on labs and labs_notes, plus all of labs_notes_cxr. +# +# Same protocol as will.sh; only the data roots and CPU tuning differ. These +# nodes are shared and run several cells at once, and torch defaults to 64 +# intra-op + 128 inter-op threads with nothing pinning them: four unpinned +# cells put ~800 threads on 128 cores and epoch time went 191s -> 8600s with +# the GPUs at 0-1%. Keep THREADS x concurrent_cells under the node's cores. +set -euo pipefail +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}" +THREADS="${THREADS:-8}" +export OMP_NUM_THREADS="$THREADS" MKL_NUM_THREADS="$THREADS" +export OPENBLAS_NUM_THREADS="$THREADS" NUMEXPR_NUM_THREADS="$THREADS" +source "$(dirname "$(readlink -f "$0")")/common.sh" +launch --num-workers "${NUM_WORKERS:-8}" \ + --loader-num-workers "${LOADER_WORKERS:-4}" --persistent-workers diff --git a/scripts/paper/will.sh b/scripts/paper/will.sh new file mode 100755 index 000000000..97d1ec8f8 --- /dev/null +++ b/scripts/paper/will.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +# Will's cells: seeds 2 and 4 on labs and labs_notes. +set -euo pipefail +EHR_ROOT="${EHR_ROOT:-/home/ubuntu/mimiciv-data/ehr}" +NOTE_ROOT="${NOTE_ROOT:-/home/ubuntu/mimiciv-data}" +CXR_ROOT="${CXR_ROOT:-/home/ubuntu/mimiciv-data/CXR-jpg}" +source "$(dirname "$(readlink -f "$0")")/common.sh" +launch --num-workers "${NUM_WORKERS:-4}" 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,))