From 0684d8e653ee398c8b2bc1ffc7729b5ef8ae89c3 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Fri, 30 Oct 2026 11:23:13 +0000 Subject: [PATCH 1/2] Add Unified_MLP and Unified_LogisticRegression for packed records Two pooling baselines over sequence_v2 fields and static tensors, built on SequenceEmbeddingModel. - Unified_LogisticRegression embeds each level, mean-pools the records of every group bottom-up (an empty group gives zeros), concatenates the field summaries with the raw static vectors, and applies one linear layer. - Unified_MLP adds an MLP at every level and on each static field before pooling; hidden_dim, num_layers and dropout are set per field. - Per-step labels: each label record uses the mean of the top-level records at or before its time (a prefix mean), so it never sees later records. Fields need a series top level. --- docs/api/models.rst | 8 + ...pyhealth.models.SequenceEmbeddingModel.rst | 3 +- ...alth.models.Unified_LogisticRegression.rst | 17 + .../models/pyhealth.models.Unified_MLP.rst | 52 ++ pyhealth/models/__init__.py | 2 + pyhealth/models/logistic_regression.py | 192 +++++- pyhealth/models/mlp.py | 103 ++- tests/core/test_unified_mlp.py | 587 ++++++++++++++++++ 8 files changed, 960 insertions(+), 4 deletions(-) create mode 100644 docs/api/models/pyhealth.models.Unified_LogisticRegression.rst create mode 100644 docs/api/models/pyhealth.models.Unified_MLP.rst create mode 100644 tests/core/test_unified_mlp.py diff --git a/docs/api/models.rst b/docs/api/models.rst index 9f5439d40..daa901ddc 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 6a19499df..60b09eec3 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.RNN_V2` builds on it. +:class:`~pyhealth.models.RNN_V2`, :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 916e1a7be..b75bf2c78 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() From 2e08ba8b7abc2e4c579fb6d8385bcdd57b0fc8b4 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Sun, 1 Nov 2026 11:12:17 +0000 Subject: [PATCH 2/2] Rename Unified_MLP and Unified_LogisticRegression to MLP_V2 and LogisticRegression_V2 --- docs/api/models.rst | 8 ++--- ...pyhealth.models.LogisticRegression_V2.rst} | 10 +++--- ...ied_MLP.rst => pyhealth.models.MLP_V2.rst} | 12 +++---- ...pyhealth.models.SequenceEmbeddingModel.rst | 4 +-- pyhealth/models/__init__.py | 4 +-- pyhealth/models/logistic_regression.py | 6 ++-- pyhealth/models/mlp.py | 8 ++--- .../{test_unified_mlp.py => test_mlp_v2.py} | 36 +++++++++---------- 8 files changed, 44 insertions(+), 44 deletions(-) rename docs/api/models/{pyhealth.models.Unified_LogisticRegression.rst => pyhealth.models.LogisticRegression_V2.rst} (59%) rename docs/api/models/{pyhealth.models.Unified_MLP.rst => pyhealth.models.MLP_V2.rst} (86%) rename tests/core/{test_unified_mlp.py => test_mlp_v2.py} (95%) diff --git a/docs/api/models.rst b/docs/api/models.rst index daa901ddc..6c48e58c3 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -36,10 +36,10 @@ 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` + * - :doc:`models/pyhealth.models.MLP_V2` - 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` + * - :doc:`models/pyhealth.models.LogisticRegression_V2` - 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`` @@ -219,8 +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.MLP_V2 + models/pyhealth.models.LogisticRegression_V2 models/pyhealth.models.CNN models/pyhealth.models.RNN models/pyhealth.models.GNN diff --git a/docs/api/models/pyhealth.models.Unified_LogisticRegression.rst b/docs/api/models/pyhealth.models.LogisticRegression_V2.rst similarity index 59% rename from docs/api/models/pyhealth.models.Unified_LogisticRegression.rst rename to docs/api/models/pyhealth.models.LogisticRegression_V2.rst index 2ac0ac57b..b1f146ae6 100644 --- a/docs/api/models/pyhealth.models.Unified_LogisticRegression.rst +++ b/docs/api/models/pyhealth.models.LogisticRegression_V2.rst @@ -1,15 +1,15 @@ -pyhealth.models.Unified_LogisticRegression -========================================= +pyhealth.models.LogisticRegression_V2 +===================================== A linear prediction head over bottom-up means of learned record embeddings -and flattened static tensors. Unlike :class:`~pyhealth.models.Unified_MLP`, +and flattened static tensors. Unlike :class:`~pyhealth.models.MLP_V2`, it introduces no MLPs at hierarchy levels or on static fields. -.. autoclass:: pyhealth.models.Unified_LogisticRegression +.. autoclass:: pyhealth.models.LogisticRegression_V2 :members: :show-inheritance: -See :doc:`pyhealth.models.Unified_MLP` for a shared input example and the +See :doc:`pyhealth.models.MLP_V2` for a shared input example and the spatial, empty-group, and per-step prediction conventions. The existing :class:`~pyhealth.models.LogisticRegression` remains available diff --git a/docs/api/models/pyhealth.models.Unified_MLP.rst b/docs/api/models/pyhealth.models.MLP_V2.rst similarity index 86% rename from docs/api/models/pyhealth.models.Unified_MLP.rst rename to docs/api/models/pyhealth.models.MLP_V2.rst index 5c7aa4ef9..85d321b96 100644 --- a/docs/api/models/pyhealth.models.Unified_MLP.rst +++ b/docs/api/models/pyhealth.models.MLP_V2.rst @@ -1,11 +1,11 @@ -pyhealth.models.Unified_MLP -========================== +pyhealth.models.MLP_V2 +====================== 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 +.. autoclass:: pyhealth.models.MLP_V2 :members: :show-inheritance: @@ -15,7 +15,7 @@ Example .. code-block:: python from pyhealth.datasets import create_sample_dataset, get_dataloader - from pyhealth.models import Unified_MLP, Unified_LogisticRegression + from pyhealth.models import MLP_V2, LogisticRegression_V2 dataset = create_sample_dataset( [ @@ -33,9 +33,9 @@ Example ) dims = {"visits": [{"tokens": (8,)}, {}]} batch = next(iter(get_dataloader(dataset, batch_size=2))) - model = Unified_MLP(dataset, dims, hidden_dim=16) + model = MLP_V2(dataset, dims, hidden_dim=16) result = model(**batch) - linear_baseline = Unified_LogisticRegression(dataset, dims) + linear_baseline = LogisticRegression_V2(dataset, dims) linear_result = linear_baseline(**batch) For per-step labels, give each sequence field a ``series`` top level. diff --git a/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst b/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst index 60b09eec3..91ae16894 100644 --- a/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst +++ b/docs/api/models/pyhealth.models.SequenceEmbeddingModel.rst @@ -5,8 +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.RNN_V2`, :class:`~pyhealth.models.Unified_MLP` and -:class:`~pyhealth.models.Unified_LogisticRegression` build on it. +:class:`~pyhealth.models.RNN_V2`, :class:`~pyhealth.models.MLP_V2` and +:class:`~pyhealth.models.LogisticRegression_V2` build on it. .. autoclass:: pyhealth.models.SequenceEmbeddingModel :members: diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index b75bf2c78..979cdd246 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -16,7 +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 .logistic_regression import LogisticRegression_V2 as LogisticRegression_V2 from .gan import GAN from .gnn import GAT, GCN from .graph_torchvision_model import Graph_TorchvisionModel @@ -26,7 +26,7 @@ from .medlink import MedLink from .micron import MICRON, MICRONLayer from .mlp import MLP -from .mlp import Unified_MLP as Unified_MLP +from .mlp import MLP_V2 as MLP_V2 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 2c95409da..49c4d1891 100644 --- a/pyhealth/models/logistic_regression.py +++ b/pyhealth/models/logistic_regression.py @@ -220,7 +220,7 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]: return results -class Unified_LogisticRegression(BaseModel): +class LogisticRegression_V2(BaseModel): """Linear prediction over bottom-up means of learned record embeddings. Accepts ``sequence_v2`` fields and static ``tensor`` fields. At each level, @@ -258,7 +258,7 @@ class Unified_LogisticRegression(BaseModel): ... for i in range(2)], ... {"age": "tensor"}, {"y": "binary"}, ... ) - >>> model = Unified_LogisticRegression(dataset, embedding_dim={}) + >>> model = LogisticRegression_V2(dataset, embedding_dim={}) >>> model(**next(iter(get_dataloader(dataset, batch_size=2))))["logit"].shape torch.Size([2, 1]) """ @@ -272,7 +272,7 @@ def __init__( ) -> None: super().__init__(dataset) if len(self.label_keys) != 1: - raise ValueError("Unified pooling models support exactly one label.") + raise ValueError(f"{type(self).__name__} supports exactly one label.") if not self.feature_keys: raise ValueError("At least one input field is required.") self.label_key = self.label_keys[0] diff --git a/pyhealth/models/mlp.py b/pyhealth/models/mlp.py index 1debf0b29..a7f8d4093 100644 --- a/pyhealth/models/mlp.py +++ b/pyhealth/models/mlp.py @@ -9,7 +9,7 @@ from pyhealth.models import BaseModel from .embedding import EmbeddingModel -from .logistic_regression import Unified_LogisticRegression +from .logistic_regression import LogisticRegression_V2 from .rnn import _per_field from .sequence_embedding import LevelEmbedding, MaskMode, PretrainedEmbedding @@ -394,11 +394,11 @@ def get_embedding_model(self) -> nn.Module | None: return self.embedding_model -class Unified_MLP(Unified_LogisticRegression): +class MLP_V2(LogisticRegression_V2): """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 + described in :class:`LogisticRegression_V2`. 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. @@ -427,7 +427,7 @@ class Unified_MLP(Unified_LogisticRegression): ... for i in range(2)], ... {"age": "tensor"}, {"y": "binary"}, ... ) - >>> model = Unified_MLP(dataset, embedding_dim={}, hidden_dim=16) + >>> model = MLP_V2(dataset, embedding_dim={}, hidden_dim=16) >>> model(**next(iter(get_dataloader(dataset, batch_size=2))))["logit"].shape torch.Size([2, 1]) """ diff --git a/tests/core/test_unified_mlp.py b/tests/core/test_mlp_v2.py similarity index 95% rename from tests/core/test_unified_mlp.py rename to tests/core/test_mlp_v2.py index cd8a9abfc..faf74ea42 100644 --- a/tests/core/test_unified_mlp.py +++ b/tests/core/test_mlp_v2.py @@ -9,7 +9,7 @@ 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 import LevelEmbedding, LogisticRegression_V2, MLP_V2 from pyhealth.models.sequence_embedding import _FieldEmbedding from pyhealth.processors import LevelConfig, SequenceProcessorV2 from pyhealth.trainer import Trainer @@ -67,7 +67,7 @@ def setUp(self) -> None: 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) + model = LogisticRegression_V2(self.dataset, self.dims) with torch.no_grad(): for level in range(2): projection = cast( @@ -90,7 +90,7 @@ def test_hierarchical_mean_and_static_flattening(self) -> None: ) def test_mlp_transforms_every_level_and_statics(self) -> None: - model = Unified_MLP( + model = MLP_V2( self.dataset, self.dims, hidden_dim=1, num_layers=1, dropout=0 ) with torch.no_grad(): @@ -116,7 +116,7 @@ def test_static_only_uses_nonlinearity(self) -> None: {"y": "binary"}, ) batch = next(iter(get_dataloader(dataset, batch_size=2))) - model = Unified_MLP(dataset, {}, hidden_dim=2, num_layers=2, dropout=0) + model = MLP_V2(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): @@ -126,7 +126,7 @@ def test_static_only_uses_nonlinearity(self) -> None: 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"], + LogisticRegression_V2(dataset, {})(**batch, embed=True)["embed"], batch["x"].float(), ) @@ -145,7 +145,7 @@ def test_order_invariance(self) -> None: for sample in original: del sample["events"]["time"] original = create_sample_dataset(original, schema, {"y": "binary"}) - for cls in (Unified_MLP, Unified_LogisticRegression): + for cls in (MLP_V2, LogisticRegression_V2): model = cls(original, self.dims).eval() torch.testing.assert_close( model(**next(iter(get_dataloader(original, batch_size=3))))["logit"], @@ -153,7 +153,7 @@ def test_order_invariance(self) -> None: ) def test_per_field_settings_and_gradients(self) -> None: - model = Unified_MLP( + model = MLP_V2( self.dataset, self.dims, hidden_dim={"events": 3, "static": 5}, @@ -184,8 +184,8 @@ def test_constructor_validation(self) -> None: ], ): with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): - Unified_MLP(self.dataset, self.dims, **kwargs) - for cls in (Unified_MLP, Unified_LogisticRegression): + MLP_V2(self.dataset, self.dims, **kwargs) + for cls in (MLP_V2, LogisticRegression_V2): with self.assertRaises(ValueError): cls(self.dataset, self.dims, mask=cast(Any, "invalid")) with self.assertRaises(ValueError): @@ -228,7 +228,7 @@ def test_per_step_means_boundaries_and_empty_samples(self) -> None: samples, self.schema, {"y": "timeseries_binary"} ) batch = next(iter(get_dataloader(dataset, batch_size=3))) - for cls in (Unified_MLP, Unified_LogisticRegression): + for cls in (MLP_V2, LogisticRegression_V2): model = cls(dataset, self.dims).eval() embed = model(**batch, embed=True)["embed"] # Predict at each individual parent time using single-record samples @@ -297,7 +297,7 @@ def test_resampled_and_all_empty_prefixes(self) -> None: {"y": "timeseries_regression"}, ) batch = next(iter(get_dataloader(dataset, batch_size=2))) - for cls in (Unified_MLP, Unified_LogisticRegression): + for cls in (MLP_V2, LogisticRegression_V2): model = cls(dataset, {"x": [{"values": 1}]}).eval() out = model(**batch, embed=True) if empty: @@ -365,7 +365,7 @@ def test_three_levels_mixed_namespaces_and_spatial_values(self) -> None: dims: dict[str, list[LevelEmbedding]] = { "x": [{"tokens": (2, 3), "values": 4}, {}, {"tokens": (2,)}] } - for cls in (Unified_MLP, Unified_LogisticRegression): + for cls in (MLP_V2, LogisticRegression_V2): model = cls(dataset, dims).eval() embedded = model.embedding({"x": batch["x"]})["x"] self.assertEqual(embedded[0]["embedding"].shape, (4, 9)) @@ -408,7 +408,7 @@ def test_sample_label_modes(self) -> None: 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): + for cls in (MLP_V2, LogisticRegression_V2): with self.subTest(mode=mode, model=cls.__name__): out = cls(dataset, self.dims)(**batch) self.assertEqual(out["logit"].shape, (3, width)) @@ -425,7 +425,7 @@ def test_per_step_static_only_and_trainer(self) -> None: 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): + for cls in (MLP_V2, LogisticRegression_V2): trainer = Trainer( model=cls(dataset, self.dims), metrics=[metric], @@ -446,7 +446,7 @@ def test_per_step_static_only_and_trainer(self) -> None: 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): + for cls in (MLP_V2, LogisticRegression_V2): model = cls(dataset, {}).eval() out = model(**batch, embed=True) torch.testing.assert_close( @@ -465,7 +465,7 @@ def test_per_step_requires_timed_top_level(self) -> None: 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): + for cls in (MLP_V2, LogisticRegression_V2): with self.assertRaisesRegex(ValueError, "series"): cls(dataset, self.dims) @@ -500,7 +500,7 @@ def test_missing_indicators_and_pretrained_tokens(self) -> None: 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): + for cls in (MLP_V2, LogisticRegression_V2): model = cls( dataset, {"x": [{"tokens": (2,), "values": 1}]}, @@ -550,7 +550,7 @@ def test_multiple_fields_sparse_labels_and_device(self) -> None: 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): + for cls in (MLP_V2, LogisticRegression_V2): model = cls(dataset, dims).eval() out = model(**batch, embed=True) torch.testing.assert_close(