From b56bee47e5164b2397a73f914126a886ff94fe3d Mon Sep 17 00:00:00 2001 From: Logiquo Date: Fri, 30 Oct 2026 11:18:49 +0000 Subject: [PATCH 1/4] Add Unified_RNN for SequenceProcessorV2 fields Unified_RNN embeds each sequence_v2 field with SequenceEmbeddingModel and encodes it bottom-up, level by level: a set level mean-pools each group, a sequence or series level runs an RNN over each group and keeps the last output. A group without records is skipped and summarized as zeros. Static tensor fields are used directly or through a projection; the summaries and statics are concatenated and a linear layer predicts the label. - Per-field settings (hidden_dim, rnn_type, num_layers, dropout, bidirectional, time_feature) take one value or a {field: value} dict. - time_feature adds the time since the previous record of the group on timed levels. - Per-step labels: one prediction per label record, from each field's top-level state after its last record at or before the label's time, so a prediction never sees later records. Fields need a series top level and the RNNs must be causal. - RNN and MultimodalRNN now point to Unified_RNN for per-step labels. --- docs/api/models.rst | 3 + docs/api/models/pyhealth.models.RNN.rst | 7 +- ...pyhealth.models.SequenceEmbeddingModel.rst | 1 + pyhealth/models/__init__.py | 1 + pyhealth/models/base_model.py | 2 +- pyhealth/models/rnn.py | 481 +++++++++++++++++- pyhealth/processors/timeseries_processor.py | 2 +- tests/core/test_per_step_labels.py | 46 +- tests/core/test_timeseries_modes.py | 74 ++- tests/core/test_unified_rnn.py | 339 ++++++++++++ tests/core/test_yaib_mimic4.py | 32 +- 11 files changed, 939 insertions(+), 49 deletions(-) create mode 100644 tests/core/test_unified_rnn.py diff --git a/docs/api/models.rst b/docs/api/models.rst index 769d9ccce..f89b9c8b9 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -39,6 +39,9 @@ routes each feature type automatically. * - ``MultimodalRNN`` - Features mix sequential codes with static tensors or multi-hot encodings - Auto-routes sequential features to RNN layers and non-sequential features to linear layers; good default for EHR + * - ``Unified_RNN`` + - Features are ``sequence_v2`` fields (nested codes, values and timestamps), plus static tensors + - Encodes each field bottom-up over packed records, without padding the hierarchy * - :doc:`models/pyhealth.models.StageNet` - You have time-stamped vital signs with irregular measurement intervals - Requires ``StageNetProcessor`` or ``StageNetTensorProcessor`` in the task schema diff --git a/docs/api/models/pyhealth.models.RNN.rst b/docs/api/models/pyhealth.models.RNN.rst index 3f264c176..c6cf31cba 100644 --- a/docs/api/models/pyhealth.models.RNN.rst +++ b/docs/api/models/pyhealth.models.RNN.rst @@ -17,4 +17,9 @@ The separate callable RNNLayer and the complete RNN model. .. autoclass:: pyhealth.models.MultimodalRNN :members: :undoc-members: - :show-inheritance: \ No newline at end of file + :show-inheritance: + +.. autoclass:: pyhealth.models.Unified_RNN + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst b/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst index 74ca799da..fee995340 100644 --- a/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst +++ b/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst @@ -5,6 +5,7 @@ Embeddings for the packed records of ``sequence_v2`` (:class:`~pyhealth.processors.SequenceProcessorV2`) fields. Each level's records get one embedding per token namespace and a projection of their values, concatenated; each level keeps its offsets and times. +:class:`~pyhealth.models.Unified_RNN` builds on it. .. autoclass:: pyhealth.models.SequenceEmbeddingModel :members: diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index 5c430d584..314dd9017 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -28,6 +28,7 @@ from .molerec import MoleRec, MoleRecLayer from .retain import MultimodalRETAIN, RETAIN, RETAINLayer from .rnn import MultimodalRNN, RNN, RNNLayer +from .rnn import Unified_RNN as Unified_RNN from .safedrug import SafeDrug, SafeDrugLayer from .sparcnet import DenseBlock, DenseLayer, SparcNet, TransitionLayer from .stagenet import StageNet, StageNetLayer diff --git a/pyhealth/models/base_model.py b/pyhealth/models/base_model.py index 62dc3768e..5a5c2d681 100644 --- a/pyhealth/models/base_model.py +++ b/pyhealth/models/base_model.py @@ -46,7 +46,7 @@ def _label_record_loss( raise ValueError( f"Expected one prediction per label record, {tuple(values.shape)}, " f"got {tuple(logits.shape)}. Per-step labels need a model that " - "predicts at each label's time." + "predicts at each label's time, such as Unified_RNN." ) if values.numel() == 0: # A batch without label records: a zero loss that keeps the graph, diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py index dc2c0a0b5..2b0741aca 100644 --- a/pyhealth/models/rnn.py +++ b/pyhealth/models/rnn.py @@ -1,4 +1,5 @@ -from typing import Optional +from collections.abc import Callable, Mapping, Sequence +from typing import Any, Literal, Optional, cast, get_args import torch import torch.nn as nn @@ -13,6 +14,7 @@ NestedFloatsProcessor, NestedSequenceProcessor, SequenceProcessor, + SequenceProcessorV2, StageNetProcessor, StageNetTensorProcessor, TensorProcessor, @@ -20,6 +22,12 @@ ) from .embedding import EmbeddingModel +from .sequence_embedding import ( + LevelEmbedding, + MaskMode, + PretrainedEmbedding, + SequenceEmbeddingModel, +) # Label modes with one prediction per label record (TimeseriesLabelProcessor). _PER_STEP_MODES = ("timeseries_binary", "timeseries_regression") @@ -230,7 +238,8 @@ def __init__( self.label_key = self.label_keys[0] if self.mode in _PER_STEP_MODES: raise ValueError( - "RNN does not support per-step (timeseries_*) labels." + "RNN does not support per-step (timeseries_*) labels; use " + "Unified_RNN with sequence_v2 inputs." ) self.embedding_model = EmbeddingModel(dataset, embedding_dim) @@ -429,7 +438,8 @@ def __init__( self.label_key = self.label_keys[0] if self.mode in _PER_STEP_MODES: raise ValueError( - "MultimodalRNN does not support per-step (timeseries_*) labels." + "MultimodalRNN does not support per-step (timeseries_*) labels; " + "use Unified_RNN with sequence_v2 inputs." ) self.embedding_model = EmbeddingModel(dataset, embedding_dim) @@ -579,3 +589,468 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]: if kwargs.get("embed", False): results["embed"] = patient_emb return results + + +# ---------------------------------------------------------------------------- +# Unified_RNN: RNNs over the packed records of SequenceProcessorV2 fields +# ---------------------------------------------------------------------------- + +RNNType = Literal["RNN", "LSTM", "GRU"] + + +def _per_field[T]( + name: str, + value: T | Mapping[str, T], + fields: Sequence[str], + default: T, + valid: Callable[[T], bool], +) -> dict[str, T]: + """One value per field, from a single value or a ``{field: value}`` dict. + + Fields a dict leaves out get ``default``. + """ + if isinstance(value, Mapping): + unknown = set(value) - set(fields) + if unknown: + raise ValueError( + f"{name} names {sorted(unknown)}, which are not among {sorted(fields)}." + ) + # T is never itself a Mapping here, so the dict holds T values. + values = {field: cast(T, value.get(field, default)) for field in fields} + else: + values = {field: cast(T, value) for field in fields} + for field, item in values.items(): + if not valid(item): + raise ValueError(f"{name} for '{field}' is invalid: {item!r}.") + return values + + +def _groups(offsets: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The group of each row and its position in the group, from CSR offsets.""" + counts = offsets.diff() + group = torch.repeat_interleave( + torch.arange(len(counts), device=offsets.device), counts + ) + return group, torch.arange(len(group), device=offsets.device) - offsets[:-1][group] + + +def _pool(x: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor: + """Mean-pools the rows of each group; an empty group gives zeros.""" + counts = offsets.diff() + group, _ = _groups(offsets) + out = x.new_zeros(len(counts), x.shape[1]) + out = out.index_add(0, group, x) + return out / counts.clamp(min=1).unsqueeze(1).to(x.dtype) + + +def _gaps(time: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor: + """Time since the previous row of the same group (0 for a group's first). + + ``time`` is ``[N]`` or ``[N, k]`` (one column per timed level). + """ + _, position = _groups(offsets) + gaps = torch.diff(time, dim=0, prepend=time[:1]) + first = position == 0 + return gaps.masked_fill(first.unsqueeze(1) if time.dim() > 1 else first, 0.0) + + +def _rnn_over_groups( + layer: RNNLayer, x: torch.Tensor, offsets: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Runs ``layer`` over each group of rows. + + Returns each group's last output ``[G, H]`` (zeros for an empty group, + which is skipped) and every row's output ``[N, H]``. + """ + counts = offsets.diff() + hidden = layer.hidden_size + summary = x.new_zeros(len(counts), hidden) + filled = counts > 0 + if not bool(filled.any()): + return summary, x.new_zeros(len(x), hidden) + group, position = _groups(offsets) + rank = torch.cumsum(filled.long(), dim=0) - 1 # index among non-empty groups + lengths = counts[filled] + padded = x.new_zeros(len(lengths), int(lengths.max()), x.shape[1]) + padded[rank[group], position] = x + steps = torch.arange(padded.shape[1], device=x.device) + outputs, last = layer(padded, (steps < lengths.unsqueeze(1)).int()) + summary[filled] = last + return summary, outputs[rank[group], position] + + +class _FieldRNN(nn.Module): + """Encodes one V2 field's level embeddings into one summary per sample.""" + + def __init__( + self, + field: str, + processor: SequenceProcessorV2, + level_dims: Sequence[int], + hidden_dim: int, + rnn_type: RNNType, + num_layers: int, + dropout: float, + bidirectional: bool, + time_feature: bool, + ): + super().__init__() + self.depth = processor.depth + orderings = processor.ordering() + # A level runs an RNN unless it is a set (a resampled series is a sequence). + self.is_rnn = [orderings[level] != "set" for level in range(self.depth + 1)] + timed = {level for level, role in processor.schema() if role == "time"} + self.gap_levels = sorted(timed) if time_feature else [] + + def rnn(input_size: int) -> nn.Module: + return RNNLayer( + input_size, + hidden_dim, + rnn_type=rnn_type, + num_layers=num_layers, + dropout=dropout, + bidirectional=bidirectional, + ) + + self.encoders = nn.ModuleList() + below = 0 # the summary size of the level below + for level in range(self.depth + 1): + width = level_dims[level] + below + (level in self.gap_levels) + if width == 0: + raise ValueError( + f"'{field}' level {level} has nothing to encode: no embedding, " + "and nothing below it." + ) + self.encoders.append(rnn(width) if self.is_rnn[level] else nn.Identity()) + below = hidden_dim if self.is_rnn[level] else width + self.output_dim = below + + # -- inputs -------------------------------------------------------------- + + def _inputs( + self, level: int, out: dict[Any, Any], parts: list[torch.Tensor] + ) -> torch.Tensor: + """A level's input rows: its embedding, the given parts, and its gaps.""" + embedding = out[level]["embedding"] + rows = [embedding] if embedding is not None else [] + rows += parts + if level in self.gap_levels: + rows.append(_gaps(out[level]["time"], out[level]["offsets"]).unsqueeze(1)) + return torch.cat(rows, dim=1) + + # -- bottom-up ------------------------------------------------------------- + + def _bottom_up( + self, out: dict[Any, Any] + ) -> tuple[torch.Tensor, torch.Tensor | None]: + summary: torch.Tensor | None = None + states = None + for level in range(self.depth + 1): + x = self._inputs(level, out, [summary] if summary is not None else []) + offsets = out[level]["offsets"] + if self.is_rnn[level]: + encoder = cast(RNNLayer, self.encoders[level]) + summary, states = _rnn_over_groups(encoder, x, offsets) + else: + summary, states = _pool(x, offsets), None + assert summary is not None + return summary, states # states: the top level's per-record outputs + + # -- forward ----------------------------------------------------------------- + + def forward( + self, out: dict[Any, Any] + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Encodes a batch of the field. + + Returns the summary per sample ``[B, S]``, the states of the + top-level records (None if the top level is pooled) and their times + (None if the top level has no time column). + """ + summary, states = self._bottom_up(out) + return summary, states, out[self.depth]["time"] + + +class Unified_RNN(BaseModel): + """RNN model over :class:`~pyhealth.processors.SequenceProcessorV2` fields. + + Each V2 field is embedded by :class:`SequenceEmbeddingModel` and encoded + into one summary per sample, independently of the other fields. Static + fields (:class:`~pyhealth.processors.TensorProcessor`) are used directly, + or through a projection. The summaries and statics are concatenated and a + linear layer predicts the label. Other processors are rejected: sequences, + sets and token fields are expressed with ``"sequence_v2"``. + + Each V2 field is encoded bottom-up, level by level, from level 0 up. + A level's input is its records' embeddings concatenated with the summary + of each record's children. A ``set`` level mean-pools each group; a + ``sequence`` / ``series`` level runs its RNN over each group and keeps + the last output, giving one summary per parent record, up to one per sample. + + A group without records (e.g. a stay without labs) is skipped by the RNN + and summarized as zeros. With ``time_feature``, each record of a timed + (``series``) level also gets the time since the previous record of its + group. + + Per-step labels (``"timeseries_binary"`` / ``"timeseries_regression"``, + a :class:`~pyhealth.processors.TimeseriesLabelProcessor`): the model + predicts once per label record. The label's times are in the unit of the + fields' top level, with the same origin, and a label at time ``t`` uses + each field's top-level state after its last record with time ``<= t`` + (zeros before the first) and the statics, so it never sees later records. + Every V2 field must have a ``series`` top level; a resampled one counts, + with record ``i`` at time ``i * step``. Wrap a ``set`` or ``sequence`` + top level in a ``series`` with ``t = 0``. Logits are ``[label records, + 1]``, packed like the label. + + Every per-field argument takes one value for all V2 fields, or a + ``{field: value}`` dict (unlisted fields get the default). + + Args: + dataset: The dataset, with fitted processors. + embedding_dim: Per V2 field, one + :class:`~pyhealth.models.LevelEmbedding` per level. + hidden_dim: The hidden size of every RNN of a field. + rnn_type: ``"RNN"``, ``"LSTM"`` or ``"GRU"`` (default), per field. + num_layers: The layers of each level's RNN, per field. + dropout: Dropout before each RNN, per field. + bidirectional: Bidirectional RNNs, per field. Not with per-step labels. + time_feature: Whether timed levels get the time since the previous + record, per field. + mask: How :class:`SequenceEmbeddingModel` uses missing-value masks: + ``"ignore"`` (default) or ``"concat"``. + pretrain: :class:`~pyhealth.models.PretrainedEmbedding` entries. + static_dim: Per static field, the size of its projection, or None + (default) to use it directly. + + Examples: + >>> from pyhealth.datasets import create_sample_dataset, get_dataloader + >>> samples = [ + ... {"patient_id": f"p{i}", "age": [50.0 + i], + ... "visits": {"time": [0.0, 1.0], "children": [{"tokens": ["A", "B"]}, {"tokens": ["C"] * i}]}, + ... "label": i % 2} + ... for i in range(4) + ... ] + >>> levels = [{"ordering": "set", "tokens": 1}, {"ordering": "series"}] + >>> dataset = create_sample_dataset( + ... samples, + ... {"age": "tensor", "visits": ("sequence_v2", {"levels": levels})}, + ... {"label": "binary"}, + ... ) + >>> model = Unified_RNN( + ... dataset, embedding_dim={"visits": [{"tokens": (8,)}, {}]}, hidden_dim=16 + ... ) + >>> out = model(**next(iter(get_dataloader(dataset, batch_size=4)))) + >>> out["logit"].shape + torch.Size([4, 1]) + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: Mapping[str, Sequence[LevelEmbedding]], + hidden_dim: int | Mapping[str, int] = 128, + rnn_type: RNNType | Mapping[str, RNNType] = "GRU", + num_layers: int | Mapping[str, int] = 1, + dropout: float | Mapping[str, float] = 0.5, + bidirectional: bool | Mapping[str, bool] = False, + time_feature: bool | Mapping[str, bool] = False, + mask: MaskMode = "ignore", + pretrain: Sequence[PretrainedEmbedding] | None = None, + static_dim: int | Mapping[str, int | None] | None = None, + ): + super().__init__(dataset) + if len(self.label_keys) != 1: + raise ValueError("Unified_RNN supports exactly one label.") + self.label_key = self.label_keys[0] + self.per_step = self.mode in _PER_STEP_MODES + + self.sequence_fields: list[str] = [] + self.static_fields: list[str] = [] + for field in self.feature_keys: + processor = dataset.input_processors[field] + if isinstance(processor, SequenceProcessorV2): + self.sequence_fields.append(field) + elif isinstance(processor, TensorProcessor): + self.static_fields.append(field) + else: + raise TypeError( + f"Field '{field}': {type(processor).__name__} is not supported. " + "Unified_RNN takes 'sequence_v2' fields and static 'tensor' " + "fields; express sequences, sets and tokens with 'sequence_v2'." + ) + + fields = self.sequence_fields + hidden = _per_field( + "hidden_dim", + hidden_dim, + fields, + 128, + lambda v: isinstance(v, int) and v > 0, + ) + rnn_types: dict[str, RNNType] = _per_field( + "rnn_type", rnn_type, fields, "GRU", lambda v: v in get_args(RNNType) + ) + layers = _per_field( + "num_layers", num_layers, fields, 1, lambda v: isinstance(v, int) and v > 0 + ) + dropouts = _per_field( + "dropout", + dropout, + fields, + 0.5, + lambda v: isinstance(v, (int, float)) and 0 <= v < 1, + ) + bidirectionals = _per_field( + "bidirectional", bidirectional, fields, False, lambda v: isinstance(v, bool) + ) + time_features = _per_field( + "time_feature", time_feature, fields, False, lambda v: isinstance(v, bool) + ) + static_dims: dict[str, int | None] = _per_field( + "static_dim", + static_dim, + self.static_fields, + None, + lambda v: v is None or (isinstance(v, int) and v > 0), + ) + + for field in fields: + if bidirectionals[field] and self.per_step: + raise ValueError( + "A per-step label needs a causal model; bidirectional=True " + "would use future time steps." + ) + self.processors = { + field: cast(SequenceProcessorV2, dataset.input_processors[field]) + for field in fields + } + if self.per_step: + for field, processor in self.processors.items(): + top = processor.levels[processor.depth] + if top.get("ordering", "sequence") != "series": + raise ValueError( + f"'{field}': per-step labels need a 'series' top level, so each " + "step can use the records before it. Wrap a 'set' or " + "'sequence' top level in a 'series' with t = 0." + ) + + self.embedding = SequenceEmbeddingModel( + dataset, + embedding_dim=embedding_dim, + pretrain=pretrain, + mask=mask, + ) + self.encoders = nn.ModuleDict( + { + field: _FieldRNN( + field, + self.processors[field], + self.embedding.embedding_dim_per_level(field), + hidden[field], + rnn_types[field], + layers[field], + dropouts[field], + bidirectionals[field], + time_features[field], + ) + for field in fields + } + ) + self.static_projections = nn.ModuleDict() + width = 0 + for field in self.static_fields: + size = int(torch.as_tensor(dataset[0][field]).numel()) + projected = static_dims[field] + if projected is None: + width += size + else: + self.static_projections[field] = nn.Linear(size, projected) + width += projected + for field in fields: + width += cast(_FieldRNN, self.encoders[field]).output_dim + self.fc = nn.Linear(width, self.get_output_size()) + + def _at_labels( + self, + field: str, + states: torch.Tensor, + time: torch.Tensor | None, + batch: torch.Tensor, + label_time: torch.Tensor, + label_sample: torch.Tensor, + ) -> torch.Tensor: + """The top-level state of ``field`` at each label record ``[ΣN, H]``. + + That is the state after the sample's last top-level record with time + ``<=`` the label's time, or zeros before the first. + """ + if time is None: # resampled top level: record i is at time i * step + processor = self.processors[field] + resample = processor.levels[processor.depth].get("resample") + assert resample is not None # __init__ allows no other untimed top level + _, position = _groups(batch) + time = position.to(states.dtype) * float(resample["step"]) + counts = batch.diff() + n, width = len(counts), int(counts.max()) if len(counts) else 0 + padded = time.new_full((n, max(width, 1)), float("inf")) + group, position = _groups(batch) + padded[group, position] = time.to(padded.dtype) + seen = torch.searchsorted( + padded[label_sample], + label_time.to(padded.dtype).unsqueeze(1), + right=True, + )[:, 0] + index = (batch[:-1][label_sample] + seen - 1).clamp(min=0) + if len(states) == 0: + return states.new_zeros(len(label_time), states.shape[1]) + picked = states[index] + return torch.where((seen > 0).unsqueeze(1), picked, torch.zeros_like(picked)) + + def forward(self, **batch: Any) -> dict[str, torch.Tensor]: + """Predicts the label of a batch of :func:`~pyhealth.datasets.get_dataloader`. + + Returns: + ``loss``, ``y_prob``, ``y_true`` and ``logit``; also ``embed``, the + concatenated summaries and statics, if the batch has ``embed=True``. + """ + embedded = self.embedding( + {field: batch[field] for field in self.sequence_fields} + ) + label: tuple[torch.Tensor, ...] = () + label_time = label_sample = torch.zeros(0, dtype=torch.long) + if self.per_step: + # Packed label records: (values [ΣN, 1], time [ΣN], offsets [B + 1]). + label = tuple(t.to(self.device) for t in batch[self.label_key]) + label_time = label[1] + label_sample, _ = _groups(label[2]) + parts = [] + for field in self.sequence_fields: + encoder = cast(_FieldRNN, self.encoders[field]) + summary, states, time = encoder(embedded[field]) + if not self.per_step: + parts.append(summary) + continue + assert states is not None # a series top level has states + out = embedded[field] + offsets = out[encoder.depth]["offsets"] + parts.append( + self._at_labels(field, states, time, offsets, label_time, label_sample) + ) + for field in self.static_fields: + x = batch[field].to(self.device).float().flatten(1) + if field in self.static_projections: + x = self.static_projections[field](x) + parts.append(x[label_sample] if self.per_step else x) + embed = torch.cat(parts, dim=-1) + logits = self.fc(embed) + y_true = label if self.per_step else batch[self.label_key].to(self.device) + results = { + "loss": self.get_loss_function()(logits, y_true), + "y_prob": self.prepare_y_prob(logits), + "y_true": y_true, + "logit": logits, + } + if batch.get("embed", False): + results["embed"] = embed + return results diff --git a/pyhealth/processors/timeseries_processor.py b/pyhealth/processors/timeseries_processor.py index c38e49cd9..92e4f5cbd 100644 --- a/pyhealth/processors/timeseries_processor.py +++ b/pyhealth/processors/timeseries_processor.py @@ -202,7 +202,7 @@ class TimeseriesLabelProcessor(SequenceProcessorV2): raises. There is no grid, resampling, imputation or mask. A model predicts one value per label, at its time, from the input records - at or before it. + at or before it (see :class:`~pyhealth.models.Unified_RNN`). Output: - ``(values, time)``: ``[N, 1]`` values and ``[N]`` times. A batch diff --git a/tests/core/test_per_step_labels.py b/tests/core/test_per_step_labels.py index feafd2da3..66e74e969 100644 --- a/tests/core/test_per_step_labels.py +++ b/tests/core/test_per_step_labels.py @@ -1,10 +1,13 @@ import unittest +from collections.abc import Callable +from itertools import product from typing import Any import numpy as np import torch -from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.datasets import SampleDataset, create_sample_dataset, get_dataloader +from pyhealth.models import BaseModel, Unified_RNN from pyhealth.processors import ( LevelConfig, SequenceProcessorV2, @@ -113,24 +116,31 @@ class TestPerStepEndToEnd(unittest.TestCase): """ def test_label_processors_train_and_evaluate(self): - for mode, metric in ( - ("timeseries_binary", "roc_auc"), - ("timeseries_regression", "mae"), + models: dict[str, Callable[[SampleDataset], BaseModel]] = { + "LabelRecordModel": LabelRecordModel, + "Unified_RNN": lambda dataset: Unified_RNN( + dataset, embedding_dim={"dynamic": [{"values": 8}]}, hidden_dim=8 + ), + } + for (mode, metric), (name, make_model) in product( + (("timeseries_binary", "roc_auc"), ("timeseries_regression", "mae")), + models.items(), ): - dataset = _dataset(mode) - model = LabelRecordModel(dataset) - self.assertEqual(model.mode, mode) - batch = next(iter(get_dataloader(dataset, batch_size=4))) - out = model(**batch) - self.assertEqual(tuple(out["y_prob"].shape), (18, 1)) - self.assertTrue(torch.isfinite(out["loss"])) - - trainer = Trainer(model=model, metrics=[metric], enable_logging=False) - loader = get_dataloader(dataset, batch_size=3) - trainer.train(train_dataloader=loader, epochs=1) - y_true, _, _ = trainer.inference(loader) - self.assertEqual(sorted(len(y) for y in y_true), [3, 4, 5, 6]) - self.assertIn(metric, trainer.evaluate(loader)) + with self.subTest(mode=mode, model=name): + dataset = _dataset(mode) + model = make_model(dataset) + self.assertEqual(model.mode, mode) + batch = next(iter(get_dataloader(dataset, batch_size=4))) + out = model(**batch) + self.assertEqual(tuple(out["y_prob"].shape), (18, 1)) + self.assertTrue(torch.isfinite(out["loss"])) + + trainer = Trainer(model=model, metrics=[metric], enable_logging=False) + loader = get_dataloader(dataset, batch_size=3) + trainer.train(train_dataloader=loader, epochs=1) + y_true, _, _ = trainer.inference(loader) + self.assertEqual(sorted(len(y) for y in y_true), [3, 4, 5, 6]) + self.assertIn(metric, trainer.evaluate(loader)) if __name__ == "__main__": diff --git a/tests/core/test_timeseries_modes.py b/tests/core/test_timeseries_modes.py index d1c5af422..e39b1e571 100644 --- a/tests/core/test_timeseries_modes.py +++ b/tests/core/test_timeseries_modes.py @@ -9,7 +9,7 @@ timeseries_binary_metrics_fn, timeseries_regression_metrics_fn, ) -from pyhealth.models import RNN, MultimodalRNN +from pyhealth.models import RNN, LevelEmbedding, MultimodalRNN, Unified_RNN from pyhealth.processors import ( LevelConfig, MultiHotProcessor, @@ -21,6 +21,7 @@ from tests.core.label_record_model_utils import LabelRecordModel DYNAMIC: list[LevelConfig] = [{"ordering": "series", "values": [2]}] +DIMS: dict[str, list[LevelEmbedding]] = {"dynamic": [{"values": 8}]} def _dataset( @@ -61,6 +62,12 @@ def _model(dataset) -> LabelRecordModel: return model +def _unified_rnn(dataset) -> Unified_RNN: + model = Unified_RNN(dataset, embedding_dim=DIMS, hidden_dim=8, dropout=0.0) + model.eval() + return model + + class TestTimeseriesModes(unittest.TestCase): def test_modes_are_separate_from_binary_and_regression(self): for mode in ("timeseries_binary", "timeseries_regression"): @@ -79,18 +86,32 @@ def test_label_processor_class_resolves_to_its_own_mode(self): def test_one_prediction_per_label_record(self): # Sample 1 has no label at hour 0: it is left out, not masked. dataset = _dataset(labeled=lambda i, h: (i, h) != (1, 0)) - model = _model(dataset) batch = next(iter(get_dataloader(dataset, batch_size=4))) + for make_model in (_model, _unified_rnn): + with self.subTest(model=make_model.__name__), torch.no_grad(): + out = make_model(dataset)(**batch) + values, time, offsets = out["y_true"] + self.assertEqual(tuple(out["logit"].shape), (17, 1)) # 5 + 2 + 6 + 4 + self.assertEqual(offsets.tolist(), [0, 5, 7, 13, 17]) + self.assertEqual(time[5:7].tolist(), [1.0, 2.0]) # sample 1 from hour 1 + loss = torch.nn.functional.binary_cross_entropy_with_logits( + out["logit"], values + ) + torch.testing.assert_close(loss, out["loss"]) + + def test_unified_rnn_predictions_do_not_see_later_records(self): + dataset = _dataset() + model = _unified_rnn(dataset) + batch = next(iter(get_dataloader(dataset, batch_size=1))) with torch.no_grad(): - out = model(**batch) - values, time, offsets = out["y_true"] - self.assertEqual(tuple(out["logit"].shape), (17, 1)) # 5 + 2 + 6 + 4 - self.assertEqual(offsets.tolist(), [0, 5, 7, 13, 17]) - self.assertEqual(time[5:7].tolist(), [1.0, 2.0]) # sample 1 from hour 1 - loss = torch.nn.functional.binary_cross_entropy_with_logits( - out["logit"], values - ) - torch.testing.assert_close(loss, out["loss"]) + before = model(**batch)["y_prob"][:, 0] + values, *rest = batch["dynamic"] + values = values.clone() + values[-1] += 100.0 # the last hour + batch["dynamic"] = (values, *rest) + after = model(**batch)["y_prob"][:, 0] + torch.testing.assert_close(before[:-1], after[:-1]) + self.assertFalse(torch.allclose(before[-1], after[-1])) def test_loss_needs_one_prediction_per_label(self): model = _model(_dataset()) @@ -142,19 +163,24 @@ def test_v1_rnns_reject_per_step_labels(self): model(dataset=dataset) def test_trainer_scores_label_records_per_sample(self): - dataset = _dataset() - trainer = Trainer( - model=_model(dataset), metrics=["roc_auc"], enable_logging=False - ) - loader = get_dataloader(dataset, batch_size=3) - y_true, y_prob, _ = trainer.inference(loader) - self.assertEqual(sorted(len(y) for y in y_true), [3, 4, 5, 6]) - self.assertEqual([len(p) for p in y_prob], [len(y) for y in y_true]) - self.assertIn("roc_auc", trainer.evaluate(loader)) - - dataset = _dataset("timeseries_regression") - trainer = Trainer(model=_model(dataset), metrics=["mae"], enable_logging=False) - self.assertIn("mae", trainer.evaluate(get_dataloader(dataset, batch_size=3))) + for make_model in (_model, _unified_rnn): + with self.subTest(model=make_model.__name__): + dataset = _dataset() + trainer = Trainer( + model=make_model(dataset), metrics=["roc_auc"], enable_logging=False + ) + loader = get_dataloader(dataset, batch_size=3) + y_true, y_prob, _ = trainer.inference(loader) + self.assertEqual(sorted(len(y) for y in y_true), [3, 4, 5, 6]) + self.assertEqual([len(p) for p in y_prob], [len(y) for y in y_true]) + self.assertIn("roc_auc", trainer.evaluate(loader)) + + dataset = _dataset("timeseries_regression") + trainer = Trainer( + model=make_model(dataset), metrics=["mae"], enable_logging=False + ) + loader = get_dataloader(dataset, batch_size=3) + self.assertIn("mae", trainer.evaluate(loader)) class TestTimeseriesMetrics(unittest.TestCase): diff --git a/tests/core/test_unified_rnn.py b/tests/core/test_unified_rnn.py new file mode 100644 index 000000000..bb2657623 --- /dev/null +++ b/tests/core/test_unified_rnn.py @@ -0,0 +1,339 @@ +import unittest +from typing import Any, cast + +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import LevelEmbedding, Unified_RNN +from pyhealth.models.rnn import _gaps, _pool +from pyhealth.processors import LevelConfig +from pyhealth.trainer import Trainer + +# visits (series, timed) -> codes (set) +VISITS: list[LevelConfig] = [ + {"ordering": "set", "tokens": 1}, + {"ordering": "series", "tokens": 1}, +] +VISIT_DIMS: dict[str, list[LevelEmbedding]] = { + "visits": [{"tokens": (4,)}, {"tokens": (3,)}] +} + + +def _visits(*visits: tuple[float, str, list[str]]) -> dict: + return { + "time": [v[0] for v in visits], + "tokens": [v[1] for v in visits], + "children": [{"tokens": v[2]} for v in visits], + } + + +VISIT_RECORDS = [ + _visits( + (0.0, "er", ["A", "B"]), (1.5, "op", []), (3.0, "er", ["C"]) + ), # an empty visit + _visits((0.0, "op", ["B"])), + {"time": [], "tokens": [], "children": []}, # no visits at all + _visits((0.5, "er", ["A", "C", "D"]), (2.0, "op", ["D"])), +] + + +def _dataset(output: str = "binary", **fields: Any): + schema = { + name: ("sequence_v2", {"levels": levels}) if levels is not None else "tensor" + for name, (levels, _) in fields.items() + } + n = len(next(iter(fields.values()))[1]) + samples = [ + { + "patient_id": f"p{i}", + **{name: records[i] for name, (_, records) in fields.items()}, + "y": i % 2, + } + for i in range(n) + ] + return create_sample_dataset(samples, schema, {"y": output}) + + +def _model(dataset, **kwargs: Any) -> Unified_RNN: + model = Unified_RNN(dataset, dropout=0.0, **kwargs) + model.eval() + return model + + +def _batch(dataset) -> dict[str, Any]: + return next(iter(get_dataloader(dataset, batch_size=len(dataset)))) + + +def _encode(model: Unified_RNN, batch: dict[str, Any], field: str): + embedded = model.embedding({field: batch[field]})[field] + return embedded, cast(Any, model.encoders[field])(embedded) + + +class TestConfiguration(unittest.TestCase): + def setUp(self): + self.dataset = _dataset(visits=(VISITS, VISIT_RECORDS)) + + def test_rejects_bad_arguments(self): + bad: list[dict[str, Any]] = [ + {"hidden_dim": 0}, + {"rnn_type": "Transformer"}, + {"num_layers": {"visits": 0}}, + {"dropout": 1.5}, + {"hidden_dim": {"other": 8}}, # not a field + ] + for kwargs in bad: + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + Unified_RNN(self.dataset, embedding_dim=VISIT_DIMS, **kwargs) + + def test_needs_exactly_one_label(self): + samples = [ + {"patient_id": f"p{i}", "visits": VISIT_RECORDS[i], "y": i % 2, "z": i % 2} + for i in range(2) + ] + dataset = create_sample_dataset( + samples, + {"visits": ("sequence_v2", {"levels": VISITS})}, + {"y": "binary", "z": "binary"}, + ) + with self.assertRaisesRegex(ValueError, "exactly one label"): + Unified_RNN(dataset, embedding_dim=VISIT_DIMS) + + def test_rejects_a_level_with_nothing_to_encode(self): + times: list[LevelConfig] = [{"ordering": "series"}] # time only + records = [{"time": [0.0, 1.0]}, {"time": [0.5]}] + dataset = _dataset(x=(times, records)) + with self.assertRaisesRegex(ValueError, "nothing to encode"): + Unified_RNN(dataset, embedding_dim={"x": [{}]}) + # The time since the previous record is something to encode. + model = _model(dataset, embedding_dim={"x": [{}]}, time_feature=True) + self.assertEqual(model(**_batch(dataset))["logit"].shape, (2, 1)) + + def test_rejects_v1_sequences_sets_and_tokens(self): + for spec, value in (("sequence", ["A", "B"]), ("multi_hot", ["A"])): + samples = [{"patient_id": f"p{i}", "x": value, "y": i} for i in range(2)] + dataset = create_sample_dataset(samples, {"x": spec}, {"y": "binary"}) + with self.subTest(spec=spec), self.assertRaises(TypeError): + Unified_RNN(dataset, embedding_dim={}) + + +class TestBottomUp(unittest.TestCase): + def test_matches_a_loop_reference(self): + dataset = _dataset(visits=(VISITS, VISIT_RECORDS)) + model = _model(dataset, embedding_dim=VISIT_DIMS, hidden_dim=5) + batch = _batch(dataset) + out, (summary, states, time) = _encode(model, batch, "visits") + encoder = cast(Any, model.encoders["visits"]) + codes, visits = out[0], out[1] + expected = [] + for b in range(len(dataset)): + lo, hi = int(visits["offsets"][b]), int(visits["offsets"][b + 1]) + if lo == hi: + expected.append(torch.zeros(5)) + continue + rows = [] + for v in range(lo, hi): + c0, c1 = int(codes["offsets"][v]), int(codes["offsets"][v + 1]) + pooled = ( + codes["embedding"][c0:c1].mean(0) if c1 > c0 else torch.zeros(4) + ) + rows.append(torch.cat([visits["embedding"][v], pooled])) + _, last = encoder.encoders[1](torch.stack(rows).unsqueeze(0)) + expected.append(last[0]) + torch.testing.assert_close(summary, torch.stack(expected)) + self.assertTrue(torch.all(summary[2] == 0)) # no visits + self.assertEqual(states.shape, (6, 5)) # one state per visit + self.assertEqual(time.tolist(), [0.0, 1.5, 3.0, 0.0, 0.5, 2.0]) + + def test_set_pooling(self): + x = torch.tensor([[1.0, -1.0], [3.0, 5.0], [2.0, 2.0]]) + offsets = torch.tensor([0, 2, 2, 3]) # the middle group is empty + self.assertEqual( + _pool(x, offsets).tolist(), [[2.0, 2.0], [0.0, 0.0], [2.0, 2.0]] + ) + + def test_time_feature(self): + time = torch.tensor([0.0, 1.5, 3.0, 0.0, 0.5, 2.0]) + offsets = torch.tensor([0, 3, 4, 4, 6]) + self.assertEqual(_gaps(time, offsets).tolist(), [0.0, 1.5, 1.5, 0.0, 0.0, 1.5]) + dataset = _dataset(visits=(VISITS, VISIT_RECORDS)) + plain = _model(dataset, embedding_dim=VISIT_DIMS, hidden_dim=5) + timed = _model( + dataset, embedding_dim=VISIT_DIMS, hidden_dim=5, time_feature=True + ) + width = lambda m: cast(Any, m.encoders["visits"]).encoders[1].input_size + self.assertEqual(width(timed), width(plain) + 1) + with torch.no_grad(): + self.assertEqual(timed(**_batch(dataset))["logit"].shape, (4, 1)) + + def test_batch_without_any_records(self): + empty = {"time": [], "tokens": [], "children": []} + dataset = _dataset(visits=(VISITS, [VISIT_RECORDS[0], empty, empty])) + model = _model(dataset, embedding_dim=VISIT_DIMS, hidden_dim=5) + batch = next(iter(get_dataloader(dataset.subset([1, 2]), batch_size=2))) + _, (summary, _, _) = _encode(model, batch, "visits") + self.assertEqual(summary.tolist(), [[0.0] * 5] * 2) + + + def test_lstm(self): + dataset = _dataset(visits=(VISITS, VISIT_RECORDS)) + model = _model( + dataset, + embedding_dim=VISIT_DIMS, + hidden_dim=5, + rnn_type="LSTM", + num_layers=2, + ) + result = model(**_batch(dataset)) + self.assertEqual(result["logit"].shape, (4, 1)) + + +class TestStatics(unittest.TestCase): + def setUp(self): + ages = [[60.0, 1.0], [70.0, 0.0], [50.0, 1.0], [80.0, 0.0]] + self.dataset = _dataset(visits=(VISITS, VISIT_RECORDS), age=(None, ages)) + + def test_used_directly_or_projected(self): + direct = _model(self.dataset, embedding_dim=VISIT_DIMS, hidden_dim=5) + self.assertEqual(direct.fc.in_features, 5 + 2) + projected = _model( + self.dataset, embedding_dim=VISIT_DIMS, hidden_dim=5, static_dim={"age": 3} + ) + self.assertEqual(projected.fc.in_features, 5 + 3) + result = projected(**_batch(self.dataset), embed=True) + self.assertEqual(result["embed"].shape, (4, 8)) + + def test_per_field_settings(self): + dataset = _dataset( + visits=(VISITS, VISIT_RECORDS), stays=(VISITS, VISIT_RECORDS) + ) + dims = {**VISIT_DIMS, "stays": VISIT_DIMS["visits"]} + model = _model( + dataset, + embedding_dim=dims, + hidden_dim={"visits": 4, "stays": 6}, + ) + self.assertEqual(model.fc.in_features, 4 + 6) + self.assertEqual(model(**_batch(dataset))["logit"].shape, (4, 1)) + + +class TestPerStep(unittest.TestCase): + def labels(self, times: list[float]) -> dict[str, Any]: + return {"time": times, "values": [i % 2 for i in range(len(times))]} + + def dataset(self, levels, records, times: tuple[float, ...] = (0.0, 1.0, 2.0)): + samples = [ + { + "patient_id": f"p{i}", + "x": records[i], + "age": [float(i)], + "y": self.labels(list(times)), + } + for i in range(len(records)) + ] + return create_sample_dataset( + samples, + {"x": ("sequence_v2", {"levels": levels}), "age": "tensor"}, + {"y": "timeseries_binary"}, + ) + + def test_top_level_must_be_a_series(self): + for top in ("set", "sequence"): + levels: list[LevelConfig] = [ + {"ordering": "set", "tokens": 1}, + {"ordering": top}, + ] + dataset = self.dataset(levels, [{"children": [{"tokens": ["A"]}]}] * 2) + with self.subTest(top=top), self.assertRaisesRegex(ValueError, "Wrap"): + Unified_RNN(dataset, embedding_dim={"x": [{"tokens": (4,)}, {}]}) + + def test_no_bidirectional(self): + records = [_visits((0.0, "er", ["A"]))] * 2 + dataset = self.dataset(VISITS, records) + with self.assertRaises(ValueError): + Unified_RNN( + dataset, embedding_dim={"x": VISIT_DIMS["visits"]}, bidirectional=True + ) + + def test_each_label_uses_the_records_at_or_before_it(self): + records = [ + # Label 0 sees nothing; label 1 sees visit 0; label 2 sees visit 1. + _visits((0.5, "er", ["A"]), (2.0, "op", ["B"])), + # Label 0 sees nothing; labels 1 and 2 see the visit. + _visits((1.0, "er", ["C"])), + ] + dataset = self.dataset(VISITS, records) + model = _model( + dataset, + embedding_dim={"x": VISIT_DIMS["visits"]}, + hidden_dim=5, + ) + batch = _batch(dataset) + result = model(**batch, embed=True) + self.assertEqual(result["logit"].shape, (6, 1)) # 3 labels each + _, (_, states, _) = _encode(model, batch, "x") + embed = result["embed"][:, :5] + self.assertTrue(torch.all(embed[0] == 0)) + torch.testing.assert_close(embed[1], states[0]) + torch.testing.assert_close(embed[2], states[1]) + self.assertTrue(torch.all(embed[3] == 0)) + torch.testing.assert_close(embed[4], states[2]) + torch.testing.assert_close(embed[5], states[2]) + # The static age goes with every label record of its sample. + self.assertEqual(result["embed"][:, 5].tolist(), [0.0] * 3 + [1.0] * 3) + + def test_labels_between_records(self): + records = [_visits((0.0, "er", ["A"]), (2.0, "op", ["B"]))] * 2 + dataset = self.dataset(VISITS, records, times=(1.5, 2.5)) + model = _model(dataset, embedding_dim={"x": VISIT_DIMS["visits"]}, hidden_dim=5) + batch = _batch(dataset) + embed = model(**batch, embed=True)["embed"][:, :5] + _, (_, states, _) = _encode(model, batch, "x") + torch.testing.assert_close(embed[0], states[0]) # 1.5: the visit at 0 + torch.testing.assert_close(embed[1], states[1]) # 2.5: the visit at 2 + + def test_batch_without_any_records(self): + empty = {"time": [], "tokens": [], "children": []} + records = [_visits((0.0, "er", ["A"])), empty, empty] + dataset = self.dataset(VISITS, records) + model = _model(dataset, embedding_dim={"x": VISIT_DIMS["visits"]}, hidden_dim=5) + batch = next(iter(get_dataloader(dataset.subset([1, 2]), batch_size=2))) + embed = model(**batch, embed=True)["embed"] + self.assertEqual(embed.shape, (6, 6)) # 3 labels each, 5 + age + self.assertTrue(torch.all(embed[:, :5] == 0)) + + def test_resampled_top_level(self): + levels: list[LevelConfig] = [ + { + "ordering": "series", + "values": [1], + "resample": {"step": 1.0, "method": "last"}, + } + ] + records = [{"time": [0.0, 1.0, 2.0], "values": [[1.0], [2.0], [3.0]]}] * 2 + dataset = self.dataset(levels, records) + model = _model(dataset, embedding_dim={"x": [{"values": 2}]}, hidden_dim=5) + batch = _batch(dataset) + result = model(**batch, embed=True) + _, (_, states, _) = _encode(model, batch, "x") + # Record i is at time i, so the label at time s sees record s. + torch.testing.assert_close(result["embed"][:3, :5], states[:3]) + + +class TestTraining(unittest.TestCase): + def test_trains(self): + dataset = _dataset(visits=(VISITS, VISIT_RECORDS)) + model = Unified_RNN( + dataset, embedding_dim=VISIT_DIMS, hidden_dim=8 + ) + loader = get_dataloader(dataset, batch_size=2, shuffle=True, seed=0) + trainer = Trainer(model=model, enable_logging=False) + trainer.train(train_dataloader=loader, epochs=2) + result = model(**_batch(dataset)) + self.assertTrue(torch.isfinite(result["loss"])) + result["loss"].backward() + grads = [p.grad for p in model.parameters() if p.requires_grad] + self.assertTrue(any(g is not None and torch.any(g != 0) for g in grads)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_yaib_mimic4.py b/tests/core/test_yaib_mimic4.py index 06454c95a..ee1f87d09 100644 --- a/tests/core/test_yaib_mimic4.py +++ b/tests/core/test_yaib_mimic4.py @@ -7,6 +7,7 @@ from pyhealth.data import Patient from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import LevelEmbedding, Unified_RNN from pyhealth.processors import LevelConfig, SequenceProcessorV2, TimeseriesProcessor from pyhealth.tasks import ( YAIBAKIMIMIC4, @@ -593,7 +594,7 @@ def test_urine24_needs_twelve_hours(self): class TestYAIBProcessing(unittest.TestCase): - """The tasks' schemas: processors on task samples.""" + """The tasks' schemas: processors, and a model, on task samples.""" def test_dynamic_matches_the_previous_timeseries_preprocessing(self): # The sequence_v2 configuration reproduces TimeseriesProcessor with @@ -620,6 +621,35 @@ def hourly(grid): # the previous input format: datetimes and values torch.testing.assert_close(values, old[:, :4]) torch.testing.assert_close(mask.float(), old[:, 4:]) + def test_tasks_train_with_unified_rnn(self): + patients = [ + _patient( + stays=[(10 + i, 1, T0, 40 + 5 * i, range(40 + 5 * i))], + admissions=[_admission(hours=60 + 5 * i)], + ) + for i in range(2) + ] + dims: dict[str, list[LevelEmbedding]] = {"dynamic": [{"values": 8}]} + for task, label in ( + (YAIBMortalityMIMIC4(), "mortality"), + (YAIBLengthOfStayMIMIC4(), "los"), + (YAIBAKIMIMIC4(), "aki"), + ): + with self.subTest(task=task.task_name): + samples = [sample for patient in patients for sample in task(patient)] + if label == "mortality": # a binary label needs both classes + samples[0]["mortality"], samples[1]["mortality"] = 0, 1 + dataset = create_sample_dataset( + samples, task.input_schema, task.output_schema + ) + model = Unified_RNN(dataset, embedding_dim=dims, hidden_dim=8, mask="concat") + out = model(**next(iter(get_dataloader(dataset, batch_size=2)))) + self.assertTrue(torch.isfinite(out["loss"])) + expected = 2 if label == "mortality" else sum( + len(s[label]["time"]) for s in samples + ) + self.assertEqual(out["logit"].shape[0], expected) + if __name__ == "__main__": unittest.main() From 66334c77d24efeba6e4aaf0734ba6c5ebc9d7f62 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Sun, 1 Nov 2026 10:42:43 +0000 Subject: [PATCH 2/4] Point the YAIB tasks to Unified_RNN for training --- pyhealth/tasks/yaib_mimic4.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pyhealth/tasks/yaib_mimic4.py b/pyhealth/tasks/yaib_mimic4.py index fdb1838ea..19f4c0b04 100644 --- a/pyhealth/tasks/yaib_mimic4.py +++ b/pyhealth/tasks/yaib_mimic4.py @@ -241,10 +241,9 @@ class _YAIBMIMIC4Task(BaseTask): standardized like the other static features. Training needs a model that reads the ``sequence_v2`` input and the - per-hour label records, which :class:`~pyhealth.models.RNN` and - :class:`~pyhealth.models.MultimodalRNN` do not. The next pull request - of this stack (#1353) adds such a model; until then these tasks build - samples, but no model in PyHealth trains on them. + per-hour label records, such as :class:`~pyhealth.models.Unified_RNN`; + :class:`~pyhealth.models.RNN` and :class:`~pyhealth.models.MultimodalRNN` + do not. Note on ``o2sat`` (deliberate deviation from YAIB): YAIB pins ricu 0.6.0, whose MIMIC-IV definition of ``o2sat`` lists the blood gas item 50817 From 31cfc46ecb31f9563682ad9f1d83ef1f4be3feb6 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Sun, 1 Nov 2026 11:08:22 +0000 Subject: [PATCH 3/4] Rename Unified_RNN to RNN_V2 --- docs/api/models.rst | 2 +- docs/api/models/pyhealth.models.RNN.rst | 2 +- ...pyhealth.models.SequenceEmbeddingModel.rst | 2 +- pyhealth/models/__init__.py | 2 +- pyhealth/models/base_model.py | 2 +- pyhealth/models/rnn.py | 14 ++++++------ pyhealth/processors/timeseries_processor.py | 2 +- pyhealth/tasks/yaib_mimic4.py | 2 +- tests/core/test_per_step_labels.py | 4 ++-- .../{test_unified_rnn.py => test_rnn_v2.py} | 22 +++++++++---------- tests/core/test_timeseries_modes.py | 14 ++++++------ tests/core/test_yaib_mimic4.py | 6 ++--- 12 files changed, 37 insertions(+), 37 deletions(-) rename tests/core/{test_unified_rnn.py => test_rnn_v2.py} (95%) diff --git a/docs/api/models.rst b/docs/api/models.rst index f89b9c8b9..9f5439d40 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -39,7 +39,7 @@ routes each feature type automatically. * - ``MultimodalRNN`` - Features mix sequential codes with static tensors or multi-hot encodings - Auto-routes sequential features to RNN layers and non-sequential features to linear layers; good default for EHR - * - ``Unified_RNN`` + * - ``RNN_V2`` - Features are ``sequence_v2`` fields (nested codes, values and timestamps), plus static tensors - Encodes each field bottom-up over packed records, without padding the hierarchy * - :doc:`models/pyhealth.models.StageNet` diff --git a/docs/api/models/pyhealth.models.RNN.rst b/docs/api/models/pyhealth.models.RNN.rst index c6cf31cba..fc4055373 100644 --- a/docs/api/models/pyhealth.models.RNN.rst +++ b/docs/api/models/pyhealth.models.RNN.rst @@ -19,7 +19,7 @@ The separate callable RNNLayer and the complete RNN model. :undoc-members: :show-inheritance: -.. autoclass:: pyhealth.models.Unified_RNN +.. autoclass:: pyhealth.models.RNN_V2 :members: :undoc-members: :show-inheritance: diff --git a/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst b/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst index fee995340..6a19499df 100644 --- a/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst +++ b/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst @@ -5,7 +5,7 @@ Embeddings for the packed records of ``sequence_v2`` (:class:`~pyhealth.processors.SequenceProcessorV2`) fields. Each level's records get one embedding per token namespace and a projection of their values, concatenated; each level keeps its offsets and times. -:class:`~pyhealth.models.Unified_RNN` builds on it. +:class:`~pyhealth.models.RNN_V2` builds on it. .. autoclass:: pyhealth.models.SequenceEmbeddingModel :members: diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index 314dd9017..916e1a7be 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -28,7 +28,7 @@ from .molerec import MoleRec, MoleRecLayer from .retain import MultimodalRETAIN, RETAIN, RETAINLayer from .rnn import MultimodalRNN, RNN, RNNLayer -from .rnn import Unified_RNN as Unified_RNN +from .rnn import RNN_V2 as RNN_V2 from .safedrug import SafeDrug, SafeDrugLayer from .sparcnet import DenseBlock, DenseLayer, SparcNet, TransitionLayer from .stagenet import StageNet, StageNetLayer diff --git a/pyhealth/models/base_model.py b/pyhealth/models/base_model.py index 5a5c2d681..5c23d0d82 100644 --- a/pyhealth/models/base_model.py +++ b/pyhealth/models/base_model.py @@ -46,7 +46,7 @@ def _label_record_loss( raise ValueError( f"Expected one prediction per label record, {tuple(values.shape)}, " f"got {tuple(logits.shape)}. Per-step labels need a model that " - "predicts at each label's time, such as Unified_RNN." + "predicts at each label's time, such as RNN_V2." ) if values.numel() == 0: # A batch without label records: a zero loss that keeps the graph, diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py index 2b0741aca..43e0a6e43 100644 --- a/pyhealth/models/rnn.py +++ b/pyhealth/models/rnn.py @@ -239,7 +239,7 @@ def __init__( if self.mode in _PER_STEP_MODES: raise ValueError( "RNN does not support per-step (timeseries_*) labels; use " - "Unified_RNN with sequence_v2 inputs." + "RNN_V2 with sequence_v2 inputs." ) self.embedding_model = EmbeddingModel(dataset, embedding_dim) @@ -439,7 +439,7 @@ def __init__( if self.mode in _PER_STEP_MODES: raise ValueError( "MultimodalRNN does not support per-step (timeseries_*) labels; " - "use Unified_RNN with sequence_v2 inputs." + "use RNN_V2 with sequence_v2 inputs." ) self.embedding_model = EmbeddingModel(dataset, embedding_dim) @@ -592,7 +592,7 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]: # ---------------------------------------------------------------------------- -# Unified_RNN: RNNs over the packed records of SequenceProcessorV2 fields +# RNN_V2: RNNs over the packed records of SequenceProcessorV2 fields # ---------------------------------------------------------------------------- RNNType = Literal["RNN", "LSTM", "GRU"] @@ -771,7 +771,7 @@ def forward( return summary, states, out[self.depth]["time"] -class Unified_RNN(BaseModel): +class RNN_V2(BaseModel): """RNN model over :class:`~pyhealth.processors.SequenceProcessorV2` fields. Each V2 field is embedded by :class:`SequenceEmbeddingModel` and encoded @@ -837,7 +837,7 @@ class Unified_RNN(BaseModel): ... {"age": "tensor", "visits": ("sequence_v2", {"levels": levels})}, ... {"label": "binary"}, ... ) - >>> model = Unified_RNN( + >>> model = RNN_V2( ... dataset, embedding_dim={"visits": [{"tokens": (8,)}, {}]}, hidden_dim=16 ... ) >>> out = model(**next(iter(get_dataloader(dataset, batch_size=4)))) @@ -861,7 +861,7 @@ def __init__( ): super().__init__(dataset) if len(self.label_keys) != 1: - raise ValueError("Unified_RNN supports exactly one label.") + raise ValueError("RNN_V2 supports exactly one label.") self.label_key = self.label_keys[0] self.per_step = self.mode in _PER_STEP_MODES @@ -876,7 +876,7 @@ def __init__( else: raise TypeError( f"Field '{field}': {type(processor).__name__} is not supported. " - "Unified_RNN takes 'sequence_v2' fields and static 'tensor' " + "RNN_V2 takes 'sequence_v2' fields and static 'tensor' " "fields; express sequences, sets and tokens with 'sequence_v2'." ) diff --git a/pyhealth/processors/timeseries_processor.py b/pyhealth/processors/timeseries_processor.py index 92e4f5cbd..1dc8c5eff 100644 --- a/pyhealth/processors/timeseries_processor.py +++ b/pyhealth/processors/timeseries_processor.py @@ -202,7 +202,7 @@ class TimeseriesLabelProcessor(SequenceProcessorV2): raises. There is no grid, resampling, imputation or mask. A model predicts one value per label, at its time, from the input records - at or before it (see :class:`~pyhealth.models.Unified_RNN`). + at or before it (see :class:`~pyhealth.models.RNN_V2`). Output: - ``(values, time)``: ``[N, 1]`` values and ``[N]`` times. A batch diff --git a/pyhealth/tasks/yaib_mimic4.py b/pyhealth/tasks/yaib_mimic4.py index 19f4c0b04..8a4920c9c 100644 --- a/pyhealth/tasks/yaib_mimic4.py +++ b/pyhealth/tasks/yaib_mimic4.py @@ -241,7 +241,7 @@ class _YAIBMIMIC4Task(BaseTask): standardized like the other static features. Training needs a model that reads the ``sequence_v2`` input and the - per-hour label records, such as :class:`~pyhealth.models.Unified_RNN`; + per-hour label records, such as :class:`~pyhealth.models.RNN_V2`; :class:`~pyhealth.models.RNN` and :class:`~pyhealth.models.MultimodalRNN` do not. diff --git a/tests/core/test_per_step_labels.py b/tests/core/test_per_step_labels.py index 66e74e969..4f3af1186 100644 --- a/tests/core/test_per_step_labels.py +++ b/tests/core/test_per_step_labels.py @@ -7,7 +7,7 @@ import torch from pyhealth.datasets import SampleDataset, create_sample_dataset, get_dataloader -from pyhealth.models import BaseModel, Unified_RNN +from pyhealth.models import BaseModel, RNN_V2 from pyhealth.processors import ( LevelConfig, SequenceProcessorV2, @@ -118,7 +118,7 @@ class TestPerStepEndToEnd(unittest.TestCase): def test_label_processors_train_and_evaluate(self): models: dict[str, Callable[[SampleDataset], BaseModel]] = { "LabelRecordModel": LabelRecordModel, - "Unified_RNN": lambda dataset: Unified_RNN( + "RNN_V2": lambda dataset: RNN_V2( dataset, embedding_dim={"dynamic": [{"values": 8}]}, hidden_dim=8 ), } diff --git a/tests/core/test_unified_rnn.py b/tests/core/test_rnn_v2.py similarity index 95% rename from tests/core/test_unified_rnn.py rename to tests/core/test_rnn_v2.py index bb2657623..fd7119187 100644 --- a/tests/core/test_unified_rnn.py +++ b/tests/core/test_rnn_v2.py @@ -4,7 +4,7 @@ import torch from pyhealth.datasets import create_sample_dataset, get_dataloader -from pyhealth.models import LevelEmbedding, Unified_RNN +from pyhealth.models import LevelEmbedding, RNN_V2 from pyhealth.models.rnn import _gaps, _pool from pyhealth.processors import LevelConfig from pyhealth.trainer import Trainer @@ -54,8 +54,8 @@ def _dataset(output: str = "binary", **fields: Any): return create_sample_dataset(samples, schema, {"y": output}) -def _model(dataset, **kwargs: Any) -> Unified_RNN: - model = Unified_RNN(dataset, dropout=0.0, **kwargs) +def _model(dataset, **kwargs: Any) -> RNN_V2: + model = RNN_V2(dataset, dropout=0.0, **kwargs) model.eval() return model @@ -64,7 +64,7 @@ def _batch(dataset) -> dict[str, Any]: return next(iter(get_dataloader(dataset, batch_size=len(dataset)))) -def _encode(model: Unified_RNN, batch: dict[str, Any], field: str): +def _encode(model: RNN_V2, batch: dict[str, Any], field: str): embedded = model.embedding({field: batch[field]})[field] return embedded, cast(Any, model.encoders[field])(embedded) @@ -83,7 +83,7 @@ def test_rejects_bad_arguments(self): ] for kwargs in bad: with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): - Unified_RNN(self.dataset, embedding_dim=VISIT_DIMS, **kwargs) + RNN_V2(self.dataset, embedding_dim=VISIT_DIMS, **kwargs) def test_needs_exactly_one_label(self): samples = [ @@ -96,14 +96,14 @@ def test_needs_exactly_one_label(self): {"y": "binary", "z": "binary"}, ) with self.assertRaisesRegex(ValueError, "exactly one label"): - Unified_RNN(dataset, embedding_dim=VISIT_DIMS) + RNN_V2(dataset, embedding_dim=VISIT_DIMS) def test_rejects_a_level_with_nothing_to_encode(self): times: list[LevelConfig] = [{"ordering": "series"}] # time only records = [{"time": [0.0, 1.0]}, {"time": [0.5]}] dataset = _dataset(x=(times, records)) with self.assertRaisesRegex(ValueError, "nothing to encode"): - Unified_RNN(dataset, embedding_dim={"x": [{}]}) + RNN_V2(dataset, embedding_dim={"x": [{}]}) # The time since the previous record is something to encode. model = _model(dataset, embedding_dim={"x": [{}]}, time_feature=True) self.assertEqual(model(**_batch(dataset))["logit"].shape, (2, 1)) @@ -113,7 +113,7 @@ def test_rejects_v1_sequences_sets_and_tokens(self): samples = [{"patient_id": f"p{i}", "x": value, "y": i} for i in range(2)] dataset = create_sample_dataset(samples, {"x": spec}, {"y": "binary"}) with self.subTest(spec=spec), self.assertRaises(TypeError): - Unified_RNN(dataset, embedding_dim={}) + RNN_V2(dataset, embedding_dim={}) class TestBottomUp(unittest.TestCase): @@ -244,13 +244,13 @@ def test_top_level_must_be_a_series(self): ] dataset = self.dataset(levels, [{"children": [{"tokens": ["A"]}]}] * 2) with self.subTest(top=top), self.assertRaisesRegex(ValueError, "Wrap"): - Unified_RNN(dataset, embedding_dim={"x": [{"tokens": (4,)}, {}]}) + RNN_V2(dataset, embedding_dim={"x": [{"tokens": (4,)}, {}]}) def test_no_bidirectional(self): records = [_visits((0.0, "er", ["A"]))] * 2 dataset = self.dataset(VISITS, records) with self.assertRaises(ValueError): - Unified_RNN( + RNN_V2( dataset, embedding_dim={"x": VISIT_DIMS["visits"]}, bidirectional=True ) @@ -322,7 +322,7 @@ def test_resampled_top_level(self): class TestTraining(unittest.TestCase): def test_trains(self): dataset = _dataset(visits=(VISITS, VISIT_RECORDS)) - model = Unified_RNN( + model = RNN_V2( dataset, embedding_dim=VISIT_DIMS, hidden_dim=8 ) loader = get_dataloader(dataset, batch_size=2, shuffle=True, seed=0) diff --git a/tests/core/test_timeseries_modes.py b/tests/core/test_timeseries_modes.py index e39b1e571..a1eddb651 100644 --- a/tests/core/test_timeseries_modes.py +++ b/tests/core/test_timeseries_modes.py @@ -9,7 +9,7 @@ timeseries_binary_metrics_fn, timeseries_regression_metrics_fn, ) -from pyhealth.models import RNN, LevelEmbedding, MultimodalRNN, Unified_RNN +from pyhealth.models import RNN, LevelEmbedding, MultimodalRNN, RNN_V2 from pyhealth.processors import ( LevelConfig, MultiHotProcessor, @@ -62,8 +62,8 @@ def _model(dataset) -> LabelRecordModel: return model -def _unified_rnn(dataset) -> Unified_RNN: - model = Unified_RNN(dataset, embedding_dim=DIMS, hidden_dim=8, dropout=0.0) +def _rnn_v2(dataset) -> RNN_V2: + model = RNN_V2(dataset, embedding_dim=DIMS, hidden_dim=8, dropout=0.0) model.eval() return model @@ -87,7 +87,7 @@ def test_one_prediction_per_label_record(self): # Sample 1 has no label at hour 0: it is left out, not masked. dataset = _dataset(labeled=lambda i, h: (i, h) != (1, 0)) batch = next(iter(get_dataloader(dataset, batch_size=4))) - for make_model in (_model, _unified_rnn): + for make_model in (_model, _rnn_v2): with self.subTest(model=make_model.__name__), torch.no_grad(): out = make_model(dataset)(**batch) values, time, offsets = out["y_true"] @@ -99,9 +99,9 @@ def test_one_prediction_per_label_record(self): ) torch.testing.assert_close(loss, out["loss"]) - def test_unified_rnn_predictions_do_not_see_later_records(self): + def test_rnn_v2_predictions_do_not_see_later_records(self): dataset = _dataset() - model = _unified_rnn(dataset) + model = _rnn_v2(dataset) batch = next(iter(get_dataloader(dataset, batch_size=1))) with torch.no_grad(): before = model(**batch)["y_prob"][:, 0] @@ -163,7 +163,7 @@ def test_v1_rnns_reject_per_step_labels(self): model(dataset=dataset) def test_trainer_scores_label_records_per_sample(self): - for make_model in (_model, _unified_rnn): + for make_model in (_model, _rnn_v2): with self.subTest(model=make_model.__name__): dataset = _dataset() trainer = Trainer( diff --git a/tests/core/test_yaib_mimic4.py b/tests/core/test_yaib_mimic4.py index ee1f87d09..845839a6e 100644 --- a/tests/core/test_yaib_mimic4.py +++ b/tests/core/test_yaib_mimic4.py @@ -7,7 +7,7 @@ from pyhealth.data import Patient from pyhealth.datasets import create_sample_dataset, get_dataloader -from pyhealth.models import LevelEmbedding, Unified_RNN +from pyhealth.models import LevelEmbedding, RNN_V2 from pyhealth.processors import LevelConfig, SequenceProcessorV2, TimeseriesProcessor from pyhealth.tasks import ( YAIBAKIMIMIC4, @@ -621,7 +621,7 @@ def hourly(grid): # the previous input format: datetimes and values torch.testing.assert_close(values, old[:, :4]) torch.testing.assert_close(mask.float(), old[:, 4:]) - def test_tasks_train_with_unified_rnn(self): + def test_tasks_train_with_rnn_v2(self): patients = [ _patient( stays=[(10 + i, 1, T0, 40 + 5 * i, range(40 + 5 * i))], @@ -642,7 +642,7 @@ def test_tasks_train_with_unified_rnn(self): dataset = create_sample_dataset( samples, task.input_schema, task.output_schema ) - model = Unified_RNN(dataset, embedding_dim=dims, hidden_dim=8, mask="concat") + model = RNN_V2(dataset, embedding_dim=dims, hidden_dim=8, mask="concat") out = model(**next(iter(get_dataloader(dataset, batch_size=2)))) self.assertTrue(torch.isfinite(out["loss"])) expected = 2 if label == "mortality" else sum( From 89c68c029165028d4179333b4aa11f1b5bd7247f Mon Sep 17 00:00:00 2001 From: Logiquo Date: Sun, 1 Nov 2026 23:41:59 +0000 Subject: [PATCH 4/4] Document that each level's time is its own and a record contains its children SequenceProcessorV2 levels each have their own time unit and origin, and a record forms one unit with its children; times are only compared within a level. Records, Resample and LevelConfig said "the task's time unit", which suggested a shared timeline. RNN_V2's per-step guarantee is now stated for top-level records: a label uses each top-level record at or before it with everything it contains, and lower levels' times are not compared to labels. A test pins this down. --- pyhealth/models/rnn.py | 10 ++++++-- pyhealth/processors/sequence_processor_v2.py | 27 +++++++++++++------- tests/core/test_rnn_v2.py | 27 ++++++++++++++++++++ 3 files changed, 53 insertions(+), 11 deletions(-) diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py index 43e0a6e43..e64fb27ba 100644 --- a/pyhealth/models/rnn.py +++ b/pyhealth/models/rnn.py @@ -786,6 +786,9 @@ class RNN_V2(BaseModel): of each record's children. A ``set`` level mean-pools each group; a ``sequence`` / ``series`` level runs its RNN over each group and keeps the last output, giving one summary per parent record, up to one per sample. + So a record and everything below it are encoded as one unit. Following + :class:`~pyhealth.processors.SequenceProcessorV2`, each level's time has + its own unit and origin, and only orders the records of a group. A group without records (e.g. a stay without labs) is skipped by the RNN and summarized as zeros. With ``time_feature``, each record of a timed @@ -796,8 +799,11 @@ class RNN_V2(BaseModel): a :class:`~pyhealth.processors.TimeseriesLabelProcessor`): the model predicts once per label record. The label's times are in the unit of the fields' top level, with the same origin, and a label at time ``t`` uses - each field's top-level state after its last record with time ``<= t`` - (zeros before the first) and the statics, so it never sees later records. + each field's top-level state after its last top-level record with time + ``<= t`` (zeros before the first) and the statics, so it never sees a + later top-level record. That state covers everything the record contains: + lower levels' times are not on the label's timeline and are not compared + to it, so data that must be cut at label times belongs at the top level. Every V2 field must have a ``series`` top level; a resampled one counts, with record ``i`` at time ``i * step``. Wrap a ``set`` or ``sequence`` top level in a ``series`` with ``t = 0``. Logits are ``[label records, diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index d5e9259f6..ca840fa63 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -48,9 +48,15 @@ class Records(TypedDict, total=False): a key it needs must be present (its column may be empty), and any other key is ignored. + A record contains its children: ``children[i]`` is part of record ``i``, + and together they are one unit at record ``i``'s place in its level. Each + level's time has its own unit and origin (e.g. hours since admission for + stays, minutes since the stay's start for its labs), so times are only + compared within a level, never across levels. + Keys: - time: ``[N]`` numbers in the task's own unit. Read on ``series`` - levels. + time: ``[N]`` numbers in this level's own unit and origin. Read on + ``series`` levels. tokens: ``[N]`` tokens, or ``[N]`` tuples of ``K`` tokens when the level has ``K`` token namespaces. values: ``[N, *S]`` numbers. Scalars are ``[N, 1]``; a flat ``[N]`` @@ -84,10 +90,10 @@ class Resample(TypedDict): Each grid step becomes one record, and the level becomes a ``sequence`` without time: record ``i`` is the step starting at ``i * step``. The grid - starts at the task's time 0, so negative time raises. + starts at the level's time 0, so negative time raises. Keys: - step: Grid step, in the task's time unit. + step: Grid step, in the level's time unit. method: How the records in one step become one record. - ``"first"`` / ``"last"``: keep the first / last record, with @@ -234,7 +240,8 @@ class LevelConfig(TypedDict, total=False): Keys: ordering: ``"set"``, ``"sequence"`` (default) or ``"series"``. - Describes the raw input; ``series`` levels carry ``time``. + Describes the raw input; ``series`` levels carry ``time``, in + the level's own unit (see :class:`Records`). tokens: Number of token namespaces ``K``; the level's records have ``tokens``, plain strings when ``K = 1`` and ``K``-tuples otherwise. values: Value shape ``S``, e.g. ``[1]`` for scalars; the level's @@ -988,10 +995,12 @@ class SequenceProcessorV2(FeatureProcessor): Reads a :class:`Records` and emits a tuple of tensors in :meth:`schema` order: by level, then by role (``tokens``, ``values``, ``mask``, ``time``, ``offsets``). ``offsets`` at level ``l`` partitions level ``l`` records - into level ``l + 1`` records, in CSR layout. Which roles a level has, the - number of token namespaces ``K`` and the value shape ``S`` come from the - configuration; :meth:`fit` only builds the vocabularies and the statistics - of imputation and normalization. + into level ``l + 1`` records, in CSR layout: a level ``l + 1`` record + contains its group of level ``l`` records, and each level's time is its + own (see :class:`Records`). Which roles a level has, the number of token + namespaces ``K`` and the value shape ``S`` come from the configuration; + :meth:`fit` only builds the vocabularies and the statistics of imputation + and normalization. Subclasses are configuration only: their ``__init__`` passes fixed ``levels`` to ``super().__init__()`` and adds no logic. diff --git a/tests/core/test_rnn_v2.py b/tests/core/test_rnn_v2.py index fd7119187..c57a3f1a7 100644 --- a/tests/core/test_rnn_v2.py +++ b/tests/core/test_rnn_v2.py @@ -281,6 +281,33 @@ def test_each_label_uses_the_records_at_or_before_it(self): # The static age goes with every label record of its sample. self.assertEqual(result["embed"][:, 5].tolist(), [0.0] * 3 + [1.0] * 3) + def test_a_top_level_record_includes_all_its_children(self): + # Stays (series) of timed labs (series); lab times are in their own unit. + levels: list[LevelConfig] = [ + {"ordering": "series", "tokens": 1}, + {"ordering": "series", "tokens": 1}, + ] + stays = { + "time": [0.0, 2.0], + "tokens": ["er", "op"], + "children": [ + {"time": [0.0, 5.0, 9.0], "tokens": ["A", "B", "C"]}, + {"time": [0.0], "tokens": ["D"]}, + ], + } + dataset = self.dataset(levels, [stays] * 2) + model = _model( + dataset, + embedding_dim={"x": [{"tokens": (4,)}, {"tokens": (3,)}]}, + hidden_dim=5, + ) + batch = _batch(dataset) + embed = model(**batch, embed=True)["embed"][:, :5] + _, (_, states, _) = _encode(model, batch, "x") + # The label at 1.0 sees the whole stay at 0, labs at 5 and 9 included. + torch.testing.assert_close(embed[1], states[0]) + torch.testing.assert_close(embed[2], states[1]) + def test_labels_between_records(self): records = [_visits((0.0, "er", ["A"]), (2.0, "op", ["B"]))] * 2 dataset = self.dataset(VISITS, records, times=(1.5, 2.5))