Fix static_inclusion_mode=INCLUDE against a cohort with no statics (fixes #16)

template/src/meds_model_base/lightning/__init__.py

@@ -2,14 +2,16 @@
 
 - :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).
 """
 
 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)
 
 

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",

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