Wire lab z-scoring, run provenance, and a named eval split

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(

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,
+)

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)

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]:

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)
+        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

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"
+}