Towards Unified Sequence Representation VI: Add MLP_V2 and LogisticRegression_V2 for packed records

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.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.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``
      - 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.MLP_V2
+    models/pyhealth.models.LogisticRegression_V2
     models/pyhealth.models.CNN
     models/pyhealth.models.RNN
     models/pyhealth.models.GNN

docs/api/models/pyhealth.models.LogisticRegression_V2.rst

@@ -0,0 +1,17 @@
+pyhealth.models.LogisticRegression_V2
+=====================================
+
+A linear prediction head over bottom-up means of learned record embeddings
+and flattened static tensors. Unlike :class:`~pyhealth.models.MLP_V2`,
+it introduces no MLPs at hierarchy levels or on static fields.
+
+.. autoclass:: pyhealth.models.LogisticRegression_V2
+    :members:
+    :show-inheritance:
+
+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
+for legacy schemas. Both versions use learned embeddings rather than
+classical logistic regression directly on raw features.

docs/api/models/pyhealth.models.MLP_V2.rst

@@ -0,0 +1,52 @@
+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.MLP_V2
+    :members:
+    :show-inheritance:
+
+Example
+-------
+
+.. code-block:: python
+
+    from pyhealth.datasets import create_sample_dataset, get_dataloader
+    from pyhealth.models import MLP_V2, LogisticRegression_V2
+
+    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 = MLP_V2(dataset, dims, hidden_dim=16)
+    result = model(**batch)
+    linear_baseline = LogisticRegression_V2(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.

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.MLP_V2` and
+:class:`~pyhealth.models.LogisticRegression_V2` build on it.
 
 .. autoclass:: pyhealth.models.SequenceEmbeddingModel
     :members:

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 LogisticRegression_V2 as LogisticRegression_V2
 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 MLP_V2 as MLP_V2
 from .molerec import MoleRec, MoleRecLayer
 from .retain import MultimodalRETAIN, RETAIN, RETAINLayer
 from .rnn import MultimodalRNN, RNN, RNNLayer

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 LogisticRegression_V2(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 = LogisticRegression_V2(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(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]
+        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