Towards Unified Sequence Representation I: Add SequenceProcessorV2 for packed hierarchical records

docs/api/processors.rst

@@ -543,6 +543,7 @@ API Reference
     processors/pyhealth.processors.SampleProcessor
     processors/pyhealth.processors.DatasetProcessor
     processors/pyhealth.processors.SequenceProcessor
+    processors/pyhealth.processors.SequenceProcessorV2
     processors/pyhealth.processors.NestedSequenceProcessor
     processors/pyhealth.processors.NestedMultiHotProcessor
     processors/pyhealth.processors.NestedFloatsProcessor

docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst

@@ -0,0 +1,63 @@
+pyhealth.processors.SequenceProcessorV2
+=======================================
+
+Processor for hierarchical packed records: codes, values and timestamps at any
+depth of nesting, stored without padding or truncation.
+
+The task returns each field as a ``Records``, a plain dict of columns.
+``children`` holds the records one level down, and the innermost records
+(level 0) have no ``children``. The processor emits a tuple of tensors in
+``schema()`` order: by level, then by role (``tokens``, ``values``, ``mask``,
+``time``, ``offsets``). ``offsets`` at level ``l`` partitions the level-``l`` records into
+level-``l + 1`` records, in CSR layout.
+
+Each level's configuration declares its roles: ``tokens`` gives the number of
+token namespaces ``K`` and ``values`` gives the value shape ``S``. ``fit()``
+only builds the vocabularies, and every record is checked against the
+configuration.
+
+.. code-block:: python
+
+    proc = SequenceProcessorV2(levels=[{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}])
+    samples = [{"visits": {"children": [{"tokens": ["A", "B"]}, {"tokens": ["C"]}]}}]
+    proc.fit(samples, "visits")
+    proc.schema()  # ((0, 'tokens'), (0, 'offsets'))
+    proc.process({"children": [{"tokens": ["A", "B"]}, {"tokens": []}, {"tokens": ["C"]}]})
+    # (tensor([2, 3, 4]), tensor([0, 2, 2, 3]))
+
+In a task schema the configuration is plain data:
+
+.. code-block:: python
+
+    input_schema = {
+        "visits": ("sequence_v2", {"levels": [{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}]}),
+        "vitals": ("sequence_v2", {"levels": [{"ordering": "series", "values": [2], "resample": {"step": 1.0, "method": "last"}}]}),
+    }
+
+Batches are packed with ``collate()``, which follows ``batch_schema()``: the
+fields of ``schema()`` plus the batch offsets ``[B + 1]``. After resampling,
+``impute`` fills NaN values (see ``Impute``) and ``normalize`` scales them (see
+``Normalize``). With ``missing_indicator`` (which needs ``impute``), a level
+also emits ``mask`` right after ``values``: a bool tensor of the same shape,
+True where a value was missing before imputation. Values still containing NaN
+or inf raise.
+
+.. autoclass:: pyhealth.processors.SequenceProcessorV2
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.processors.Records
+    :members:
+
+.. autoclass:: pyhealth.processors.LevelConfig
+    :members:
+
+.. autoclass:: pyhealth.processors.Resample
+    :members:
+
+.. autoclass:: pyhealth.processors.Impute
+    :members:
+
+.. autoclass:: pyhealth.processors.Normalize
+    :members:

pyhealth/datasets/sample_dataset.py

@@ -353,7 +353,8 @@ def encode(self, transformed: dict[str, Any]) -> dict[str, Any]:
         litdata infers one serializer per value from the first sample it
         writes and expects every later sample to have the same structure, so a
         list whose length varies between samples cannot be written as-is.
-        Fields whose processor declares ``stores_tensor`` stay tensors; every
+        Fields whose processor declares ``stores_tensor`` stay tensors (or
+        tuples of tensors); every
         other field (including keys not in the schema, such as ``patient_id``)
         goes into one pickled payload, which has the same structure in every
         sample whatever the values look like. :class:`SampleDataset` decodes it
@@ -366,20 +367,24 @@ def encode(self, transformed: dict[str, Any]) -> dict[str, Any]:
             A record with the tensor fields plus one bytes payload.
 
         Raises:
-            ValueError: If a field declared ``stores_tensor`` is not a tensor.
+            ValueError: If a field declared ``stores_tensor`` is not a tensor
+                or a tuple of tensors.
         """
         tensor_fields = self._tensor_fields()
         record: dict[str, Any] = {}
         objects: dict[str, Any] = {}
         for key, value in transformed.items():
             if key in tensor_fields:
-                if not isinstance(value, torch.Tensor):
+                is_tensors = isinstance(value, tuple) and all(
+                    isinstance(item, torch.Tensor) for item in value
+                )
+                if not (isinstance(value, torch.Tensor) or is_tensors):
                     processor = self._input_processors.get(key) or self._output_processors.get(key)
                     raise ValueError(
                         f"Field '{key}': {type(processor).__name__} declares "
                         f"stores_tensor=True but process() returned "
                         f"{type(value).__name__}. Set stores_tensor = False on the "
-                        "processor, or return a torch.Tensor."
+                        "processor, or return a torch.Tensor or a tuple of them."
                     )
                 record[key] = value
             else:

pyhealth/datasets/utils.py

@@ -1,11 +1,13 @@
+import functools
 import hashlib
 import os
 import pickle
+from collections.abc import Callable, Mapping
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Tuple
 
-import torch
 import litdata
+import torch
 from dateutil.parser import parse as dateutil_parse
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import DataLoader
@@ -250,25 +252,42 @@ def collate_fn_dict(batch: List[dict]) -> dict:
     return {key: [d[key] for d in batch] for key in batch[0]}
 
 
-def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
+def collate_fn_dict_with_padding(
+    batch: list[dict], collators: Mapping[str, Callable] | None = None
+) -> dict:
     """Collates a batch of data into a dictionary with padding for tensor values.
 
     Args:
         batch: List of dictionaries, where each dictionary represents a data sample.
