diff --git a/docs/api/models.rst b/docs/api/models.rst index f89b9c8b9..c48d439cf 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -36,6 +36,12 @@ routes each feature type automatically. * - :doc:`models/pyhealth.models.MLP` - Features are static numeric vectors (aggregated lab values, demographics) - Fully connected; no notion of sequence order + * - :doc:`models/pyhealth.models.Unified_MLP` + - Features are nested ``sequence_v2`` records and static tensors + - MLPs at every hierarchy level and on static fields, with bottom-up mean pooling + * - :doc:`models/pyhealth.models.Unified_LogisticRegression` + - You want a linear-head baseline on the same packed records + - Learned embeddings, bottom-up means, and raw static vectors; supports per-step prefix means * - ``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 @@ -213,6 +219,8 @@ API Reference models/pyhealth.models.BaseModel models/pyhealth.models.LogisticRegression models/pyhealth.models.MLP + models/pyhealth.models.Unified_MLP + models/pyhealth.models.Unified_LogisticRegression models/pyhealth.models.CNN models/pyhealth.models.RNN models/pyhealth.models.GNN diff --git a/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst b/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst index fee995340..4d87fd39c 100644 --- a/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst +++ b/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst @@ -5,7 +5,8 @@ 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.Unified_RNN`, :class:`~pyhealth.models.Unified_MLP` and +:class:`~pyhealth.models.Unified_LogisticRegression` build on it. .. autoclass:: pyhealth.models.SequenceEmbeddingModel :members: diff --git a/docs/api/models/pyhealth.models.Unified_LogisticRegression.rst b/docs/api/models/pyhealth.models.Unified_LogisticRegression.rst new file mode 100644 index 000000000..2ac0ac57b --- /dev/null +++ b/docs/api/models/pyhealth.models.Unified_LogisticRegression.rst @@ -0,0 +1,17 @@ +pyhealth.models.Unified_LogisticRegression +========================================= + +A linear prediction head over bottom-up means of learned record embeddings +and flattened static tensors. Unlike :class:`~pyhealth.models.Unified_MLP`, +it introduces no MLPs at hierarchy levels or on static fields. + +.. autoclass:: pyhealth.models.Unified_LogisticRegression + :members: + :show-inheritance: + +See :doc:`pyhealth.models.Unified_MLP` for a shared input example and the +spatial, empty-group, and per-step prediction conventions. + +The existing :class:`~pyhealth.models.LogisticRegression` remains available +for legacy schemas. Both versions use learned embeddings rather than +classical logistic regression directly on raw features. diff --git a/docs/api/models/pyhealth.models.Unified_MLP.rst b/docs/api/models/pyhealth.models.Unified_MLP.rst new file mode 100644 index 000000000..5c7aa4ef9 --- /dev/null +++ b/docs/api/models/pyhealth.models.Unified_MLP.rst @@ -0,0 +1,52 @@ +pyhealth.models.Unified_MLP +========================== + +An MLP baseline for packed hierarchical records and static tensors. +Each level transforms its records before mean pooling into their parents. +Static tensors are flattened per sample and passed through independent MLPs. + +.. autoclass:: pyhealth.models.Unified_MLP + :members: + :show-inheritance: + +Example +------- + +.. code-block:: python + + from pyhealth.datasets import create_sample_dataset, get_dataloader + from pyhealth.models import Unified_MLP, Unified_LogisticRegression + + dataset = create_sample_dataset( + [ + {"patient_id": "p0", "age": [50.0], "visits": { + "children": [{"tokens": ["A", "B"]}, {"tokens": []}] + }, "label": 0}, + {"patient_id": "p1", "age": [65.0], "visits": { + "children": [{"tokens": ["B"]}] + }, "label": 1}, + ], + {"age": "tensor", "visits": ("sequence_v2", {"levels": [ + {"ordering": "set", "tokens": 1}, {"ordering": "sequence"} + ]})}, + {"label": "binary"}, + ) + dims = {"visits": [{"tokens": (8,)}, {}]} + batch = next(iter(get_dataloader(dataset, batch_size=2))) + model = Unified_MLP(dataset, dims, hidden_dim=16) + result = model(**batch) + linear_baseline = Unified_LogisticRegression(dataset, dims) + linear_result = linear_baseline(**batch) + +For per-step labels, give each sequence field a ``series`` top level. +At label time ``t``, both models mean-pool the available top-level record +representations with timestamps ``<= t``. No record yields a zero summary. +Resampled record ``i`` has timestamp ``i * step``. A top-level timestamp +must indicate when the complete record, including its children, is available. + +Record ordering and timestamps are not input features. Spatial arrays are +flattened and projected within each record; these models do not learn spatial +locality. Nested means weight parent records equally, so a parent with many +children does not receive more weight than one with few children. + +The existing :class:`~pyhealth.models.MLP` remains available for legacy schemas. diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index 314dd9017..02f885c1c 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -16,6 +16,7 @@ from .gamenet import GAMENet, GAMENetLayer from .jamba_ehr import JambaEHR, JambaLayer from .logistic_regression import LogisticRegression +from .logistic_regression import Unified_LogisticRegression as Unified_LogisticRegression from .gan import GAN from .gnn import GAT, GCN from .graph_torchvision_model import Graph_TorchvisionModel @@ -25,6 +26,7 @@ from .medlink import MedLink from .micron import MICRON, MICRONLayer from .mlp import MLP +from .mlp import Unified_MLP as Unified_MLP from .molerec import MoleRec, MoleRecLayer from .retain import MultimodalRETAIN, RETAIN, RETAINLayer from .rnn import MultimodalRNN, RNN, RNNLayer diff --git a/pyhealth/models/logistic_regression.py b/pyhealth/models/logistic_regression.py index 8155d101f..2c95409da 100644 --- a/pyhealth/models/logistic_regression.py +++ b/pyhealth/models/logistic_regression.py @@ -1,12 +1,22 @@ -from typing import Dict +from collections.abc import Mapping, Sequence +from typing import Any, cast import torch import torch.nn as nn from pyhealth.datasets import SampleDataset from pyhealth.models import BaseModel +from pyhealth.processors import SequenceProcessorV2, TensorProcessor +from pyhealth.processors.sequence_processor_v2 import Resample from .embedding import EmbeddingModel +from .rnn import _groups, _pool +from .sequence_embedding import ( + LevelEmbedding, + MaskMode, + PretrainedEmbedding, + SequenceEmbeddingModel, +) class LogisticRegression(BaseModel): @@ -114,7 +124,7 @@ def mean_pooling(x, mask): """ return x.sum(dim=1) / mask.sum(dim=1, keepdim=True) - def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + def forward(self, **kwargs) -> dict[str, torch.Tensor]: """Forward propagation. Args: @@ -210,6 +220,184 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: return results +class Unified_LogisticRegression(BaseModel): + """Linear prediction over bottom-up means of learned record embeddings. + + Accepts ``sequence_v2`` fields and static ``tensor`` fields. At each level, + a record's embedding is concatenated with its children's mean summary, + then records are mean-pooled within their parent. Levels may contain + multiple token namespaces, numeric values, both, or only child groups. + Spatial value axes are flattened and projected within each record by + :class:`SequenceEmbeddingModel`. Sequence order and timestamps do not + become input features. Empty groups give zero summaries; parents with + empty children still contribute their own embeddings. + + Field summaries and flattened static tensors are concatenated and passed + to a single linear head. This is a linear head over learned embeddings, + rather than classical logistic regression on raw inputs. Nested means + weight parents equally, not individual leaves across different parents. + + Per-step binary/regression labels use the cumulative mean of top-level + records with time ``<=`` the label time, or zeros before the first record. + Every sequence field must then have a ``series`` top level. Resampled + record ``i`` has time ``i * step``. Each top-level timestamp must indicate + when the entire record, including its children, is available. Static + vectors are repeated for each label record of their sample. + + Args: + dataset: Dataset with fitted processors and exactly one label. + embedding_dim: One :class:`LevelEmbedding` per level of each V2 + field, innermost first. Use ``{}`` for static-only datasets. + mask: ``"ignore"`` or ``"concat"`` for numeric missing indicators. + pretrain: Pretrained token namespace embeddings. + + Examples: + >>> from pyhealth.datasets import create_sample_dataset, get_dataloader + >>> dataset = create_sample_dataset( + ... [{"patient_id": str(i), "age": [50.0 + i], "y": i} + ... for i in range(2)], + ... {"age": "tensor"}, {"y": "binary"}, + ... ) + >>> model = Unified_LogisticRegression(dataset, embedding_dim={}) + >>> model(**next(iter(get_dataloader(dataset, batch_size=2))))["logit"].shape + torch.Size([2, 1]) + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: Mapping[str, Sequence[LevelEmbedding]], + mask: MaskMode = "ignore", + pretrain: Sequence[PretrainedEmbedding] | None = None, + ) -> None: + super().__init__(dataset) + if len(self.label_keys) != 1: + raise ValueError("Unified pooling models support exactly one label.") + if not self.feature_keys: + raise ValueError("At least one input field is required.") + self.label_key = self.label_keys[0] + self.per_step = self.mode in ("timeseries_binary", "timeseries_regression") + self.processors: dict[str, SequenceProcessorV2] = {} + self.static_fields: list[str] = [] + for field in self.feature_keys: + if isinstance( + processor := dataset.input_processors[field], SequenceProcessorV2 + ): + self.processors[field] = processor + if ( + self.per_step + and self.processors[field].levels[-1].get("ordering", "sequence") + != "series" + ): + raise ValueError( + f"'{field}': per-step labels require a 'series' top level." + ) + elif isinstance(dataset.input_processors[field], TensorProcessor): + self.static_fields.append(field) + else: + raise TypeError( + f"'{field}': use 'sequence_v2' or static 'tensor' inputs." + ) + self.embedding = SequenceEmbeddingModel( + dataset, embedding_dim=embedding_dim, mask=mask, pretrain=pretrain + ) + self.encoders = nn.ModuleDict() + self.static_encoders = nn.ModuleDict() + self.field_dims: dict[str, int] = {} + for field, processor in self.processors.items(): + self.encoders[field] = nn.ModuleList( + nn.Identity() for _ in range(processor.depth + 1) + ) + self.field_dims[field] = sum(self.embedding.embedding_dim_per_level(field)) + if self.embedding.embedding_dim_per_level(field)[0] == 0: + raise ValueError( + f"'{field}': the innermost level has nothing to encode." + ) + for field in self.static_fields: + self.static_encoders[field] = nn.Identity() + self.field_dims[field] = int(torch.as_tensor(dataset[0][field]).numel()) + self.fc = nn.Linear(sum(self.field_dims.values()), self.get_output_size()) + + def forward( + self, **batch: torch.Tensor | tuple[torch.Tensor, ...] | bool + ) -> dict[str, Any]: + """Return loss, probabilities, targets, logits, and optionally ``embed``.""" + embedded = self.embedding({field: batch[field] for field in self.processors}) + label: tuple[torch.Tensor, ...] = () + label_sample = torch.empty(0, dtype=torch.long, device=self.device) + if self.per_step: + label = tuple( + t.to(self.device) + for t in cast(tuple[torch.Tensor, ...], batch[self.label_key]) + ) + label_sample, _ = _groups(label[2]) + parts: list[torch.Tensor] = [] + for field, processor in self.processors.items(): + summary: torch.Tensor | None = None + for level in range(processor.depth + 1): + rows: list[torch.Tensor] = [] + if embedded[field][level]["embedding"] is not None: + rows.append(embedded[field][level]["embedding"]) + if summary is not None: + rows.append(summary) + x = cast(nn.ModuleList, self.encoders[field])[level]( + torch.cat(rows, dim=-1) + ) + offsets = embedded[field][level]["offsets"] + if self.per_step and level == processor.depth: + # Search each sample's times, then subtract cumulative sums + # at its start to obtain an independent causal prefix mean. + group, position = _groups(offsets) + time = embedded[field][level]["time"] + if time is None: + time = position.to(x.dtype) * float( + cast(Resample, processor.levels[level].get("resample"))[ + "step" + ] + ) + padded = time.new_full( + (len(offsets) - 1, max(int(offsets.diff().max()), 1)), + float("inf"), + ) + padded[group, position] = time + seen = torch.searchsorted( + padded[label_sample], + cast(tuple[torch.Tensor, ...], label)[1] + .to(time.dtype) + .unsqueeze(1), + right=True, + )[:, 0] + prefix = torch.cat([x.new_zeros(1, x.shape[1]), x.cumsum(0)]) + summary = ( + prefix[offsets[:-1][label_sample] + seen] + - prefix[offsets[:-1][label_sample]] + ) / seen.clamp(min=1).unsqueeze(1).to(x.dtype) + else: + summary = _pool(x, offsets) + parts.append(cast(torch.Tensor, summary)) + for field in self.static_fields: + x = self.static_encoders[field]( + cast(torch.Tensor, batch[field]).to(self.device).float().flatten(1) + ) + 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 cast(torch.Tensor, batch[self.label_key]).to(self.device) + ) + results: dict[str, torch.Tensor | tuple[torch.Tensor, ...]] = { + "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 + + if __name__ == "__main__": from pyhealth.datasets import create_sample_dataset diff --git a/pyhealth/models/mlp.py b/pyhealth/models/mlp.py index 299dc151e..1debf0b29 100644 --- a/pyhealth/models/mlp.py +++ b/pyhealth/models/mlp.py @@ -1,13 +1,17 @@ +from collections.abc import Mapping, Sequence from typing import Dict, cast import torch import torch.nn as nn from pyhealth.datasets import SampleDataset -from pyhealth.models import BaseModel from pyhealth.interpret.api import Interpretable +from pyhealth.models import BaseModel from .embedding import EmbeddingModel +from .logistic_regression import Unified_LogisticRegression +from .rnn import _per_field +from .sequence_embedding import LevelEmbedding, MaskMode, PretrainedEmbedding class MLP(BaseModel, Interpretable): @@ -390,6 +394,103 @@ def get_embedding_model(self) -> nn.Module | None: return self.embedding_model +class Unified_MLP(Unified_LogisticRegression): + """Bottom-up record MLPs and mean pooling, plus MLPs for static fields. + + Uses the input, hierarchy, spatial flattening, and per-step prefix means + described in :class:`Unified_LogisticRegression`. Each hierarchy level + has its own MLP, applied to records before pooling. Static tensors are + flattened per sample and transformed by their own MLPs. All field outputs + are concatenated and passed to a linear prediction head. + + Args: + dataset: Dataset with fitted processors and exactly one label. + embedding_dim: Per V2 field, one :class:`LevelEmbedding` per level. + hidden_dim: MLP output width at every level, including static MLPs. + One value for all fields or a field mapping; default 128. + num_layers: Linear layers per MLP, at least one. One value or a + field mapping; default 2. ReLU and dropout occur between layers. + dropout: Dropout between linear layers. One value or a field mapping; + default 0.5. Has no effect with one layer. + mask: ``"ignore"`` or ``"concat"`` for numeric missing indicators. + pretrain: Pretrained token namespace embeddings. + + Unlisted fields in mappings use the default. Token/value embedding sizes + are independent of MLP hidden sizes. Neither static nor sequence fields + bypass the MLP, so static-only datasets also receive nonlinear modeling + with the default two layers. + + Examples: + >>> from pyhealth.datasets import create_sample_dataset, get_dataloader + >>> dataset = create_sample_dataset( + ... [{"patient_id": str(i), "age": [50.0 + i], "y": i} + ... for i in range(2)], + ... {"age": "tensor"}, {"y": "binary"}, + ... ) + >>> model = Unified_MLP(dataset, embedding_dim={}, hidden_dim=16) + >>> model(**next(iter(get_dataloader(dataset, batch_size=2))))["logit"].shape + torch.Size([2, 1]) + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: Mapping[str, Sequence[LevelEmbedding]], + hidden_dim: int | Mapping[str, int] = 128, + num_layers: int | Mapping[str, int] = 2, + dropout: float | Mapping[str, float] = 0.5, + mask: MaskMode = "ignore", + pretrain: Sequence[PretrainedEmbedding] | None = None, + ) -> None: + super().__init__(dataset, embedding_dim, mask=mask, pretrain=pretrain) + hidden: dict[str, int] = _per_field( + "hidden_dim", + hidden_dim, + self.feature_keys, + 128, + lambda v: isinstance(v, int) and not isinstance(v, bool) and v > 0, + ) + depths: dict[str, int] = _per_field( + "num_layers", + num_layers, + self.feature_keys, + 2, + lambda v: isinstance(v, int) and not isinstance(v, bool) and v > 0, + ) + dropouts: dict[str, float] = _per_field( + "dropout", + dropout, + self.feature_keys, + 0.5, + lambda v: isinstance(v, (int, float)) and 0 <= v < 1, + ) + for field in self.feature_keys: + below = 0 + for level, width in enumerate( + self.embedding.embedding_dim_per_level(field) + if field in self.processors + else [self.field_dims[field]] + ): + layers: list[nn.Module] = [nn.Linear(width + below, hidden[field])] + for _ in range(depths[field] - 1): + layers.extend( + [ + nn.ReLU(), + nn.Dropout(dropouts[field]), + nn.Linear(hidden[field], hidden[field]), + ] + ) + if field in self.processors: + cast(nn.ModuleList, self.encoders[field])[level] = nn.Sequential( + *layers + ) + else: + self.static_encoders[field] = nn.Sequential(*layers) + below = hidden[field] + self.field_dims[field] = hidden[field] + self.fc = nn.Linear(sum(self.field_dims.values()), self.get_output_size()) + + if __name__ == "__main__": from pyhealth.datasets import create_sample_dataset diff --git a/tests/core/test_unified_mlp.py b/tests/core/test_unified_mlp.py new file mode 100644 index 000000000..cd8a9abfc --- /dev/null +++ b/tests/core/test_unified_mlp.py @@ -0,0 +1,587 @@ +import copy +import tempfile +import unittest +from collections.abc import Sequence, Sized +from pathlib import Path +from typing import Any, cast + +import torch +from torch import nn + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import LevelEmbedding, Unified_LogisticRegression, Unified_MLP +from pyhealth.models.sequence_embedding import _FieldEmbedding +from pyhealth.processors import LevelConfig, SequenceProcessorV2 +from pyhealth.trainer import Trainer + + +class TestUnifiedPooling(unittest.TestCase): + def setUp(self) -> None: + self.samples: list[dict[str, Any]] = [ + { + "patient_id": "p0", + "events": { + "time": [1.0, 3.0], + "values": [[10.0], [20.0]], + "children": [ + {"values": [[2.0], [6.0]]}, + {"values": []}, + ], + }, + "static": [[1.0, 2.0], [3.0, 4.0]], + "y": 0, + }, + { + "patient_id": "p1", + "events": {"time": [], "values": [], "children": []}, + "static": [[5.0, 6.0], [7.0, 8.0]], + "y": 1, + }, + { + "patient_id": "p2", + "events": { + "time": [2.0], + "values": [[30.0]], + "children": [{"values": [[8.0]]}], + }, + "static": [[9.0, 10.0], [11.0, 12.0]], + "y": 0, + }, + ] + self.schema: dict[str, Any] = { + "events": ( + "sequence_v2", + { + "levels": [ + {"ordering": "set", "values": [1]}, + {"ordering": "series", "values": [1]}, + ] + }, + ), + "static": "tensor", + } + self.dims: dict[str, list[LevelEmbedding]] = { + "events": [{"values": 1}, {"values": 1}] + } + self.dataset = create_sample_dataset(self.samples, self.schema, {"y": "binary"}) + self.batch = next(iter(get_dataloader(self.dataset, batch_size=3))) + + def test_hierarchical_mean_and_static_flattening(self) -> None: + model = Unified_LogisticRegression(self.dataset, self.dims) + with torch.no_grad(): + for level in range(2): + projection = cast( + nn.Linear, + cast( + _FieldEmbedding, model.embedding.embedding_layers["events"] + ).layer(level, "value_proj"), + ) + projection.weight.fill_(1) + cast(torch.Tensor, projection.bias).zero_() + torch.testing.assert_close( + model(**self.batch, embed=True)["embed"], + torch.tensor( + [ + [15.0, 2.0, 1.0, 2.0, 3.0, 4.0], + [0.0, 0.0, 5.0, 6.0, 7.0, 8.0], + [30.0, 8.0, 9.0, 10.0, 11.0, 12.0], + ] + ), + ) + + def test_mlp_transforms_every_level_and_statics(self) -> None: + model = Unified_MLP( + self.dataset, self.dims, hidden_dim=1, num_layers=1, dropout=0 + ) + with torch.no_grad(): + for module in model.modules(): + if isinstance(module, nn.Linear): + module.weight.fill_(1) + module.bias.fill_(2) + # Leaves: [6, 10] -> mean 8; empty children -> 0. + # Parents: [12 + 8 + 2, 22 + 0 + 2] -> mean 23. + # An empty sample stays zero despite the MLP biases. + torch.testing.assert_close( + model(**self.batch, embed=True)["embed"], + torch.tensor([[23.0, 12.0], [0.0, 28.0], [46.0, 44.0]]), + ) + + def test_static_only_uses_nonlinearity(self) -> None: + dataset = create_sample_dataset( + [ + {"patient_id": "p0", "x": [-2.0, 3.0], "y": 0}, + {"patient_id": "p1", "x": [2.0, -3.0], "y": 1}, + ], + {"x": "tensor"}, + {"y": "binary"}, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + model = Unified_MLP(dataset, {}, hidden_dim=2, num_layers=2, dropout=0) + with torch.no_grad(): + for module in cast(nn.Sequential, model.static_encoders["x"]): + if isinstance(module, nn.Linear): + module.weight.copy_(torch.eye(2)) + module.bias.zero_() + torch.testing.assert_close( + model(**batch, embed=True)["embed"], torch.tensor([[0.0, 3.0], [2.0, 0.0]]) + ) + torch.testing.assert_close( + Unified_LogisticRegression(dataset, {})(**batch, embed=True)["embed"], + batch["x"].float(), + ) + + def test_order_invariance(self) -> None: + samples = copy.deepcopy(self.samples) + samples[0]["events"]["children"][0]["values"].reverse() + for role in ("time", "values", "children"): + samples[0]["events"][role].reverse() + # An untimed sequence permits testing permutation independently of sorting. + schema = copy.deepcopy(self.schema) + schema["events"][1]["levels"][1]["ordering"] = "sequence" + for sample in samples: + del sample["events"]["time"] + dataset = create_sample_dataset(samples, schema, {"y": "binary"}) + original = copy.deepcopy(self.samples) + for sample in original: + del sample["events"]["time"] + original = create_sample_dataset(original, schema, {"y": "binary"}) + for cls in (Unified_MLP, Unified_LogisticRegression): + model = cls(original, self.dims).eval() + torch.testing.assert_close( + model(**next(iter(get_dataloader(original, batch_size=3))))["logit"], + model(**next(iter(get_dataloader(dataset, batch_size=3))))["logit"], + ) + + def test_per_field_settings_and_gradients(self) -> None: + model = Unified_MLP( + self.dataset, + self.dims, + hidden_dim={"events": 3, "static": 5}, + num_layers={"events": 1}, + dropout={"static": 0}, + ) + self.assertEqual(model.fc.in_features, 8) + self.assertEqual(len(cast(nn.Sequential, model.static_encoders["static"])), 4) + model(**self.batch)["loss"].backward() + self.assertIsNotNone( + cast(nn.Sequential, model.static_encoders["static"])[0].weight.grad + ) + self.assertIsNotNone( + cast(nn.Sequential, cast(nn.ModuleList, model.encoders["events"])[0])[ + 0 + ].weight.grad + ) + + def test_constructor_validation(self) -> None: + for kwargs in cast( + list[dict[str, Any]], + [ + {"hidden_dim": 0}, + {"num_layers": 0}, + {"dropout": 1}, + {"hidden_dim": {"missing": 2}}, + {"num_layers": True}, + ], + ): + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + Unified_MLP(self.dataset, self.dims, **kwargs) + for cls in (Unified_MLP, Unified_LogisticRegression): + with self.assertRaises(ValueError): + cls(self.dataset, self.dims, mask=cast(Any, "invalid")) + with self.assertRaises(ValueError): + cls(self.dataset, {}) + dataset = create_sample_dataset( + [{"patient_id": str(i), "x": ["a"], "y": i} for i in range(2)], + {"x": "sequence"}, + {"y": "binary"}, + ) + with self.assertRaises(TypeError): + cls(dataset, {}) + two_labels = create_sample_dataset( + [{**sample, "z": 1 - sample["y"]} for sample in self.samples], + self.schema, + {"y": "binary", "z": "binary"}, + ) + with self.assertRaisesRegex(ValueError, "exactly one label"): + cls(two_labels, self.dims) + no_inputs = create_sample_dataset( + [{"patient_id": str(i), "y": i} for i in range(2)], {}, {"y": "binary"} + ) + with self.assertRaisesRegex(ValueError, "At least one input field"): + cls(no_inputs, {}) + times_only = create_sample_dataset( + [ + {"patient_id": str(i), "t": {"time": [0.0, 1.0]}, "y": i} + for i in range(2) + ], + {"t": ("sequence_v2", {"levels": [{"ordering": "series"}]})}, + {"y": "binary"}, + ) + with self.assertRaisesRegex(ValueError, "nothing to encode"): + cls(times_only, {"t": [{}]}) + + def test_per_step_means_boundaries_and_empty_samples(self) -> None: + samples = copy.deepcopy(self.samples) + for sample in samples: + sample["y"] = {"time": [0.0, 1.0, 2.0, 3.0, 5.0], "values": [0, 1, 0, 1, 0]} + dataset = create_sample_dataset( + samples, self.schema, {"y": "timeseries_binary"} + ) + batch = next(iter(get_dataloader(dataset, batch_size=3))) + for cls in (Unified_MLP, Unified_LogisticRegression): + model = cls(dataset, self.dims).eval() + embed = model(**batch, embed=True)["embed"] + # Predict at each individual parent time using single-record samples + # to construct a reference independent of prefix implementation. + reference = copy.deepcopy(self.samples) + reference[0]["events"]["time"] = [1.0] + reference[0]["events"]["values"] = [[10.0]] + reference[0]["events"]["children"] = [{"values": [[2.0], [6.0]]}] + reference[1]["events"] = { + "time": [3.0], + "values": [[20.0]], + "children": [{"values": []}], + } + reference = create_sample_dataset(reference, self.schema, {"y": "binary"}) + reference_model = cls(reference, self.dims).eval() + reference_model.load_state_dict(model.state_dict()) + records = reference_model( + **next(iter(get_dataloader(reference, batch_size=3))), embed=True + )["embed"] + width = model.field_dims["events"] + torch.testing.assert_close(embed[0, :width], torch.zeros(width)) + torch.testing.assert_close( + embed[1:3, :width], records[0, :width].expand(2, -1) + ) + torch.testing.assert_close( + embed[3:5, :width], records[:2, :width].mean(0).expand(2, -1) + ) + torch.testing.assert_close(embed[5:10, :width], torch.zeros(5, width)) + torch.testing.assert_close(embed[10:12, :width], torch.zeros(2, width)) + torch.testing.assert_close( + embed[12:15, :width], records[2, :width].expand(3, -1) + ) + torch.testing.assert_close( + embed[:5, width:], records[0, width:].expand(5, -1) + ) + changed = copy.deepcopy(batch) + # Changing a future parent cannot affect the first three predictions. + changed["events"] = tuple(t.clone() for t in batch["events"]) + schema = cast( + SequenceProcessorV2, dataset.input_processors["events"] + ).batch_schema() + changed["events"][schema.index((1, "values"))][1] += 1000 + torch.testing.assert_close( + model(**changed)["logit"][:3], model(**batch)["logit"][:3] + ) + + def test_resampled_and_all_empty_prefixes(self) -> None: + for empty in (False, True): + samples = [ + { + "patient_id": str(i), + "x": { + "time": [] if empty else [0.0, 1.0, 2.0], + "values": [] if empty else [[2.0], [4.0], [9.0]], + }, + "y": {"time": [0.0, 1.0, 2.0, 4.0], "values": [1.0, 2.0, 3.0, 4.0]}, + } + for i in range(2) + ] + levels: list[LevelConfig] = [{"ordering": "series", "values": [1]}] + if not empty: + levels[0]["resample"] = {"step": 1.0, "method": "last"} + dataset = create_sample_dataset( + samples, + {"x": ("sequence_v2", {"levels": levels})}, + {"y": "timeseries_regression"}, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + for cls in (Unified_MLP, Unified_LogisticRegression): + model = cls(dataset, {"x": [{"values": 1}]}).eval() + out = model(**batch, embed=True) + if empty: + torch.testing.assert_close( + out["embed"], torch.zeros_like(out["embed"]) + ) + else: + x = cast(nn.ModuleList, model.encoders["x"])[0]( + model.embedding({"x": batch["x"]})["x"][0]["embedding"] + ) + torch.testing.assert_close(out["embed"][0], x[0]) + torch.testing.assert_close(out["embed"][1], x[:2].mean(0)) + torch.testing.assert_close( + out["embed"][2:4], x[:3].mean(0).expand(2, -1) + ) + self.assertTrue(torch.isfinite(out["loss"])) + out["loss"].backward() + + def test_three_levels_mixed_namespaces_and_spatial_values(self) -> None: + dataset = create_sample_dataset( + [ + { + "patient_id": str(i), + "x": { + "tokens": ["stay"], + "children": [ + { + "children": [ + { + "tokens": [("a", "unit"), ("b", "unit")], + "values": [ + [[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0]], + ], + }, + {"tokens": [], "values": []}, + ] + } + ], + }, + "y": i, + } + for i in range(2) + ], + { + "x": ( + "sequence_v2", + { + "levels": [ + { + "ordering": "set", + "tokens": 2, + "values": [2, 2], + "spatial": [False, True], + }, + {"ordering": "sequence"}, + {"ordering": "set", "tokens": 1}, + ] + }, + ) + }, + {"y": "binary"}, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + dims: dict[str, list[LevelEmbedding]] = { + "x": [{"tokens": (2, 3), "values": 4}, {}, {"tokens": (2,)}] + } + for cls in (Unified_MLP, Unified_LogisticRegression): + model = cls(dataset, dims).eval() + embedded = model.embedding({"x": batch["x"]})["x"] + self.assertEqual(embedded[0]["embedding"].shape, (4, 9)) + leaves = cast(nn.ModuleList, model.encoders["x"])[0]( + embedded[0]["embedding"] + ) + middle = cast(nn.ModuleList, model.encoders["x"])[1]( + torch.stack( + [ + leaves[:2].mean(0), + torch.zeros(leaves.shape[1]), + leaves[2:].mean(0), + torch.zeros(leaves.shape[1]), + ] + ) + ) + expected = cast(nn.ModuleList, model.encoders["x"])[2]( + torch.cat( + [embedded[2]["embedding"], middle.reshape(2, 2, -1).mean(1)], dim=1 + ) + ) + out = model(**batch, embed=True) + torch.testing.assert_close(out["embed"], expected) + out["loss"].backward() + self.assertIsNotNone( + cast(_FieldEmbedding, model.embedding.embedding_layers["x"]) + .layer(0, "token_1") + .weight.grad + ) + + def test_sample_label_modes(self) -> None: + for mode, labels, width in ( + ("binary", [0, 1, 0], 1), + ("multiclass", ["a", "b", "c"], 3), + ("multilabel", [["a"], ["b"], ["a", "b"]], 2), + ("regression", [1.0, 2.0, 3.0], 1), + ): + samples = copy.deepcopy(self.samples) + for sample, label in zip(samples, labels): + sample["y"] = label + dataset = create_sample_dataset(samples, self.schema, {"y": mode}) + batch = next(iter(get_dataloader(dataset, batch_size=3))) + for cls in (Unified_MLP, Unified_LogisticRegression): + with self.subTest(mode=mode, model=cls.__name__): + out = cls(dataset, self.dims)(**batch) + self.assertEqual(out["logit"].shape, (3, width)) + self.assertTrue(torch.isfinite(out["loss"])) + out["loss"].backward() + + def test_per_step_static_only_and_trainer(self) -> None: + for mode, metric in ( + ("timeseries_binary", "roc_auc"), + ("timeseries_regression", "mae"), + ): + samples = copy.deepcopy(self.samples) + for sample in samples: + sample["y"] = {"time": [0.0, 1.0, 3.0], "values": [0.0, 1.0, 0.0]} + dataset = create_sample_dataset(samples, self.schema, {"y": mode}) + loader = get_dataloader(dataset, batch_size=2) + for cls in (Unified_MLP, Unified_LogisticRegression): + trainer = Trainer( + model=cls(dataset, self.dims), + metrics=[metric], + device="cpu", + enable_logging=False, + ) + trainer.train(train_dataloader=loader, epochs=1) + self.assertEqual( + [ + len(y) + for y in cast( + Sequence[Sized], + cast(list[object], trainer.inference(loader))[0], + ) + ], + [3, 3, 3], + ) + self.assertIn(metric, trainer.evaluate(loader)) + dataset = create_sample_dataset(samples, {"static": "tensor"}, {"y": mode}) + batch = next(iter(get_dataloader(dataset, batch_size=3))) + for cls in (Unified_MLP, Unified_LogisticRegression): + model = cls(dataset, {}).eval() + out = model(**batch, embed=True) + torch.testing.assert_close( + out["embed"], + model.static_encoders["static"]( + batch["static"].flatten(1) + ).repeat_interleave(3, dim=0), + ) + + def test_per_step_requires_timed_top_level(self) -> None: + samples = copy.deepcopy(self.samples) + for sample in samples: + del sample["events"]["time"] + sample["y"] = {"time": [0.0], "values": [1.0]} + for ordering in ("set", "sequence"): + schema = copy.deepcopy(self.schema) + schema["events"][1]["levels"][1]["ordering"] = ordering + dataset = create_sample_dataset(samples, schema, {"y": "timeseries_binary"}) + for cls in (Unified_MLP, Unified_LogisticRegression): + with self.assertRaisesRegex(ValueError, "series"): + cls(dataset, self.dims) + + def test_missing_indicators_and_pretrained_tokens(self) -> None: + dataset = create_sample_dataset( + [ + { + "patient_id": str(i), + "x": {"tokens": ["a", "b"], "values": [[float("nan")], [float(i)]]}, + "y": i, + } + for i in range(2) + ], + { + "x": ( + "sequence_v2", + { + "levels": [ + { + "ordering": "set", + "tokens": 1, + "values": [1], + "impute": {"method": "zero"}, + "missing_indicator": True, + } + ] + }, + ) + }, + {"y": "binary"}, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + with tempfile.TemporaryDirectory() as directory: + Path(directory, "tokens.txt").write_text("a 1 2\nb 3 4\n", encoding="utf-8") + for cls in (Unified_MLP, Unified_LogisticRegression): + model = cls( + dataset, + {"x": [{"tokens": (2,), "values": 1}]}, + mask="concat", + pretrain=[ + { + "field": "x", + "level": 0, + "namespace": 0, + "path": str(Path(directory, "tokens.txt")), + "freeze": True, + } + ], + ) + self.assertEqual( + cast(_FieldEmbedding, model.embedding.embedding_layers["x"]) + .layer(0, "value_proj") + .in_features, + 2, + ) + torch.testing.assert_close( + model.embedding({"x": batch["x"]})["x"][0]["embedding"][:, :2], + torch.tensor([[1.0, 2.0], [3.0, 4.0], [1.0, 2.0], [3.0, 4.0]]), + ) + self.assertFalse( + cast(_FieldEmbedding, model.embedding.embedding_layers["x"]) + .layer(0, "token_0") + .weight.requires_grad + ) + out = model(**batch) + self.assertTrue(torch.isfinite(out["loss"])) + out["loss"].backward() + + def test_multiple_fields_sparse_labels_and_device(self) -> None: + samples = copy.deepcopy(self.samples) + for i, sample in enumerate(samples): + sample["other"] = {"time": [0.0, 2.0], "values": [[1.0], [5.0]]} + sample["y"] = { + "time": [0.0, 2.0, 4.0][: i + 1], + "values": [0.0, 1.0, 0.0][: i + 1], + } + schema = copy.deepcopy(self.schema) + schema["other"] = ( + "sequence_v2", + {"levels": [{"ordering": "series", "values": [1]}]}, + ) + dataset = create_sample_dataset(samples, schema, {"y": "timeseries_binary"}) + batch = next(iter(get_dataloader(dataset, batch_size=3))) + dims: dict[str, list[LevelEmbedding]] = {**self.dims, "other": [{"values": 1}]} + for cls in (Unified_MLP, Unified_LogisticRegression): + model = cls(dataset, dims).eval() + out = model(**batch, embed=True) + torch.testing.assert_close( + out["loss"], + nn.functional.binary_cross_entropy_with_logits( + out["logit"], batch["y"][0] + ), + ) + self.assertEqual(out["logit"].shape, (6, 1)) + width = model.field_dims["events"] + x = cast(nn.ModuleList, model.encoders["other"])[0]( + model.embedding({"other": batch["other"]})["other"][0]["embedding"] + ) + torch.testing.assert_close( + out["embed"][0, width : width + model.field_dims["other"]], x[0] + ) + torch.testing.assert_close( + out["embed"][-1, width : width + model.field_dims["other"]], + x[-2:].mean(0), + ) + if torch.cuda.is_available(): + model.cuda() + cuda_out = model(**batch, embed=True) + torch.testing.assert_close( + cuda_out["embed"].cpu(), out["embed"], rtol=1e-4, atol=1e-5 + ) + torch.testing.assert_close( + cuda_out["logit"].cpu(), out["logit"], rtol=1e-4, atol=1e-5 + ) + cuda_out["loss"].backward() + + +if __name__ == "__main__": + unittest.main()