diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2ddad4ca..b49edf7c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,6 @@ -# Third-party LaTeX template files under paper/ are kept verbatim. -exclude: ^paper/(sn-jnl\.cls|sn-nature\.bst|template-.*|main\.pdf)$ +# Kept verbatim, never rewritten by hooks: third-party LaTeX template files +# under paper/, and scripts/gemini/out/ (GEMINI's own exported reports). +exclude: ^(paper/(sn-jnl\.cls|sn-nature\.bst|template-.*|main\.pdf)|scripts/gemini/out/.*)$ repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v6.0.0 # Use the ref you want to point at @@ -42,7 +43,10 @@ repos: entry: python3 -m mypy --config-file pyproject.toml language: system types: [python] - exclude: 'tests|scripts/cohort/cohort_check_(mimic|eicu)\.py' + # tests, plus the cohort producers committed verbatim as they ran + # (scripts/cohort/README.md); pyproject's mypy exclude does not apply + # to files pre-commit passes by name. + exclude: ^(tests/|scripts/cohort/cohort_check_(mimic|eicu)\.py$) - repo: https://github.com/crate-ci/typos rev: v1.50.2 diff --git a/_typos.toml b/_typos.toml index d2699e0a..32e84fb5 100644 --- a/_typos.toml +++ b/_typos.toml @@ -43,4 +43,10 @@ get_rolling_window_indicies = "get_rolling_window_indicies" # ours to rename and will keep showing up as new false positives every time # a fresh report is committed, so the whole directory is excluded rather # than growing an extend-words entry per column. -extend-exclude = ["scripts/gemini/out/", "scripts/cohort/cohort_check_mimic.py", "scripts/cohort/cohort_check_eicu.py"] +# scripts/cohort/cohort_check_*.py are committed verbatim as they ran +# (scripts/cohort/README.md); their `evn` (eICU visit number) stays. +extend-exclude = [ + "scripts/gemini/out/", + "scripts/cohort/cohort_check_mimic.py", + "scripts/cohort/cohort_check_eicu.py", +] diff --git a/odyssey/inference/run_inference.py b/odyssey/inference/run_inference.py index ed69d62e..f8c02fb8 100644 --- a/odyssey/inference/run_inference.py +++ b/odyssey/inference/run_inference.py @@ -410,6 +410,18 @@ def load_run( config.event_head_hidden = ( int(first_layer.shape[0]) if first_layer is not None else 0 ) + # Summary head (training-only): rebuild it at the checkpoint's own width + # so the weights load; it is never read at inference. + summary_out = state.get("summary_head.proj.weight") + summary_hidden = state.get("summary_head.proj.0.weight") + if summary_out is None and summary_hidden is not None: + summary_out = state.get("summary_head.proj.2.weight") + config.summary_num_targets = ( + int(summary_out.shape[0]) if summary_out is not None else 0 + ) + config.summary_head_hidden = ( + int(summary_hidden.shape[0]) if summary_hidden is not None else 0 + ) concepts = concepts_for_source( getattr(config, "source", "mimic_iv"), diff --git a/odyssey/models/sequence_model.py b/odyssey/models/sequence_model.py index 30a063ee..87f648b0 100644 --- a/odyssey/models/sequence_model.py +++ b/odyssey/models/sequence_model.py @@ -57,6 +57,7 @@ fold_in_bottleneck_losses, ) from odyssey.models.injection import stream_injection +from odyssey.models.summary_head import SummaryHead, masked_huber_loss from odyssey.models.time_to_event import ( EventHazardHeads, TimeToEventHead, @@ -74,6 +75,7 @@ if TYPE_CHECKING: # the targets live in training; avoid a runtime import cycle from odyssey.training.event_targets import EventHazardTargets + from odyssey.training.summary_targets import SummaryTargets @dataclass @@ -122,6 +124,12 @@ class ForecastObjective: time_weight: float = 0.0 event_hazard_weight: float = 0.0 value_head_weight: float = 0.0 + summary_weight: float = 0.0 + """Weight of the self-supervised window-summary loss + (:mod:`odyssey.models.summary_head`), if the model has a summary head; + the targets come per chunk from :mod:`odyssey.training.summary_targets`.""" + summary_target_weights: torch.Tensor | None = None + """Optional ``(K,)`` per-target weights inside the summary loss.""" class ForwardWithFeatures(NamedTuple): @@ -421,6 +429,23 @@ def _streaming_event_loss( event_heads.edges, ) + def _streaming_summary_loss( + self, + summary_head: SummaryHead | None, + features: torch.Tensor, + summary_targets: Optional["SummaryTargets"], + target_weights: torch.Tensor | None = None, + ) -> torch.Tensor: + """Masked Huber loss of the window-summary head (zero-graph if absent).""" + if summary_head is None or summary_targets is None: + return features.sum() * 0.0 + return masked_huber_loss( + summary_head(features), + summary_targets.values, + summary_targets.mask, + target_weights=target_weights, + ) + def _streaming_time_loss( self, time_head: TimeToEventHead | None, @@ -486,6 +511,8 @@ def __init__( value_head: bool = False, value_head_hidden: int = 0, source: str = "mimic_iv", + summary_targets: int = 0, + summary_head_hidden: int = 0, ) -> None: """Initialize the baseline sequence model. @@ -524,6 +551,11 @@ def __init__( if value_head else None ) + self.summary_head: SummaryHead | None = ( + SummaryHead(head_in, summary_targets, hidden_size=summary_head_hidden) + if summary_targets > 0 + else None + ) def forward_features( self, @@ -580,6 +612,7 @@ def compute_streaming_loss( *, objective: ForecastObjective | None = None, event_targets: Optional["EventHazardTargets"] = None, + summary_targets: Optional["SummaryTargets"] = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor], TimeAwareState]: """Compute the forecasting loss over one packed, chunked training step.""" objective = objective or ForecastObjective() @@ -590,11 +623,18 @@ def compute_streaming_loss( time_loss, _ = self._streaming_time_loss(self.time_head, hidden, chunk) event_loss = self._streaming_event_loss(self.event_heads, hidden, event_targets) value_loss, _ = self._streaming_value_loss(self.value_head, hidden, chunk) + summary_loss = self._streaming_summary_loss( + self.summary_head, + hidden, + summary_targets, + target_weights=objective.summary_target_weights, + ) total = ( task_loss + objective.time_weight * time_loss + objective.event_hazard_weight * event_loss + objective.value_head_weight * value_loss + + objective.summary_weight * summary_loss ) return ( total, @@ -603,6 +643,7 @@ def compute_streaming_loss( "time_loss": time_loss.detach(), "event_loss": event_loss.detach(), "value_loss": value_loss.detach(), + "summary_loss": summary_loss.detach(), }, new_state, ) @@ -632,6 +673,8 @@ def __init__( value_head: bool = False, value_head_hidden: int = 0, source: str = "mimic_iv", + summary_targets: int = 0, + summary_head_hidden: int = 0, ) -> None: """Initialize the concept-bottleneck sequence model. @@ -700,6 +743,11 @@ def __init__( if value_head else None ) + self.summary_head: SummaryHead | None = ( + SummaryHead(head_in, summary_targets, hidden_size=summary_head_hidden) + if summary_targets > 0 + else None + ) def forward( self, @@ -815,6 +863,7 @@ def compute_steering_loss( lifted_ids: torch.Tensor, objective: ForecastObjective | None = None, event_targets: Optional["EventHazardTargets"] = None, + summary_targets: Optional["SummaryTargets"] = None, respond_weight: float = 1.0, express_weight: float = 1.0, forecast_at_injected: bool = True, @@ -869,11 +918,18 @@ def compute_steering_loss( self.event_heads, head_feats, event_targets ) value_loss, _ = self._streaming_value_loss(self.value_head, head_feats, scored) + summary_loss = self._streaming_summary_loss( + self.summary_head, + head_feats, + summary_targets, + target_weights=objective.summary_target_weights, + ) forecast_loss = ( next_token_loss + objective.time_weight * time_loss + objective.event_hazard_weight * event_loss + objective.value_head_weight * value_loss + + objective.summary_weight * summary_loss ) at = injected & chunk.real_mask if bool(at.any()): @@ -895,6 +951,7 @@ def compute_steering_loss( "time_loss": time_loss.detach(), "event_loss": event_loss.detach(), "value_loss": value_loss.detach(), + "summary_loss": summary_loss.detach(), "respond_loss": respond.detach(), "express_loss": express.detach(), "n_injected": at.sum().detach(), @@ -976,6 +1033,7 @@ def compute_streaming_loss( intervention: BottleneckIntervention | None = None, objective: ForecastObjective | None = None, event_targets: Optional["EventHazardTargets"] = None, + summary_targets: Optional["SummaryTargets"] = None, teacher_alpha_known: float = 0.0, teacher_alpha_unknown: float = 0.0, ) -> tuple[torch.Tensor, dict[str, torch.Tensor], TimeAwareState]: @@ -1043,11 +1101,18 @@ def compute_streaming_loss( self.event_heads, head_feats, event_targets ) value_loss, _ = self._streaming_value_loss(self.value_head, head_feats, chunk) + summary_loss = self._streaming_summary_loss( + self.summary_head, + head_feats, + summary_targets, + target_weights=objective.summary_target_weights, + ) forecast_loss = ( next_token_loss + objective.time_weight * time_loss + objective.event_hazard_weight * event_loss + objective.value_head_weight * value_loss + + objective.summary_weight * summary_loss ) pool_mask = chunk.patient_end if supervision == "stay" else chunk.visit_end @@ -1058,6 +1123,7 @@ def compute_streaming_loss( "time_loss": time_loss.detach(), "event_loss": event_loss.detach(), "value_loss": value_loss.detach(), + "summary_loss": summary_loss.detach(), "concept_loss": zero, "orthogonality_loss": zero, "observability_loss": zero, @@ -1127,4 +1193,5 @@ def compute_streaming_loss( components["time_loss"] = time_loss.detach() components["event_loss"] = event_loss.detach() components["value_loss"] = value_loss.detach() + components["summary_loss"] = summary_loss.detach() return total, components, new_state diff --git a/odyssey/models/summary_head.py b/odyssey/models/summary_head.py new file mode 100644 index 00000000..bb327e0c --- /dev/null +++ b/odyssey/models/summary_head.py @@ -0,0 +1,74 @@ +"""Summary head: report window statistics of the chart from the state. + +An auxiliary, training-only readout. At landmark positions it predicts the +standardized window-summary panel of +:mod:`odyssey.training.summary_targets` (per-signal window min/max/mean, +change from the visit's first value, per-family occurrence counts) from +the same features the hazard heads read. The loss is a masked Huber loss: +robust to the residual heavy tails of clinical values, and zero (with a +live graph) where no target applies, so it can be summed unconditionally. +The head is never used at inference; it exists to shape the state. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import nn + + +class SummaryHead(nn.Module): + """Features -> one standardized value per summary target.""" + + def __init__( + self, in_features: int, num_targets: int, hidden_size: int = 0 + ) -> None: + """Initialize a linear readout, or a GELU MLP when ``hidden_size`` > 0.""" + super().__init__() + self.num_targets = int(num_targets) + self.hidden_size = int(hidden_size) + self.proj: nn.Module = ( + nn.Sequential( + nn.Linear(in_features, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, self.num_targets), + ) + if self.hidden_size > 0 + else nn.Linear(in_features, self.num_targets) + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + """Return ``(..., num_targets)`` predictions.""" + out: torch.Tensor = self.proj(features) + return out + + +def masked_huber_loss( + prediction: torch.Tensor, + target: torch.Tensor, + mask: torch.Tensor, + *, + delta: float = 1.0, + target_weights: torch.Tensor | None = None, +) -> torch.Tensor: + """Huber loss averaged over the ``True`` entries of ``mask``. + + ``prediction``/``target``/``mask`` share the shape ``(..., K)``. + ``target_weights`` (``(K,)``, optional) reweights the per-target terms + inside the average, so targets the state finds hardest (changes from + baseline, counts) can carry more of the gradient than the levels it + learns anyway. Returns ``0 * prediction.sum()`` when the mask is empty, + so the result always carries a graph and can be added unconditionally. + """ + if not bool(mask.any()): + return prediction.sum() * 0.0 + per_entry = F.huber_loss( + prediction.float(), target.float(), reduction="none", delta=delta + ) + weights = mask.to(per_entry.dtype) + if target_weights is not None: + weights = weights * target_weights.to(per_entry.dtype).to(weights.device) + return (per_entry * weights).sum() / weights.sum() + + +__all__ = ["SummaryHead", "masked_huber_loss"] diff --git a/odyssey/training/summary_targets.py b/odyssey/training/summary_targets.py new file mode 100644 index 00000000..ca026ac7 --- /dev/null +++ b/odyssey/training/summary_targets.py @@ -0,0 +1,413 @@ +"""Self-supervised window-summary targets for the summary head. + +The tuned GBM the hazard heads are compared with is hand-fed windowed +statistics of the chart: per-signal minimum, maximum and mean over 6 h and +24 h, the change from the visit's first value, and per-family occurrence +counts. Frozen probes (2026-09-13, ``scripts/probe_summary_signal.py``) +showed the backbone's state holds current levels but not changes from +baseline and not fine window counts, linearly or otherwise. This module +turns those statistics into an auxiliary training target: at 4-hourly +landmark positions the model must *report* them from its own state. The +values are computed from the input stream by the GBM's own feature code, +never fed in as inputs, and the head that predicts them is discarded at +inference. The state is shaped to hold what the statistics need; nothing +is handed to it. + +Pipeline: :func:`compute_summary_targets` runs once per shard (offline, +``scripts/build_summary_targets.py``) and writes one parquet per shard; +:func:`fit_summary_stats` standardizes columns over the training set +(counts through ``log1p``, everything winsorized at the 0.5/99.5 +percentiles so sentinel values cannot dominate); :class:`SummaryTargetTables` +loads the parquets and :func:`summary_targets_for_chunk` looks up, for +every streaming chunk, which positions carry a target -- the last token of +the bundle at each landmark row -- and returns a ``(lanes, T, K)`` tensor +with a mask. Times are hours on the sequence origin, the same clock as +chunk time stamps and the alert harness' landmark rows. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import polars as pl +import torch + +from odyssey.data.alert_events import origin_hours +from odyssey.data.signal_panel import SIGNAL_PANEL +from odyssey.data.streaming import StreamingChunk +from odyssey.inference.baseline_features import ( + DRUG_CLASSES, + FAMILY_LABELS, + StrongFeatureBuilder, + feature_names, +) + + +# Window statistics asked for per panel signal: the GBM's ``summary_stats`` +# group without ``last``/``hours_since_last`` (recency probes already showed +# the state holds those) and without ``delta_prev``/``ratio_visit_min`` +# (functions of the others). +SIGNAL_STATS: tuple[str, ...] = ( + "min_6h", + "max_6h", + "min_24h", + "max_24h", + "mean_24h", + "delta_visit_first", +) + +# Occurrence counts asked for per drug class and per code family: the GBM's +# ``counts_occurrence`` group at its two window lengths. +COUNT_STATS: tuple[str, ...] = ("n_6h", "n_24h") + +DEFAULT_LANDMARK_HOURS = 4.0 +WINSOR_PERCENTILES = (0.5, 99.5) +KEY_COLUMNS = ("subject_id", "visit_id", "time_hours") + + +def summary_target_names() -> list[str]: + """Column names of the target panel, in the order the head predicts them.""" + names = [f"{label}.{stat}" for label, _ in SIGNAL_PANEL for stat in SIGNAL_STATS] + names += [ + f"drug.{label}.{stat}" for label, _ in DRUG_CLASSES for stat in COUNT_STATS + ] + names += [ + f"family.{label}.{stat}" for label in FAMILY_LABELS for stat in COUNT_STATS + ] + return names + + +def summary_target_weights(change_weight: float = 1.0) -> torch.Tensor: + """Return ``(K,)`` per-target loss weights. + + ``change_weight`` on the ``delta_visit_first`` and count targets, 1 on + the window-level targets. + """ + names = summary_target_names() + hard = [ + n.rsplit(".", 1)[-1] == "delta_visit_first" + or n.rsplit(".", 1)[-1] in COUNT_STATS + for n in names + ] + return torch.where( + torch.tensor(hard), torch.tensor(float(change_weight)), torch.tensor(1.0) + ) + + +def count_target_mask(names: Sequence[str] | None = None) -> np.ndarray: + """Boolean mask over the panel: which targets are occurrence counts.""" + names = list(names) if names is not None else summary_target_names() + return np.array([n.rsplit(".", 1)[-1] in COUNT_STATS for n in names]) + + +def landmark_rows( + events: pl.DataFrame, landmark_hours: float = DEFAULT_LANDMARK_HOURS +) -> tuple[list[int], list[int], list[float]]: + """``(subject_ids, visit_ids, times)`` of the landmark rows of every visit. + + Every ``landmark_hours`` from a visit's first event to its last, the + row is the visit's last event time at or before that instant, so a row + time always coincides with a real token's time stamp (hours since the + subject's first event, like :class:`~odyssey.data.sequences.PatientSequence`). + """ + origins = origin_hours(events) + timed = ( + events.filter(pl.col("time").is_not_null() & pl.col("hadm_id").is_not_null()) + .join(origins, on="subject_id", how="left") + .with_columns( + ((pl.col("time") - pl.col("_origin")).dt.total_seconds() / 3600.0).alias( + "_hours" + ) + ) + .group_by("subject_id", "hadm_id") + .agg(pl.col("_hours").unique().sort().alias("_hours")) + ) + sids: list[int] = [] + vids: list[int] = [] + times: list[float] = [] + for sid, vid, hours in zip( + timed["subject_id"].to_list(), + timed["hadm_id"].to_list(), + timed["_hours"].to_list(), + ): + arr = np.asarray(hours, dtype=np.float64) + grid = np.arange(arr[0], arr[-1] + 1e-9, landmark_hours) + idx = np.searchsorted(arr, grid, side="right") - 1 + picked = np.unique(arr[idx]) + sids.extend([int(sid)] * len(picked)) + vids.extend([int(vid)] * len(picked)) + times.extend(float(t) for t in picked) + return sids, vids, times + + +def compute_summary_targets( + events: pl.DataFrame, + *, + source: str = "mimic_iv", + landmark_hours: float = DEFAULT_LANDMARK_HOURS, +) -> pl.DataFrame: + """Raw (unstandardized) targets at every landmark row of ``events``. + + Columns: ``subject_id``, ``visit_id``, ``time_hours`` and one column per + :func:`summary_target_names` entry; NaN where a window is empty. + """ + names = summary_target_names() + all_names = feature_names() + idx = [all_names.index(n) for n in names] + sids, vids, times = landmark_rows(events, landmark_hours) + if not sids: + return pl.DataFrame( + schema={ + "subject_id": pl.Int64, + "visit_id": pl.Int64, + "time_hours": pl.Float64, + **dict.fromkeys(names, pl.Float32), + } + ) + builder = StrongFeatureBuilder(events, source=source) + values = builder.features(sids, vids, times)[:, idx].astype(np.float32) + frame = pl.DataFrame( + { + "subject_id": pl.Series(sids, dtype=pl.Int64), + "visit_id": pl.Series(vids, dtype=pl.Int64), + "time_hours": pl.Series(times, dtype=pl.Float64), + } + ) + return frame.with_columns( + [pl.Series(n, values[:, i], dtype=pl.Float32) for i, n in enumerate(names)] + ) + + +@dataclass(frozen=True) +class SummaryTargetStats: + """Per-column standardization fitted on the training targets.""" + + names: list[str] + lo: np.ndarray + """Winsorizing floor per column (after ``log1p`` for counts).""" + hi: np.ndarray + mean: np.ndarray + std: np.ndarray + + def transform(self, raw: np.ndarray) -> np.ndarray: + """Standardize raw targets ``(n, K)``; NaN stays NaN.""" + x = np.asarray(raw, dtype=np.float64).copy() + counts = count_target_mask(self.names) + x[:, counts] = np.log1p(np.clip(x[:, counts], 0.0, None)) + x = np.clip(x, self.lo, self.hi) + out: np.ndarray = ((x - self.mean) / self.std).astype(np.float32) + return out + + def save(self, path: str | Path) -> None: + """Write the stats as JSON.""" + Path(path).write_text( + json.dumps( + { + "names": self.names, + "lo": self.lo.tolist(), + "hi": self.hi.tolist(), + "mean": self.mean.tolist(), + "std": self.std.tolist(), + }, + indent=1, + ) + ) + + @classmethod + def load(cls, path: str | Path) -> SummaryTargetStats: + """Read stats written by :meth:`save`.""" + d = json.loads(Path(path).read_text()) + return cls( + names=list(d["names"]), + lo=np.asarray(d["lo"], dtype=np.float64), + hi=np.asarray(d["hi"], dtype=np.float64), + mean=np.asarray(d["mean"], dtype=np.float64), + std=np.asarray(d["std"], dtype=np.float64), + ) + + +def fit_summary_stats(frames: Iterable[pl.DataFrame]) -> SummaryTargetStats: + """Fit :class:`SummaryTargetStats` over the raw target frames of a split. + + Counts go through ``log1p``; every column is winsorized at + :data:`WINSOR_PERCENTILES` of its finite values before mean/std, so a + handful of sentinel readings cannot set the scale. A column with no + finite value, or no spread, gets mean 0 and std 1. + """ + names = summary_target_names() + counts = count_target_mask(names) + parts = [f.select(names).to_numpy().astype(np.float64) for f in frames] + x = np.concatenate(parts, axis=0) if parts else np.zeros((0, len(names))) + x[:, counts] = np.log1p(np.clip(x[:, counts], 0.0, None)) + lo = np.zeros(len(names)) + hi = np.ones(len(names)) + mean = np.zeros(len(names)) + std = np.ones(len(names)) + for j in range(len(names)): + col = x[:, j] + col = col[np.isfinite(col)] + if col.size == 0: + continue + lo[j], hi[j] = np.percentile(col, WINSOR_PERCENTILES) + clipped = np.clip(col, lo[j], hi[j]) + mean[j] = clipped.mean() + s = clipped.std() + std[j] = s if s > 0 else 1.0 + return SummaryTargetStats(names=names, lo=lo, hi=hi, mean=mean, std=std) + + +class SummaryTargetTables: + """Standardized targets per ``(subject, visit)``, queried per chunk. + + Each key maps to ``(times, values)``: sorted landmark row times and the + matching ``(rows, K)`` standardized targets, stored as float16 to keep + a whole split (millions of rows) in memory. + """ + + def __init__(self, stats: SummaryTargetStats) -> None: + self.stats = stats + self.num_targets = len(stats.names) + self._rows: dict[tuple[int, int], tuple[np.ndarray, np.ndarray]] = {} + + def __len__(self) -> int: + """Return the number of visits with targets.""" + return len(self._rows) + + def add_frame(self, frame: pl.DataFrame) -> None: + """Ingest one raw target frame from :func:`compute_summary_targets`.""" + if frame.height == 0: + return + frame = frame.sort("subject_id", "visit_id", "time_hours") + values = self.stats.transform(frame.select(self.stats.names).to_numpy()) + keys = frame.select("subject_id", "visit_id").to_numpy() + times = frame["time_hours"].to_numpy().astype(np.float64) + change = np.flatnonzero(np.any(keys[1:] != keys[:-1], axis=1)) + 1 + starts = np.concatenate([[0], change]) + ends = np.concatenate([change, [len(keys)]]) + for a, b in zip(starts, ends): + key = (int(keys[a, 0]), int(keys[a, 1])) + self._rows[key] = (times[a:b], values[a:b].astype(np.float16)) + + @classmethod + def from_dir( + cls, directory: str | Path, stats: SummaryTargetStats + ) -> SummaryTargetTables: + """Load every ``*.parquet`` under ``directory``.""" + tables = cls(stats) + for path in sorted(Path(directory).glob("*.parquet")): + tables.add_frame(pl.read_parquet(path)) + return tables + + def lookup( + self, subject_id: int, visit_id: int + ) -> tuple[np.ndarray, np.ndarray] | None: + """Return ``(times, values)`` for one visit, or ``None``.""" + return self._rows.get((subject_id, visit_id)) + + +@dataclass +class SummaryTargets: + """``(lanes, T, K)`` standardized targets and where they apply.""" + + values: torch.Tensor + mask: torch.Tensor + + +def _bundle_end_mask( + time_stamps: np.ndarray, subject_ids: np.ndarray, real: np.ndarray +) -> np.ndarray: + """Mark the last real token of each same-time bundle, per lane.""" + lanes, chunk_len = time_stamps.shape + end = np.zeros_like(real) + if chunk_len == 0: + return end + same_next = np.zeros((lanes, chunk_len), dtype=bool) + same_next[:, :-1] = ( + (time_stamps[:, 1:] == time_stamps[:, :-1]) + & (subject_ids[:, 1:] == subject_ids[:, :-1]) + & real[:, 1:] + ) + return real & ~same_next + + +def summary_targets_for_chunk( + chunk: StreamingChunk, tables: SummaryTargetTables +) -> SummaryTargets | None: + """Targets for one chunk, or ``None`` when no position carries one. + + A position carries a target when its ``(subject, visit, time)`` is a + landmark row of the tables and it is the last real token of its + bundle, so the state has read the whole bundle the row summarizes. + """ + sids = chunk.subject_ids.detach().cpu().numpy() + vids = chunk.visit_ids.detach().cpu().numpy() + times = chunk.batch.aux.time_stamps.detach().cpu().numpy().astype(np.float64) + real = chunk.real_mask.detach().cpu().numpy().astype(bool) + lanes, chunk_len = sids.shape + k = tables.num_targets + values = np.zeros((lanes, chunk_len, k), dtype=np.float32) + mask = np.zeros((lanes, chunk_len), dtype=bool) + ends = _bundle_end_mask(times, sids, real) + keys = np.stack([sids.reshape(-1), vids.reshape(-1)], axis=1) + unique_keys, inverse = np.unique(keys, axis=0, return_inverse=True) + inverse = inverse.reshape(-1) + for i, (s, v) in enumerate(unique_keys.tolist()): + rows = tables.lookup(int(s), int(v)) + if rows is None: + continue + row_times, row_values = rows + flat = np.flatnonzero((inverse == i) & ends.reshape(-1)) + if flat.size == 0: + continue + t = times.reshape(-1)[flat] + pos = np.searchsorted(row_times, t) + pos = np.clip(pos, 0, len(row_times) - 1) + hit = np.isclose(row_times[pos], t, rtol=0.0, atol=1e-6) + if not hit.any(): + continue + lane_idx, pos_idx = np.divmod(flat[hit], chunk_len) + values[lane_idx, pos_idx] = row_values[pos[hit]].astype(np.float32) + mask[lane_idx, pos_idx] = True + if not mask.any(): + return None + device = chunk.batch.concept_ids.device + finite = np.isfinite(values) + full_mask = finite & mask[:, :, None] + values = np.where(full_mask, values, 0.0).astype(np.float32) + return SummaryTargets( + values=torch.from_numpy(values).to(device), + mask=torch.from_numpy(full_mask).to(device), + ) + + +def load_summary_tables( + directory: str | Path, stats_path: str | Path | None = None +) -> SummaryTargetTables: + """Load a split's tables; the stats default to ``/../stats.json``.""" + directory = Path(directory) + stats = SummaryTargetStats.load( + stats_path if stats_path is not None else directory.parent / "stats.json" + ) + return SummaryTargetTables.from_dir(directory, stats) + + +__all__ = [ + "COUNT_STATS", + "DEFAULT_LANDMARK_HOURS", + "SIGNAL_STATS", + "SummaryTargetStats", + "SummaryTargetTables", + "SummaryTargets", + "compute_summary_targets", + "count_target_mask", + "fit_summary_stats", + "landmark_rows", + "load_summary_tables", + "summary_target_names", + "summary_target_weights", + "summary_targets_for_chunk", +] diff --git a/odyssey/training/train.py b/odyssey/training/train.py index e3744fce..1ee6a1db 100644 --- a/odyssey/training/train.py +++ b/odyssey/training/train.py @@ -113,6 +113,13 @@ SteeringSchedule, choose_injection, ) +from odyssey.training.summary_targets import ( + SummaryTargetTables, + load_summary_tables, + summary_target_names, + summary_target_weights, + summary_targets_for_chunk, +) from odyssey.utils.env_fingerprint import write_run_provenance @@ -454,6 +461,41 @@ class TrainingConfig: value_head -- lets an A/B separate "better input encoding" from "better output objective".""" + summary_targets_dir: str | None = None + """Directory of precomputed self-supervised window-summary targets + (scripts/build_summary_targets.py: ``/train/*.parquet``, + ``/tuning/*.parquet``, ``/stats.json``). When set, a summary + head (odyssey.models.summary_head) is trained to report, at 4-hourly + landmark positions, the GBM's window statistics of the chart so far + (per-signal window min/max/mean, change from the visit's first value, + per-family occurrence counts) from the same features the hazard heads + read. The statistics are computed from the input and never fed in; + the head is unused at inference. Frozen probes on 2026-09-13 showed the + state lacks exactly these (change from baseline R^2 0.00), which is + the content of the GBM's margin on AKI and the counting events. None + (the default) reproduces every existing run.""" + + summary_weight: float = 0.5 + """Weight of the summary head's masked Huber loss.""" + + summary_head_hidden: int = 0 + """Hidden width of the summary head; 0 = linear, which asks for the + same linear readability the probes measure.""" + + summary_num_targets: int = 0 + """Width of the summary head, recorded at training time so a checkpoint + rebuilds the same head even if the target panel definition changes.""" + + summary_change_weight: float = 1.0 + """Multiplier, inside the summary loss, on the change-from-baseline + (``delta_visit_first``) and occurrence-count targets relative to the + window-level targets. The first arm (weight 0.5, all targets equal) + tripled the readability of window levels but left changes and counts + almost where they were (creatinine change from admission R^2 0.00 -> + 0.23 post-bottleneck, ~0 pre; vasopressor 6 h count unchanged); levels + are 241 of the 328 targets and won the average. > 1 shifts the gradient + to the targets the state does not hold.""" + randint_prob: float = 0.25 """Intervention-aware training (CEM's RandInt): at every training position, each observed concept's mixing probability is replaced by @@ -691,6 +733,8 @@ def build_model( value_head=bool(getattr(config, "value_head", False)), value_head_hidden=int(getattr(config, "value_head_hidden", 0) or 0), source=getattr(config, "source", "mimic_iv"), + summary_targets=int(getattr(config, "summary_num_targets", 0) or 0), + summary_head_hidden=int(getattr(config, "summary_head_hidden", 0) or 0), ) return ConceptBottleneckSequenceModel( backbone=backbone, @@ -710,6 +754,8 @@ def build_model( value_head=bool(getattr(config, "value_head", False)), value_head_hidden=int(getattr(config, "value_head_hidden", 0) or 0), source=getattr(config, "source", "mimic_iv"), + summary_targets=int(getattr(config, "summary_num_targets", 0) or 0), + summary_head_hidden=int(getattr(config, "summary_head_hidden", 0) or 0), ) @@ -978,6 +1024,18 @@ def build_objective( value_head_weight=( config.value_head_weight if getattr(config, "value_head", False) else 0.0 ), + summary_weight=( + config.summary_weight + if getattr(config, "summary_targets_dir", None) + else 0.0 + ), + summary_target_weights=( + summary_target_weights( + float(getattr(config, "summary_change_weight", 1.0)) + ).to(device) + if getattr(config, "summary_targets_dir", None) + else None + ), ) @@ -992,6 +1050,7 @@ def evaluate_streaming( supervision: ConceptSupervision = "stay", objective: ForecastObjective | None = None, event_tables: EventTimeTables | None = None, + summary_tables: SummaryTargetTables | None = None, ) -> dict[str, float]: """Average loss components over one (partial), gradient-free sampler pass.""" model.eval() @@ -1009,9 +1068,18 @@ def evaluate_streaming( if event_tables is not None else None ) + summary_targets = ( + summary_targets_for_chunk(chunk, summary_tables) + if summary_tables is not None + else None + ) if isinstance(model, BaselineSequenceModel): _, components, state = model.compute_streaming_loss( - chunk, state=state, objective=objective, event_targets=event_targets + chunk, + state=state, + objective=objective, + event_targets=event_targets, + summary_targets=summary_targets, ) else: _, components, state = model.compute_streaming_loss( @@ -1022,6 +1090,7 @@ def evaluate_streaming( supervision=supervision, objective=objective, event_targets=event_targets, + summary_targets=summary_targets, ) state = _detach_state(state) for key, value in components.items(): @@ -1042,6 +1111,8 @@ def _combined_val_loss( time_weight: float = 0.0, event_hazard_weight: float = 0.0, value_head_weight: float = 0.0, + *, + summary_weight: float = 0.0, ) -> float: """Compute the same task + weighted-auxiliary combination the training loss uses. @@ -1056,6 +1127,7 @@ def _combined_val_loss( + time_weight * components.get("time_loss", 0.0) + event_hazard_weight * components.get("event_loss", 0.0) + value_head_weight * components.get("value_loss", 0.0) + + summary_weight * components.get("summary_loss", 0.0) + weights.concept * components.get("concept_loss", 0.0) + weights.orthogonality * components.get("orthogonality_loss", 0.0) + weights.observability * components.get("observability_loss", 0.0) @@ -1078,6 +1150,9 @@ def train(config: TrainingConfig) -> Path: # noqa: PLR0912, PLR0915 config.event_hazards = any(k.startswith("event_heads.") for k in resume_keys) config.value_head = any(k.startswith("value_head.") for k in resume_keys) del resume_keys + config.summary_num_targets = ( + len(summary_target_names()) if config.summary_targets_dir else 0 + ) (output_dir / "config.json").write_text(json.dumps(asdict(config), indent=2)) _activate_run_sidecars(config) @@ -1224,6 +1299,7 @@ def make_train_patients(epoch: int) -> Iterator[PatientSequence]: shuffle_seed=config.seed + epoch, ) + train_summary_tables, tuning_summary_tables = _summary_tables(config) corpus = PreparedCorpus( vocab=vocab, concepts=concepts, @@ -1237,6 +1313,8 @@ def make_train_patients(epoch: int) -> Iterator[PatientSequence]: tuning_event_tables=tuning_event_tables, tuning_events_binned=tuning_events_binned, make_train_patients=make_train_patients, + train_summary_tables=train_summary_tables, + tuning_summary_tables=tuning_summary_tables, ) return _run_training(config, output_dir, device, corpus) @@ -1258,6 +1336,28 @@ class PreparedCorpus: tuning_events_binned: pl.DataFrame make_train_patients: Callable[[int], Iterator[PatientSequence]] """epoch -> the training patient stream for that epoch.""" + train_summary_tables: SummaryTargetTables | None = None + tuning_summary_tables: SummaryTargetTables | None = None + + +def _summary_tables( + config: TrainingConfig, +) -> tuple[SummaryTargetTables | None, SummaryTargetTables | None]: + """Load the precomputed window-summary targets for both splits, if configured.""" + directory = getattr(config, "summary_targets_dir", None) + if not directory: + return None, None + root = Path(directory) + logger.info("[data] loading window-summary targets from %s", root) + train = load_summary_tables(root / "train", root / "stats.json") + tuning = load_summary_tables(root / "tuning", root / "stats.json") + logger.info( + "[data] summary targets: %d train visits, %d tuning visits, %d targets", + len(train), + len(tuning), + train.num_targets, + ) + return train, tuning def _train_streaming(config: TrainingConfig, output_dir: Path, device: str) -> Path: @@ -1360,6 +1460,7 @@ def make_train_patients(epoch: int) -> Iterator[PatientSequence]: shuffle_seed=config.seed + epoch, ) + train_summary_tables, tuning_summary_tables = _summary_tables(config) corpus = PreparedCorpus( vocab=vocab, concepts=concepts, @@ -1373,6 +1474,8 @@ def make_train_patients(epoch: int) -> Iterator[PatientSequence]: tuning_event_tables=tuning_event_tables, tuning_events_binned=tuning_events_binned, make_train_patients=make_train_patients, + train_summary_tables=train_summary_tables, + tuning_summary_tables=tuning_summary_tables, ) return _run_training(config, output_dir, device, corpus) @@ -1420,6 +1523,8 @@ def _run_training( # noqa: PLR0912, PLR0915 tuning_labels, tuning_masks = corpus.tuning_labels, corpus.tuning_masks train_event_tables = corpus.train_event_tables tuning_event_tables = corpus.tuning_event_tables + train_summary_tables = corpus.train_summary_tables + tuning_summary_tables = corpus.tuning_summary_tables tuning_events_binned = corpus.tuning_events_binned model = build_model(config, vocab_size=len(vocab), num_concepts=len(concepts)).to( @@ -1616,9 +1721,18 @@ def make_tuning_sampler() -> StreamingSampler: if train_event_tables is not None else None ) + summary_targets = ( + summary_targets_for_chunk(chunk, train_summary_tables) + if train_summary_tables is not None + else None + ) if isinstance(model, BaselineSequenceModel): total, components, state = model.compute_streaming_loss( - chunk, state=state, objective=objective, event_targets=event_targets + chunk, + state=state, + objective=objective, + event_targets=event_targets, + summary_targets=summary_targets, ) elif ( steering is not None @@ -1639,6 +1753,7 @@ def make_tuning_sampler() -> StreamingSampler: lifted_ids=steering.lifted[injection.concept_index], objective=objective, event_targets=event_targets, + summary_targets=summary_targets, respond_weight=config.respond_weight, express_weight=config.express_weight, forecast_at_injected=config.steering_forecast_at_injected, @@ -1664,6 +1779,7 @@ def make_tuning_sampler() -> StreamingSampler: intervention=intervention, objective=objective, event_targets=event_targets, + summary_targets=summary_targets, teacher_alpha_known=annealed_alpha( global_step, config.teacher_anneal_steps, @@ -1709,6 +1825,7 @@ def make_tuning_sampler() -> StreamingSampler: supervision=config.concept_supervision, # type: ignore[arg-type] objective=objective, event_tables=tuning_event_tables, + summary_tables=tuning_summary_tables, ) fields = { "step": global_step, @@ -1727,6 +1844,7 @@ def make_tuning_sampler() -> StreamingSampler: objective.time_weight, objective.event_hazard_weight, objective.value_head_weight, + summary_weight=objective.summary_weight, ) if val_loss < best_val_loss: best_val_loss = val_loss diff --git a/scripts/build_summary_targets.py b/scripts/build_summary_targets.py new file mode 100644 index 00000000..0103e23a --- /dev/null +++ b/scripts/build_summary_targets.py @@ -0,0 +1,100 @@ +"""Precompute the self-supervised window-summary targets for a data root. + +Writes ``/train/.parquet`` and ``/tuning/.parquet`` +(one raw target frame per shard, see +:func:`odyssey.training.summary_targets.compute_summary_targets`) and +``/stats.json`` (standardization fitted on the train frames). Point a +training run at ```` with ``summary_targets_dir``. + +Shards are independent, so they are processed in parallel. The prepare +step (code normalization, history recap) is the training run's own, taken +from a run config or the defaults, so the target rows' time stamps match +the tokens the model streams. + + uv run python scripts/build_summary_targets.py \\ + --data-root ~/data/mimiciv_3.1_v1/data \\ + --out ~/data/mimiciv_3.1_v1/summary_targets \\ + --config ~/runs/full_run_v10/config.json --workers 12 +""" + +from __future__ import annotations + +import argparse +import json +import logging +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +import polars as pl + +from odyssey.training.data import load_meds_shard +from odyssey.training.shard_stream import make_preparer, shard_paths +from odyssey.training.summary_targets import ( + DEFAULT_LANDMARK_HOURS, + compute_summary_targets, + fit_summary_stats, +) + + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("build_summary_targets") + + +def _one(args: tuple[str, str, str, bool, bool, float]) -> str: + path, out, source, normalize, recap, landmark_hours = args + out_path = Path(out) / (Path(path).stem + ".parquet") + if out_path.exists(): + return f"skip {out_path.name}" + prepare = make_preparer( + source=source, normalize_medications=normalize, history_recap=recap + ) + events = prepare(load_meds_shard(Path(path))) + frame = compute_summary_targets( + events, source=source, landmark_hours=landmark_hours + ) + frame.write_parquet(out_path) + return f"{out_path.name}: {frame.height} rows" + + +def main() -> None: + """Build targets for the train and tuning splits and fit the stats.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-root", required=True) + parser.add_argument("--out", required=True) + parser.add_argument( + "--config", default=None, help="a run's config.json for source/prepare flags" + ) + parser.add_argument("--max-train-shards", type=int, default=None) + parser.add_argument("--max-tuning-shards", type=int, default=None) + parser.add_argument("--landmark-hours", type=float, default=DEFAULT_LANDMARK_HOURS) + parser.add_argument("--workers", type=int, default=4) + args = parser.parse_args() + + cfg = json.loads(Path(args.config).read_text()) if args.config else {} + source = str(cfg.get("source", "mimic_iv")) + normalize = bool(cfg.get("normalize_medications", True)) + recap = bool(cfg.get("history_recap", False)) + out = Path(args.out) + for split, cap in ( + ("train", args.max_train_shards), + ("tuning", args.max_tuning_shards), + ): + (out / split).mkdir(parents=True, exist_ok=True) + paths = shard_paths(Path(args.data_root) / split, cap) + logger.info("%s: %d shards -> %s", split, len(paths), out / split) + jobs = [ + (str(p), str(out / split), source, normalize, recap, args.landmark_hours) + for p in paths + ] + with ProcessPoolExecutor(max_workers=args.workers) as pool: + for msg in pool.map(_one, jobs): + logger.info(" %s", msg) + logger.info("fitting standardization on the train frames") + frames = (pl.read_parquet(p) for p in sorted((out / "train").glob("*.parquet"))) + stats = fit_summary_stats(frames) + stats.save(out / "stats.json") + logger.info("wrote %s (%d targets)", out / "stats.json", len(stats.names)) + + +if __name__ == "__main__": + main() diff --git a/scripts/probe_summary_signal.py b/scripts/probe_summary_signal.py new file mode 100644 index 00000000..76e06cc1 --- /dev/null +++ b/scripts/probe_summary_signal.py @@ -0,0 +1,389 @@ +"""Frozen pre/post-bottleneck probes for the GBM's window-statistic features. + +Companion to scripts/probe_counting_signal.py. That script asks whether the +backbone's own hidden state already encodes the GBM's occurrence COUNTS +(the `counts_occurrence` group). This one asks the same question for the +`summary_stats` group -- per-signal window minimum/maximum/mean over 6 h and +24 h, the change from the previous value, the change from the visit's first +value and the ratio to the visit's minimum -- which the 2026-08-24 +feature-group ablation found leads the GBM's margin on acute kidney injury +(and is literally what KDIGO defines AKI on). Same frozen Ridge probe, same +row selection (4-hourly landmarks), same GBM feature code for the targets. + +Two differences from the counting script, both forced by the targets: + +- A window statistic is NaN when the signal was never measured in that + window, so each target is probed on its own finite rows and the row + count is reported beside its R^2. Multi-output fitting would discard + almost every row. +- Ratios and deltas have heavy tails, so Spearman rank correlation is + reported next to R^2 as a scale-free companion. + +`--feature-group both` also runs the counting group on the same extracted +embeddings, so one invocation gives a run both baselines at once. + +`--probe hgb` fits a HistGradientBoosting regressor beside the Ridge probe +(same rows, same targets), the nonlinear check the counting work used on +2026-08-29: a target the linear probe cannot read but the tree probe can is +"encoded, but not linearly". `--targets change` restricts to the change +statistics (delta_prev, delta_visit_first, ratio_visit_min) plus the level +anchors of a few key signals. Targets are winsorized at the train 0.5/99.5 +percentiles for BOTH probes when `--winsorize` is set, so artifact values +(sentinel MAPs) cannot dominate R^2. + +Not wired into any CI/registry path. Run directly: + + uv run python scripts/probe_summary_signal.py \ + --run-dir ~/runs/full_run_v10 \ + --train-shard-dir ~/data/mimiciv_3.1_v1/data/train \ + --held-out-shard-dir ~/data/mimiciv_3.1_v1/data/held_out \ + --max-train-shards 5 --max-held-out-shards 4 --feature-group both +""" + +from __future__ import annotations + +import argparse +import logging +from collections import defaultdict + +import numpy as np +import polars as pl +import torch +from scipy.stats import spearmanr +from sklearn.ensemble import HistGradientBoostingRegressor +from sklearn.linear_model import Ridge +from sklearn.metrics import r2_score +from sklearn.preprocessing import StandardScaler + +from odyssey.data.alert_events import alert_events_for +from odyssey.data.sidecars import activate_sidecars +from odyssey.data.value_binning import add_value_tokens +from odyssey.inference.alerts import _load_prepared_raw, _visit_starts +from odyssey.inference.baseline_features import ( + CONTEXT_FEATURES, + StrongFeatureBuilder, + feature_names, +) +from odyssey.inference.embedding_probe import collect_embeddings +from odyssey.inference.run_inference import load_run +from odyssey.models.sequence_model import ConceptBottleneckSequenceModel + + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("probe_summary_signal") + +SUMMARY_STATS: tuple[str, ...] = ( + "last", + "mean_24h", + "min_24h", + "max_24h", + "min_6h", + "max_6h", + "delta_prev", + "delta_visit_first", + "ratio_visit_min", +) +# `last` is included as an anchor: the most recent value is the easiest +# statistic to read off a sequence model, so its R^2 bounds the others. + +MIN_TRAIN_ROWS = 500 +MIN_HELD_ROWS = 200 + +CHANGE_STATS: tuple[str, ...] = ("delta_prev", "delta_visit_first", "ratio_visit_min") +ANCHOR_SIGNALS: tuple[str, ...] = ( + "creatinine", + "bun", + "lactate", + "map_noninvasive", + "sbp_noninvasive", + "heart_rate", + "platelets", + "hemoglobin", + "urine_output", +) +ANCHOR_STATS: tuple[str, ...] = ("last", "min_6h", "mean_24h") + + +def _is_change_target(name: str) -> bool: + """Change statistics for every signal, plus level anchors for key signals.""" + signal, stat = name.rsplit(".", 1) + return stat in CHANGE_STATS or (signal in ANCHOR_SIGNALS and stat in ANCHOR_STATS) + + +def _counting_columns(names: list[str]) -> list[int]: + """Occurrence-count columns, exactly as scripts/probe_counting_signal.py.""" + keep_suffixes = (".n_6h", ".n_24h", ".n_visit", ".ever_visit") + idx = [ + i + for i, n in enumerate(names) + if (n.startswith("drug.") or n.startswith("family.")) + and n.endswith(keep_suffixes) + ] + idx += [names.index("n_prior_visits"), names.index("n_events_visit")] + return sorted(idx) + + +def _summary_columns(names: list[str]) -> list[int]: + """Per-signal window-statistic columns (the GBM's `summary_stats` group).""" + context = set(CONTEXT_FEATURES) + return [ + i + for i, n in enumerate(names) + if n not in context + and not n.startswith(("drug.", "family.")) + and n.rsplit(".", 1)[-1] in SUMMARY_STATS + ] + + +def probe_one( + train_x: np.ndarray, + train_y: np.ndarray, + test_x: np.ndarray, + test_y: np.ndarray, + *, + alpha: float, + kind: str = "ridge", +) -> tuple[float, float]: + """Fit one frozen probe (Ridge or HistGradientBoosting); return (R^2, rho).""" + x_scaler = StandardScaler().fit(train_x) + y_mean, y_std = float(train_y.mean()), float(train_y.std() or 1.0) + reg: Ridge | HistGradientBoostingRegressor + if kind == "hgb": + reg = HistGradientBoostingRegressor( + max_iter=300, + learning_rate=0.1, + max_leaf_nodes=31, + early_stopping=True, + random_state=0, + ) + else: + reg = Ridge(alpha=alpha) + reg.fit(x_scaler.transform(train_x), (train_y - y_mean) / y_std) + pred = reg.predict(x_scaler.transform(test_x)) + r2 = float(r2_score((test_y - y_mean) / y_std, pred)) + rho = float(spearmanr(test_y, pred).correlation) + return r2, rho + + +def probe_group( # noqa: PLR0913 + label: str, + idx: list[int], + names: list[str], + *, + train_y_all: np.ndarray, + held_y_all: np.ndarray, + train_pre: np.ndarray, + train_post: np.ndarray, + held_pre: np.ndarray, + held_post: np.ndarray, + alpha: float, + kinds: tuple[str, ...] = ("ridge",), + winsorize: bool = False, + train_subsample: int = 0, + seed: int = 0, +) -> list[dict[str, object]]: + """Probe every column of a group on its own finite rows, per probe kind.""" + rng = np.random.default_rng(seed) + rows: list[dict[str, object]] = [] + for i in idx: + ty, hy = train_y_all[:, i], held_y_all[:, i] + m_tr, m_he = np.isfinite(ty), np.isfinite(hy) + if ( + m_tr.sum() < MIN_TRAIN_ROWS + or m_he.sum() < MIN_HELD_ROWS + or np.nanstd(hy) == 0 + ): + logger.info( + " skip %-40s (train %d, held %d rows)", + names[i], + m_tr.sum(), + m_he.sum(), + ) + continue + tr = np.flatnonzero(m_tr) + if train_subsample and len(tr) > train_subsample: + tr = rng.choice(tr, size=train_subsample, replace=False) + ty_use, hy_use = ty[tr], hy[m_he] + if winsorize: + lo, hi = np.percentile(ty_use, [0.5, 99.5]) + ty_use, hy_use = np.clip(ty_use, lo, hi), np.clip(hy_use, lo, hi) + row: dict[str, object] = {"feature": names[i], "n_held": int(m_he.sum())} + for kind in kinds: + pre = probe_one( + train_pre[tr], ty_use, held_pre[m_he], hy_use, alpha=alpha, kind=kind + ) + post = probe_one( + train_post[tr], ty_use, held_post[m_he], hy_use, alpha=alpha, kind=kind + ) + row[f"{kind}_pre_r2"], row[f"{kind}_pre_rho"] = pre + row[f"{kind}_post_r2"], row[f"{kind}_post_rho"] = post + rows.append(row) + logger.info( + " %-40s " + + " ".join( + f"{k}: pre={row[f'{k}_pre_r2']:.3f} post={row[f'{k}_post_r2']:.3f} rho={row[f'{k}_pre_rho']:.2f}" + for k in kinds + ), + names[i], + ) + logger.info("[%s] %d targets probed", label, len(rows)) + return rows + + +def report(label: str, rows: list[dict[str, object]], kinds: tuple[str, ...]) -> None: + """Log per-statistic medians for every probe kind, then print the CSV block.""" + if not rows: + logger.info("[%s] nothing to report", label) + return + cols = [ + f"{k}_{side}_{m}" + for k in kinds + for side in ("pre", "post") + for m in ("r2", "rho") + ] + by_stat: dict[str, list[dict[str, object]]] = defaultdict(list) + for r in rows: + by_stat[str(r["feature"]).rsplit(".", 1)[-1]].append(r) + logger.info("[%s] MEDIAN by statistic; columns: %s", label, " ".join(cols)) + for stat, rs in sorted(by_stat.items()): + meds = [float(np.median([float(r[c]) for r in rs])) for c in cols] # type: ignore[arg-type] + logger.info( + " %-20s %s (n=%d)", stat, " ".join(f"{m:6.3f}" for m in meds), len(rs) + ) + print(f"\n# group={label}") + print("feature,n_held," + ",".join(cols)) + for r in rows: + print(f"{r['feature']},{r['n_held']}," + ",".join(str(r[c]) for c in cols)) + + +def main() -> None: # noqa: PLR0915 + """Extract pre/post-bottleneck embeddings and report window-stat recovery.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", required=True) + parser.add_argument("--train-shard-dir", required=True) + parser.add_argument("--held-out-shard-dir", required=True) + parser.add_argument("--max-train-shards", type=int, default=5) + parser.add_argument("--max-held-out-shards", type=int, default=4) + parser.add_argument("--landmark-hours", type=float, default=4.0) + parser.add_argument("--num-lanes", type=int, default=64) + parser.add_argument("--chunk-size", type=int, default=512) + parser.add_argument("--ridge-alpha", type=float, default=1.0) + parser.add_argument( + "--feature-group", choices=("summary", "counts", "both"), default="summary" + ) + parser.add_argument("--probe", choices=("ridge", "hgb", "both"), default="ridge") + parser.add_argument( + "--targets", + choices=("all", "change"), + default="all", + help="'change' keeps delta_prev/delta_visit_first/ratio_visit_min plus level anchors", + ) + parser.add_argument( + "--winsorize", + action="store_true", + help="clip targets at train 0.5/99.5 percentiles", + ) + parser.add_argument( + "--train-subsample", + type=int, + default=0, + help="rows per target for fitting (0 = all)", + ) + args = parser.parse_args() + kinds: tuple[str, ...] = ("ridge", "hgb") if args.probe == "both" else (args.probe,) + + device = "cuda" if torch.cuda.is_available() else "cpu" + model, vocab, binner, config = load_run(args.run_dir, device=device) + if not isinstance(model, ConceptBottleneckSequenceModel): + raise ValueError(f"{args.run_dir} is not a concept-bottleneck run") + source = getattr(config, "source", "mimic_iv") + task_set = getattr(config, "task_set", "v1") + landmark_alerts = [a for a in alert_events_for(task_set) if not a.next_visit] + + def load_split( + shard_dir: str, max_shards: int + ) -> tuple[pl.DataFrame, dict[tuple[int, int], float], StrongFeatureBuilder]: + activate_sidecars(shard_dir) + raw = _load_prepared_raw(shard_dir, max_shards, config, source) + visit_start = _visit_starts(raw) + binned = add_value_tokens(raw, binner, source=source) + builder = StrongFeatureBuilder(raw, source=source) + del raw + return binned, visit_start, builder + + logger.info( + "loading %d train shard(s) from %s", args.max_train_shards, args.train_shard_dir + ) + train_binned, train_visit_start, train_builder = load_split( + args.train_shard_dir, args.max_train_shards + ) + logger.info( + "loading %d held-out shard(s) from %s", + args.max_held_out_shards, + args.held_out_shard_dir, + ) + held_binned, held_visit_start, held_builder = load_split( + args.held_out_shard_dir, args.max_held_out_shards + ) + + def embed( + binned: pl.DataFrame, visit_start: dict[tuple[int, int], float] + ) -> tuple[list[tuple[int, int, float]], np.ndarray, np.ndarray]: + keys, pre, post, _, _, _ = collect_embeddings( + model, + binned, + vocab, + landmark_alerts=landmark_alerts, + visit_end_alerts=[], + visit_start=visit_start, + landmark_hours=args.landmark_hours, + num_lanes=args.num_lanes, + chunk_size=args.chunk_size, + device=device, + ) + return keys, pre, post + + logger.info("extracting train embeddings") + train_keys, train_pre, train_post = embed(train_binned, train_visit_start) + logger.info("train: %d landmark rows", len(train_keys)) + logger.info("extracting held-out embeddings") + held_keys, held_pre, held_post = embed(held_binned, held_visit_start) + logger.info("held-out: %d landmark rows", len(held_keys)) + + names = feature_names() + train_y = train_builder.features( + [k[0] for k in train_keys], + [k[1] for k in train_keys], + [k[2] for k in train_keys], + ) + held_y = held_builder.features( + [k[0] for k in held_keys], [k[1] for k in held_keys], [k[2] for k in held_keys] + ) + + groups = {"summary": _summary_columns, "counts": _counting_columns} + chosen = list(groups) if args.feature_group == "both" else [args.feature_group] + for label in chosen: + idx = groups[label](names) + if args.targets == "change" and label == "summary": + idx = [i for i in idx if _is_change_target(names[i])] + logger.info("[%s] %d candidate targets", label, len(idx)) + rows = probe_group( + label, + idx, + names, + train_y_all=train_y, + held_y_all=held_y, + train_pre=train_pre, + train_post=train_post, + held_pre=held_pre, + held_post=held_post, + alpha=args.ridge_alpha, + kinds=kinds, + winsorize=args.winsorize, + train_subsample=args.train_subsample, + ) + report(label, rows, kinds) + + +if __name__ == "__main__": + main() diff --git a/tests/odyssey/models/test_forecast_objective.py b/tests/odyssey/models/test_forecast_objective.py index ce5c60a1..b921f99f 100644 --- a/tests/odyssey/models/test_forecast_objective.py +++ b/tests/odyssey/models/test_forecast_objective.py @@ -219,7 +219,13 @@ def test_baseline_model_shares_objective_and_time_head() -> None: total, comp, _ = model.compute_streaming_loss( chunk, objective=ForecastObjective(bundle_invariant=True, time_weight=0.5) ) - assert set(comp) == {"task_loss", "time_loss", "event_loss", "value_loss"} + assert set(comp) == { + "task_loss", + "time_loss", + "event_loss", + "value_loss", + "summary_loss", + } total.backward() diff --git a/tests/odyssey/models/test_streaming_training.py b/tests/odyssey/models/test_streaming_training.py index 07759b70..03758588 100644 --- a/tests/odyssey/models/test_streaming_training.py +++ b/tests/odyssey/models/test_streaming_training.py @@ -102,6 +102,7 @@ def test_streaming_loss_is_finite_when_a_patient_ends_in_chunk() -> None: "time_loss", "event_loss", "value_loss", + "summary_loss", "concept_loss", "orthogonality_loss", "observability_loss", diff --git a/tests/odyssey/models/test_summary_head.py b/tests/odyssey/models/test_summary_head.py new file mode 100644 index 00000000..40a6ab56 --- /dev/null +++ b/tests/odyssey/models/test_summary_head.py @@ -0,0 +1,161 @@ +"""Tests for the summary head, its masked loss, and its wiring into the models.""" + +import torch + +from odyssey.data.sequences import PatientSequence +from odyssey.data.streaming import PackedLaneSampler +from odyssey.models.backbones.tiny_gru import TinyGRUBackbone +from odyssey.models.sequence_model import ( + BaselineSequenceModel, + ConceptBottleneckSequenceModel, + ForecastObjective, +) +from odyssey.models.summary_head import SummaryHead, masked_huber_loss +from odyssey.training.summary_targets import ( + SummaryTargets, + summary_target_names, + summary_target_weights, +) + + +VOCAB = 40 +HIDDEN = 8 +K = 5 + + +def _seq(sid: int, n: int) -> PatientSequence: + return PatientSequence( + subject_id=sid, + concept_ids=[1 + (i % (VOCAB - 1)) for i in range(n)], + type_ids=[1] * n, + time_stamps=[float(i) for i in range(n)], + ages=[40.0] * n, + visit_orders=[0] * n, + visit_segments=[0] * n, + ) + + +def _backbone() -> TinyGRUBackbone: + return TinyGRUBackbone( + vocab_size=VOCAB, hidden_size=HIDDEN, num_layers=1, padding_idx=0 + ) + + +def test_masked_huber_loss_averages_over_the_mask_only() -> None: + pred = torch.zeros(1, 3, 2) + target = torch.tensor([[[1.0, 100.0], [0.5, 0.0], [0.0, 0.0]]]) + mask = torch.tensor([[[True, False], [True, False], [False, False]]]) + loss = masked_huber_loss(pred, target, mask) + # huber(1.0) = 0.5, huber(0.5) = 0.125 -> mean 0.3125; the 100 is masked + assert loss.item() == 0.3125 + + +def test_masked_huber_loss_is_zero_with_a_graph_when_mask_is_empty() -> None: + pred = torch.zeros(1, 2, 2, requires_grad=True) + loss = masked_huber_loss( + pred, torch.ones(1, 2, 2), torch.zeros(1, 2, 2, dtype=torch.bool) + ) + assert loss.item() == 0.0 + loss.backward() + assert pred.grad is not None + + +def test_summary_head_shapes_linear_and_mlp() -> None: + x = torch.randn(2, 3, HIDDEN) + assert SummaryHead(HIDDEN, K)(x).shape == (2, 3, K) + mlp = SummaryHead(HIDDEN, K, hidden_size=6) + assert mlp(x).shape == (2, 3, K) + assert mlp.hidden_size == 6 + + +def test_models_without_the_head_report_a_zero_summary_loss() -> None: + model = ConceptBottleneckSequenceModel( + backbone=_backbone(), + vocab_size=VOCAB, + num_concepts=2, + embedding_dim=4, + padding_idx=0, + ) + assert model.summary_head is None + chunk = PackedLaneSampler( + iter([_seq(1, 6)]), num_lanes=1, chunk_size=6 + ).next_chunk() + total, components, _ = model.compute_streaming_loss(chunk, {1: torch.zeros(2)}) + assert components["summary_loss"].item() == 0.0 + assert torch.isfinite(total) + + +def test_summary_loss_trains_the_state_and_is_weighted() -> None: + torch.manual_seed(0) + model = ConceptBottleneckSequenceModel( + backbone=_backbone(), + vocab_size=VOCAB, + num_concepts=2, + embedding_dim=4, + padding_idx=0, + summary_targets=K, + ) + assert model.summary_head is not None + model.eval() # dropout would otherwise make the two passes differ + chunk = PackedLaneSampler( + iter([_seq(1, 6)]), num_lanes=1, chunk_size=6 + ).next_chunk() + values = torch.randn(1, 6, K) + mask = torch.zeros(1, 6, K, dtype=torch.bool) + mask[0, 3] = True + targets = SummaryTargets(values=values, mask=mask) + zero_w = ForecastObjective(summary_weight=0.0) + one_w = ForecastObjective(summary_weight=1.0) + total0, comp0, _ = model.compute_streaming_loss( + chunk, {1: torch.zeros(2)}, objective=zero_w, summary_targets=targets + ) + total1, comp1, _ = model.compute_streaming_loss( + chunk, {1: torch.zeros(2)}, objective=one_w, summary_targets=targets + ) + assert comp0["summary_loss"].item() > 0.0 + assert total1.item() > total0.item() + assert abs((total1 - total0).item() - comp1["summary_loss"].item()) < 1e-5 + total1.backward() + assert model.backbone.embeddings.embeddings.word_embeddings.weight.grad is not None + + +def test_baseline_model_accepts_summary_targets_too() -> None: + model = BaselineSequenceModel( + backbone=_backbone(), vocab_size=VOCAB, padding_idx=0, summary_targets=K + ) + chunk = PackedLaneSampler( + iter([_seq(1, 4)]), num_lanes=1, chunk_size=4 + ).next_chunk() + targets = SummaryTargets( + values=torch.zeros(1, 4, K), mask=torch.ones(1, 4, K, dtype=torch.bool) + ) + _, components, _ = model.compute_streaming_loss( + chunk, objective=ForecastObjective(summary_weight=1.0), summary_targets=targets + ) + assert "summary_loss" in components + + +def test_masked_huber_loss_target_weights_reweight_inside_the_average() -> None: + pred = torch.zeros(1, 1, 2) + target = torch.tensor([[[1.0, 0.5]]]) # huber 0.5 and 0.125 + mask = torch.ones(1, 1, 2, dtype=torch.bool) + plain = masked_huber_loss(pred, target, mask) + assert plain.item() == (0.5 + 0.125) / 2 + weighted = masked_huber_loss( + pred, target, mask, target_weights=torch.tensor([3.0, 1.0]) + ) + assert abs(weighted.item() - (3 * 0.5 + 0.125) / 4) < 1e-6 + # a weight of one everywhere is the plain average + same = masked_huber_loss(pred, target, mask, target_weights=torch.ones(2)) + assert same.item() == plain.item() + + +def test_summary_target_weights_mark_changes_and_counts() -> None: + names = summary_target_names() + w = summary_target_weights(4.0) + assert w.shape == (len(names),) + assert w[names.index("creatinine.delta_visit_first")].item() == 4.0 + assert w[names.index("drug.vasopressor.n_6h")].item() == 4.0 + assert w[names.index("family.lab.n_24h")].item() == 4.0 + assert w[names.index("creatinine.min_6h")].item() == 1.0 + assert summary_target_weights(1.0).eq(1.0).all() diff --git a/tests/odyssey/training/test_summary_targets.py b/tests/odyssey/training/test_summary_targets.py new file mode 100644 index 00000000..0824d07b --- /dev/null +++ b/tests/odyssey/training/test_summary_targets.py @@ -0,0 +1,234 @@ +"""Tests for the self-supervised window-summary targets and their chunk lookup.""" + +from datetime import datetime, timedelta + +import numpy as np +import polars as pl +import pytest +import torch + +from odyssey.data.streaming import StreamingChunk +from odyssey.data.types import AuxiliaryInputs, ClinicalSequenceBatch +from odyssey.training.summary_targets import ( + COUNT_STATS, + SIGNAL_STATS, + SummaryTargetStats, + SummaryTargetTables, + compute_summary_targets, + count_target_mask, + fit_summary_stats, + landmark_rows, + load_summary_tables, + summary_target_names, + summary_targets_for_chunk, +) + + +T0 = datetime(2024, 1, 1) + + +def _frame(rows): + return pl.DataFrame( + rows, + schema={ + "subject_id": pl.Int64, + "code": pl.Utf8, + "time": pl.Datetime, + "numeric_value": pl.Float32, + "hadm_id": pl.Int64, + }, + orient="row", + ) + + +def _creat(h: float, v: float, sid: int = 1, hadm: int = 10): + return (sid, "LAB//RESULT//50912//mg/dL::HIGH", T0 + timedelta(hours=h), v, hadm) + + +def _events() -> pl.DataFrame: + return _frame( + [ + (1, "MEDS_BIRTH", T0 - timedelta(days=365.25 * 60), None, None), + (1, "GENDER//F", None, None, None), + _creat(0.0, 1.0), + _creat(5.0, 1.2), + _creat(9.0, 2.2), + (1, "MEDICATION//norepinephrine", T0 + timedelta(hours=8.5), None, 10), + (1, "MEDICATION//norepinephrine", T0 + timedelta(hours=9.0), None, 10), + ] + ) + + +def test_target_names_cover_signals_and_counts() -> None: + names = summary_target_names() + assert len(names) == len(set(names)) + assert "creatinine.delta_visit_first" in names + assert "drug.vasopressor.n_6h" in names + assert "family.lab.n_24h" in names + assert all(n.rsplit(".", 1)[-1] in SIGNAL_STATS + COUNT_STATS for n in names) + counts = count_target_mask(names) + assert counts.sum() == sum(n.rsplit(".", 1)[-1] in COUNT_STATS for n in names) + + +def test_landmark_rows_are_every_4h_at_the_last_event_time() -> None: + sids, vids, times = landmark_rows(_events(), landmark_hours=4.0) + # visit spans 0..9 h: landmarks at 0, 4, 8 -> last event at or before: 0, 0, 5 + assert sids == [1, 1] # the 0 h and 4 h landmarks both map to the 0 h event + assert vids == [10, 10] + assert times == [0.0, 5.0] + + +def test_compute_summary_targets_reports_baseline_change_and_counts() -> None: + frame = compute_summary_targets(_events(), landmark_hours=4.0) + assert frame.columns[:3] == ["subject_id", "visit_id", "time_hours"] + at5 = frame.filter(pl.col("time_hours") == 5.0) + assert at5.height == 1 + assert at5["creatinine.delta_visit_first"][0] == pytest.approx(0.2) + assert at5["creatinine.max_24h"][0] == pytest.approx(1.2) + assert at5["creatinine.min_24h"][0] == pytest.approx(1.0) + assert at5["drug.vasopressor.n_6h"][0] == 0.0 + assert at5["family.lab.n_24h"][0] == 2.0 + at0 = frame.filter(pl.col("time_hours") == 0.0) + assert at0["creatinine.delta_visit_first"][0] == pytest.approx(0.0) + # a signal never measured is NaN, never a fake zero + assert np.isnan(at0["lactate.min_6h"][0]) + + +def test_compute_summary_targets_on_empty_frame_has_the_full_schema() -> None: + frame = compute_summary_targets(_events().head(0)) + assert frame.height == 0 + assert set(summary_target_names()) <= set(frame.columns) + + +def test_stats_standardize_counts_through_log1p_and_winsorize() -> None: + names = summary_target_names() + j_count = names.index("drug.vasopressor.n_6h") + j_delta = names.index("creatinine.delta_visit_first") + n = 1000 + raw = np.full((n, len(names)), np.nan) + rng = np.random.default_rng(0) + raw[:, j_count] = rng.poisson(2.0, n) + raw[:, j_delta] = rng.normal(0.0, 1.0, n) + raw[0, j_delta] = 1e6 # a sentinel that must not set the scale + frame = pl.DataFrame( + { + "subject_id": [1] * n, + "visit_id": [1] * n, + "time_hours": np.arange(n, dtype=float), + } + ) + frame = frame.with_columns( + [pl.Series(name, raw[:, i], dtype=pl.Float32) for i, name in enumerate(names)] + ) + stats = fit_summary_stats([frame]) + assert stats.std[j_delta] < 2.0 + z = stats.transform(raw) + assert abs(float(np.nanmean(z[:, j_delta]))) < 0.2 + assert float(z[0, j_delta]) < 5.0 # clipped, not 1e6 standard deviations + assert np.isnan(z[:, names.index("lactate.min_6h")]).all() + # counts: log1p(0) -> the smallest standardized value + zero = stats.transform( + np.where(np.arange(len(names)) == j_count, 0.0, np.nan)[None, :] + ) + assert float(zero[0, j_count]) < 0.0 + + +def test_stats_round_trip_through_json(tmp_path) -> None: + stats = fit_summary_stats([compute_summary_targets(_events())]) + stats.save(tmp_path / "stats.json") + back = SummaryTargetStats.load(tmp_path / "stats.json") + assert back.names == stats.names + np.testing.assert_allclose(back.mean, stats.mean) + np.testing.assert_allclose(back.std, stats.std) + + +def _chunk(subject_ids, visit_ids, times, real=None) -> StreamingChunk: + lanes, length = np.asarray(times).shape + real_mask = ( + torch.ones(lanes, length, dtype=torch.bool) + if real is None + else torch.tensor(real) + ) + return StreamingChunk( + batch=ClinicalSequenceBatch( + concept_ids=torch.ones(lanes, length, dtype=torch.long), + aux=AuxiliaryInputs( + type_ids=torch.ones(lanes, length, dtype=torch.long), + time_stamps=torch.tensor(times, dtype=torch.float32), + ages=torch.full((lanes, length), 40.0), + visit_orders=torch.zeros(lanes, length, dtype=torch.long), + visit_segments=torch.zeros(lanes, length, dtype=torch.long), + ), + ), + targets=torch.ones(lanes, length, dtype=torch.long), + reset_mask=torch.zeros(lanes, length, dtype=torch.bool), + real_mask=real_mask, + subject_ids=torch.tensor(subject_ids), + patient_end=torch.zeros(lanes, length, dtype=torch.bool), + visit_ids=torch.tensor(visit_ids), + visit_end=torch.zeros(lanes, length, dtype=torch.bool), + ) + + +def _tables() -> SummaryTargetTables: + frame = compute_summary_targets(_events()) + stats = fit_summary_stats([frame]) + tables = SummaryTargetTables(stats) + tables.add_frame(frame) + return tables + + +def test_chunk_targets_land_on_bundle_ends_of_landmark_rows() -> None: + tables = _tables() + assert len(tables) == 1 + # lane 0: patient 1, visit 10; the 5.0 h bundle has two tokens (a + # panel), the target must sit on the LAST of them. lane 1: unknown patient. + chunk = _chunk( + subject_ids=[[1, 1, 1, 1], [7, 7, 7, 7]], + visit_ids=[[10, 10, 10, 10], [1, 1, 1, 1]], + times=[[0.0, 5.0, 5.0, 9.0], [0.0, 4.0, 8.0, 12.0]], + ) + out = summary_targets_for_chunk(chunk, tables) + assert out is not None + hit = out.mask.any(dim=-1) + assert hit.tolist() == [[True, False, True, False], [False, False, False, False]] + k = tables.num_targets + assert out.values.shape == (2, 4, k) + j = tables.stats.names.index("creatinine.delta_visit_first") + # the 5 h row's standardized delta is above the 0 h row's (0.2 vs 0.0) + assert out.values[0, 2, j] > out.values[0, 0, j] + # NaN targets are masked out and zero-filled, never passed as NaN + assert torch.isfinite(out.values).all() + assert not out.mask[0, 2, tables.stats.names.index("lactate.min_6h")] + + +def test_chunk_targets_ignore_padding_and_return_none_when_nothing_matches() -> None: + tables = _tables() + chunk = _chunk( + subject_ids=[[1, 1]], + visit_ids=[[10, 10]], + times=[[5.0, 5.0]], + real=[[True, False]], + ) + out = summary_targets_for_chunk(chunk, tables) + assert out is not None + assert out.mask.any(dim=-1).tolist() == [[True, False]] + assert ( + summary_targets_for_chunk( + _chunk(subject_ids=[[1]], visit_ids=[[10]], times=[[3.0]]), tables + ) + is None + ) + + +def test_tables_load_from_a_directory_of_parquets(tmp_path) -> None: + frame = compute_summary_targets(_events()) + (tmp_path / "train").mkdir() + frame.write_parquet(tmp_path / "train" / "shard_0.parquet") + fit_summary_stats([frame]).save(tmp_path / "stats.json") + tables = load_summary_tables(tmp_path / "train") + assert len(tables) == 1 + times, values = tables.lookup(1, 10) + assert times.tolist() == [0.0, 5.0] + assert values.shape == (2, tables.num_targets) + assert tables.lookup(1, 11) is None