diff --git a/template/src/meds_model_base/lightning/__init__.py b/template/src/meds_model_base/lightning/__init__.py index 6765c50..54eee31 100644 --- a/template/src/meds_model_base/lightning/__init__.py +++ b/template/src/meds_model_base/lightning/__init__.py @@ -2,7 +2,8 @@ - :func:`register_structured_configs` registers ``MEDSTorchDataConfig`` with Hydra's ConfigStore so a ``datamodule.config`` group is a type-checked structured config (enums are UPPERCASE on the CLI). -- :func:`build_datamodule` builds the meds-torch-data ``Datamodule`` from a resolved ``cfg.datamodule``. +- :func:`build_datamodule` builds the meds-torch-data ``Datamodule`` from a resolved ``cfg.datamodule``, + after :func:`require_statics_if_requested` checks the cohort can satisfy the static data it asks for. - :func:`build_trainer` builds the ``Trainer``, directing checkpoints at a training run's work directory. - :mod:`meds_model_base.lightning.modules` holds reusable ``nn.Module`` blocks + ``BaseLightningModule``. - :mod:`meds_model_base.lightning.probe` holds the frozen-embedding probe (dataset assembly + head). @@ -10,6 +11,7 @@ from __future__ import annotations +from pathlib import Path from typing import TYPE_CHECKING if TYPE_CHECKING: # pragma: no cover - typing only @@ -34,14 +36,77 @@ def register_structured_configs(group: str = "datamodule/config") -> None: _STRUCTURED_CONFIGS_REGISTERED = True +#: Static-inclusion modes that make the batch carry ``static_code`` / ``static_numeric_value``. +_STATIC_MODES = frozenset({"include", "prepend"}) + + +def require_statics_if_requested(cfg: DictConfig) -> None: + """Refuse a cohort with no static measurements when the datamodule asks for them. + + meds-torch-data builds the static tensors from whatever the cohort holds. When that is *nothing at + all* — no subject has a single null-time measurement — ``JointNestedRaggedTensorDict`` has no values + to infer a dtype from, and collation dies inside the dataloader with:: + + ValueError: Cannot infer dtype from empty values; provide an explicit `schema=`. + + That error names neither the config key that asked for static data nor the cohort that lacks it, and it + arrives once training has already started. This is the same precondition check as + ``_require_split_sharded``: knowable in a directory listing, so it belongs before the work rather than + several minutes into it. + + A cohort with no static measurements is perfectly legitimate MEDS — the mismatch is between it and the + request, so either side is a valid fix, and the message says so. + + Raises: + ValueError: if static data is requested and the tensorized cohort has none. + """ + node = cfg.get("datamodule") + config = node.get("config") if node else None + if not config: + return + # `StaticInclusionMode` is a StrEnum, but the Hydra structured config spells it UPPERCASE. + if str(config.get("static_inclusion_mode", "")).lower() not in _STATIC_MODES: + return + + cohort = config.get("tensorized_cohort_dir") + if not cohort: + return + schemas = sorted(Path(cohort).glob("tokenization/schemas/*/*.parquet")) + if not schemas: + return # not a tensorized cohort at all; meds-torch-data's own error is the clearer one here + + import polars as pl + + frame = pl.scan_parquet(schemas) + if "static_code" not in frame.collect_schema().names(): + return + if frame.select(pl.col("static_code").list.len().sum()).collect().item(): + return + + raise ValueError( + f"datamodule.config.static_inclusion_mode is " + f"{config.get('static_inclusion_mode')!s}, but no subject in {cohort} has any static " + "measurement, so meds-torch-data would fail to collate a batch at all ('Cannot infer dtype from " + "empty values').\n\n" + "Static measurements are MEDS rows with a null `time` — baseline variables such as age, sex or " + "ethnicity. Either:\n" + " * the source dataset genuinely has none, in which case set " + "`datamodule.config.static_inclusion_mode=OMIT` and drop them from the model; or\n" + " * they were lost in preprocessing — check that `external_meds_dir` carries null-time rows and " + "that any `pipeline=` you passed preserves them." + ) + + def build_datamodule(cfg: DictConfig) -> pl.LightningDataModule: """Instantiate the meds-torch-data ``Datamodule`` from a resolved ``cfg.datamodule`` node. Thin wrapper around ``hydra.utils.instantiate`` kept as a single choke-point so every command - constructs the datamodule the same way (and so tests can monkeypatch it). + constructs the datamodule the same way (and so tests can monkeypatch it). Being the one place every + command builds a datamodule is also why :func:`require_statics_if_requested` runs here. """ from hydra.utils import instantiate + require_statics_if_requested(cfg) return instantiate(cfg.datamodule) diff --git a/template/src/meds_model_base/testing/__init__.py b/template/src/meds_model_base/testing/__init__.py index 8e19331..00e1836 100644 --- a/template/src/meds_model_base/testing/__init__.py +++ b/template/src/meds_model_base/testing/__init__.py @@ -7,6 +7,12 @@ - :func:`build_signal_dataset` — a classifier signal: a marker code deterministically predicts the label. - :func:`build_pattern_dataset` — a generative signal: a fixed repeating code pattern. + +Both give every subject the static (baseline) measurements in :data:`STATIC_CODES`, drawn independently of +the labels. That is not decoration: a cohort with no static measurements at all cannot be collated by +meds-torch-data under ``static_inclusion_mode: INCLUDE`` — it raises ``Cannot infer dtype from empty +values`` before the model is called — so without them this suite could not exercise any model that reads +baseline data. - :func:`binary_auroc`, :func:`assert_learns_signal` — learnability assertions (+ negative control). - :func:`skip_if_stub` — skip a conformance test while ``model.py`` is still the generated stub. - :func:`run_chain`, :func:`build_workspace` — drive a model's DAG from its ``COMMANDS`` registry. @@ -15,10 +21,20 @@ from .harness import build_workspace, run_chain, run_cli, supported_sources from .property import assert_learns_signal, binary_auroc from .stub import is_stub, skip_if_stub -from .synthetic import SIGNAL_CODE, build_pattern_dataset, build_signal_dataset +from .synthetic import ( + SIGNAL_CODE, + STATIC_CODES, + STATIC_GROUPS, + STATIC_NUMERIC, + build_pattern_dataset, + build_signal_dataset, +) __all__ = [ "SIGNAL_CODE", + "STATIC_CODES", + "STATIC_GROUPS", + "STATIC_NUMERIC", "assert_learns_signal", "binary_auroc", "build_workspace", diff --git a/template/src/meds_model_base/testing/synthetic.py b/template/src/meds_model_base/testing/synthetic.py index 39997bd..6419b41 100644 --- a/template/src/meds_model_base/testing/synthetic.py +++ b/template/src/meds_model_base/testing/synthetic.py @@ -32,6 +32,32 @@ #: Fixed programs for a generative "pattern" dataset (a tiny grammar; cf. MEDS-EIC-AR). PATTERN_PROGRAMS = {"A": ("P//A0", "P//A1", "P//A2"), "B": ("P//B0", "P//B1")} +#: Static (baseline) measurements every synthetic subject gets: one value-less code drawn from +#: :data:`STATIC_GROUPS` and one numeric :data:`STATIC_NUMERIC`. Both branches of a model's static handling +#: are therefore covered, and a cohort built here is a cohort meds-torch-data can collate under +#: ``static_inclusion_mode: INCLUDE``. +STATIC_GROUPS = ("BASELINE//GROUP_A", "BASELINE//GROUP_B") +STATIC_NUMERIC = "BASELINE//AGE" +STATIC_CODES = (*STATIC_GROUPS, STATIC_NUMERIC) + + +def _static_rows(subject_id: int, rng: random.Random) -> list[dict]: + """Static measurements for one subject: MEDS marks them with a null ``time``. + + ``rng`` must be independent of the generator that decides the labels. A static that correlated with the + outcome would be a second signal, which would make the designed-signal test pass for the wrong reason + and give the negative control something real to learn. + """ + return [ + {"subject_id": subject_id, "time": None, "code": rng.choice(STATIC_GROUPS), "numeric_value": None}, + { + "subject_id": subject_id, + "time": None, + "code": STATIC_NUMERIC, + "numeric_value": rng.uniform(20.0, 90.0), + }, + ] + def _write_meds( root: Path, @@ -82,6 +108,7 @@ def build_signal_dataset( signal_rate: float = 0.5, seed: int = 0, shuffle_labels: bool = False, + with_statics: bool = True, ) -> Path: """Write a classifier-signal MEDS dataset: the presence of :data:`SIGNAL_CODE` determines the label. @@ -91,13 +118,24 @@ def build_signal_dataset( control: nothing to learn, AUROC ≈ 0.5). **Only the signal code is informative.** Sequence length and the signal's position within the sequence - are both label-independent by construction; see the comment in the loop for why that matters. + are both label-independent by construction; see the comment in the loop for why that matters. The + static measurements are drawn from a separate generator for the same reason. + + ``with_statics`` (default on) gives every subject the baseline measurements described by + :data:`STATIC_CODES`. It defaults on because a cohort *without* them is not merely a simpler cohort: a + model whose datamodule sets ``static_inclusion_mode: INCLUDE`` cannot collate it at all — + meds-torch-data raises ``ValueError: Cannot infer dtype from empty values`` when the static tensors are + empty for a whole batch, before the model is ever called. A learnability suite that cannot exercise a + model's static path is not testing that model. Turn it off only to reproduce that case. Returns ``root`` (a MEDS dataset with a ``signal_task`` task-labels directory). """ rng = random.Random(seed) + # Independent of `rng`, so nothing about a subject's baseline variables can carry information about + # its label — the signal code stays the only thing there is to learn. + static_rng = random.Random(seed + 9973) counts = {train_split: n_train, tuning_split: n_tuning, held_out_split: n_held_out} - codes = [*_BACKGROUND, SIGNAL_CODE] + codes = [*_BACKGROUND, SIGNAL_CODE, *(STATIC_CODES if with_statics else ())] per_split: dict[str, list[dict]] = {} task_labels: dict[str, pl.DataFrame] = {} @@ -106,6 +144,9 @@ def build_signal_dataset( rows: list[dict] = [] labels: list[dict] = [] for _ in range(n): + if with_statics: + # MEDS puts a subject's static measurements first, marked by a null time. + rows.extend(_static_rows(subject_id, static_rng)) has_signal = rng.random() < signal_rate t = _BASE_TIME + timedelta(days=int(rng.random() * 100)) # The signal is *inserted at a random position* and a negative subject gets a filler event, so @@ -150,16 +191,30 @@ def build_signal_dataset( def build_pattern_dataset( - root: Path, *, n_train: int = 200, n_tuning: int = 40, n_held_out: int = 40, seed: int = 0 + root: Path, + *, + n_train: int = 200, + n_tuning: int = 40, + n_held_out: int = 40, + seed: int = 0, + with_statics: bool = True, ) -> Path: """Write a generative-signal MEDS dataset: each subject is a run of fixed code *programs*. Used to test autoregressive generation: a model that learns the grammar should generate valid programs. Also emits a ``pattern_task`` index (one prediction_time per subject) for the generation entry point. + + ``with_statics`` carries the same meaning and the same default as in :func:`build_signal_dataset`: the + baseline measurements are what makes the cohort collatable under ``static_inclusion_mode: INCLUDE``. + They are not part of the grammar and are never emitted into the dynamic sequence, so a model generating + programs does not have to generate them. """ rng = random.Random(seed) + static_rng = random.Random(seed + 9973) counts = {train_split: n_train, tuning_split: n_tuning, held_out_split: n_held_out} codes = sorted({c for prog in PATTERN_PROGRAMS.values() for c in prog}) + if with_statics: + codes += list(STATIC_CODES) per_split: dict[str, list[dict]] = {} task_labels: dict[str, pl.DataFrame] = {} @@ -168,6 +223,8 @@ def build_pattern_dataset( rows: list[dict] = [] labels: list[dict] = [] for _ in range(n): + if with_statics: + rows.extend(_static_rows(subject_id, static_rng)) t = _BASE_TIME + timedelta(days=int(rng.random() * 100)) seq: list[str] = [] for _ in range(rng.randint(3, 6)): diff --git a/template/tests/test_synthetic_statics.py b/template/tests/test_synthetic_statics.py new file mode 100644 index 0000000..4ac890e --- /dev/null +++ b/template/tests/test_synthetic_statics.py @@ -0,0 +1,153 @@ +"""The synthetic cohorts carry static measurements, and a batch built from them can be collated. + +Needs no model, so this runs from the moment a repository is generated. + +The bug it pins is quiet in the worst way: it does not show up until the *slow* tier, and only for models +that read baseline data. `build_signal_dataset` used to emit nothing but timestamped events, so a cohort +built from it had no static measurements at all — and meds-torch-data cannot collate that under +`static_inclusion_mode: INCLUDE`: + + File ".../meds_torchdata/pytorch_dataset.py", line 1822, in collate + static_data = JointNestedRaggedTensorDict( + ValueError: Cannot infer dtype from empty values; provide an explicit `schema=`. + +The model is never called, so the failure names the collator rather than the fixture that caused it. A +model whose datamodule sets `INCLUDE` — anything that reads age, sex, or any other baseline variable — +would pass every fast test and then fail the one tier that proves it learns. +""" + +import polars as pl +import pytest +from meds_model_base.schemas import code_metadata_filepath +from meds_model_base.testing import ( + STATIC_CODES, + STATIC_GROUPS, + STATIC_NUMERIC, + build_signal_dataset, + build_workspace, +) + +#: Small enough to tensorize in a couple of seconds; large enough for both static groups to appear on both +#: sides of the label, which is what the leakage check below needs. +COHORT = {"n_train": 24, "n_tuning": 8, "n_held_out": 8} + + +@pytest.fixture(scope="module") +def raw(tmp_path_factory): + return build_signal_dataset(tmp_path_factory.mktemp("statics") / "raw", seed=0, **COHORT) + + +def _events(root) -> pl.DataFrame: + return pl.concat([pl.read_parquet(p) for p in sorted((root / "data").rglob("*.parquet"))]) + + +def test_every_subject_has_both_kinds_of_static_measurement(raw): + """One value-less code and one numeric one, so both branches of a model's static handling are covered.""" + events = _events(raw) + statics = events.filter(pl.col("time").is_null()) + assert len(statics) > 0, "the cohort has no static measurements at all" + + subjects = set(events["subject_id"].to_list()) + assert set(statics["subject_id"].to_list()) == subjects, "not every subject has static measurements" + + by_subject = statics.group_by("subject_id").agg(pl.col("code")) + for codes in by_subject["code"].to_list(): + assert any(code in STATIC_GROUPS for code in codes), f"no categorical static among {codes}" + assert STATIC_NUMERIC in codes, f"no numeric static among {codes}" + + numeric = statics.filter(pl.col("code") == STATIC_NUMERIC) + assert numeric["numeric_value"].null_count() == 0, "the numeric static carries no value" + assert statics.filter(pl.col("code").is_in(STATIC_GROUPS))["numeric_value"].null_count() == len( + subjects + ), "the categorical static should carry no value" + + +def test_static_codes_are_declared_in_the_metadata(raw): + """A code absent from `metadata/codes.parquet` is a code the vocabulary will not have.""" + declared = set(pl.read_parquet(raw / code_metadata_filepath)["code"].to_list()) + assert set(STATIC_CODES) <= declared + + +def test_statics_carry_no_information_about_the_label(raw): + """The baseline variables must not be a second signal. + + The designed-signal dataset exists to check that a model reads `SIGNAL_CODE`. A static that correlated + with the outcome would let a model score well without ever reading it — and would give the negative + control something real to learn, which is exactly the vacuous pass the control exists to rule out. + """ + labels = pl.concat( + [pl.read_parquet(p) for p in sorted((raw / "task_labels" / "signal_task").glob("*.parquet"))] + ) + groups = _events(raw).filter(pl.col("code").is_in(STATIC_GROUPS)).select("subject_id", "code") + joined = labels.join(groups, on="subject_id", how="inner") + + # Perfect separation would mean the group *is* the label. Both groups appearing on both sides is a + # cheap, non-flaky refutation of that. + seen = {(row["code"], row["boolean_value"]) for row in joined.iter_rows(named=True)} + for code in STATIC_GROUPS: + assert (code, True) in seen, f"{code} never occurs with a positive label" + assert (code, False) in seen, f"{code} never occurs with a negative label" + + +def test_a_cohort_without_statics_is_refused_before_training(tmp_path_factory): + """The half the fixture change does not fix: a *real* cohort with no static measurements. + + Giving the synthetic builders statics stops this suite tripping the collate crash, but it does nothing + for a user whose dataset genuinely has no baseline variables — they still hit + ``Cannot infer dtype from empty values`` inside the dataloader, naming neither the config key that + asked for static data nor the cohort that lacks it. `build_datamodule` is the one place every command + constructs a datamodule, so the mismatch is refused there instead, with both ways out named. + """ + from meds_model_base.lightning import build_datamodule, require_statics_if_requested + from omegaconf import OmegaConf + + root = tmp_path_factory.mktemp("no_statics") + bare = build_signal_dataset(root / "raw", seed=0, with_statics=False, **COHORT) + patients = build_workspace(bare, root / "data") / "patients" + + cfg = OmegaConf.create( + { + "datamodule": { + "config": {"tensorized_cohort_dir": str(patients), "static_inclusion_mode": "INCLUDE"} + } + } + ) + with pytest.raises(ValueError, match=r"no subject in .* has any static measurement"): + require_statics_if_requested(cfg) + # …and it is reached through the function every command actually calls, not only directly. + with pytest.raises(ValueError, match=r"static_inclusion_mode=OMIT"): + build_datamodule(cfg) + + # OMIT is the documented way out, and must not be blocked by the check. + cfg.datamodule.config.static_inclusion_mode = "OMIT" + require_statics_if_requested(cfg) + + +def test_a_batch_collates_with_statics_included(raw, tmp_path): + """The regression itself: `INCLUDE` must survive collation, and the tensors must be populated. + + This is the assertion the crash would fail. It runs the real preprocessing, so it also proves the + static measurements survive tensorization rather than merely existing in the raw parquet. + """ + import torch + from meds_torchdata import MEDSPytorchDataset, MEDSTorchDataConfig + from meds_torchdata.types import StaticInclusionMode, SubsequenceSamplingStrategy + + patients = build_workspace(raw, tmp_path / "data") / "patients" + dataset = MEDSPytorchDataset( + MEDSTorchDataConfig( + tensorized_cohort_dir=str(patients), + max_seq_len=32, + seq_sampling_strategy=SubsequenceSamplingStrategy.TO_END, + static_inclusion_mode=StaticInclusionMode.INCLUDE, + ), + split="train", + ) + loader = torch.utils.data.DataLoader(dataset, batch_size=4, collate_fn=dataset.collate, num_workers=0) + batch = next(iter(loader)) + + assert batch.static_code is not None, "the batch carries no static tensors" + assert batch.static_code.numel() > 0, "the static tensors are empty" + assert (batch.static_code != 0).any(), "every static code is padding" + # The numeric baseline variable must arrive as a value, not merely as a code. + assert batch.static_numeric_value_mask.any(), "no static measurement carries a numeric value" diff --git a/tests/test_render.py b/tests/test_render.py index 374f4bb..797a8ab 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -114,6 +114,7 @@ def test_render_profile(tmp_path, profile, commands): "src/meds_model_base/commands/base.py", "tests/test_smoke_pipeline.py", "tests/test_meds_dev_e2e.py", + "tests/test_synthetic_statics.py", ]: assert (dst / rel).exists(), f"missing rendered file: {rel}"