+        collators: Optional collate function per key, given the key's values
+            in the batch. Keys with one are collated only by it.
 
     Returns:
         A dictionary where each key corresponds to a list of values from the batch.
         Tensor values are padded to the same shape.
         Tuples of (time, values) from temporal processors are collated separately.
+
+    Examples:
+        >>> import torch
+        >>> batch = [{"x": torch.tensor([1, 2]), "id": "a"}, {"x": torch.tensor([3]), "id": "b"}]
+        >>> collate_fn_dict_with_padding(batch)["x"]
+        tensor([[1, 2],
+                [3, 0]])
+        >>> collate_fn_dict_with_padding(batch)["id"]
+        ['a', 'b']
+        >>> collate_fn_dict_with_padding(batch, collators={"x": len})["x"]
+        2
     """
     collated = {}
     keys = batch[0].keys()
 
     for key in keys:
         values = [sample[key] for sample in batch]
 
+        if collators is not None and key in collators:
+            collated[key] = collators[key](values)
         # Check if this is a temporal feature tuple (time, values)
-        if isinstance(values[0], tuple) and len(values[0]) == 2:
+        elif isinstance(values[0], tuple) and len(values[0]) == 2:
             # Handle (time, values) tuples from processors
             time_tensors = [v[0] for v in values]
             value_tensors = [v[1] for v in values]
@@ -403,13 +422,25 @@ def get_dataloader(
     from .sample_dataset import SampleDataset
 
     if isinstance(dataset, SampleDataset):
+        # Fields whose processor defines collate() are batched by it.
+        processors = {**dataset.input_processors, **dataset.output_processors}
+        collators = {
+            key: processor.collate
+            for key, processor in processors.items()
+            if callable(getattr(processor, "collate", None))
+        }
+        collate_fn = (
+            functools.partial(collate_fn_dict_with_padding, collators=collators)
+            if collators
+            else collate_fn_dict_with_padding
+        )
         return _SeededDataLoader(
             dataset,
             shuffle_samples=shuffle,
             seed=seed,
             batch_size=batch_size,
             generator=generator,
-            collate_fn=collate_fn_dict_with_padding,
+            collate_fn=collate_fn,
         )
     if hasattr(dataset, "set_shuffle"):  # another iterable dataset, e.g. litdata's
         dataset.set_shuffle(shuffle)

pyhealth/medcode/cross_map.py

@@ -187,6 +187,19 @@ def unmapped_codes(self):
         """Source codes seen so far that had no target in this mapping."""
         return getattr(self.mapping, "unmapped_codes", frozenset())
 
+    @property
+    def has_many_targets(self):
+        """Whether some source code maps to more than one target code.
+
+        False means the mapping is single-valued (x-to-one): `map()` returns
+        at most one code. This describes the table, so it is True even if only
+        a few codes have several targets.
+        """
+        if isinstance(self.mapping, _ICDMappingsBackend):
+            # icd-mappings returns a single (primary) target for every code.
+            return False
+        return any(len(targets) > 1 for targets in self.mapping.values())
+
     def __repr__(self):
         return f"CrossMap(source_vocabulary={self.s_vocab}, source_class={self.s_class} target_vocabulary={self.t_vocab}, target_class={self.t_class})"
 

pyhealth/processors/__init__.py

@@ -37,6 +37,14 @@ def get_processor(name: str):
 )
 from .raw_processor import RawProcessor
 from .sequence_processor import SequenceProcessor
+from .sequence_processor_v2 import (
+    Impute,
+    LevelConfig,
+    Normalize,
+    Records,
+    Resample,
+    SequenceProcessorV2,
+)
 from .signal_processor import SignalProcessor
 from .stagenet_processor import (
     StageNetProcessor,
@@ -76,6 +84,12 @@ def get_processor(name: str):
     "NestedSequenceProcessor",
     "RawProcessor",
     "SequenceProcessor",
+    "SequenceProcessorV2",
+    "Records",
+    "LevelConfig",
+    "Impute",
+    "Normalize",
+    "Resample",
     "SignalProcessor",
     "StageNetProcessor",
     "StageNetTensorProcessor",

pyhealth/processors/base_processor.py

@@ -70,11 +70,16 @@ class FeatureProcessor(Processor):
     warns unless it is fitted on a training split
     (``set_task(task, split=PatientSplit(...))``).
 
-    Set ``stores_tensor = True`` only if ``process()`` always returns a single
-    ``torch.Tensor``. Such fields are written to the disk cache as tensors;
-    every other output (lists, tuples, dicts, strings, None, objects) is
-    stored as one pickled value per sample, so it may vary freely in length
-    and shape between samples.
+    Set ``stores_tensor = True`` only if ``process()`` always returns a
+    ``torch.Tensor``, or a tuple of tensors with the same length for every
+    sample (each tensor may change shape). Such fields are written to the disk
+    cache as tensors; every other output (lists, other tuples, dicts, strings,
+    None, objects) is stored as one pickled value per sample, so it may vary
+    freely in length and shape between samples.
+
+    A processor whose batches need more than padding can define
+    ``collate(items)``, taking the ``process()`` outputs of one batch;
+    :func:`~pyhealth.datasets.get_dataloader` then uses it for that field.
 
     Examples:
         >>> import torch