From 8e64f6fba41ccb992f237c78833b519e3b01e2c9 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Fri, 30 Oct 2026 11:18:49 +0000 Subject: [PATCH] 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 8c76693d8..5547551ec 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." ) return loss_fn(logits, values) 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 c9d9b340a..3832e0bb7 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, TimeseriesBinaryLabelProcessor, @@ -20,6 +20,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( @@ -60,6 +61,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"): @@ -78,18 +85,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()) @@ -122,19 +143,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()